Merge pull request #15576 from [BEAM-12950] Not delete orphaned files to avoid missing events
[BEAM-12950] Not delete orphaned files to avoid missing events
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index cef59cc..2f12a3a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -36,7 +36,7 @@
- repo: https://github.com/pycqa/pylint
# this rev is a release tag in the repo above and corresponds with a pylint
# version. make sure this matches the version of pylint in tox.ini.
- rev: pylint-2.4.3
+ rev: v2.11.1
hooks:
- id: pylint
args: ["--rcfile=sdks/python/.pylintrc"]
diff --git a/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy b/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy
index a41760e..20f9030 100644
--- a/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy
@@ -39,7 +39,7 @@
// Gradle goals for this job.
steps {
- shell('echo "*** RUN NEXMARK IN BATCH MODE USING SPARK RUNNER ***"')
+ shell('echo "*** RUN NEXMARK IN BATCH MODE USING SPARK 2 RUNNER ***"')
gradle {
rootBuildScriptDir(commonJobProperties.checkoutDir)
tasks(':sdks:java:testing:nexmark:run')
@@ -57,7 +57,7 @@
'--monitorJobs=true'
].join(' '))
}
- shell('echo "*** RUN NEXMARK SQL IN BATCH MODE USING SPARK RUNNER ***"')
+ shell('echo "*** RUN NEXMARK SQL IN BATCH MODE USING SPARK 2 RUNNER ***"')
gradle {
rootBuildScriptDir(commonJobProperties.checkoutDir)
tasks(':sdks:java:testing:nexmark:run')
@@ -76,7 +76,44 @@
'--monitorJobs=true'
].join(' '))
}
- shell('echo "*** RUN NEXMARK IN BATCH MODE USING SPARK STRUCTURED STREAMING RUNNER ***"')
+ shell('echo "*** RUN NEXMARK IN BATCH MODE USING SPARK 3 RUNNER ***"')
+ gradle {
+ rootBuildScriptDir(commonJobProperties.checkoutDir)
+ tasks(':sdks:java:testing:nexmark:run')
+ commonJobProperties.setGradleSwitches(delegate)
+ switches('-Pnexmark.runner=":runners:spark:3"' +
+ ' -Pnexmark.args="' +
+ [
+ commonJobProperties.mapToArgString(nexmarkBigQueryArgs),
+ commonJobProperties.mapToArgString(nexmarkInfluxDBArgs),
+ '--runner=SparkRunner',
+ '--streaming=false',
+ '--suite=SMOKE',
+ '--streamTimeout=60' ,
+ '--manageResources=false',
+ '--monitorJobs=true'
+ ].join(' '))
+ }
+ shell('echo "*** RUN NEXMARK SQL IN BATCH MODE USING SPARK 3 RUNNER ***"')
+ gradle {
+ rootBuildScriptDir(commonJobProperties.checkoutDir)
+ tasks(':sdks:java:testing:nexmark:run')
+ commonJobProperties.setGradleSwitches(delegate)
+ switches('-Pnexmark.runner=":runners:spark:3"' +
+ ' -Pnexmark.args="' +
+ [
+ commonJobProperties.mapToArgString(nexmarkBigQueryArgs),
+ commonJobProperties.mapToArgString(nexmarkInfluxDBArgs),
+ '--runner=SparkRunner',
+ '--queryLanguage=sql',
+ '--streaming=false',
+ '--suite=SMOKE',
+ '--streamTimeout=60' ,
+ '--manageResources=false',
+ '--monitorJobs=true'
+ ].join(' '))
+ }
+ shell('echo "*** RUN NEXMARK IN BATCH MODE USING SPARK 2 STRUCTURED STREAMING RUNNER ***"')
gradle {
rootBuildScriptDir(commonJobProperties.checkoutDir)
tasks(':sdks:java:testing:nexmark:run')
@@ -96,7 +133,7 @@
'--monitorJobs=true'
].join(' '))
}
- shell('echo "*** RUN NEXMARK SQL IN BATCH MODE USING SPARK STRUCTURED STREAMING RUNNER ***"')
+ shell('echo "*** RUN NEXMARK SQL IN BATCH MODE USING SPARK 2 STRUCTURED STREAMING RUNNER ***"')
gradle {
rootBuildScriptDir(commonJobProperties.checkoutDir)
tasks(':sdks:java:testing:nexmark:run')
@@ -115,6 +152,45 @@
'--monitorJobs=true'
].join(' '))
}
+ shell('echo "*** RUN NEXMARK IN BATCH MODE USING SPARK 3 STRUCTURED STREAMING RUNNER ***"')
+ gradle {
+ rootBuildScriptDir(commonJobProperties.checkoutDir)
+ tasks(':sdks:java:testing:nexmark:run')
+ commonJobProperties.setGradleSwitches(delegate)
+ switches('-Pnexmark.runner=":runners:spark:3"' +
+ ' -Pnexmark.args="' +
+ [
+ commonJobProperties.mapToArgString(nexmarkBigQueryArgs),
+ commonJobProperties.mapToArgString(nexmarkInfluxDBArgs),
+ '--runner=SparkStructuredStreamingRunner',
+ '--streaming=false',
+ '--suite=SMOKE',
+ // Skip query 3 (SparkStructuredStreamingRunner does not support State/Timers yet)
+ '--skipQueries=3',
+ '--streamTimeout=60' ,
+ '--manageResources=false',
+ '--monitorJobs=true'
+ ].join(' '))
+ }
+ shell('echo "*** RUN NEXMARK SQL IN BATCH MODE USING SPARK 3 STRUCTURED STREAMING RUNNER ***"')
+ gradle {
+ rootBuildScriptDir(commonJobProperties.checkoutDir)
+ tasks(':sdks:java:testing:nexmark:run')
+ commonJobProperties.setGradleSwitches(delegate)
+ switches('-Pnexmark.runner=":runners:spark:3"' +
+ ' -Pnexmark.args="' +
+ [
+ commonJobProperties.mapToArgString(nexmarkBigQueryArgs),
+ commonJobProperties.mapToArgString(nexmarkInfluxDBArgs),
+ '--runner=SparkStructuredStreamingRunner',
+ '--queryLanguage=sql',
+ '--streaming=false',
+ '--suite=SMOKE',
+ '--streamTimeout=60' ,
+ '--manageResources=false',
+ '--monitorJobs=true'
+ ].join(' '))
+ }
}
}
diff --git a/.test-infra/metrics/grafana/provisioning/datasources/beamgithubjavatests-api.yaml b/.test-infra/metrics/grafana/provisioning/datasources/beamgithubjavatests-api.yaml
index a3a0e9b..34b4d3c 100644
--- a/.test-infra/metrics/grafana/provisioning/datasources/beamgithubjavatests-api.yaml
+++ b/.test-infra/metrics/grafana/provisioning/datasources/beamgithubjavatests-api.yaml
@@ -21,14 +21,13 @@
deleteDatasources:
datasources:
- - name: Java Tests
+ - name: "Java Tests"
type: marcusolsson-json-datasource
access: proxy
orgId: 1
url: https://api.github.com/repos/apache/beam/actions/workflows/java_tests.yml/runs
jsonData:
httpHeaderName1: "accept"
- customQueryParameters: "per_page=100"
secureJsonData:
httpHeaderValue1: "application/vnd.github.v3+json"
editable: false
diff --git a/.test-infra/metrics/grafana/provisioning/datasources/beamgithubpythontests-api.yaml b/.test-infra/metrics/grafana/provisioning/datasources/beamgithubpythontests-api.yaml
index abcd060..4feac1b 100644
--- a/.test-infra/metrics/grafana/provisioning/datasources/beamgithubpythontests-api.yaml
+++ b/.test-infra/metrics/grafana/provisioning/datasources/beamgithubpythontests-api.yaml
@@ -21,14 +21,13 @@
deleteDatasources:
datasources:
- - name: Python Tests
+ - name: "Python Tests"
type: marcusolsson-json-datasource
access: proxy
orgId: 1
url: https://api.github.com/repos/apache/beam/actions/workflows/python_tests.yml/runs
jsonData:
httpHeaderName1: "accept"
- customQueryParameters: "per_page=100"
secureJsonData:
httpHeaderValue1: "application/vnd.github.v3+json"
editable: false
diff --git a/.test-infra/metrics/kubernetes/beamgrafana-deploy.yaml b/.test-infra/metrics/kubernetes/beamgrafana-deploy.yaml
index 9ff402d..ad2de5f 100644
--- a/.test-infra/metrics/kubernetes/beamgrafana-deploy.yaml
+++ b/.test-infra/metrics/kubernetes/beamgrafana-deploy.yaml
@@ -46,6 +46,8 @@
value: "true"
- name: GF_AUTH_ANONYMOUS_ORG_NAME
value: Beam
+ - name: GF_INSTALL_PLUGINS
+ value: marcusolsson-json-datasource
- name: GF_SECURITY_ADMIN_PASSWORD
valueFrom:
secretKeyRef:
@@ -89,6 +91,7 @@
volumeMounts:
- mountPath: /var/lib/grafana
name: beam-grafana-libdata
+ readOnly: false
- mountPath: /etc/grafana
name: beam-grafana-etcdata
- mountPath: /var/log/grafana
diff --git a/CHANGES.md b/CHANGES.md
index a94d6cd..87c1c50 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -49,6 +49,36 @@
* ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
-->
+# [2.35.0] - Unreleased
+
+## Highlights
+
+* New highly anticipated feature X added to Python SDK ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+* New highly anticipated feature Y added to Java SDK ([BEAM-Y](https://issues.apache.org/jira/browse/BEAM-Y)).
+
+## I/Os
+
+* Support for X source added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+
+## New Features / Improvements
+
+* X feature added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+
+## Breaking Changes
+
+* X behavior was changed ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+
+## Deprecations
+
+* X behavior is deprecated and will be removed in X versions ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+
+## Bugfixes
+
+* Fixed X (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+
+## Known Issues
+
+* ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
# [2.34.0] - Unreleased
* Add an [example](https://github.com/cometta/python-apache-beam-spark) of deploying Python Apache Beam job with Spark Cluster
@@ -61,7 +91,9 @@
## I/Os
* Support for X source added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
-* `ReadFromBigQuery` now runs queries with BATCH priority by default. The `query_priority` parameter is introduced to the same transform to allow configuring the query priority (Python) ([BEAM-12913](https://issues.apache.org/jira/browse/BEAM-12913)).
+* `ReadFromBigQuery` and `ReadAllFromBigQuery` now run queries with BATCH priority by default. The `query_priority` parameter is introduced to the same transforms to allow configuring the query priority (Python) ([BEAM-12913](https://issues.apache.org/jira/browse/BEAM-12913)).
+* [EXPERIMENTAL] Support for [BigQuery Storage Read API](https://cloud.google.com/bigquery/docs/reference/storage) added to `ReadFromBigQuery`. The newly introduced `method` parameter can be set as `DIRECT_READ` to use the Storage Read API. The default is `EXPORT` which invokes a BigQuery export request. (Python) ([BEAM-10917](https://issues.apache.org/jira/browse/BEAM-10917)).
+* [EXPERIMENTAL] Added `use_native_datetime` parameter to `ReadFromBigQuery` to configure the return type of [DATETIME](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#datetime_type) fields when using `ReadFromBigQuery`. This parameter can *only* be used when `method = DIRECT_READ`(Python) ([BEAM-10917](https://issues.apache.org/jira/browse/BEAM-10917)).
## New Features / Improvements
@@ -72,6 +104,8 @@
* X behavior was changed ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
* SQL Rows are no longer flattened ([BEAM-5505](https://issues.apache.org/jira/browse/BEAM-5505)).
+* [Go SDK] beam.TryCrossLanguage's signature now matches beam.CrossLanguage. Like other Try functions it returns an error instead of panicking. ([BEAM-9918](https://issues.apache.org/jira/browse/BEAM-9918)).
+* [BEAM-12925](https://jira.apache.org/jira/browse/BEAM-12925) was fixed. It used to silently pass incorrect null data read from JdbcIO. Pipelines affected by this will now start throwing failures instead of silently passing incorrect data.
## Deprecations
@@ -104,9 +138,11 @@
* Easier path to contribute to the Go SDK, no need to set up a GO_PATH.
* Minimum Go version is now Go v1.16
* See the announcement blogpost for full information (TODO(lostluck): Add link once published.)
+* Python's ParDo (Map, FlatMap, etc.) transforms now suport a `with_exception_handling` option for easily ignoring bad records and implementing the dead letter pattern.
## I/Os
+* `DebeziumIO` now uses Debezium 1.7.0.Final ([BEAM-12993](https://issues.apache.org/jira/browse/BEAM-12993)).
* Support for X source added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
## New Features / Improvements
diff --git a/build.gradle.kts b/build.gradle.kts
index 72a7725..bee53bd 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -42,6 +42,7 @@
val exclusions = mutableListOf(
// Ignore files we track but do not distribute
"**/.github/**/*",
+ "**/.gitkeep",
"gradlew",
"gradlew.bat",
"gradle/wrapper/gradle-wrapper.properties",
@@ -117,7 +118,14 @@
"sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/yarn.lock",
// Sample text file for Java quickstart
- "sdks/java/maven-archetypes/examples/sample.txt"
+ "sdks/java/maven-archetypes/examples/sample.txt",
+
+ // Ignore Flutter autogenerated files for Playground
+ "playground/frontend/.metadata",
+ "playground/frontend/pubspec.lock",
+
+ // Ignore .gitkeep file
+ "**/.gitkeep"
)
// Add .gitignore excludes to the Apache Rat exclusion list. We re-create the behavior
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
index 293ea23..3f01717 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -378,7 +378,7 @@
// Automatically use the official release version if we are performing a release
// otherwise append '-SNAPSHOT'
- project.version = '2.34.0'
+ project.version = '2.35.0'
if (!isRelease(project)) {
project.version += '-SNAPSHOT'
}
@@ -475,7 +475,8 @@
def protobuf_version = "3.17.3"
def quickcheck_version = "0.8"
def slf4j_version = "1.7.30"
- def spark_version = "2.4.8"
+ def spark2_version = "2.4.8"
+ def spark3_version = "3.1.2"
def spotbugs_version = "4.0.6"
def testcontainers_version = "1.15.1"
def arrow_version = "5.0.0"
@@ -655,10 +656,14 @@
slf4j_jdk14 : "org.slf4j:slf4j-jdk14:$slf4j_version",
slf4j_log4j12 : "org.slf4j:slf4j-log4j12:$slf4j_version",
snappy_java : "org.xerial.snappy:snappy-java:1.1.8.4",
- spark_core : "org.apache.spark:spark-core_2.11:$spark_version",
- spark_network_common : "org.apache.spark:spark-network-common_2.11:$spark_version",
- spark_sql : "org.apache.spark:spark-sql_2.11:$spark_version",
- spark_streaming : "org.apache.spark:spark-streaming_2.11:$spark_version",
+ spark_core : "org.apache.spark:spark-core_2.11:$spark2_version",
+ spark_network_common : "org.apache.spark:spark-network-common_2.11:$spark2_version",
+ spark_sql : "org.apache.spark:spark-sql_2.11:$spark2_version",
+ spark_streaming : "org.apache.spark:spark-streaming_2.11:$spark2_version",
+ spark3_core : "org.apache.spark:spark-core_2.12:$spark3_version",
+ spark3_network_common : "org.apache.spark:spark-network-common_2.12:$spark3_version",
+ spark3_sql : "org.apache.spark:spark-sql_2.12:$spark3_version",
+ spark3_streaming : "org.apache.spark:spark-streaming_2.12:$spark3_version",
stax2_api : "org.codehaus.woodstox:stax2-api:4.2.1",
testcontainers_clickhouse : "org.testcontainers:clickhouse:$testcontainers_version",
testcontainers_elasticsearch : "org.testcontainers:elasticsearch:$testcontainers_version",
@@ -2337,7 +2342,7 @@
def distTarBall = "${pythonRootDir}/build/apache-beam.tar.gz"
project.exec {
executable 'sh'
- args '-c', ". ${project.ext.envdir}/bin/activate && pip install --retries 10 ${distTarBall}[gcp,test,aws,azure]"
+ args '-c', ". ${project.ext.envdir}/bin/activate && pip install --retries 10 ${distTarBall}[gcp,test,aws,azure,dataframe]"
}
}
}
diff --git a/examples/notebooks/tour-of-beam/dataframes.ipynb b/examples/notebooks/tour-of-beam/dataframes.ipynb
index c19d991..ac0ad10 100644
--- a/examples/notebooks/tour-of-beam/dataframes.ipynb
+++ b/examples/notebooks/tour-of-beam/dataframes.ipynb
@@ -65,7 +65,7 @@
"[Beam DataFrames overview](https://beam.apache.org/documentation/dsls/dataframes/overview) page.\n",
"\n",
"First, we need to install Apache Beam with the `interactive` extra for the Interactive runner.",
- "We also need `pandas` for this notebook, but the Interactive runner already depends on it."
+ "We also need to install a version of `pandas` supported by the DataFrame API, which we can get with the `dataframe` extra in Beam 2.34.0 and newer."
],
"metadata": {
"id": "hDuXLLSZnI1D"
@@ -75,7 +75,7 @@
"cell_type": "code",
"execution_count": null,
"source": [
- "%pip install --quiet apache-beam[interactive]"
+ "%pip install --quiet apache-beam[interactive,dataframe]"
],
"outputs": [],
"metadata": {
@@ -663,7 +663,7 @@
"\n",
"> ℹ️ It's recommended to **only** do this if you need to use a Pandas operation that is\n",
"> [not supported in Beam DataFrames](https://beam.apache.org/documentation/dsls/dataframes/differences-from-pandas/#classes-of-unsupported-operations).\n",
- "> Converting a PCollection into a Pandas DataFrame consolidates elements from potentially multiple workers into a single worker, which could create a performance bottleneck.\n"
+ "> Converting a PCollection into a Pandas DataFrame consolidates elements from potentially multiple workers into a single worker, which could create a performance bottleneck.\n",
"\n",
"> ⚠️ Pandas DataFrames are in-memory data structures, so make sure all the elements in the PCollection fit into memory.\n",
"> If they don't fit into memory, consider yielding multiple DataFrame elements via\n",
diff --git a/gradle.properties b/gradle.properties
index e3a005d..c5e3f76 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -24,8 +24,8 @@
signing.gnupg.executable=gpg
signing.gnupg.useLegacyGpg=true
-version=2.34.0-SNAPSHOT
-sdk_version=2.34.0.dev
+version=2.35.0-SNAPSHOT
+sdk_version=2.35.0.dev
javaVersion=1.8
diff --git a/learning/katas/python/log_elements.py b/learning/katas/python/log_elements.py
index 6989cee..4477256 100644
--- a/learning/katas/python/log_elements.py
+++ b/learning/katas/python/log_elements.py
@@ -22,7 +22,7 @@
class _LoggingFn(beam.DoFn):
def __init__(self, prefix='', with_timestamp=False, with_window=False):
- super(LogElements._LoggingFn, self).__init__()
+ super().__init__()
self.prefix = prefix
self.with_timestamp = with_timestamp
self.with_window = with_window
@@ -43,7 +43,7 @@
def __init__(self, label=None, prefix='',
with_timestamp=False, with_window=False):
- super(LogElements, self).__init__(label)
+ super().__init__(label)
self.prefix = prefix
self.with_timestamp = with_timestamp
self.with_window = with_window
diff --git a/model/pipeline/src/main/proto/beam_runner_api.proto b/model/pipeline/src/main/proto/beam_runner_api.proto
index 80370c5..39fd9d8 100644
--- a/model/pipeline/src/main/proto/beam_runner_api.proto
+++ b/model/pipeline/src/main/proto/beam_runner_api.proto
@@ -1718,7 +1718,14 @@
string string_value = 2;
bool bool_value = 3;
double double_value = 4;
+ int64 int_value = 5;
}
+
+ // (Required) The key identifies the actual content of the metadata.
+ string key = 6;
+
+ // (Required) The namespace describes the context that specified the key.
+ string namespace = 7;
}
// Static display data associated with a pipeline component. Display data is
diff --git a/model/pipeline/src/main/proto/metrics.proto b/model/pipeline/src/main/proto/metrics.proto
index 8f819b6..913b2d0 100644
--- a/model/pipeline/src/main/proto/metrics.proto
+++ b/model/pipeline/src/main/proto/metrics.proto
@@ -420,6 +420,10 @@
BIGTABLE_PROJECT_ID = 20 [(label_props) = { name: "BIGTABLE_PROJECT_ID"}];
INSTANCE_ID = 21 [(label_props) = { name: "INSTANCE_ID"}];
TABLE_ID = 22 [(label_props) = { name: "TABLE_ID"}];
+ SPANNER_PROJECT_ID = 23 [(label_props) = { name: "SPANNER_PROJECT_ID"}];
+ SPANNER_DATABASE_ID = 24 [(label_props) = { name: "SPANNER_DATABASE_ID"}];
+ SPANNER_INSTANCE_ID = 25 [(label_props) = { name: "SPANNER_INSTANCE_ID" }];
+ SPANNER_QUERY_NAME = 26 [(label_props) = { name: "SPANNER_QUERY_NAME" }];
}
// A set of key and value labels which define the scope of the metric. For
diff --git a/playground/backend/.gitkeep b/playground/backend/.gitkeep
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/playground/backend/.gitkeep
diff --git a/playground/backend/README.md b/playground/backend/README.md
new file mode 100644
index 0000000..e845e3b
--- /dev/null
+++ b/playground/backend/README.md
@@ -0,0 +1,101 @@
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT 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 Playground
+
+## About
+
+Apache Beam is an open-source, unified model for defining parallel processing pipelines for batch and streaming data.
+It provides a portable API layer for building sophisticated data-parallel processing pipelines that may be executed
+across a diversity of execution engines, or runners.
+
+Beam Playground helps facilitate trying out and adopting Apache Beam by providing a very quick way for prospective Beam
+users to see and run examples of Apache Beam pipelines in a web interface that requires no setup.
+
+## Requirements
+
+- [Go](https://golang.org/dl/) installed (1.16 version).
+- [Protocol Buffer Compiler](https://grpc.io/docs/protoc-installation/)
+
+## Getting Started
+
+This section describes what is needed to run the backend application.
+- Generating models from proto file
+- Go commands to run/test application locally
+
+### Generating models from proto file
+
+Beam Playground uses gRPC for communication between client and server. For using gRPC there is a
+`playground/proto/playground.proto` file which contains the definition of all models. To use models in the code they
+should be generated firstly using the `playground/proto/playground.proto` file.
+
+Run the following command from `/playground` directory:
+
+```
+protoc --go_out=backend/pkg/api --go_opt=paths=source_relative \
+--go-grpc_out=backend/pkg/api --go-grpc_opt=paths=source_relative --proto_path=playground/v1 \
+playground.proto
+```
+
+As a result you will receive 2 files into `backend/pkg/api` folder with models (`playground.pb.go`) and client/server
+structures (`playground_grpc.pb.go`).
+More information about using gRPC you can find [here](https://grpc.io/docs/languages/).
+
+### Go commands to run/test application locally
+
+The following command is used to build and serve the backend locally:
+
+```
+go run ./cmd/server/server.go
+```
+
+Run the following command to generate a release build file:
+
+```
+go build ./cmd/server/server.go
+```
+
+Playground tests may be run using this command:
+
+```
+go test ./test/... -v
+```
+
+The full list of commands can be found [here](https://pkg.go.dev/cmd/go).
+
+## Running the server app
+
+To run the server using Docker images there are `Docker` files in the `containers` folder for Go and Java languages.
+Each of them processes the corresponding SDK, so the backend with Go SDK will work with Go examples/katas/tests only.
+
+One more way to run the server is to run it locally how it is described above.
+
+## Calling the server from another client
+
+To call the server from another client – models and client code should be generated using the
+`playground/playground/v1/playground.proto` file. More information about generating models and client's code using `.proto`
+files for each language can be found [here](https://grpc.io/docs/languages/).
+
+## Deployment
+
+TBD
+
+## How to Contribute
+
+TBD
diff --git a/playground/backend/cmd/server/server.go b/playground/backend/cmd/server/server.go
new file mode 100644
index 0000000..bbeb155
--- /dev/null
+++ b/playground/backend/cmd/server/server.go
@@ -0,0 +1,24 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 main
+
+import (
+ "beam.apache.org/playground/backend/internal/executors"
+)
+
+func main() {
+ _ = executors.GoExecutor{}
+}
diff --git a/playground/backend/containers/go/Dockerfile b/playground/backend/containers/go/Dockerfile
new file mode 100644
index 0000000..5786f50
--- /dev/null
+++ b/playground/backend/containers/go/Dockerfile
@@ -0,0 +1,19 @@
+###############################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+###############################################################################
+
+#Dokerfile to set up the Beam Go SDK
\ No newline at end of file
diff --git a/playground/backend/containers/java/Dockerfile b/playground/backend/containers/java/Dockerfile
new file mode 100644
index 0000000..6684267
--- /dev/null
+++ b/playground/backend/containers/java/Dockerfile
@@ -0,0 +1,19 @@
+###############################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+###############################################################################
+
+#Dokerfile to set up the Beam Java SDK
\ No newline at end of file
diff --git a/playground/backend/go.mod b/playground/backend/go.mod
new file mode 100644
index 0000000..7bc8c78
--- /dev/null
+++ b/playground/backend/go.mod
@@ -0,0 +1,23 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 beam.apache.org/playground/backend
+
+go 1.16
+
+require (
+ google.golang.org/grpc v1.41.0
+ google.golang.org/protobuf v1.27.1
+)
diff --git a/playground/backend/internal/api/.gitkeep b/playground/backend/internal/api/.gitkeep
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/playground/backend/internal/api/.gitkeep
diff --git a/playground/backend/internal/errors/grpc_errors.go b/playground/backend/internal/errors/grpc_errors.go
new file mode 100644
index 0000000..6919622
--- /dev/null
+++ b/playground/backend/internal/errors/grpc_errors.go
@@ -0,0 +1,36 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 errors
+
+import (
+ "google.golang.org/grpc/codes"
+ "google.golang.org/grpc/status"
+)
+
+// InvalidArgumentError Returns error with InvalidArgument code error and message like "title: message"
+func InvalidArgumentError(title string, message string) error {
+ return status.Errorf(codes.InvalidArgument, "%s: %s", title, message)
+}
+
+// NotFoundError Returns error with NotFound code error and message like "title: message"
+func NotFoundError(title string, message string) error {
+ return status.Errorf(codes.NotFound, "%s: %s", title, message)
+}
+
+// InternalError Returns error with Internal code error and message like "title: message"
+func InternalError(title string, message string) error {
+ return status.Errorf(codes.Internal, "%s: %s", title, message)
+}
diff --git a/playground/backend/internal/errors/grpc_errors_test.go b/playground/backend/internal/errors/grpc_errors_test.go
new file mode 100644
index 0000000..ef9720e
--- /dev/null
+++ b/playground/backend/internal/errors/grpc_errors_test.go
@@ -0,0 +1,102 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 errors
+
+import (
+ "strings"
+ "testing"
+)
+
+func TestInternalError(t *testing.T) {
+ type args struct {
+ title string
+ message string
+ }
+ tests := []struct {
+ name string
+ args args
+ expected string
+ wantErr bool
+ }{
+ {name: "TestInternalError", args: args{title: "TEST_TITLE", message: "TEST_MESSAGE"},
+ expected: "rpc error: code = Internal desc = TEST_TITLE: TEST_MESSAGE", wantErr: true},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ err := InternalError(tt.args.title, tt.args.message)
+ if (err != nil) != tt.wantErr {
+ t.Errorf("InternalError() error = %v, wantErr %v", err, tt.wantErr)
+ }
+ if !strings.EqualFold(err.Error(), tt.expected) {
+ t.Errorf("InternalError() error = %v, wantErr %v", err.Error(), tt.expected)
+ }
+ })
+ }
+}
+
+func TestInvalidArgumentError(t *testing.T) {
+ type args struct {
+ title string
+ message string
+ }
+ tests := []struct {
+ name string
+ args args
+ expected string
+ wantErr bool
+ }{
+ {name: "TestInvalidArgumentError", args: args{title: "TEST_TITLE", message: "TEST_MESSAGE"},
+ expected: "rpc error: code = InvalidArgument desc = TEST_TITLE: TEST_MESSAGE", wantErr: true},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ err := InvalidArgumentError(tt.args.title, tt.args.message)
+ if (err != nil) != tt.wantErr {
+ t.Errorf("InvalidArgumentError() error = %v, wantErr %v", err, tt.wantErr)
+ }
+ if !strings.EqualFold(err.Error(), tt.expected) {
+ t.Errorf("InvalidArgumentError() error = %v, wantErr %v", err.Error(), tt.expected)
+ }
+ })
+ }
+}
+
+func TestNotFoundError(t *testing.T) {
+ type args struct {
+ title string
+ message string
+ }
+ tests := []struct {
+ name string
+ args args
+ expected string
+ wantErr bool
+ }{
+ {name: "TestNotFoundError", args: args{title: "TEST_TITLE", message: "TEST_MESSAGE"},
+ expected: "rpc error: code = NotFound desc = TEST_TITLE: TEST_MESSAGE", wantErr: true},
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ err := NotFoundError(tt.args.title, tt.args.message)
+ if (err != nil) != tt.wantErr {
+ t.Errorf("NotFoundError() error = %v, wantErr %v", err, tt.wantErr)
+ }
+ if !strings.EqualFold(err.Error(), tt.expected) {
+ t.Errorf("NotFoundError() error = %v, wantErr %v", err.Error(), tt.expected)
+ }
+ })
+ }
+}
diff --git a/playground/backend/internal/executors/executor.go b/playground/backend/internal/executors/executor.go
new file mode 100644
index 0000000..ecc655f
--- /dev/null
+++ b/playground/backend/internal/executors/executor.go
@@ -0,0 +1,31 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Interface for all executors (Java/Python/Go/SCIO)
+package executors
+
+type executor interface {
+ // Validate validates executable file.
+ // Return result of validation (true/false) and error if it occurs
+ Validate(filePath string) (bool, error)
+
+ // Compile compiles executable file.
+ // Return error if it occurs
+ Compile(filePath string) error
+
+ // Run runs executable file.
+ // Return logs and error if it occurs
+ Run(filePath string) (string, error)
+}
diff --git a/playground/backend/internal/executors/goexecutor.go b/playground/backend/internal/executors/goexecutor.go
new file mode 100644
index 0000000..339d09f
--- /dev/null
+++ b/playground/backend/internal/executors/goexecutor.go
@@ -0,0 +1,31 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Executor for Go
+package executors
+
+type GoExecutor struct{}
+
+func (goExec GoExecutor) Validate(filePath string) (bool, error) {
+ return true, nil
+}
+
+func (goExec GoExecutor) Compile(filePath string) error {
+ return nil
+}
+
+func (goExec GoExecutor) Run(filePath string) (string, error) {
+ return "", nil
+}
diff --git a/playground/backend/internal/executors/javaexecutor.go b/playground/backend/internal/executors/javaexecutor.go
new file mode 100644
index 0000000..e67f715
--- /dev/null
+++ b/playground/backend/internal/executors/javaexecutor.go
@@ -0,0 +1,31 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Executor for Java
+package executors
+
+type JavaExecutor struct{}
+
+func (javaExec JavaExecutor) Validate(filePath string) (bool, error) {
+ return true, nil
+}
+
+func (javaExec JavaExecutor) Compile(filePath string) error {
+ return nil
+}
+
+func (javaExec JavaExecutor) Run(filePath string) (string, error) {
+ return "", nil
+}
diff --git a/playground/frontend/.metadata b/playground/frontend/.metadata
new file mode 100644
index 0000000..0f055bf
--- /dev/null
+++ b/playground/frontend/.metadata
@@ -0,0 +1,10 @@
+# This file tracks properties of this Flutter project.
+# Used by Flutter tool to assess capabilities and perform upgrades etc.
+#
+# This file should be version controlled and should not be manually edited.
+
+version:
+ revision: ffb2ecea5223acdd139a5039be2f9c796962833d
+ channel: stable
+
+project_type: app
diff --git a/playground/frontend/README.md b/playground/frontend/README.md
new file mode 100644
index 0000000..661945f
--- /dev/null
+++ b/playground/frontend/README.md
@@ -0,0 +1,47 @@
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT 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 Playground
+
+## About
+
+Apache Beam is an open-source, unified model for defining parallel processing pipelines for batch and streaming data.
+It provides a portable API layer for building sophisticated data-parallel processing pipelines that may be executed across a diversity of execution engines, or runners.
+
+## Getting Started
+
+Website development requires [Flutter](https://flutter.dev/docs/get-started/install) installed.
+
+The following command is used to build and serve the website locally:
+
+`$ flutter run`
+
+Run the following command to generate a release build:
+
+`flutter build web`
+
+Playground tests may be run using this command:
+
+`flutter test`
+
+Dart code should follow next [code style](https://dart-lang.github.io/linter/lints/index.html). Code may be analyzed using this command:
+
+`flutter analyze`
+
+The full list of command can be found [here](https://flutter.dev/docs/reference/flutter-cli)
diff --git a/playground/frontend/analysis_options.yaml b/playground/frontend/analysis_options.yaml
new file mode 100644
index 0000000..3080e1b
--- /dev/null
+++ b/playground/frontend/analysis_options.yaml
@@ -0,0 +1,46 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT 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 configures the analyzer, which statically analyzes Dart code to
+# check for errors, warnings, and lints.
+#
+# The issues identified by the analyzer are surfaced in the UI of Dart-enabled
+# IDEs (https://dart.dev/tools#ides-and-editors). The analyzer can also be
+# invoked from the command line by running `flutter analyze`.
+
+# The following line activates a set of recommended lints for Flutter apps,
+# packages, and plugins designed to encourage good coding practices.
+include: package:flutter_lints/flutter.yaml
+
+linter:
+ # The lint rules applied to this project can be customized in the
+ # section below to disable rules from the `package:flutter_lints/flutter.yaml`
+ # included above or to enable additional rules. A list of all available lints
+ # and their documentation is published at
+ # https://dart-lang.github.io/linter/lints/index.html.
+ #
+ # Instead of disabling a lint rule for the entire project in the
+ # section below, it can also be suppressed for a single line of code
+ # or a specific dart file by using the `// ignore: name_of_lint` and
+ # `// ignore_for_file: name_of_lint` syntax on the line or in the file
+ # producing the lint.
+ rules:
+ # avoid_print: false # Uncomment to disable the `avoid_print` rule
+ # prefer_single_quotes: true # Uncomment to enable the `prefer_single_quotes` rule
+
+# Additional information about this file can be found at
+# https://dart.dev/guides/language/analysis-options
diff --git a/playground/frontend/lib/components/logo/logo_component.dart b/playground/frontend/lib/components/logo/logo_component.dart
new file mode 100644
index 0000000..7c6bc1b
--- /dev/null
+++ b/playground/frontend/lib/components/logo/logo_component.dart
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+
+class Logo extends StatelessWidget {
+ const Logo({Key? key}) : super(key: key);
+
+ @override
+ Widget build(BuildContext context) {
+ final theme = Theme.of(context);
+ return RichText(
+ text: TextSpan(
+ style: theme.textTheme.headline6,
+ children: <TextSpan>[
+ TextSpan(text: 'Beam', style: TextStyle(color: theme.primaryColor)),
+ const TextSpan(text: ' Playground'),
+ ],
+ ),
+ );
+ }
+}
diff --git a/playground/frontend/lib/components/toggle_theme_button/toggle_theme_button.dart b/playground/frontend/lib/components/toggle_theme_button/toggle_theme_button.dart
new file mode 100644
index 0000000..e6ca372
--- /dev/null
+++ b/playground/frontend/lib/components/toggle_theme_button/toggle_theme_button.dart
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+import 'package:playground/config/theme.dart';
+import 'package:playground/constants/sizes.dart';
+import 'package:provider/provider.dart';
+
+const kLightMode = "Light Mode";
+const kDartMode = "Dark Mode";
+
+class ToggleThemeButton extends StatelessWidget {
+ const ToggleThemeButton({Key? key}) : super(key: key);
+
+ @override
+ Widget build(BuildContext context) {
+ return Consumer<ThemeProvider>(builder: (context, theme, child) {
+ final text = theme.isDarkMode ? kLightMode : kDartMode;
+ final icon = theme.isDarkMode
+ ? Icons.light_mode_outlined
+ : Icons.mode_night_outlined;
+ return Padding(
+ padding: const EdgeInsets.symmetric(
+ vertical: kSmPadding,
+ horizontal: kMdPadding,
+ ),
+ child: TextButton.icon(
+ icon: Icon(icon),
+ label: Text(text),
+ onPressed: theme.toggleTheme,
+ ),
+ );
+ });
+ }
+}
diff --git a/playground/frontend/lib/config/theme.dart b/playground/frontend/lib/config/theme.dart
new file mode 100644
index 0000000..4763452
--- /dev/null
+++ b/playground/frontend/lib/config/theme.dart
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+import 'package:playground/constants/colors.dart';
+import 'package:playground/constants/sizes.dart';
+import 'package:provider/provider.dart';
+
+class ThemeProvider extends ChangeNotifier {
+ ThemeMode themeMode = ThemeMode.light;
+
+ bool get isDarkMode {
+ return themeMode == ThemeMode.dark;
+ }
+
+ void toggleTheme() {
+ themeMode = themeMode == ThemeMode.light ? ThemeMode.dark : ThemeMode.light;
+ notifyListeners();
+ }
+}
+
+final kLightTheme = ThemeData(
+ brightness: Brightness.light,
+ primaryColor: kLightPrimary,
+ backgroundColor: kLightPrimaryBackground,
+ appBarTheme: const AppBarTheme(
+ color: kLightSecondaryBackground,
+ elevation: 1,
+ centerTitle: false,
+ ),
+ textButtonTheme: TextButtonThemeData(
+ style: TextButton.styleFrom(
+ primary: kLightText,
+ shape: const RoundedRectangleBorder(
+ borderRadius: BorderRadius.all(Radius.circular(kBorderRadius)),
+ ),
+ ),
+ ),
+);
+
+final kDarkTheme = ThemeData(
+ brightness: Brightness.dark,
+ primaryColor: kDarkPrimary,
+ backgroundColor: kDarkGrey,
+ appBarTheme: const AppBarTheme(
+ color: kDarkSecondaryBackground,
+ elevation: 1,
+ centerTitle: false,
+ ),
+ textButtonTheme: TextButtonThemeData(
+ style: TextButton.styleFrom(
+ primary: kDarkText,
+ shape: const RoundedRectangleBorder(
+ borderRadius: BorderRadius.all(Radius.circular(kBorderRadius)),
+ ),
+ ),
+ ),
+);
+
+class ThemeColors {
+ final bool isDark;
+
+ static ThemeColors of(BuildContext context) {
+ final theme = Provider.of<ThemeProvider>(context);
+ return ThemeColors(theme.isDarkMode);
+ }
+
+ ThemeColors(this.isDark);
+
+ Color get greyColor => isDark ? kDarkGrey : kLightGrey;
+}
diff --git a/playground/frontend/lib/constants/colors.dart b/playground/frontend/lib/constants/colors.dart
new file mode 100644
index 0000000..f382ea6
--- /dev/null
+++ b/playground/frontend/lib/constants/colors.dart
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+
+const Color kLightPrimaryBackground = Colors.white;
+const Color kLightSecondaryBackground = Color(0xFFFCFCFC);
+const Color kLightGrey = Color(0xFFE5E5E5);
+const Color kLightGrey1 = Color(0xFFA0A4AB);
+const Color kLightText = Color(0xFF242639);
+const Color kLightPrimary = Color(0xFFE74D1A);
+
+const Color kDarkPrimaryBackground = Color(0xFF18181B);
+const Color kDarkSecondaryBackground = Color(0xFF2E2E34);
+const Color kDarkGrey = Color(0xFF3F3F46);
+const Color kDarkGrey1 = Color(0xFF606772);
+const Color kDarkText = Color(0xFFFFFFFF);
+const Color kDarkPrimary = Color(0xFFF26628);
diff --git a/playground/frontend/lib/constants/sizes.dart b/playground/frontend/lib/constants/sizes.dart
new file mode 100644
index 0000000..5e2ccd0
--- /dev/null
+++ b/playground/frontend/lib/constants/sizes.dart
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// paddings
+const double kZeroPadding = 0.0;
+const double kSmPadding = 4.0;
+const double kMdPadding = 8.0;
+const double kLgPadding = 16.0;
+
+// border radius
+const double kBorderRadius = 8.0;
+
+// elevation
+const int kElevation = 1;
+
+// icon sizes
+const double kIconSizeMd = 24.0;
diff --git a/playground/frontend/lib/main.dart b/playground/frontend/lib/main.dart
new file mode 100644
index 0000000..3ce8c37
--- /dev/null
+++ b/playground/frontend/lib/main.dart
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+import 'package:playground/playground_app.dart';
+
+void main() {
+ runApp(const PlaygroundApp());
+}
diff --git a/playground/frontend/lib/modules/editor/components/editor_textarea.dart b/playground/frontend/lib/modules/editor/components/editor_textarea.dart
new file mode 100644
index 0000000..89412f4
--- /dev/null
+++ b/playground/frontend/lib/modules/editor/components/editor_textarea.dart
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+
+class EditorTextArea extends StatelessWidget {
+ const EditorTextArea({Key? key}) : super(key: key);
+
+ @override
+ Widget build(BuildContext context) {
+ return const Center(child: Text('Editor Text Area'));
+ }
+}
diff --git a/playground/frontend/lib/modules/output/components/output_area.dart b/playground/frontend/lib/modules/output/components/output_area.dart
new file mode 100644
index 0000000..26b2033
--- /dev/null
+++ b/playground/frontend/lib/modules/output/components/output_area.dart
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+
+class OutputArea extends StatelessWidget {
+ const OutputArea({Key? key}) : super(key: key);
+
+ @override
+ Widget build(BuildContext context) {
+ return const Center(
+ child: Text('Output'),
+ );
+ }
+}
diff --git a/playground/frontend/lib/modules/sdk/components/sdk_selector.dart b/playground/frontend/lib/modules/sdk/components/sdk_selector.dart
new file mode 100644
index 0000000..37f9e79
--- /dev/null
+++ b/playground/frontend/lib/modules/sdk/components/sdk_selector.dart
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+import 'package:playground/config/theme.dart';
+import 'package:playground/constants/sizes.dart';
+import 'package:playground/modules/sdk/models/sdk.dart';
+
+typedef SetSdk = void Function(SDK sdk);
+
+class SDKSelector extends StatelessWidget {
+ final SDK sdk;
+ final SetSdk setSdk;
+
+ const SDKSelector({Key? key, required this.sdk, required this.setSdk})
+ : super(key: key);
+
+ @override
+ Widget build(BuildContext context) {
+ return Container(
+ padding: const EdgeInsets.symmetric(
+ vertical: kZeroPadding,
+ horizontal: kLgPadding,
+ ),
+ decoration: BoxDecoration(
+ color: ThemeColors.of(context).greyColor,
+ borderRadius: BorderRadius.circular(kBorderRadius),
+ ),
+ child: DropdownButtonHideUnderline(
+ child: DropdownButton<SDK>(
+ value: sdk,
+ icon: const Icon(Icons.keyboard_arrow_down),
+ iconSize: kIconSizeMd,
+ elevation: kElevation,
+ borderRadius: BorderRadius.circular(kBorderRadius),
+ alignment: Alignment.bottomCenter,
+ onChanged: (SDK? newSdk) {
+ if (newSdk != null) {
+ setSdk(newSdk);
+ }
+ },
+ items: SDK.values.map<DropdownMenuItem<SDK>>((SDK value) {
+ return DropdownMenuItem<SDK>(
+ value: value,
+ child: Text(value.displayName),
+ );
+ }).toList(),
+ ),
+ ),
+ );
+ }
+}
diff --git a/playground/frontend/lib/modules/sdk/models/sdk.dart b/playground/frontend/lib/modules/sdk/models/sdk.dart
new file mode 100644
index 0000000..1406a03
--- /dev/null
+++ b/playground/frontend/lib/modules/sdk/models/sdk.dart
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+enum SDK {
+ java,
+ go,
+ python,
+ scio,
+}
+
+extension SDKToString on SDK {
+ String get displayName {
+ switch (this) {
+ case SDK.go:
+ return "Go";
+ case SDK.java:
+ return "Java";
+ case SDK.python:
+ return "Python";
+ case SDK.scio:
+ return "SCIO";
+ }
+ }
+}
diff --git a/playground/frontend/lib/pages/playground/playground_page.dart b/playground/frontend/lib/pages/playground/playground_page.dart
new file mode 100644
index 0000000..fb053db
--- /dev/null
+++ b/playground/frontend/lib/pages/playground/playground_page.dart
@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+import 'package:playground/components/toggle_theme_button/toggle_theme_button.dart';
+import 'package:provider/provider.dart';
+import 'package:playground/modules/output/components/output_area.dart';
+import 'package:playground/pages/playground/playground_state.dart';
+import 'package:playground/modules/editor/components/editor_textarea.dart';
+import 'package:playground/modules/sdk/components/sdk_selector.dart';
+import 'package:playground/components/logo/logo_component.dart';
+
+class PlaygroundPage extends StatelessWidget {
+ const PlaygroundPage({Key? key}) : super(key: key);
+
+ @override
+ Widget build(BuildContext context) {
+ return ChangeNotifierProvider<PlaygroundState>(
+ create: (context) => PlaygroundState(),
+ child: Scaffold(
+ appBar: AppBar(
+ title: Wrap(
+ crossAxisAlignment: WrapCrossAlignment.center,
+ spacing: 16.0,
+ children: [
+ const Logo(),
+ Consumer<PlaygroundState>(
+ builder: (context, state, child) {
+ return SDKSelector(
+ sdk: state.sdk,
+ setSdk: state.setSdk,
+ );
+ },
+ ),
+ ],
+ ),
+ actions: const [ToggleThemeButton()],
+ ),
+ body: Column(
+ children: [
+ const Expanded(child: EditorTextArea()),
+ Container(height: 16.0, color: Theme.of(context).backgroundColor),
+ const Expanded(child: OutputArea()),
+ ],
+ ),
+ ),
+ );
+ }
+}
diff --git a/playground/frontend/lib/pages/playground/playground_state.dart b/playground/frontend/lib/pages/playground/playground_state.dart
new file mode 100644
index 0000000..d50929b
--- /dev/null
+++ b/playground/frontend/lib/pages/playground/playground_state.dart
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+import 'package:playground/modules/sdk/models/sdk.dart';
+
+class PlaygroundState with ChangeNotifier {
+ SDK sdk = SDK.java;
+
+ setSdk(SDK sdk) {
+ this.sdk = sdk;
+ notifyListeners();
+ }
+}
diff --git a/playground/frontend/lib/playground_app.dart b/playground/frontend/lib/playground_app.dart
new file mode 100644
index 0000000..8ce11d1
--- /dev/null
+++ b/playground/frontend/lib/playground_app.dart
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter/material.dart';
+import 'package:playground/config/theme.dart';
+import 'package:playground/pages/playground/playground_page.dart';
+import 'package:provider/provider.dart';
+
+class PlaygroundApp extends StatelessWidget {
+ const PlaygroundApp({Key? key}) : super(key: key);
+
+ @override
+ Widget build(BuildContext context) {
+ return ChangeNotifierProvider(
+ create: (context) => ThemeProvider(),
+ builder: (context, _) {
+ final themeProvider = Provider.of<ThemeProvider>(context);
+ return MaterialApp(
+ title: 'Apache Beam Playground',
+ themeMode: themeProvider.themeMode,
+ theme: kLightTheme,
+ darkTheme: kDarkTheme,
+ home: const PlaygroundPage(),
+ debugShowCheckedModeBanner: false,
+ );
+ },
+ );
+ }
+}
diff --git a/playground/frontend/pubspec.lock b/playground/frontend/pubspec.lock
new file mode 100644
index 0000000..9a31ab2
--- /dev/null
+++ b/playground/frontend/pubspec.lock
@@ -0,0 +1,175 @@
+# Generated by pub
+# See https://dart.dev/tools/pub/glossary#lockfile
+packages:
+ async:
+ dependency: transitive
+ description:
+ name: async
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "2.8.1"
+ boolean_selector:
+ dependency: transitive
+ description:
+ name: boolean_selector
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "2.1.0"
+ characters:
+ dependency: transitive
+ description:
+ name: characters
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.1.0"
+ charcode:
+ dependency: transitive
+ description:
+ name: charcode
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.3.1"
+ clock:
+ dependency: transitive
+ description:
+ name: clock
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.1.0"
+ collection:
+ dependency: transitive
+ description:
+ name: collection
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.15.0"
+ fake_async:
+ dependency: transitive
+ description:
+ name: fake_async
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.2.0"
+ flutter:
+ dependency: "direct main"
+ description: flutter
+ source: sdk
+ version: "0.0.0"
+ flutter_lints:
+ dependency: "direct dev"
+ description:
+ name: flutter_lints
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.0.4"
+ flutter_test:
+ dependency: "direct dev"
+ description: flutter
+ source: sdk
+ version: "0.0.0"
+ lints:
+ dependency: transitive
+ description:
+ name: lints
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.0.1"
+ matcher:
+ dependency: transitive
+ description:
+ name: matcher
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "0.12.10"
+ meta:
+ dependency: transitive
+ description:
+ name: meta
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.7.0"
+ nested:
+ dependency: transitive
+ description:
+ name: nested
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.0.0"
+ path:
+ dependency: transitive
+ description:
+ name: path
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.8.0"
+ provider:
+ dependency: "direct main"
+ description:
+ name: provider
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "6.0.1"
+ sky_engine:
+ dependency: transitive
+ description: flutter
+ source: sdk
+ version: "0.0.99"
+ source_span:
+ dependency: transitive
+ description:
+ name: source_span
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.8.1"
+ stack_trace:
+ dependency: transitive
+ description:
+ name: stack_trace
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.10.0"
+ stream_channel:
+ dependency: transitive
+ description:
+ name: stream_channel
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "2.1.0"
+ string_scanner:
+ dependency: transitive
+ description:
+ name: string_scanner
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.1.0"
+ term_glyph:
+ dependency: transitive
+ description:
+ name: term_glyph
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.2.0"
+ test_api:
+ dependency: transitive
+ description:
+ name: test_api
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "0.4.2"
+ typed_data:
+ dependency: transitive
+ description:
+ name: typed_data
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "1.3.0"
+ vector_math:
+ dependency: transitive
+ description:
+ name: vector_math
+ url: "https://pub.dartlang.org"
+ source: hosted
+ version: "2.1.0"
+sdks:
+ dart: ">=2.12.0 <3.0.0"
+ flutter: ">=1.16.0"
diff --git a/playground/frontend/pubspec.yaml b/playground/frontend/pubspec.yaml
new file mode 100644
index 0000000..3374e9e
--- /dev/null
+++ b/playground/frontend/pubspec.yaml
@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT 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: playground
+description: A new Flutter project.
+
+# The following line prevents the package from being accidentally published to
+# pub.dev using `flutter pub publish`. This is preferred for private packages.
+publish_to: 'none' # Remove this line if you wish to publish to pub.dev
+
+# The following defines the version and build number for your application.
+# A version number is three numbers separated by dots, like 1.2.43
+# followed by an optional build number separated by a +.
+# Both the version and the builder number may be overridden in flutter
+# build by specifying --build-name and --build-number, respectively.
+# In Android, build-name is used as versionName while build-number used as versionCode.
+# Read more about Android versioning at https://developer.android.com/studio/publish/versioning
+# In iOS, build-name is used as CFBundleShortVersionString while build-number used as CFBundleVersion.
+# Read more about iOS versioning at
+# https://developer.apple.com/library/archive/documentation/General/Reference/InfoPlistKeyReference/Articles/CoreFoundationKeys.html
+version: 1.0.0+1
+
+environment:
+ sdk: ">=2.12.0 <3.0.0"
+
+# Dependencies specify other packages that your package needs in order to work.
+# To automatically upgrade your package dependencies to the latest versions
+# consider running `flutter pub upgrade --major-versions`. Alternatively,
+# dependencies can be manually updated by changing the version numbers below to
+# the latest version available on pub.dev. To see which dependencies have newer
+# versions available, run `flutter pub outdated`.
+dependencies:
+ flutter:
+ sdk: flutter
+ provider: ^6.0.0
+
+dev_dependencies:
+ flutter_test:
+ sdk: flutter
+
+ # The "flutter_lints" package below contains a set of recommended lints to
+ # encourage good coding practices. The lint set provided by the package is
+ # activated in the `analysis_options.yaml` file located at the root of your
+ # package. See that file for information about deactivating specific lint
+ # rules and activating additional ones.
+ flutter_lints: ^1.0.0
+
+# For information on the generic Dart part of this file, see the
+# following page: https://dart.dev/tools/pub/pubspec
+
+# The following section is specific to Flutter.
+flutter:
+
+ # The following line ensures that the Material Icons font is
+ # included with your application, so that you can use the icons in
+ # the material Icons class.
+ uses-material-design: true
+
+ # To add assets to your application, add an assets section, like this:
+ # assets:
+ # - images/a_dot_burr.jpeg
+ # - images/a_dot_ham.jpeg
+
+ # An image asset can refer to one or more resolution-specific "variants", see
+ # https://flutter.dev/assets-and-images/#resolution-aware.
+
+ # For details regarding adding assets from package dependencies, see
+ # https://flutter.dev/assets-and-images/#from-packages
+
+ # To add custom fonts to your application, add a fonts section here,
+ # in this "flutter" section. Each entry in this list should have a
+ # "family" key with the font family name, and a "fonts" key with a
+ # list giving the asset and other descriptors for the font. For
+ # example:
+ # fonts:
+ # - family: Schyler
+ # fonts:
+ # - asset: fonts/Schyler-Regular.ttf
+ # - asset: fonts/Schyler-Italic.ttf
+ # style: italic
+ # - family: Trajan Pro
+ # fonts:
+ # - asset: fonts/TrajanPro.ttf
+ # - asset: fonts/TrajanPro_Bold.ttf
+ # weight: 700
+ #
+ # For details regarding fonts from package dependencies,
+ # see https://flutter.dev/custom-fonts/#from-packages
diff --git a/playground/frontend/test/pages/playground/playground_state_test.dart b/playground/frontend/test/pages/playground/playground_state_test.dart
new file mode 100644
index 0000000..1c29871
--- /dev/null
+++ b/playground/frontend/test/pages/playground/playground_state_test.dart
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:playground/modules/sdk/models/sdk.dart';
+import 'package:flutter_test/flutter_test.dart';
+import 'package:playground/pages/playground/playground_state.dart';
+
+void main() {
+ test('Playground State initial value should be java', () {
+ final state = PlaygroundState();
+ expect(state.sdk, equals(SDK.java));
+ });
+
+ test('Playground state should notify all listeners about sdk change', () {
+ final state = PlaygroundState();
+ state.addListener(() {
+ expect(state.sdk, SDK.go);
+ });
+ state.setSdk(SDK.go);
+ });
+}
diff --git a/playground/frontend/test/widget_test.dart b/playground/frontend/test/widget_test.dart
new file mode 100644
index 0000000..2e17b1c
--- /dev/null
+++ b/playground/frontend/test/widget_test.dart
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 'package:flutter_test/flutter_test.dart';
+import 'package:playground/playground_app.dart';
+
+void main() {
+ testWidgets('Home Page', (WidgetTester tester) async {
+ // Build our app and trigger a frame.
+ await tester.pumpWidget(const PlaygroundApp());
+
+ // Verify that Playground text is displayed
+ expect(find.text('Playground'), findsOneWidget);
+ });
+}
diff --git a/playground/frontend/web/favicon.png b/playground/frontend/web/favicon.png
new file mode 100644
index 0000000..8aaa46a
--- /dev/null
+++ b/playground/frontend/web/favicon.png
Binary files differ
diff --git a/playground/frontend/web/icons/Icon-192.png b/playground/frontend/web/icons/Icon-192.png
new file mode 100644
index 0000000..b749bfe
--- /dev/null
+++ b/playground/frontend/web/icons/Icon-192.png
Binary files differ
diff --git a/playground/frontend/web/icons/Icon-512.png b/playground/frontend/web/icons/Icon-512.png
new file mode 100644
index 0000000..88cfd48
--- /dev/null
+++ b/playground/frontend/web/icons/Icon-512.png
Binary files differ
diff --git a/playground/frontend/web/icons/Icon-maskable-192.png b/playground/frontend/web/icons/Icon-maskable-192.png
new file mode 100644
index 0000000..eb9b4d7
--- /dev/null
+++ b/playground/frontend/web/icons/Icon-maskable-192.png
Binary files differ
diff --git a/playground/frontend/web/icons/Icon-maskable-512.png b/playground/frontend/web/icons/Icon-maskable-512.png
new file mode 100644
index 0000000..d69c566
--- /dev/null
+++ b/playground/frontend/web/icons/Icon-maskable-512.png
Binary files differ
diff --git a/playground/frontend/web/index.html b/playground/frontend/web/index.html
new file mode 100644
index 0000000..3273786
--- /dev/null
+++ b/playground/frontend/web/index.html
@@ -0,0 +1,120 @@
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+-->
+
+<!DOCTYPE html>
+<html>
+<head>
+ <!--
+ If you are serving your web app in a path other than the root, change the
+ href value below to reflect the base path you are serving from.
+
+ The path provided below has to start and end with a slash "/" in order for
+ it to work correctly.
+
+ For more details:
+ * https://developer.mozilla.org/en-US/docs/Web/HTML/Element/base
+
+ This is a placeholder for base href that will be replaced by the value of
+ the `--base-href` argument provided to `flutter build`.
+ -->
+ <base href="$FLUTTER_BASE_HREF">
+
+ <meta charset="UTF-8">
+ <meta content="IE=Edge" http-equiv="X-UA-Compatible">
+ <meta name="description" content="A new Flutter project.">
+
+ <!-- iOS meta tags & icons -->
+ <meta name="apple-mobile-web-app-capable" content="yes">
+ <meta name="apple-mobile-web-app-status-bar-style" content="black">
+ <meta name="apple-mobile-web-app-title" content="playground">
+ <link rel="apple-touch-icon" href="icons/Icon-192.png">
+
+ <title>playground</title>
+ <link rel="manifest" href="manifest.json">
+</head>
+<body>
+ <!-- This script installs service_worker.js to provide PWA functionality to
+ application. For more information, see:
+ https://developers.google.com/web/fundamentals/primers/service-workers -->
+ <script>
+ var serviceWorkerVersion = null;
+ var scriptLoaded = false;
+ function loadMainDartJs() {
+ if (scriptLoaded) {
+ return;
+ }
+ scriptLoaded = true;
+ var scriptTag = document.createElement('script');
+ scriptTag.src = 'main.dart.js';
+ scriptTag.type = 'application/javascript';
+ document.body.append(scriptTag);
+ }
+
+ if ('serviceWorker' in navigator) {
+ // Service workers are supported. Use them.
+ window.addEventListener('load', function () {
+ // Wait for registration to finish before dropping the <script> tag.
+ // Otherwise, the browser will load the script multiple times,
+ // potentially different versions.
+ var serviceWorkerUrl = 'flutter_service_worker.js?v=' + serviceWorkerVersion;
+ navigator.serviceWorker.register(serviceWorkerUrl)
+ .then((reg) => {
+ function waitForActivation(serviceWorker) {
+ serviceWorker.addEventListener('statechange', () => {
+ if (serviceWorker.state == 'activated') {
+ console.log('Installed new service worker.');
+ loadMainDartJs();
+ }
+ });
+ }
+ if (!reg.active && (reg.installing || reg.waiting)) {
+ // No active web worker and we have installed or are installing
+ // one for the first time. Simply wait for it to activate.
+ waitForActivation(reg.installing || reg.waiting);
+ } else if (!reg.active.scriptURL.endsWith(serviceWorkerVersion)) {
+ // When the app updates the serviceWorkerVersion changes, so we
+ // need to ask the service worker to update.
+ console.log('New service worker available.');
+ reg.update();
+ waitForActivation(reg.installing);
+ } else {
+ // Existing service worker is still good.
+ console.log('Loading app from service worker.');
+ loadMainDartJs();
+ }
+ });
+
+ // If service worker doesn't succeed in a reasonable amount of time,
+ // fallback to plaint <script> tag.
+ setTimeout(() => {
+ if (!scriptLoaded) {
+ console.warn(
+ 'Failed to load app from service worker. Falling back to plain <script> tag.',
+ );
+ loadMainDartJs();
+ }
+ }, 4000);
+ });
+ } else {
+ // Service workers not supported. Just drop the <script> tag.
+ loadMainDartJs();
+ }
+ </script>
+</body>
+</html>
diff --git a/playground/frontend/web/manifest.json b/playground/frontend/web/manifest.json
new file mode 100644
index 0000000..3274aac
--- /dev/null
+++ b/playground/frontend/web/manifest.json
@@ -0,0 +1,35 @@
+{
+ "name": "playground",
+ "short_name": "playground",
+ "start_url": ".",
+ "display": "standalone",
+ "background_color": "#0175C2",
+ "theme_color": "#0175C2",
+ "description": "A new Flutter project.",
+ "orientation": "portrait-primary",
+ "prefer_related_applications": false,
+ "icons": [
+ {
+ "src": "icons/Icon-192.png",
+ "sizes": "192x192",
+ "type": "image/png"
+ },
+ {
+ "src": "icons/Icon-512.png",
+ "sizes": "512x512",
+ "type": "image/png"
+ },
+ {
+ "src": "icons/Icon-maskable-192.png",
+ "sizes": "192x192",
+ "type": "image/png",
+ "purpose": "maskable"
+ },
+ {
+ "src": "icons/Icon-maskable-512.png",
+ "sizes": "512x512",
+ "type": "image/png",
+ "purpose": "maskable"
+ }
+ ]
+}
diff --git a/playground/playground/v1/playground.proto b/playground/playground/v1/playground.proto
new file mode 100644
index 0000000..7c26e00
--- /dev/null
+++ b/playground/playground/v1/playground.proto
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 = "proto3";
+
+option go_package = "github.com/apache/beam/playground/v1;playground";
+package playground.v1;
+
+enum Sdk {
+ SDK_UNSPECIFIED = 0;
+ SDK_JAVA = 1;
+ SDK_GO = 2;
+ SDK_PYTHON = 3;
+ SDK_SCIO = 4;
+}
+
+enum Status {
+ STATUS_UNSPECIFIED = 0;
+ STATUS_EXECUTING = 1;
+ STATUS_FINISHED = 2;
+ STATUS_ERROR = 3;
+}
+
+// RunCodeRequest represents a code text and options of SDK which executes the code.
+message RunCodeRequest {
+ string code = 1;
+ Sdk sdk = 2;
+}
+
+// RunCodeResponse contains information of the pipeline uuid.
+message RunCodeResponse {
+ string pipeline_uuid = 1;
+}
+
+// CheckStatusRequest contains information of the pipeline uuid.
+message CheckStatusRequest {
+ string pipeline_uuid = 1;
+}
+
+// StatusInfo contains information about the status of the code execution.
+message CheckStatusResponse {
+ Status status = 1;
+}
+
+// GetCompileOutputRequest contains information of the pipeline uuid.
+message GetCompileOutputRequest {
+ string pipeline_uuid = 1;
+}
+
+// GetCompileOutputResponse represents the result of the compiled code.
+message GetCompileOutputResponse {
+ string output = 1;
+ Status compilation_status = 2;
+}
+
+// GetRunOutputRequest contains information of the pipeline uuid.
+message GetRunOutputRequest {
+ string pipeline_uuid = 1;
+}
+
+// RunOutputResponse represents the result of the executed code.
+message GetRunOutputResponse {
+ string output = 1;
+ Status compilation_status = 2;
+}
+
+service PlaygroundService {
+
+ // Submit the job for an execution and get the pipeline uuid.
+ rpc RunCode(RunCodeRequest) returns (RunCodeResponse);
+
+ // Get the status of pipeline execution.
+ rpc CheckStatus(CheckStatusRequest) returns (CheckStatusResponse);
+
+ // Get the result of pipeline execution.
+ rpc GetRunOutput(GetRunOutputRequest) returns (GetRunOutputResponse);
+
+ // Get the result of pipeline compilation.
+ rpc GetCompileOutput(GetCompileOutputRequest) returns (GetCompileOutputResponse);
+}
\ No newline at end of file
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DisplayDataTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DisplayDataTranslation.java
index b776556..8677bf9 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DisplayDataTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DisplayDataTranslation.java
@@ -41,8 +41,7 @@
}
private static final Map<String, Function<DisplayData.Item, ByteString>>
- WELL_KNOWN_URN_TRANSLATORS =
- ImmutableMap.of(LABELLED, DisplayDataTranslation::translateStringUtf8);
+ WELL_KNOWN_URN_TRANSLATORS = ImmutableMap.of(LABELLED, DisplayDataTranslation::translate);
public static List<RunnerApi.DisplayData> toProto(DisplayData displayData) {
ImmutableList.Builder<RunnerApi.DisplayData> builder = ImmutableList.builder();
@@ -54,7 +53,7 @@
urn = item.getKey();
} else {
urn = LABELLED;
- translator = DisplayDataTranslation::translateStringUtf8;
+ translator = DisplayDataTranslation::translate;
}
builder.add(
RunnerApi.DisplayData.newBuilder()
@@ -65,13 +64,24 @@
return builder.build();
}
- private static ByteString translateStringUtf8(DisplayData.Item item) {
- String value = String.valueOf(item.getValue() == null ? item.getShortValue() : item.getValue());
+ private static ByteString translate(DisplayData.Item item) {
String label = item.getLabel() == null ? item.getKey() : item.getLabel();
- return RunnerApi.LabelledPayload.newBuilder()
- .setLabel(label)
- .setStringValue(value)
- .build()
- .toByteString();
+ String namespace = item.getNamespace() == null ? "" : item.getNamespace().getName();
+ RunnerApi.LabelledPayload.Builder builder =
+ RunnerApi.LabelledPayload.newBuilder()
+ .setKey(item.getKey())
+ .setLabel(label)
+ .setNamespace(namespace);
+ Object valueObj = item.getValue() == null ? item.getShortValue() : item.getValue();
+ if (valueObj instanceof Boolean) {
+ builder.setBoolValue((Boolean) valueObj);
+ } else if (valueObj instanceof Integer || valueObj instanceof Long) {
+ builder.setIntValue((Long) valueObj);
+ } else if (valueObj instanceof Number) {
+ builder.setDoubleValue(((Number) valueObj).doubleValue());
+ } else {
+ builder.setStringValue(String.valueOf(valueObj));
+ }
+ return builder.build().toByteString();
}
}
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DisplayDataTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DisplayDataTranslationTest.java
index 9fe9b4b..502e2fe 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DisplayDataTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DisplayDataTranslationTest.java
@@ -42,6 +42,11 @@
public void populateDisplayData(DisplayData.Builder builder) {
builder.add(DisplayData.item("foo", "value"));
builder.add(DisplayData.item("foo2", "value2").withLabel("label"));
+ builder.add(DisplayData.item("foo3", true).withLabel("label"));
+ builder.add(DisplayData.item("foo4", 123.4).withLabel("label"));
+ builder.add(DisplayData.item("foo5", 4.321f).withLabel("label"));
+ builder.add(DisplayData.item("foo6", 321).withLabel("label"));
+ builder.add(DisplayData.item("foo7", 123L).withLabel("label"));
}
})),
containsInAnyOrder(
@@ -51,6 +56,9 @@
RunnerApi.LabelledPayload.newBuilder()
.setLabel("foo")
.setStringValue("value")
+ .setKey("foo")
+ .setNamespace(
+ "org.apache.beam.runners.core.construction.DisplayDataTranslationTest$1")
.build()
.toByteString())
.build(),
@@ -60,6 +68,69 @@
RunnerApi.LabelledPayload.newBuilder()
.setLabel("label")
.setStringValue("value2")
+ .setKey("foo2")
+ .setNamespace(
+ "org.apache.beam.runners.core.construction.DisplayDataTranslationTest$1")
+ .build()
+ .toByteString())
+ .build(),
+ RunnerApi.DisplayData.newBuilder()
+ .setUrn(DisplayDataTranslation.LABELLED)
+ .setPayload(
+ RunnerApi.LabelledPayload.newBuilder()
+ .setLabel("label")
+ .setBoolValue(true)
+ .setKey("foo3")
+ .setNamespace(
+ "org.apache.beam.runners.core.construction.DisplayDataTranslationTest$1")
+ .build()
+ .toByteString())
+ .build(),
+ RunnerApi.DisplayData.newBuilder()
+ .setUrn(DisplayDataTranslation.LABELLED)
+ .setPayload(
+ RunnerApi.LabelledPayload.newBuilder()
+ .setLabel("label")
+ .setDoubleValue(123.4)
+ .setKey("foo4")
+ .setNamespace(
+ "org.apache.beam.runners.core.construction.DisplayDataTranslationTest$1")
+ .build()
+ .toByteString())
+ .build(),
+ RunnerApi.DisplayData.newBuilder()
+ .setUrn(DisplayDataTranslation.LABELLED)
+ .setPayload(
+ RunnerApi.LabelledPayload.newBuilder()
+ .setLabel("label")
+ .setDoubleValue(4.321000099182129)
+ .setKey("foo5")
+ .setNamespace(
+ "org.apache.beam.runners.core.construction.DisplayDataTranslationTest$1")
+ .build()
+ .toByteString())
+ .build(),
+ RunnerApi.DisplayData.newBuilder()
+ .setUrn(DisplayDataTranslation.LABELLED)
+ .setPayload(
+ RunnerApi.LabelledPayload.newBuilder()
+ .setLabel("label")
+ .setIntValue(321)
+ .setKey("foo6")
+ .setNamespace(
+ "org.apache.beam.runners.core.construction.DisplayDataTranslationTest$1")
+ .build()
+ .toByteString())
+ .build(),
+ RunnerApi.DisplayData.newBuilder()
+ .setUrn(DisplayDataTranslation.LABELLED)
+ .setPayload(
+ RunnerApi.LabelledPayload.newBuilder()
+ .setLabel("label")
+ .setIntValue(123)
+ .setKey("foo7")
+ .setNamespace(
+ "org.apache.beam.runners.core.construction.DisplayDataTranslationTest$1")
.build()
.toByteString())
.build()));
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SchemaTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SchemaTranslationTest.java
index 3b083f6..15ba6a4 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SchemaTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SchemaTranslationTest.java
@@ -18,7 +18,9 @@
package org.apache.beam.runners.core.construction;
import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.equalTo;
+import static org.junit.Assert.assertThrows;
import java.util.ArrayList;
import java.util.HashMap;
@@ -44,6 +46,7 @@
import org.junit.Test;
import org.junit.experimental.runners.Enclosed;
import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
import org.junit.runners.Parameterized;
import org.junit.runners.Parameterized.Parameter;
import org.junit.runners.Parameterized.Parameters;
@@ -284,4 +287,32 @@
assertThat(reencodedSchemaProto, equalTo(schemaProto));
}
}
+
+ /** Tests that we raise helpful errors when decoding bad {@link Schema} protos. */
+ @RunWith(JUnit4.class)
+ public static class DecodeErrorTest {
+
+ @Test
+ public void typeInfoNotSet() {
+ SchemaApi.Schema.Builder builder = SchemaApi.Schema.newBuilder();
+
+ builder.addFields(
+ SchemaApi.Field.newBuilder()
+ .setName("field_no_typeInfo")
+ .setType(SchemaApi.FieldType.newBuilder())
+ .setId(0)
+ .setEncodingPosition(0)
+ .build());
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> {
+ SchemaTranslation.schemaFromProto(builder.build());
+ });
+
+ assertThat(exception.getMessage(), containsString("field_no_typeInfo"));
+ assertThat(exception.getCause().getMessage(), containsString("TYPEINFO_NOT_SET"));
+ }
+ }
}
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/GcpResourceIdentifiers.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/GcpResourceIdentifiers.java
index 4c388bf..336f08d 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/GcpResourceIdentifiers.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/GcpResourceIdentifiers.java
@@ -51,4 +51,13 @@
return String.format(
"//bigtable.googleapis.com/projects/%s/namespaces/%s", projectId, namespace);
}
+
+ public static String spannerTable(String projectId, String databaseId, String tableId) {
+ return String.format(
+ "//spanner.googleapis.com/projects/%s/topics/%s/tables/%s", projectId, databaseId, tableId);
+ }
+
+ public static String spannerQuery(String projectId, String queryName) {
+ return String.format("//spanner.googleapis.com/projects/%s/queries/%s", projectId, queryName);
+ }
}
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MonitoringInfoConstants.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MonitoringInfoConstants.java
index c792719..03496fe 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MonitoringInfoConstants.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MonitoringInfoConstants.java
@@ -85,6 +85,10 @@
public static final String TABLE_ID = "TABLE_ID";
public static final String GCS_BUCKET = "GCS_BUCKET";
public static final String GCS_PROJECT_ID = "GCS_PROJECT_ID";
+ public static final String SPANNER_PROJECT_ID = "SPANNER_PROJECT_ID";
+ public static final String SPANNER_DATABASE_ID = "SPANNER_DATABASE_ID";
+ public static final String SPANNER_INSTANCE_ID = "SPANNER_INSTANCE_ID";
+ public static final String SPANNER_QUERY_NAME = "SPANNER_QUERY_NAME";
static {
// Note: One benefit of defining these strings above, instead of pulling them in from
@@ -120,6 +124,14 @@
checkArgument(TABLE_ID.equals(extractLabel(MonitoringInfoLabels.TABLE_ID)));
checkArgument(GCS_BUCKET.equals(extractLabel(MonitoringInfoLabels.GCS_BUCKET)));
checkArgument(GCS_PROJECT_ID.equals(extractLabel(MonitoringInfoLabels.GCS_PROJECT_ID)));
+ checkArgument(
+ SPANNER_PROJECT_ID.equals(extractLabel(MonitoringInfoLabels.SPANNER_PROJECT_ID)));
+ checkArgument(
+ SPANNER_DATABASE_ID.equals(extractLabel(MonitoringInfoLabels.SPANNER_DATABASE_ID)));
+ checkArgument(
+ SPANNER_INSTANCE_ID.equals(extractLabel(MonitoringInfoLabels.SPANNER_INSTANCE_ID)));
+ checkArgument(
+ SPANNER_QUERY_NAME.equals(extractLabel(MonitoringInfoLabels.SPANNER_QUERY_NAME)));
}
}
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java
index f0c16ff..5653d27 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineDebugOptions.java
@@ -230,6 +230,28 @@
void setReaderCacheTimeoutSec(Integer value);
+ /** The max amount of time an UnboundedReader is consumed before checkpointing. */
+ @Description(
+ "The max amount of time before an UnboundedReader is consumed before checkpointing, in seconds.")
+ @Default.Integer(10)
+ Integer getUnboundedReaderMaxReadTimeSec();
+
+ void setUnboundedReaderMaxReadTimeSec(Integer value);
+
+ /** The max elements read from an UnboundedReader before checkpointing. */
+ @Description("The max elements read from an UnboundedReader before checkpointing. ")
+ @Default.Integer(10 * 1000)
+ Integer getUnboundedReaderMaxElements();
+
+ void setUnboundedReaderMaxElements(Integer value);
+
+ /** The max amount of time waiting for elements when reading from UnboundedReader. */
+ @Description("The max amount of time waiting for elements when reading from UnboundedReader.")
+ @Default.Integer(1000)
+ Integer getUnboundedReaderMaxWaitForElementsMs();
+
+ void setUnboundedReaderMaxWaitForElementsMs(Integer value);
+
/**
* CAUTION: This option implies dumpHeapOnOOM, and has similar caveats. Specifically, heap dumps
* can of comparable size to the default boot disk. Consider increasing the boot disk size before
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternals.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternals.java
index 9ca72cc..ac59769d 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternals.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternals.java
@@ -70,6 +70,7 @@
import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.state.OrderedListState;
import org.apache.beam.sdk.state.ReadableState;
+import org.apache.beam.sdk.state.ReadableStates;
import org.apache.beam.sdk.state.SetState;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateContext;
@@ -1513,43 +1514,31 @@
@Override
public @UnknownKeyFor @NonNull @Initialized ReadableState<V> computeIfAbsent(
K key, Function<? super K, ? extends V> mappingFunction) {
- return new ReadableState<V>() {
- @Override
- public @Nullable V read() {
- Future<V> persistedData = getFutureForKey(key);
- try (Closeable scope = scopedReadState()) {
- if (localRemovals.contains(key) || negativeCache.contains(key)) {
- return null;
- }
- @Nullable V cachedValue = cachedValues.get(key);
- if (cachedValue != null || complete) {
- return cachedValue;
- }
-
- V persistedValue = persistedData.get();
- if (persistedValue == null) {
- // This is a new value. Add it to the map and return null.
- put(key, mappingFunction.apply(key));
- return null;
- }
- // TODO: Don't do this if it was already in cache.
- cachedValues.put(key, persistedValue);
- return persistedValue;
- } catch (InterruptedException | ExecutionException | IOException e) {
- if (e instanceof InterruptedException) {
- Thread.currentThread().interrupt();
- }
- throw new RuntimeException("Unable to read state", e);
- }
+ Future<V> persistedData = getFutureForKey(key);
+ try (Closeable scope = scopedReadState()) {
+ if (localRemovals.contains(key) || negativeCache.contains(key)) {
+ return ReadableStates.immediate(null);
+ }
+ @Nullable V cachedValue = cachedValues.get(key);
+ if (cachedValue != null || complete) {
+ return ReadableStates.immediate(cachedValue);
}
- @Override
- @SuppressWarnings("FutureReturnValueIgnored")
- public @UnknownKeyFor @NonNull @Initialized ReadableState<V> readLater() {
- WindmillMap.this.getFutureForKey(key);
- return this;
+ V persistedValue = persistedData.get();
+ if (persistedValue == null) {
+ // This is a new value. Add it to the map and return null.
+ put(key, mappingFunction.apply(key));
+ return ReadableStates.immediate(null);
}
- };
+ // TODO: Don't do this if it was already in cache.
+ cachedValues.put(key, persistedValue);
+ return ReadableStates.immediate(persistedValue);
+ } catch (InterruptedException | ExecutionException | IOException e) {
+ if (e instanceof InterruptedException) {
+ Thread.currentThread().interrupt();
+ }
+ throw new RuntimeException("Unable to read state", e);
+ }
}
@Override
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java
index 0c4a508..1236c61 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java
@@ -47,6 +47,7 @@
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.apache.beam.runners.dataflow.internal.CustomSources;
+import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.runners.dataflow.util.CloudObject;
import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader;
@@ -435,7 +436,7 @@
context.setActiveReader(reader);
- return new UnboundedReaderIterator<>(reader, context, started);
+ return new UnboundedReaderIterator<>(reader, context, started, options);
}
@Override
@@ -754,34 +755,34 @@
}
}
- // Commit at least once every 10 seconds or 10k records. This keeps the watermark advancing
- // smoothly, and ensures that not too much work will have to be reprocessed in the event of
- // a crash.
- @VisibleForTesting static int maxUnboundedBundleSize = 10000;
-
- @VisibleForTesting
- static final Duration MAX_UNBOUNDED_BUNDLE_READ_TIME = Duration.standardSeconds(10);
- // Backoff starting at 100ms, for approximately 1s total. 100+150+225+337.5~=1000.
- private static final FluentBackoff BACKOFF_FACTORY =
- FluentBackoff.DEFAULT.withMaxRetries(4).withInitialBackoff(Duration.millis(100));
-
private static class UnboundedReaderIterator<T>
extends NativeReader.NativeReaderIterator<WindowedValue<ValueWithRecordId<T>>> {
private final UnboundedSource.UnboundedReader<T> reader;
private final StreamingModeExecutionContext context;
private final boolean started;
private final Instant endTime;
- private int elemsRead;
+ private final int maxElems;
+ private final FluentBackoff backoffFactory;
+ private int elemsRead = 0;
private UnboundedReaderIterator(
UnboundedSource.UnboundedReader<T> reader,
StreamingModeExecutionContext context,
- boolean started) {
+ boolean started,
+ PipelineOptions options) {
this.reader = reader;
this.context = context;
- this.endTime = Instant.now().plus(MAX_UNBOUNDED_BUNDLE_READ_TIME);
- this.elemsRead = 0;
this.started = started;
+ DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class);
+ this.endTime =
+ Instant.now()
+ .plus(Duration.standardSeconds(debugOptions.getUnboundedReaderMaxReadTimeSec()));
+ this.maxElems = debugOptions.getUnboundedReaderMaxElements();
+ this.backoffFactory =
+ FluentBackoff.DEFAULT
+ .withInitialBackoff(Duration.millis(10))
+ .withMaxCumulativeBackoff(
+ Duration.millis(debugOptions.getUnboundedReaderMaxWaitForElementsMs()));
}
@Override
@@ -806,14 +807,16 @@
@Override
public boolean advance() throws IOException {
- if (elemsRead >= maxUnboundedBundleSize
- || Instant.now().isAfter(endTime)
- || context.isSinkFullHintSet()) {
- return false;
- }
-
- BackOff backoff = BACKOFF_FACTORY.backoff();
+ // Limits are placed on how much data we allow to return, how long we process the input
+ // before checkpointing and how long we block for input to be available. This ensures
+ // that there are regular checkpoints and that state does not become too large.
+ BackOff backoff = backoffFactory.backoff();
while (true) {
+ if (elemsRead >= maxElems
+ || Instant.now().isAfter(endTime)
+ || context.isSinkFullHintSet()) {
+ return false;
+ }
try {
if (reader.advance()) {
elemsRead++;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
index 17b59ec..5c62cf9 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
@@ -76,6 +76,7 @@
import org.apache.beam.runners.core.construction.SdkComponents;
import org.apache.beam.runners.core.construction.WindowingStrategyTranslation;
import org.apache.beam.runners.dataflow.internal.CustomSources;
+import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.runners.dataflow.util.CloudObject;
import org.apache.beam.runners.dataflow.util.CloudObjects;
@@ -2802,209 +2803,206 @@
@Test
public void testExceptionInvalidatesCache() throws Exception {
- DataflowPipelineOptions options =
- PipelineOptionsFactory.create().as(DataflowPipelineOptions.class);
- options.setNumWorkers(1);
-
// We'll need to force the system to limit bundles to one message at a time.
- int originalMaxUnboundedBundleSize = WorkerCustomSources.maxUnboundedBundleSize;
- WorkerCustomSources.maxUnboundedBundleSize = 1;
- try {
- // Sequence is as follows:
- // 01. GetWork[0] (token 0)
- // 02. Create counter reader
- // 03. Counter yields 0
- // 04. GetData[0] (state as null)
- // 05. Read state as null
- // 06. Set state as 42
- // 07. THROW on taking counter reader checkpoint
- // 08. Create counter reader
- // 09. Counter yields 0
- // 10. GetData[1] (state as null)
- // 11. Read state as null (*** not 42 ***)
- // 12. Take counter reader checkpoint as 0
- // 13. CommitWork[0] (message 0:0, state 42, checkpoint 0)
- // 14. GetWork[1] (token 1, checkpoint as 0)
- // 15. Counter yields 1
- // 16. Read (cached) state as 42
- // 17. Take counter reader checkpoint 1
- // 18. CommitWork[1] (message 0:1, checkpoint 1)
- // 19. GetWork[2] (token 2, checkpoint as 1)
- // 20. Counter yields 2
- // 21. THROW on processElement
- // 22. Recreate reader from checkpoint 1
- // 23. Counter yields 2 (*** not eof ***)
- // 24. GetData[2] (state as 42)
- // 25. Read state as 42
- // 26. Take counter reader checkpoint 2
- // 27. CommitWork[2] (message 0:2, checkpoint 2)
+ // Sequence is as follows:
+ // 01. GetWork[0] (token 0)
+ // 02. Create counter reader
+ // 03. Counter yields 0
+ // 04. GetData[0] (state as null)
+ // 05. Read state as null
+ // 06. Set state as 42
+ // 07. THROW on taking counter reader checkpoint
+ // 08. Create counter reader
+ // 09. Counter yields 0
+ // 10. GetData[1] (state as null)
+ // 11. Read state as null (*** not 42 ***)
+ // 12. Take counter reader checkpoint as 0
+ // 13. CommitWork[0] (message 0:0, state 42, checkpoint 0)
+ // 14. GetWork[1] (token 1, checkpoint as 0)
+ // 15. Counter yields 1
+ // 16. Read (cached) state as 42
+ // 17. Take counter reader checkpoint 1
+ // 18. CommitWork[1] (message 0:1, checkpoint 1)
+ // 19. GetWork[2] (token 2, checkpoint as 1)
+ // 20. Counter yields 2
+ // 21. THROW on processElement
+ // 22. Recreate reader from checkpoint 1
+ // 23. Counter yields 2 (*** not eof ***)
+ // 24. GetData[2] (state as 42)
+ // 25. Read state as 42
+ // 26. Take counter reader checkpoint 2
+ // 27. CommitWork[2] (message 0:2, checkpoint 2)
+ FakeWindmillServer server = new FakeWindmillServer(errorCollector);
+ server.setExpectedExceptionCount(2);
- CloudObject codec =
- CloudObjects.asCloudObject(
- WindowedValue.getFullCoder(
- ValueWithRecordId.ValueWithRecordIdCoder.of(
- KvCoder.of(VarIntCoder.of(), VarIntCoder.of())),
- GlobalWindow.Coder.INSTANCE),
- /*sdkComponents=*/ null);
+ DataflowPipelineOptions options = createTestingPipelineOptions(server);
+ options.setNumWorkers(1);
+ DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class);
+ debugOptions.setUnboundedReaderMaxElements(1);
- TestCountingSource counter = new TestCountingSource(3).withThrowOnFirstSnapshot(true);
- List<ParallelInstruction> instructions =
- Arrays.asList(
- new ParallelInstruction()
- .setOriginalName("OriginalReadName")
- .setSystemName("Read")
- .setName(DEFAULT_PARDO_USER_NAME)
- .setRead(
- new ReadInstruction()
- .setSource(
- CustomSources.serializeToCloudSource(counter, options)
- .setCodec(codec)))
- .setOutputs(
- Arrays.asList(
- new InstructionOutput()
- .setName("read_output")
- .setOriginalName(DEFAULT_OUTPUT_ORIGINAL_NAME)
- .setSystemName(DEFAULT_OUTPUT_SYSTEM_NAME)
- .setCodec(codec))),
- makeDoFnInstruction(
- new TestExceptionInvalidatesCacheFn(),
- 0,
- StringUtf8Coder.of(),
- WindowingStrategy.globalDefault()),
- makeSinkInstruction(StringUtf8Coder.of(), 1, GlobalWindow.Coder.INSTANCE));
+ CloudObject codec =
+ CloudObjects.asCloudObject(
+ WindowedValue.getFullCoder(
+ ValueWithRecordId.ValueWithRecordIdCoder.of(
+ KvCoder.of(VarIntCoder.of(), VarIntCoder.of())),
+ GlobalWindow.Coder.INSTANCE),
+ /*sdkComponents=*/ null);
- FakeWindmillServer server = new FakeWindmillServer(errorCollector);
- server.setExpectedExceptionCount(2);
- StreamingDataflowWorker worker =
- makeWorker(
- instructions, createTestingPipelineOptions(server), true /* publishCounters */);
- worker.setRetryLocallyDelayMs(100);
- worker.start();
+ TestCountingSource counter = new TestCountingSource(3).withThrowOnFirstSnapshot(true);
- // Three GetData requests
- for (int i = 0; i < 3; i++) {
- ByteString state;
- if (i == 0 || i == 1) {
- state = ByteString.EMPTY;
- } else {
- state = ByteString.copyFrom(new byte[] {42});
- }
- Windmill.GetDataResponse.Builder dataResponse = Windmill.GetDataResponse.newBuilder();
- dataResponse
- .addDataBuilder()
- .setComputationId(DEFAULT_COMPUTATION_ID)
- .addDataBuilder()
- .setKey(ByteString.copyFromUtf8("0000000000000001"))
- .setShardingKey(1)
- .addValuesBuilder()
- .setTag(ByteString.copyFromUtf8("//+uint"))
- .setStateFamily(DEFAULT_PARDO_STATE_FAMILY)
- .getValueBuilder()
- .setTimestamp(0)
- .setData(state);
- server.addDataToOffer(dataResponse.build());
+ List<ParallelInstruction> instructions =
+ Arrays.asList(
+ new ParallelInstruction()
+ .setOriginalName("OriginalReadName")
+ .setSystemName("Read")
+ .setName(DEFAULT_PARDO_USER_NAME)
+ .setRead(
+ new ReadInstruction()
+ .setSource(
+ CustomSources.serializeToCloudSource(counter, options).setCodec(codec)))
+ .setOutputs(
+ Arrays.asList(
+ new InstructionOutput()
+ .setName("read_output")
+ .setOriginalName(DEFAULT_OUTPUT_ORIGINAL_NAME)
+ .setSystemName(DEFAULT_OUTPUT_SYSTEM_NAME)
+ .setCodec(codec))),
+ makeDoFnInstruction(
+ new TestExceptionInvalidatesCacheFn(),
+ 0,
+ StringUtf8Coder.of(),
+ WindowingStrategy.globalDefault()),
+ makeSinkInstruction(StringUtf8Coder.of(), 1, GlobalWindow.Coder.INSTANCE));
+
+ StreamingDataflowWorker worker =
+ makeWorker(
+ instructions,
+ options.as(StreamingDataflowWorkerOptions.class),
+ true /* publishCounters */);
+ worker.setRetryLocallyDelayMs(100);
+ worker.start();
+
+ // Three GetData requests
+ for (int i = 0; i < 3; i++) {
+ ByteString state;
+ if (i == 0 || i == 1) {
+ state = ByteString.EMPTY;
+ } else {
+ state = ByteString.copyFrom(new byte[] {42});
}
+ Windmill.GetDataResponse.Builder dataResponse = Windmill.GetDataResponse.newBuilder();
+ dataResponse
+ .addDataBuilder()
+ .setComputationId(DEFAULT_COMPUTATION_ID)
+ .addDataBuilder()
+ .setKey(ByteString.copyFromUtf8("0000000000000001"))
+ .setShardingKey(1)
+ .addValuesBuilder()
+ .setTag(ByteString.copyFromUtf8("//+uint"))
+ .setStateFamily(DEFAULT_PARDO_STATE_FAMILY)
+ .getValueBuilder()
+ .setTimestamp(0)
+ .setData(state);
+ server.addDataToOffer(dataResponse.build());
+ }
- // Three GetWork requests and commits
- for (int i = 0; i < 3; i++) {
- StringBuilder sb = new StringBuilder();
- sb.append("work {\n");
- sb.append(" computation_id: \"computation\"\n");
- sb.append(" input_data_watermark: 0\n");
- sb.append(" work {\n");
- sb.append(" key: \"0000000000000001\"\n");
- sb.append(" sharding_key: 1\n");
- sb.append(" work_token: ");
- sb.append(i);
- sb.append(" cache_token: 1");
- sb.append("\n");
- if (i > 0) {
- int previousCheckpoint = i - 1;
- sb.append(" source_state {\n");
- sb.append(" state: \"");
- sb.append((char) previousCheckpoint);
- sb.append("\"\n");
- // We'll elide the finalize ids since it's not necessary to trigger the finalizer
- // for this test.
- sb.append(" }\n");
- }
- sb.append(" }\n");
- sb.append("}\n");
-
- server.addWorkToOffer(buildInput(sb.toString(), null));
-
- Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1);
-
- Windmill.WorkItemCommitRequest commit = result.get((long) i);
- UnsignedLong finalizeId =
- UnsignedLong.fromLongBits(commit.getSourceStateUpdates().getFinalizeIds(0));
-
- sb = new StringBuilder();
- sb.append("key: \"0000000000000001\"\n");
- sb.append("sharding_key: 1\n");
- sb.append("work_token: ");
- sb.append(i);
- sb.append("\n");
- sb.append("cache_token: 1\n");
- sb.append("output_messages {\n");
- sb.append(" destination_stream_id: \"out\"\n");
- sb.append(" bundles {\n");
- sb.append(" key: \"0000000000000001\"\n");
-
- int messageNum = i;
- sb.append(" messages {\n");
- sb.append(" timestamp: ");
- sb.append(messageNum * 1000);
- sb.append("\n");
- sb.append(" data: \"0:");
- sb.append(messageNum);
+ // Three GetWork requests and commits
+ for (int i = 0; i < 3; i++) {
+ StringBuilder sb = new StringBuilder();
+ sb.append("work {\n");
+ sb.append(" computation_id: \"computation\"\n");
+ sb.append(" input_data_watermark: 0\n");
+ sb.append(" work {\n");
+ sb.append(" key: \"0000000000000001\"\n");
+ sb.append(" sharding_key: 1\n");
+ sb.append(" work_token: ");
+ sb.append(i);
+ sb.append(" cache_token: 1");
+ sb.append("\n");
+ if (i > 0) {
+ int previousCheckpoint = i - 1;
+ sb.append(" source_state {\n");
+ sb.append(" state: \"");
+ sb.append((char) previousCheckpoint);
sb.append("\"\n");
+ // We'll elide the finalize ids since it's not necessary to trigger the finalizer
+ // for this test.
sb.append(" }\n");
-
- sb.append(" messages_ids: \"\"\n");
- sb.append(" }\n");
- sb.append("}\n");
- if (i == 0) {
- sb.append("value_updates {\n");
- sb.append(" tag: \"//+uint\"\n");
- sb.append(" value {\n");
- sb.append(" timestamp: 0\n");
- sb.append(" data: \"");
- sb.append((char) 42);
- sb.append("\"\n");
- sb.append(" }\n");
- sb.append(" state_family: \"parDoStateFamily\"\n");
- sb.append("}\n");
- }
-
- int sourceState = i;
- sb.append("source_state_updates {\n");
- sb.append(" state: \"");
- sb.append((char) sourceState);
- sb.append("\"\n");
- sb.append(" finalize_ids: ");
- sb.append(finalizeId);
- sb.append("}\n");
- sb.append("source_watermark: ");
- sb.append((sourceState + 1) * 1000);
- sb.append("\n");
- sb.append("source_backlog_bytes: 7\n");
-
- assertThat(
- // The commit will include a timer to clean up state - this timer is irrelevant
- // for the current test.
- setValuesTimestamps(commit.toBuilder().clearOutputTimers()).build(),
- equalTo(
- setMessagesMetadata(
- PaneInfo.NO_FIRING,
- CoderUtils.encodeToByteArray(
- CollectionCoder.of(GlobalWindow.Coder.INSTANCE),
- ImmutableList.of(GlobalWindow.INSTANCE)),
- parseCommitRequest(sb.toString()))
- .build()));
}
- } finally {
- WorkerCustomSources.maxUnboundedBundleSize = originalMaxUnboundedBundleSize;
+ sb.append(" }\n");
+ sb.append("}\n");
+
+ server.addWorkToOffer(buildInput(sb.toString(), null));
+
+ Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1);
+
+ Windmill.WorkItemCommitRequest commit = result.get((long) i);
+ UnsignedLong finalizeId =
+ UnsignedLong.fromLongBits(commit.getSourceStateUpdates().getFinalizeIds(0));
+
+ sb = new StringBuilder();
+ sb.append("key: \"0000000000000001\"\n");
+ sb.append("sharding_key: 1\n");
+ sb.append("work_token: ");
+ sb.append(i);
+ sb.append("\n");
+ sb.append("cache_token: 1\n");
+ sb.append("output_messages {\n");
+ sb.append(" destination_stream_id: \"out\"\n");
+ sb.append(" bundles {\n");
+ sb.append(" key: \"0000000000000001\"\n");
+
+ int messageNum = i;
+ sb.append(" messages {\n");
+ sb.append(" timestamp: ");
+ sb.append(messageNum * 1000);
+ sb.append("\n");
+ sb.append(" data: \"0:");
+ sb.append(messageNum);
+ sb.append("\"\n");
+ sb.append(" }\n");
+
+ sb.append(" messages_ids: \"\"\n");
+ sb.append(" }\n");
+ sb.append("}\n");
+ if (i == 0) {
+ sb.append("value_updates {\n");
+ sb.append(" tag: \"//+uint\"\n");
+ sb.append(" value {\n");
+ sb.append(" timestamp: 0\n");
+ sb.append(" data: \"");
+ sb.append((char) 42);
+ sb.append("\"\n");
+ sb.append(" }\n");
+ sb.append(" state_family: \"parDoStateFamily\"\n");
+ sb.append("}\n");
+ }
+
+ int sourceState = i;
+ sb.append("source_state_updates {\n");
+ sb.append(" state: \"");
+ sb.append((char) sourceState);
+ sb.append("\"\n");
+ sb.append(" finalize_ids: ");
+ sb.append(finalizeId);
+ sb.append("}\n");
+ sb.append("source_watermark: ");
+ sb.append((sourceState + 1) * 1000);
+ sb.append("\n");
+ sb.append("source_backlog_bytes: 7\n");
+
+ assertThat(
+ // The commit will include a timer to clean up state - this timer is irrelevant
+ // for the current test.
+ setValuesTimestamps(commit.toBuilder().clearOutputTimers()).build(),
+ equalTo(
+ setMessagesMetadata(
+ PaneInfo.NO_FIRING,
+ CoderUtils.encodeToByteArray(
+ CollectionCoder.of(GlobalWindow.Coder.INSTANCE),
+ ImmutableList.of(GlobalWindow.INSTANCE)),
+ parseCommitRequest(sb.toString()))
+ .build()));
}
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternalsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternalsTest.java
index 727a086..8ba5b3a 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternalsTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateInternalsTest.java
@@ -440,6 +440,69 @@
}
@Test
+ public void testMapPutIfAbsentNoReadSucceeds() throws Exception {
+ StateTag<MapState<String, Integer>> addr =
+ StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of());
+ MapState<String, Integer> mapState = underTest.state(NAMESPACE, addr);
+
+ final String tag1 = "tag1";
+ SettableFuture<Integer> future = SettableFuture.create();
+ when(mockReader.valueFuture(
+ protoKeyFromUserKey(tag1, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of()))
+ .thenReturn(future);
+ waitAndSet(future, null, 50);
+ ReadableState<Integer> readableState = mapState.putIfAbsent(tag1, 42);
+ assertEquals(42, (int) mapState.get(tag1).read());
+ assertNull(readableState.read());
+ }
+
+ @Test
+ public void testMapPutIfAbsentNoReadFails() throws Exception {
+ StateTag<MapState<String, Integer>> addr =
+ StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of());
+ MapState<String, Integer> mapState = underTest.state(NAMESPACE, addr);
+
+ final String tag1 = "tag1";
+ mapState.put(tag1, 1);
+ ReadableState<Integer> readableState = mapState.putIfAbsent(tag1, 42);
+ assertEquals(1, (int) mapState.get(tag1).read());
+ assertEquals(1, (int) readableState.read());
+
+ final String tag2 = "tag2";
+ SettableFuture<Integer> future = SettableFuture.create();
+ when(mockReader.valueFuture(
+ protoKeyFromUserKey(tag2, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of()))
+ .thenReturn(future);
+ waitAndSet(future, 2, 50);
+ readableState = mapState.putIfAbsent(tag2, 42);
+ assertEquals(2, (int) mapState.get(tag2).read());
+ assertEquals(2, (int) readableState.read());
+ }
+
+ @Test
+ public void testMapMultiplePutIfAbsentNoRead() throws Exception {
+ StateTag<MapState<String, Integer>> addr =
+ StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of());
+ MapState<String, Integer> mapState = underTest.state(NAMESPACE, addr);
+
+ final String tag1 = "tag1";
+ SettableFuture<Integer> future = SettableFuture.create();
+ when(mockReader.valueFuture(
+ protoKeyFromUserKey(tag1, StringUtf8Coder.of()), STATE_FAMILY, VarIntCoder.of()))
+ .thenReturn(future);
+ waitAndSet(future, null, 50);
+ ReadableState<Integer> readableState = mapState.putIfAbsent(tag1, 42);
+ assertEquals(42, (int) mapState.get(tag1).read());
+ ReadableState<Integer> readableState2 = mapState.putIfAbsent(tag1, 43);
+ mapState.put(tag1, 1);
+ ReadableState<Integer> readableState3 = mapState.putIfAbsent(tag1, 44);
+ assertEquals(1, (int) mapState.get(tag1).read());
+ assertNull(readableState.read());
+ assertEquals(42, (int) readableState2.read());
+ assertEquals(1, (int) readableState3.read());
+ }
+
+ @Test
public void testMapNegativeCache() throws Exception {
StateTag<MapState<String, Integer>> addr =
StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of());
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java
index 6806774..61c4ad7 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java
@@ -76,6 +76,7 @@
import org.apache.beam.runners.core.metrics.ExecutionStateSampler;
import org.apache.beam.runners.dataflow.DataflowPipelineTranslator;
import org.apache.beam.runners.dataflow.DataflowRunner;
+import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.runners.dataflow.util.CloudObject;
import org.apache.beam.runners.dataflow.util.PropertyNames;
@@ -519,9 +520,12 @@
Long.MAX_VALUE);
options.setNumWorkers(5);
+ int maxElements = 10;
+ DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class);
+ debugOptions.setUnboundedReaderMaxElements(maxElements);
ByteString state = ByteString.EMPTY;
- for (int i = 0; i < 10 * WorkerCustomSources.maxUnboundedBundleSize;
+ for (int i = 0; i < 10 * maxElements;
/* Incremented in inner loop */ ) {
// Initialize streaming context with state from previous iteration.
context.start(
@@ -565,12 +569,12 @@
numReadOnThisIteration++;
}
Instant afterReading = Instant.now();
+ long maxReadSec = debugOptions.getUnboundedReaderMaxReadTimeSec();
assertThat(
new Duration(beforeReading, afterReading).getStandardSeconds(),
- lessThanOrEqualTo(
- WorkerCustomSources.MAX_UNBOUNDED_BUNDLE_READ_TIME.getStandardSeconds() + 1));
+ lessThanOrEqualTo(maxReadSec + 1));
assertThat(
- numReadOnThisIteration, lessThanOrEqualTo(WorkerCustomSources.maxUnboundedBundleSize));
+ numReadOnThisIteration, lessThanOrEqualTo(debugOptions.getUnboundedReaderMaxElements()));
// Extract and verify state modifications.
context.flushState();
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineRunner.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineRunner.java
index bbdd1f5..6bfec1e 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineRunner.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaPipelineRunner.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.samza;
import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.graph.ExecutableStage;
import org.apache.beam.runners.core.construction.graph.GreedyPipelineFuser;
@@ -41,16 +40,16 @@
private final SamzaPipelineOptions options;
@Override
- public PortablePipelineResult run(final Pipeline pipeline, JobInfo jobInfo) {
+ public PortablePipelineResult run(final RunnerApi.Pipeline pipeline, JobInfo jobInfo) {
// Expand any splittable DoFns within the graph to enable sizing and splitting of bundles.
- Pipeline pipelineWithSdfExpanded =
+ RunnerApi.Pipeline pipelineWithSdfExpanded =
ProtoOverrides.updateTransform(
PTransformTranslation.PAR_DO_TRANSFORM_URN,
pipeline,
SplittableParDoExpander.createSizedReplacement());
// Don't let the fuser fuse any subcomponents of native transforms.
- Pipeline trimmedPipeline =
+ RunnerApi.Pipeline trimmedPipeline =
TrivialNativeTransformExpander.forKnownUrns(
pipelineWithSdfExpanded, SamzaPortablePipelineTranslator.knownUrns());
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/ConfigBuilder.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/ConfigBuilder.java
index 08bef16..4b5989a 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/ConfigBuilder.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/ConfigBuilder.java
@@ -190,10 +190,15 @@
|| BeamJobCoordinatorRunner.class.getName().equals(appRunner)
|| RemoteApplicationRunner.class.getName().equals(appRunner)
|| BeamContainerRunner.class.getName().equals(appRunner),
- "Config %s must be set to %s for %s Deployment",
+ "Config %s must be set to %s for %s Deployment, but found %s",
APP_RUNNER_CLASS,
- RemoteApplicationRunner.class.getName(),
- SamzaExecutionEnvironment.YARN);
+ String.format(
+ "[%s, %s or %s]",
+ BeamJobCoordinatorRunner.class.getName(),
+ RemoteApplicationRunner.class.getName(),
+ BeamContainerRunner.class.getName()),
+ SamzaExecutionEnvironment.YARN,
+ appRunner);
checkArgument(
config.containsKey(JOB_FACTORY_CLASS),
"Config %s not found for %s Deployment",
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/ReshuffleTranslator.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/ReshuffleTranslator.java
new file mode 100644
index 0000000..c7f8acc
--- /dev/null
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/ReshuffleTranslator.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 org.apache.beam.runners.samza.translation;
+
+import com.google.auto.service.AutoService;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.core.construction.NativeTransforms;
+import org.apache.beam.runners.core.construction.graph.PipelineNode;
+import org.apache.beam.runners.core.construction.graph.QueryablePipeline;
+import org.apache.beam.runners.samza.runtime.OpMessage;
+import org.apache.beam.runners.samza.util.SamzaCoders;
+import org.apache.beam.runners.samza.util.WindowUtils;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.samza.operators.MessageStream;
+import org.apache.samza.serializers.KVSerde;
+
+/**
+ * Translates Reshuffle transform into Samza's native partitionBy operator, which will partition
+ * each incoming message by the key into a Task corresponding to that key.
+ */
+public class ReshuffleTranslator<K, InT, OutT>
+ implements TransformTranslator<PTransform<PCollection<KV<K, InT>>, PCollection<KV<K, OutT>>>> {
+
+ @Override
+ public void translate(
+ PTransform<PCollection<KV<K, InT>>, PCollection<KV<K, OutT>>> transform,
+ TransformHierarchy.Node node,
+ TranslationContext ctx) {
+
+ final PCollection<KV<K, InT>> input = ctx.getInput(transform);
+ final PCollection<KV<K, OutT>> output = ctx.getOutput(transform);
+ final MessageStream<OpMessage<KV<K, InT>>> inputStream = ctx.getMessageStream(input);
+ // input will be OpMessage of Windowed<KV<K, Iterable<V>>>
+ final KvCoder<K, InT> inputCoder = (KvCoder<K, InT>) input.getCoder();
+ final Coder<WindowedValue<KV<K, InT>>> elementCoder = SamzaCoders.of(input);
+
+ final MessageStream<OpMessage<KV<K, InT>>> outputStream =
+ doTranslate(
+ inputStream,
+ inputCoder.getKeyCoder(),
+ elementCoder,
+ "rshfl-" + ctx.getTransformId(),
+ ctx.getPipelineOptions().getMaxSourceParallelism() > 1);
+
+ ctx.registerMessageStream(output, outputStream);
+ }
+
+ @Override
+ public void translatePortable(
+ PipelineNode.PTransformNode transform,
+ QueryablePipeline pipeline,
+ PortableTranslationContext ctx) {
+
+ final String inputId = ctx.getInputId(transform);
+ final MessageStream<OpMessage<KV<K, InT>>> inputStream = ctx.getMessageStreamById(inputId);
+ final WindowedValue.WindowedValueCoder<KV<K, InT>> windowedInputCoder =
+ WindowUtils.instantiateWindowedCoder(inputId, pipeline.getComponents());
+ final String outputId = ctx.getOutputId(transform);
+
+ final MessageStream<OpMessage<KV<K, InT>>> outputStream =
+ doTranslate(
+ inputStream,
+ ((KvCoder<K, InT>) windowedInputCoder.getValueCoder()).getKeyCoder(),
+ windowedInputCoder,
+ "rshfl-" + ctx.getTransformId(),
+ ctx.getSamzaPipelineOptions().getMaxSourceParallelism() > 1);
+
+ ctx.registerMessageStream(outputId, outputStream);
+ }
+
+ private static <K, InT> MessageStream<OpMessage<KV<K, InT>>> doTranslate(
+ MessageStream<OpMessage<KV<K, InT>>> inputStream,
+ Coder<K> keyCoder,
+ Coder<WindowedValue<KV<K, InT>>> valueCoder,
+ String partitionById, // will be used in the intermediate stream name
+ boolean needRepartition) {
+
+ return needRepartition
+ ? inputStream
+ .filter(op -> OpMessage.Type.ELEMENT == op.getType())
+ .partitionBy(
+ opMessage -> opMessage.getElement().getValue().getKey(),
+ OpMessage::getElement, // windowed value
+ KVSerde.of(SamzaCoders.toSerde(keyCoder), SamzaCoders.toSerde(valueCoder)),
+ partitionById)
+ // convert back to OpMessage
+ .map(kv -> OpMessage.ofElement(kv.getValue()))
+ : inputStream.filter(op -> OpMessage.Type.ELEMENT == op.getType());
+ }
+
+ /** Predicate to determine whether a URN is a Samza native transform. */
+ @AutoService(NativeTransforms.IsNativeTransform.class)
+ public static class IsSamzaNativeTransform implements NativeTransforms.IsNativeTransform {
+ @Override
+ public boolean test(RunnerApi.PTransform pTransform) {
+ return false;
+ // Re-enable after BEAM-12999 is completed
+ // return PTransformTranslation.RESHUFFLE_URN.equals(
+ // PTransformTranslation.urnForTransformOrNull(pTransform));
+ }
+ }
+}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPipelineTranslator.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPipelineTranslator.java
index 3514b4f..1dd779d 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPipelineTranslator.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPipelineTranslator.java
@@ -33,8 +33,6 @@
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
/** This class knows all the translators from a primitive BEAM transform to a Samza operator. */
@SuppressWarnings({
@@ -42,7 +40,6 @@
"nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
})
public class SamzaPipelineTranslator {
- private static final Logger LOG = LoggerFactory.getLogger(SamzaPipelineTranslator.class);
private static final Map<String, TransformTranslator<?>> TRANSLATORS = loadTranslators();
@@ -117,7 +114,7 @@
}
private static class SamzaPipelineVisitor extends Pipeline.PipelineVisitor.Defaults {
- private TransformVisitorFn visitorFn;
+ private final TransformVisitorFn visitorFn;
private SamzaPipelineVisitor(TransformVisitorFn visitorFn) {
this.visitorFn = visitorFn;
@@ -177,6 +174,7 @@
public Map<String, TransformTranslator<?>> getTransformTranslators() {
return ImmutableMap.<String, TransformTranslator<?>>builder()
.put(PTransformTranslation.READ_TRANSFORM_URN, new ReadTranslator<>())
+ .put(PTransformTranslation.RESHUFFLE_URN, new ReshuffleTranslator<>())
.put(PTransformTranslation.PAR_DO_TRANSFORM_URN, new ParDoBoundMultiTranslator<>())
.put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator<>())
.put(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN, new GroupByKeyTranslator<>())
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPortablePipelineTranslator.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPortablePipelineTranslator.java
index 9158cd4..e634be6 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPortablePipelineTranslator.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/SamzaPortablePipelineTranslator.java
@@ -103,6 +103,8 @@
@Override
public Map<String, TransformTranslator<?>> getTransformTranslators() {
return ImmutableMap.<String, TransformTranslator<?>>builder()
+ // Re-enable after BEAM-12999 is completed
+ // .put(PTransformTranslation.RESHUFFLE_URN, new ReshuffleTranslator<>())
.put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator<>())
.put(PTransformTranslation.FLATTEN_TRANSFORM_URN, new FlattenPCollectionsTranslator<>())
.put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator())
diff --git a/sdks/go.mod b/sdks/go.mod
index 9cb3c90..2c7238f 100644
--- a/sdks/go.mod
+++ b/sdks/go.mod
@@ -30,17 +30,14 @@
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/protobuf v1.5.2 // TODO(danoliveira): Fully replace this with google.golang.org/protobuf
github.com/golang/snappy v0.0.4 // indirect
- github.com/google/btree v1.0.0 // indirect
github.com/google/go-cmp v0.5.6
github.com/google/martian/v3 v3.2.1 // indirect
github.com/google/uuid v1.3.0
- github.com/hashicorp/golang-lru v0.5.1 // indirect
github.com/kr/text v0.2.0 // indirect
github.com/linkedin/goavro v2.1.0+incompatible
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
github.com/nightlyone/lockfile v1.0.0
github.com/spf13/cobra v1.2.1
- golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6 // indirect
golang.org/x/net v0.0.0-20210423184538-5f58ad60dda6
golang.org/x/oauth2 v0.0.0-20210628180205-a41e5a781914
golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c // indirect
diff --git a/sdks/go/examples/snippets/04transforms.go b/sdks/go/examples/snippets/04transforms.go
index 4f07f04..d920acd 100644
--- a/sdks/go/examples/snippets/04transforms.go
+++ b/sdks/go/examples/snippets/04transforms.go
@@ -300,7 +300,7 @@
var cutOff float64
ok := lengthCutOffIter(&cutOff)
if !ok {
- return fmt.Errorf("No length cutoff provided.")
+ return fmt.Errorf("no length cutoff provided")
}
if float64(len(word)) > cutOff {
emitAboveCutoff(word)
diff --git a/sdks/go/examples/snippets/06schemas.go b/sdks/go/examples/snippets/06schemas.go
new file mode 100644
index 0000000..d68bc7c
--- /dev/null
+++ b/sdks/go/examples/snippets/06schemas.go
@@ -0,0 +1,143 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 snippets
+
+import (
+ "fmt"
+ "io"
+ "reflect"
+ "time"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
+)
+
+// [START schema_define]
+
+type Purchase struct {
+ // ID of the user who made the purchase.
+ UserID string `beam:"userId"`
+ // Identifier of the item that was purchased.
+ ItemID int64 `beam:"itemId"`
+ // The shipping address, a nested type.
+ ShippingAddress ShippingAddress `beam:"shippingAddress"`
+ // The cost of the item in cents.
+ Cost int64 `beam:"cost"`
+ // The transactions that paid for this purchase.
+ // A slice since the purchase might be spread out over multiple
+ // credit cards.
+ Transactions []Transaction `beam:"transactions"`
+}
+
+type ShippingAddress struct {
+ StreetAddress string `beam:"streetAddress"`
+ City string `beam:"city"`
+ State *string `beam:"state"`
+ Country string `beam:"country"`
+ PostCode string `beam:"postCode"`
+}
+
+type Transaction struct {
+ Bank string `beam:"bank"`
+ PurchaseAmount float64 `beam:"purchaseAmount"`
+}
+
+// [END schema_define]
+
+// Validate that the interface is being implemented.
+var _ beam.SchemaProvider = &TimestampNanosProvider{}
+
+// [START schema_logical_provider]
+
+// TimestampNanos is a logical type using time.Time, but
+// encodes as a schema type.
+type TimestampNanos time.Time
+
+func (tn TimestampNanos) Seconds() int64 {
+ return time.Time(tn).Unix()
+}
+func (tn TimestampNanos) Nanos() int32 {
+ return int32(time.Time(tn).UnixNano() % 1000000000)
+}
+
+// tnStorage is the storage schema for TimestampNanos.
+type tnStorage struct {
+ Seconds int64 `beam:"seconds"`
+ Nanos int32 `beam:"nanos"`
+}
+
+var (
+ // reflect.Type of the Value type of TimestampNanos
+ tnType = reflect.TypeOf((*TimestampNanos)(nil)).Elem()
+ tnStorageType = reflect.TypeOf((*tnStorage)(nil)).Elem()
+)
+
+// TimestampNanosProvider implements the beam.SchemaProvider interface.
+type TimestampNanosProvider struct{}
+
+// FromLogicalType converts checks if the given type is TimestampNanos, and if so
+// returns the storage type.
+func (p *TimestampNanosProvider) FromLogicalType(rt reflect.Type) (reflect.Type, error) {
+ if rt != tnType {
+ return nil, fmt.Errorf("unable to provide schema.LogicalType for type %v, want %v", rt, tnType)
+ }
+ return tnStorageType, nil
+}
+
+// BuildEncoder builds a Beam schema encoder for the TimestampNanos type.
+func (p *TimestampNanosProvider) BuildEncoder(rt reflect.Type) (func(interface{}, io.Writer) error, error) {
+ if _, err := p.FromLogicalType(rt); err != nil {
+ return nil, err
+ }
+ enc, err := coder.RowEncoderForStruct(tnStorageType)
+ if err != nil {
+ return nil, err
+ }
+ return func(iface interface{}, w io.Writer) error {
+ v := iface.(TimestampNanos)
+ return enc(tnStorage{
+ Seconds: v.Seconds(),
+ Nanos: v.Nanos(),
+ }, w)
+ }, nil
+}
+
+// BuildDecoder builds a Beam schema decoder for the TimestampNanos type.
+func (p *TimestampNanosProvider) BuildDecoder(rt reflect.Type) (func(io.Reader) (interface{}, error), error) {
+ if _, err := p.FromLogicalType(rt); err != nil {
+ return nil, err
+ }
+ dec, err := coder.RowDecoderForStruct(tnStorageType)
+ if err != nil {
+ return nil, err
+ }
+ return func(r io.Reader) (interface{}, error) {
+ s, err := dec(r)
+ if err != nil {
+ return nil, err
+ }
+ tn := s.(tnStorage)
+ return TimestampNanos(time.Unix(tn.Seconds, int64(tn.Nanos))), nil
+ }, nil
+}
+
+// [END schema_logical_provider]
+
+func LogicalTypeExample() {
+ // [START schema_logical_register]
+ beam.RegisterSchemaProvider(tnType, &TimestampNanosProvider{})
+ // [END schema_logical_register]
+}
diff --git a/sdks/go/examples/snippets/06schemas_test.go b/sdks/go/examples/snippets/06schemas_test.go
new file mode 100644
index 0000000..1353c97
--- /dev/null
+++ b/sdks/go/examples/snippets/06schemas_test.go
@@ -0,0 +1,170 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 snippets
+
+import (
+ "fmt"
+ "reflect"
+ "testing"
+ "time"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder/testutil"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/schema"
+ pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
+ "github.com/google/go-cmp/cmp"
+ "google.golang.org/protobuf/proto"
+ "google.golang.org/protobuf/testing/protocmp"
+)
+
+func atomicSchemaField(name string, typ pipepb.AtomicType) *pipepb.Field {
+ return &pipepb.Field{
+ Name: name,
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: typ,
+ },
+ },
+ }
+}
+
+func rowSchemaField(name string, typ *pipepb.Schema) *pipepb.Field {
+ return &pipepb.Field{
+ Name: name,
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_RowType{
+ RowType: &pipepb.RowType{
+ Schema: typ,
+ },
+ },
+ },
+ }
+}
+
+func listSchemaField(name string, typ *pipepb.Field) *pipepb.Field {
+ return &pipepb.Field{
+ Name: name,
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_ArrayType{
+ ArrayType: &pipepb.ArrayType{
+ ElementType: typ.GetType(),
+ },
+ },
+ },
+ }
+}
+
+func nillable(f *pipepb.Field) *pipepb.Field {
+ f.Type.Nullable = true
+ return f
+}
+
+func TestSchemaTypes(t *testing.T) {
+ transactionSchema := &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ atomicSchemaField("bank", pipepb.AtomicType_STRING),
+ atomicSchemaField("purchaseAmount", pipepb.AtomicType_DOUBLE),
+ },
+ }
+ shippingAddressSchema := &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ atomicSchemaField("streetAddress", pipepb.AtomicType_STRING),
+ atomicSchemaField("city", pipepb.AtomicType_STRING),
+ nillable(atomicSchemaField("state", pipepb.AtomicType_STRING)),
+ atomicSchemaField("country", pipepb.AtomicType_STRING),
+ atomicSchemaField("postCode", pipepb.AtomicType_STRING),
+ },
+ }
+
+ tests := []struct {
+ rt reflect.Type
+ st *pipepb.Schema
+ preReg func(reg *schema.Registry)
+ }{{
+ rt: reflect.TypeOf(Transaction{}),
+ st: transactionSchema,
+ }, {
+ rt: reflect.TypeOf(ShippingAddress{}),
+ st: shippingAddressSchema,
+ }, {
+ rt: reflect.TypeOf(Purchase{}),
+ st: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ atomicSchemaField("userId", pipepb.AtomicType_STRING),
+ atomicSchemaField("itemId", pipepb.AtomicType_INT64),
+ rowSchemaField("shippingAddress", shippingAddressSchema),
+ atomicSchemaField("cost", pipepb.AtomicType_INT64),
+ listSchemaField("transactions",
+ rowSchemaField("n/a", transactionSchema)),
+ },
+ },
+ }, {
+ rt: tnType,
+ st: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ atomicSchemaField("seconds", pipepb.AtomicType_INT64),
+ atomicSchemaField("nanos", pipepb.AtomicType_INT32),
+ },
+ },
+ preReg: func(reg *schema.Registry) {
+ reg.RegisterLogicalType(schema.ToLogicalType(tnType.Name(), tnType, tnStorageType))
+ },
+ }}
+ for _, test := range tests {
+ t.Run(fmt.Sprintf("%v", test.rt), func(t *testing.T) {
+ reg := schema.NewRegistry()
+ if test.preReg != nil {
+ test.preReg(reg)
+ }
+ {
+ got, err := reg.FromType(test.rt)
+ if err != nil {
+ t.Fatalf("error FromType(%v) = %v", test.rt, err)
+ }
+ if d := cmp.Diff(test.st, got,
+ protocmp.Transform(),
+ protocmp.IgnoreFields(proto.Message(&pipepb.Schema{}), "id", "options"),
+ ); d != "" {
+ t.Errorf("diff (-want, +got): %v", d)
+ }
+ }
+ })
+ }
+}
+
+func TestSchema_validate(t *testing.T) {
+ tests := []struct {
+ rt reflect.Type
+ p beam.SchemaProvider
+ logical, storage interface{}
+ }{
+ {
+ rt: tnType,
+ p: &TimestampNanosProvider{},
+ logical: TimestampNanos(time.Unix(2300003, 456789)),
+ storage: tnStorage{},
+ },
+ }
+ for _, test := range tests {
+ sc := &testutil.SchemaCoder{
+ CmpOptions: cmp.Options{
+ cmp.Comparer(func(a, b TimestampNanos) bool {
+ return a.Seconds() == b.Seconds() && a.Nanos() == b.Nanos()
+ })},
+ }
+ sc.Validate(t, test.rt, test.p.BuildEncoder, test.p.BuildDecoder, test.storage, test.logical)
+ }
+}
diff --git a/sdks/go/examples/snippets/08windowing.go b/sdks/go/examples/snippets/08windowing.go
new file mode 100644
index 0000000..8e8e1b6
--- /dev/null
+++ b/sdks/go/examples/snippets/08windowing.go
@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 snippets
+
+import (
+ "time"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
+)
+
+func settingWindows(s beam.Scope, items beam.PCollection) {
+ // [START setting_fixed_windows]
+ fixedWindowedItems := beam.WindowInto(s,
+ window.NewFixedWindows(60*time.Second),
+ items)
+ // [END setting_fixed_windows]
+
+ // [START setting_sliding_windows]
+ slidingWindowedItems := beam.WindowInto(s,
+ window.NewSlidingWindows(5*time.Second, 30*time.Second),
+ items)
+ // [END setting_sliding_windows]
+
+ // [START setting_session_windows]
+ sessionWindowedItems := beam.WindowInto(s,
+ window.NewSessions(600*time.Second),
+ items)
+ // [END setting_session_windows]
+
+ // [START setting_global_window]
+ globalWindowedItems := beam.WindowInto(s,
+ window.NewGlobalWindows(),
+ items)
+ // [END setting_global_window]
+
+ // [START setting_allowed_lateness]
+ windowedItems := beam.WindowInto(s,
+ window.NewFixedWindows(1*time.Minute), items,
+ beam.AllowedLateness(2*24*time.Hour), // 2 days
+ )
+ // [END setting_allowed_lateness]
+
+ _ = []beam.PCollection{
+ fixedWindowedItems,
+ slidingWindowedItems,
+ sessionWindowedItems,
+ globalWindowedItems,
+ windowedItems,
+ }
+}
+
+// LogEntry is a dummy type for documentation purposes.
+type LogEntry int
+
+func extractEventTime(LogEntry) time.Time {
+ // Note: Returning time.Now() is always going to be processing time
+ // not EventTime. For true event time, one needs to extract the
+ // time from the element itself.
+ return time.Now()
+}
+
+// [START setting_timestamp]
+
+// AddTimestampDoFn extracts an event time from a LogEntry.
+func AddTimestampDoFn(element LogEntry, emit func(beam.EventTime, LogEntry)) {
+ et := extractEventTime(element)
+ // Defining an emitter with beam.EventTime as the first parameter
+ // allows the DoFn to set the event time for the emitted element.
+ emit(mtime.FromTime(et), element)
+}
+
+// [END setting_timestamp]
+
+func timestampedCollection(s beam.Scope, unstampedLogs beam.PCollection) {
+ // [START setting_timestamp_pipeline]
+ stampedLogs := beam.ParDo(s, AddTimestampDoFn, unstampedLogs)
+ // [END setting_timestamp_pipeline]
+ _ = stampedLogs
+}
diff --git a/sdks/go/examples/snippets/09triggers.go b/sdks/go/examples/snippets/09triggers.go
new file mode 100644
index 0000000..0af6d84
--- /dev/null
+++ b/sdks/go/examples/snippets/09triggers.go
@@ -0,0 +1,98 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 snippets contains code used in the Beam Programming Guide
+// as examples for the Apache Beam Go SDK. These snippets are compiled
+// and their tests run to ensure correctness. However, due to their
+// piecemeal pedagogical use, they may not be the best example of
+// production code.
+//
+// The Beam Programming Guide can be found at https://beam.apache.org/documentation/programming-guide/.
+package snippets
+
+import (
+ "time"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
+)
+
+func TriggerAfterEndOfWindow(s beam.Scope, pCollection beam.PCollection) {
+ // [START after_window_trigger]
+ trigger := window.TriggerAfterEndOfWindow().
+ EarlyFiring(window.TriggerAfterProcessingTime().
+ PlusDelay(60 * time.Second)).
+ LateFiring(window.TriggerRepeat(window.TriggerAfterCount(1)))
+ // [END after_window_trigger]
+ beam.WindowInto(s, window.NewFixedWindows(10*time.Second), pCollection, beam.Trigger(trigger), beam.PanesDiscard())
+}
+
+func TriggerAlways(s beam.Scope, pCollection beam.PCollection) {
+ // [START always_trigger]
+ beam.WindowInto(s, window.NewFixedWindows(10*time.Second), pCollection,
+ beam.Trigger(window.TriggerAlways()),
+ beam.PanesDiscard(),
+ )
+ // [END always_trigger]
+}
+
+func ComplexTriggers(s beam.Scope, pcollection beam.PCollection) {
+ // [START setting_a_trigger]
+ windowedItems := beam.WindowInto(s,
+ window.NewFixedWindows(1*time.Minute), pcollection,
+ beam.Trigger(window.TriggerAfterProcessingTime().
+ PlusDelay(1*time.Minute)),
+ beam.AllowedLateness(30*time.Minute),
+ beam.PanesDiscard(),
+ )
+ // [END setting_a_trigger]
+
+ // [START setting_allowed_lateness]
+ allowedToBeLateItems := beam.WindowInto(s,
+ window.NewFixedWindows(1*time.Minute), pcollection,
+ beam.Trigger(window.TriggerAfterProcessingTime().
+ PlusDelay(1*time.Minute)),
+ beam.AllowedLateness(30*time.Minute),
+ )
+ // [END setting_allowed_lateness]
+
+ // [START model_composite_triggers]
+ compositeTriggerItems := beam.WindowInto(s,
+ window.NewFixedWindows(1*time.Minute), pcollection,
+ beam.Trigger(window.TriggerAfterEndOfWindow().
+ LateFiring(window.TriggerAfterProcessingTime().
+ PlusDelay(10*time.Minute))),
+ beam.AllowedLateness(2*24*time.Hour),
+ )
+ // [END model_composite_triggers]
+
+ // TODO(BEAM-3304) AfterAny is not yet implemented.
+ // Implement so the following compiles when no longer commented out.
+
+ // [START other_composite_trigger]
+ // beam.Trigger(
+ // window.TriggerAfterAny(
+ // window.TriggerAfterCount(100),
+ // window.TriggerAfterProcessingTime().
+ // PlusDelay(1*time.Minute)),
+ // )
+ // [END other_composite_trigger]
+
+ _ = []beam.PCollection{
+ windowedItems,
+ allowedToBeLateItems,
+ compositeTriggerItems,
+ }
+}
diff --git a/sdks/go/examples/snippets/10metrics.go b/sdks/go/examples/snippets/10metrics.go
new file mode 100644
index 0000000..c16e95a
--- /dev/null
+++ b/sdks/go/examples/snippets/10metrics.go
@@ -0,0 +1,84 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 snippets
+
+import (
+ "context"
+ "fmt"
+ "reflect"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/metrics"
+)
+
+// [START metrics_query]
+
+func queryMetrics(pr beam.PipelineResult, ns, n string) metrics.QueryResults {
+ return pr.Metrics().Query(func(r metrics.SingleResult) bool {
+ return r.Namespace() == ns && r.Name() == n
+ })
+}
+
+// [END metrics_query]
+
+var runner = "direct"
+
+// [START metrics_pipeline]
+
+func addMetricDoFnToPipeline(s beam.Scope, input beam.PCollection) beam.PCollection {
+ return beam.ParDo(s, &MyMetricsDoFn{}, input)
+}
+
+func executePipelineAndGetMetrics(ctx context.Context, p *beam.Pipeline) (metrics.QueryResults, error) {
+ pr, err := beam.Run(ctx, runner, p)
+ if err != nil {
+ return metrics.QueryResults{}, err
+ }
+
+ // Request the metric called "counter1" in namespace called "namespace"
+ ms := pr.Metrics().Query(func(r metrics.SingleResult) bool {
+ return r.Namespace() == "namespace" && r.Name() == "counter1"
+ })
+
+ // Print the metric value - there should be only one line because there is
+ // only one metric called "counter1" in the namespace called "namespace"
+ for _, c := range ms.Counters() {
+ fmt.Println(c.Namespace(), "-", c.Name(), ":", c.Committed)
+ }
+ return ms, nil
+}
+
+type MyMetricsDoFn struct {
+ counter beam.Counter
+}
+
+func init() {
+ beam.RegisterType(reflect.TypeOf((*MyMetricsDoFn)(nil)))
+}
+
+func (fn *MyMetricsDoFn) Setup() {
+ // While metrics can be defined in package scope or dynamically
+ // it's most efficient to include them in the DoFn.
+ fn.counter = beam.NewCounter("namespace", "counter1")
+}
+
+func (fn *MyMetricsDoFn) ProcessElement(ctx context.Context, v beam.V, emit func(beam.V)) {
+ // count the elements
+ fn.counter.Inc(ctx, 1)
+ emit(v)
+}
+
+// [END metrics_pipeline]
diff --git a/sdks/go/examples/snippets/10metrics_test.go b/sdks/go/examples/snippets/10metrics_test.go
new file mode 100644
index 0000000..9a91b31
--- /dev/null
+++ b/sdks/go/examples/snippets/10metrics_test.go
@@ -0,0 +1,40 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 snippets
+
+import (
+ "context"
+ "testing"
+
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+)
+
+func TestMetricsPipeline(t *testing.T) {
+ p, s := beam.NewPipelineWithRoot()
+ input := beam.Create(s, 1, 2, 3, 4, 5, 6, 7)
+ addMetricDoFnToPipeline(s, input)
+ ms, err := executePipelineAndGetMetrics(context.Background(), p)
+ if err != nil {
+ t.Errorf("error executing pipeline: %v", err)
+ }
+ if got, want := len(ms.Counters()), 1; got != want {
+ t.Errorf("got %v counters, want %v: %+v", got, want, ms.Counters())
+ }
+ c := ms.Counters()[0]
+ if got, want := c.Committed, int64(7); got != want {
+ t.Errorf("Attempted Counter: got %v, want %v: %+v", got, want, c)
+ }
+}
diff --git a/sdks/go/examples/wordcount/wordcount.go b/sdks/go/examples/wordcount/wordcount.go
index 08d9cb5..4d54db9 100644
--- a/sdks/go/examples/wordcount/wordcount.go
+++ b/sdks/go/examples/wordcount/wordcount.go
@@ -110,16 +110,16 @@
}
var (
- wordRE = regexp.MustCompile(`[a-zA-Z]+('[a-z])?`)
- empty = beam.NewCounter("extract", "emptyLines")
- small_word_length = flag.Int("small_word_length", 9, "small_word_length")
- small_words = beam.NewCounter("extract", "small_words")
- lineLen = beam.NewDistribution("extract", "lineLenDistro")
+ wordRE = regexp.MustCompile(`[a-zA-Z]+('[a-z])?`)
+ empty = beam.NewCounter("extract", "emptyLines")
+ smallWordLength = flag.Int("small_word_length", 9, "length of small words (default: 9)")
+ smallWords = beam.NewCounter("extract", "smallWords")
+ lineLen = beam.NewDistribution("extract", "lineLenDistro")
)
// extractFn is a DoFn that emits the words in a given line and keeps a count for small words.
type extractFn struct {
- SmallWordLength int `json:"min_length"`
+ SmallWordLength int `json:"smallWordLength"`
}
func (f *extractFn) ProcessElement(ctx context.Context, line string, emit func(string)) {
@@ -131,7 +131,7 @@
// increment the counter for small words if length of words is
// less than small_word_length
if len(word) < f.SmallWordLength {
- small_words.Inc(ctx, 1)
+ smallWords.Inc(ctx, 1)
}
emit(word)
}
@@ -160,7 +160,7 @@
s = s.Scope("CountWords")
// Convert lines of text into individual words.
- col := beam.ParDo(s, &extractFn{SmallWordLength: *small_word_length}, lines)
+ col := beam.ParDo(s, &extractFn{SmallWordLength: *smallWordLength}, lines)
// Count the number of times each word occurs.
return stats.Count(s, col)
diff --git a/sdks/go/pkg/beam/core/core.go b/sdks/go/pkg/beam/core/core.go
index 74e42b4..b405bc6 100644
--- a/sdks/go/pkg/beam/core/core.go
+++ b/sdks/go/pkg/beam/core/core.go
@@ -27,5 +27,5 @@
// SdkName is the human readable name of the SDK for UserAgents.
SdkName = "Apache Beam SDK for Go"
// SdkVersion is the current version of the SDK.
- SdkVersion = "2.34.0.dev"
+ SdkVersion = "2.35.0.dev"
)
diff --git a/sdks/go/pkg/beam/core/graph/window/trigger.go b/sdks/go/pkg/beam/core/graph/window/trigger.go
index 9068076..f465ff5 100644
--- a/sdks/go/pkg/beam/core/graph/window/trigger.go
+++ b/sdks/go/pkg/beam/core/graph/window/trigger.go
@@ -15,17 +15,54 @@
package window
-import "fmt"
+import (
+ "fmt"
+ "time"
+)
+// Trigger describes when to emit new aggregations.
+// Fields are exported for use by the framework, and not intended
+// to be set by end users.
+//
+// This API is experimental and subject to change.
type Trigger struct {
- Kind string
- SubTriggers []Trigger
- Delay int64 // in milliseconds
- ElementCount int32
- EarlyTrigger *Trigger
- LateTrigger *Trigger
+ Kind string
+ SubTriggers []Trigger // Repeat, OrFinally, Any, All
+ TimestampTransforms []TimestampTransform // AfterProcessingTime
+ ElementCount int32 // ElementCount
+ EarlyTrigger *Trigger // AfterEndOfWindow
+ LateTrigger *Trigger // AfterEndOfWindow
}
+// TimestampTransform describes how an after processing time trigger
+// time is transformed to determine when to fire an aggregation.const
+// The base timestamp is always the when the first element of the pane
+// is received.
+//
+// A series of these transforms will be applied in order emit at regular intervals.
+type TimestampTransform interface {
+ timestampTransform()
+}
+
+// DelayTransform takes the timestamp and adds the given delay to it.
+type DelayTransform struct {
+ Delay int64 // in milliseconds
+}
+
+func (DelayTransform) timestampTransform() {}
+
+// AlignToTransform takes the timestamp and transforms it to the lowest
+// multiple of the period starting from the offset.
+//
+// Eg. A period of 20 with an offset of 45 would have alignments at 5,25,45,65 etc.
+// Timestamps would be transformed as follows: 0 to 5 would be transformed to 5,
+// 6 to 25 would be transformed to 25, 26 to 45 would be transformed to 45, and so on.
+type AlignToTransform struct {
+ Period, Offset int64 // in milliseconds
+}
+
+func (AlignToTransform) timestampTransform() {}
+
const (
DefaultTrigger string = "Trigger_Default_"
AlwaysTrigger string = "Trigger_Always_"
@@ -46,49 +83,101 @@
return Trigger{Kind: DefaultTrigger}
}
-// TriggerAlways constructs an always trigger that keeps firing immediately after an element is processed.
+// TriggerAlways constructs a trigger that fires immediately
+// whenever an element is received.
+//
// Equivalent to window.TriggerRepeat(window.TriggerAfterCount(1))
func TriggerAlways() Trigger {
return Trigger{Kind: AlwaysTrigger}
}
-// TriggerAfterCount constructs an element count trigger that fires after atleast `count` number of elements are processed.
+// TriggerAfterCount constructs a trigger that fires after
+// at least `count` number of elements are processed.
func TriggerAfterCount(count int32) Trigger {
return Trigger{Kind: ElementCountTrigger, ElementCount: count}
}
-// TriggerAfterProcessingTime constructs an after processing time trigger that fires after 'delay' milliseconds of processing time have passed.
-func TriggerAfterProcessingTime(delay int64) Trigger {
- return Trigger{Kind: AfterProcessingTimeTrigger, Delay: delay}
+// TriggerAfterProcessingTime constructs a trigger that fires relative to
+// when input first arrives.
+//
+// Must be configured with calls to PlusDelay, or AlignedTo. May be
+// configured with additional delay.
+func TriggerAfterProcessingTime() Trigger {
+ return Trigger{Kind: AfterProcessingTimeTrigger}
}
-// TriggerRepeat constructs a repeat trigger that fires a trigger repeatedly once the condition has been met.
+// PlusDelay configures an AfterProcessingTime trigger to fire after a specified delay,
+// no smaller than a millisecond.
+func (tr Trigger) PlusDelay(delay time.Duration) Trigger {
+ if tr.Kind != AfterProcessingTimeTrigger {
+ panic(fmt.Errorf("can't apply processing delay to %s, want: AfterProcessingTimeTrigger", tr.Kind))
+ }
+ if delay < time.Millisecond {
+ panic(fmt.Errorf("can't apply processing delay of less than a millisecond. Got: %v", delay))
+ }
+ tr.TimestampTransforms = append(tr.TimestampTransforms, DelayTransform{Delay: int64(delay / time.Millisecond)})
+ return tr
+}
+
+// AlignedTo configures an AfterProcessingTime trigger to fire
+// at the smallest multiple of period since the offset greater than the first element timestamp.
+//
+// * Period may not be smaller than a millisecond.
+// * Offset may be a zero time (time.Time{}).
+func (tr Trigger) AlignedTo(period time.Duration, offset time.Time) Trigger {
+ if tr.Kind != AfterProcessingTimeTrigger {
+ panic(fmt.Errorf("can't apply processing delay to %s, want: AfterProcessingTimeTrigger", tr.Kind))
+ }
+ if period < time.Millisecond {
+ panic(fmt.Errorf("can't apply an alignment period of less than a millisecond. Got: %v", period))
+ }
+ offsetMillis := int64(0)
+ if !offset.IsZero() {
+ // TODO: Change to call UnixMilli() once we move to only supporting a go version > 1.17.
+ offsetMillis = offset.Unix()*1e3 + int64(offset.Nanosecond())/1e6
+ }
+ tr.TimestampTransforms = append(tr.TimestampTransforms, AlignToTransform{
+ Period: int64(period / time.Millisecond),
+ Offset: offsetMillis,
+ })
+ return tr
+}
+
+// TriggerRepeat constructs a trigger that fires a trigger repeatedly
+// once the condition has been met.
+//
// Ex: window.TriggerRepeat(window.TriggerAfterCount(1)) is same as window.TriggerAlways().
func TriggerRepeat(tr Trigger) Trigger {
return Trigger{Kind: RepeatTrigger, SubTriggers: []Trigger{tr}}
}
-// TriggerAfterEndOfWindow constructs an end of window trigger that is configurable for early firing trigger(before the end of window)
-// and late firing trigger(after the end of window).
-// Default Options are: Default Trigger for EarlyFiring and No LateFiring. Override it with EarlyFiring and LateFiring methods on this trigger.
+// TriggerAfterEndOfWindow constructs a trigger that is configurable for early firing
+// (before the end of window) and late firing (after the end of window).
+//
+// Default Options are: Default Trigger for EarlyFiring and No LateFiring.
+// Override it with EarlyFiring and LateFiring methods on this trigger.
func TriggerAfterEndOfWindow() Trigger {
defaultEarly := TriggerDefault()
return Trigger{Kind: AfterEndOfWindowTrigger, EarlyTrigger: &defaultEarly, LateTrigger: nil}
}
-// EarlyFiring configures AfterEndOfWindow trigger with an early firing trigger.
+// EarlyFiring configures an AfterEndOfWindow trigger with an implicitly
+// repeated trigger that applies before the end of the window.
func (tr Trigger) EarlyFiring(early Trigger) Trigger {
if tr.Kind != AfterEndOfWindowTrigger {
- panic(fmt.Errorf("can't apply early firing to %s, got: %s, want: AfterEndOfWindowTrigger", tr.Kind, tr.Kind))
+ panic(fmt.Errorf("can't apply early firing to %s, want: AfterEndOfWindowTrigger", tr.Kind))
}
tr.EarlyTrigger = &early
return tr
}
-// LateFiring configures AfterEndOfWindow trigger with a late firing trigger
+// LateFiring configures an AfterEndOfWindow trigger with an implicitly
+// repeated trigger that applies after the end of the window.
+//
+// Not setting a late firing trigger means elements are discarded.
func (tr Trigger) LateFiring(late Trigger) Trigger {
if tr.Kind != AfterEndOfWindowTrigger {
- panic(fmt.Errorf("can't apply late firing to %s, got: %s, want: AfterEndOfWindowTrigger", tr.Kind, tr.Kind))
+ panic(fmt.Errorf("can't apply late firing to %s, want: AfterEndOfWindowTrigger", tr.Kind))
}
tr.LateTrigger = &late
return tr
diff --git a/sdks/go/pkg/beam/core/metrics/metrics.go b/sdks/go/pkg/beam/core/metrics/metrics.go
index 3633e0c..32fdd35 100644
--- a/sdks/go/pkg/beam/core/metrics/metrics.go
+++ b/sdks/go/pkg/beam/core/metrics/metrics.go
@@ -482,7 +482,40 @@
return QueryResults{mr.counters, mr.distributions, mr.gauges}
}
-// TODO(BEAM-11217): Implement Query(Filter) and metrics filtering
+// TODO(BEAM-11217): Implement querying metrics by DoFn
+
+// SingleResult interface facilitates metrics query filtering methods.
+type SingleResult interface {
+ Name() string
+ Namespace() string
+}
+
+// Query allows metrics querying with filter. The filter takes the form of predicate function. Example:
+// qr = pr.Metrics().Query(func(sr metrics.SingleResult) bool {
+// return sr.Namespace() == test.namespace
+// })
+func (mr Results) Query(f func(SingleResult) bool) QueryResults {
+ counters := []CounterResult{}
+ distributions := []DistributionResult{}
+ gauges := []GaugeResult{}
+
+ for _, counter := range mr.counters {
+ if f(counter) {
+ counters = append(counters, counter)
+ }
+ }
+ for _, distribution := range mr.distributions {
+ if f(distribution) {
+ distributions = append(distributions, distribution)
+ }
+ }
+ for _, gauge := range mr.gauges {
+ if f(gauge) {
+ gauges = append(gauges, gauge)
+ }
+ }
+ return QueryResults{counters, distributions, gauges}
+}
// QueryResults is the result of a query. Allows accessing all of the
// metrics that matched the filter.
@@ -529,6 +562,16 @@
return r.Attempted
}
+// Name returns the Name of this Counter.
+func (r CounterResult) Name() string {
+ return r.Key.Name
+}
+
+// Namespace returns the Namespace of this Counter.
+func (r CounterResult) Namespace() string {
+ return r.Key.Namespace
+}
+
// MergeCounters combines counter metrics that share a common key.
func MergeCounters(
attempted map[StepKey]int64,
@@ -571,6 +614,16 @@
return r.Attempted
}
+// Name returns the Name of this Distribution.
+func (r DistributionResult) Name() string {
+ return r.Key.Name
+}
+
+// Namespace returns the Namespace of this Distribution.
+func (r DistributionResult) Namespace() string {
+ return r.Key.Namespace
+}
+
// MergeDistributions combines distribution metrics that share a common key.
func MergeDistributions(
attempted map[StepKey]DistributionValue,
@@ -613,6 +666,16 @@
return r.Attempted
}
+// Name returns the Name of this Gauge.
+func (r GaugeResult) Name() string {
+ return r.Key.Name
+}
+
+// Namespace returns the Namespace of this Gauge.
+func (r GaugeResult) Namespace() string {
+ return r.Key.Namespace
+}
+
// StepKey uniquely identifies a metric within a pipeline graph.
type StepKey struct {
Step, Name, Namespace string
@@ -643,9 +706,9 @@
return res
}
-// MetricsExtractor extracts the metrics.Results from Store using ctx.
+// ResultsExtractor extracts the metrics.Results from Store using ctx.
// This is same as what metrics.dumperExtractor and metrics.dumpTo would do together.
-func MetricsExtractor(ctx context.Context) Results {
+func ResultsExtractor(ctx context.Context) Results {
store := GetStore(ctx)
m := make(map[Labels]interface{})
e := &Extractor{
diff --git a/sdks/go/pkg/beam/core/runtime/exec/data.go b/sdks/go/pkg/beam/core/runtime/exec/data.go
index 2b00a37..ed671ee 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/data.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/data.go
@@ -44,6 +44,16 @@
State StateReader
}
+// SideCache manages cached ReStream values for side inputs that can be re-used across
+// bundles.
+type SideCache interface {
+ // QueryCache checks the cache for a ReStream corresponding to the transform and
+ // side input being used.
+ QueryCache(transformID, sideInputID string) ReStream
+ // SetCache places a ReStream into the cache for a transform and side input.
+ SetCache(transformID, sideInputID string, input ReStream)
+}
+
// DataManager manages external data byte streams. Each data stream can be
// opened by one consumer only.
type DataManager interface {
@@ -59,6 +69,8 @@
OpenSideInput(ctx context.Context, id StreamID, sideInputID string, key, w []byte) (io.ReadCloser, error)
// OpenIterable opens a byte stream for reading unwindowed iterables from the runner.
OpenIterable(ctx context.Context, id StreamID, key []byte) (io.ReadCloser, error)
+ // GetSideInputCache returns the SideInputCache being used at the harness level.
+ GetSideInputCache() SideCache
}
// TODO(herohde) 7/20/2018: user state management
diff --git a/sdks/go/pkg/beam/core/runtime/exec/pardo.go b/sdks/go/pkg/beam/core/runtime/exec/pardo.go
index fdf20b0..6a41e9d 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/pardo.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/pardo.go
@@ -257,6 +257,7 @@
if err != nil {
return err
}
+ n.cache.key = w
n.cache.sideinput = sideinput
for i := 0; i < len(n.Side); i++ {
n.cache.extra[i] = sideinput[i].Value()
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/translate.go b/sdks/go/pkg/beam/core/runtime/graphx/translate.go
index c74b645..b31a40b 100644
--- a/sdks/go/pkg/beam/core/runtime/graphx/translate.go
+++ b/sdks/go/pkg/beam/core/runtime/graphx/translate.go
@@ -1042,16 +1042,34 @@
},
}
case window.AfterProcessingTimeTrigger:
- // TODO(BEAM-3304) Right now would work only for single delay value.
- // could be configured to take more than one delay values later.
- ttd := &pipepb.TimestampTransform{
- TimestampTransform: &pipepb.TimestampTransform_Delay_{
- Delay: &pipepb.TimestampTransform_Delay{DelayMillis: t.Delay},
- }}
- tt := []*pipepb.TimestampTransform{ttd}
+ if len(t.TimestampTransforms) == 0 {
+ panic("AfterProcessingTime trigger set without a delay or alignment.")
+ }
+ tts := []*pipepb.TimestampTransform{}
+ for _, tt := range t.TimestampTransforms {
+ var ttp *pipepb.TimestampTransform
+ switch tt := tt.(type) {
+ case window.DelayTransform:
+ ttp = &pipepb.TimestampTransform{
+ TimestampTransform: &pipepb.TimestampTransform_Delay_{
+ Delay: &pipepb.TimestampTransform_Delay{DelayMillis: tt.Delay},
+ }}
+ case window.AlignToTransform:
+ ttp = &pipepb.TimestampTransform{
+ TimestampTransform: &pipepb.TimestampTransform_AlignTo_{
+ AlignTo: &pipepb.TimestampTransform_AlignTo{
+ Period: tt.Period,
+ Offset: tt.Offset,
+ },
+ }}
+ }
+ tts = append(tts, ttp)
+ }
return &pipepb.Trigger{
Trigger: &pipepb.Trigger_AfterProcessingTime_{
- AfterProcessingTime: &pipepb.Trigger_AfterProcessingTime{TimestampTransforms: tt},
+ AfterProcessingTime: &pipepb.Trigger_AfterProcessingTime{
+ TimestampTransforms: tts,
+ },
},
}
case window.ElementCountTrigger:
diff --git a/sdks/go/pkg/beam/core/runtime/harness/harness.go b/sdks/go/pkg/beam/core/runtime/harness/harness.go
index a46f2c3..c0b1295 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/harness.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/harness.go
@@ -317,8 +317,9 @@
return fail(ctx, instID, "Failed: %v", err)
}
+ // TODO(BEAM-11097): Get and set valid tokens in cache
data := NewScopedDataManager(c.data, instID)
- state := NewScopedStateReader(c.state, instID)
+ state := NewScopedStateReaderWithCache(c.state, instID, c.cache)
err = plan.Execute(ctx, string(instID), exec.DataContext{Data: data, State: state})
data.Close()
state.Close()
diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go
index 5496d8b..529c21e 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go
@@ -44,12 +44,13 @@
type SideInputCache struct {
capacity int
mu sync.Mutex
- cache map[token]exec.ReusableInput
+ cache map[token]exec.ReStream
idsToTokens map[string]token
validTokens map[token]int8 // Maps tokens to active bundle counts
metrics CacheMetrics
}
+// CacheMetrics stores metrics for the cache across a pipeline run.
type CacheMetrics struct {
Hits int64
Misses int64
@@ -66,7 +67,7 @@
}
c.mu.Lock()
defer c.mu.Unlock()
- c.cache = make(map[token]exec.ReusableInput, cap)
+ c.cache = make(map[token]exec.ReStream, cap)
c.idsToTokens = make(map[string]token)
c.validTokens = make(map[token]int8)
c.capacity = cap
@@ -77,7 +78,7 @@
// transform and side input IDs to cache tokens in the process. Should be called at the start of every
// new ProcessBundleRequest. If the runner does not support caching, the passed cache token values
// should be empty and all get/set requests will silently be no-ops.
-func (c *SideInputCache) SetValidTokens(cacheTokens ...fnpb.ProcessBundleRequest_CacheToken) {
+func (c *SideInputCache) SetValidTokens(cacheTokens ...*fnpb.ProcessBundleRequest_CacheToken) {
c.mu.Lock()
defer c.mu.Unlock()
for _, tok := range cacheTokens {
@@ -109,7 +110,7 @@
// CompleteBundle takes the cache tokens passed to set the valid tokens and decrements their
// usage count for the purposes of maintaining a valid count of whether or not a value is
// still in use. Should be called once ProcessBundle has completed.
-func (c *SideInputCache) CompleteBundle(cacheTokens ...fnpb.ProcessBundleRequest_CacheToken) {
+func (c *SideInputCache) CompleteBundle(cacheTokens ...*fnpb.ProcessBundleRequest_CacheToken) {
c.mu.Lock()
defer c.mu.Unlock()
for _, tok := range cacheTokens {
@@ -148,7 +149,7 @@
// input has been cached. A query having a bad token (e.g. one that doesn't make a known
// token or one that makes a known but currently invalid token) is treated the same as a
// cache miss.
-func (c *SideInputCache) QueryCache(transformID, sideInputID string) exec.ReusableInput {
+func (c *SideInputCache) QueryCache(transformID, sideInputID string) exec.ReStream {
c.mu.Lock()
defer c.mu.Unlock()
tok, ok := c.makeAndValidateToken(transformID, sideInputID)
@@ -170,7 +171,7 @@
// with its corresponding transform ID and side input ID. If the IDs do not pair with a known, valid token
// then we silently do not cache the input, as this is an indication that the runner is treating that input
// as uncacheable.
-func (c *SideInputCache) SetCache(transformID, sideInputID string, input exec.ReusableInput) {
+func (c *SideInputCache) SetCache(transformID, sideInputID string, input exec.ReStream) {
c.mu.Lock()
defer c.mu.Unlock()
tok, ok := c.makeAndValidateToken(transformID, sideInputID)
diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go
index b9970c3..83a47a5 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go
@@ -22,32 +22,34 @@
fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
)
-// TestReusableInput implements the ReusableInput interface for the purposes
-// of testing.
-type TestReusableInput struct {
- transformID string
- sideInputID string
- value interface{}
+type TestReStream struct {
+ value interface{}
}
-func makeTestReusableInput(transformID, sideInputID string, value interface{}) exec.ReusableInput {
- return &TestReusableInput{transformID: transformID, sideInputID: sideInputID, value: value}
+func (r *TestReStream) Open() (exec.Stream, error) {
+ return &TestFixedStream{value: r.value}, nil
}
-// Init is a ReusableInput interface method, this is a no-op.
-func (r *TestReusableInput) Init() error {
+type TestFixedStream struct {
+ value interface{}
+}
+
+func (s *TestFixedStream) Close() error {
return nil
}
-// Value returns the stored value in the TestReusableInput.
-func (r *TestReusableInput) Value() interface{} {
- return r.value
+func (s *TestFixedStream) Read() (*exec.FullValue, error) {
+ return &exec.FullValue{Elm: s.value}, nil
}
-// Reset clears the value in the TestReusableInput.
-func (r *TestReusableInput) Reset() error {
- r.value = nil
- return nil
+func makeTestReStream(value interface{}) exec.ReStream {
+ return &TestReStream{value: value}
+}
+
+func getValue(rs exec.ReStream) interface{} {
+ stream, _ := rs.Open()
+ fullVal, _ := stream.Read()
+ return fullVal.Elm
}
func TestInit(t *testing.T) {
@@ -84,7 +86,7 @@
if err != nil {
t.Fatalf("cache init failed, got %v", err)
}
- input := makeTestReusableInput("t1", "s1", 10)
+ input := makeTestReStream(10)
s.SetCache("t1", "s1", input)
output := s.QueryCache("t1", "s1")
if output != nil {
@@ -102,31 +104,31 @@
sideID := "s1"
tok := token("tok1")
s.setValidToken(transID, sideID, tok)
- input := makeTestReusableInput(transID, sideID, 10)
+ input := makeTestReStream(10)
s.SetCache(transID, sideID, input)
output := s.QueryCache(transID, sideID)
if output == nil {
t.Fatalf("call to query cache missed when should have hit")
}
- val, ok := output.Value().(int)
+ val, ok := getValue(output).(int)
if !ok {
- t.Errorf("failed to convert value to integer, got %v", output.Value())
+ t.Errorf("failed to convert value to integer, got %v", getValue(output))
}
if val != 10 {
t.Errorf("element mismatch, expected 10, got %v", val)
}
}
-func makeRequest(transformID, sideInputID string, t token) fnpb.ProcessBundleRequest_CacheToken {
- var tok fnpb.ProcessBundleRequest_CacheToken
- var wrap fnpb.ProcessBundleRequest_CacheToken_SideInput_
- var side fnpb.ProcessBundleRequest_CacheToken_SideInput
- side.TransformId = transformID
- side.SideInputId = sideInputID
- wrap.SideInput = &side
- tok.Type = &wrap
- tok.Token = []byte(t)
- return tok
+func makeRequest(transformID, sideInputID string, t token) *fnpb.ProcessBundleRequest_CacheToken {
+ return &fnpb.ProcessBundleRequest_CacheToken{
+ Token: []byte(t),
+ Type: &fnpb.ProcessBundleRequest_CacheToken_SideInput_{
+ SideInput: &fnpb.ProcessBundleRequest_CacheToken_SideInput{
+ TransformId: transformID,
+ SideInputId: sideInputID,
+ },
+ },
+ }
}
func TestSetValidTokens(t *testing.T) {
@@ -158,7 +160,7 @@
t.Fatalf("cache init failed, got %v", err)
}
- var tokens []fnpb.ProcessBundleRequest_CacheToken
+ var tokens []*fnpb.ProcessBundleRequest_CacheToken
for _, input := range inputs {
t := makeRequest(input.transformID, input.sideInputID, input.tok)
tokens = append(tokens, t)
@@ -244,14 +246,14 @@
}
tokOne := makeRequest("t1", "s1", "tok1")
- inOne := makeTestReusableInput("t1", "s1", 10)
+ inOne := makeTestReStream(10)
s.SetValidTokens(tokOne)
s.SetCache("t1", "s1", inOne)
// Mark bundle as complete, drop count for tokOne to 0
s.CompleteBundle(tokOne)
tokTwo := makeRequest("t2", "s2", "tok2")
- inTwo := makeTestReusableInput("t2", "s2", 20)
+ inTwo := makeTestReStream(20)
s.SetValidTokens(tokTwo)
s.SetCache("t2", "s2", inTwo)
@@ -271,10 +273,10 @@
}
tokOne := makeRequest("t1", "s1", "tok1")
- inOne := makeTestReusableInput("t1", "s1", 10)
+ inOne := makeTestReStream(10)
tokTwo := makeRequest("t2", "s2", "tok2")
- inTwo := makeTestReusableInput("t2", "s2", 20)
+ inTwo := makeTestReStream(20)
s.SetValidTokens(tokOne, tokTwo)
s.SetCache("t1", "s1", inOne)
diff --git a/sdks/go/pkg/beam/core/runtime/harness/statemgr.go b/sdks/go/pkg/beam/core/runtime/harness/statemgr.go
index 09daa11..2b7ea73 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/statemgr.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/statemgr.go
@@ -24,6 +24,7 @@
"time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/harness/statecache"
"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
@@ -39,11 +40,18 @@
opened []io.Closer // track open readers to force close all
closed bool
mu sync.Mutex
+
+ cache *statecache.SideInputCache
}
// NewScopedStateReader returns a ScopedStateReader for the given instruction.
func NewScopedStateReader(mgr *StateChannelManager, instID instructionID) *ScopedStateReader {
- return &ScopedStateReader{mgr: mgr, instID: instID}
+ return &ScopedStateReader{mgr: mgr, instID: instID, cache: nil}
+}
+
+// NewScopedStateReaderWithCache returns a ScopedState reader for the given instruction with a pointer to a SideInputCache.
+func NewScopedStateReaderWithCache(mgr *StateChannelManager, instID instructionID, cache *statecache.SideInputCache) *ScopedStateReader {
+ return &ScopedStateReader{mgr: mgr, instID: instID, cache: cache}
}
// OpenSideInput opens a byte stream for reading iterable side input.
@@ -60,6 +68,11 @@
})
}
+// GetSideInputCache returns a pointer to the SideInputCache being used by the SDK harness.
+func (s *ScopedStateReader) GetSideInputCache() exec.SideCache {
+ return s.cache
+}
+
func (s *ScopedStateReader) openReader(ctx context.Context, id exec.StreamID, readerFn func(*StateChannel) *stateKeyReader) (*stateKeyReader, error) {
ch, err := s.open(ctx, id.Port)
if err != nil {
diff --git a/sdks/go/pkg/beam/metrics.go b/sdks/go/pkg/beam/metrics.go
index 07dbc8a..f4b9da4 100644
--- a/sdks/go/pkg/beam/metrics.go
+++ b/sdks/go/pkg/beam/metrics.go
@@ -34,12 +34,14 @@
*metrics.Counter
}
-// Inc increments the counter within by the given amount.
+// Inc increments the counter within by the given amount. The context must be
+// provided by the framework, or the value will not be recorded.
func (c Counter) Inc(ctx context.Context, v int64) {
c.Counter.Inc(ctx, v)
}
-// Dec decrements the counter within by the given amount.
+// Dec decrements the counter within by the given amount. The context must be
+// provided by the framework, or the value will not be recorded.
func (c Counter) Dec(ctx context.Context, v int64) {
c.Counter.Dec(ctx, v)
}
@@ -59,7 +61,8 @@
*metrics.Distribution
}
-// Update adds an observation to this distribution.
+// Update adds an observation to this distribution. The context must be
+// provided by the framework, or the value will not be recorded.
func (c Distribution) Update(ctx context.Context, v int64) {
c.Distribution.Update(ctx, v)
}
@@ -79,7 +82,8 @@
*metrics.Gauge
}
-// Set sets the current value for this gauge.
+// Set sets the current value for this gauge. The context must be
+// provided by the framework, or the value will not be recorded.
func (c Gauge) Set(ctx context.Context, v int64) {
c.Gauge.Set(ctx, v)
}
diff --git a/sdks/go/pkg/beam/option.go b/sdks/go/pkg/beam/option.go
index 5d8df6f..db8f433 100644
--- a/sdks/go/pkg/beam/option.go
+++ b/sdks/go/pkg/beam/option.go
@@ -54,11 +54,11 @@
var infer []TypeDefinition
for _, opt := range opts {
- switch opt.(type) {
+ switch opt := opt.(type) {
case SideInput:
- side = append(side, opt.(SideInput))
+ side = append(side, opt)
case TypeDefinition:
- infer = append(infer, opt.(TypeDefinition))
+ infer = append(infer, opt)
default:
panic(fmt.Sprintf("Unexpected opt: %v", opt))
}
diff --git a/sdks/go/pkg/beam/runners/direct/direct.go b/sdks/go/pkg/beam/runners/direct/direct.go
index 172915a..399a74c 100644
--- a/sdks/go/pkg/beam/runners/direct/direct.go
+++ b/sdks/go/pkg/beam/runners/direct/direct.go
@@ -84,7 +84,7 @@
}
func newDirectPipelineResult(ctx context.Context) (*directPipelineResult, error) {
- metrics := metrics.MetricsExtractor(ctx)
+ metrics := metrics.ResultsExtractor(ctx)
return &directPipelineResult{metrics: &metrics}, nil
}
diff --git a/sdks/go/pkg/beam/schema.go b/sdks/go/pkg/beam/schema.go
index 95811bf..b25a3e2 100644
--- a/sdks/go/pkg/beam/schema.go
+++ b/sdks/go/pkg/beam/schema.go
@@ -16,6 +16,7 @@
package beam
import (
+ "fmt"
"io"
"reflect"
@@ -55,7 +56,24 @@
// is called in a package init() function.
func RegisterSchemaProvider(rt reflect.Type, provider interface{}) {
p := provider.(SchemaProvider)
- schema.RegisterLogicalTypeProvider(rt, p.FromLogicalType)
+ switch rt.Kind() {
+ case reflect.Interface:
+ schema.RegisterLogicalTypeProvider(rt, p.FromLogicalType)
+ case reflect.Ptr:
+ if rt.Elem().Kind() != reflect.Struct {
+ panic(fmt.Sprintf("beam.RegisterSchemaProvider: unsupported type kind for schema provider %v is a %v, must be interface, struct or *struct.", rt, rt.Kind()))
+ }
+ fallthrough
+ case reflect.Struct:
+ st, err := p.FromLogicalType(rt)
+ if err != nil {
+ panic(fmt.Sprintf("beam.RegisterSchemaProvider: schema type provider for %v, doesn't support that type", rt))
+ }
+ schema.RegisterLogicalType(schema.ToLogicalType(rt.Name(), rt, st))
+ default:
+ panic(fmt.Sprintf("beam.RegisterSchemaProvider: unsupported type kind for schema provider %v is a %v, must be interface, struct or *struct.", rt, rt.Kind()))
+ }
+
coder.RegisterSchemaProviders(rt, p.BuildEncoder, p.BuildDecoder)
}
diff --git a/sdks/go/pkg/beam/windowing.go b/sdks/go/pkg/beam/windowing.go
index 58254ac..90cf299 100644
--- a/sdks/go/pkg/beam/windowing.go
+++ b/sdks/go/pkg/beam/windowing.go
@@ -17,6 +17,7 @@
import (
"fmt"
+ "time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
@@ -27,31 +28,50 @@
windowIntoOption()
}
-type WindowTrigger struct {
- Name window.Trigger
+type windowTrigger struct {
+ trigger window.Trigger
}
-func (t WindowTrigger) windowIntoOption() {}
+func (t windowTrigger) windowIntoOption() {}
-// Trigger applies `tr` trigger to the window.
-func Trigger(tr window.Trigger) WindowTrigger {
- return WindowTrigger{Name: tr}
+// Trigger applies the given trigger to the window.
+//
+// Trigger support in the Go SDK is currently experimental
+// and may have breaking changes made to it.
+// Use at your own discretion.
+func Trigger(tr window.Trigger) WindowIntoOption {
+ return windowTrigger{trigger: tr}
}
-type AccumulationMode struct {
- Mode window.AccumulationMode
+type accumulationMode struct {
+ mode window.AccumulationMode
}
-func (m AccumulationMode) windowIntoOption() {}
+func (m accumulationMode) windowIntoOption() {}
// PanesAccumulate applies an Accumulating AccumulationMode to the window.
-func PanesAccumulate() AccumulationMode {
- return AccumulationMode{Mode: window.Accumulating}
+// After a pane fires, already processed elements will accumulate and
+// elements will be repeated in subseqent firings for the window.
+func PanesAccumulate() WindowIntoOption {
+ return accumulationMode{mode: window.Accumulating}
}
// PanesDiscard applies a Discarding AccumulationMode to the window.
-func PanesDiscard() AccumulationMode {
- return AccumulationMode{Mode: window.Discarding}
+// After a pane fires, already processed elements will be discarded
+// and not included in later firings for the window.
+func PanesDiscard() WindowIntoOption {
+ return accumulationMode{mode: window.Discarding}
+}
+
+type allowedLateness struct {
+ delay time.Duration
+}
+
+func (m allowedLateness) windowIntoOption() {}
+
+// AllowedLateness configures for how long data may arrive after the end of a window.
+func AllowedLateness(delay time.Duration) WindowIntoOption {
+ return allowedLateness{delay: delay}
}
// WindowInto applies the windowing strategy to each element.
@@ -70,10 +90,15 @@
ws := window.WindowingStrategy{Fn: wfn, Trigger: window.Trigger{}}
for _, opt := range opts {
switch opt := opt.(type) {
- case WindowTrigger:
- ws.Trigger = opt.Name
- case AccumulationMode:
- ws.AccumulationMode = opt.Mode
+ case windowTrigger:
+ // TODO(BEAM-3304): call validation on trigger construction here
+ // so local errors can be returned to the user in their pipeline
+ // context instead of at pipeline translation time.
+ ws.Trigger = opt.trigger
+ case accumulationMode:
+ ws.AccumulationMode = opt.mode
+ case allowedLateness:
+ ws.AllowedLateness = int(opt.delay / time.Millisecond)
default:
panic(fmt.Sprintf("Unknown WindowInto option type: %T: %v", opt, opt))
}
diff --git a/sdks/go/pkg/beam/xlang.go b/sdks/go/pkg/beam/xlang.go
index b9c6ec7..2d36265 100644
--- a/sdks/go/pkg/beam/xlang.go
+++ b/sdks/go/pkg/beam/xlang.go
@@ -143,6 +143,23 @@
namedInputs map[string]PCollection,
namedOutputTypes map[string]FullType,
) map[string]PCollection {
+ namedOutputs, err := TryCrossLanguage(s, urn, payload, expansionAddr, namedInputs, namedOutputTypes)
+ if err != nil {
+ panic(errors.WithContextf(err, "tried cross-language for %v against %v and failed", urn, expansionAddr))
+ }
+ return namedOutputs
+}
+
+// TryCrossLanguage coordinates the core functions required to execute the cross-language transform.
+// See CrossLanguage for user documentation.
+func TryCrossLanguage(
+ s Scope,
+ urn string,
+ payload []byte,
+ expansionAddr string,
+ namedInputs map[string]PCollection,
+ namedOutputTypes map[string]FullType,
+) (map[string]PCollection, error) {
if !s.IsValid() {
panic(errors.New("invalid scope"))
}
@@ -150,41 +167,36 @@
inputsMap, inboundLinks := graph.NamedInboundLinks(mapPCollectionToNode(namedInputs))
outputsMap, outboundLinks := graph.NamedOutboundLinks(s.real, namedOutputTypes)
+ // Set the coder for outbound links for downstream validation.
+ for n, i := range outputsMap {
+ c := NewCoder(namedOutputTypes[n])
+ outboundLinks[i].To.Coder = c.coder
+ }
+
ext := graph.ExternalTransform{
Urn: urn,
Payload: payload,
ExpansionAddr: expansionAddr,
}.WithNamedInputs(inputsMap).WithNamedOutputs(outputsMap)
- namedOutputs, err := TryCrossLanguage(s, &ext, inboundLinks, outboundLinks)
- if err != nil {
- panic(errors.WithContextf(err, "tried cross-language and failed"))
- }
- return mapNodeToPCollection(namedOutputs)
-}
-
-// TryCrossLanguage coordinates the core functions required to execute the cross-language transform.
-// This is mainly intended for internal use. For the general-use entry point, see
-// beam.CrossLanguage.
-func TryCrossLanguage(s Scope, ext *graph.ExternalTransform, ins []*graph.Inbound, outs []*graph.Outbound) (map[string]*graph.Node, error) {
// Adding an edge in the graph corresponding to the ExternalTransform
- edge, isBoundedUpdater := graph.NewCrossLanguage(s.real, s.scope, ext, ins, outs)
+ edge, isBoundedUpdater := graph.NewCrossLanguage(s.real, s.scope, &ext, inboundLinks, outboundLinks)
// Once the appropriate input and output nodes are added to the edge, a
// unique namespace can be requested.
ext.Namespace = graph.NewNamespace()
// Expand the transform into ext.Expanded.
- if err := xlangx.Expand(edge, ext); err != nil {
+ if err := xlangx.Expand(edge, &ext); err != nil {
return nil, errors.WithContext(err, "expanding external transform")
}
// Ensures the expected named outputs are present
- graphx.VerifyNamedOutputs(ext)
+ graphx.VerifyNamedOutputs(&ext)
// Using the expanded outputs, the graph's counterpart outputs are updated with bounded values
graphx.ResolveOutputIsBounded(edge, isBoundedUpdater)
- return graphx.ExternalOutputs(edge), nil
+ return mapNodeToPCollection(graphx.ExternalOutputs(edge)), nil
}
// Wrapper functions to handle beam <-> graph boundaries
diff --git a/sdks/go/test/integration/integration.go b/sdks/go/test/integration/integration.go
index 853b248..6132841 100644
--- a/sdks/go/test/integration/integration.go
+++ b/sdks/go/test/integration/integration.go
@@ -95,6 +95,8 @@
"TestTestStream.*",
// The trigger tests uses TestStream
"TestTrigger.*",
+ // TODO(BEAM-13006): Samza doesn't yet support post job metrics, used by WordCount
+ "TestWordCount.*",
}
var sparkFilters = []string{
diff --git a/sdks/go/test/integration/primitives/windowinto.go b/sdks/go/test/integration/primitives/windowinto.go
index 1c2286d..afa1500 100644
--- a/sdks/go/test/integration/primitives/windowinto.go
+++ b/sdks/go/test/integration/primitives/windowinto.go
@@ -93,8 +93,8 @@
WindowSums(s.Scope("Lifted"), stats.SumPerKey)
}
-func validateEquals(s beam.Scope, wfn *window.Fn, in beam.PCollection, tr window.Trigger, m beam.AccumulationMode, expected ...interface{}) {
- windowed := beam.WindowInto(s, wfn, in, beam.Trigger(tr), m)
+func validateEquals(s beam.Scope, wfn *window.Fn, in beam.PCollection, opts []beam.WindowIntoOption, expected ...interface{}) {
+ windowed := beam.WindowInto(s, wfn, in, opts...)
sums := stats.Sum(s, windowed)
sums = beam.WindowInto(s, window.NewGlobalWindows(), sums)
passert.Equals(s, sums, expected...)
@@ -110,7 +110,10 @@
col := teststream.Create(s, con)
windowSize := 10 * time.Second
- validateEquals(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col, window.TriggerDefault(), beam.PanesDiscard(), 6.0, 9.0)
+ validateEquals(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col,
+ []beam.WindowIntoOption{
+ beam.Trigger(window.TriggerDefault()),
+ }, 6.0, 9.0)
}
// TriggerAlways tests the Always trigger, it is expected to receive every input value as the output.
@@ -121,11 +124,16 @@
col := teststream.Create(s, con)
windowSize := 10 * time.Second
- validateEquals(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col, window.TriggerAlways(), beam.PanesDiscard(), 1.0, 2.0, 3.0)
+ validateEquals(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col,
+ []beam.WindowIntoOption{
+ beam.Trigger(window.TriggerAlways()),
+ }, 1.0, 2.0, 3.0)
}
-func validateCount(s beam.Scope, wfn *window.Fn, in beam.PCollection, tr window.Trigger, m beam.AccumulationMode, expected int) {
- windowed := beam.WindowInto(s, wfn, in, beam.Trigger(tr), m)
+// validateCount handles cases where we can only be sure of the count of elements
+// and not their ordering.
+func validateCount(s beam.Scope, wfn *window.Fn, in beam.PCollection, opts []beam.WindowIntoOption, expected int) {
+ windowed := beam.WindowInto(s, wfn, in, opts...)
sums := stats.Sum(s, windowed)
sums = beam.WindowInto(s, window.NewGlobalWindows(), sums)
passert.Count(s, sums, "total collections", expected)
@@ -147,7 +155,10 @@
// waits only for two elements to arrive and fires output after that and never fires that.
// For the trigger to fire every 2 elements, combine it with Repeat Trigger
- validateCount(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col, window.TriggerAfterCount(2), beam.PanesDiscard(), 2)
+ validateCount(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col,
+ []beam.WindowIntoOption{
+ beam.Trigger(window.TriggerAfterCount(2)),
+ }, 2)
}
// TriggerAfterProcessingTime tests the AfterProcessingTime Trigger, it fires output panes once 't' processing time has passed
@@ -162,7 +173,10 @@
col := teststream.Create(s, con)
- validateEquals(s.Scope("Global"), window.NewGlobalWindows(), col, window.TriggerAfterProcessingTime(5000), beam.PanesDiscard(), 6.0)
+ validateEquals(s.Scope("Global"), window.NewGlobalWindows(), col,
+ []beam.WindowIntoOption{
+ beam.Trigger(window.TriggerAfterProcessingTime().PlusDelay(5 * time.Second)),
+ }, 6.0)
}
// TriggerRepeat tests the repeat trigger. As of now is it is configure to take only one trigger as a subtrigger.
@@ -177,7 +191,10 @@
col := teststream.Create(s, con)
- validateCount(s.Scope("Global"), window.NewGlobalWindows(), col, window.TriggerRepeat(window.TriggerAfterCount(2)), beam.PanesDiscard(), 3)
+ validateCount(s.Scope("Global"), window.NewGlobalWindows(), col,
+ []beam.WindowIntoOption{
+ beam.Trigger(window.TriggerRepeat(window.TriggerAfterCount(2))),
+ }, 3)
}
// TriggerAfterEndOfWindow tests the AfterEndOfWindow Trigger. With AfterCount(2) as the early firing trigger and AfterCount(1) as late firing trigger.
@@ -189,7 +206,12 @@
col := teststream.Create(s, con)
windowSize := 10 * time.Second
- trigger := window.TriggerAfterEndOfWindow().EarlyFiring(window.TriggerAfterCount(2)).LateFiring(window.TriggerAfterCount(1))
+ trigger := window.TriggerAfterEndOfWindow().
+ EarlyFiring(window.TriggerAfterCount(2)).
+ LateFiring(window.TriggerAfterCount(1))
- validateCount(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col, trigger, beam.PanesDiscard(), 2)
+ validateCount(s.Scope("Fixed"), window.NewFixedWindows(windowSize), col,
+ []beam.WindowIntoOption{
+ beam.Trigger(trigger),
+ }, 2)
}
diff --git a/sdks/go/test/integration/wordcount/wordcount.go b/sdks/go/test/integration/wordcount/wordcount.go
index bcc40b5..6be5901 100644
--- a/sdks/go/test/integration/wordcount/wordcount.go
+++ b/sdks/go/test/integration/wordcount/wordcount.go
@@ -30,9 +30,10 @@
)
var (
- wordRE = regexp.MustCompile(`[a-zA-Z]+('[a-z])?`)
- empty = beam.NewCounter("extract", "emptyLines")
- lineLen = beam.NewDistribution("extract", "lineLenDistro")
+ wordRE = regexp.MustCompile(`[a-zA-Z]+('[a-z])?`)
+ empty = beam.NewCounter("extract", "emptyLines")
+ lineLen = beam.NewDistribution("extract", "lineLenDistro")
+ smallWords = beam.NewCounter("extract", "smallWords")
)
// CountWords is a composite transform that counts the words of a PCollection
@@ -56,6 +57,9 @@
empty.Inc(ctx, 1)
}
for _, word := range wordRE.FindAllString(line, -1) {
+ if len(word) < 6 {
+ smallWords.Inc(ctx, 1)
+ }
emit(word)
}
}
@@ -74,8 +78,12 @@
p, s := beam.NewPipelineWithRoot()
in := textio.Read(s, glob)
+ WordCountFromPCol(s, in, hash, size)
+ return p
+}
+
+// WordCountFromPCol counts the words from a PCollection and validates it.
+func WordCountFromPCol(s beam.Scope, in beam.PCollection, hash string, size int) {
out := Format(s, CountWords(s, in))
passert.Hash(s, out, "out", hash, size)
-
- return p
}
diff --git a/sdks/go/test/integration/wordcount/wordcount_test.go b/sdks/go/test/integration/wordcount/wordcount_test.go
index 482d9a3..805b0ec 100644
--- a/sdks/go/test/integration/wordcount/wordcount_test.go
+++ b/sdks/go/test/integration/wordcount/wordcount_test.go
@@ -19,7 +19,8 @@
"strings"
"testing"
- "github.com/apache/beam/sdks/v2/go/pkg/beam/io/filesystem/memfs"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam"
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/core/metrics"
_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow"
_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/flink"
_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/samza"
@@ -30,9 +31,11 @@
func TestWordCount(t *testing.T) {
tests := []struct {
- lines []string
- words int
- hash string
+ lines []string
+ words int
+ hash string
+ smallWordsCount int64
+ lineLen metrics.DistributionValue
}{
{
[]string{
@@ -40,6 +43,8 @@
},
1,
"6zZtmVTet7aIhR3wmPE8BA==",
+ 1,
+ metrics.DistributionValue{Count: 1, Sum: 3, Min: 3, Max: 3},
},
{
[]string{
@@ -49,6 +54,8 @@
},
1,
"jAk8+k4BOH7vQDUiUZdfWg==",
+ 6,
+ metrics.DistributionValue{Count: 3, Sum: 21, Min: 3, Max: 11},
},
{
[]string{
@@ -56,6 +63,8 @@
},
2,
"Nz70m/sn3Ep9o484r7MalQ==",
+ 6,
+ metrics.DistributionValue{Count: 1, Sum: 23, Min: 23, Max: 23},
},
{
[]string{
@@ -63,6 +72,8 @@
},
2,
"Nz70m/sn3Ep9o484r7MalQ==", // ordering doesn't matter: same hash as above
+ 6,
+ metrics.DistributionValue{Count: 1, Sum: 23, Min: 23, Max: 23},
},
{
[]string{
@@ -75,19 +86,42 @@
},
2,
"Nz70m/sn3Ep9o484r7MalQ==", // whitespace doesn't matter: same hash as above
+ 6,
+ metrics.DistributionValue{Count: 6, Sum: 37, Min: 0, Max: 11},
},
}
for _, test := range tests {
integration.CheckFilters(t)
- const filename = "memfs://input"
- memfs.Write(filename, []byte(strings.Join(test.lines, "\n")))
-
- p := WordCount(filename, test.hash, test.words)
- _, err := ptest.RunWithMetrics(p)
+ p, s := beam.NewPipelineWithRoot()
+ lines := beam.CreateList(s, test.lines)
+ WordCountFromPCol(s, lines, test.hash, test.words)
+ pr, err := ptest.RunWithMetrics(p)
if err != nil {
t.Errorf("WordCount(\"%v\") failed: %v", strings.Join(test.lines, "|"), err)
}
+
+ qr := pr.Metrics().Query(func(sr metrics.SingleResult) bool {
+ return sr.Name() == "smallWords"
+ })
+ counter := metrics.CounterResult{}
+ if len(qr.Counters()) != 0 {
+ counter = qr.Counters()[0]
+ }
+ if counter.Result() != test.smallWordsCount {
+ t.Errorf("Metrics().Query(by Name) failed. Got %d counters, Want %d counters", counter.Result(), test.smallWordsCount)
+ }
+
+ qr = pr.Metrics().Query(func(sr metrics.SingleResult) bool {
+ return sr.Name() == "lineLenDistro"
+ })
+ distribution := metrics.DistributionResult{}
+ if len(qr.Distributions()) != 0 {
+ distribution = qr.Distributions()[0]
+ }
+ if distribution.Result() != test.lineLen {
+ t.Errorf("Metrics().Query(by Name) failed. Got %v distribution, Want %v distribution", distribution.Result(), test.lineLen)
+ }
}
}
diff --git a/sdks/go/test/run_validatesrunner_tests.sh b/sdks/go/test/run_validatesrunner_tests.sh
index 9c9c16f..6c34c30 100755
--- a/sdks/go/test/run_validatesrunner_tests.sh
+++ b/sdks/go/test/run_validatesrunner_tests.sh
@@ -70,7 +70,8 @@
# jar from the appropriate gradle module, which may not succeed.
set -e
-set -v
+trap '! [[ "$BASH_COMMAND" =~ ^(echo|read|if|ARGS|shift|SOCKET_SCRIPT|\[\[) ]] && \
+cmd=`eval echo "$BASH_COMMAND" 2>/dev/null` && echo "\$ $cmd"' DEBUG
# Default test targets.
TESTS="./test/integration/... ./test/regression"
@@ -395,7 +396,7 @@
cd sdks/go
echo ">>> RUNNING $RUNNER integration tests with pipeline options: $ARGS"
-go test -v $TESTS $ARGS \
+go test -v $TESTS $ARGS 1>&2 \
|| TEST_EXIT_CODE=$? # don't fail fast here; clean up environment before exiting
cd ../..
diff --git a/sdks/java/container/license_scripts/dep_urls_java.yaml b/sdks/java/container/license_scripts/dep_urls_java.yaml
index 14dbb3e..fd4126a 100644
--- a/sdks/java/container/license_scripts/dep_urls_java.yaml
+++ b/sdks/java/container/license_scripts/dep_urls_java.yaml
@@ -56,3 +56,7 @@
'2.12.4':
license: "https://raw.githubusercontent.com/FasterXML/jackson-bom/master/LICENSE"
type: "Apache License 2.0"
+junit-dep:
+ '4.11':
+ license: "https://opensource.org/licenses/cpl1.0.txt"
+ type: "Common Public License Version 1.0"
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/AvroCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/AvroCoder.java
index b236783..8d43fdd 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/AvroCoder.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/AvroCoder.java
@@ -117,9 +117,19 @@
* @param <T> the element type
*/
public static <T> AvroCoder<T> of(TypeDescriptor<T> type) {
+ return of(type, true);
+ }
+
+ /**
+ * Returns an {@code AvroCoder} instance for the provided element type, respecting whether to use
+ * Avro's Reflect* or Specific* suite for encoding and decoding.
+ *
+ * @param <T> the element type
+ */
+ public static <T> AvroCoder<T> of(TypeDescriptor<T> type, boolean useReflectApi) {
@SuppressWarnings("unchecked")
Class<T> clazz = (Class<T>) type.getRawType();
- return of(clazz);
+ return of(clazz, useReflectApi);
}
/**
@@ -128,7 +138,7 @@
* @param <T> the element type
*/
public static <T> AvroCoder<T> of(Class<T> clazz) {
- return of(clazz, false);
+ return of(clazz, true);
}
/**
@@ -140,8 +150,8 @@
}
/**
- * Returns an {@code AvroCoder} instance for the given class using Avro's Reflection API for
- * encoding and decoding.
+ * Returns an {@code AvroCoder} instance for the given class, respecting whether to use Avro's
+ * Reflect* or Specific* suite for encoding and decoding.
*
* @param <T> the element type
*/
@@ -158,12 +168,12 @@
* @param <T> the element type
*/
public static <T> AvroCoder<T> of(Class<T> type, Schema schema) {
- return of(type, schema, false);
+ return of(type, schema, true);
}
/**
- * Returns an {@code AvroCoder} instance for the given class and schema using Avro's Reflection
- * API for encoding and decoding.
+ * Returns an {@code AvroCoder} instance for the given class and schema, respecting whether to use
+ * Avro's Reflect* or Specific* suite for encoding and decoding.
*
* @param <T> the element type
*/
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadableFileCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadableFileCoder.java
index 51bb83e..bdf10c4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadableFileCoder.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadableFileCoder.java
@@ -20,31 +20,66 @@
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
-import org.apache.beam.sdk.coders.AtomicCoder;
+import java.util.Collections;
+import java.util.List;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.coders.StructuredCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.io.fs.MatchResult;
+import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
import org.apache.beam.sdk.io.fs.MetadataCoder;
+import org.checkerframework.checker.initialization.qual.Initialized;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.UnknownKeyFor;
-/** A {@link Coder} for {@link FileIO.ReadableFile}. */
-public class ReadableFileCoder extends AtomicCoder<FileIO.ReadableFile> {
- private static final ReadableFileCoder INSTANCE = new ReadableFileCoder();
+/** A {@link Coder} for {@link org.apache.beam.sdk.io.FileIO.ReadableFile}. */
+public class ReadableFileCoder extends StructuredCoder<FileIO.ReadableFile> {
- /** Returns the instance of {@link ReadableFileCoder}. */
+ private final Coder<Metadata> metadataCoder;
+
+ public static ReadableFileCoder of(Coder<Metadata> metadataCoder) {
+ return new ReadableFileCoder(metadataCoder);
+ }
+
public static ReadableFileCoder of() {
- return INSTANCE;
+ return new ReadableFileCoder(MetadataCoder.of());
+ }
+
+ public Coder<Metadata> getMetadataCoder() {
+ return metadataCoder;
+ }
+
+ private ReadableFileCoder(Coder<Metadata> metadataCoder) {
+ this.metadataCoder = metadataCoder;
}
@Override
- public void encode(FileIO.ReadableFile value, OutputStream os) throws IOException {
- MetadataCoder.of().encode(value.getMetadata(), os);
- VarIntCoder.of().encode(value.getCompression().ordinal(), os);
+ public void encode(
+ FileIO.ReadableFile value, @UnknownKeyFor @NonNull @Initialized OutputStream outStream)
+ throws CoderException, IOException {
+ getMetadataCoder().encode(value.getMetadata(), outStream);
+ VarIntCoder.of().encode(value.getCompression().ordinal(), outStream);
}
@Override
- public FileIO.ReadableFile decode(InputStream is) throws IOException {
- MatchResult.Metadata metadata = MetadataCoder.of().decode(is);
- Compression compression = Compression.values()[VarIntCoder.of().decode(is)];
+ public FileIO.ReadableFile decode(@UnknownKeyFor @NonNull @Initialized InputStream inStream)
+ throws CoderException, IOException {
+ MatchResult.Metadata metadata = getMetadataCoder().decode(inStream);
+ Compression compression = Compression.values()[VarIntCoder.of().decode(inStream)];
return new FileIO.ReadableFile(metadata, compression);
}
+
+ @Override
+ public @UnknownKeyFor @NonNull @Initialized List<? extends Coder<?>> getCoderArguments() {
+ return Collections.singletonList(metadataCoder);
+ }
+
+ @Override
+ public void verifyDeterministic() throws NonDeterministicException {
+ // ignore the default Metadata coder for backward compatible
+ if (!getMetadataCoder().equals(MetadataCoder.of())) {
+ verifyDeterministic(this, "Metadata coder must be deterministic", getMetadataCoder());
+ }
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaTranslation.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaTranslation.java
index 2215f3d..e26a3e9 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaTranslation.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaTranslation.java
@@ -226,7 +226,13 @@
Schema.Builder builder = Schema.builder();
Map<String, Integer> encodingLocationMap = Maps.newHashMap();
for (SchemaApi.Field protoField : protoSchema.getFieldsList()) {
- Field field = fieldFromProto(protoField);
+ Field field;
+ try {
+ field = fieldFromProto(protoField);
+ } catch (Exception e) {
+ throw new IllegalArgumentException(
+ "Failed to decode Schema due to an error decoding Field proto:\n\n" + protoField, e);
+ }
builder.addField(field);
encodingLocationMap.put(protoField.getName(), protoField.getEncodingPosition());
}
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/AvroCoderTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/AvroCoderTest.java
index d7886c3..9443aad 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/AvroCoderTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/AvroCoderTest.java
@@ -323,7 +323,8 @@
@Test
public void testSpecificRecordEncoding() throws Exception {
- AvroCoder<TestAvro> coder = AvroCoder.of(TestAvro.class, AVRO_SPECIFIC_RECORD.getSchema());
+ AvroCoder<TestAvro> coder =
+ AvroCoder.of(TestAvro.class, AVRO_SPECIFIC_RECORD.getSchema(), false);
assertTrue(SpecificRecord.class.isAssignableFrom(coder.getType()));
CoderProperties.coderDecodeEncodeEqual(coder, AVRO_SPECIFIC_RECORD);
@@ -415,8 +416,8 @@
}
@Test
- public void testAvroReflectCoderIsSerializable() throws Exception {
- AvroCoder<Pojo> coder = AvroCoder.of(Pojo.class, true);
+ public void testAvroSpecificCoderIsSerializable() throws Exception {
+ AvroCoder<Pojo> coder = AvroCoder.of(Pojo.class, false);
// Check that the coder is serializable using the regular JSON approach.
SerializableUtils.ensureSerializable(coder);
diff --git a/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java b/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java
index 3c1189e..7e4e1bf 100644
--- a/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java
+++ b/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java
@@ -434,19 +434,21 @@
+ "native Read transform, your Pipeline will fail during Pipeline submission.");
}
- ClassLoader classLoader = Environments.class.getClassLoader();
- if (classLoader == null) {
- throw new RuntimeException(
- "Cannot detect classpath: classloader is null (is it the bootstrap classloader?)");
- }
+ List<String> filesToStage = pipelineOptions.as(PortablePipelineOptions.class).getFilesToStage();
+ if (filesToStage == null || filesToStage.isEmpty()) {
+ ClassLoader classLoader = Environments.class.getClassLoader();
+ if (classLoader == null) {
+ throw new RuntimeException(
+ "Cannot detect classpath: classloader is null (is it the bootstrap classloader?)");
+ }
- List<String> classpathResources =
- detectClassPathResourcesToStage(classLoader, pipeline.getOptions());
- if (classpathResources.isEmpty()) {
- throw new IllegalArgumentException("No classpath elements found.");
+ filesToStage = detectClassPathResourcesToStage(classLoader, pipeline.getOptions());
+ if (filesToStage.isEmpty()) {
+ throw new IllegalArgumentException("No classpath elements found.");
+ }
+ LOG.debug("Staging to files from the classpath: {}", filesToStage.size());
}
- LOG.debug("Staging to files from the classpath: {}", classpathResources.size());
- pipeline.getOptions().as(PortablePipelineOptions.class).setFilesToStage(classpathResources);
+ pipeline.getOptions().as(PortablePipelineOptions.class).setFilesToStage(filesToStage);
RehydratedComponents rehydratedComponents =
RehydratedComponents.forComponents(request.getComponents()).withPipeline(pipeline);
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java
index c435695..543d322 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.extensions.sql.impl.rel;
+import static org.apache.beam.sdk.schemas.Schema.Field;
import static org.apache.beam.sdk.schemas.Schema.FieldType;
import static org.apache.beam.vendor.calcite.v1_26_0.com.google.common.base.Preconditions.checkArgument;
@@ -40,6 +41,7 @@
import java.util.Map;
import java.util.Set;
import java.util.TimeZone;
+import java.util.TreeSet;
import java.util.stream.Collectors;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlPipelineOptions;
import org.apache.beam.sdk.extensions.sql.impl.JavaUdfLoader;
@@ -48,6 +50,7 @@
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.CharType;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.TimeWithLocalTzType;
+import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
import org.apache.beam.sdk.transforms.DoFn;
@@ -157,15 +160,11 @@
final RelOptPredicateList predicates = mq.getPulledUpPredicates(getInput());
final RexSimplify simplify = new RexSimplify(rexBuilder, predicates, RexUtil.EXECUTOR);
final RexProgram program = getProgram().normalize(rexBuilder, simplify);
+ final InputGetterImpl inputGetter = new InputGetterImpl(rowParam, upstream.getSchema());
Expression condition =
RexToLixTranslator.translateCondition(
- program,
- typeFactory,
- builder,
- new InputGetterImpl(rowParam, upstream.getSchema()),
- null,
- conformance);
+ program, typeFactory, builder, inputGetter, null, conformance);
List<Expression> expressions =
RexToLixTranslator.translateProjects(
@@ -175,7 +174,7 @@
builder,
physType,
DataContext.ROOT,
- new InputGetterImpl(rowParam, upstream.getSchema()),
+ inputGetter,
null);
builder.add(
@@ -192,10 +191,8 @@
builder.toBlock().toString(),
outputSchema,
options.getVerifyRowValues(),
- getJarPaths(program));
-
- // validate generated code
- calcFn.compile();
+ getJarPaths(program),
+ inputGetter.getFieldAccess());
return upstream.apply(ParDo.of(calcFn)).setRowSchema(outputSchema);
}
@@ -207,20 +204,29 @@
private final Schema outputSchema;
private final boolean verifyRowValues;
private final List<String> jarPaths;
+
+ @FieldAccess("row")
+ private final FieldAccessDescriptor fieldAccess;
+
private transient @Nullable ScriptEvaluator se = null;
public CalcFn(
String processElementBlock,
Schema outputSchema,
boolean verifyRowValues,
- List<String> jarPaths) {
+ List<String> jarPaths,
+ FieldAccessDescriptor fieldAccess) {
this.processElementBlock = processElementBlock;
this.outputSchema = outputSchema;
this.verifyRowValues = verifyRowValues;
this.jarPaths = jarPaths;
+ this.fieldAccess = fieldAccess;
+
+ // validate generated code
+ compile(processElementBlock, jarPaths);
}
- ScriptEvaluator compile() {
+ private static ScriptEvaluator compile(String processElementBlock, List<String> jarPaths) {
ScriptEvaluator se = new ScriptEvaluator();
if (!jarPaths.isEmpty()) {
try {
@@ -246,22 +252,22 @@
@Setup
public void setup() {
- this.se = compile();
+ this.se = compile(processElementBlock, jarPaths);
}
@ProcessElement
- public void processElement(ProcessContext c) {
+ public void processElement(@FieldAccess("row") Row row, OutputReceiver<Row> r) {
assert se != null;
final Object[] v;
try {
- v = (Object[]) se.evaluate(new Object[] {c.element(), CONTEXT_INSTANCE});
+ v = (Object[]) se.evaluate(new Object[] {row, CONTEXT_INSTANCE});
} catch (InvocationTargetException e) {
throw new RuntimeException(
"CalcFn failed to evaluate: " + processElementBlock, e.getCause());
}
if (v != null) {
- Row row = toBeamRow(Arrays.asList(v), outputSchema, verifyRowValues);
- c.output(row);
+ final Row output = toBeamRow(Arrays.asList(v), outputSchema, verifyRowValues);
+ r.output(output);
}
}
}
@@ -411,14 +417,21 @@
private final Expression input;
private final Schema inputSchema;
+ private final Set<Integer> referencedColumns;
private InputGetterImpl(Expression input, Schema inputSchema) {
this.input = input;
this.inputSchema = inputSchema;
+ this.referencedColumns = new TreeSet<>();
+ }
+
+ FieldAccessDescriptor getFieldAccess() {
+ return FieldAccessDescriptor.withFieldIds(this.referencedColumns);
}
@Override
public Expression field(BlockBuilder list, int index, Type storageType) {
+ this.referencedColumns.add(index);
return getBeamField(list, index, input, inputSchema);
}
@@ -431,64 +444,66 @@
final Expression expression = list.append(list.newName("current"), input);
- FieldType fieldType = schema.getField(index).getType();
- Expression value;
+ final Field field = schema.getField(index);
+ final FieldType fieldType = field.getType();
+ final Expression fieldName = Expressions.constant(field.getName());
+ final Expression value;
switch (fieldType.getTypeName()) {
case BYTE:
- value = Expressions.call(expression, "getByte", Expressions.constant(index));
+ value = Expressions.call(expression, "getByte", fieldName);
break;
case INT16:
- value = Expressions.call(expression, "getInt16", Expressions.constant(index));
+ value = Expressions.call(expression, "getInt16", fieldName);
break;
case INT32:
- value = Expressions.call(expression, "getInt32", Expressions.constant(index));
+ value = Expressions.call(expression, "getInt32", fieldName);
break;
case INT64:
- value = Expressions.call(expression, "getInt64", Expressions.constant(index));
+ value = Expressions.call(expression, "getInt64", fieldName);
break;
case DECIMAL:
- value = Expressions.call(expression, "getDecimal", Expressions.constant(index));
+ value = Expressions.call(expression, "getDecimal", fieldName);
break;
case FLOAT:
- value = Expressions.call(expression, "getFloat", Expressions.constant(index));
+ value = Expressions.call(expression, "getFloat", fieldName);
break;
case DOUBLE:
- value = Expressions.call(expression, "getDouble", Expressions.constant(index));
+ value = Expressions.call(expression, "getDouble", fieldName);
break;
case STRING:
- value = Expressions.call(expression, "getString", Expressions.constant(index));
+ value = Expressions.call(expression, "getString", fieldName);
break;
case DATETIME:
- value = Expressions.call(expression, "getDateTime", Expressions.constant(index));
+ value = Expressions.call(expression, "getDateTime", fieldName);
break;
case BOOLEAN:
- value = Expressions.call(expression, "getBoolean", Expressions.constant(index));
+ value = Expressions.call(expression, "getBoolean", fieldName);
break;
case BYTES:
- value = Expressions.call(expression, "getBytes", Expressions.constant(index));
+ value = Expressions.call(expression, "getBytes", fieldName);
break;
case ARRAY:
- value = Expressions.call(expression, "getArray", Expressions.constant(index));
+ value = Expressions.call(expression, "getArray", fieldName);
break;
case MAP:
- value = Expressions.call(expression, "getMap", Expressions.constant(index));
+ value = Expressions.call(expression, "getMap", fieldName);
break;
case ROW:
- value = Expressions.call(expression, "getRow", Expressions.constant(index));
+ value = Expressions.call(expression, "getRow", fieldName);
break;
case LOGICAL_TYPE:
String identifier = fieldType.getLogicalType().getIdentifier();
if (CharType.IDENTIFIER.equals(identifier)) {
- value = Expressions.call(expression, "getString", Expressions.constant(index));
+ value = Expressions.call(expression, "getString", fieldName);
} else if (TimeWithLocalTzType.IDENTIFIER.equals(identifier)) {
- value = Expressions.call(expression, "getDateTime", Expressions.constant(index));
+ value = Expressions.call(expression, "getDateTime", fieldName);
} else if (SqlTypes.DATE.getIdentifier().equals(identifier)) {
value =
Expressions.convert_(
Expressions.call(
expression,
"getLogicalTypeValue",
- Expressions.constant(index),
+ fieldName,
Expressions.constant(LocalDate.class)),
LocalDate.class);
} else if (SqlTypes.TIME.getIdentifier().equals(identifier)) {
@@ -497,7 +512,7 @@
Expressions.call(
expression,
"getLogicalTypeValue",
- Expressions.constant(index),
+ fieldName,
Expressions.constant(LocalTime.class)),
LocalTime.class);
} else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) {
@@ -506,7 +521,7 @@
Expressions.call(
expression,
"getLogicalTypeValue",
- Expressions.constant(index),
+ fieldName,
Expressions.constant(LocalDateTime.class)),
LocalDateTime.class);
} else {
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRelTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRelTest.java
index 27baad3..656a514 100644
--- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRelTest.java
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRelTest.java
@@ -18,20 +18,34 @@
package org.apache.beam.sdk.extensions.sql.impl.rel;
import java.math.BigDecimal;
+import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.extensions.sql.impl.BeamTableStatistics;
import org.apache.beam.sdk.extensions.sql.impl.planner.NodeStats;
import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestBoundedTable;
import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestUnboundedTable;
+import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.calcite.v1_26_0.org.apache.calcite.rel.RelNode;
import org.joda.time.DateTime;
import org.joda.time.Duration;
import org.junit.Assert;
import org.junit.BeforeClass;
+import org.junit.Rule;
import org.junit.Test;
/** Tests related to {@code BeamCalcRel}. */
public class BeamCalcRelTest extends BaseRelTest {
+
+ @Rule public final TestPipeline pipeline = TestPipeline.create();
+
private static final DateTime FIRST_DATE = new DateTime(1);
private static final DateTime SECOND_DATE = new DateTime(1 + 3600 * 1000);
@@ -160,4 +174,74 @@
Assert.assertTrue(doubleEqualEstimate.getRowCount() < equalEstimate.getRowCount());
Assert.assertTrue(doubleEqualEstimate.getWindow() < equalEstimate.getWindow());
}
+
+ private static class NodeGetter extends Pipeline.PipelineVisitor.Defaults {
+
+ private final PValue target;
+ private TransformHierarchy.Node producer;
+
+ private NodeGetter(PValue target) {
+ this.target = target;
+ }
+
+ @Override
+ public void visitValue(PValue value, TransformHierarchy.Node producer) {
+ if (value == target) {
+ assert this.producer == null;
+ this.producer = producer;
+ }
+ }
+ }
+
+ @Test
+ public void testSingleFieldAccess() throws IllegalAccessException {
+ String sql = "SELECT order_id FROM ORDER_DETAILS_BOUNDED";
+
+ PCollection<Row> rows = compilePipeline(sql, pipeline);
+
+ final NodeGetter nodeGetter = new NodeGetter(rows);
+ pipeline.traverseTopologically(nodeGetter);
+
+ ParDo.MultiOutput<Row, Row> pardo =
+ (ParDo.MultiOutput<Row, Row>) nodeGetter.producer.getTransform();
+ DoFnSignature sig = DoFnSignatures.getSignature(pardo.getFn().getClass());
+
+ Assert.assertEquals(1, sig.fieldAccessDeclarations().size());
+ DoFnSignature.FieldAccessDeclaration dec =
+ sig.fieldAccessDeclarations().values().iterator().next();
+ FieldAccessDescriptor fieldAccess = (FieldAccessDescriptor) dec.field().get(pardo.getFn());
+
+ Assert.assertTrue(fieldAccess.referencesSingleField());
+
+ fieldAccess =
+ fieldAccess.resolve(nodeGetter.producer.getInputs().values().iterator().next().getSchema());
+ Assert.assertEquals("order_id", fieldAccess.fieldNamesAccessed().iterator().next());
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testNoFieldAccess() throws IllegalAccessException {
+ String sql = "SELECT 1 FROM ORDER_DETAILS_BOUNDED";
+
+ PCollection<Row> rows = compilePipeline(sql, pipeline);
+
+ final NodeGetter nodeGetter = new NodeGetter(rows);
+ pipeline.traverseTopologically(nodeGetter);
+
+ ParDo.MultiOutput<Row, Row> pardo =
+ (ParDo.MultiOutput<Row, Row>) nodeGetter.producer.getTransform();
+ DoFnSignature sig = DoFnSignatures.getSignature(pardo.getFn().getClass());
+
+ Assert.assertEquals(1, sig.fieldAccessDeclarations().size());
+ DoFnSignature.FieldAccessDeclaration dec =
+ sig.fieldAccessDeclarations().values().iterator().next();
+ FieldAccessDescriptor fieldAccess = (FieldAccessDescriptor) dec.field().get(pardo.getFn());
+
+ Assert.assertFalse(fieldAccess.getAllFields());
+ Assert.assertTrue(fieldAccess.getFieldsAccessed().isEmpty());
+ Assert.assertTrue(fieldAccess.getNestedFieldsAccessed().isEmpty());
+
+ pipeline.run().waitUntilFinish();
+ }
}
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRel.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRel.java
index 38c604e..1d93ed3 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRel.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRel.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.extensions.sql.zetasql;
+import static org.apache.beam.sdk.schemas.Schema.Field;
import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
import com.google.auto.value.AutoValue;
@@ -39,6 +40,7 @@
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
import org.apache.beam.sdk.extensions.sql.meta.provider.bigquery.BeamBigQuerySqlDialect;
import org.apache.beam.sdk.extensions.sql.meta.provider.bigquery.BeamSqlUnparseContext;
+import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
@@ -142,9 +144,6 @@
options.getZetaSqlDefaultTimezone(),
options.getVerifyRowValues());
- // validate prepared expressions
- calcFn.setup();
-
return upstream.apply(ParDo.of(calcFn)).setRowSchema(outputSchema);
}
}
@@ -171,7 +170,11 @@
private final Schema outputSchema;
private final String defaultTimezone;
private final boolean verifyRowValues;
- private transient List<Integer> referencedColumns = ImmutableList.of();
+ private final List<Integer> referencedColumns;
+
+ @FieldAccess("row")
+ private final FieldAccessDescriptor fieldAccess;
+
private transient Map<BoundedWindow, Queue<TimestampedFuture>> pending = new HashMap<>();
private transient PreparedExpression exp;
private transient PreparedExpression.@Nullable Stream stream;
@@ -190,10 +193,21 @@
this.outputSchema = outputSchema;
this.defaultTimezone = defaultTimezone;
this.verifyRowValues = verifyRowValues;
+
+ try (PreparedExpression exp =
+ prepareExpression(sql, nullParams, inputSchema, defaultTimezone)) {
+ ImmutableList.Builder<Integer> columns = new ImmutableList.Builder<>();
+ for (String c : exp.getReferencedColumns()) {
+ columns.add(Integer.parseInt(c.substring(1)));
+ }
+ this.referencedColumns = columns.build();
+ this.fieldAccess = FieldAccessDescriptor.withFieldIds(this.referencedColumns);
+ }
}
/** exp cannot be reused and is transient so needs to be reinitialized. */
- private void prepareExpression() {
+ private static PreparedExpression prepareExpression(
+ String sql, Map<String, Value> nullParams, Schema inputSchema, String defaultTimezone) {
AnalyzerOptions options =
SqlAnalyzer.getAnalyzerOptions(QueryParameters.ofNamed(nullParams), defaultTimezone);
for (int i = 0; i < inputSchema.getFieldCount(); i++) {
@@ -202,21 +216,15 @@
ZetaSqlBeamTranslationUtils.toZetaSqlType(inputSchema.getField(i).getType()));
}
- exp = new PreparedExpression(sql);
+ PreparedExpression exp = new PreparedExpression(sql);
exp.prepare(options);
+ return exp;
}
@Setup
public void setup() {
- prepareExpression();
-
- ImmutableList.Builder<Integer> columns = new ImmutableList.Builder<>();
- for (String c : exp.getReferencedColumns()) {
- columns.add(Integer.parseInt(c.substring(1)));
- }
- referencedColumns = columns.build();
-
- stream = exp.stream();
+ this.exp = prepareExpression(sql, nullParams, inputSchema, defaultTimezone);
+ this.stream = exp.stream();
}
@StartBundle
@@ -231,14 +239,15 @@
@ProcessElement
public void processElement(
- @Element Row row, @Timestamp Instant t, BoundedWindow w, OutputReceiver<Row> r)
+ @FieldAccess("row") Row row, @Timestamp Instant t, BoundedWindow w, OutputReceiver<Row> r)
throws InterruptedException {
Map<String, Value> columns = new HashMap<>();
for (int i : referencedColumns) {
+ final Field field = inputSchema.getField(i);
columns.put(
columnName(i),
ZetaSqlBeamTranslationUtils.toZetaSqlValue(
- row.getBaseValue(i, Object.class), inputSchema.getField(i).getType()));
+ row.getBaseValue(field.getName(), Object.class), field.getType()));
}
@NonNull
diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRelTest.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRelTest.java
new file mode 100644
index 0000000..352e83a
--- /dev/null
+++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRelTest.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 org.apache.beam.sdk.extensions.sql.zetasql;
+
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
+import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
+import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils;
+import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.Row;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+/** Tests related to {@code BeamZetaSqlCalcRel}. */
+public class BeamZetaSqlCalcRelTest extends ZetaSqlTestBase {
+
+ private PCollection<Row> compile(String sql) {
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql, QueryParameters.ofNone());
+ return BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+ }
+
+ @Rule public final TestPipeline pipeline = TestPipeline.create();
+
+ @Before
+ public void setUp() {
+ initialize();
+ }
+
+ private static class NodeGetter extends Pipeline.PipelineVisitor.Defaults {
+
+ private final PValue target;
+ private TransformHierarchy.Node producer;
+
+ private NodeGetter(PValue target) {
+ this.target = target;
+ }
+
+ @Override
+ public void visitValue(PValue value, TransformHierarchy.Node producer) {
+ if (value == target) {
+ assert this.producer == null;
+ this.producer = producer;
+ }
+ }
+ }
+
+ @Test
+ public void testSingleFieldAccess() throws IllegalAccessException {
+ String sql = "SELECT Key FROM KeyValue";
+
+ PCollection<Row> rows = compile(sql);
+
+ final NodeGetter nodeGetter = new NodeGetter(rows);
+ pipeline.traverseTopologically(nodeGetter);
+
+ ParDo.MultiOutput<Row, Row> pardo =
+ (ParDo.MultiOutput<Row, Row>) nodeGetter.producer.getTransform();
+ DoFnSignature sig = DoFnSignatures.getSignature(pardo.getFn().getClass());
+
+ Assert.assertEquals(1, sig.fieldAccessDeclarations().size());
+ DoFnSignature.FieldAccessDeclaration dec =
+ sig.fieldAccessDeclarations().values().iterator().next();
+ FieldAccessDescriptor fieldAccess = (FieldAccessDescriptor) dec.field().get(pardo.getFn());
+
+ Assert.assertTrue(fieldAccess.referencesSingleField());
+
+ fieldAccess =
+ fieldAccess.resolve(nodeGetter.producer.getInputs().values().iterator().next().getSchema());
+ Assert.assertEquals("Key", fieldAccess.fieldNamesAccessed().iterator().next());
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testNoFieldAccess() throws IllegalAccessException {
+ String sql = "SELECT 1 FROM KeyValue";
+
+ PCollection<Row> rows = compile(sql);
+
+ final NodeGetter nodeGetter = new NodeGetter(rows);
+ pipeline.traverseTopologically(nodeGetter);
+
+ ParDo.MultiOutput<Row, Row> pardo =
+ (ParDo.MultiOutput<Row, Row>) nodeGetter.producer.getTransform();
+ DoFnSignature sig = DoFnSignatures.getSignature(pardo.getFn().getClass());
+
+ Assert.assertEquals(1, sig.fieldAccessDeclarations().size());
+ DoFnSignature.FieldAccessDeclaration dec =
+ sig.fieldAccessDeclarations().values().iterator().next();
+ FieldAccessDescriptor fieldAccess = (FieldAccessDescriptor) dec.field().get(pardo.getFn());
+
+ Assert.assertFalse(fieldAccess.getAllFields());
+ Assert.assertTrue(fieldAccess.getFieldsAccessed().isEmpty());
+ Assert.assertTrue(fieldAccess.getNestedFieldsAccessed().isEmpty());
+
+ pipeline.run().waitUntilFinish();
+ }
+}
diff --git a/sdks/java/fn-execution/build.gradle b/sdks/java/fn-execution/build.gradle
index 8cb5f0a..2d495ee 100644
--- a/sdks/java/fn-execution/build.gradle
+++ b/sdks/java/fn-execution/build.gradle
@@ -34,8 +34,10 @@
compile library.java.slf4j_api
compile library.java.joda_time
provided library.java.junit
+ testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
testCompile library.java.junit
testCompile library.java.mockito_core
testCompile library.java.commons_lang3
+ testCompile "com.github.stefanbirkner:system-rules:1.19.0"
testRuntimeOnly library.java.slf4j_jdk14
}
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/JvmInitializers.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/JvmInitializers.java
index 221ebe7..ad40817 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/JvmInitializers.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/JvmInitializers.java
@@ -20,6 +20,8 @@
import org.apache.beam.sdk.harness.JvmInitializer;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.common.ReflectHelpers;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/** Helpers for executing {@link JvmInitializer} implementations. */
public class JvmInitializers {
@@ -30,6 +32,8 @@
*/
public static void runOnStartup() {
for (JvmInitializer initializer : ReflectHelpers.loadServicesOrdered(JvmInitializer.class)) {
+ // We write to standard out since logging has yet to be initialized.
+ System.out.format("Running JvmInitializer#onStartup for %s%n", initializer);
initializer.onStartup();
}
}
@@ -42,7 +46,11 @@
* @param options The pipeline options passed to the worker.
*/
public static void runBeforeProcessing(PipelineOptions options) {
+ // We load the logger in the the method to minimize the amount of class loading that happens
+ // during class initialization.
+ Logger logger = LoggerFactory.getLogger(JvmInitializers.class);
for (JvmInitializer initializer : ReflectHelpers.loadServicesOrdered(JvmInitializer.class)) {
+ logger.info("Running JvmInitializer#beforeProcessing for {}", initializer);
initializer.beforeProcessing(options);
}
}
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/JvmInitializersTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/JvmInitializersTest.java
index a20b3a9..e1da00a 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/JvmInitializersTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/JvmInitializersTest.java
@@ -17,15 +17,20 @@
*/
package org.apache.beam.sdk.fn;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.core.StringContains.containsString;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import com.google.auto.service.AutoService;
import org.apache.beam.sdk.harness.JvmInitializer;
import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.testing.ExpectedLogs;
import org.apache.beam.sdk.testing.TestPipeline;
import org.junit.Before;
+import org.junit.Rule;
import org.junit.Test;
+import org.junit.contrib.java.lang.system.SystemOutRule;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -33,6 +38,9 @@
@RunWith(JUnit4.class)
public final class JvmInitializersTest {
+ @Rule public ExpectedLogs expectedLogs = ExpectedLogs.none(JvmInitializers.class);
+ @Rule public SystemOutRule systemOutRule = new SystemOutRule().enableLog();
+
private static Boolean onStartupRan;
private static Boolean beforeProcessingRan;
private static PipelineOptions receivedOptions;
@@ -64,6 +72,7 @@
JvmInitializers.runOnStartup();
assertTrue(onStartupRan);
+ assertThat(systemOutRule.getLog(), containsString("Running JvmInitializer#onStartup"));
}
@Test
@@ -74,5 +83,6 @@
assertTrue(beforeProcessingRan);
assertEquals(options, receivedOptions);
+ expectedLogs.verifyInfo("Running JvmInitializer#beforeProcessing");
}
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java
index 90b5fed..1aa5d52 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/BeamFnControlClient.java
@@ -23,6 +23,7 @@
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
+import org.apache.beam.fn.harness.logging.BeamFnLoggingMDC;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnControlGrpc;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
@@ -102,18 +103,28 @@
}
@Override
- public void onNext(BeamFnApi.InstructionRequest value) {
- LOG.debug("Received InstructionRequest {}", value);
- executor.execute(
- () -> {
- try {
- BeamFnApi.InstructionResponse response = delegateOnInstructionRequestType(value);
- sendInstructionResponse(response);
- } catch (Error e) {
- sendErrorResponse(e);
- throw e;
- }
- });
+ public void onNext(BeamFnApi.InstructionRequest request) {
+ try {
+ BeamFnLoggingMDC.setInstructionId(request.getInstructionId());
+ LOG.debug("Received InstructionRequest {}", request);
+ executor.execute(
+ () -> {
+ try {
+ // Ensure that we set and clear the MDC since processing the request will occur
+ // in a separate thread.
+ BeamFnLoggingMDC.setInstructionId(request.getInstructionId());
+ BeamFnApi.InstructionResponse response = delegateOnInstructionRequestType(request);
+ sendInstructionResponse(response);
+ } catch (Error e) {
+ sendErrorResponse(e);
+ throw e;
+ } finally {
+ BeamFnLoggingMDC.setInstructionId(null);
+ }
+ });
+ } finally {
+ BeamFnLoggingMDC.setInstructionId(null);
+ }
}
@Override
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
index 9cceb72..74ddab5 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
@@ -48,7 +48,6 @@
import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.data.QueueingBeamFnDataClient;
-import org.apache.beam.fn.harness.logging.BeamFnLoggingMDC;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache;
import org.apache.beam.fn.harness.state.CachingBeamFnStateClient;
@@ -335,7 +334,6 @@
}
});
try {
- BeamFnLoggingMDC.setInstructionId(request.getInstructionId());
PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry();
PTransformFunctionRegistry finishFunctionRegistry =
bundleProcessor.getFinishFunctionRegistry();
@@ -390,8 +388,6 @@
// Make sure we clean-up from the active set of bundle processors.
bundleProcessorCache.discard(bundleProcessor);
throw e;
- } finally {
- BeamFnLoggingMDC.setInstructionId(null);
}
}
@@ -852,13 +848,13 @@
@Override
@SuppressWarnings("FutureReturnValueIgnored") // async arriveAndDeregister task doesn't need
// monitoring.
- public void handle(
- StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> response) {
+ public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
// Register each request with the phaser and arrive and deregister each time a request
// completes.
+ CompletableFuture<StateResponse> response = beamFnStateClient.handle(requestBuilder);
phaser.register();
response.whenComplete((stateResponse, throwable) -> phaser.arriveAndDeregister());
- beamFnStateClient.handle(requestBuilder, response);
+ return response;
}
}
@@ -879,7 +875,7 @@
}
@Override
- public void handle(Builder requestBuilder, CompletableFuture<StateResponse> response) {
+ public CompletableFuture<StateResponse> handle(Builder requestBuilder) {
throw new IllegalStateException(
String.format(
"State API calls are unsupported because the "
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java
index 777036a..76664ba 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java
@@ -21,7 +21,6 @@
import java.util.ArrayList;
import java.util.Collections;
-import java.util.concurrent.CompletableFuture;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateAppendRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateClearRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
@@ -115,6 +114,7 @@
newValues = new ArrayList<>();
}
+ @SuppressWarnings("FutureReturnValueIgnored")
public void asyncClose() throws Exception {
checkState(
!isClosed,
@@ -122,8 +122,7 @@
request.getStateKey());
if (oldValues == null) {
beamFnStateClient.handle(
- request.toBuilder().setClear(StateClearRequest.getDefaultInstance()),
- new CompletableFuture<>());
+ request.toBuilder().setClear(StateClearRequest.getDefaultInstance()));
}
if (!newValues.isEmpty()) {
ByteString.Output out = ByteString.newOutput();
@@ -134,8 +133,7 @@
beamFnStateClient.handle(
request
.toBuilder()
- .setAppend(StateAppendRequest.newBuilder().setData(out.toByteString())),
- new CompletableFuture<>());
+ .setAppend(StateAppendRequest.newBuilder().setData(out.toByteString())));
}
isClosed = true;
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateClient.java
index ffef28d..4cf03c4 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateClient.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateClient.java
@@ -31,9 +31,6 @@
* Consumes a state request populating a unique id returning a future to the response.
*
* @param requestBuilder A partially completed state request. The id will be populated the client.
- * @param response A future containing a corresponding {@link StateResponse} for the supplied
- * request.
*/
- void handle(
- BeamFnApi.StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> response);
+ CompletableFuture<StateResponse> handle(BeamFnApi.StateRequest.Builder requestBuilder);
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
index 633db1d..629c1fd 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCache.java
@@ -93,15 +93,16 @@
}
@Override
- public void handle(
- StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> response) {
+ public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
requestBuilder.setId(idGenerator.getId());
StateRequest request = requestBuilder.build();
+ CompletableFuture<StateResponse> response = new CompletableFuture<>();
outstandingRequests.put(request.getId(), response);
// If the server closes, gRPC will throw an error if onNext is called.
LOG.debug("Sending StateRequest {}", request);
outboundObserver.onNext(request);
+ return response;
}
private synchronized void closeAndCleanUp(RuntimeException cause) {
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClient.java
index 888d231..0f39b45 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClient.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClient.java
@@ -76,16 +76,14 @@
*/
@Override
@SuppressWarnings("FutureReturnValueIgnored")
- public void handle(
- StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> response) {
+ public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
StateKey stateKey = requestBuilder.getStateKey();
ByteString cacheToken = getCacheToken(stateKey);
// If state is not cacheable proceed as normal.
if (ByteString.EMPTY.equals(cacheToken)) {
- beamFnStateClient.handle(requestBuilder, response);
- return;
+ return beamFnStateClient.handle(requestBuilder);
}
switch (requestBuilder.getRequestCase()) {
@@ -98,37 +96,38 @@
// If data is not cached, add callback to add response to cache on completion.
// Otherwise, complete the response with the cached data.
+ CompletableFuture<StateResponse> response;
if (cachedPage == null) {
+ response = beamFnStateClient.handle(requestBuilder);
response.thenAccept(
stateResponse ->
stateCache.getUnchecked(stateKey).put(cacheKey, stateResponse.getGet()));
- beamFnStateClient.handle(requestBuilder, response);
} else {
- response.complete(
+ return CompletableFuture.completedFuture(
StateResponse.newBuilder().setId(requestBuilder.getId()).setGet(cachedPage).build());
}
- return;
+ return response;
case APPEND:
// TODO(BEAM-12637): Support APPEND in CachingBeamFnStateClient.
- beamFnStateClient.handle(requestBuilder, response);
+ response = beamFnStateClient.handle(requestBuilder);
// Invalidate last page of cached values (entry with a blank continuation token response)
Map<StateCacheKey, StateGetResponse> map = stateCache.getUnchecked(stateKey);
map.entrySet()
.removeIf(entry -> (entry.getValue().getContinuationToken().equals(ByteString.EMPTY)));
- return;
+ return response;
case CLEAR:
// Remove all state key data and replace with an empty response.
- beamFnStateClient.handle(requestBuilder, response);
+ response = beamFnStateClient.handle(requestBuilder);
Map<StateCacheKey, StateGetResponse> clearedData = new HashMap<>();
StateCacheKey newKey = StateCacheKey.create(cacheToken, ByteString.EMPTY);
clearedData.put(newKey, StateGetResponse.getDefaultInstance());
stateCache.put(stateKey, clearedData);
- return;
+ return response;
default:
throw new IllegalStateException(
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java
index 49efa35..f679608 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java
@@ -28,7 +28,6 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
-import java.util.concurrent.CompletableFuture;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateAppendRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateClearRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
@@ -195,6 +194,7 @@
}
@SuppressWarnings({
+ "FutureReturnValueIgnored",
"nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-12687)
})
// Update data in persistent store
@@ -211,27 +211,30 @@
// Clear currently persisted key-values
if (isCleared) {
- beamFnStateClient.handle(
- keysStateRequest.toBuilder().setClear(StateClearRequest.getDefaultInstance()),
- new CompletableFuture<>());
+ beamFnStateClient
+ .handle(keysStateRequest.toBuilder().setClear(StateClearRequest.getDefaultInstance()))
+ .get();
} else if (!pendingRemoves.isEmpty()) {
for (K key : pendingRemoves) {
- beamFnStateClient.handle(
- createUserStateRequest(key)
- .toBuilder()
- .setClear(StateClearRequest.getDefaultInstance()),
- new CompletableFuture<>());
+ beamFnStateClient
+ .handle(
+ createUserStateRequest(key)
+ .toBuilder()
+ .setClear(StateClearRequest.getDefaultInstance()))
+ .get();
}
}
// Persist pending key-values
if (!pendingAdds.isEmpty()) {
for (Map.Entry<K, List<V>> entry : pendingAdds.entrySet()) {
- beamFnStateClient.handle(
- createUserStateRequest(entry.getKey())
- .toBuilder()
- .setAppend(StateAppendRequest.newBuilder().setData(encodeValues(entry.getValue()))),
- new CompletableFuture<>());
+ beamFnStateClient
+ .handle(
+ createUserStateRequest(entry.getKey())
+ .toBuilder()
+ .setAppend(
+ StateAppendRequest.newBuilder().setData(encodeValues(entry.getValue()))))
+ .get();
}
}
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
index 1026ba5..1b8d43a 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
@@ -207,10 +207,9 @@
private void prefetchFirstPage() {
if (firstPageResponseFuture == null) {
- firstPageResponseFuture = new CompletableFuture<>();
- beamFnStateClient.handle(
- stateRequestForFirstChunk.toBuilder().setGet(stateRequestForFirstChunk.getGet()),
- firstPageResponseFuture);
+ firstPageResponseFuture =
+ beamFnStateClient.handle(
+ stateRequestForFirstChunk.toBuilder().setGet(stateRequestForFirstChunk.getGet()));
}
}
}
@@ -256,12 +255,11 @@
@Override
public void prefetch() {
if (currentState == State.READ_REQUIRED && prefetchedResponse == null) {
- prefetchedResponse = new CompletableFuture<>();
- beamFnStateClient.handle(
- stateRequestForFirstChunk
- .toBuilder()
- .setGet(StateGetRequest.newBuilder().setContinuationToken(continuationToken)),
- prefetchedResponse);
+ prefetchedResponse =
+ beamFnStateClient.handle(
+ stateRequestForFirstChunk
+ .toBuilder()
+ .setGet(StateGetRequest.newBuilder().setContinuationToken(continuationToken)));
}
}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java
index 404c79c..e13c934 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/BeamFnControlClientTest.java
@@ -33,6 +33,8 @@
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.beam.fn.harness.logging.BeamFnLoggingMDC;
+import org.apache.beam.fn.harness.logging.RestoreBeamFnLoggingMDC;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RegisterRequest;
@@ -47,7 +49,9 @@
import org.apache.beam.vendor.grpc.v1p36p0.io.grpc.stub.CallStreamObserver;
import org.apache.beam.vendor.grpc.v1p36p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
+import org.junit.Rule;
import org.junit.Test;
+import org.junit.rules.TestRule;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -85,6 +89,8 @@
.setError(getStackTraceAsString(FAILURE))
.build();
+ @Rule public TestRule restoreMDCAfterTest = new RestoreBeamFnLoggingMDC();
+
@Test
public void testDelegation() throws Exception {
AtomicBoolean clientClosedStream = new AtomicBoolean();
@@ -120,12 +126,15 @@
handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
handlers.put(
BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE,
- value ->
- BeamFnApi.InstructionResponse.newBuilder()
- .setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance()));
+ value -> {
+ assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
+ return BeamFnApi.InstructionResponse.newBuilder()
+ .setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance());
+ });
handlers.put(
BeamFnApi.InstructionRequest.RequestCase.REGISTER,
value -> {
+ assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
throw FAILURE;
});
@@ -199,6 +208,7 @@
handlers.put(
BeamFnApi.InstructionRequest.RequestCase.REGISTER,
value -> {
+ assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
throw new Error("Test Error");
});
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
index f947cd6..a82e104 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
@@ -960,8 +960,8 @@
.build();
Map<String, Message> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
- CompletableFuture<StateResponse> successfulResponse = new CompletableFuture<>();
- CompletableFuture<StateResponse> unsuccessfulResponse = new CompletableFuture<>();
+ CompletableFuture<StateResponse>[] successfulResponse = new CompletableFuture[1];
+ CompletableFuture<StateResponse>[] unsuccessfulResponse = new CompletableFuture[1];
BeamFnStateGrpcClientCache mockBeamFnStateGrpcClient =
Mockito.mock(BeamFnStateGrpcClientCache.class);
@@ -973,8 +973,7 @@
invocation -> {
StateRequest.Builder stateRequestBuilder =
(StateRequest.Builder) invocation.getArguments()[0];
- CompletableFuture<StateResponse> completableFuture =
- (CompletableFuture<StateResponse>) invocation.getArguments()[1];
+ CompletableFuture<StateResponse> completableFuture = new CompletableFuture<>();
new Thread(
() -> {
// Simulate sleeping which introduces a race which most of the time requires
@@ -990,10 +989,10 @@
}
})
.start();
- return null;
+ return completableFuture;
})
.when(mockBeamFnStateClient)
- .handle(any(), any());
+ .handle(any());
ProcessBundleHandler handler =
new ProcessBundleHandler(
@@ -1034,10 +1033,12 @@
}
private void doStateCalls(BeamFnStateClient beamFnStateClient) {
- beamFnStateClient.handle(
- StateRequest.newBuilder().setInstructionId("SUCCESS"), successfulResponse);
- beamFnStateClient.handle(
- StateRequest.newBuilder().setInstructionId("FAIL"), unsuccessfulResponse);
+ successfulResponse[0] =
+ beamFnStateClient.handle(
+ StateRequest.newBuilder().setInstructionId("SUCCESS"));
+ unsuccessfulResponse[0] =
+ beamFnStateClient.handle(
+ StateRequest.newBuilder().setInstructionId("FAIL"));
}
}),
new BundleProcessorCache());
@@ -1047,8 +1048,8 @@
BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L"))
.build());
- assertTrue(successfulResponse.isDone());
- assertTrue(unsuccessfulResponse.isDone());
+ assertTrue(successfulResponse[0].isDone());
+ assertTrue(unsuccessfulResponse[0].isDone());
}
@Test
@@ -1101,10 +1102,9 @@
return null;
}
+ @SuppressWarnings("FutureReturnValueIgnored")
private void doStateCalls(BeamFnStateClient beamFnStateClient) {
- beamFnStateClient.handle(
- StateRequest.newBuilder().setInstructionId("SUCCESS"),
- new CompletableFuture<>());
+ beamFnStateClient.handle(StateRequest.newBuilder().setInstructionId("SUCCESS"));
}
}),
new BundleProcessorCache());
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java
index 75a871a..4badcf9 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BeamFnStateGrpcClientCacheTest.java
@@ -120,11 +120,10 @@
public void testRequestResponses() throws Exception {
BeamFnStateClient client = clientCache.forApiServiceDescriptor(apiServiceDescriptor);
- CompletableFuture<StateResponse> successfulResponse = new CompletableFuture<>();
- CompletableFuture<StateResponse> unsuccessfulResponse = new CompletableFuture<>();
-
- client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), successfulResponse);
- client.handle(StateRequest.newBuilder().setInstructionId(FAIL), unsuccessfulResponse);
+ CompletableFuture<StateResponse> successfulResponse =
+ client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS));
+ CompletableFuture<StateResponse> unsuccessfulResponse =
+ client.handle(StateRequest.newBuilder().setInstructionId(FAIL));
// Wait for the client to connect.
StreamObserver<StateResponse> outboundServerObserver = outboundServerObservers.take();
@@ -149,8 +148,8 @@
public void testServerErrorCausesPendingAndFutureCallsToFail() throws Exception {
BeamFnStateClient client = clientCache.forApiServiceDescriptor(apiServiceDescriptor);
- CompletableFuture<StateResponse> inflight = new CompletableFuture<>();
- client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), inflight);
+ CompletableFuture<StateResponse> inflight =
+ client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS));
// Wait for the client to connect.
StreamObserver<StateResponse> outboundServerObserver = outboundServerObservers.take();
@@ -166,8 +165,8 @@
}
// Send a response after the client will have received an error.
- CompletableFuture<StateResponse> late = new CompletableFuture<>();
- client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), late);
+ CompletableFuture<StateResponse> late =
+ client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS));
try {
inflight.get();
@@ -181,8 +180,8 @@
public void testServerCompletionCausesPendingAndFutureCallsToFail() throws Exception {
BeamFnStateClient client = clientCache.forApiServiceDescriptor(apiServiceDescriptor);
- CompletableFuture<StateResponse> inflight = new CompletableFuture<>();
- client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), inflight);
+ CompletableFuture<StateResponse> inflight =
+ client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS));
// Wait for the client to connect.
StreamObserver<StateResponse> outboundServerObserver = outboundServerObservers.take();
@@ -197,8 +196,8 @@
}
// Send a response after the client will have received an error.
- CompletableFuture<StateResponse> late = new CompletableFuture<>();
- client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), late);
+ CompletableFuture<StateResponse> late =
+ client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS));
try {
inflight.get();
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClientTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClientTest.java
index 8604934..8f14f3b 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClientTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/CachingBeamFnStateClientTest.java
@@ -84,19 +84,16 @@
CachingBeamFnStateClient cachingClient =
new CachingBeamFnStateClient(fakeClient, stateCache, cacheTokenList);
- CompletableFuture<BeamFnApi.StateResponse> response1 = new CompletableFuture<>();
- CompletableFuture<BeamFnApi.StateResponse> response2 = new CompletableFuture<>();
-
StateRequest.Builder request =
StateRequest.newBuilder()
.setStateKey(key("A"))
.setGet(BeamFnApi.StateGetRequest.newBuilder().build());
- cachingClient.handle(request, response1);
+ CompletableFuture<BeamFnApi.StateResponse> response1 = cachingClient.handle(request);
assertEquals(1, fakeClient.getCallCount());
request.clearId();
- cachingClient.handle(request, response2);
+ CompletableFuture<BeamFnApi.StateResponse> response2 = cachingClient.handle(request);
assertEquals(2, fakeClient.getCallCount());
}
@@ -301,8 +298,7 @@
.setStateKey(key)
.setAppend(StateAppendRequest.newBuilder().setData(data));
- CompletableFuture<StateResponse> appendResponse = new CompletableFuture<>();
- cachingClient.handle(appendRequestBuilder, appendResponse);
+ CompletableFuture<StateResponse> appendResponse = cachingClient.handle(appendRequestBuilder);
appendResponse.get();
}
@@ -310,8 +306,7 @@
StateRequest.Builder clearRequestBuilder =
StateRequest.newBuilder().setStateKey(key).setClear(StateClearRequest.getDefaultInstance());
- CompletableFuture<StateResponse> clearResponse = new CompletableFuture<>();
- cachingClient.handle(clearRequestBuilder, clearResponse);
+ CompletableFuture<StateResponse> clearResponse = cachingClient.handle(clearRequestBuilder);
clearResponse.get();
}
@@ -324,8 +319,7 @@
requestBuilder
.clearId()
.setGet(StateGetRequest.newBuilder().setContinuationToken(continuationToken));
- CompletableFuture<StateResponse> getResponse = new CompletableFuture<>();
- cachingClient.handle(requestBuilder, getResponse);
+ CompletableFuture<StateResponse> getResponse = cachingClient.handle(requestBuilder);
continuationToken = getResponse.get().getGet().getContinuationToken();
allData = allData.concat(getResponse.get().getGet().getData());
} while (!continuationToken.equals(ByteString.EMPTY));
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/FakeBeamFnStateClient.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/FakeBeamFnStateClient.java
index 64c4292..04f05ae 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/FakeBeamFnStateClient.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/FakeBeamFnStateClient.java
@@ -73,8 +73,8 @@
}
@Override
- public void handle(
- StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> responseFuture) {
+ public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
+
// The id should never be filled out
assertEquals("", requestBuilder.getId());
requestBuilder.setId(generateId());
@@ -138,7 +138,8 @@
String.format("Unknown request type %s", request.getRequestCase()));
}
- responseFuture.complete(response.setId(requestBuilder.getId()).build());
+ CompletableFuture<StateResponse> responseFuture = new CompletableFuture<>();
+ return CompletableFuture.completedFuture(response.setId(requestBuilder.getId()).build());
}
private String generateId() {
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java
index 384d2df..ae897cb 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java
@@ -52,15 +52,14 @@
private static BeamFnStateClient fakeStateClient(
AtomicInteger callCount, ByteString... expected) {
- return (requestBuilder, response) -> {
+ return requestBuilder -> {
callCount.incrementAndGet();
if (expected.length == 0) {
- response.complete(
+ return CompletableFuture.completedFuture(
StateResponse.newBuilder()
.setId(requestBuilder.getId())
.setGet(StateGetResponse.newBuilder())
.build());
- return;
}
ByteString continuationToken = requestBuilder.getGet().getContinuationToken();
@@ -75,7 +74,7 @@
if (requestedPosition != expected.length - 1) {
newContinuationToken = ByteString.copyFromUtf8(Integer.toString(requestedPosition + 1));
}
- response.complete(
+ return CompletableFuture.completedFuture(
StateResponse.newBuilder()
.setId(requestBuilder.getId())
.setGet(
@@ -121,15 +120,49 @@
ByteString.EMPTY);
}
+ private BeamFnStateClient fakeStateClient(AtomicInteger callCount, ByteString... expected) {
+ return (requestBuilder) -> {
+ callCount.incrementAndGet();
+ if (expected.length == 0) {
+ return CompletableFuture.completedFuture(
+ StateResponse.newBuilder()
+ .setId(requestBuilder.getId())
+ .setGet(StateGetResponse.newBuilder())
+ .build());
+ }
+
+ ByteString continuationToken = requestBuilder.getGet().getContinuationToken();
+
+ int requestedPosition = 0; // Default position is 0
+ if (!ByteString.EMPTY.equals(continuationToken)) {
+ requestedPosition = Integer.parseInt(continuationToken.toStringUtf8());
+ }
+
+ // Compute the new continuation token
+ ByteString newContinuationToken = ByteString.EMPTY;
+ if (requestedPosition != expected.length - 1) {
+ newContinuationToken = ByteString.copyFromUtf8(Integer.toString(requestedPosition + 1));
+ }
+ return CompletableFuture.completedFuture(
+ StateResponse.newBuilder()
+ .setId(requestBuilder.getId())
+ .setGet(
+ StateGetResponse.newBuilder()
+ .setData(expected[requestedPosition])
+ .setContinuationToken(newContinuationToken))
+ .build());
+ };
+ }
+
@Test
public void testPrefetchIgnoredWhenExistingPrefetchOngoing() throws Exception {
AtomicInteger callCount = new AtomicInteger();
BeamFnStateClient fakeStateClient =
new BeamFnStateClient() {
@Override
- public void handle(
- StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> response) {
+ public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
callCount.incrementAndGet();
+ return new CompletableFuture<StateResponse>();
}
};
PrefetchableIterator<ByteString> byteStrings =
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpointTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpointTest.java
index aab5c1d..ca23fd7 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpointTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpointTest.java
@@ -27,7 +27,7 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
/** * */
@RunWith(MockitoJUnitRunner.class)
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/RecordFilterTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/RecordFilterTest.java
index 0058e2e..05304bc 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/RecordFilterTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/RecordFilterTest.java
@@ -26,7 +26,7 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
/** * */
@RunWith(MockitoJUnitRunner.class)
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardCheckpointTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardCheckpointTest.java
index aa4b21b..1c7407a 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardCheckpointTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardCheckpointTest.java
@@ -36,7 +36,7 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber;
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java
index 83b3e54..f7fd2ff 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java
@@ -35,7 +35,7 @@
import org.junit.runner.RunWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
import org.mockito.stubbing.Answer;
import software.amazon.awssdk.core.SdkBytes;
import software.amazon.awssdk.core.exception.SdkClientException;
diff --git a/sdks/java/io/debezium/build.gradle b/sdks/java/io/debezium/build.gradle
index 722528f..692736b 100644
--- a/sdks/java/io/debezium/build.gradle
+++ b/sdks/java/io/debezium/build.gradle
@@ -54,8 +54,8 @@
permitUnusedDeclared "org.apache.kafka:connect-json:2.5.0" // BEAM-11761
// Debezium dependencies
- compile group: 'io.debezium', name: 'debezium-core', version: '1.3.1.Final'
- testCompile group: 'io.debezium', name: 'debezium-connector-mysql', version: '1.3.1.Final'
+ compile group: 'io.debezium', name: 'debezium-core', version: '1.7.0.Final'
+ testCompile group: 'io.debezium', name: 'debezium-connector-mysql', version: '1.7.0.Final'
}
test {
diff --git a/sdks/java/io/debezium/expansion-service/build.gradle b/sdks/java/io/debezium/expansion-service/build.gradle
index a183c91..f3a0a30 100644
--- a/sdks/java/io/debezium/expansion-service/build.gradle
+++ b/sdks/java/io/debezium/expansion-service/build.gradle
@@ -38,7 +38,7 @@
runtime library.java.slf4j_jdk14
// Debezium runtime dependencies
- def debezium_version = '1.3.1.Final'
+ def debezium_version = '1.7.0.Final'
runtimeOnly group: 'io.debezium', name: 'debezium-connector-mysql', version: debezium_version
runtimeOnly group: 'io.debezium', name: 'debezium-connector-postgres', version: debezium_version
runtimeOnly group: 'io.debezium', name: 'debezium-connector-sqlserver', version: debezium_version
diff --git a/sdks/java/io/debezium/src/README.md b/sdks/java/io/debezium/src/README.md
index 4cf9be8..14356d6 100644
--- a/sdks/java/io/debezium/src/README.md
+++ b/sdks/java/io/debezium/src/README.md
@@ -25,7 +25,7 @@
### Getting Started
-DebeziumIO uses [Debezium Connectors v1.3](https://debezium.io/documentation/reference/1.3/connectors/) to connect to Apache Beam. All you need to do is choose the Debezium Connector that suits your Debezium setup and pick a [Serializable Function](https://beam.apache.org/releases/javadoc/2.23.0/org/apache/beam/sdk/transforms/SerializableFunction.html), then you will be able to connect to Apache Beam and start building your own Pipelines.
+DebeziumIO uses [Debezium Connectors v1.7](https://debezium.io/documentation/reference/1.7/connectors/) to connect to Apache Beam. All you need to do is choose the Debezium Connector that suits your Debezium setup and pick a [Serializable Function](https://beam.apache.org/releases/javadoc/2.23.0/org/apache/beam/sdk/transforms/SerializableFunction.html), then you will be able to connect to Apache Beam and start building your own Pipelines.
These connectors have been successfully tested and are known to work fine:
* MySQL Connector
@@ -65,7 +65,7 @@
|Method|Params|Description|
|-|-|-|
|`.withConnectionProperty(propName, propValue)`|_String_, _String_|Adds a custom property to the connector.|
-> **Note:** For more information on custom properties, see your [Debezium Connector](https://debezium.io/documentation/reference/1.3/connectors/) specific documentation.
+> **Note:** For more information on custom properties, see your [Debezium Connector](https://debezium.io/documentation/reference/1.7/connectors/) specific documentation.
Example of a MySQL Debezium Connector setup:
```
@@ -165,7 +165,7 @@
### Requirements and Supported versions
- JDK v8
-- Debezium Connectors v1.3
+- Debezium Connectors v1.7
- Apache Beam 2.25
## Running Unit Tests
diff --git a/sdks/java/io/debezium/src/test/java/org/apache/beam/io/debezium/DebeziumIOMySqlConnectorIT.java b/sdks/java/io/debezium/src/test/java/org/apache/beam/io/debezium/DebeziumIOMySqlConnectorIT.java
index 6056ca0..c9b1baf 100644
--- a/sdks/java/io/debezium/src/test/java/org/apache/beam/io/debezium/DebeziumIOMySqlConnectorIT.java
+++ b/sdks/java/io/debezium/src/test/java/org/apache/beam/io/debezium/DebeziumIOMySqlConnectorIT.java
@@ -89,7 +89,7 @@
.withMaxNumberOfRecords(30)
.withCoder(StringUtf8Coder.of()));
String expected =
- "{\"metadata\":{\"connector\":\"mysql\",\"version\":\"1.3.1.Final\",\"name\":\"dbserver1\","
+ "{\"metadata\":{\"connector\":\"mysql\",\"version\":\"1.7.0.Final\",\"name\":\"dbserver1\","
+ "\"database\":\"inventory\",\"schema\":\"mysql-bin.000003\",\"table\":\"addresses\"},\"before\":null,"
+ "\"after\":{\"fields\":{\"zip\":\"76036\",\"city\":\"Euless\","
+ "\"street\":\"3183 Moore Avenue\",\"id\":10,\"state\":\"Texas\",\"customer_id\":1001,"
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
index 16b96bf..aca0c69 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
@@ -696,7 +696,9 @@
kmsKey,
rowWriterFactory.getSourceFormat(),
useAvroLogicalTypes,
- schemaUpdateOptions))
+ // Note that we can't pass through the schema update options when creating temporary
+ // tables. They also shouldn't be needed. See BEAM-12482 for additional details.
+ Collections.emptySet()))
.setCoder(KvCoder.of(tableDestinationCoder, WriteTables.ResultCoder.INSTANCE));
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
index 8b9b705..66f77ab 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -2627,7 +2627,7 @@
if (getJsonTableRef() != null) {
dynamicDestinations =
DynamicDestinationsHelpers.ConstantTableDestinations.fromJsonTableRef(
- getJsonTableRef(), getTableDescription());
+ getJsonTableRef(), getTableDescription(), getClustering() != null);
} else if (getTableFunction() != null) {
dynamicDestinations =
new TableFunctionDestinations<>(getTableFunction(), getClustering() != null);
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java
index fb13ba8..0f6b4cc 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java
@@ -60,21 +60,28 @@
static class ConstantTableDestinations<T> extends DynamicDestinations<T, TableDestination> {
private final ValueProvider<String> tableSpec;
private final @Nullable String tableDescription;
+ private final boolean clusteringEnabled;
- ConstantTableDestinations(ValueProvider<String> tableSpec, @Nullable String tableDescription) {
+ ConstantTableDestinations(
+ ValueProvider<String> tableSpec,
+ @Nullable String tableDescription,
+ boolean clusteringEnabled) {
this.tableSpec = tableSpec;
this.tableDescription = tableDescription;
+ this.clusteringEnabled = clusteringEnabled;
}
static <T> ConstantTableDestinations<T> fromTableSpec(
- ValueProvider<String> tableSpec, String tableDescription) {
- return new ConstantTableDestinations<>(tableSpec, tableDescription);
+ ValueProvider<String> tableSpec, String tableDescription, boolean clusteringEnabled) {
+ return new ConstantTableDestinations<>(tableSpec, tableDescription, clusteringEnabled);
}
static <T> ConstantTableDestinations<T> fromJsonTableRef(
- ValueProvider<String> jsonTableRef, String tableDescription) {
+ ValueProvider<String> jsonTableRef, String tableDescription, boolean clusteringEnabled) {
return new ConstantTableDestinations<>(
- NestedValueProvider.of(jsonTableRef, new JsonTableRefToTableSpec()), tableDescription);
+ NestedValueProvider.of(jsonTableRef, new JsonTableRefToTableSpec()),
+ tableDescription,
+ clusteringEnabled);
}
@Override
@@ -96,7 +103,11 @@
@Override
public Coder<TableDestination> getDestinationCoder() {
- return TableDestinationCoderV2.of();
+ if (clusteringEnabled) {
+ return TableDestinationCoderV3.of();
+ } else {
+ return TableDestinationCoderV2.of();
+ }
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
index 32ed1fe..2a2697a 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
@@ -332,7 +332,7 @@
@FinishBundle
public void finishBundle(FinishBundleContext c) throws Exception {
DatasetService datasetService =
- bqServices.getDatasetService(c.getPipelineOptions().as(BigQueryOptions.class));
+ getDatasetService(c.getPipelineOptions().as(BigQueryOptions.class));
PendingJobManager jobManager = new PendingJobManager();
for (final PendingJobData pendingJob : pendingJobs) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
index 2b6955f..3ac87bd 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
@@ -82,6 +82,7 @@
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
import org.apache.beam.sdk.transforms.display.HasDisplayData;
@@ -92,7 +93,9 @@
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
@@ -100,6 +103,7 @@
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Duration;
+import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -1004,7 +1008,7 @@
* using {@link DatastoreV1.Write#withProjectId}.
*/
public Write write() {
- return new Write(null, null, true, DEFAULT_HINT_NUM_WORKERS);
+ return new Write(null, null, true, StaticValueProvider.of(DEFAULT_HINT_NUM_WORKERS));
}
/**
@@ -1012,7 +1016,7 @@
* using {@link DeleteEntity#withProjectId}.
*/
public DeleteEntity deleteEntity() {
- return new DeleteEntity(null, null, true, DEFAULT_HINT_NUM_WORKERS);
+ return new DeleteEntity(null, null, true, StaticValueProvider.of(DEFAULT_HINT_NUM_WORKERS));
}
/**
@@ -1020,7 +1024,7 @@
* {@link DeleteKey#withProjectId}.
*/
public DeleteKey deleteKey() {
- return new DeleteKey(null, null, true, DEFAULT_HINT_NUM_WORKERS);
+ return new DeleteKey(null, null, true, StaticValueProvider.of(DEFAULT_HINT_NUM_WORKERS));
}
/**
@@ -1038,7 +1042,7 @@
@Nullable ValueProvider<String> projectId,
@Nullable String localhost,
boolean throttleRampup,
- int hintNumWorkers) {
+ ValueProvider<Integer> hintNumWorkers) {
super(projectId, localhost, new UpsertFn(), throttleRampup, hintNumWorkers);
}
@@ -1073,7 +1077,12 @@
* is ignored if ramp-up throttling is disabled.
*/
public Write withHintNumWorkers(int hintNumWorkers) {
- checkArgument(hintNumWorkers > 0, "hintNumWorkers must be positive");
+ return withHintNumWorkers(StaticValueProvider.of(hintNumWorkers));
+ }
+
+ /** Same as {@link Write#withHintNumWorkers(int)} but with a {@link ValueProvider}. */
+ public Write withHintNumWorkers(ValueProvider<Integer> hintNumWorkers) {
+ checkArgument(hintNumWorkers != null, "hintNumWorkers can not be null");
return new Write(projectId, localhost, throttleRampup, hintNumWorkers);
}
}
@@ -1093,7 +1102,7 @@
@Nullable ValueProvider<String> projectId,
@Nullable String localhost,
boolean throttleRampup,
- int hintNumWorkers) {
+ ValueProvider<Integer> hintNumWorkers) {
super(projectId, localhost, new DeleteEntityFn(), throttleRampup, hintNumWorkers);
}
@@ -1132,6 +1141,12 @@
*/
public DeleteEntity withHintNumWorkers(int hintNumWorkers) {
checkArgument(hintNumWorkers > 0, "hintNumWorkers must be positive");
+ return withHintNumWorkers(StaticValueProvider.of(hintNumWorkers));
+ }
+
+ /** Same as {@link DeleteEntity#withHintNumWorkers(int)} but with a {@link ValueProvider}. */
+ public DeleteEntity withHintNumWorkers(ValueProvider<Integer> hintNumWorkers) {
+ checkArgument(hintNumWorkers != null, "hintNumWorkers can not be null");
return new DeleteEntity(projectId, localhost, throttleRampup, hintNumWorkers);
}
}
@@ -1152,7 +1167,7 @@
@Nullable ValueProvider<String> projectId,
@Nullable String localhost,
boolean throttleRampup,
- int hintNumWorkers) {
+ ValueProvider<Integer> hintNumWorkers) {
super(projectId, localhost, new DeleteKeyFn(), throttleRampup, hintNumWorkers);
}
@@ -1191,6 +1206,12 @@
*/
public DeleteKey withHintNumWorkers(int hintNumWorkers) {
checkArgument(hintNumWorkers > 0, "hintNumWorkers must be positive");
+ return withHintNumWorkers(StaticValueProvider.of(hintNumWorkers));
+ }
+
+ /** Same as {@link DeleteKey#withHintNumWorkers(int)} but with a {@link ValueProvider}. */
+ public DeleteKey withHintNumWorkers(ValueProvider<Integer> hintNumWorkers) {
+ checkArgument(hintNumWorkers != null, "hintNumWorkers can not be null");
return new DeleteKey(projectId, localhost, throttleRampup, hintNumWorkers);
}
}
@@ -1208,11 +1229,11 @@
protected ValueProvider<String> projectId;
protected @Nullable String localhost;
protected boolean throttleRampup;
- protected int hintNumWorkers;
+ protected ValueProvider<Integer> hintNumWorkers;
/** A function that transforms each {@code T} into a mutation. */
private final SimpleFunction<T, Mutation> mutationFn;
- private final RampupThrottlingFn<Mutation> rampupThrottlingFn;
+ private RampupThrottlingFn<Mutation> rampupThrottlingFn;
/**
* Note that {@code projectId} is only {@code @Nullable} as a matter of build order, but if it
@@ -1223,13 +1244,12 @@
@Nullable String localhost,
SimpleFunction<T, Mutation> mutationFn,
boolean throttleRampup,
- int hintNumWorkers) {
+ ValueProvider<Integer> hintNumWorkers) {
this.projectId = projectId;
this.localhost = localhost;
this.throttleRampup = throttleRampup;
this.hintNumWorkers = hintNumWorkers;
this.mutationFn = checkNotNull(mutationFn);
- this.rampupThrottlingFn = new RampupThrottlingFn<>(hintNumWorkers);
}
@Override
@@ -1243,9 +1263,28 @@
PCollection<Mutation> intermediateOutput =
input.apply("Convert to Mutation", MapElements.via(mutationFn));
if (throttleRampup) {
+ PCollectionView<Instant> startTimestampView =
+ input
+ .getPipeline()
+ .apply(
+ "Generate start timestamp",
+ new PTransform<PBegin, PCollectionView<Instant>>() {
+ @Override
+ public PCollectionView<Instant> expand(PBegin input) {
+ return input
+ .apply(Create.of("side input"))
+ .apply(
+ MapElements.into(TypeDescriptor.of(Instant.class))
+ .via((s) -> Instant.now()))
+ .apply(View.asSingleton());
+ }
+ });
+ rampupThrottlingFn = new RampupThrottlingFn<>(hintNumWorkers, startTimestampView);
+
intermediateOutput =
intermediateOutput.apply(
- "Enforce ramp-up through throttling", ParDo.of(rampupThrottlingFn));
+ "Enforce ramp-up through throttling",
+ ParDo.of(rampupThrottlingFn).withSideInputs(startTimestampView));
}
intermediateOutput.apply(
"Write Mutation to Datastore", ParDo.of(new DatastoreWriterFn(projectId, localhost)));
@@ -1267,7 +1306,7 @@
builder
.addIfNotNull(DisplayData.item("projectId", projectId).withLabel("Output Project"))
.include("mutationFn", mutationFn);
- if (throttleRampup) {
+ if (rampupThrottlingFn != null) {
builder.include("rampupThrottlingFn", rampupThrottlingFn);
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFn.java
index e58814c..34922356c 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFn.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFn.java
@@ -21,6 +21,8 @@
import java.io.Serializable;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Sum;
import org.apache.beam.sdk.transforms.display.DisplayData;
@@ -28,6 +30,7 @@
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.util.MovingFunction;
import org.apache.beam.sdk.util.Sleeper;
+import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -46,19 +49,19 @@
private static final Duration RAMP_UP_INTERVAL = Duration.standardMinutes(5);
private static final FluentBackoff fluentBackoff = FluentBackoff.DEFAULT;
- private final int numWorkers;
+ private final ValueProvider<Integer> numWorkers;
+ private final PCollectionView<Instant> firstInstantSideInput;
@VisibleForTesting
Counter throttlingMsecs = Metrics.counter(RampupThrottlingFn.class, "throttling-msecs");
// Initialized on every setup.
private transient MovingFunction successfulOps;
- // Initialized once in constructor.
- private Instant firstInstant;
@VisibleForTesting transient Sleeper sleeper;
- public RampupThrottlingFn(int numWorkers) {
+ public RampupThrottlingFn(
+ ValueProvider<Integer> numWorkers, PCollectionView<Instant> firstInstantSideInput) {
this.numWorkers = numWorkers;
this.sleeper = Sleeper.DEFAULT;
this.successfulOps =
@@ -68,18 +71,22 @@
1 /* numSignificantBuckets */,
1 /* numSignificantSamples */,
Sum.ofLongs());
- this.firstInstant = Instant.now();
+ this.firstInstantSideInput = firstInstantSideInput;
+ }
+
+ public RampupThrottlingFn(int numWorkers, PCollectionView<Instant> timestampSideInput) {
+ this(StaticValueProvider.of(numWorkers), timestampSideInput);
}
// 500 / numWorkers * 1.5^max(0, (x-5)/5), or "+50% every 5 minutes"
- private int calcMaxOpsBudget(Instant first, Instant instant) {
+ private int calcMaxOpsBudget(Instant first, Instant instant, int hintNumWorkers) {
double rampUpIntervalMinutes = (double) RAMP_UP_INTERVAL.getStandardMinutes();
Duration durationSinceFirst = new Duration(first, instant);
double calculatedGrowth =
(durationSinceFirst.getStandardMinutes() - rampUpIntervalMinutes) / rampUpIntervalMinutes;
double growth = Math.max(0, calculatedGrowth);
- double maxOpsBudget = BASE_BUDGET / this.numWorkers * Math.pow(1.5, growth);
+ double maxOpsBudget = BASE_BUDGET / hintNumWorkers * Math.pow(1.5, growth);
return (int) Math.min(Integer.MAX_VALUE, Math.max(1, maxOpsBudget));
}
@@ -98,13 +105,13 @@
/** Emit only as many elements as the exponentially increasing budget allows. */
@ProcessElement
public void processElement(ProcessContext c) throws IOException, InterruptedException {
- Instant nonNullableFirstInstant = firstInstant;
+ Instant firstInstant = c.sideInput(firstInstantSideInput);
T element = c.element();
BackOff backoff = fluentBackoff.backoff();
while (true) {
Instant instant = Instant.now();
- int maxOpsBudget = calcMaxOpsBudget(nonNullableFirstInstant, instant);
+ int maxOpsBudget = calcMaxOpsBudget(firstInstant, instant, this.numWorkers.get());
long currentOpCount = successfulOps.get(instant.getMillis());
long availableOps = maxOpsBudget - currentOpCount;
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
index 2ade129..8046ce0 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
@@ -130,6 +130,12 @@
private static final int PUBSUB_NAME_MIN_LENGTH = 3;
private static final int PUBSUB_NAME_MAX_LENGTH = 255;
+ // See https://cloud.google.com/pubsub/quotas#resource_limits.
+ private static final int PUBSUB_MESSAGE_DATA_MAX_LENGTH = 10 << 20;
+ private static final int PUBSUB_MESSAGE_MAX_ATTRIBUTES = 100;
+ private static final int PUBSUB_MESSAGE_ATTRIBUTE_MAX_KEY_LENGTH = 256;
+ private static final int PUBSUB_MESSAGE_ATTRIBUTE_MAX_VALUE_LENGTH = 1024;
+
private static final String SUBSCRIPTION_RANDOM_TEST_PREFIX = "_random/";
private static final String SUBSCRIPTION_STARTING_SIGNAL = "_starting_signal/";
private static final String TOPIC_DEV_NULL_TEST_NAME = "/topics/dev/null";
@@ -165,6 +171,48 @@
}
}
+ private static void validatePubsubMessage(PubsubMessage message)
+ throws SizeLimitExceededException {
+ if (message.getPayload().length > PUBSUB_MESSAGE_DATA_MAX_LENGTH) {
+ throw new SizeLimitExceededException(
+ "Pubsub message data field of length "
+ + message.getPayload().length
+ + " exceeds maximum of "
+ + PUBSUB_MESSAGE_DATA_MAX_LENGTH
+ + ". See https://cloud.google.com/pubsub/quotas#resource_limits");
+ }
+ @Nullable Map<String, String> attributes = message.getAttributeMap();
+ if (attributes != null) {
+ if (attributes.size() > PUBSUB_MESSAGE_MAX_ATTRIBUTES) {
+ throw new SizeLimitExceededException(
+ "Pubsub message contains "
+ + attributes.size()
+ + " attributes which exceeds the maximum of "
+ + PUBSUB_MESSAGE_MAX_ATTRIBUTES
+ + ". See https://cloud.google.com/pubsub/quotas#resource_limits");
+ }
+ for (Map.Entry<String, String> attribute : attributes.entrySet()) {
+ if (attribute.getKey().length() > PUBSUB_MESSAGE_ATTRIBUTE_MAX_KEY_LENGTH) {
+ throw new SizeLimitExceededException(
+ "Pubsub message attribute key "
+ + attribute.getKey()
+ + " exceeds the maximum of "
+ + PUBSUB_MESSAGE_ATTRIBUTE_MAX_KEY_LENGTH
+ + ". See https://cloud.google.com/pubsub/quotas#resource_limits");
+ }
+ String value = attribute.getValue();
+ if (value.length() > PUBSUB_MESSAGE_ATTRIBUTE_MAX_VALUE_LENGTH) {
+ throw new SizeLimitExceededException(
+ "Pubsub message attribute value starting with "
+ + value.substring(0, Math.min(256, value.length()))
+ + " exceeds the maximum of "
+ + PUBSUB_MESSAGE_ATTRIBUTE_MAX_VALUE_LENGTH
+ + ". See https://cloud.google.com/pubsub/quotas#resource_limits");
+ }
+ }
+ }
+ }
+
/** Populate common {@link DisplayData} between Pubsub source and sink. */
private static void populateCommonDisplayData(
DisplayData.Builder builder,
@@ -1171,7 +1219,18 @@
return PDone.in(input.getPipeline());
case UNBOUNDED:
return input
- .apply(MapElements.into(new TypeDescriptor<PubsubMessage>() {}).via(getFormatFn()))
+ .apply(
+ MapElements.into(new TypeDescriptor<PubsubMessage>() {})
+ .via(
+ elem -> {
+ PubsubMessage message = getFormatFn().apply(elem);
+ try {
+ validatePubsubMessage(message);
+ } catch (SizeLimitExceededException e) {
+ throw new IllegalArgumentException(e);
+ }
+ return message;
+ }))
.apply(
new PubsubUnboundedSink(
getPubsubClientFactory(),
@@ -1233,6 +1292,7 @@
public void processElement(ProcessContext c) throws IOException, SizeLimitExceededException {
byte[] payload;
PubsubMessage message = getFormatFn().apply(c.element());
+ validatePubsubMessage(message);
payload = message.getPayload();
Map<String, String> attributes = message.getAttributeMap();
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/BatchSpannerRead.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/BatchSpannerRead.java
index fc24c8f..5393c7d 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/BatchSpannerRead.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/BatchSpannerRead.java
@@ -21,9 +21,14 @@
import com.google.cloud.spanner.BatchReadOnlyTransaction;
import com.google.cloud.spanner.Partition;
import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.SpannerException;
import com.google.cloud.spanner.Struct;
import com.google.cloud.spanner.TimestampBound;
+import java.util.HashMap;
import java.util.List;
+import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
+import org.apache.beam.runners.core.metrics.ServiceCallMetric;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
@@ -158,18 +163,43 @@
@ProcessElement
public void processElement(ProcessContext c) throws Exception {
+ ServiceCallMetric serviceCallMetric =
+ createServiceCallMetric(
+ this.config.getProjectId().toString(),
+ this.config.getDatabaseId().toString(),
+ this.config.getInstanceId().toString());
Transaction tx = c.sideInput(txView);
BatchReadOnlyTransaction batchTx =
spannerAccessor.getBatchClient().batchReadOnlyTransaction(tx.transactionId());
+ serviceCallMetric.call("ok");
Partition p = c.element();
try (ResultSet resultSet = batchTx.execute(p)) {
while (resultSet.next()) {
Struct s = resultSet.getCurrentRowAsStruct();
c.output(s);
}
+ } catch (SpannerException e) {
+ serviceCallMetric.call(e.getErrorCode().getGrpcStatusCode().toString());
}
}
+
+ private ServiceCallMetric createServiceCallMetric(
+ String projectId, String databaseId, String tableId) {
+ HashMap<String, String> baseLabels = new HashMap<>();
+ baseLabels.put(MonitoringInfoConstants.Labels.PTRANSFORM, "");
+ baseLabels.put(MonitoringInfoConstants.Labels.SERVICE, "Spanner");
+ baseLabels.put(MonitoringInfoConstants.Labels.METHOD, "Read");
+ baseLabels.put(
+ MonitoringInfoConstants.Labels.RESOURCE,
+ GcpResourceIdentifiers.spannerTable(projectId, databaseId, tableId));
+ baseLabels.put(MonitoringInfoConstants.Labels.SPANNER_PROJECT_ID, projectId);
+ baseLabels.put(MonitoringInfoConstants.Labels.SPANNER_DATABASE_ID, databaseId);
+ baseLabels.put(MonitoringInfoConstants.Labels.SPANNER_INSTANCE_ID, tableId);
+ ServiceCallMetric serviceCallMetric =
+ new ServiceCallMetric(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, baseLabels);
+ return serviceCallMetric;
+ }
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadOperation.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadOperation.java
index 1066115..0c9c42d 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadOperation.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/ReadOperation.java
@@ -42,6 +42,8 @@
public abstract @Nullable Statement getQuery();
+ public abstract @Nullable String getQueryName();
+
public abstract @Nullable String getTable();
public abstract @Nullable String getIndex();
@@ -57,6 +59,8 @@
abstract Builder setQuery(Statement statement);
+ abstract Builder setQueryName(String queryName);
+
abstract Builder setTable(String table);
abstract Builder setIndex(String index);
@@ -92,6 +96,10 @@
return withQuery(Statement.of(sql));
}
+ public ReadOperation withQueryName(String queryName) {
+ return toBuilder().setQueryName(queryName).build();
+ }
+
public ReadOperation withKeySet(KeySet keySet) {
return toBuilder().setKeySet(keySet).build();
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
index c34e21c..faff06e 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
@@ -17,10 +17,8 @@
*/
package org.apache.beam.sdk.io.gcp.spanner;
-import com.google.api.gax.core.ExecutorProvider;
-import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.api.gax.retrying.RetrySettings;
-import com.google.api.gax.rpc.HeaderProvider;
+import com.google.api.gax.rpc.FixedHeaderProvider;
import com.google.api.gax.rpc.ServerStreamingCallSettings;
import com.google.api.gax.rpc.UnaryCallSettings;
import com.google.cloud.NoCredentials;
@@ -31,7 +29,6 @@
import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.Spanner;
import com.google.cloud.spanner.SpannerOptions;
-import com.google.cloud.spanner.spi.v1.SpannerInterceptorProvider;
import com.google.spanner.v1.CommitRequest;
import com.google.spanner.v1.CommitResponse;
import com.google.spanner.v1.ExecuteSqlRequest;
@@ -41,21 +38,11 @@
import io.grpc.ClientCall;
import io.grpc.ClientInterceptor;
import io.grpc.MethodDescriptor;
-import java.net.MalformedURLException;
-import java.net.URL;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
-import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.util.ReleaseInfo;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -88,12 +75,6 @@
private final DatabaseAdminClient databaseAdminClient;
private final SpannerConfig spannerConfig;
- private static final int MAX_MESSAGE_SIZE = 100 * 1024 * 1024;
- private static final int MAX_METADATA_SIZE = 32 * 1024; // bytes
- private static final int NUM_CHANNELS = 4;
- public static final org.threeten.bp.Duration GRPC_KEEP_ALIVE_SECONDS =
- org.threeten.bp.Duration.ofSeconds(120);
-
private SpannerAccessor(
Spanner spanner,
DatabaseClient databaseClient,
@@ -161,23 +142,6 @@
.setTotalTimeout(org.threeten.bp.Duration.ofMinutes(120))
.build());
- ManagedInstantiatingExecutorProvider executorProvider =
- new ManagedInstantiatingExecutorProvider(
- new ThreadFactoryBuilder()
- .setDaemon(true)
- .setNameFormat("Cloud-Spanner-TransportChannel-%d")
- .build());
-
- InstantiatingGrpcChannelProvider.Builder instantiatingGrpcChannelProvider =
- InstantiatingGrpcChannelProvider.newBuilder()
- .setMaxInboundMessageSize(MAX_MESSAGE_SIZE)
- .setMaxInboundMetadataSize(MAX_METADATA_SIZE)
- .setPoolSize(NUM_CHANNELS)
- .setExecutorProvider(executorProvider)
- .setKeepAliveTime(GRPC_KEEP_ALIVE_SECONDS)
- .setInterceptorProvider(SpannerInterceptorProvider.createDefault())
- .setAttemptDirectPath(true);
-
ValueProvider<String> projectId = spannerConfig.getProjectId();
if (projectId != null) {
builder.setProjectId(projectId.get());
@@ -189,34 +153,14 @@
ValueProvider<String> host = spannerConfig.getHost();
if (host != null) {
builder.setHost(host.get());
- instantiatingGrpcChannelProvider.setEndpoint(getEndpoint(host.get()));
}
ValueProvider<String> emulatorHost = spannerConfig.getEmulatorHost();
if (emulatorHost != null) {
builder.setEmulatorHost(emulatorHost.get());
builder.setCredentials(NoCredentials.getInstance());
- } else {
- String userAgentString = USER_AGENT_PREFIX + "/" + ReleaseInfo.getReleaseInfo().getVersion();
- /* Workaround to setup user-agent string.
- * InstantiatingGrpcChannelProvider will override the settings provided.
- * The section below and all associated artifacts will be removed once the bug
- * that prevents setting user-agent is fixed.
- * https://github.com/googleapis/java-spanner/pull/871
- *
- * Code to be replaced:
- * builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", userAgentString));
- */
- instantiatingGrpcChannelProvider.setHeaderProvider(
- new HeaderProvider() {
- @Override
- public Map<String, String> getHeaders() {
- final Map<String, String> headers = new HashMap<>();
- headers.put("user-agent", userAgentString);
- return headers;
- }
- });
- builder.setChannelProvider(instantiatingGrpcChannelProvider.build());
}
+ String userAgentString = USER_AGENT_PREFIX + "/" + ReleaseInfo.getReleaseInfo().getVersion();
+ builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", userAgentString));
SpannerOptions options = builder.build();
Spanner spanner = options.getService();
@@ -232,17 +176,6 @@
spanner, databaseClient, databaseAdminClient, batchClient, spannerConfig);
}
- private static String getEndpoint(String host) {
- URL url;
- try {
- url = new URL(host);
- } catch (MalformedURLException e) {
- throw new IllegalArgumentException("Invalid host: " + host, e);
- }
- return String.format(
- "%s:%s", url.getHost(), url.getPort() < 0 ? url.getDefaultPort() : url.getPort());
- }
-
public DatabaseClient getDatabaseClient() {
return databaseClient;
}
@@ -291,32 +224,4 @@
return next.newCall(method, callOptions);
}
}
-
- private static final class ManagedInstantiatingExecutorProvider implements ExecutorProvider {
- // 4 Gapic clients * 4 channels per client.
- private static final int DEFAULT_MIN_THREAD_COUNT = 16;
- private final List<ScheduledExecutorService> executors = new ArrayList<>();
- private final ThreadFactory threadFactory;
-
- private ManagedInstantiatingExecutorProvider(ThreadFactory threadFactory) {
- this.threadFactory = threadFactory;
- }
-
- @Override
- public boolean shouldAutoClose() {
- return false;
- }
-
- @Override
- public ScheduledExecutorService getExecutor() {
- int numCpus = Runtime.getRuntime().availableProcessors();
- int numThreads = Math.max(DEFAULT_MIN_THREAD_COUNT, numCpus);
- ScheduledExecutorService executor =
- new ScheduledThreadPoolExecutor(numThreads, threadFactory);
- synchronized (this) {
- executors.add(executor);
- }
- return executor;
- }
- }
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
index 07ff216..1244b90 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
@@ -43,9 +43,13 @@
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
+import java.util.HashMap;
import java.util.List;
import java.util.OptionalInt;
import java.util.concurrent.TimeUnit;
+import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
+import org.apache.beam.runners.core.metrics.ServiceCallMetric;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.coders.SerializableCoder;
@@ -661,6 +665,10 @@
return withQuery(Statement.of(sql));
}
+ public Read withQueryName(String queryName) {
+ return withReadOperation(getReadOperation().withQueryName(queryName));
+ }
+
public Read withKeySet(KeySet keySet) {
return withReadOperation(getReadOperation().withKeySet(keySet));
}
@@ -1638,10 +1646,18 @@
private void spannerWriteWithRetryIfSchemaChange(Iterable<Mutation> batch)
throws SpannerException {
for (int retry = 1; ; retry++) {
+ ServiceCallMetric serviceCallMetric =
+ createServiceCallMetric(
+ this.spannerConfig.getProjectId().toString(),
+ this.spannerConfig.getDatabaseId().toString(),
+ this.spannerConfig.getInstanceId().toString(),
+ "Write");
try {
spannerAccessor.getDatabaseClient().writeAtLeastOnce(batch);
+ serviceCallMetric.call("ok");
return;
} catch (AbortedException e) {
+ serviceCallMetric.call(e.getErrorCode().getGrpcStatusCode().toString());
if (retry >= ABORTED_RETRY_ATTEMPTS) {
throw e;
}
@@ -1649,10 +1665,30 @@
continue;
}
throw e;
+ } catch (SpannerException e) {
+ serviceCallMetric.call(e.getErrorCode().getGrpcStatusCode().toString());
+ throw e;
}
}
}
+ private ServiceCallMetric createServiceCallMetric(
+ String projectId, String databaseId, String tableId, String method) {
+ HashMap<String, String> baseLabels = new HashMap<>();
+ baseLabels.put(MonitoringInfoConstants.Labels.PTRANSFORM, "");
+ baseLabels.put(MonitoringInfoConstants.Labels.SERVICE, "Spanner");
+ baseLabels.put(MonitoringInfoConstants.Labels.METHOD, method);
+ baseLabels.put(
+ MonitoringInfoConstants.Labels.RESOURCE,
+ GcpResourceIdentifiers.spannerTable(projectId, databaseId, tableId));
+ baseLabels.put(MonitoringInfoConstants.Labels.SPANNER_PROJECT_ID, projectId);
+ baseLabels.put(MonitoringInfoConstants.Labels.SPANNER_DATABASE_ID, databaseId);
+ baseLabels.put(MonitoringInfoConstants.Labels.SPANNER_INSTANCE_ID, tableId);
+ ServiceCallMetric serviceCallMetric =
+ new ServiceCallMetric(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, baseLabels);
+ return serviceCallMetric;
+ }
+
/** Write the Mutations to Spanner, handling DEADLINE_EXCEEDED with backoff/retries. */
private void writeMutations(Iterable<Mutation> mutations) throws SpannerException, IOException {
BackOff backoff = bundleWriteBackoff.backoff();
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java
index 6799c67..09b1791 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java
@@ -1735,7 +1735,7 @@
boolean isSingleton = numTables == 1 && numFilesPerTable == 0;
DynamicDestinations<String, TableDestination> dynamicDestinations =
new DynamicDestinationsHelpers.ConstantTableDestinations<>(
- ValueProvider.StaticValueProvider.of("SINGLETON"), "");
+ ValueProvider.StaticValueProvider.of("SINGLETON"), "", false);
List<ShardedKey<TableDestination>> expectedPartitions = Lists.newArrayList();
if (isSingleton) {
expectedPartitions.add(ShardedKey.of(new TableDestination("SINGLETON", ""), 1));
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java
index 604b5cb..940884b 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java
@@ -22,12 +22,18 @@
import java.util.Map;
import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFnTester;
import org.apache.beam.sdk.transforms.DoFnTester.CloningBehavior;
+import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.Sleeper;
+import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.joda.time.DateTimeUtils;
import org.joda.time.Duration;
+import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -50,10 +56,13 @@
@Before
public void setUp() throws Exception {
MockitoAnnotations.openMocks(this);
-
DateTimeUtils.setCurrentMillisFixed(0);
+
+ TestPipeline pipeline = TestPipeline.create();
+ PCollectionView<Instant> startTimeView =
+ pipeline.apply(Create.of(Instant.now())).apply(View.asSingleton());
RampupThrottlingFn<Void> rampupThrottlingFn =
- new RampupThrottlingFn<Void>(1) {
+ new RampupThrottlingFn<Void>(1, startTimeView) {
@Override
@Setup
public void setup() {
@@ -62,6 +71,7 @@
}
};
rampupThrottlingFnTester = DoFnTester.of(rampupThrottlingFn);
+ rampupThrottlingFnTester.setSideInput(startTimeView, GlobalWindow.INSTANCE, Instant.now());
rampupThrottlingFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
rampupThrottlingFnTester.startBundle();
rampupThrottlingFn.throttlingMsecs = mockCounter;
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/ReadWriteIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/ReadWriteIT.java
index 80c362a..bd15310 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/ReadWriteIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/ReadWriteIT.java
@@ -18,7 +18,6 @@
package org.apache.beam.sdk.io.gcp.pubsublite;
import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
-import static org.junit.Assert.fail;
import com.google.cloud.pubsublite.AdminClient;
import com.google.cloud.pubsublite.AdminClientSettings;
@@ -36,19 +35,18 @@
import com.google.cloud.pubsublite.proto.Subscription.DeliveryConfig.DeliveryRequirement;
import com.google.cloud.pubsublite.proto.Topic;
import com.google.cloud.pubsublite.proto.Topic.PartitionConfig.Capacity;
-import com.google.errorprone.annotations.concurrent.GuardedBy;
import com.google.protobuf.ByteString;
import java.util.ArrayDeque;
-import java.util.ArrayList;
import java.util.Deque;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.io.gcp.pubsub.TestPubsubSignal;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -57,12 +55,11 @@
import org.apache.beam.sdk.transforms.FlatMapElements;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.joda.time.Duration;
import org.junit.After;
-import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -71,12 +68,12 @@
import org.slf4j.LoggerFactory;
@RunWith(JUnit4.class)
-@Ignore("https://issues.apache.org/jira/browse/BEAM-12908")
public class ReadWriteIT {
private static final Logger LOG = LoggerFactory.getLogger(ReadWriteIT.class);
private static final CloudZone ZONE = CloudZone.parse("us-central1-b");
private static final int MESSAGE_COUNT = 90;
+ @Rule public transient TestPubsubSignal signal = TestPubsubSignal.create();
@Rule public transient TestPipeline pipeline = TestPipeline.create();
private static ProjectId getProject(PipelineOptions options) {
@@ -208,28 +205,21 @@
"dedupeMessages", PubsubLiteIO.deduplicate(UuidDeduplicationOptions.newBuilder().build()));
}
- // This static out of band communication is needed to retain serializability.
- @GuardedBy("ReadWriteIT.class")
- private static final List<SequencedMessage> received = new ArrayList<>();
-
- private static synchronized void addMessageReceived(SequencedMessage message) {
- received.add(message);
+ public static SimpleFunction<SequencedMessage, Integer> extractIds() {
+ return new SimpleFunction<SequencedMessage, Integer>() {
+ @Override
+ public Integer apply(SequencedMessage input) {
+ return Integer.parseInt(input.getMessage().getData().toStringUtf8());
+ }
+ };
}
- private static synchronized List<SequencedMessage> getTestQuickstartReceived() {
- return ImmutableList.copyOf(received);
- }
-
- private static PTransform<PCollection<? extends SequencedMessage>, PCollection<Void>>
- collectTestQuickstart() {
- return MapElements.via(
- new SimpleFunction<SequencedMessage, Void>() {
- @Override
- public Void apply(SequencedMessage input) {
- addMessageReceived(input);
- return null;
- }
- });
+ public static SerializableFunction<Set<Integer>, Boolean> testIds() {
+ return ids -> {
+ LOG.info("Ids are: {}", ids);
+ Set<Integer> target = IntStream.range(0, MESSAGE_COUNT).boxed().collect(Collectors.toSet());
+ return target.equals(ids);
+ };
}
@Test
@@ -260,37 +250,17 @@
// Read some messages. They should be deduplicated by the time we see them, so there should be
// exactly numMessages, one for every index in [0,MESSAGE_COUNT).
PCollection<SequencedMessage> messages = readMessages(subscription, pipeline);
- messages.apply("messageReceiver", collectTestQuickstart());
- pipeline.run();
+ PCollection<Integer> ids = messages.apply(MapElements.via(extractIds()));
+ ids.apply("PubsubSignalTest", signal.signalSuccessWhen(BigEndianIntegerCoder.of(), testIds()));
+ pipeline.apply(signal.signalStart());
+ PipelineResult job = pipeline.run();
LOG.info("Running!");
- for (int round = 0; round < 120; ++round) {
- Thread.sleep(1000);
- Map<Integer, Integer> receivedCounts = new HashMap<>();
- for (SequencedMessage message : getTestQuickstartReceived()) {
- int id = Integer.parseInt(message.getMessage().getData().toStringUtf8());
- receivedCounts.put(id, receivedCounts.getOrDefault(id, 0) + 1);
- }
- LOG.info("Performing comparison round {}.\n", round);
- boolean done = true;
- List<Integer> missing = new ArrayList<>();
- for (int id = 0; id < MESSAGE_COUNT; id++) {
- int idCount = receivedCounts.getOrDefault(id, 0);
- if (idCount == 0) {
- missing.add(id);
- done = false;
- }
- if (idCount > 1) {
- fail(String.format("Failed to deduplicate message with id %s.", id));
- }
- }
- LOG.info("Still messing messages: {}.\n", missing);
- if (done) {
- return;
- }
+ signal.waitForSuccess(Duration.standardMinutes(5));
+ // A runner may not support cancel
+ try {
+ job.cancel();
+ } catch (UnsupportedOperationException exc) {
+ // noop
}
- fail(
- String.format(
- "Failed to receive all messages after 2 minutes. Received %s messages.",
- getTestQuickstartReceived().size()));
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java
index 5977c2e..0610596 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.gcp.spanner;
+import static org.junit.Assert.assertEquals;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.when;
@@ -24,12 +25,14 @@
import com.google.cloud.Timestamp;
import com.google.cloud.spanner.BatchReadOnlyTransaction;
import com.google.cloud.spanner.BatchTransactionId;
+import com.google.cloud.spanner.ErrorCode;
import com.google.cloud.spanner.FakeBatchTransactionId;
import com.google.cloud.spanner.FakePartitionFactory;
import com.google.cloud.spanner.KeySet;
import com.google.cloud.spanner.Partition;
import com.google.cloud.spanner.PartitionOptions;
import com.google.cloud.spanner.ResultSets;
+import com.google.cloud.spanner.SpannerExceptionFactory;
import com.google.cloud.spanner.Statement;
import com.google.cloud.spanner.Struct;
import com.google.cloud.spanner.TimestampBound;
@@ -38,12 +41,19 @@
import com.google.protobuf.ByteString;
import java.io.Serializable;
import java.util.Arrays;
+import java.util.HashMap;
import java.util.List;
+import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers;
+import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
+import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
+import org.checkerframework.checker.nullness.qual.Nullable;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -83,6 +93,9 @@
public void setUp() throws Exception {
serviceFactory = new FakeServiceFactory();
mockBatchTx = Mockito.mock(BatchReadOnlyTransaction.class);
+ // Setup the ProcessWideContainer for testing metrics are set.
+ MetricsContainerImpl container = new MetricsContainerImpl(null);
+ MetricsEnvironment.setProcessWideContainer(container);
}
@Test
@@ -90,12 +103,7 @@
Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
- SpannerConfig spannerConfig =
- SpannerConfig.create()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withServiceFactory(serviceFactory);
+ SpannerConfig spannerConfig = getSpannerConfig();
PCollection<Struct> one =
pipeline.apply(
@@ -129,17 +137,20 @@
pipeline.run();
}
+ private SpannerConfig getSpannerConfig() {
+ return SpannerConfig.create()
+ .withProjectId("test")
+ .withInstanceId("123")
+ .withDatabaseId("aaa")
+ .withServiceFactory(serviceFactory);
+ }
+
@Test
public void runRead() throws Exception {
Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
- SpannerConfig spannerConfig =
- SpannerConfig.create()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withServiceFactory(serviceFactory);
+ SpannerConfig spannerConfig = getSpannerConfig();
PCollection<Struct> one =
pipeline.apply(
@@ -179,16 +190,137 @@
}
@Test
+ public void testQueryMetrics() throws Exception {
+ Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
+ TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
+
+ SpannerConfig spannerConfig = getSpannerConfig();
+
+ PCollection<Struct> one =
+ pipeline.apply(
+ "read q",
+ SpannerIO.read()
+ .withSpannerConfig(spannerConfig)
+ .withQuery("SELECT * FROM users")
+ .withQueryName("queryName")
+ .withTimestampBound(timestampBound));
+
+ FakeBatchTransactionId id = new FakeBatchTransactionId("runQueryTest");
+ when(mockBatchTx.getBatchTransactionId()).thenReturn(id);
+
+ when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
+ .thenReturn(mockBatchTx);
+ when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
+ .thenReturn(mockBatchTx);
+
+ Partition fakePartition =
+ FakePartitionFactory.createFakeQueryPartition(ByteString.copyFromUtf8("one"));
+
+ when(mockBatchTx.partitionQuery(
+ any(PartitionOptions.class), eq(Statement.of("SELECT * FROM users"))))
+ .thenReturn(Arrays.asList(fakePartition, fakePartition));
+ when(mockBatchTx.execute(any(Partition.class)))
+ .thenThrow(
+ SpannerExceptionFactory.newSpannerException(
+ ErrorCode.DEADLINE_EXCEEDED, "Simulated Timeout 1"))
+ .thenThrow(
+ SpannerExceptionFactory.newSpannerException(
+ ErrorCode.DEADLINE_EXCEEDED, "Simulated Timeout 2"))
+ .thenReturn(
+ ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)),
+ ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 6)));
+
+ pipeline.run();
+ verifyMetricWasSet("test", "aaa", "123", "deadline_exceeded", null, 2);
+ verifyMetricWasSet("test", "aaa", "123", "ok", null, 2);
+ }
+
+ @Test
+ public void testReadMetrics() throws Exception {
+ Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
+ TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
+
+ SpannerConfig spannerConfig = getSpannerConfig();
+
+ PCollection<Struct> one =
+ pipeline.apply(
+ "read q",
+ SpannerIO.read()
+ .withSpannerConfig(spannerConfig)
+ .withTable("users")
+ .withColumns("id", "name")
+ .withTimestampBound(timestampBound));
+
+ FakeBatchTransactionId id = new FakeBatchTransactionId("runReadTest");
+ when(mockBatchTx.getBatchTransactionId()).thenReturn(id);
+
+ when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(timestampBound))
+ .thenReturn(mockBatchTx);
+ when(serviceFactory.mockBatchClient().batchReadOnlyTransaction(any(BatchTransactionId.class)))
+ .thenReturn(mockBatchTx);
+
+ Partition fakePartition =
+ FakePartitionFactory.createFakeReadPartition(ByteString.copyFromUtf8("one"));
+
+ when(mockBatchTx.partitionRead(
+ any(PartitionOptions.class),
+ eq("users"),
+ eq(KeySet.all()),
+ eq(Arrays.asList("id", "name"))))
+ .thenReturn(Arrays.asList(fakePartition, fakePartition, fakePartition));
+ when(mockBatchTx.execute(any(Partition.class)))
+ .thenThrow(
+ SpannerExceptionFactory.newSpannerException(
+ ErrorCode.DEADLINE_EXCEEDED, "Simulated Timeout 1"))
+ .thenThrow(
+ SpannerExceptionFactory.newSpannerException(
+ ErrorCode.DEADLINE_EXCEEDED, "Simulated Timeout 2"))
+ .thenReturn(
+ ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(0, 2)),
+ ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(2, 4)),
+ ResultSets.forRows(FAKE_TYPE, FAKE_ROWS.subList(4, 6)));
+
+ pipeline.run();
+ verifyMetricWasSet("test", "aaa", "123", "deadline_exceeded", null, 2);
+ verifyMetricWasSet("test", "aaa", "123", "ok", null, 3);
+ }
+
+ private void verifyMetricWasSet(
+ String projectId,
+ String databaseId,
+ String tableId,
+ String status,
+ @Nullable String queryName,
+ long count) {
+ // Verify the metric was reported.
+ HashMap<String, String> labels = new HashMap<>();
+ labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, "");
+ labels.put(MonitoringInfoConstants.Labels.SERVICE, "Spanner");
+ labels.put(MonitoringInfoConstants.Labels.METHOD, "Read");
+ labels.put(
+ MonitoringInfoConstants.Labels.RESOURCE,
+ GcpResourceIdentifiers.spannerTable(projectId, databaseId, tableId));
+ labels.put(MonitoringInfoConstants.Labels.SPANNER_PROJECT_ID, projectId);
+ labels.put(MonitoringInfoConstants.Labels.SPANNER_DATABASE_ID, databaseId);
+ labels.put(MonitoringInfoConstants.Labels.SPANNER_INSTANCE_ID, tableId);
+ if (queryName != null) {
+ labels.put(MonitoringInfoConstants.Labels.SPANNER_QUERY_NAME, queryName);
+ }
+ labels.put(MonitoringInfoConstants.Labels.STATUS, status);
+
+ MonitoringInfoMetricName name =
+ MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels);
+ MetricsContainerImpl container =
+ (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer();
+ assertEquals(count, (long) container.getCounter(name).getCumulative());
+ }
+
+ @Test
public void runReadUsingIndex() throws Exception {
Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
- SpannerConfig spannerConfig =
- SpannerConfig.create()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withServiceFactory(serviceFactory);
+ SpannerConfig spannerConfig = getSpannerConfig();
PCollection<Struct> one =
pipeline.apply(
@@ -237,12 +369,7 @@
Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
- SpannerConfig spannerConfig =
- SpannerConfig.create()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withServiceFactory(serviceFactory);
+ SpannerConfig spannerConfig = getSpannerConfig();
PCollection<Struct> one =
pipeline.apply(
@@ -281,12 +408,7 @@
Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
TimestampBound timestampBound = TimestampBound.ofReadTimestamp(timestamp);
- SpannerConfig spannerConfig =
- SpannerConfig.create()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withServiceFactory(serviceFactory);
+ SpannerConfig spannerConfig = getSpannerConfig();
PCollectionView<Transaction> tx =
pipeline.apply(
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java
index 58ce514..d0239fc 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java
@@ -50,13 +50,19 @@
import com.google.cloud.spanner.Type;
import java.io.Serializable;
import java.util.Arrays;
+import java.util.HashMap;
import java.util.List;
+import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers;
+import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
+import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName;
import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO.FailureMode;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO.GatherSortCreateBatchesFn;
import org.apache.beam.sdk.io.gcp.spanner.SpannerIO.WriteToSpannerFn;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -120,6 +126,10 @@
// Simplest schema: a table with int64 key
preparePkMetadata(tx, Arrays.asList(pkMetadata("tEsT", "key", "ASC")));
prepareColumnMetadata(tx, Arrays.asList(columnMetadata("tEsT", "key", "INT64", CELLS_PER_KEY)));
+
+ // Setup the ProcessWideContainer for testing metrics are set.
+ MetricsContainerImpl container = new MetricsContainerImpl(null);
+ MetricsEnvironment.setProcessWideContainer(container);
}
private SpannerSchema getSchema() {
@@ -408,6 +418,62 @@
}
@Test
+ public void testSpannerWriteMetricIsSet() {
+ Mutation mutation = m(2L);
+ PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));
+
+ // respond with 2 error codes and a success.
+ when(serviceFactory.mockDatabaseClient().writeAtLeastOnce(any()))
+ .thenThrow(
+ SpannerExceptionFactory.newSpannerException(
+ ErrorCode.DEADLINE_EXCEEDED, "Simulated Timeout 1"))
+ .thenThrow(
+ SpannerExceptionFactory.newSpannerException(
+ ErrorCode.DEADLINE_EXCEEDED, "Simulated Timeout 2"))
+ .thenReturn(Timestamp.now());
+
+ mutations.apply(
+ SpannerIO.write()
+ .withProjectId("test-project")
+ .withInstanceId("test-instance")
+ .withDatabaseId("test-database")
+ .withFailureMode(FailureMode.FAIL_FAST)
+ .withServiceFactory(serviceFactory));
+ pipeline.run();
+
+ verifyMetricWasSet(
+ "test-project", "test-database", "test-instance", "Write", "deadline_exceeded", 2);
+ verifyMetricWasSet("test-project", "test-database", "test-instance", "Write", "ok", 1);
+ }
+
+ private void verifyMetricWasSet(
+ String projectId,
+ String databaseId,
+ String tableId,
+ String method,
+ String status,
+ long count) {
+ // Verify the metric was reported.
+ HashMap<String, String> labels = new HashMap<>();
+ labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, "");
+ labels.put(MonitoringInfoConstants.Labels.SERVICE, "Spanner");
+ labels.put(MonitoringInfoConstants.Labels.METHOD, method);
+ labels.put(
+ MonitoringInfoConstants.Labels.RESOURCE,
+ GcpResourceIdentifiers.spannerTable(projectId, databaseId, tableId));
+ labels.put(MonitoringInfoConstants.Labels.SPANNER_PROJECT_ID, projectId);
+ labels.put(MonitoringInfoConstants.Labels.SPANNER_DATABASE_ID, databaseId);
+ labels.put(MonitoringInfoConstants.Labels.SPANNER_INSTANCE_ID, tableId);
+ labels.put(MonitoringInfoConstants.Labels.STATUS, status);
+
+ MonitoringInfoMetricName name =
+ MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels);
+ MetricsContainerImpl container =
+ (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer();
+ assertEquals(count, (long) container.getCounter(name).getCumulative());
+ }
+
+ @Test
public void deadlineExceededRetries() throws InterruptedException {
List<Mutation> mutationList = Arrays.asList(m((long) 1));
diff --git a/sdks/java/io/hadoop-format/src/test/java/org/apache/beam/sdk/io/hadoop/format/HadoopFormatIOWriteTest.java b/sdks/java/io/hadoop-format/src/test/java/org/apache/beam/sdk/io/hadoop/format/HadoopFormatIOWriteTest.java
index 9c41bcd..6508fd7 100644
--- a/sdks/java/io/hadoop-format/src/test/java/org/apache/beam/sdk/io/hadoop/format/HadoopFormatIOWriteTest.java
+++ b/sdks/java/io/hadoop-format/src/test/java/org/apache/beam/sdk/io/hadoop/format/HadoopFormatIOWriteTest.java
@@ -47,7 +47,7 @@
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
/** Unit tests for {@link HadoopFormatIO.Write}. */
@RunWith(MockitoJUnitRunner.class)
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/SchemaUtil.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/SchemaUtil.java
index 0cc7c26..cf0dbfc 100644
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/SchemaUtil.java
+++ b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/SchemaUtil.java
@@ -351,7 +351,12 @@
public Row mapRow(ResultSet rs) throws Exception {
Row.Builder rowBuilder = Row.withSchema(schema);
for (int i = 0; i < schema.getFieldCount(); i++) {
- rowBuilder.addValue(fieldExtractors.get(i).extract(rs, i + 1));
+ Object value = fieldExtractors.get(i).extract(rs, i + 1);
+ if (rs.wasNull()) {
+ rowBuilder.addValue(null);
+ } else {
+ rowBuilder.addValue(value);
+ }
}
return rowBuilder.build();
}
diff --git a/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/SchemaUtilTest.java b/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/SchemaUtilTest.java
index 18acf04..fe8c32d 100644
--- a/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/SchemaUtilTest.java
+++ b/sdks/java/io/jdbc/src/test/java/org/apache/beam/sdk/io/jdbc/SchemaUtilTest.java
@@ -36,6 +36,7 @@
import java.sql.Time;
import java.sql.Timestamp;
import java.sql.Types;
+import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.utils.AvroUtils;
import org.apache.beam.sdk.values.Row;
@@ -157,6 +158,14 @@
@Test
public void testBeamRowMapperPrimitiveTypes() throws Exception {
ResultSet mockResultSet = mock(ResultSet.class);
+ AtomicBoolean isNull = new AtomicBoolean(false);
+ when(mockResultSet.wasNull())
+ .thenAnswer(
+ x -> {
+ boolean val = isNull.get();
+ isNull.set(false);
+ return val;
+ });
when(mockResultSet.getLong(eq(1))).thenReturn(42L);
when(mockResultSet.getBytes(eq(2))).thenReturn("binary".getBytes(Charset.forName("UTF-8")));
when(mockResultSet.getBoolean(eq(3))).thenReturn(true);
@@ -175,6 +184,18 @@
when(mockResultSet.getShort(eq(15))).thenReturn((short) 4);
when(mockResultSet.getBytes(eq(16))).thenReturn("varbinary".getBytes(Charset.forName("UTF-8")));
when(mockResultSet.getString(eq(17))).thenReturn("varchar");
+ when(mockResultSet.getBoolean(eq(18)))
+ .thenAnswer(
+ x -> {
+ isNull.set(true);
+ return false;
+ });
+ when(mockResultSet.getInt(eq(19)))
+ .thenAnswer(
+ x -> {
+ isNull.set(true);
+ return 0;
+ });
Schema wantSchema =
Schema.builder()
@@ -195,6 +216,8 @@
.addField("tinyint_col", Schema.FieldType.INT16)
.addField("varbinary_col", Schema.FieldType.BYTES)
.addField("varchar_col", Schema.FieldType.STRING)
+ .addField("nullable_boolean_col", Schema.FieldType.BOOLEAN.withNullable(true))
+ .addField("another_int_col", Schema.FieldType.INT32.withNullable(true))
.build();
Row wantRow =
Row.withSchema(wantSchema)
@@ -215,7 +238,9 @@
(short) 8,
(short) 4,
"varbinary".getBytes(Charset.forName("UTF-8")),
- "varchar")
+ "varchar",
+ null,
+ null)
.build();
SchemaUtil.BeamRowMapper beamRowMapper = SchemaUtil.BeamRowMapper.of(wantSchema);
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java
index 437fd86..2fb5209 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java
@@ -27,7 +27,7 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
/** * */
@RunWith(MockitoJUnitRunner.class)
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java
index 1653daf..9ce4b70 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java
@@ -27,7 +27,7 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
/** * */
@RunWith(MockitoJUnitRunner.class)
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java
index e17fa86..429024e 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java
@@ -26,7 +26,7 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
/** * */
@RunWith(MockitoJUnitRunner.class)
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java
index 5abe605..227542c 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java
@@ -38,7 +38,7 @@
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
/** */
@RunWith(MockitoJUnitRunner.class)
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java
index 9c8ea29..4a7fed2 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java
@@ -62,7 +62,7 @@
import org.junit.runner.RunWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
+import org.mockito.junit.MockitoJUnitRunner;
import org.mockito.stubbing.Answer;
/** * */
diff --git a/sdks/java/testing/nexmark/build.gradle b/sdks/java/testing/nexmark/build.gradle
index 3345b2d..8d2c4da 100644
--- a/sdks/java/testing/nexmark/build.gradle
+++ b/sdks/java/testing/nexmark/build.gradle
@@ -38,7 +38,8 @@
?: ":runners:direct-java"
def nexmarkRunnerVersionProperty = "nexmark.runner.version"
def nexmarkRunnerVersion = project.findProperty(nexmarkRunnerVersionProperty)
-def shouldProvideSpark = ":runners:spark:2".equals(nexmarkRunnerDependency)
+def shouldProvideSpark2 = ":runners:spark:2".equals(nexmarkRunnerDependency)
+def shouldProvideSpark3 = ":runners:spark:3".equals(nexmarkRunnerDependency)
def isDataflowRunner = ":runners:google-cloud-dataflow-java".equals(nexmarkRunnerDependency)
def isDataflowRunnerV2 = isDataflowRunner && "V2".equals(nexmarkRunnerVersion)
def runnerConfiguration = ":runners:direct-java".equals(nexmarkRunnerDependency) ? "shadow" : null
@@ -93,22 +94,35 @@
// The Spark runner requires the user to provide a Spark dependency. For self-contained
// runs with the Spark runner, we can provide such a dependency. This is deliberately phrased
// to not hardcode any runner other than :runners:direct-java
- if (shouldProvideSpark) {
+ if (shouldProvideSpark2) {
gradleRun library.java.spark_core, {
exclude group:"org.slf4j", module:"jul-to-slf4j"
}
gradleRun library.java.spark_sql
gradleRun library.java.spark_streaming
}
+ if (shouldProvideSpark3) {
+ gradleRun library.java.spark3_core, {
+ exclude group:"org.slf4j", module:"jul-to-slf4j"
+ }
+
+ gradleRun library.java.spark3_sql
+ gradleRun library.java.spark3_streaming
+ }
}
-if (shouldProvideSpark) {
+if (shouldProvideSpark2) {
configurations.gradleRun {
// Using Spark runner causes a StackOverflowError if slf4j-jdk14 is on the classpath
exclude group: "org.slf4j", module: "slf4j-jdk14"
}
}
-
+if (shouldProvideSpark3) {
+ configurations.gradleRun {
+ // Using Spark runner causes a StackOverflowError if slf4j-jdk14 is on the classpath
+ exclude group: "org.slf4j", module: "slf4j-jdk14"
+ }
+}
def getNexmarkArgs = {
def nexmarkArgsStr = project.findProperty(nexmarkArgsProperty) ?: ""
def nexmarkArgsList = new ArrayList<String>()
diff --git a/sdks/python/.pylintrc b/sdks/python/.pylintrc
index 1f92812..bead17d 100644
--- a/sdks/python/.pylintrc
+++ b/sdks/python/.pylintrc
@@ -80,7 +80,9 @@
[MESSAGES CONTROL]
disable =
abstract-method,
+ abstract-class-instantiated,
arguments-differ,
+ arguments-renamed,
attribute-defined-outside-init,
bad-builtin,
bad-super-call,
@@ -88,8 +90,10 @@
broad-except,
comparison-with-callable,
consider-using-enumerate,
+ consider-using-f-string,
consider-using-in,
consider-using-sys-exit,
+ consider-using-with,
cyclic-import,
design,
fixme,
@@ -122,6 +126,7 @@
not-callable,
pointless-statement,
protected-access,
+ raise-missing-from, #TODO(BEAM-12991) Enable and fix warnings
raising-format-tuple,
raising-non-exception,
redefined-builtin,
@@ -143,6 +148,7 @@
unnecessary-pass,
unneeded-not,
unsubscriptable-object,
+ unspecified-encoding, #TODO(BEAM-12992) Enable explicit encoding
unused-argument,
unused-wildcard-import,
useless-object-inheritance,
diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py
index 618ee55..668c56c 100644
--- a/sdks/python/apache_beam/coders/coder_impl.py
+++ b/sdks/python/apache_beam/coders/coder_impl.py
@@ -747,7 +747,7 @@
def decode_from_stream(self, in_, nested):
# type: (create_InputStream, bool) -> IntervalWindow
if not TYPE_CHECKING:
- global IntervalWindow
+ 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
@@ -1390,8 +1390,7 @@
and pane info values during decoding when reconstructing the windowed
value."""
def __init__(self, value_coder, window_coder, payload):
- super(ParamWindowedValueCoderImpl,
- self).__init__(value_coder, TimestampCoderImpl(), window_coder)
+ super().__init__(value_coder, TimestampCoderImpl(), window_coder)
self._timestamp, self._windows, self._pane_info = self._from_proto(
payload, window_coder)
diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py
index 05f7a9d..1299a4a 100644
--- a/sdks/python/apache_beam/coders/coders.py
+++ b/sdks/python/apache_beam/coders/coders.py
@@ -730,7 +730,7 @@
return False
def as_cloud_object(self, coders_context=None, is_pair_like=True):
- value = super(_PickleCoderBase, self).as_cloud_object(coders_context)
+ value = super().as_cloud_object(coders_context)
# We currently use this coder in places where we cannot infer the coder to
# use for the value type in a more granular way. In places where the
# service expects a pair, it checks for the "is_pair_like" key, in which
@@ -767,7 +767,7 @@
class _MemoizingPickleCoder(_PickleCoderBase):
"""Coder using Python's pickle functionality with memoization."""
def __init__(self, cache_size=16):
- super(_MemoizingPickleCoder, self).__init__()
+ super().__init__()
self.cache_size = cache_size
def _create_impl(self):
@@ -867,7 +867,7 @@
return Any
def as_cloud_object(self, coders_context=None, is_pair_like=True):
- value = super(FastCoder, self).as_cloud_object(coders_context)
+ value = super().as_cloud_object(coders_context)
# We currently use this coder in places where we cannot infer the coder to
# use for the value type in a more granular way. In places where the
# service expects a pair, it checks for the "is_pair_like" key, in which
@@ -1088,7 +1088,7 @@
],
}
- return super(TupleCoder, self).as_cloud_object(coders_context)
+ return super().as_cloud_object(coders_context)
def _get_component_coders(self):
# type: () -> Tuple[Coder, ...]
@@ -1250,7 +1250,7 @@
"""Coder for global windows."""
def __init__(self):
from apache_beam.transforms import window
- super(GlobalWindowCoder, self).__init__(window.GlobalWindow())
+ super().__init__(window.GlobalWindow())
def as_cloud_object(self, coders_context=None):
return {
@@ -1357,7 +1357,7 @@
class ParamWindowedValueCoder(WindowedValueCoder):
"""A coder used for parameterized windowed values."""
def __init__(self, payload, components):
- super(ParamWindowedValueCoder, self).__init__(components[0], components[1])
+ super().__init__(components[0], components[1])
self.payload = payload
def _create_impl(self):
diff --git a/sdks/python/apache_beam/coders/coders_test.py b/sdks/python/apache_beam/coders/coders_test.py
index 42fb3a3e..0eeb75d 100644
--- a/sdks/python/apache_beam/coders/coders_test.py
+++ b/sdks/python/apache_beam/coders/coders_test.py
@@ -121,7 +121,7 @@
"""
def __init__(self):
- super(AvroTestCoder, self).__init__(self.SCHEMA)
+ super().__init__(self.SCHEMA)
class AvroTestRecord(AvroRecord):
diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py
index 44ba749..dbe2453 100644
--- a/sdks/python/apache_beam/coders/coders_test_common.py
+++ b/sdks/python/apache_beam/coders/coders_test_common.py
@@ -128,12 +128,12 @@
False,
]
test_values = test_values_deterministic + [
- dict(),
+ {},
{
'a': 'b'
},
{
- 0: dict(), 1: len
+ 0: {}, 1: len
},
set(),
{'a', 'b'},
@@ -223,13 +223,12 @@
tuple(self.test_values_deterministic))
with self.assertRaises(TypeError):
- self.check_coder(deterministic_coder, dict())
+ self.check_coder(deterministic_coder, {})
with self.assertRaises(TypeError):
- self.check_coder(deterministic_coder, [1, dict()])
+ self.check_coder(deterministic_coder, [1, {}])
self.check_coder(
- coders.TupleCoder((deterministic_coder, coder)), (1, dict()),
- ('a', [dict()]))
+ coders.TupleCoder((deterministic_coder, coder)), (1, {}), ('a', [{}]))
self.check_coder(deterministic_coder, test_message.MessageA(field1='value'))
@@ -260,7 +259,7 @@
with self.assertRaises(TypeError):
self.check_coder(deterministic_coder, DefinesGetState(1))
with self.assertRaises(TypeError):
- self.check_coder(deterministic_coder, DefinesGetAndSetState(dict()))
+ self.check_coder(deterministic_coder, DefinesGetAndSetState({}))
def test_dill_coder(self):
cell_value = (lambda x: lambda: x)(0).__closure__[0]
diff --git a/sdks/python/apache_beam/coders/row_coder.py b/sdks/python/apache_beam/coders/row_coder.py
index 4d67f8b..ec3778d 100644
--- a/sdks/python/apache_beam/coders/row_coder.py
+++ b/sdks/python/apache_beam/coders/row_coder.py
@@ -59,6 +59,10 @@
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
@@ -71,7 +75,7 @@
return all(c.is_deterministic() for c in self.components)
def to_type_hint(self):
- return named_tuple_from_schema(self.schema)
+ return self._type_hint
def __hash__(self):
return hash(self.schema.SerializeToString())
diff --git a/sdks/python/apache_beam/coders/row_coder_test.py b/sdks/python/apache_beam/coders/row_coder_test.py
index b36aee7..331f824 100644
--- a/sdks/python/apache_beam/coders/row_coder_test.py
+++ b/sdks/python/apache_beam/coders/row_coder_test.py
@@ -185,6 +185,7 @@
)
# 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))
@@ -255,6 +256,17 @@
self.assertEqual(value, coder.decode(coder.encode(value)))
+ 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())
+ ])
+
+ # Should raise an exception referencing the problem field
+ self.assertRaisesRegex(
+ ValueError, "type_with_no_typeinfo", lambda: RowCoder(schema_proto))
+
if __name__ == "__main__":
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/coders/slow_stream.py b/sdks/python/apache_beam/coders/slow_stream.py
index 23dc0ee..cf71c3e 100644
--- a/sdks/python/apache_beam/coders/slow_stream.py
+++ b/sdks/python/apache_beam/coders/slow_stream.py
@@ -92,7 +92,7 @@
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(ByteCountingOutputStream, self).__init__()
+ super().__init__()
self.count = 0
def write(self, byte_array, nested=False):
diff --git a/sdks/python/apache_beam/coders/standard_coders_test.py b/sdks/python/apache_beam/coders/standard_coders_test.py
index 454939f..acec22a 100644
--- a/sdks/python/apache_beam/coders/standard_coders_test.py
+++ b/sdks/python/apache_beam/coders/standard_coders_test.py
@@ -151,13 +151,13 @@
window_parser: windowed_value.create(
value_parser(x['value']),
x['timestamp'] * 1000,
- tuple([window_parser(w) for w in x['windows']])),
+ 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']]),
+ tuple(window_parser(w) for w in x['windows']),
PaneInfo(
x['pane']['is_first'],
x['pane']['is_last'],
@@ -170,7 +170,7 @@
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']]),
+ windows=tuple(window_parser(w) for w in x['windows']),
fire_timestamp=None,
hold_timestamp=None,
paneinfo=None) if x['clearBit'] else userstate.Timer(
@@ -179,7 +179,7 @@
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']]),
+ windows=tuple(window_parser(w) for w in x['windows']),
paneinfo=PaneInfo(
x['pane']['is_first'],
x['pane']['is_last'],
diff --git a/sdks/python/apache_beam/dataframe/doctests.py b/sdks/python/apache_beam/dataframe/doctests.py
index 45171db..a7cff5a 100644
--- a/sdks/python/apache_beam/dataframe/doctests.py
+++ b/sdks/python/apache_beam/dataframe/doctests.py
@@ -68,7 +68,7 @@
def __call__(self, *args, **kwargs):
result = self._pandas_obj(*args, **kwargs)
- if type(result) in DeferredBase._pandas_type_map.keys():
+ 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)
@@ -322,8 +322,7 @@
self.reset()
want, got = self.fix(want, got)
- return super(_DeferrredDataframeOutputChecker,
- self).check_output(want, got, optionflags)
+ return super().check_output(want, got, optionflags)
def output_difference(self, example, got, optionflags):
want, got = self.fix(example.want, got)
@@ -335,8 +334,7 @@
example.lineno,
example.indent,
example.options)
- return super(_DeferrredDataframeOutputChecker,
- self).output_difference(example, got, optionflags)
+ return super().output_difference(example, got, optionflags)
class BeamDataframeDoctestRunner(doctest.DocTestRunner):
@@ -374,7 +372,7 @@
for test,
examples in (skip or {}).items()
}
- super(BeamDataframeDoctestRunner, self).__init__(
+ super().__init__(
checker=_DeferrredDataframeOutputChecker(self._test_env, use_beam),
**kwargs)
self.success = 0
@@ -412,7 +410,7 @@
# Don't fail doctests that raise this error.
example.exc_msg = '|'.join(allowed_exceptions)
with self._test_env.context():
- result = super(BeamDataframeDoctestRunner, self).run(test, **kwargs)
+ result = super().run(test, **kwargs)
# Can't add attributes to builtin result.
result = AugmentedTestResults(result.failed, result.attempted)
result.summary = self.summary()
@@ -444,14 +442,13 @@
# use the wrong previous value.
del test.globs[var]
- return super(BeamDataframeDoctestRunner,
- self).report_success(out, test, example, got)
+ return super().report_success(out, test, example, got)
def fake_pandas_module(self):
return self._test_env.fake_pandas_module()
def summarize(self):
- super(BeamDataframeDoctestRunner, self).summarize()
+ super().summarize()
self.summary().summarize()
def summary(self):
diff --git a/sdks/python/apache_beam/dataframe/doctests_test.py b/sdks/python/apache_beam/dataframe/doctests_test.py
index 1adff65..df24213 100644
--- a/sdks/python/apache_beam/dataframe/doctests_test.py
+++ b/sdks/python/apache_beam/dataframe/doctests_test.py
@@ -234,6 +234,7 @@
def test_rst_ipython(self):
try:
+ # pylint: disable=unused-import
import IPython
except ImportError:
raise unittest.SkipTest('IPython not available')
diff --git a/sdks/python/apache_beam/dataframe/expressions.py b/sdks/python/apache_beam/dataframe/expressions.py
index c8960c3..97997a4 100644
--- a/sdks/python/apache_beam/dataframe/expressions.py
+++ b/sdks/python/apache_beam/dataframe/expressions.py
@@ -69,7 +69,7 @@
if expr not in self._bindings:
if is_scalar(expr) or not expr.args():
- result = super(PartitioningSession, self).evaluate(expr)
+ result = super().evaluate(expr)
else:
scaler_args = [arg for arg in expr.args() if is_scalar(arg)]
@@ -260,7 +260,7 @@
proxy: A proxy object with the type expected to be bound to this
expression. Used for type checking at pipeline construction time.
"""
- super(PlaceholderExpression, self).__init__('placeholder', proxy)
+ super().__init__('placeholder', proxy)
self._reference = reference
def placeholders(self):
@@ -296,7 +296,7 @@
"""
if proxy is None:
proxy = value
- super(ConstantExpression, self).__init__('constant', proxy)
+ super().__init__('constant', proxy)
self._value = value
def placeholders(self):
@@ -357,7 +357,7 @@
args = tuple(args)
if proxy is None:
proxy = func(*(arg.proxy() for arg in args))
- super(ComputedExpression, self).__init__(name, proxy, _id)
+ super().__init__(name, proxy, _id)
self._func = func
self._args = args
self._requires_partition_by = requires_partition_by
@@ -409,5 +409,5 @@
class NonParallelOperation(Exception):
def __init__(self, msg):
- super(NonParallelOperation, self).__init__(self, msg)
+ super().__init__(self, msg)
self.msg = msg
diff --git a/sdks/python/apache_beam/dataframe/frame_base.py b/sdks/python/apache_beam/dataframe/frame_base.py
index 4bb9ddf..b1b0b85 100644
--- a/sdks/python/apache_beam/dataframe/frame_base.py
+++ b/sdks/python/apache_beam/dataframe/frame_base.py
@@ -64,7 +64,7 @@
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Singleton())
- return tuple([cls.wrap(get(ix)) for ix in range(len(expr.proxy()))])
+ 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:
@@ -641,4 +641,4 @@
if 'url' in reason_data:
msg = f"{msg}\nFor more information see {reason_data['url']}."
- super(WontImplementError, self).__init__(msg)
+ super().__init__(msg)
diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py
index b834d9c..5a0e826 100644
--- a/sdks/python/apache_beam/dataframe/frames.py
+++ b/sdks/python/apache_beam/dataframe/frames.py
@@ -55,6 +55,9 @@
'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):
@@ -294,10 +297,14 @@
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=requires))
- ffill = _fillna_alias('ffill')
- bfill = _fillna_alias('bfill')
- backfill = _fillna_alias('backfill')
- pad = _fillna_alias('pad')
+ 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):
@@ -1932,7 +1939,7 @@
else:
column = self
- result = column.groupby(column).size()
+ result = column.groupby(column, dropna=dropna).size()
# groupby.size() names the index, which we don't need
result.index.name = None
@@ -2392,8 +2399,8 @@
if func in ('quantile',):
return getattr(self, func)(*args, axis=axis, **kwargs)
- # Maps to a property, args are ignored
- if func in ('size',):
+ # 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
@@ -3390,25 +3397,32 @@
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Singleton()))
- @frame_base.with_docs_from(pd.DataFrame)
- def value_counts(self, subset=None, sort=False, normalize=False,
- ascending=False):
- """``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 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)
- result = self.groupby(columns).size()
+ 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 normalize:
- return result/self.dropna().length()
- else:
- return result
+ if dropna:
+ dropped = self.dropna()
+ else:
+ dropped = self
+
+ result = dropped.groupby(columns, dropna=dropna).size()
+
+ if normalize:
+ return result/dropped.length()
+ else:
+ return result
for io_func in dir(io):
@@ -3449,7 +3463,7 @@
: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(DeferredGroupBy, self).__init__(expr)
+ super().__init__(expr)
self._ungrouped = ungrouped
self._ungrouped_with_index = ungrouped_with_index
self._projection = projection
@@ -4239,6 +4253,9 @@
return func
for method in ELEMENTWISE_STRING_METHODS:
+ if not hasattr(pd.core.strings.StringMethods, 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),
@@ -4382,6 +4399,9 @@
]
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(
diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py
index c3972ad..a2703d8 100644
--- a/sdks/python/apache_beam/dataframe/frames_test.py
+++ b/sdks/python/apache_beam/dataframe/frames_test.py
@@ -25,7 +25,8 @@
from apache_beam.dataframe import frame_base
from apache_beam.dataframe import frames
-PD_VERSION = tuple(map(int, pd.__version__.split('.')))
+# 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)],
@@ -235,6 +236,17 @@
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(BEAM-12495): 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'],
@@ -369,10 +381,15 @@
nonparallel=True)
def test_combine_Series(self):
- with expressions.allow_non_parallel_operations():
- 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)
+ 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]})
@@ -587,8 +604,27 @@
self._run_test(lambda df: df.value_counts(), df)
self._run_test(lambda df: df.value_counts(normalize=True), df)
+ if PD_VERSION >= (1, 3):
+ # dropna=False is new in pandas 1.3
+ # TODO(BEAM-12495): 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({
diff --git a/sdks/python/apache_beam/dataframe/io.py b/sdks/python/apache_beam/dataframe/io.py
index 5ef8e2c..9cf66ee 100644
--- a/sdks/python/apache_beam/dataframe/io.py
+++ b/sdks/python/apache_beam/dataframe/io.py
@@ -353,7 +353,7 @@
'for splittable csv reads.')
if kwargs.get('skipfooter', 0):
raise ValueError('Splittablility incompatible with skipping footers.')
- super(_CsvSplitter, self).__init__(
+ super().__init__(
_maybe_encode(kwargs.get('lineterminator', b'\n')),
_DEFAULT_BYTES_CHUNKSIZE)
self._kwargs = kwargs
diff --git a/sdks/python/apache_beam/dataframe/io_test.py b/sdks/python/apache_beam/dataframe/io_test.py
index 060eebf..d525b40 100644
--- a/sdks/python/apache_beam/dataframe/io_test.py
+++ b/sdks/python/apache_beam/dataframe/io_test.py
@@ -230,7 +230,7 @@
min(len(s) for s in splits), len(numbers) * 0.9**20 * 0.1)
@parameterized.expand([
- ('defaults', dict()),
+ ('defaults', {}),
('header', dict(header=1)),
('multi_header', dict(header=[0, 1])),
('multi_header', dict(header=[0, 1, 4])),
diff --git a/sdks/python/apache_beam/dataframe/pandas_doctests_test.py b/sdks/python/apache_beam/dataframe/pandas_doctests_test.py
index edc42f1..755e4e5 100644
--- a/sdks/python/apache_beam/dataframe/pandas_doctests_test.py
+++ b/sdks/python/apache_beam/dataframe/pandas_doctests_test.py
@@ -20,6 +20,7 @@
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
@@ -68,7 +69,8 @@
"df.replace(regex={r'^ba.$': 'new', 'foo': 'xyz'})"
],
'pandas.core.generic.NDFrame.fillna': [
- "df.fillna(method='ffill')",
+ 'df.fillna(method=\'ffill\')',
+ 'df.fillna(method="ffill")',
'df.fillna(value=values, limit=1)',
],
'pandas.core.generic.NDFrame.sort_values': ['*'],
@@ -164,7 +166,8 @@
'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(method="ffill")',
'df.fillna(value=values, limit=1)',
],
'pandas.core.frame.DataFrame.items': ['*'],
@@ -237,6 +240,8 @@
# 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': [
@@ -244,6 +249,8 @@
# frames_test.py::DeferredFrameTest::test_groupby_transform_sum
"df.groupby('Date')['Data'].transform('sum')",
],
+ 'pandas.core.frame.DataFrame.swaplevel': ['*'],
+ 'pandas.core.frame.DataFrame.melt': ['*'],
'pandas.core.frame.DataFrame.reindex_axis': ['*'],
'pandas.core.frame.DataFrame.round': [
'df.round(decimals)',
@@ -267,6 +274,11 @@
'pandas.core.frame.DataFrame.set_index': [
"df.set_index([s, s**2])",
],
+
+ # TODO(BEAM-12495)
+ 'pandas.core.frame.DataFrame.value_counts': [
+ 'df.value_counts(dropna=False)'
+ ],
},
skip={
# s2 created with reindex
@@ -274,6 +286,8 @@
'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.
@@ -302,6 +316,14 @@
'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"<S{df.index.str.len().max()}"', # 1.x
@@ -385,7 +407,8 @@
's.dot(arr)', # non-deferred result
],
'pandas.core.series.Series.fillna': [
- "df.fillna(method='ffill')",
+ 'df.fillna(method=\'ffill\')',
+ 'df.fillna(method="ffill")',
'df.fillna(value=values, limit=1)',
],
'pandas.core.series.Series.items': ['*'],
@@ -434,11 +457,11 @@
's.drop_duplicates()',
"s.drop_duplicates(keep='last')",
],
- 'pandas.core.series.Series.repeat': [
- 's.repeat([1, 2, 3])'
- ],
'pandas.core.series.Series.reindex': ['*'],
'pandas.core.series.Series.autocorr': ['*'],
+ 'pandas.core.series.Series.repeat': ['s.repeat([1, 2, 3])'],
+ 'pandas.core.series.Series.resample': ['*'],
+ 'pandas.core.series.Series': ['ser.iloc[0] = 999'],
},
not_implemented_ok={
'pandas.core.series.Series.transform': [
@@ -451,8 +474,11 @@
'ser.groupby(["a", "b", "a", np.nan]).mean()',
'ser.groupby(["a", "b", "a", np.nan], dropna=False).mean()',
],
+ 'pandas.core.series.Series.swaplevel' :['*']
},
skip={
+ # Relies on setting values with iloc
+ 'pandas.core.series.Series': ['ser', 'r'],
'pandas.core.series.Series.groupby': [
# TODO(BEAM-11393): This example requires aligning two series
# with non-unique indexes. It only works in pandas because
@@ -460,6 +486,7 @@
# alignment.
'ser.groupby(ser > 100).mean()',
],
+ 'pandas.core.series.Series.asfreq': ['*'],
# error formatting
'pandas.core.series.Series.append': [
's1.append(s2, verify_integrity=True)',
@@ -491,12 +518,12 @@
# Inspection after modification.
's'
],
+ 'pandas.core.series.Series.resample': ['df'],
})
self.assertEqual(result.failed, 0)
def test_string_tests(self):
- PD_VERSION = tuple(int(v) for v in pd.__version__.split('.'))
- if PD_VERSION < (1, 2, 0):
+ if PD_VERSION < (1, 2):
module = pd.core.strings
else:
# Definitions were moved to accessor in pandas 1.2.0
@@ -668,11 +695,13 @@
'pandas.core.groupby.generic.SeriesGroupBy.diff': ['*'],
'pandas.core.groupby.generic.DataFrameGroupBy.hist': ['*'],
'pandas.core.groupby.generic.DataFrameGroupBy.fillna': [
- "df.fillna(method='ffill')",
+ 'df.fillna(method=\'ffill\')',
+ 'df.fillna(method="ffill")',
'df.fillna(value=values, limit=1)',
],
'pandas.core.groupby.generic.SeriesGroupBy.fillna': [
- "df.fillna(method='ffill')",
+ 'df.fillna(method=\'ffill\')',
+ 'df.fillna(method="ffill")',
'df.fillna(value=values, limit=1)',
],
},
@@ -682,6 +711,7 @@
'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': [
@@ -698,6 +728,14 @@
# These examples rely on grouping by a list
'pandas.core.groupby.generic.SeriesGroupBy.aggregate': ['*'],
'pandas.core.groupby.generic.DataFrameGroupBy.aggregate': ['*'],
+ 'pandas.core.groupby.generic.SeriesGroupBy.transform': [
+ # Dropping invalid columns during a transform is unsupported.
+ 'grouped.transform(lambda x: (x - x.mean()) / x.std())'
+ ],
+ 'pandas.core.groupby.generic.DataFrameGroupBy.transform': [
+ # Dropping invalid columns during a transform is unsupported.
+ 'grouped.transform(lambda x: (x - x.mean()) / x.std())'
+ ],
})
self.assertEqual(result.failed, 0)
diff --git a/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py b/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py
index 443843e..39df3f2 100644
--- a/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py
+++ b/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py
@@ -38,7 +38,7 @@
def _maybe_wrap_constant_expr(res):
- if type(res) in frame_base.DeferredBase._pandas_type_map.keys():
+ if type(res) in frame_base.DeferredBase._pandas_type_map:
return frame_base.DeferredBase.wrap(
expressions.ConstantExpression(res, res[0:0]))
else:
diff --git a/sdks/python/apache_beam/dataframe/partitionings.py b/sdks/python/apache_beam/dataframe/partitionings.py
index 9891e71..bb8c994 100644
--- a/sdks/python/apache_beam/dataframe/partitionings.py
+++ b/sdks/python/apache_beam/dataframe/partitionings.py
@@ -195,7 +195,6 @@
random.shuffle(seq)
return seq
- # pylint: disable=range-builtin-not-iterating
part = pd.Series(shuffled(range(len(df))), index=df.index) % num_partitions
for k in range(num_partitions):
yield k, df[part == k]
diff --git a/sdks/python/apache_beam/dataframe/partitionings_test.py b/sdks/python/apache_beam/dataframe/partitionings_test.py
index b60aa67..c0fa8a9 100644
--- a/sdks/python/apache_beam/dataframe/partitionings_test.py
+++ b/sdks/python/apache_beam/dataframe/partitionings_test.py
@@ -24,7 +24,7 @@
class PartitioningsTest(unittest.TestCase):
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
multi_index_df = pd.DataFrame({
'shape': ['dodecahedron', 'icosahedron'] * 12,
diff --git a/sdks/python/apache_beam/dataframe/transforms.py b/sdks/python/apache_beam/dataframe/transforms.py
index b698448..9b0c4fb 100644
--- a/sdks/python/apache_beam/dataframe/transforms.py
+++ b/sdks/python/apache_beam/dataframe/transforms.py
@@ -109,7 +109,7 @@
input_dict = _flatten(input_pcolls) # type: Dict[Any, PCollection]
proxies = _flatten(self._proxy) if self._proxy is not None else {
tag: None
- for tag in input_dict.keys()
+ for tag in input_dict
}
input_frames = {
k: convert.to_dataframe(pc, proxies[k])
diff --git a/sdks/python/apache_beam/examples/avro_bitcoin.py b/sdks/python/apache_beam/examples/avro_bitcoin.py
index df7fad6..9b851a8 100644
--- a/sdks/python/apache_beam/examples/avro_bitcoin.py
+++ b/sdks/python/apache_beam/examples/avro_bitcoin.py
@@ -43,7 +43,7 @@
"""Count inputs and outputs per transaction"""
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(BitcoinTxnCountDoFn, self).__init__()
+ # super().__init__()
beam.DoFn.__init__(self)
self.txn_counter = Metrics.counter(self.__class__, 'txns')
self.inputs_dist = Metrics.distribution(self.__class__, 'inputs_per_txn')
diff --git a/sdks/python/apache_beam/examples/complete/autocomplete.py b/sdks/python/apache_beam/examples/complete/autocomplete.py
index 3de8ff7..4e4c514 100644
--- a/sdks/python/apache_beam/examples/complete/autocomplete.py
+++ b/sdks/python/apache_beam/examples/complete/autocomplete.py
@@ -59,7 +59,7 @@
class TopPerPrefix(beam.PTransform):
def __init__(self, count):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(TopPerPrefix, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self._count = count
diff --git a/sdks/python/apache_beam/examples/complete/game/game_stats.py b/sdks/python/apache_beam/examples/complete/game/game_stats.py
index e6b6e35..d6f5aab 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats.py
@@ -106,7 +106,7 @@
"""
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ParseGameEventFn, self).__init__()
+ # super().__init__()
beam.DoFn.__init__(self)
self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors')
@@ -132,7 +132,7 @@
"""
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ExtractAndSumScore, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.field = field
@@ -172,7 +172,7 @@
project: Name of the Cloud project containing BigQuery table.
"""
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(WriteToBigQuery, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.table_name = table_name
self.dataset = dataset
diff --git a/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py b/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
index 621a88c..48a105a 100644
--- a/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
+++ b/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
@@ -106,7 +106,7 @@
"""
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ParseGameEventFn, self).__init__()
+ # super().__init__()
beam.DoFn.__init__(self)
self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors')
@@ -132,7 +132,7 @@
"""
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ExtractAndSumScore, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.field = field
@@ -172,7 +172,7 @@
project: Name of the Cloud project containing BigQuery table.
"""
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(WriteToBigQuery, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.table_name = table_name
self.dataset = dataset
@@ -197,7 +197,7 @@
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(HourlyTeamScore, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.start_timestamp = str2timestamp(start_min)
self.stop_timestamp = str2timestamp(stop_min)
diff --git a/sdks/python/apache_beam/examples/complete/game/leader_board.py b/sdks/python/apache_beam/examples/complete/game/leader_board.py
index 5e04d60..308e1e1 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board.py
@@ -115,7 +115,7 @@
"""
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ParseGameEventFn, self).__init__()
+ # super().__init__()
beam.DoFn.__init__(self)
self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors')
@@ -141,7 +141,7 @@
"""
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ExtractAndSumScore, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.field = field
@@ -181,7 +181,7 @@
project: Name of the Cloud project containing BigQuery table.
"""
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(WriteToBigQuery, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.table_name = table_name
self.dataset = dataset
@@ -211,7 +211,7 @@
"""
def __init__(self, team_window_duration, allowed_lateness):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(CalculateTeamScores, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.team_window_duration = team_window_duration * 60
self.allowed_lateness_seconds = allowed_lateness * 60
@@ -243,7 +243,7 @@
"""
def __init__(self, allowed_lateness):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(CalculateUserScores, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.allowed_lateness_seconds = allowed_lateness * 60
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score.py b/sdks/python/apache_beam/examples/complete/game/user_score.py
index a87f221..6a97d7e 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score.py
@@ -79,7 +79,7 @@
"""
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ParseGameEventFn, self).__init__()
+ # super().__init__()
beam.DoFn.__init__(self)
self.num_parse_errors = Metrics.counter(self.__class__, 'num_parse_errors')
@@ -106,7 +106,7 @@
"""
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(ExtractAndSumScore, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.field = field
diff --git a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
index 5727758..7064a5a 100644
--- a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
+++ b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
@@ -117,7 +117,7 @@
"""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(ComputeTopSessions, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.sampling_threshold = sampling_threshold
diff --git a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
index f12c945..e3ea144 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
@@ -68,7 +68,7 @@
"""
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(WriteToBigTable, self).__init__()
+ # super().__init__()
beam.PTransform.__init__(self)
self.number = number
self.rand = random.choice(string.ascii_letters + string.digits)
@@ -179,8 +179,8 @@
with beam.Pipeline(options=pipeline_options) as pipeline:
config_data = {
'project_id': self.project,
- 'instance_id': self.instance,
- 'table_id': self.table
+ 'instance_id': self.instance_id,
+ 'table_id': self.table_id
}
_ = (
pipeline
diff --git a/sdks/python/apache_beam/examples/dataframe/README.md b/sdks/python/apache_beam/examples/dataframe/README.md
index 25fca9e..cba84c0 100644
--- a/sdks/python/apache_beam/examples/dataframe/README.md
+++ b/sdks/python/apache_beam/examples/dataframe/README.md
@@ -26,12 +26,10 @@
You must have `apache-beam>=2.30.0` installed in order to run these pipelines,
because the `apache_beam.examples.dataframe` module was added in that release.
-Additionally using the DataFrame API requires `pandas>=1.0.0` to be installed
-in your local Python session. The _same_ version should be installed on workers
-when executing DataFrame API pipelines on distributed runners. Reference
-[`base_image_requirements.txt`](../../../container/base_image_requirements.txt)
-for the Beam release you are using to see what version of pandas will be used
-by default on distributed workers.
+Using the DataFrame API also requires a compatible pandas version to be
+installed, see the
+[documentation](https://beam.apache.org/documentation/dsls/dataframes/overview/#pre-requisites)
+for details.
## Wordcount Pipeline
diff --git a/sdks/python/apache_beam/examples/fastavro_it_test.py b/sdks/python/apache_beam/examples/fastavro_it_test.py
index c9bb988..f25db8e 100644
--- a/sdks/python/apache_beam/examples/fastavro_it_test.py
+++ b/sdks/python/apache_beam/examples/fastavro_it_test.py
@@ -109,11 +109,11 @@
batch_size = self.test_pipeline.get_option('batch-size')
batch_size = int(batch_size) if batch_size else 10000
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
batches = range(int(num_records / batch_size))
def batch_indices(start):
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
return range(start * batch_size, (start + 1) * batch_size)
# A `PCollection` with `num_records` avro records
diff --git a/sdks/python/apache_beam/examples/flink/flink_streaming_impulse.py b/sdks/python/apache_beam/examples/flink/flink_streaming_impulse.py
index badf401..fdd209a 100644
--- a/sdks/python/apache_beam/examples/flink/flink_streaming_impulse.py
+++ b/sdks/python/apache_beam/examples/flink/flink_streaming_impulse.py
@@ -27,9 +27,9 @@
import sys
import apache_beam as beam
-import apache_beam.transforms.window as window
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
diff --git a/sdks/python/apache_beam/examples/snippets/snippets.py b/sdks/python/apache_beam/examples/snippets/snippets.py
index b665e85..936d06e 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets.py
@@ -801,7 +801,7 @@
# [START model_custom_source_new_ptransform]
class ReadFromCountingSource(PTransform):
def __init__(self, count):
- super(ReadFromCountingSource, self).__init__()
+ super().__init__()
self._count = count
def expand(self, pcoll):
@@ -923,7 +923,7 @@
class WriteToKVSink(PTransform):
def __init__(self, simplekv, url, final_table_name):
self._simplekv = simplekv
- super(WriteToKVSink, self).__init__()
+ super().__init__()
self._url = url
self._final_table_name = final_table_name
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py b/sdks/python/apache_beam/examples/snippets/snippets_test.py
index 8f215a3..940e9fe 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_test.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_test.py
@@ -35,7 +35,6 @@
import parameterized
import apache_beam as beam
-import apache_beam.transforms.combiners as combiners
from apache_beam import WindowInto
from apache_beam import coders
from apache_beam import pvalue
@@ -51,6 +50,7 @@
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
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount.py b/sdks/python/apache_beam/examples/streaming_wordcount.py
index d276cfc..9ae763d 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount.py
@@ -24,11 +24,11 @@
import logging
import apache_beam as beam
-import apache_beam.transforms.window as window
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):
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py b/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
index 2df87f4..f64e6fe 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
@@ -40,13 +40,13 @@
import time
import apache_beam as beam
-import apache_beam.transforms.window as window
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
diff --git a/sdks/python/apache_beam/examples/windowed_wordcount.py b/sdks/python/apache_beam/examples/windowed_wordcount.py
index 861a147..7889f61 100644
--- a/sdks/python/apache_beam/examples/windowed_wordcount.py
+++ b/sdks/python/apache_beam/examples/windowed_wordcount.py
@@ -27,7 +27,7 @@
import logging
import apache_beam as beam
-import apache_beam.transforms.window as window
+from apache_beam.transforms import window
TABLE_SCHEMA = (
'word:STRING, count:INTEGER, '
diff --git a/sdks/python/apache_beam/examples/wordcount_debugging.py b/sdks/python/apache_beam/examples/wordcount_debugging.py
index d4195df..404c123 100644
--- a/sdks/python/apache_beam/examples/wordcount_debugging.py
+++ b/sdks/python/apache_beam/examples/wordcount_debugging.py
@@ -60,7 +60,7 @@
"""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(FilterTextFn, self).__init__()
+ # super().__init__()
beam.DoFn.__init__(self)
self.pattern = pattern
# A custom metric can track values in your pipeline as it runs. Those
diff --git a/sdks/python/apache_beam/examples/wordcount_it_test.py b/sdks/python/apache_beam/examples/wordcount_it_test.py
index 8ee49c7..be8bbbf 100644
--- a/sdks/python/apache_beam/examples/wordcount_it_test.py
+++ b/sdks/python/apache_beam/examples/wordcount_it_test.py
@@ -97,8 +97,7 @@
# Register clean up before pipeline execution
self.addCleanup(delete_files, [test_output + '*'])
- publish_to_bq = bool(
- test_pipeline.get_option('publish_to_big_query') or False)
+ publish_to_bq = bool(test_pipeline.get_option('publish_to_big_query'))
# Start measure time for performance test
start_time = time.time()
diff --git a/sdks/python/apache_beam/examples/wordcount_with_metrics.py b/sdks/python/apache_beam/examples/wordcount_with_metrics.py
index bf61476..8e1dd05 100644
--- a/sdks/python/apache_beam/examples/wordcount_with_metrics.py
+++ b/sdks/python/apache_beam/examples/wordcount_with_metrics.py
@@ -36,7 +36,7 @@
"""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(WordExtractingDoFn, self).__init__()
+ # super().__init__()
beam.DoFn.__init__(self)
self.words_counter = Metrics.counter(self.__class__, 'words')
self.word_lengths_counter = Metrics.counter(self.__class__, 'word_lengths')
diff --git a/sdks/python/apache_beam/internal/gcp/auth.py b/sdks/python/apache_beam/internal/gcp/auth.py
index b3dbe32..ec8a4cd 100644
--- a/sdks/python/apache_beam/internal/gcp/auth.py
+++ b/sdks/python/apache_beam/internal/gcp/auth.py
@@ -52,8 +52,7 @@
@retry.with_exponential_backoff(
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
def _do_refresh_request(self, http_request):
- return super(_GceAssertionCredentials,
- self)._do_refresh_request(http_request)
+ return super()._do_refresh_request(http_request)
def set_running_in_gce(worker_executing_project):
diff --git a/sdks/python/apache_beam/internal/metrics/metric.py b/sdks/python/apache_beam/internal/metrics/metric.py
index 0510674..b76e895 100644
--- a/sdks/python/apache_beam/internal/metrics/metric.py
+++ b/sdks/python/apache_beam/internal/metrics/metric.py
@@ -105,7 +105,7 @@
"""Metrics Histogram that Delegates functionality to MetricsEnvironment."""
def __init__(self, metric_name, bucket_type, logger):
# type: (MetricName, BucketType, Optional[MetricLogger]) -> None
- super(Metrics.DelegatingHistogram, self).__init__(metric_name)
+ super().__init__(metric_name)
self.metric_name = metric_name
self.cell_type = HistogramCellFactory(bucket_type)
self.logger = logger
@@ -125,7 +125,7 @@
"""
def __init__(self):
# type: () -> None
- self._metric = dict() # type: Dict[MetricName, MetricCell]
+ self._metric = {} # type: Dict[MetricName, MetricCell]
self._lock = threading.Lock()
self._last_logging_millis = int(time.time() * 1000)
self.minimum_logging_frequency_msec = 180000
@@ -158,7 +158,7 @@
logging_metric_info.append('%s: %s' % (name, cell.get_cumulative()))
_LOGGER.info('\n'.join(logging_metric_info))
if reset_after_logging:
- self._metric = dict()
+ self._metric = {}
self._last_logging_millis = current_millis
finally:
self._lock.release()
@@ -223,3 +223,43 @@
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):
+ # type: (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 not None and
+ grpc_status_code in grpc_to_canonical_gcp_status):
+ return grpc_to_canonical_gcp_status[grpc_status_code]
+ return str(grpc_status_code)
diff --git a/sdks/python/apache_beam/io/avroio.py b/sdks/python/apache_beam/io/avroio.py
index 45a619c..e861d53 100644
--- a/sdks/python/apache_beam/io/avroio.py
+++ b/sdks/python/apache_beam/io/avroio.py
@@ -143,7 +143,7 @@
use_fastavro (bool); when set, use the `fastavro` library for IO, which
is significantly faster, and is now the default.
"""
- super(ReadFromAvro, self).__init__()
+ super().__init__()
self._source = _create_avro_source(
file_pattern,
min_bundle_size,
@@ -578,7 +578,7 @@
num_shards,
shard_name_template,
mime_type):
- super(_BaseAvroSink, self).__init__(
+ super().__init__(
file_path_prefix,
file_name_suffix=file_name_suffix,
num_shards=num_shards,
@@ -592,7 +592,7 @@
self._codec = codec
def display_data(self):
- res = super(_BaseAvroSink, self).display_data()
+ res = super().display_data()
res['codec'] = str(self._codec)
res['schema'] = str(self._schema)
return res
@@ -601,7 +601,7 @@
class _AvroSink(_BaseAvroSink):
"""A sink for avro files using Avro. """
def open(self, temp_path):
- file_handle = super(_AvroSink, self).open(temp_path)
+ file_handle = super().open(temp_path)
return avro.datafile.DataFileWriter(
file_handle, avro.io.DatumWriter(), self._schema, self._codec)
@@ -612,7 +612,7 @@
class _FastAvroSink(_BaseAvroSink):
"""A sink for avro files using FastAvro. """
def open(self, temp_path):
- file_handle = super(_FastAvroSink, self).open(temp_path)
+ file_handle = super().open(temp_path)
return Writer(file_handle, self._schema, self._codec)
def write_record(self, writer, value):
diff --git a/sdks/python/apache_beam/io/avroio_test.py b/sdks/python/apache_beam/io/avroio_test.py
index 2cb5c5c..dcd1cf7 100644
--- a/sdks/python/apache_beam/io/avroio_test.py
+++ b/sdks/python/apache_beam/io/avroio_test.py
@@ -75,7 +75,7 @@
_temp_files = [] # type: List[str]
def __init__(self, methodName='runTest'):
- super(AvroBase, self).__init__(methodName)
+ super().__init__(methodName)
self.RECORDS = RECORDS
self.SCHEMA_STRING = '''
{"namespace": "example.avro",
@@ -447,7 +447,7 @@
'See: BEAM-6522.')
class TestAvro(AvroBase, unittest.TestCase):
def __init__(self, methodName='runTest'):
- super(TestAvro, self).__init__(methodName)
+ super().__init__(methodName)
self.use_fastavro = False
self.SCHEMA = Parse(self.SCHEMA_STRING)
@@ -477,7 +477,7 @@
class TestFastAvro(AvroBase, unittest.TestCase):
def __init__(self, methodName='runTest'):
- super(TestFastAvro, self).__init__(methodName)
+ super().__init__(methodName)
self.use_fastavro = True
self.SCHEMA = parse_schema(json.loads(self.SCHEMA_STRING))
diff --git a/sdks/python/apache_beam/io/aws/s3filesystem.py b/sdks/python/apache_beam/io/aws/s3filesystem.py
index fa26bad..8a5e94e 100644
--- a/sdks/python/apache_beam/io/aws/s3filesystem.py
+++ b/sdks/python/apache_beam/io/aws/s3filesystem.py
@@ -42,7 +42,7 @@
Connection configuration is done by passing pipeline options.
See :class:`~apache_beam.options.pipeline_options.S3Options`.
"""
- super(S3FileSystem, self).__init__(pipeline_options)
+ super().__init__(pipeline_options)
self._options = pipeline_options
@classmethod
diff --git a/sdks/python/apache_beam/io/concat_source.py b/sdks/python/apache_beam/io/concat_source.py
index 3872ccb..35ae6fe 100644
--- a/sdks/python/apache_beam/io/concat_source.py
+++ b/sdks/python/apache_beam/io/concat_source.py
@@ -91,7 +91,7 @@
# to produce the same coder.
return self._source_bundles[0].source.default_output_coder()
else:
- return super(ConcatSource, self).default_output_coder()
+ return super().default_output_coder()
class ConcatRangeTracker(iobase.RangeTracker):
@@ -106,7 +106,7 @@
end: end position, a tuple of (source_index, source_position)
source_bundles: the list of source bundles in the ConcatSource
"""
- super(ConcatRangeTracker, self).__init__()
+ super().__init__()
self._start = start
self._end = end
self._source_bundles = source_bundles
diff --git a/sdks/python/apache_beam/io/external/generate_sequence.py b/sdks/python/apache_beam/io/external/generate_sequence.py
index e7f56e7..1f94a8d 100644
--- a/sdks/python/apache_beam/io/external/generate_sequence.py
+++ b/sdks/python/apache_beam/io/external/generate_sequence.py
@@ -54,7 +54,7 @@
elements_per_period=None,
max_read_time=None,
expansion_service=None):
- super(GenerateSequence, self).__init__(
+ super().__init__(
self.URN,
ImplicitSchemaPayloadBuilder({
'start': start,
diff --git a/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py b/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
index 6da4567..7f75e2b 100644
--- a/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
@@ -73,7 +73,7 @@
def build_write_pipeline(self, pipeline):
_ = (
pipeline
- | 'Generate' >> beam.Create(range(NUM_RECORDS)) # pylint: disable=range-builtin-not-iterating
+ | 'Generate' >> beam.Create(range(NUM_RECORDS)) # pylint: disable=bad-option-value
| 'MakeKV' >> beam.Map(lambda x:
(b'', str(x).encode())).with_output_types(
typing.Tuple[bytes, bytes])
diff --git a/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py b/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py
index 01ff279..2817ea9 100644
--- a/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py
@@ -104,7 +104,7 @@
_ = (
p
| 'Impulse' >> beam.Impulse()
- | 'Generate' >> beam.FlatMap(lambda x: range(NUM_RECORDS)) # pylint: disable=range-builtin-not-iterating
+ | '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(
diff --git a/sdks/python/apache_beam/io/external/xlang_parquetio_test.py b/sdks/python/apache_beam/io/external/xlang_parquetio_test.py
index e6c06be..10470c1 100644
--- a/sdks/python/apache_beam/io/external/xlang_parquetio_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_parquetio_test.py
@@ -78,7 +78,7 @@
"""
def __init__(self):
- super(AvroTestCoder, self).__init__(self.SCHEMA)
+ super().__init__(self.SCHEMA)
coders.registry.register_coder(AvroRecord, AvroTestCoder)
diff --git a/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py b/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py
index 052ea31..f78175a 100644
--- a/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py
@@ -109,7 +109,7 @@
_ = (
p
| 'Impulse' >> beam.Impulse()
- | 'Generate' >> beam.FlatMap(lambda x: range(NUM_RECORDS)) # pylint: disable=range-builtin-not-iterating
+ | '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()))
diff --git a/sdks/python/apache_beam/io/filebasedsink_test.py b/sdks/python/apache_beam/io/filebasedsink_test.py
index c75958c..121bc47 100644
--- a/sdks/python/apache_beam/io/filebasedsink_test.py
+++ b/sdks/python/apache_beam/io/filebasedsink_test.py
@@ -81,7 +81,7 @@
class MyFileBasedSink(filebasedsink.FileBasedSink):
def open(self, temp_path):
# TODO: Fix main session pickling.
- # file_handle = super(MyFileBasedSink, self).open(temp_path)
+ # file_handle = super().open(temp_path)
file_handle = filebasedsink.FileBasedSink.open(self, temp_path)
file_handle.write(b'[start]')
return file_handle
@@ -94,7 +94,7 @@
def close(self, file_handle):
file_handle.write(b'[end]')
# TODO: Fix main session pickling.
- # file_handle = super(MyFileBasedSink, self).close(file_handle)
+ # file_handle = super().close(file_handle)
file_handle = filebasedsink.FileBasedSink.close(self, file_handle)
diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py
index e890bbe..f21fb8d 100644
--- a/sdks/python/apache_beam/io/fileio_test.py
+++ b/sdks/python/apache_beam/io/fileio_test.py
@@ -102,7 +102,7 @@
'%s%s' % (self._new_tempdir(), os.sep)
]
- files = list()
+ files = []
files.append(self._create_temp_file(dir=directories[0]))
files.append(self._create_temp_file(dir=directories[0]))
@@ -122,7 +122,7 @@
'%s%s' % (self._new_tempdir(), os.sep)
]
- files = list()
+ files = []
files.append(self._create_temp_file(dir=directories[0]))
files.append(self._create_temp_file(dir=directories[0]))
diff --git a/sdks/python/apache_beam/io/filesystem.py b/sdks/python/apache_beam/io/filesystem.py
index 687ab04..92d5592 100644
--- a/sdks/python/apache_beam/io/filesystem.py
+++ b/sdks/python/apache_beam/io/filesystem.py
@@ -468,7 +468,7 @@
the current state of the system.
"""
message = "%s with exceptions %s" % (msg, exception_details)
- super(BeamIOError, self).__init__(message)
+ super().__init__(message)
self.exception_details = exception_details
diff --git a/sdks/python/apache_beam/io/filesystem_test.py b/sdks/python/apache_beam/io/filesystem_test.py
index 83453a9..a246335 100644
--- a/sdks/python/apache_beam/io/filesystem_test.py
+++ b/sdks/python/apache_beam/io/filesystem_test.py
@@ -42,7 +42,7 @@
class TestingFileSystem(FileSystem):
def __init__(self, pipeline_options, has_dirs=False):
- super(TestingFileSystem, self).__init__(pipeline_options)
+ super().__init__(pipeline_options)
self._has_dirs = has_dirs
self._files = {}
@@ -499,7 +499,7 @@
return b''.join(byte_list)
def create_test_file(compression_type, lines):
- filenames = list()
+ filenames = []
file_name = self._create_temp_file()
if compression_type == CompressionTypes.BZIP2:
compress_factory = bz2.BZ2File
diff --git a/sdks/python/apache_beam/io/filesystemio.py b/sdks/python/apache_beam/io/filesystemio.py
index 70f39da..571d1f2 100644
--- a/sdks/python/apache_beam/io/filesystemio.py
+++ b/sdks/python/apache_beam/io/filesystemio.py
@@ -214,7 +214,7 @@
if not self.closed:
self._uploader.finish()
- super(UploaderStream, self).close()
+ super().close()
def writable(self):
return True
diff --git a/sdks/python/apache_beam/io/gcp/__init__.py b/sdks/python/apache_beam/io/gcp/__init__.py
index 87dc1c3..f88a011 100644
--- a/sdks/python/apache_beam/io/gcp/__init__.py
+++ b/sdks/python/apache_beam/io/gcp/__init__.py
@@ -21,8 +21,8 @@
try:
# pylint: disable=wrong-import-order, wrong-import-position
# pylint: disable=ungrouped-imports
- import apitools.base.py.transfer as transfer
import email.generator as email_generator
+ from apitools.base.py import transfer
class _WrapperNamespace(object):
class BytesGenerator(email_generator.BytesGenerator):
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py
index 584b32e..e52ac18 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -270,6 +270,7 @@
# pytype: skip-file
import collections
+import io
import itertools
import json
import logging
@@ -277,13 +278,18 @@
import time
import uuid
from typing import Dict
+from typing import List
+from typing import Optional
from typing import Union
+import fastavro
+
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
@@ -292,6 +298,7 @@
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
@@ -328,6 +335,13 @@
DatasetReference = None
TableReference = None
+_LOGGER = logging.getLogger(__name__)
+
+try:
+ import google.cloud.bigquery_storage_v1 as bq_storage
+except ImportError:
+ _LOGGER.warning('ERROR: ', exc_info=True)
+
__all__ = [
'TableRowJsonCoder',
'BigQueryDisposition',
@@ -340,8 +354,6 @@
'ReadAllFromBigQuery',
'SCHEMA_AUTODETECT',
]
-
-_LOGGER = logging.getLogger(__name__)
"""
Template for BigQuery jobs created by BigQueryIO. This template is:
`"beam_bq_job_{job_type}_{job_id}_{step_id}_{random}"`, where:
@@ -645,6 +657,7 @@
class _CustomBigQuerySource(BoundedSource):
def __init__(
self,
+ method,
gcs_location=None,
table=None,
dataset=None,
@@ -682,6 +695,7 @@
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
@@ -707,6 +721,7 @@
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),
@@ -899,6 +914,388 @@
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. The ID must contain only
+ letters ``a-z``, ``A-Z``, numbers ``0-9``, or underscores ``_`` 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):
+
+ 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._job_name = job_name or 'BQ_DIRECT_READ_JOB'
+ self._step_name = step_name
+ self._source_uuid = unique_id
+
+ 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)
+
+ @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)
+ }
+
+ def estimate_size(self):
+ # Returns the pre-filtering size of the (temporary) table being read.
+ bq = bigquery_tools.BigQueryWrapper()
+ if self.table_reference is not None:
+ return self._get_table_size(bq, self.table_reference)
+ 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))
+ size = int(job.statistics.totalBytesProcessed)
+ return size
+ 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))
+
+ if self.query is not None:
+ self._setup_temporary_dataset(bq)
+ self.table_reference = self._execute_query(bq)
+
+ 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)
+
+ 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)
+ _LOGGER.info(
+ 'Sent BigQuery Storage API CreateReadSession request: \n %s \n'
+ 'Received response \n %s.',
+ requested_session,
+ read_session)
+
+ self.split_result = [
+ _CustomBigQueryStorageStreamSource(
+ stream.name, self.use_native_datetime)
+ 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."""
+ def __init__(
+ self, read_stream_name: str, use_native_datetime: Optional[bool] = True):
+ self.read_stream_name = read_stream_name
+ self.use_native_datetime = use_native_datetime
+
+ 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(BEAM-12990): 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(BEAM-12989): 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()
+
+ def read_arrow(self):
+ storage_client = bq_storage.BigQueryReadClient()
+ row_iter = iter(storage_client.read_rows(self.read_stream_name).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_iterator = iter(storage_client.read_rows(self.read_stream_name))
+ # 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:
+ 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")
class BigQuerySink(dataflow_io.NativeSink):
"""A sink based on a BigQuery table.
@@ -1860,6 +2257,17 @@
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.
+ 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. The ID must contain only letters ``a-z``, ``A-Z``,
numbers ``0-9``, or underscores ``_``. If dataset argument is
@@ -1933,11 +2341,30 @@
reading from a table rather than a query. To learn more about query
priority, see: https://cloud.google.com/bigquery/docs/running-queries
"""
+ class Method(object):
+ EXPORT = 'EXPORT' # This is currently the default.
+ DIRECT_READ = 'DIRECT_READ'
COUNTER = 0
- def __init__(self, gcs_location=None, *args, **kwargs):
- if gcs_location:
+ def __init__(
+ self,
+ gcs_location=None,
+ method=None,
+ use_native_datetime=False,
+ *args,
+ **kwargs):
+ self.method = method or ReadFromBigQuery.Method.EXPORT
+ self.use_native_datetime = use_native_datetime
+
+ if self.method is 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 is ReadFromBigQuery.Method.EXPORT:
if not isinstance(gcs_location, (str, ValueProvider)):
raise TypeError(
'%s: gcs_location must be of type string'
@@ -1947,10 +2374,24 @@
gcs_location = StaticValueProvider(str, gcs_location)
self.gcs_location = gcs_location
+ self.bigquery_dataset_labels = {
+ 'type': 'bq_direct_read_' + str(uuid.uuid4())[0:10]
+ }
self._args = args
self._kwargs = kwargs
def expand(self, pcoll):
+ if self.method is ReadFromBigQuery.Method.EXPORT:
+ return self._expand_export(pcoll)
+ elif self.method is ReadFromBigQuery.Method.DIRECT_READ:
+ return self._expand_direct_read(pcoll)
+ else:
+ raise ValueError(
+ 'The method to read from BigQuery must be either EXPORT'
+ 'or DIRECT_READ.')
+
+ def _expand_export(self, pcoll):
+ # TODO(BEAM-11115): 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
@@ -1978,6 +2419,7 @@
_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,
@@ -1985,6 +2427,45 @@
**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
+
+ def _get_pipeline_details(unused_elm):
+ 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
+ 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:
"""
@@ -2068,7 +2549,8 @@
validate: bool = False,
kms_key: str = None,
temp_dataset: Union[str, DatasetReference] = None,
- bigquery_job_labels: Dict[str, str] = 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(
@@ -2081,6 +2563,7 @@
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
@@ -2105,7 +2588,8 @@
unique_id=unique_id,
kms_key=self.kms_key,
project=project,
- temp_dataset=self.temp_dataset)).with_outputs(
+ temp_dataset=self.temp_dataset,
+ query_priority=self.query_priority)).with_outputs(
"location_to_cleanup", main="files_to_read")
)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
index f5b9593..25ffd35 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
@@ -132,6 +132,52 @@
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):
+ 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()
+ 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.
@@ -149,7 +195,8 @@
unique_id: str = None,
kms_key: str = None,
project: str = None,
- temp_dataset: Union[str, DatasetReference] = None):
+ temp_dataset: Union[str, DatasetReference] = None,
+ query_priority: Optional[str] = None):
self.options = options
self.use_json_exports = use_json_exports
self.gcs_location = gcs_location
@@ -160,6 +207,7 @@
self.kms_key = kms_key
self.project = project
self.temp_dataset = temp_dataset or 'bq_read_all_%s' % uuid.uuid4().hex
+ self.query_priority = query_priority
self.bq_io_metadata = None
def display_data(self):
@@ -254,6 +302,7 @@
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))
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py b/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py
index 9216a9c..9101039 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py
@@ -26,12 +26,14 @@
import random
import time
import unittest
+import uuid
from decimal import Decimal
from functools import wraps
import pytest
import apache_beam as beam
+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
@@ -175,6 +177,279 @@
assert_that(result, equal_to(self.TABLE_DATA))
+class ReadUsingStorageApiTests(BigQueryReadIntegrationTests):
+ TABLE_DATA = [{
+ 'number': 1,
+ 'string': u'你好',
+ 'time': '12:44:31',
+ 'datetime': '2018-12-31 12:44:31',
+ 'rec': None
+ },
+ {
+ 'number': 4,
+ 'string': u'привет',
+ '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 = 'python_read_table'
+ 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)
+ 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,
+ '%s_%s' % (int(time.time()), random.randint(0, 1000)))
+ 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)
+
+ def test_iobase_source(self):
+ EXPECTED_TABLE_DATA = [
+ {
+ 'number': 1,
+ 'string': u'你好',
+ 'time': datetime.time(12, 44, 31),
+ 'datetime': '2018-12-31T12:44:31',
+ 'rec': None,
+ },
+ {
+ 'number': 4,
+ 'string': u'привет',
+ '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))
+
+ def test_iobase_source_with_native_datetime(self):
+ EXPECTED_TABLE_DATA = [
+ {
+ 'number': 1,
+ 'string': u'你好',
+ 'time': datetime.time(12, 44, 31),
+ 'datetime': datetime.datetime(2018, 12, 31, 12, 44, 31),
+ 'rec': None,
+ },
+ {
+ 'number': 4,
+ 'string': u'привет',
+ '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))
+
+ 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))
+
+ def test_iobase_source_with_row_restriction(self):
+ EXPECTED_TABLE_DATA = [{
+ 'number': 1,
+ 'string': u'你好',
+ '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))
+
+ def test_iobase_source_with_column_selection_and_row_restriction(self):
+ EXPECTED_TABLE_DATA = [{'string': u'привет'}]
+ 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))
+
+ 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([]))
+
+ def test_iobase_source_with_query(self):
+ EXPECTED_TABLE_DATA = [
+ {
+ 'number': 1,
+ 'string': u'你好',
+ 'time': datetime.time(12, 44, 31),
+ 'datetime': datetime.datetime(2018, 12, 31, 12, 44, 31),
+ 'rec': None,
+ },
+ {
+ 'number': 4,
+ 'string': u'привет',
+ '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))
+
+ def test_iobase_source_with_query_and_filters(self):
+ EXPECTED_TABLE_DATA = [{'string': u'привет'}]
+ 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):
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py b/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py
index 957028c..0b4cfe2 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py
@@ -80,7 +80,7 @@
class BigQueryReadPerfTest(LoadTest):
def __init__(self):
- super(BigQueryReadPerfTest, self).__init__()
+ 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()
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py
index f6e014f..078b6e9 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py
@@ -464,6 +464,7 @@
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',
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index 1c4b62a..c1bc2d9 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -286,6 +286,27 @@
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
+
+
class BigQueryWrapper(object):
"""BigQuery client wrapper with utilities for querying.
@@ -302,7 +323,7 @@
HISTOGRAM_METRIC_LOGGER = MetricLogger()
- def __init__(self, client=None, temp_dataset_id=None):
+ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
self.client = client or bigquery.BigqueryV2(
http=get_new_http(),
credentials=auth.get_service_credentials(),
@@ -312,17 +333,30 @@
# 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._temporary_table_suffix = uuid.uuid4().hex
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)
- self.temp_dataset_id = temp_dataset_id or self._get_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
@@ -341,12 +375,17 @@
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_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(
@@ -734,7 +773,8 @@
@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):
+ def get_or_create_dataset(
+ self, project_id, dataset_id, location=None, labels=None):
# Check if dataset already exists otherwise create it
try:
dataset = self.client.datasets.Get(
@@ -749,6 +789,8 @@
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)
request = bigquery.BigqueryDatasetsInsertRequest(
projectId=project_id, dataset=dataset)
response = self.client.datasets.Insert(request)
@@ -820,7 +862,7 @@
@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):
+ def create_temporary_dataset(self, project_id, location, labels=None):
# Check if dataset exists to make sure that the temporary id is unique
try:
self.client.datasets.Get(
@@ -841,7 +883,7 @@
self.temp_dataset_id,
location)
self.get_or_create_dataset(
- project_id, self.temp_dataset_id, location=location)
+ project_id, self.temp_dataset_id, location=location, labels=labels)
else:
raise
@@ -883,6 +925,48 @@
@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
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py b/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py
index 3936923..1aafb1b 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py
@@ -68,7 +68,7 @@
class BigQueryWritePerfTest(LoadTest):
def __init__(self):
- super(BigQueryWritePerfTest, self).__init__()
+ super().__init__()
self.output_dataset = self.pipeline.get_option('output_dataset')
self.output_table = self.pipeline.get_option('output_table')
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py
index 4887c11..f41f802 100644
--- a/sdks/python/apache_beam/io/gcp/bigtableio.py
+++ b/sdks/python/apache_beam/io/gcp/bigtableio.py
@@ -40,13 +40,38 @@
import logging
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.transforms.display import DisplayDataItem
_LOGGER = logging.getLogger(__name__)
try:
from google.cloud.bigtable import Client
+ from google.cloud.bigtable.batcher import MutationsBatcher
+
+ FLUSH_COUNT = 1000
+ MAX_ROW_BYTES = 5242880 # 5MB
+
+ class _MutationsBatcher(MutationsBatcher):
+ def __init__(
+ self, table, flush_count=FLUSH_COUNT, max_row_bytes=MAX_ROW_BYTES):
+ super().__init__(table, flush_count, max_row_bytes)
+ self.rows = []
+
+ def set_flush_callback(self, callback_fn):
+ self.callback_fn = callback_fn
+
+ def flush(self):
+ if len(self.rows) != 0:
+ rows = self.table.mutate_rows(self.rows)
+ self.callback_fn(rows)
+ self.total_mutation_count = 0
+ self.total_size = 0
+ self.rows = []
+
except ImportError:
_LOGGER.warning(
'ImportError: from google.cloud.bigtable import Client', exc_info=True)
@@ -70,7 +95,7 @@
instance_id(str): GCP Instance to write the Rows
table_id(str): GCP Table to write the `DirectRows`
"""
- super(_BigTableWriteFn, self).__init__()
+ super().__init__()
self.beam_options = {
'project_id': project_id,
'instance_id': instance_id,
@@ -78,6 +103,7 @@
}
self.table = None
self.batcher = None
+ self.service_call_metric = None
self.written = Metrics.counter(self.__class__, 'Written Row')
def __getstate__(self):
@@ -87,14 +113,44 @@
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, rows):
+ for status in rows:
+ grpc_status_string = (
+ ServiceCallMetric.bigtable_error_code_to_grpc_status_string(
+ status.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.batcher = self.table.mutations_batcher()
+ 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)
+ self.batcher.set_flush_callback(self.write_mutate_metrics)
def process(self, row):
self.written.inc()
@@ -136,7 +192,7 @@
instance_id(str): GCP Instance to write the Rows
table_id(str): GCP Table to write the `DirectRows`
"""
- super(WriteToBigTable, self).__init__()
+ super().__init__()
self.beam_options = {
'project_id': project_id,
'instance_id': instance_id,
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py
new file mode 100644
index 0000000..29703ea
--- /dev/null
+++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py
@@ -0,0 +1,137 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT 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."""
+
+# pytype: skip-file
+import datetime
+import string
+import unittest
+import uuid
+from random import choice
+
+from mock import MagicMock
+from mock import patch
+
+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
+
+# Protect against environments where bigtable library is not available.
+try:
+ from google.cloud.bigtable import client, row
+ from google.cloud.bigtable.instance import Instance
+ from google.cloud.bigtable.table import Table
+ from google.rpc.code_pb2 import OK, ALREADY_EXISTS
+ from google.rpc.status_pb2 import Status
+except ImportError as e:
+ client = None
+
+
+@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_metrics(self):
+ MetricsEnvironment.process_wide_container().reset()
+ write_fn = bigtableio._BigTableWriteFn(
+ self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID)
+ 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)
+ write_fn.finish_bundle()
+ 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 = 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.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)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
index 34539ec..4ac2803 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
@@ -114,7 +114,7 @@
used to fetch entities.
num_splits: (:class:`int`) (optional) Number of splits for the query.
"""
- super(ReadFromDatastore, self).__init__()
+ super().__init__()
if not query.project:
raise ValueError("query.project cannot be empty")
@@ -168,7 +168,7 @@
class _SplitQueryFn(DoFn):
"""A `DoFn` that splits a given query into multiple sub-queries."""
def __init__(self, num_splits):
- super(ReadFromDatastore._SplitQueryFn, self).__init__()
+ super().__init__()
self._num_splits = num_splits
def process(self, query, *args, **kwargs):
@@ -529,8 +529,7 @@
estimate appropriate limits during ramp-up throttling.
"""
mutate_fn = WriteToDatastore._DatastoreWriteFn(project)
- super(WriteToDatastore,
- self).__init__(mutate_fn, throttle_rampup, hint_num_workers)
+ super().__init__(mutate_fn, throttle_rampup, hint_num_workers)
class _DatastoreWriteFn(_Mutate.DatastoreMutateFn):
def element_to_client_batch_item(self, element):
@@ -583,8 +582,7 @@
estimate appropriate limits during ramp-up throttling.
"""
mutate_fn = DeleteFromDatastore._DatastoreDeleteFn(project)
- super(DeleteFromDatastore,
- self).__init__(mutate_fn, throttle_rampup, hint_num_workers)
+ super().__init__(mutate_fn, throttle_rampup, hint_num_workers)
class _DatastoreDeleteFn(_Mutate.DatastoreMutateFn):
def element_to_client_batch_item(self, element):
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py
index 8ac5b93..842579d 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py
@@ -146,7 +146,7 @@
def __lt__(self, other):
if not isinstance(other, IdOrName):
- return super(IdOrName, self).__lt__(other)
+ return super().__lt__(other)
if self.id is not None:
if other.id is None:
@@ -161,7 +161,7 @@
def __eq__(self, other):
if not isinstance(other, IdOrName):
- return super(IdOrName, self).__eq__(other)
+ return super().__eq__(other)
return self.id == other.id and self.name == other.name
def __hash__(self):
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn.py
index bf54401..034e230 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn.py
@@ -51,7 +51,7 @@
num_workers: A hint for the expected number of workers, used to derive
the local rate limit.
"""
- super(RampupThrottlingFn, self).__init__(*unused_args, **unused_kwargs)
+ 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()
diff --git a/sdks/python/apache_beam/io/gcp/experimental/spannerio.py b/sdks/python/apache_beam/io/gcp/experimental/spannerio.py
index d50b3a8..fdd08a8 100644
--- a/sdks/python/apache_beam/io/gcp/experimental/spannerio.py
+++ b/sdks/python/apache_beam/io/gcp/experimental/spannerio.py
@@ -669,7 +669,7 @@
return p
def display_data(self):
- res = dict()
+ res = {}
sql = []
table = []
if self._read_operations is not None:
diff --git a/sdks/python/apache_beam/io/gcp/experimental/spannerio_read_perf_test.py b/sdks/python/apache_beam/io/gcp/experimental/spannerio_read_perf_test.py
index ae7c1ea..18f6c29 100644
--- a/sdks/python/apache_beam/io/gcp/experimental/spannerio_read_perf_test.py
+++ b/sdks/python/apache_beam/io/gcp/experimental/spannerio_read_perf_test.py
@@ -80,7 +80,7 @@
class SpannerReadPerfTest(LoadTest):
def __init__(self):
- super(SpannerReadPerfTest, self).__init__()
+ 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')
diff --git a/sdks/python/apache_beam/io/gcp/experimental/spannerio_write_perf_test.py b/sdks/python/apache_beam/io/gcp/experimental/spannerio_write_perf_test.py
index 707db81..c61608f 100644
--- a/sdks/python/apache_beam/io/gcp/experimental/spannerio_write_perf_test.py
+++ b/sdks/python/apache_beam/io/gcp/experimental/spannerio_write_perf_test.py
@@ -76,7 +76,7 @@
TEST_DATABASE = None
def __init__(self):
- super(SpannerWritePerfTest, self).__init__()
+ 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')
diff --git a/sdks/python/apache_beam/io/gcp/gcsfilesystem.py b/sdks/python/apache_beam/io/gcp/gcsfilesystem.py
index 0b20718..e53ceef 100644
--- a/sdks/python/apache_beam/io/gcp/gcsfilesystem.py
+++ b/sdks/python/apache_beam/io/gcp/gcsfilesystem.py
@@ -328,6 +328,7 @@
match_result = self.match([path_to_use])[0]
statuses = gcsio.GcsIO().delete_batch(
[m.path for m in match_result.metadata_list])
+ # pylint: disable=used-before-assignment
failures = [e for (_, e) in statuses if e is not None]
if failures:
raise failures[0]
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py
index b9fb15f..1b31284 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -53,9 +53,9 @@
try:
# pylint: disable=wrong-import-order, wrong-import-position
# pylint: disable=ungrouped-imports
- import apitools.base.py.transfer as transfer
from apitools.base.py.batch import BatchApiRequest
from apitools.base.py.exceptions import HttpError
+ from apitools.base.py import transfer
from apache_beam.internal.gcp import auth
from apache_beam.io.gcp.internal.clients import storage
except ImportError:
diff --git a/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py b/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py
index b06e374..e13d993 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py
@@ -141,7 +141,7 @@
max_bytes_rewritten_per_call=50 * 1024 * 1024,
src=self.INPUT_FILE_LARGE)
# Verify that there was a multi-part rewrite.
- self.assertTrue(any([not r.done for r in rewrite_responses]))
+ self.assertTrue(any(not r.done for r in rewrite_responses))
def _test_copy_batch(
self,
@@ -195,7 +195,7 @@
max_bytes_rewritten_per_call=50 * 1024 * 1024,
src=self.INPUT_FILE_LARGE)
# Verify that there was a multi-part rewrite.
- self.assertTrue(any([not r.done for r in rewrite_responses]))
+ self.assertTrue(any(not r.done for r in rewrite_responses))
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_client.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_client.py
index b695843..146872f 100644
--- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_client.py
+++ b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_client.py
@@ -48,7 +48,7 @@
additional_http_headers=None, response_encoding=None):
"""Create a new bigquery handle."""
url = url or self.BASE_URL
- super(BigqueryV2, self).__init__(
+ super().__init__(
url, credentials=credentials,
get_credentials=get_credentials, http=http, model=model,
log_request=log_request, log_response=log_response,
@@ -71,7 +71,7 @@
_NAME = 'datasets'
def __init__(self, client):
- super(BigqueryV2.DatasetsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -237,7 +237,7 @@
_NAME = 'jobs'
def __init__(self, client):
- super(BigqueryV2.JobsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
'Insert': base_api.ApiUploadInfo(
accept=['*/*'],
@@ -415,7 +415,7 @@
_NAME = 'models'
def __init__(self, client):
- super(BigqueryV2.ModelsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -533,7 +533,7 @@
_NAME = 'projects'
def __init__(self, client):
- super(BigqueryV2.ProjectsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -595,7 +595,7 @@
_NAME = 'routines'
def __init__(self, client):
- super(BigqueryV2.RoutinesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -740,7 +740,7 @@
_NAME = 'rowAccessPolicies'
def __init__(self, client):
- super(BigqueryV2.RowAccessPoliciesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -777,7 +777,7 @@
_NAME = 'tabledata'
def __init__(self, client):
- super(BigqueryV2.TabledataService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -839,7 +839,7 @@
_NAME = 'tables'
def __init__(self, client):
- super(BigqueryV2.TablesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/storage/storage_v1_client.py b/sdks/python/apache_beam/io/gcp/internal/clients/storage/storage_v1_client.py
index 9a6e426..cd84001 100644
--- a/sdks/python/apache_beam/io/gcp/internal/clients/storage/storage_v1_client.py
+++ b/sdks/python/apache_beam/io/gcp/internal/clients/storage/storage_v1_client.py
@@ -48,7 +48,7 @@
additional_http_headers=None, response_encoding=None):
"""Create a new storage handle."""
url = url or self.BASE_URL
- super(StorageV1, self).__init__(
+ super().__init__(
url, credentials=credentials,
get_credentials=get_credentials, http=http, model=model,
log_request=log_request, log_response=log_response, num_retries=20,
@@ -73,7 +73,7 @@
_NAME = u'bucketAccessControls'
def __init__(self, client):
- super(StorageV1.BucketAccessControlsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -239,7 +239,7 @@
_NAME = u'buckets'
def __init__(self, client):
- super(StorageV1.BucketsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -509,7 +509,7 @@
_NAME = u'channels'
def __init__(self, client):
- super(StorageV1.ChannelsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -545,7 +545,7 @@
_NAME = u'defaultObjectAccessControls'
def __init__(self, client):
- super(StorageV1.DefaultObjectAccessControlsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -711,7 +711,7 @@
_NAME = u'notifications'
def __init__(self, client):
- super(StorageV1.NotificationsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -825,7 +825,7 @@
_NAME = u'objectAccessControls'
def __init__(self, client):
- super(StorageV1.ObjectAccessControlsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -991,7 +991,7 @@
_NAME = u'objects'
def __init__(self, client):
- super(StorageV1.ObjectsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
'Insert': base_api.ApiUploadInfo(
accept=['*/*'],
@@ -1354,7 +1354,7 @@
_NAME = u'projects_serviceAccount'
def __init__(self, client):
- super(StorageV1.ProjectsServiceAccountService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1390,6 +1390,6 @@
_NAME = u'projects'
def __init__(self, client):
- super(StorageV1.ProjectsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py b/sdks/python/apache_beam/io/gcp/pubsub.py
index 39312bf..0048f04 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub.py
@@ -219,7 +219,7 @@
timestamp is optional, and digits beyond the first three (i.e., time
units smaller than milliseconds) may be ignored.
"""
- super(ReadFromPubSub, self).__init__()
+ super().__init__()
self.with_attributes = with_attributes
self._source = _PubSubSource(
topic=topic,
@@ -250,7 +250,7 @@
class _ReadStringsFromPubSub(PTransform):
"""This class is deprecated. Use ``ReadFromPubSub`` instead."""
def __init__(self, topic=None, subscription=None, id_label=None):
- super(_ReadStringsFromPubSub, self).__init__()
+ super().__init__()
self.topic = topic
self.subscription = subscription
self.id_label = id_label
@@ -278,7 +278,7 @@
Attributes:
topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
"""
- super(_WriteStringsToPubSub, self).__init__()
+ super().__init__()
self.topic = topic
def expand(self, pcoll):
@@ -315,7 +315,7 @@
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(WriteToPubSub, self).__init__()
+ super().__init__()
self.with_attributes = with_attributes
self.id_label = id_label
self.timestamp_attribute = timestamp_attribute
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_io_perf_test.py b/sdks/python/apache_beam/io/gcp/pubsub_io_perf_test.py
index 674f4e4..8cefab3 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_io_perf_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_io_perf_test.py
@@ -111,7 +111,7 @@
class PubsubWritePerfTest(PubsubIOPerfTest):
def __init__(self):
- super(PubsubWritePerfTest, self).__init__(WRITE_METRICS_NAMESPACE)
+ super().__init__(WRITE_METRICS_NAMESPACE)
self._setup_env()
self._setup_pubsub()
self._setup_pipeline()
@@ -144,7 +144,7 @@
self.pipeline = TestPipeline(options=options)
def _setup_pubsub(self):
- super(PubsubWritePerfTest, self)._setup_pubsub()
+ super()._setup_pubsub()
_ = self.pub_client.create_topic(self.topic_name)
_ = self.sub_client.create_subscription(
@@ -155,7 +155,7 @@
class PubsubReadPerfTest(PubsubIOPerfTest):
def __init__(self):
- super(PubsubReadPerfTest, self).__init__(READ_METRICS_NAMESPACE)
+ super().__init__(READ_METRICS_NAMESPACE)
self._setup_env()
self._setup_pubsub()
self._setup_pipeline()
@@ -183,7 +183,7 @@
| 'Write to Pubsub' >> beam.io.WriteToPubSub(self.matcher_topic_name))
def _setup_pubsub(self):
- super(PubsubReadPerfTest, self)._setup_pubsub()
+ super()._setup_pubsub()
_ = self.pub_client.create_topic(self.matcher_topic_name)
_ = self.sub_client.create_subscription(
diff --git a/sdks/python/apache_beam/io/gcp/resource_identifiers.py b/sdks/python/apache_beam/io/gcp/resource_identifiers.py
index a89a9e1..573d8ac 100644
--- a/sdks/python/apache_beam/io/gcp/resource_identifiers.py
+++ b/sdks/python/apache_beam/io/gcp/resource_identifiers.py
@@ -42,3 +42,8 @@
def DatastoreNamespace(project_id, namespace_id):
return '//bigtable.googleapis.com/projects/%s/namespaces/%s' % (
project_id, namespace_id)
+
+
+def BigtableTable(project_id, instance_id, table_id):
+ return '//bigtable.googleapis.com/projects/%s/instances/%s/tables/%s' % (
+ project_id, instance_id, table_id)
diff --git a/sdks/python/apache_beam/io/gcp/spanner.py b/sdks/python/apache_beam/io/gcp/spanner.py
index 60fae04..a5d3d14 100644
--- a/sdks/python/apache_beam/io/gcp/spanner.py
+++ b/sdks/python/apache_beam/io/gcp/spanner.py
@@ -246,7 +246,7 @@
assert timestamp_bound_mode is TimestampBoundMode.MIN_READ_TIMESTAMP\
or timestamp_bound_mode is TimestampBoundMode.READ_TIMESTAMP
- super(ReadFromSpanner, self).__init__(
+ super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
ReadFromSpannerSchema(
diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
index f30baad..4504ba4 100644
--- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
+++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
@@ -148,8 +148,7 @@
query: The query (string) to perform.
data: List of tuples with the expected data.
"""
- super(BigqueryFullResultMatcher,
- self).__init__(project, query, 'unused_checksum')
+ super().__init__(project, query, 'unused_checksum')
self.expected_data = data
self.actual_data = None
@@ -191,8 +190,7 @@
DEFAULT_TIMEOUT = 5 * 60
def __init__(self, project, query, data, timeout=DEFAULT_TIMEOUT):
- super(BigqueryFullResultStreamingMatcher,
- self).__init__(project, query, data)
+ super().__init__(project, query, data)
self.timeout = timeout
def _get_query_result(self):
diff --git a/sdks/python/apache_beam/io/gcp/tests/xlang_spannerio_it_test.py b/sdks/python/apache_beam/io/gcp/tests/xlang_spannerio_it_test.py
index eb0c71c..5d70105 100644
--- a/sdks/python/apache_beam/io/gcp/tests/xlang_spannerio_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/tests/xlang_spannerio_it_test.py
@@ -227,7 +227,7 @@
_ = (
p
| 'Impulse' >> beam.Impulse()
- | 'Generate' >> beam.FlatMap(lambda x: range(num_rows)) # pylint: disable=range-builtin-not-iterating
+ | '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,
diff --git a/sdks/python/apache_beam/io/hadoopfilesystem.py b/sdks/python/apache_beam/io/hadoopfilesystem.py
index 041b9c5..046908d 100644
--- a/sdks/python/apache_beam/io/hadoopfilesystem.py
+++ b/sdks/python/apache_beam/io/hadoopfilesystem.py
@@ -105,7 +105,7 @@
Connection configuration is done by passing pipeline options.
See :class:`~apache_beam.options.pipeline_options.HadoopFileSystemOptions`.
"""
- super(HadoopFileSystem, self).__init__(pipeline_options)
+ super().__init__(pipeline_options)
logging.getLogger('hdfs.client').setLevel(logging.WARN)
if pipeline_options is None:
raise ValueError('pipeline_options is not set')
diff --git a/sdks/python/apache_beam/io/iobase.py b/sdks/python/apache_beam/io/iobase.py
index 8e81da3..c0d34d8 100644
--- a/sdks/python/apache_beam/io/iobase.py
+++ b/sdks/python/apache_beam/io/iobase.py
@@ -874,7 +874,7 @@
Args:
source: Data source to read from.
"""
- super(Read, self).__init__()
+ super().__init__()
self.source = source
@staticmethod
@@ -1048,7 +1048,7 @@
Args:
sink: Data sink to write to.
"""
- super(Write, self).__init__()
+ super().__init__()
self.sink = sink
def display_data(self):
@@ -1083,7 +1083,7 @@
timestamp_attribute=self.sink.timestamp_attribute)
return (common_urns.composites.PUBSUB_WRITE.urn, payload)
else:
- return super(Write, self).to_runner_api_parameter(context)
+ return super().to_runner_api_parameter(context)
@staticmethod
@ptransform.PTransform.register_urn(
@@ -1117,7 +1117,7 @@
"""Implements the writing of custom sinks."""
def __init__(self, sink):
# type: (Sink) -> None
- super(WriteImpl, self).__init__()
+ super().__init__()
self.sink = sink
def expand(self, pcoll):
@@ -1647,7 +1647,7 @@
"""
def __init__(self, data_to_display=None):
self._data_to_display = data_to_display or {}
- super(SDFBoundedSourceReader, self).__init__()
+ super().__init__()
def _create_sdf_bounded_source_dofn(self):
class SDFBoundedSourceDoFn(core.DoFn):
diff --git a/sdks/python/apache_beam/io/iobase_test.py b/sdks/python/apache_beam/io/iobase_test.py
index bde0566..eb9617c 100644
--- a/sdks/python/apache_beam/io/iobase_test.py
+++ b/sdks/python/apache_beam/io/iobase_test.py
@@ -80,10 +80,9 @@
split_bundles = list(
self.sdf_restriction_provider.split(element, restriction))
self.assertTrue(
- all([
+ all(
isinstance(bundle._source_bundle, SourceBundle)
- for bundle in split_bundles
- ]))
+ for bundle in split_bundles))
splits = ([(
bundle._source_bundle.start_position,
@@ -101,10 +100,9 @@
sdf_concat_restriction_provider.split(
initial_concat_source, restriction))
self.assertTrue(
- all([
+ all(
isinstance(bundle._source_bundle, SourceBundle)
- for bundle in split_bundles
- ]))
+ for bundle in split_bundles))
splits = ([(
bundle._source_bundle.start_position,
bundle._source_bundle.stop_position) for bundle in split_bundles])
diff --git a/sdks/python/apache_beam/io/jdbc.py b/sdks/python/apache_beam/io/jdbc.py
index 060a4d8..afd39e0 100644
--- a/sdks/python/apache_beam/io/jdbc.py
+++ b/sdks/python/apache_beam/io/jdbc.py
@@ -183,7 +183,7 @@
:param expansion_service: The address (host:port) of the ExpansionService.
"""
- super(WriteToJdbc, self).__init__(
+ super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
JdbcConfigSchema(
@@ -269,7 +269,7 @@
passed as list of strings
:param expansion_service: The address (host:port) of the ExpansionService.
"""
- super(ReadFromJdbc, self).__init__(
+ super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
JdbcConfigSchema(
diff --git a/sdks/python/apache_beam/io/kafka.py b/sdks/python/apache_beam/io/kafka.py
index 3e28cfe..8d58dc3 100644
--- a/sdks/python/apache_beam/io/kafka.py
+++ b/sdks/python/apache_beam/io/kafka.py
@@ -171,7 +171,7 @@
'timestamp_policy should be one of '
'[ProcessingTime, CreateTime, LogAppendTime]')
- super(ReadFromKafka, self).__init__(
+ super().__init__(
self.URN_WITH_METADATA if with_metadata else self.URN_WITHOUT_METADATA,
NamedTupleBasedPayloadBuilder(
ReadFromKafkaSchema(
@@ -234,7 +234,7 @@
Default: 'org.apache.kafka.common.serialization.ByteArraySerializer'.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
- super(WriteToKafka, self).__init__(
+ super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
WriteToKafkaSchema(
diff --git a/sdks/python/apache_beam/io/kinesis.py b/sdks/python/apache_beam/io/kinesis.py
index 4a70f76..aca0dc1 100644
--- a/sdks/python/apache_beam/io/kinesis.py
+++ b/sdks/python/apache_beam/io/kinesis.py
@@ -153,7 +153,7 @@
Example: {'CollectionMaxCount': '1000', 'ConnectTimeout': '10000'}
:param expansion_service: The address (host:port) of the ExpansionService.
"""
- super(WriteToKinesis, self).__init__(
+ super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
WriteToKinesisSchema(
@@ -277,7 +277,7 @@
):
logging.warning('Provided timestamp emplaced not in the past.')
- super(ReadDataFromKinesis, self).__init__(
+ super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
ReadFromKinesisSchema(
diff --git a/sdks/python/apache_beam/io/mongodbio.py b/sdks/python/apache_beam/io/mongodbio.py
index c6f7d97..f6c1bdf 100644
--- a/sdks/python/apache_beam/io/mongodbio.py
+++ b/sdks/python/apache_beam/io/mongodbio.py
@@ -299,8 +299,7 @@
# for desired bundle size, if desired chunk size smaller than 1mb, use
# MongoDB default split size of 1mb.
- if desired_bundle_size_in_mb < 1:
- desired_bundle_size_in_mb = 1
+ 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(
@@ -765,7 +764,7 @@
self.batch = []
def display_data(self):
- res = super(_WriteMongoFn, self).display_data()
+ res = super().display_data()
res["database"] = self.db
res["collection"] = self.coll
res["batch_size"] = self.batch_size
diff --git a/sdks/python/apache_beam/io/parquetio.py b/sdks/python/apache_beam/io/parquetio.py
index 188aab7..872140d 100644
--- a/sdks/python/apache_beam/io/parquetio.py
+++ b/sdks/python/apache_beam/io/parquetio.py
@@ -124,7 +124,7 @@
'a.b', 'a.c', and 'a.d.e'
"""
- super(ReadFromParquetBatched, self).__init__()
+ super().__init__()
self._source = _create_parquet_source(
file_pattern,
min_bundle_size,
@@ -190,7 +190,7 @@
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(ReadFromParquet, self).__init__()
+ super().__init__()
self._source = _create_parquet_source(
file_pattern,
min_bundle_size,
@@ -237,7 +237,7 @@
name and the value being the actual data. If False, it only returns
the data.
"""
- super(ReadAllFromParquetBatched, self).__init__()
+ super().__init__()
source_from_file = partial(
_create_parquet_source,
min_bundle_size=min_bundle_size,
@@ -305,7 +305,7 @@
"""A source for reading Parquet files.
"""
def __init__(self, file_pattern, min_bundle_size, validate, columns):
- super(_ParquetSource, self).__init__(
+ super().__init__(
file_pattern=file_pattern,
min_bundle_size=min_bundle_size,
validate=validate)
@@ -447,7 +447,7 @@
Returns:
A WriteToParquet transform usable for writing.
"""
- super(WriteToParquet, self).__init__()
+ super().__init__()
self._sink = \
_create_parquet_sink(
file_path_prefix,
@@ -509,7 +509,7 @@
num_shards,
shard_name_template,
mime_type):
- super(_ParquetSink, self).__init__(
+ super().__init__(
file_path_prefix,
file_name_suffix=file_name_suffix,
num_shards=num_shards,
@@ -535,7 +535,7 @@
self._file_handle = None
def open(self, temp_path):
- self._file_handle = super(_ParquetSink, self).open(temp_path)
+ self._file_handle = super().open(temp_path)
return pq.ParquetWriter(
self._file_handle,
self._schema,
@@ -565,7 +565,7 @@
self._file_handle = None
def display_data(self):
- res = super(_ParquetSink, self).display_data()
+ res = super().display_data()
res['codec'] = str(self._codec)
res['schema'] = str(self._schema)
res['row_group_buffer_size'] = str(self._row_group_buffer_size)
diff --git a/sdks/python/apache_beam/io/parquetio_it_test.py b/sdks/python/apache_beam/io/parquetio_it_test.py
index 0d3cd0d..052b54f 100644
--- a/sdks/python/apache_beam/io/parquetio_it_test.py
+++ b/sdks/python/apache_beam/io/parquetio_it_test.py
@@ -118,7 +118,7 @@
class ProducerFn(DoFn):
def __init__(self, number):
- super(ProducerFn, self).__init__()
+ super().__init__()
self._number = number
self._string_index = 0
self._number_index = 0
diff --git a/sdks/python/apache_beam/io/range_trackers.py b/sdks/python/apache_beam/io/range_trackers.py
index 33b15d5..adadc1a 100644
--- a/sdks/python/apache_beam/io/range_trackers.py
+++ b/sdks/python/apache_beam/io/range_trackers.py
@@ -48,7 +48,7 @@
OFFSET_INFINITY = float('inf')
def __init__(self, start, end):
- super(OffsetRangeTracker, self).__init__()
+ super().__init__()
if start is None:
raise ValueError('Start offset must not be \'None\'')
diff --git a/sdks/python/apache_beam/io/restriction_trackers.py b/sdks/python/apache_beam/io/restriction_trackers.py
index f2b3e1d..06b06fa 100644
--- a/sdks/python/apache_beam/io/restriction_trackers.py
+++ b/sdks/python/apache_beam/io/restriction_trackers.py
@@ -170,6 +170,6 @@
# stubs in the baseclass.
def __getattribute__(self, name):
if name.startswith('_') or name in ('try_split', ):
- return super(UnsplittableRestrictionTracker, self).__getattribute__(name)
+ return super().__getattribute__(name)
else:
return getattr(self._underling_tracker, name)
diff --git a/sdks/python/apache_beam/io/source_test_utils_test.py b/sdks/python/apache_beam/io/source_test_utils_test.py
index 6d3f2e3..081a6fc 100644
--- a/sdks/python/apache_beam/io/source_test_utils_test.py
+++ b/sdks/python/apache_beam/io/source_test_utils_test.py
@@ -21,7 +21,7 @@
import tempfile
import unittest
-import apache_beam.io.source_test_utils as source_test_utils
+from apache_beam.io import source_test_utils
from apache_beam.io.filebasedsource_test import LineSource
diff --git a/sdks/python/apache_beam/io/textio.py b/sdks/python/apache_beam/io/textio.py
index 277b993..bcec439 100644
--- a/sdks/python/apache_beam/io/textio.py
+++ b/sdks/python/apache_beam/io/textio.py
@@ -118,7 +118,7 @@
Please refer to documentation in class `ReadFromText` for the rest
of the arguments.
"""
- super(_TextSource, self).__init__(
+ super().__init__(
file_pattern,
min_bundle_size,
compression_type=compression_type,
@@ -139,7 +139,7 @@
self._header_matcher, self._header_processor = header_processor_fns
def display_data(self):
- parent_dd = super(_TextSource, self).display_data()
+ parent_dd = super().display_data()
parent_dd['strip_newline'] = DisplayDataItem(
self._strip_trailing_newlines, label='Strip Trailing New Lines')
parent_dd['buffer_size'] = DisplayDataItem(
@@ -327,8 +327,7 @@
class _TextSourceWithFilename(_TextSource):
def read_records(self, file_name, range_tracker):
- records = super(_TextSourceWithFilename,
- self).read_records(file_name, range_tracker)
+ records = super().read_records(file_name, range_tracker)
for record in records:
yield (file_name, record)
@@ -383,7 +382,7 @@
Returns:
A _TextSink object usable for writing.
"""
- super(_TextSink, self).__init__(
+ super().__init__(
file_path_prefix,
file_name_suffix=file_name_suffix,
num_shards=num_shards,
@@ -396,7 +395,7 @@
self._footer = footer
def open(self, temp_path):
- file_handle = super(_TextSink, self).open(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:
@@ -408,10 +407,10 @@
file_handle.write(coders.ToBytesCoder().encode(self._footer))
if self._append_trailing_newlines:
file_handle.write(b'\n')
- super(_TextSink, self).close(file_handle)
+ super().close(file_handle)
def display_data(self):
- dd_parent = super(_TextSink, self).display_data()
+ dd_parent = super().display_data()
dd_parent['append_newline'] = DisplayDataItem(
self._append_trailing_newlines, label='Append Trailing New Lines')
return dd_parent
@@ -493,7 +492,7 @@
name and the value being the actual data. If False, it only returns
the data.
"""
- super(ReadAllFromText, self).__init__(**kwargs)
+ super().__init__(**kwargs)
source_from_file = partial(
_create_text_source,
min_bundle_size=min_bundle_size,
@@ -564,7 +563,7 @@
coder (~apache_beam.coders.coders.Coder): Coder used to decode each line.
"""
- super(ReadFromText, self).__init__(**kwargs)
+ super().__init__(**kwargs)
self._source = self._source_class(
file_pattern,
min_bundle_size,
diff --git a/sdks/python/apache_beam/io/textio_test.py b/sdks/python/apache_beam/io/textio_test.py
index 6be4370..3818d36 100644
--- a/sdks/python/apache_beam/io/textio_test.py
+++ b/sdks/python/apache_beam/io/textio_test.py
@@ -29,10 +29,10 @@
import zlib
import apache_beam as beam
-import apache_beam.io.source_test_utils as source_test_utils
from apache_beam import coders
from apache_beam.io import ReadAllFromText
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.textio import _TextSink as TextSink
from apache_beam.io.textio import _TextSource as TextSource
@@ -1018,7 +1018,7 @@
class TextSinkTest(unittest.TestCase):
def setUp(self):
- super(TextSinkTest, self).setUp()
+ super().setUp()
self.lines = [b'Line %d' % d for d in range(100)]
self.tempdir = tempfile.mkdtemp()
self.path = self._create_temp_file()
diff --git a/sdks/python/apache_beam/io/tfrecordio.py b/sdks/python/apache_beam/io/tfrecordio.py
index e699756..d3bb0f8 100644
--- a/sdks/python/apache_beam/io/tfrecordio.py
+++ b/sdks/python/apache_beam/io/tfrecordio.py
@@ -168,7 +168,7 @@
"""
def __init__(self, file_pattern, coder, compression_type, validate):
"""Initialize a TFRecordSource. See ReadFromTFRecord for details."""
- super(_TFRecordSource, self).__init__(
+ super().__init__(
file_pattern=file_pattern,
compression_type=compression_type,
splittable=False,
@@ -218,7 +218,7 @@
name and the value being the actual data. If False, it only returns
the data.
"""
- super(ReadAllFromTFRecord, self).__init__()
+ super().__init__()
source_from_file = partial(
_create_tfrecordio_source,
compression_type=compression_type,
@@ -259,7 +259,7 @@
Returns:
A ReadFromTFRecord transform object.
"""
- super(ReadFromTFRecord, self).__init__()
+ super().__init__()
self._source = _TFRecordSource(
file_pattern, coder, compression_type, validate)
@@ -283,7 +283,7 @@
compression_type):
"""Initialize a TFRecordSink. See WriteToTFRecord for details."""
- super(_TFRecordSink, self).__init__(
+ super().__init__(
file_path_prefix=file_path_prefix,
coder=coder,
file_name_suffix=file_name_suffix,
@@ -330,7 +330,7 @@
Returns:
A WriteToTFRecord transform object.
"""
- super(WriteToTFRecord, self).__init__()
+ super().__init__()
self._sink = _TFRecordSink(
file_path_prefix,
coder,
diff --git a/sdks/python/apache_beam/metrics/cells.py b/sdks/python/apache_beam/metrics/cells.py
index 0c6b8f3..6402989 100644
--- a/sdks/python/apache_beam/metrics/cells.py
+++ b/sdks/python/apache_beam/metrics/cells.py
@@ -106,7 +106,7 @@
This class is thread safe.
"""
def __init__(self, *args):
- super(CounterCell, self).__init__(*args)
+ super().__init__(*args)
self.value = CounterAggregator.identity_element()
def reset(self):
@@ -170,7 +170,7 @@
This class is thread safe.
"""
def __init__(self, *args):
- super(DistributionCell, self).__init__(*args)
+ super().__init__(*args)
self.data = DistributionAggregator.identity_element()
def reset(self):
@@ -229,7 +229,7 @@
This class is thread safe.
"""
def __init__(self, *args):
- super(GaugeCell, self).__init__(*args)
+ super().__init__(*args)
self.data = GaugeAggregator.identity_element()
def reset(self):
diff --git a/sdks/python/apache_beam/metrics/execution.py b/sdks/python/apache_beam/metrics/execution.py
index 9400860..0b404de 100644
--- a/sdks/python/apache_beam/metrics/execution.py
+++ b/sdks/python/apache_beam/metrics/execution.py
@@ -78,7 +78,7 @@
"""
self.step = step
self.metric = metric
- self.labels = labels if labels else dict()
+ self.labels = labels if labels else {}
def __eq__(self, other):
return (
@@ -239,7 +239,7 @@
def __init__(self, step_name):
self.step_name = step_name
self.lock = threading.Lock()
- self.metrics = dict() # type: Dict[_TypedMetricName, MetricCell]
+ self.metrics = {} # type: Dict[_TypedMetricName, MetricCell]
def get_counter(self, metric_name):
# type: (MetricName) -> CounterCell
diff --git a/sdks/python/apache_beam/metrics/metric.py b/sdks/python/apache_beam/metrics/metric.py
index f4896e9..fca1fd0 100644
--- a/sdks/python/apache_beam/metrics/metric.py
+++ b/sdks/python/apache_beam/metrics/metric.py
@@ -122,7 +122,7 @@
"""Metrics Counter that Delegates functionality to MetricsEnvironment."""
def __init__(self, metric_name, process_wide=False):
# type: (MetricName, bool) -> None
- super(Metrics.DelegatingCounter, self).__init__(metric_name)
+ super().__init__(metric_name)
self.inc = MetricUpdater( # type: ignore[assignment]
cells.CounterCell,
metric_name,
@@ -133,14 +133,14 @@
"""Metrics Distribution Delegates functionality to MetricsEnvironment."""
def __init__(self, metric_name):
# type: (MetricName) -> None
- super(Metrics.DelegatingDistribution, self).__init__(metric_name)
+ super().__init__(metric_name)
self.update = MetricUpdater(cells.DistributionCell, metric_name) # type: ignore[assignment]
class DelegatingGauge(Gauge):
"""Metrics Gauge that Delegates functionality to MetricsEnvironment."""
def __init__(self, metric_name):
# type: (MetricName) -> None
- super(Metrics.DelegatingGauge, self).__init__(metric_name)
+ super().__init__(metric_name)
self.set = MetricUpdater(cells.GaugeCell, metric_name) # type: ignore[assignment]
diff --git a/sdks/python/apache_beam/metrics/monitoring_infos.py b/sdks/python/apache_beam/metrics/monitoring_infos.py
index 2d8faa8..ba3100c 100644
--- a/sdks/python/apache_beam/metrics/monitoring_infos.py
+++ b/sdks/python/apache_beam/metrics/monitoring_infos.py
@@ -103,6 +103,11 @@
common_urns.monitoring_info_labels.DATASTORE_PROJECT.label_props.name)
DATASTORE_NAMESPACE_LABEL = (
common_urns.monitoring_info_labels.DATASTORE_NAMESPACE.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):
@@ -184,7 +189,7 @@
ptransform: The ptransform id used as a label.
pcollection: The pcollection id used as a label.
"""
- labels = labels or dict()
+ labels = labels or {}
labels.update(create_labels(ptransform=ptransform, pcollection=pcollection))
if isinstance(metric, int):
metric = coders.VarIntCoder().encode(metric)
@@ -287,7 +292,7 @@
labels: The label dictionary to use in the MonitoringInfo.
"""
return metrics_pb2.MonitoringInfo(
- urn=urn, type=type_urn, labels=labels or dict(), payload=payload)
+ urn=urn, type=type_urn, labels=labels or {}, payload=payload)
def is_counter(monitoring_info_proto):
diff --git a/sdks/python/apache_beam/ml/gcp/videointelligenceml.py b/sdks/python/apache_beam/ml/gcp/videointelligenceml.py
index bc0aa08..fb0d7f0 100644
--- a/sdks/python/apache_beam/ml/gcp/videointelligenceml.py
+++ b/sdks/python/apache_beam/ml/gcp/videointelligenceml.py
@@ -95,7 +95,7 @@
videointelligenceml.AnnotateVideo(features,
context_side_input=beam.pvalue.AsDict(context_side_input)))
"""
- super(AnnotateVideo, self).__init__()
+ super().__init__()
self.features = features
self.location_id = location_id
self.metadata = metadata
@@ -120,7 +120,7 @@
(``google.cloud.videointelligence_v1.types.AnnotateVideoResponse``).
"""
def __init__(self, features, location_id, metadata, timeout):
- super(_VideoAnnotateFn, self).__init__()
+ super().__init__()
self._client = None
self.features = features
self.location_id = location_id
@@ -186,7 +186,7 @@
The time in seconds to wait for the response from the
Video Intelligence API
"""
- super(AnnotateVideoWithContext, self).__init__(
+ super().__init__(
features=features,
location_id=location_id,
metadata=metadata,
@@ -210,7 +210,7 @@
(``google.cloud.videointelligence_v1.types.AnnotateVideoResponse``).
"""
def __init__(self, features, location_id, metadata, timeout):
- super(_VideoAnnotateFnWithContext, self).__init__(
+ super().__init__(
features=features,
location_id=location_id,
metadata=metadata,
diff --git a/sdks/python/apache_beam/ml/gcp/visionml.py b/sdks/python/apache_beam/ml/gcp/visionml.py
index 0fb45ce..3e556b9 100644
--- a/sdks/python/apache_beam/ml/gcp/visionml.py
+++ b/sdks/python/apache_beam/ml/gcp/visionml.py
@@ -122,7 +122,7 @@
metadata: (Optional[Sequence[Tuple[str, str]]]): Optional.
Additional metadata that is provided to the method.
"""
- super(AnnotateImage, self).__init__()
+ super().__init__()
self.features = features
self.retry = retry
self.timeout = timeout
@@ -219,7 +219,7 @@
metadata: (Optional[Sequence[Tuple[str, str]]]): Optional.
Additional metadata that is provided to the method.
"""
- super(AnnotateImageWithContext, self).__init__(
+ super().__init__(
features=features,
retry=retry,
timeout=timeout,
@@ -265,7 +265,7 @@
Returns ``google.cloud.vision.types.BatchAnnotateImagesResponse``.
"""
def __init__(self, features, retry, timeout, client_options, metadata):
- super(_ImageAnnotateFn, self).__init__()
+ super().__init__()
self._client = None
self.features = features
self.retry = retry
diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py
index bba56ef..c8f31e3 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -125,7 +125,7 @@
def error(self, message):
if message.startswith('ambiguous option: '):
return
- super(_BeamArgumentParser, self).error(message)
+ super().error(message)
class PipelineOptions(HasDisplayData):
@@ -391,7 +391,7 @@
def __setattr__(self, name, value):
if name in ('_flags', '_all_options', '_visible_options'):
- super(PipelineOptions, self).__setattr__(name, value)
+ super().__setattr__(name, value)
elif name in self._visible_option_list():
self._all_options[name] = value
else:
diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py
index c83bc5c..3cf0a51 100644
--- a/sdks/python/apache_beam/pipeline.py
+++ b/sdks/python/apache_beam/pipeline.py
@@ -394,6 +394,12 @@
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.
@@ -458,15 +464,15 @@
self.visit(InputOutputUpdater(self))
- for transform in output_replacements:
- for tag, output in output_replacements[transform]:
+ for transform, output_replacement in output_replacements.items():
+ for tag, output in output_replacement:
transform.replace_output(output, tag=tag)
- for transform in input_replacements:
- transform.replace_inputs(input_replacements[transform])
+ for transform, input_replacement in input_replacements.items():
+ transform.replace_inputs(input_replacement)
- for transform in side_input_replacements:
- transform.replace_side_inputs(side_input_replacements[transform])
+ for transform, side_input_replacement in side_input_replacements.items():
+ transform.replace_side_inputs(side_input_replacement)
def _check_replacement(self, override):
# type: (PTransformOverride) -> None
diff --git a/sdks/python/apache_beam/pipeline_test.py b/sdks/python/apache_beam/pipeline_test.py
index 731d5e3..8f8f44b 100644
--- a/sdks/python/apache_beam/pipeline_test.py
+++ b/sdks/python/apache_beam/pipeline_test.py
@@ -430,7 +430,7 @@
self, applied_ptransform):
return ToStringParDo().with_input_types(int).with_output_types(str)
- for override, expected_type in [(NoTypeHintOverride(), typehints.Any),
+ for override, expected_type in [(NoTypeHintOverride(), int),
(WithTypeHintOverride(), str)]:
p = TestPipeline()
pcoll = (
@@ -997,7 +997,7 @@
return p | beam.Create([None])
def display_data(self): # type: () -> dict
- parent_dd = super(MyParentTransform, self).display_data()
+ 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')
@@ -1011,12 +1011,12 @@
return p | beam.Create([None])
def display_data(self): # type: () -> dict
- parent_dd = super(MyPTransform, self).display_data()
+ 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_int'] = DisplayDataItem(1.1, label='dd_int_label')
+ parent_dd['dd_double'] = DisplayDataItem(1.1, label='dd_double_label')
return parent_dd
p = beam.Pipeline()
@@ -1037,41 +1037,57 @@
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',
- double_value=1).SerializeToString()),
+ 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_int_label',
+ label='dd_double_label',
+ key='dd_double',
+ namespace='apache_beam.pipeline_test.MyPTransform',
double_value=1.1).SerializeToString()),
])
diff --git a/sdks/python/apache_beam/pvalue.py b/sdks/python/apache_beam/pvalue.py
index 2b593e4..df3caee 100644
--- a/sdks/python/apache_beam/pvalue.py
+++ b/sdks/python/apache_beam/pvalue.py
@@ -238,12 +238,15 @@
pipeline, # type: Pipeline
transform, # type: ParDo
tags, # type: Sequence[str]
- main_tag # type: Optional[str]
+ main_tag, # type: Optional[str]
+ allow_unknown_tags=None, # type: Optional[bool]
):
self._pipeline = pipeline
self._tags = tags
self._main_tag = main_tag
self._transform = transform
+ 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.
@@ -288,7 +291,7 @@
tag = str(tag)
if tag == self._main_tag:
tag = None
- elif self._tags and tag not in self._tags:
+ 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))
@@ -492,14 +495,14 @@
def __init__(self, pcoll, default_value=_NO_DEFAULT):
# type: (PCollection, Any) -> None
- super(AsSingleton, self).__init__(pcoll)
+ super().__init__(pcoll)
self.default_value = default_value
def __repr__(self):
return 'AsSingleton(%s)' % self.pvalue
def _view_options(self):
- base = super(AsSingleton, self)._view_options()
+ base = super()._view_options()
if self.default_value != AsSingleton._NO_DEFAULT:
return dict(base, default=self.default_value)
return base
diff --git a/sdks/python/apache_beam/runners/common.py b/sdks/python/apache_beam/runners/common.py
index 30004d1..7d5e7e3 100644
--- a/sdks/python/apache_beam/runners/common.py
+++ b/sdks/python/apache_beam/runners/common.py
@@ -112,7 +112,7 @@
user_name: The full user-given name of the step (e.g. Foo/Bar/ParDo(Far)).
system_name: The step name in the optimized graph (e.g. s2-1).
"""
- super(DataflowNameContext, self).__init__(step_name)
+ super().__init__(step_name)
self.user_name = user_name
self.system_name = system_name
@@ -557,7 +557,7 @@
signature # type: DoFnSignature
):
# type: (...) -> None
- super(SimpleInvoker, self).__init__(output_processor, signature)
+ super().__init__(output_processor, signature)
self.process_method = signature.process_method.method_value
def invoke_process(self,
@@ -567,9 +567,10 @@
additional_args=None,
additional_kwargs=None
):
- # type: (...) -> None
+ # type: (...) -> Iterable[SplitResultResidual]
self.output_processor.process_outputs(
windowed_value, self.process_method(windowed_value.value))
+ return []
class PerWindowInvoker(DoFnInvoker):
@@ -585,7 +586,7 @@
user_state_context, # type: Optional[userstate.UserStateContext]
bundle_finalizer_param # type: Optional[core._BundleFinalizerParam]
):
- super(PerWindowInvoker, self).__init__(output_processor, signature)
+ super().__init__(output_processor, signature)
self.side_inputs = side_inputs
self.context = context
self.process_method = signature.process_method.method_value
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py b/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py
index e5728e8..41c13b2 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py
@@ -75,7 +75,7 @@
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(DataflowMetrics, self).__init__()
+ super().__init__()
self._dataflow_client = dataflow_client
self.job_result = job_result
self._queried_after_termination = False
@@ -122,7 +122,7 @@
"""Populate the MetricKey object for a queried metric result."""
step = ""
name = metric.name.name # Always extract a name
- labels = dict()
+ 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:
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
index 29a6016..e78b973 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -138,7 +138,7 @@
def apply(self, transform, input, options):
self._maybe_add_unified_worker_missing_options(options)
- return super(DataflowRunner, self).apply(transform, input, options)
+ return super().apply(transform, input, options)
def _get_unique_step_name(self):
self._unique_step_id += 1
@@ -508,13 +508,6 @@
# in the proto representation of the graph.
pipeline.replace_all(DataflowRunner._NON_PORTABLE_PTRANSFORM_OVERRIDES)
- # Always upload graph out-of-band when explicitly using runner v2 with
- # use_portable_job_submission to avoid irrelevant large graph limits.
- if (apiclient._use_unified_worker(debug_options) and
- debug_options.lookup_experiment('use_portable_job_submission') and
- not debug_options.lookup_experiment('upload_graph')):
- debug_options.add_experiment("upload_graph")
-
# Add setup_options for all the BeamPlugin imports
setup_options = options.view_as(SetupOptions)
plugins = BeamPlugin.get_all_plugin_paths()
@@ -601,9 +594,16 @@
return result
def _maybe_add_unified_worker_missing_options(self, options):
+ debug_options = options.view_as(DebugOptions)
+ # Streaming is always portable, default to runner v2.
+ if (options.view_as(StandardOptions).streaming and
+ not options.view_as(GoogleCloudOptions).dataflow_kms_key):
+ if not debug_options.lookup_experiment('disable_runner_v2'):
+ debug_options.add_experiment('beam_fn_api')
+ debug_options.add_experiment('use_runner_v2')
+ debug_options.add_experiment('use_portable_job_submission')
# set default beam_fn_api experiment if use unified
# worker experiment flag exists, no-op otherwise.
- debug_options = options.view_as(DebugOptions)
from apache_beam.runners.dataflow.internal import apiclient
if apiclient._use_unified_worker(options):
if not debug_options.lookup_experiment('beam_fn_api'):
@@ -1674,5 +1674,5 @@
class DataflowRuntimeException(Exception):
"""Indicates an error has occurred in running this pipeline."""
def __init__(self, msg, result):
- super(DataflowRuntimeException, self).__init__(msg)
+ super().__init__(msg)
self.result = result
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
index e7ce71b..d5b60c7 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
@@ -71,7 +71,7 @@
# composite transforms support display data.
class SpecialParDo(beam.ParDo):
def __init__(self, fn, now):
- super(SpecialParDo, self).__init__(fn)
+ super().__init__(fn)
self.fn = fn
self.now = now
@@ -256,6 +256,7 @@
def test_streaming_create_translation(self):
remote_runner = DataflowRunner()
self.default_properties.append("--streaming")
+ self.default_properties.append("--experiments=disable_runner_v2")
with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p:
p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
job_dict = json.loads(str(remote_runner.job))
@@ -839,7 +840,8 @@
'Runner determined sharding not available in Dataflow for '
'GroupIntoBatches for jobs not using Runner V2'):
_ = self._run_group_into_batches_and_get_step_properties(
- True, ['--enable_streaming_engine'])
+ True,
+ ['--enable_streaming_engine', '--experiments=disable_runner_v2'])
# JRH
with self.assertRaisesRegex(
@@ -847,7 +849,12 @@
'Runner determined sharding not available in Dataflow for '
'GroupIntoBatches for jobs not using Runner V2'):
_ = self._run_group_into_batches_and_get_step_properties(
- True, ['--enable_streaming_engine', '--experiments=beam_fn_api'])
+ True,
+ [
+ '--enable_streaming_engine',
+ '--experiments=beam_fn_api',
+ '--experiments=disable_runner_v2'
+ ])
def test_pack_combiners(self):
class PackableCombines(beam.PTransform):
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
index c60e88e..7e2ba13 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
@@ -163,7 +163,7 @@
self.proto.userAgent = dataflow.Environment.UserAgentValue()
self.local = 'localhost' in self.google_cloud_options.dataflow_endpoint
self._proto_pipeline = proto_pipeline
- self._sdk_image_overrides = _sdk_image_overrides or dict()
+ self._sdk_image_overrides = _sdk_image_overrides or {}
if self.google_cloud_options.service_account_email:
self.proto.serviceAccountEmail = (
@@ -405,7 +405,7 @@
# 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), encoding='shortstrings'),
+ json.loads(encoding.MessageToJson(self.proto)),
indent=2,
sort_keys=True)
@@ -554,8 +554,7 @@
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 dict())
+ dict(s.split(',', 1) for s in sdk_overrides) if sdk_overrides else {})
@retry.with_exponential_backoff(
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
@@ -1031,7 +1030,7 @@
class _LegacyDataflowStager(Stager):
def __init__(self, dataflow_application_client):
- super(_LegacyDataflowStager, self).__init__()
+ super().__init__()
self._dataflow_application_client = dataflow_application_client
def stage_artifact(self, local_path_to_artifact, artifact_name):
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
index b63da70..59eaf58 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
@@ -269,7 +269,7 @@
# Accessing non-public method for testing.
apiclient.DataflowApplicationClient._apply_sdk_environment_overrides(
- proto_pipeline, dict(), pipeline_options)
+ proto_pipeline, {}, pipeline_options)
from apache_beam.utils import proto_utils
found_override = False
@@ -300,7 +300,7 @@
# Accessing non-public method for testing.
apiclient.DataflowApplicationClient._apply_sdk_environment_overrides(
- proto_pipeline, dict(), pipeline_options)
+ proto_pipeline, {}, pipeline_options)
self.assertIsNotNone(2, len(proto_pipeline.components.environments))
@@ -336,7 +336,7 @@
# Accessing non-public method for testing.
apiclient.DataflowApplicationClient._apply_sdk_environment_overrides(
- proto_pipeline, dict(), pipeline_options)
+ proto_pipeline, {}, pipeline_options)
self.assertIsNotNone(2, len(proto_pipeline.components.environments))
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_client.py b/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_client.py
index a48b9ea..985934d 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_client.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/dataflow_v1b3_client.py
@@ -47,7 +47,7 @@
additional_http_headers=None, response_encoding=None):
"""Create a new dataflow handle."""
url = url or self.BASE_URL
- super(DataflowV1b3, self).__init__(
+ super().__init__(
url, credentials=credentials,
get_credentials=get_credentials, http=http, model=model,
log_request=log_request, log_response=log_response,
@@ -83,7 +83,7 @@
_NAME = 'projects_catalogTemplates_templateVersions'
def __init__(self, client):
- super(DataflowV1b3.ProjectsCatalogTemplatesTemplateVersionsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -120,7 +120,7 @@
_NAME = 'projects_catalogTemplates'
def __init__(self, client):
- super(DataflowV1b3.ProjectsCatalogTemplatesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -265,7 +265,7 @@
_NAME = 'projects_jobs_debug'
def __init__(self, client):
- super(DataflowV1b3.ProjectsJobsDebugService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -327,7 +327,7 @@
_NAME = 'projects_jobs_messages'
def __init__(self, client):
- super(DataflowV1b3.ProjectsJobsMessagesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -363,7 +363,7 @@
_NAME = 'projects_jobs_workItems'
def __init__(self, client):
- super(DataflowV1b3.ProjectsJobsWorkItemsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -425,7 +425,7 @@
_NAME = 'projects_jobs'
def __init__(self, client):
- super(DataflowV1b3.ProjectsJobsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -617,7 +617,7 @@
_NAME = 'projects_locations_flexTemplates'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsFlexTemplatesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -653,7 +653,7 @@
_NAME = 'projects_locations_jobs_debug'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsJobsDebugService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -715,7 +715,7 @@
_NAME = 'projects_locations_jobs_messages'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsJobsMessagesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -751,7 +751,7 @@
_NAME = 'projects_locations_jobs_snapshots'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsJobsSnapshotsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -787,7 +787,7 @@
_NAME = 'projects_locations_jobs_stages'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsJobsStagesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -823,7 +823,7 @@
_NAME = 'projects_locations_jobs_workItems'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsJobsWorkItemsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -885,7 +885,7 @@
_NAME = 'projects_locations_jobs'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsJobsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1077,7 +1077,7 @@
_NAME = 'projects_locations_snapshots'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsSnapshotsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1165,7 +1165,7 @@
_NAME = 'projects_locations_sql'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsSqlService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1201,7 +1201,7 @@
_NAME = 'projects_locations_templates'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsTemplatesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1289,7 +1289,7 @@
_NAME = 'projects_locations'
def __init__(self, client):
- super(DataflowV1b3.ProjectsLocationsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1325,7 +1325,7 @@
_NAME = 'projects_snapshots'
def __init__(self, client):
- super(DataflowV1b3.ProjectsSnapshotsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1387,7 +1387,7 @@
_NAME = 'projects_templateVersions'
def __init__(self, client):
- super(DataflowV1b3.ProjectsTemplateVersionsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1424,7 +1424,7 @@
_NAME = 'projects_templates'
def __init__(self, client):
- super(DataflowV1b3.ProjectsTemplatesService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
@@ -1512,7 +1512,7 @@
_NAME = 'projects'
def __init__(self, client):
- super(DataflowV1b3.ProjectsService, self).__init__(client)
+ super().__init__(client)
self._upload_configs = {
}
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/names.py b/sdks/python/apache_beam/runners/dataflow/internal/names.py
index 837989a..7101a1e 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/names.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/names.py
@@ -36,10 +36,10 @@
# Update this version to the next version whenever there is a change that will
# require changes to legacy Dataflow worker execution environment.
-BEAM_CONTAINER_VERSION = 'beam-master-20210809'
+BEAM_CONTAINER_VERSION = 'beam-master-20210920'
# Update this version to the next version whenever there is a change that
# requires changes to SDK harness container or SDK harness launcher.
-BEAM_FNAPI_CONTAINER_VERSION = 'beam-master-20210809'
+BEAM_FNAPI_CONTAINER_VERSION = 'beam-master-20210920'
DATAFLOW_CONTAINER_IMAGE_REPOSITORY = 'gcr.io/cloud-dataflow/v1beta3'
diff --git a/sdks/python/apache_beam/runners/dataflow/native_io/iobase.py b/sdks/python/apache_beam/runners/dataflow/native_io/iobase.py
index c545ecd..3d1afe5 100644
--- a/sdks/python/apache_beam/runners/dataflow/native_io/iobase.py
+++ b/sdks/python/apache_beam/runners/dataflow/native_io/iobase.py
@@ -334,7 +334,7 @@
Args:
sink: Sink to use for the write
"""
- super(_NativeWrite, self).__init__()
+ super().__init__()
self.sink = sink
def expand(self, pcoll):
diff --git a/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py b/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
index add8885..e8a660c 100644
--- a/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
+++ b/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
@@ -198,7 +198,7 @@
class WriteToBigQueryPTransformOverride(PTransformOverride):
def __init__(self, pipeline, options):
- super(WriteToBigQueryPTransformOverride, self).__init__()
+ super().__init__()
self.options = options
self.outputs = []
@@ -214,7 +214,7 @@
gives a user-friendsly error.
"""
# Imported here to avoid circular dependencies.
- # pylint: disable=wrong-import-order, wrong-import-position
+ # pylint: disable=wrong-import-order, wrong-import-position, unused-import
from apache_beam.pipeline import PipelineVisitor
from apache_beam.io import WriteToBigQuery
diff --git a/sdks/python/apache_beam/runners/dataflow/test_dataflow_runner.py b/sdks/python/apache_beam/runners/dataflow/test_dataflow_runner.py
index b632490..d4743a5 100644
--- a/sdks/python/apache_beam/runners/dataflow/test_dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/test_dataflow_runner.py
@@ -50,8 +50,7 @@
# send this option to remote executors.
test_options.on_success_matcher = None
- self.result = super(TestDataflowRunner,
- self).run_pipeline(pipeline, options)
+ self.result = super().run_pipeline(pipeline, options)
if self.result.has_job:
# TODO(markflyhigh)(BEAM-1890): Use print since Nose dosen't show logs
# in some cases.
diff --git a/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor_test.py b/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor_test.py
index 5d8f21c..7eba868 100644
--- a/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor_test.py
+++ b/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor_test.py
@@ -147,13 +147,15 @@
# Convert to string to assert they are equal.
out_of_order_labels = {
- str(k): [str(t) for t in v_out_of_order.value_to_consumers[k]]
- for k in v_out_of_order.value_to_consumers
+ 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 v_original.value_to_consumers[k]]
- for k in v_original.value_to_consumers
+ 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)
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py
index 4149ea2..3b40ad1 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -154,7 +154,7 @@
class _GroupAlsoByWindow(ParDo):
"""The GroupAlsoByWindow transform."""
def __init__(self, windowing):
- super(_GroupAlsoByWindow, self).__init__(_GroupAlsoByWindowDoFn(windowing))
+ super().__init__(_GroupAlsoByWindowDoFn(windowing))
self.windowing = windowing
def expand(self, pcoll):
@@ -166,7 +166,7 @@
# TODO(robertwb): Support combiner lifting.
def __init__(self, windowing):
- super(_GroupAlsoByWindowDoFn, self).__init__()
+ super().__init__()
self.windowing = windowing
def infer_output_type(self, input_type):
@@ -254,7 +254,7 @@
value_type]]])
gbk_output_type = typehints.KV[key_type, typehints.Iterable[value_type]]
- # pylint: disable=bad-continuation
+ # pylint: disable=bad-option-value
return (
pcoll
| 'ReifyWindows' >> (
@@ -565,7 +565,7 @@
class DirectPipelineResult(PipelineResult):
"""A DirectPipelineResult provides access to info about a pipeline."""
def __init__(self, executor, evaluation_context):
- super(DirectPipelineResult, self).__init__(PipelineState.RUNNING)
+ super().__init__(PipelineState.RUNNING)
self._executor = executor
self._evaluation_context = evaluation_context
diff --git a/sdks/python/apache_beam/runners/direct/direct_userstate.py b/sdks/python/apache_beam/runners/direct/direct_userstate.py
index 715355c..196a9a0 100644
--- a/sdks/python/apache_beam/runners/direct/direct_userstate.py
+++ b/sdks/python/apache_beam/runners/direct/direct_userstate.py
@@ -62,8 +62,7 @@
class ReadModifyWriteRuntimeState(DirectRuntimeState,
userstate.ReadModifyWriteRuntimeState):
def __init__(self, state_spec, state_tag, current_value_accessor):
- super(ReadModifyWriteRuntimeState,
- self).__init__(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
@@ -96,8 +95,7 @@
class BagRuntimeState(DirectRuntimeState, userstate.BagRuntimeState):
def __init__(self, state_spec, state_tag, current_value_accessor):
- super(BagRuntimeState,
- self).__init__(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 = []
@@ -122,8 +120,7 @@
class SetRuntimeState(DirectRuntimeState, userstate.SetRuntimeState):
def __init__(self, state_spec, state_tag, current_value_accessor):
- super(SetRuntimeState,
- self).__init__(state_spec, state_tag, current_value_accessor)
+ super().__init__(state_spec, state_tag, current_value_accessor)
self._current_accumulator = UNREAD_VALUE
self._modified = False
@@ -155,8 +152,7 @@
userstate.CombiningValueRuntimeState):
"""Combining value state interface object passed to user code."""
def __init__(self, state_spec, state_tag, current_value_accessor):
- super(CombiningValueRuntimeState,
- self).__init__(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)
diff --git a/sdks/python/apache_beam/runners/direct/evaluation_context.py b/sdks/python/apache_beam/runners/direct/evaluation_context.py
index 8d50d68..fbe59b0 100644
--- a/sdks/python/apache_beam/runners/direct/evaluation_context.py
+++ b/sdks/python/apache_beam/runners/direct/evaluation_context.py
@@ -458,7 +458,7 @@
class DirectUnmergedState(InMemoryUnmergedState):
"""UnmergedState implementation for the DirectRunner."""
def __init__(self):
- super(DirectUnmergedState, self).__init__(defensive_copy=False)
+ super().__init__(defensive_copy=False)
class DirectStepContext(object):
diff --git a/sdks/python/apache_beam/runners/direct/executor.py b/sdks/python/apache_beam/runners/direct/executor.py
index 8b47b0b..0ab3033 100644
--- a/sdks/python/apache_beam/runners/direct/executor.py
+++ b/sdks/python/apache_beam/runners/direct/executor.py
@@ -67,7 +67,7 @@
self,
queue, # type: queue.Queue[_ExecutorService.CallableTask]
index):
- super(_ExecutorService._ExecutorServiceWorker, self).__init__()
+ super().__init__()
self.queue = queue
self._index = index
self._default_name = 'ExecutorServiceWorker-' + str(index)
@@ -188,14 +188,14 @@
_GroupByKeyOnly.
"""
def __init__(self, executor_service, scheduled):
- super(_SerialEvaluationState, self).__init__(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(_SerialEvaluationState, self).complete(completed_work)
+ super().complete(completed_work)
def schedule(self, new_work):
self._update_currently_evaluating(new_work, None)
@@ -210,7 +210,7 @@
if self.serial_queue and not self.currently_evaluating:
next_work = self.serial_queue.pop()
self.currently_evaluating = next_work
- super(_SerialEvaluationState, self).schedule(next_work)
+ super().schedule(next_work)
class _TransformExecutorServices(object):
diff --git a/sdks/python/apache_beam/runners/direct/sdf_direct_runner_test.py b/sdks/python/apache_beam/runners/direct/sdf_direct_runner_test.py
index f60d50c..246d180 100644
--- a/sdks/python/apache_beam/runners/direct/sdf_direct_runner_test.py
+++ b/sdks/python/apache_beam/runners/direct/sdf_direct_runner_test.py
@@ -146,7 +146,7 @@
class SDFDirectRunnerTest(unittest.TestCase):
def setUp(self):
- super(SDFDirectRunnerTest, self).setUp()
+ super().setUp()
# Importing following for DirectRunner SDF implemenation for testing.
from apache_beam.runners.direct import transform_evaluator
self._old_default_max_num_outputs = (
diff --git a/sdks/python/apache_beam/runners/direct/test_direct_runner.py b/sdks/python/apache_beam/runners/direct/test_direct_runner.py
index 507ec1a..084820e 100644
--- a/sdks/python/apache_beam/runners/direct/test_direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/test_direct_runner.py
@@ -39,7 +39,7 @@
# send this option to remote executors.
test_options.on_success_matcher = None
- self.result = super(TestDirectRunner, self).run_pipeline(pipeline, options)
+ self.result = super().run_pipeline(pipeline, options)
try:
if not is_streaming:
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
index c77e442..2c4e1d9 100644
--- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
@@ -31,8 +31,8 @@
from typing import Tuple
from typing import Type
-import apache_beam.io as io
from apache_beam import coders
+from apache_beam import io
from apache_beam import pvalue
from apache_beam.internal import pickler
from apache_beam.runners import common
@@ -345,7 +345,7 @@
assert not side_inputs
self._source = applied_ptransform.transform.source
self._source.pipeline_options = evaluation_context.pipeline_options
- super(_BoundedReadEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -394,7 +394,7 @@
side_inputs):
assert not side_inputs
self.transform = applied_ptransform.transform
- super(_WatermarkControllerEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -464,7 +464,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_PairWithTimingEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -512,7 +512,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_TestStreamEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -602,7 +602,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_PubSubReadEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -736,7 +736,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_FlattenEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -770,7 +770,7 @@
def __init__(self, evaluation_context):
self._evaluation_context = evaluation_context
self._null_receiver = None
- super(_TaggedReceivers, self).__init__()
+ super().__init__()
class NullReceiver(common.Receiver):
"""Ignores undeclared outputs, default execution mode."""
@@ -804,7 +804,7 @@
side_inputs,
perform_dofn_pickle_test=True
):
- super(_ParDoEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -904,7 +904,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_GroupByKeyOnlyEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -1006,7 +1006,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_StreamingGroupByKeyOnlyEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -1061,7 +1061,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_StreamingGroupAlsoByWindowEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -1132,7 +1132,7 @@
input_committed_bundle,
side_inputs):
assert not side_inputs
- super(_NativeWriteEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -1207,7 +1207,7 @@
applied_ptransform,
input_committed_bundle,
side_inputs):
- super(_ProcessElementsEvaluator, self).__init__(
+ super().__init__(
evaluation_context,
applied_ptransform,
input_committed_bundle,
@@ -1273,6 +1273,6 @@
par_do_result.counters,
par_do_result.keyed_watermark_holds,
par_do_result.undeclared_tag_values)
- for key in self.keyed_holds:
- transform_result.keyed_watermark_holds[key] = self.keyed_holds[key]
+ for key, keyed_hold in self.keyed_holds.items():
+ transform_result.keyed_watermark_holds[key] = keyed_hold
return transform_result
diff --git a/sdks/python/apache_beam/runners/interactive/augmented_pipeline.py b/sdks/python/apache_beam/runners/interactive/augmented_pipeline.py
index 37f914b..1cfc5bc 100644
--- a/sdks/python/apache_beam/runners/interactive/augmented_pipeline.py
+++ b/sdks/python/apache_beam/runners/interactive/augmented_pipeline.py
@@ -52,8 +52,8 @@
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)
+ 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(
diff --git a/sdks/python/apache_beam/runners/interactive/background_caching_job.py b/sdks/python/apache_beam/runners/interactive/background_caching_job.py
index 5219538..bb94b3b 100644
--- a/sdks/python/apache_beam/runners/interactive/background_caching_job.py
+++ b/sdks/python/apache_beam/runners/interactive/background_caching_job.py
@@ -88,7 +88,7 @@
time.sleep(0.5)
def _should_end_condition_checker(self):
- return any([l.is_triggered() for l in self._limiters])
+ return any(l.is_triggered() for l in self._limiters)
def is_done(self):
with self._result_lock:
diff --git a/sdks/python/apache_beam/runners/interactive/cache_manager.py b/sdks/python/apache_beam/runners/interactive/cache_manager.py
index 9ed0b25..a697494 100644
--- a/sdks/python/apache_beam/runners/interactive/cache_manager.py
+++ b/sdks/python/apache_beam/runners/interactive/cache_manager.py
@@ -345,7 +345,7 @@
class SafeFastPrimitivesCoder(coders.Coder):
"""This class add an quote/unquote step to escape special characters."""
- # pylint: disable=deprecated-urllib-function
+ # pylint: disable=bad-option-value
def encode(self, value):
return quote(
diff --git a/sdks/python/apache_beam/runners/interactive/display/interactive_pipeline_graph.py b/sdks/python/apache_beam/runners/interactive/display/interactive_pipeline_graph.py
index 48c926f..5a0943e 100644
--- a/sdks/python/apache_beam/runners/interactive/display/interactive_pipeline_graph.py
+++ b/sdks/python/apache_beam/runners/interactive/display/interactive_pipeline_graph.py
@@ -71,7 +71,7 @@
self._referenced_pcollections = referenced_pcollections or set()
self._cached_pcollections = cached_pcollections or set()
- super(InteractivePipelineGraph, self).__init__(
+ super().__init__(
pipeline=pipeline,
default_vertex_attrs={
'color': 'gray', 'fontcolor': 'gray'
diff --git a/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization_test.py b/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization_test.py
index 5b0b51d..d34b966 100644
--- a/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization_test.py
+++ b/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization_test.py
@@ -59,7 +59,7 @@
ib.options.display_timezone = pytz.timezone('US/Pacific')
self._p = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
self._pcoll = self._p | 'Create' >> beam.Create(range(5))
ib.watch(self)
diff --git a/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_test.py b/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_test.py
index 9a0c7a9..419cd50 100644
--- a/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_test.py
+++ b/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_test.py
@@ -28,7 +28,7 @@
from apache_beam.runners.interactive.display import pipeline_graph
from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython
-# pylint: disable=range-builtin-not-iterating,unused-variable,possibly-unused-variable
+# 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
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py b/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
index da948fd..feb9092 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
@@ -98,7 +98,7 @@
@unittest.skipIf(sys.platform == "win32", "[BEAM-10627]")
def test_show_always_watch_given_pcolls(self):
p = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # 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())
@@ -111,7 +111,7 @@
@unittest.skipIf(sys.platform == "win32", "[BEAM-10627]")
def test_show_mark_pcolls_computed_when_done(self):
p = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # 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.
@@ -125,7 +125,7 @@
'visualize_computed_pcoll'))
def test_show_handles_dict_of_pcolls(self, mocked_visualize):
p = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
pcoll = p | 'Create' >> beam.Create(range(10))
ib.watch(locals())
ie.current_env().track_user_pipelines()
@@ -140,7 +140,7 @@
'visualize_computed_pcoll'))
def test_show_handles_iterable_of_pcolls(self, mocked_visualize):
p = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
pcoll = p | 'Create' >> beam.Create(range(10))
ib.watch(locals())
ie.current_env().track_user_pipelines()
@@ -172,7 +172,7 @@
self._pcoll = pcoll
p = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # 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
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_environment.py b/sdks/python/apache_beam/runners/interactive/interactive_environment.py
index fe10ab1..9f3d66e 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_environment.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_environment.py
@@ -36,6 +36,7 @@
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.recording_manager import RecordingManager
+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 register_ipython_log_handler
from apache_beam.utils.interactive_utils import is_in_ipython
@@ -206,6 +207,8 @@
self._inspector_with_synthetic = InteractiveEnvironmentInspector(
ignore_synthetic=False)
+ self.sql_chain = {}
+
@property
def options(self):
"""A reference to the global interactive options.
@@ -651,3 +654,17 @@
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
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_environment_test.py b/sdks/python/apache_beam/runners/interactive/interactive_environment_test.py
index f08db01..15a0b5c 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_environment_test.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_environment_test.py
@@ -27,6 +27,7 @@
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
# The module name is also a variable in module.
_module_name = 'apache_beam.runners.interactive.interactive_environment_test'
@@ -45,8 +46,8 @@
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()])
+ 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
@@ -303,6 +304,37 @@
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)
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_runner.py b/sdks/python/apache_beam/runners/interactive/interactive_runner.py
index 4778737..e19b85b 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_runner.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_runner.py
@@ -203,7 +203,7 @@
main_job_result.wait_until_finish()
if main_job_result.state is beam.runners.runner.PipelineState.DONE:
- # pylint: disable=dict-values-not-iterating
+ # pylint: disable=bad-option-value
ie.current_env().mark_pcollection_computed(
pipeline_instrument.cached_pcolls)
@@ -222,7 +222,7 @@
the pipeline being executed with interactivity applied and related
metadata including where the interactivity-backing cache lies.
"""
- super(PipelineResult, self).__init__(underlying_result.state)
+ super().__init__(underlying_result.state)
self._underlying_result = underlying_result
self._pipeline_instrument = pipeline_instrument
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_runner_test.py b/sdks/python/apache_beam/runners/interactive/interactive_runner_test.py
index 69551de..ea2b18f 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_runner_test.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_runner_test.py
@@ -267,7 +267,7 @@
ib.watch({'p': p})
with cell: # Cell 2
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
init = p | 'Init' >> beam.Create(range(5))
with cell: # Cell 3
diff --git a/sdks/python/apache_beam/runners/interactive/messaging/interactive_environment_inspector_test.py b/sdks/python/apache_beam/runners/interactive/messaging/interactive_environment_inspector_test.py
index 4edb0d4..2eb1004 100644
--- a/sdks/python/apache_beam/runners/interactive/messaging/interactive_environment_inspector_test.py
+++ b/sdks/python/apache_beam/runners/interactive/messaging/interactive_environment_inspector_test.py
@@ -47,13 +47,13 @@
pipeline = beam.Pipeline(ir.InteractiveRunner())
# Early watch the pipeline so that cell re-execution can be handled.
ib.watch({'pipeline': pipeline})
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
pcoll = pipeline | 'Create' >> beam.Create(range(10))
with cell: # Cell 2
# Re-executes the line that created the pcoll causing the original
# pcoll no longer inspectable.
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
pcoll = pipeline | 'Create' >> beam.Create(range(10))
ib.watch(locals())
@@ -106,7 +106,7 @@
def test_list_inspectables(self, cell):
with cell: # Cell 1
pipeline = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
pcoll_1 = pipeline | 'Create' >> beam.Create(range(10))
pcoll_2 = pcoll_1 | 'Square' >> beam.Map(lambda x: x * x)
@@ -144,7 +144,7 @@
def test_get_val(self, cell):
with cell: # Cell 1
pipeline = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
pcoll = pipeline | 'Create' >> beam.Create(range(10))
with cell: # Cell 2
@@ -167,7 +167,7 @@
def test_get_pcoll_data(self):
pipeline = beam.Pipeline(ir.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
pcoll = pipeline | 'Create' >> beam.Create(list(range(10)))
counts = pcoll | beam.combiners.Count.PerElement()
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_fragment.py b/sdks/python/apache_beam/runners/interactive/pipeline_fragment.py
index 7564a76..91c9b51 100644
--- a/sdks/python/apache_beam/runners/interactive/pipeline_fragment.py
+++ b/sdks/python/apache_beam/runners/interactive/pipeline_fragment.py
@@ -190,9 +190,21 @@
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=map-builtin-not-iterating
+ # 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)
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_fragment_test.py b/sdks/python/apache_beam/runners/interactive/pipeline_fragment_test.py
index 6e9d327..f1f423f 100644
--- a/sdks/python/apache_beam/runners/interactive/pipeline_fragment_test.py
+++ b/sdks/python/apache_beam/runners/interactive/pipeline_fragment_test.py
@@ -49,7 +49,7 @@
ib.watch(locals())
with cell: # Cell 2
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
init = p | 'Init' >> beam.Create(range(10))
init_expected = p_expected | 'Init' >> beam.Create(range(10))
@@ -71,7 +71,7 @@
ib.watch({'p': p})
with cell: # Cell 2
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
init = p | 'Init' >> beam.Create(range(10))
with cell: # Cell 3
@@ -100,7 +100,7 @@
ib.watch({'p': p})
with cell: # Cell 2
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
init = p | 'Init' >> beam.Create(range(5))
with cell: # Cell 3
@@ -129,6 +129,41 @@
# 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: n)
+
+ @beam.ptransform_fn
+ def Foo(pcoll):
+ p1 = pcoll | beam.Map(lambda n: n)
+ p2 = pcoll | beam.Map(str)
+ bar = pcoll | Bar()
+ return {'pc1': p1, 'pc2': p2, 'bar': bar}
+
+ res = init | Foo()
+
+ ib.watch(locals())
+ pc = res['pc1']
+
+ result = pf.PipelineFragment([pc]).run()
+ self.assertEqual([0, 1, 2, 3, 4], list(result.get(pc)))
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py b/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py
index bba315d..893603d 100644
--- a/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py
+++ b/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py
@@ -51,7 +51,7 @@
def test_pcoll_to_pcoll_id(self):
p = beam.Pipeline(interactive_runner.InteractiveRunner())
ie.current_env().set_cache_manager(InMemoryCache(), p)
- # pylint: disable=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
init_pcoll = p | 'Init Create' >> beam.Impulse()
_, ctx = p.to_runner_api(return_context=True)
self.assertEqual(
@@ -81,7 +81,7 @@
# in the original instance and if the evaluation has changed since last
# execution.
p2_id_runner = beam.Pipeline(interactive_runner.InteractiveRunner())
- # pylint: disable=range-builtin-not-iterating
+ # 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)
@@ -94,7 +94,7 @@
def test_cache_key(self):
p = beam.Pipeline(interactive_runner.InteractiveRunner())
ie.current_env().set_cache_manager(InMemoryCache(), p)
- # pylint: disable=range-builtin-not-iterating
+ # 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)
@@ -114,7 +114,7 @@
def test_cacheables(self):
p_cacheables = beam.Pipeline(interactive_runner.InteractiveRunner())
ie.current_env().set_cache_manager(InMemoryCache(), p_cacheables)
- # pylint: disable=range-builtin-not-iterating
+ # 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)
@@ -185,7 +185,7 @@
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=range-builtin-not-iterating
+ # pylint: disable=bad-option-value
if bounded:
source = beam.Create(range(10))
else:
@@ -263,8 +263,8 @@
def visit_transform(self, transform_node):
if transform_node.inputs:
main_inputs = dict(transform_node.main_inputs)
- for tag in main_inputs.keys():
- if main_inputs[tag] == init_pcoll:
+ 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
@@ -780,7 +780,7 @@
# 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=range-builtin-not-iterating,expression-not-assigned
+ # 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)
diff --git a/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics.py b/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics.py
index bd40f13..d27fc61 100644
--- a/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics.py
+++ b/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics.py
@@ -32,24 +32,27 @@
import apache_beam as beam
from apache_beam.pvalue import PValue
-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.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 is_namedtuple
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
from IPython.core.magic import Magics
from IPython.core.magic import line_cell_magic
from IPython.core.magic import magics_class
@@ -58,11 +61,11 @@
_EXAMPLE_USAGE = """beam_sql magic to execute Beam SQL in notebooks
---------------------------------------------------------
-%%beam_sql [-o OUTPUT_NAME] query
+%%beam_sql [-o OUTPUT_NAME] [-v] [-r RUNNER] query
---------------------------------------------------------
Or
---------------------------------------------------------
-%%%%beam_sql [-o OUTPUT_NAME] query-line#1
+%%%%beam_sql [-o OUTPUT_NAME] [-v] [-r RUNNER] query-line#1
query-line#2
...
query-line#N
@@ -82,6 +85,8 @@
to build Beam pipelines in a non-interactive manner.
"""
+_SUPPORTED_RUNNERS = ['DirectRunner', 'DataflowRunner']
+
class BeamSqlParser:
"""A parser to parse beam_sql inputs."""
@@ -100,6 +105,14 @@
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='*',
@@ -157,8 +170,9 @@
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, otherwise a PValue as if a
- SqlTransform is applied.
+ 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:
@@ -170,6 +184,7 @@
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):
@@ -181,11 +196,18 @@
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)
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 is_namedtuple(pcoll.element_type):
+ 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 '
@@ -194,45 +216,93 @@
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)
- output_name, output = apply_sql(query, output_name, found)
- cache_output(output_name, output)
- return output
+ 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:
+ _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]) -> Tuple[str, PValue]:
+ 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[str, PValue]. First str value is the output variable name in
- __main__ module (auto-generated if not provided). Second PValue is
- most likely a PCollection, depending on the query.
+ 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 = _build_query_components(query, found)
- 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.
- setattr(importlib.import_module('__main__'), output_name, output)
- ib.watch({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
- except (KeyboardInterrupt, SystemExit):
- raise
- except Exception as e:
- on_error('Error when applying the Beam SQL: %s', e)
+ 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 Exception as e:
+ on_error('Error when applying the Beam SQL: %s', e)
+ else:
+ return output_name, chain.current, chain
def pcolls_from_streaming_cache(
@@ -304,19 +374,26 @@
def _build_query_components(
- query: str, found: Dict[str, beam.PCollection]
+ query: str,
+ found: Dict[str, beam.PCollection],
+ output_name: str,
+ run: bool = True
) -> Tuple[str,
- Union[Dict[str, beam.PCollection], beam.PCollection, beam.Pipeline]]:
+ 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 and a source to apply the
+ 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.
+ 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(
@@ -324,26 +401,38 @@
sql_pipeline = beam.Pipeline(options=user_pipeline._options)
ie.current_env().add_derived_pipeline(user_pipeline, sql_pipeline)
sql_source = {}
- if has_source_to_cache(user_pipeline):
- sql_source = pcolls_from_streaming_cache(
- user_pipeline, sql_pipeline, found)
+ 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:
- 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)
+ 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()))
- else:
+
+ 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)
- return query, 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
diff --git a/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics_test.py b/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics_test.py
index 538abbb..3d843a0 100644
--- a/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics_test.py
+++ b/sdks/python/apache_beam/runners/interactive/sql/beam_sql_magics_test.py
@@ -59,9 +59,13 @@
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 = _build_query_components(query, {})
+ 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()
@@ -76,10 +80,14 @@
cache_key,
cache_manager,
element_type: target):
- processed_query, sql_source = _build_query_components(query, found)
-
- self.assertEqual(processed_query, 'SELECT * FROM PCOLLECTION where a=1')
+ 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()
@@ -95,12 +103,17 @@
cache_key,
cache_manager,
element_type: pcoll_1):
- processed_query, sql_source = _build_query_components(query, found)
+ 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()
@@ -115,8 +128,11 @@
lambda a,
b,
c: found):
- _, sql_source = _build_query_components(query, 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()
diff --git a/sdks/python/apache_beam/runners/interactive/sql/sql_chain.py b/sdks/python/apache_beam/runners/interactive/sql/sql_chain.py
new file mode 100644
index 0000000..a6f4866
--- /dev/null
+++ b/sdks/python/apache_beam/runners/interactive/sql/sql_chain.py
@@ -0,0 +1,226 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT 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 Dict
+from typing import Optional
+from typing import Set
+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)
diff --git a/sdks/python/apache_beam/runners/interactive/sql/sql_chain_test.py b/sdks/python/apache_beam/runners/interactive/sql/sql_chain_test.py
new file mode 100644
index 0000000..42d0804
--- /dev/null
+++ b/sdks/python/apache_beam/runners/interactive/sql/sql_chain_test.py
@@ -0,0 +1,109 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT 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()
diff --git a/sdks/python/apache_beam/runners/interactive/sql/utils.py b/sdks/python/apache_beam/runners/interactive/sql/utils.py
index fb4e57d..b2e75c8 100644
--- a/sdks/python/apache_beam/runners/interactive/sql/utils.py
+++ b/sdks/python/apache_beam/runners/interactive/sql/utils.py
@@ -23,29 +23,39 @@
# pytype: skip-file
import logging
+import os
+import tempfile
+from dataclasses import dataclass
+from typing import Any
+from typing import Callable
from typing import Dict
from typing import NamedTuple
+from typing import Optional
+from typing import Type
+from typing import Union
import apache_beam as beam
-from apache_beam.runners.interactive import interactive_beam as ib
+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 is_namedtuple(cls: type) -> bool:
- """Determines if a class is built from typing.NamedTuple."""
- return (
- isinstance(cls, type) and issubclass(cls, tuple) and
- hasattr(cls, '_fields') and hasattr(cls, '__annotations__'))
-
-
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 is_namedtuple(schema), (
+ 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):
@@ -77,21 +87,6 @@
if verbose:
_LOGGER.info('Found PCollections used in the magic: %s.', found)
_LOGGER.info('Collecting data...')
- for name, pcoll in found.items():
- try:
- _ = ib.collect(pcoll)
- except (KeyboardInterrupt, SystemExit):
- raise
- 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,
- sql)
- raise
return found
@@ -123,3 +118,314 @@
'{}: {}'.format(k, v.__name__) 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):
+ self.options = PipelineOptions()
+ 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 <a href="%s" target="_new">here</a> 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)
diff --git a/sdks/python/apache_beam/runners/interactive/sql/utils_test.py b/sdks/python/apache_beam/runners/interactive/sql/utils_test.py
index 01a54c3..16d03f5 100644
--- a/sdks/python/apache_beam/runners/interactive/sql/utils_test.py
+++ b/sdks/python/apache_beam/runners/interactive/sql/utils_test.py
@@ -23,9 +23,15 @@
from typing import NamedTuple
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 is_namedtuple
+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
@@ -37,19 +43,6 @@
class UtilsTest(unittest.TestCase):
- def test_is_namedtuple(self):
- class AType:
- pass
-
- a_type = AType
- a_tuple = type((1, 2, 3))
-
- a_namedtuple = ANamedTuple
-
- self.assertTrue(is_namedtuple(a_namedtuple))
- self.assertFalse(is_namedtuple(a_type))
- self.assertFalse(is_namedtuple(a_tuple))
-
def test_register_coder_for_schema(self):
self.assertNotIsInstance(
beam.coders.registry.get_coder(ANamedTuple), beam.coders.RowCoder)
@@ -80,6 +73,35 @@
self.assertEqual(
'ANamedTuple(a: int, b: str)', pformat_namedtuple(ANamedTuple))
+ def test_pformat_dict(self):
+ self.assertEqual('{\na: 1,\nb: 2\n}', pformat_dict({'a': 1, 'b': '2'}))
+
+
+@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()
diff --git a/sdks/python/apache_beam/runners/interactive/testing/integration/notebook_executor.py b/sdks/python/apache_beam/runners/interactive/testing/integration/notebook_executor.py
index ecff3bd..6a80639 100644
--- a/sdks/python/apache_beam/runners/interactive/testing/integration/notebook_executor.py
+++ b/sdks/python/apache_beam/runners/interactive/testing/integration/notebook_executor.py
@@ -151,7 +151,7 @@
"""A parser to extract iframe content from given HTML."""
def __init__(self):
self._srcdocs = []
- super(IFrameParser, self).__init__()
+ super().__init__()
def handle_starttag(self, tag, attrs):
if tag == 'iframe':
diff --git a/sdks/python/apache_beam/runners/interactive/testing/integration/screen_diff.py b/sdks/python/apache_beam/runners/interactive/testing/integration/screen_diff.py
index 5e0c855..a1c9971 100644
--- a/sdks/python/apache_beam/runners/interactive/testing/integration/screen_diff.py
+++ b/sdks/python/apache_beam/runners/interactive/testing/integration/screen_diff.py
@@ -175,7 +175,7 @@
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(BaseTestCase, self).__init__(*args, **kwargs)
+ super().__init__(*args, **kwargs)
@classmethod
def get_web_driver(cls):
@@ -195,7 +195,7 @@
self._golden_dir,
self._cleanup) as test_env:
self._test_env = test_env
- super(BaseTestCase, self).run(result)
+ super().run(result)
def explicit_wait(self):
"""Wait for common elements to be visible."""
diff --git a/sdks/python/apache_beam/runners/interactive/testing/integration/tests/screen_diff_test.py b/sdks/python/apache_beam/runners/interactive/testing/integration/tests/screen_diff_test.py
index 0d36c88..a3f8ace 100644
--- a/sdks/python/apache_beam/runners/interactive/testing/integration/tests/screen_diff_test.py
+++ b/sdks/python/apache_beam/runners/interactive/testing/integration/tests/screen_diff_test.py
@@ -29,7 +29,7 @@
class DataFramesTest(BaseTestCase):
def __init__(self, *args, **kwargs):
kwargs['golden_size'] = (1024, 10000)
- super(DataFramesTest, self).__init__(*args, **kwargs)
+ super().__init__(*args, **kwargs)
def explicit_wait(self):
try:
@@ -39,6 +39,7 @@
WebDriverWait(self.driver, 5).until(
expected_conditions.presence_of_element_located((By.ID, 'test-done')))
+ # pylint: disable=bare-except
except:
pass # The test will be ignored.
@@ -50,7 +51,7 @@
class InitSquareCubeTest(BaseTestCase):
def __init__(self, *args, **kwargs):
kwargs['golden_size'] = (1024, 10000)
- super(InitSquareCubeTest, self).__init__(*args, **kwargs)
+ super().__init__(*args, **kwargs)
def test_init_square_cube_notebook(self):
self.assert_notebook('init_square_cube')
diff --git a/sdks/python/apache_beam/runners/interactive/testing/pipeline_assertion.py b/sdks/python/apache_beam/runners/interactive/testing/pipeline_assertion.py
index 9f4bdbc..9b07342 100644
--- a/sdks/python/apache_beam/runners/interactive/testing/pipeline_assertion.py
+++ b/sdks/python/apache_beam/runners/interactive/testing/pipeline_assertion.py
@@ -88,10 +88,9 @@
pipeline_proto.root_transform_ids[0]].subtransforms
test_case.assertEqual(
contain,
- any([
+ any(
transform_label in top_level_transform_label
- for top_level_transform_label in top_level_transform_labels
- ]))
+ for top_level_transform_label in top_level_transform_labels))
def _assert_transform_equal(
diff --git a/sdks/python/apache_beam/runners/interactive/utils.py b/sdks/python/apache_beam/runners/interactive/utils.py
index 49b87ba..7a75e52 100644
--- a/sdks/python/apache_beam/runners/interactive/utils.py
+++ b/sdks/python/apache_beam/runners/interactive/utils.py
@@ -20,9 +20,12 @@
import functools
import hashlib
+import importlib
import json
import logging
+from typing import Any
from typing import Dict
+from typing import Tuple
import pandas as pd
@@ -150,8 +153,8 @@
# 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]):
+ if any(isinstance(h, IPythonLogHandler)
+ for h in interactive_root_logger.handlers):
return
interactive_root_logger.setLevel(logging.INFO)
interactive_root_logger.addHandler(IPythonLogHandler())
@@ -405,3 +408,22 @@
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
diff --git a/sdks/python/apache_beam/runners/interactive/utils_test.py b/sdks/python/apache_beam/runners/interactive/utils_test.py
index 784081e..0915ff2 100644
--- a/sdks/python/apache_beam/runners/interactive/utils_test.py
+++ b/sdks/python/apache_beam/runners/interactive/utils_test.py
@@ -15,6 +15,7 @@
# limitations under the License.
#
+import importlib
import json
import logging
import tempfile
@@ -318,6 +319,13 @@
})
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)
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/runners/portability/abstract_job_service.py b/sdks/python/apache_beam/runners/portability/abstract_job_service.py
index 55224f4..a369af9 100644
--- a/sdks/python/apache_beam/runners/portability/abstract_job_service.py
+++ b/sdks/python/apache_beam/runners/portability/abstract_job_service.py
@@ -321,7 +321,7 @@
pipeline,
options,
artifact_port=0):
- super(UberJarBeamJob, self).__init__(job_id, job_name, pipeline, options)
+ super().__init__(job_id, job_name, pipeline, options)
self._executable_jar = executable_jar
self._jar_uploaded = False
self._artifact_port = artifact_port
@@ -344,7 +344,10 @@
for (env_id,
env) in self._pipeline_proto.components.environments.items()
})
- self._artifact_staging_server = grpc.server(futures.ThreadPoolExecutor())
+ 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(
diff --git a/sdks/python/apache_beam/runners/portability/artifact_service.py b/sdks/python/apache_beam/runners/portability/artifact_service.py
index 64023eb..bec9317 100644
--- a/sdks/python/apache_beam/runners/portability/artifact_service.py
+++ b/sdks/python/apache_beam/runners/portability/artifact_service.py
@@ -290,6 +290,7 @@
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:
diff --git a/sdks/python/apache_beam/runners/portability/flink_runner.py b/sdks/python/apache_beam/runners/portability/flink_runner.py
index 6486d3d..efa17cd 100644
--- a/sdks/python/apache_beam/runners/portability/flink_runner.py
+++ b/sdks/python/apache_beam/runners/portability/flink_runner.py
@@ -42,7 +42,7 @@
not portable_options.environment_type and
not portable_options.output_executable_path):
portable_options.environment_type = 'LOOPBACK'
- return super(FlinkRunner, self).run_pipeline(pipeline, options)
+ return super().run_pipeline(pipeline, options)
def default_job_server(self, options):
flink_options = options.view_as(pipeline_options.FlinkRunnerOptions)
@@ -82,7 +82,7 @@
class FlinkJarJobServer(job_server.JavaJarJobServer):
def __init__(self, options):
- super(FlinkJarJobServer, self).__init__(options)
+ super().__init__(options)
options = options.view_as(pipeline_options.FlinkRunnerOptions)
self._jar = options.flink_job_server_jar
self._master_url = options.flink_master
diff --git a/sdks/python/apache_beam/runners/portability/flink_runner_test.py b/sdks/python/apache_beam/runners/portability/flink_runner_test.py
index 7d26a70..cb6345e 100644
--- a/sdks/python/apache_beam/runners/portability/flink_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/flink_runner_test.py
@@ -70,7 +70,7 @@
flink_job_server_jar = None
def __init__(self, *args, **kwargs):
- super(FlinkRunnerTest, self).__init__(*args, **kwargs)
+ super().__init__(*args, **kwargs)
self.environment_type = None
self.environment_config = None
@@ -123,7 +123,7 @@
if cls.conf_dir and exists(cls.conf_dir):
_LOGGER.info("removing conf dir: %s" % cls.conf_dir)
rmtree(cls.conf_dir)
- super(FlinkRunnerTest, cls).tearDownClass()
+ super().tearDownClass()
@classmethod
def _create_conf_dir(cls):
@@ -195,7 +195,7 @@
cls.flink_job_server_jar = flink_job_server_jar
def create_options(self):
- options = super(FlinkRunnerTest, self).create_options()
+ 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
@@ -291,11 +291,10 @@
def test_flattened_side_input(self):
# Blocked on support for transcoding
# https://jira.apache.org/jira/browse/BEAM-6523
- super(FlinkRunnerTest,
- self).test_flattened_side_input(with_transcoding=False)
+ super().test_flattened_side_input(with_transcoding=False)
def test_metrics(self):
- super(FlinkRunnerTest, self).test_metrics(check_gauge=False)
+ super().test_metrics(check_gauge=False)
def test_flink_metrics(self):
"""Run a simple DoFn that increments a counter and verifies state
@@ -405,7 +404,7 @@
# TODO: Remove these tests after resolving BEAM-7248 and enabling
# PortableRunnerOptimized
def create_options(self):
- options = super(FlinkRunnerTestOptimized, self).create_options()
+ options = super().create_options()
options.view_as(DebugOptions).experiments = [
'pre_optimize=all'
] + options.view_as(DebugOptions).experiments
@@ -432,14 +431,14 @@
class FlinkRunnerTestStreaming(FlinkRunnerTest):
def __init__(self, *args, **kwargs):
- super(FlinkRunnerTestStreaming, self).__init__(*args, **kwargs)
+ super().__init__(*args, **kwargs)
self.enable_commit = False
def setUp(self):
self.enable_commit = False
def create_options(self):
- options = super(FlinkRunnerTestStreaming, self).create_options()
+ options = super().create_options()
options.view_as(StandardOptions).streaming = True
if self.enable_commit:
options._all_options['checkpointing_interval'] = 3000
@@ -448,11 +447,11 @@
def test_callbacks_with_exception(self):
self.enable_commit = True
- super(FlinkRunnerTest, self).test_callbacks_with_exception()
+ super().test_callbacks_with_exception()
def test_register_finalizations(self):
self.enable_commit = True
- super(FlinkRunnerTest, self).test_register_finalizations()
+ super().test_register_finalizations()
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py
index 9b0d6ff..9a40a55 100644
--- a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py
+++ b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py
@@ -43,7 +43,7 @@
the pipeline artifacts.
"""
def __init__(self, master_url, options):
- super(FlinkUberJarJobServer, self).__init__()
+ super().__init__()
self._master_url = master_url
self._executable_jar = (
options.view_as(
@@ -116,7 +116,7 @@
pipeline,
options,
artifact_port=0):
- super(FlinkBeamJob, self).__init__(
+ super().__init__(
executable_jar,
job_id,
job_name,
@@ -208,7 +208,7 @@
state, timestamp = self._get_state()
if timestamp is None:
# state has not changed since it was last checked: use previous timestamp
- return super(FlinkBeamJob, self).get_state()
+ return super().get_state()
else:
return state, timestamp
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py
index be8fe60..4b77772 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py
@@ -115,7 +115,7 @@
waits before requesting progress from the SDK.
is_drain: identify whether expand the sdf graph in the drain mode.
"""
- super(FnApiRunner, self).__init__()
+ super().__init__()
self._default_environment = (
default_environment or environments.EmbeddedPythonEnvironment.default())
self._bundle_repeat = bundle_repeat
@@ -1140,7 +1140,7 @@
cache_token_generator=None,
**kwargs):
# type: (...) -> None
- super(ParallelBundleManager, self).__init__(
+ super().__init__(
bundle_context_manager,
progress_frequency,
cache_token_generator=cache_token_generator)
@@ -1184,7 +1184,7 @@
dry_run)
with thread_pool_executor.shared_unbounded_instance() as executor:
- for result, split_result in executor.map(execute, zip(part_inputs, # pylint: disable=zip-builtin-not-iterating
+ 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:
@@ -1214,7 +1214,7 @@
callback=None
):
# type: (...) -> None
- super(ProgressRequester, self).__init__()
+ super().__init__()
self._worker_handler = worker_handler
self._instruction_id = instruction_id
self._frequency = frequency
@@ -1299,7 +1299,7 @@
class RunnerResult(runner.PipelineResult):
def __init__(self, state, monitoring_infos_by_stage):
- super(RunnerResult, self).__init__(state)
+ super().__init__(state)
self._monitoring_infos_by_stage = monitoring_infos_by_stage
self._metrics = None
self._monitoring_metrics = None
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
index c002963..87d03ca 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
@@ -1086,10 +1086,10 @@
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]))
+ 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]))
+ any(re.match(packed_step_name_regex, s) for s in step_names))
@retry(stop=stop_after_attempt(3))
def test_pack_combiners(self):
@@ -1970,8 +1970,7 @@
class CheckpointOnlyOffsetRestrictionTracker(
restriction_trackers.OffsetRestrictionTracker):
def try_split(self, unused_fraction_of_remainder):
- return super(CheckpointOnlyOffsetRestrictionTracker,
- self).try_split(0.0)
+ return super().try_split(0.0)
return CheckpointOnlyOffsetRestrictionTracker(restriction)
if self.use_bounded_offset_range:
@@ -1987,7 +1986,7 @@
class OffsetRangeProviderWithTruncate(OffsetRangeProvider):
def __init__(self):
- super(OffsetRangeProviderWithTruncate, self).__init__(True)
+ super().__init__(True)
def truncate(self, element, restriction):
return restriction_trackers.OffsetRange(
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py
index dc536d2..03b4f7e 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/translations.py
@@ -24,7 +24,9 @@
import functools
import itertools
import logging
+import operator
from typing import Callable
+from typing import Collection
from typing import Container
from typing import DefaultDict
from typing import Dict
@@ -36,6 +38,7 @@
from typing import Set
from typing import Tuple
from typing import TypeVar
+from typing import Union
from apache_beam import coders
from apache_beam.internal import pickler
@@ -357,12 +360,29 @@
class TransformContext(object):
- _KNOWN_CODER_URNS = set(
+ _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
@@ -373,6 +393,14 @@
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 all 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
@@ -456,7 +484,7 @@
coder = self.components.coders[coder_id]
if coder.spec.urn == common_urns.coders.LENGTH_PREFIX.urn:
return coder_id, self.bytes_coder_id
- elif coder.spec.urn in self._KNOWN_CODER_URNS:
+ elif coder.spec.urn in self._known_coder_urns:
new_component_ids = [
self.maybe_length_prefixed_coder(c) for c in coder.component_coder_ids
]
@@ -765,6 +793,27 @@
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:
@@ -803,7 +852,7 @@
def _eliminate_common_key_with_none(stages, context, can_pack=lambda s: True):
- # type: (Iterable[Stage], TransformContext, Callable[[str], bool]) -> Iterable[Stage]
+ # type: (Iterable[Stage], TransformContext, Callable[[str], Union[bool, int]]) -> Iterable[Stage]
"""Runs common subexpression elimination for sibling KeyWithNone stages.
@@ -866,8 +915,11 @@
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], bool]) -> Iterator[Stage]
+ # type: (Iterable[Stage], TransformContext, Callable[[str], Union[bool, int]]) -> Iterator[Stage]
"""Packs sibling CombinePerKey stages into a single CombinePerKey.
@@ -919,6 +971,13 @@
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):
@@ -939,7 +998,12 @@
for stage in ineligible_stages:
yield stage
- for stage_key, packable_stages in grouped_eligible_stages.items():
+ 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:
@@ -1063,18 +1127,22 @@
def pack_combiners(stages, context, can_pack=None):
- # type: (Iterable[Stage], TransformContext, Optional[Callable[[str], bool]]) -> Iterator[Stage]
+ # type: (Iterable[Stage], TransformContext, Optional[Callable[[str], Union[bool, int]]]) -> Iterator[Stage]
if can_pack is None:
- can_pack_names = {} # type: Dict[str, bool]
+ can_pack_names = {} # type: Dict[str, Union[bool, int]]
parents = context.parents_map()
- def can_pack_fn(name: str) -> bool:
+ 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:
- result = True
+ 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:
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/translations_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/translations_test.py
index 37882dc..144f067 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/translations_test.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/translations_test.py
@@ -307,6 +307,83 @@
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_conditionally_packed_combiners(self):
class RecursiveCombine(beam.PTransform):
def __init__(self, labels):
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/trigger_manager.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/trigger_manager.py
index 0fd74b3..021f595 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/trigger_manager.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/trigger_manager.py
@@ -96,7 +96,7 @@
class TriggerMergeContext(WindowFn.MergeContext):
def __init__(
self, all_windows, context: 'FnRunnerStatefulTriggerContext', windowing):
- super(TriggerMergeContext, self).__init__(all_windows)
+ super().__init__(all_windows)
self.trigger_context = context
self.windowing = windowing
self.merged_away: typing.Dict[BoundedWindow, BoundedWindow] = {}
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/worker_handlers.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/worker_handlers.py
index b967e01..e886473 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/worker_handlers.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/worker_handlers.py
@@ -106,7 +106,7 @@
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 = dict() # type: Dict[str, ControlFuture]
+ 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
@@ -354,7 +354,7 @@
worker_manager, # type: WorkerHandlerManager
):
# type: (...) -> None
- super(EmbeddedWorkerHandler, self).__init__(
+ 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
@@ -457,26 +457,29 @@
worker_manager, # type: WorkerHandlerManager
):
# type: (...) -> None
- self.state = state
- self.provision_info = provision_info
- self.control_server = grpc.server(
- thread_pool_executor.shared_unbounded_instance())
- self.control_port = self.control_server.add_insecure_port('[::]:0')
- self.control_address = 'localhost:%s' % self.control_port
# 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.
- no_max_message_sizes = [("grpc.max_receive_message_length", -1),
- ("grpc.max_send_message_length", -1)]
+ # 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)]
+
+ 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=no_max_message_sizes)
+ 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=no_max_message_sizes)
+ thread_pool_executor.shared_unbounded_instance(), options=options)
self.state_port = self.state_server.add_insecure_port('[::]:0')
self.control_handler = BeamFnControlServicer(worker_manager)
@@ -510,8 +513,7 @@
GrpcStateServicer(state), self.state_server)
self.logging_server = grpc.server(
- thread_pool_executor.shared_unbounded_instance(),
- options=no_max_message_sizes)
+ 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)
@@ -548,7 +550,7 @@
):
# type: (...) -> None
self._grpc_server = grpc_server
- super(GrpcWorkerHandler, self).__init__(
+ super().__init__(
self._grpc_server.control_handler,
self._grpc_server.data_plane_handler,
state,
@@ -591,7 +593,7 @@
# type: () -> None
self.control_conn.close()
self.data_conn.close()
- super(GrpcWorkerHandler, self).close()
+ super().close()
def port_from_worker(self, port):
# type: (int) -> str
@@ -612,8 +614,7 @@
grpc_server # type: GrpcServer
):
# type: (...) -> None
- super(ExternalWorkerHandler,
- self).__init__(state, provision_info, grpc_server)
+ super().__init__(state, provision_info, grpc_server)
self._external_payload = external_payload
def start_worker(self):
@@ -657,8 +658,7 @@
grpc_server # type: GrpcServer
):
# type: (...) -> None
- super(EmbeddedGrpcWorkerHandler,
- self).__init__(state, provision_info, grpc_server)
+ super().__init__(state, provision_info, grpc_server)
from apache_beam.transforms.environments import EmbeddedPythonGrpcEnvironment
config = EmbeddedPythonGrpcEnvironment.parse_config(payload.decode('utf-8'))
@@ -697,8 +697,7 @@
grpc_server # type: GrpcServer
):
# type: (...) -> None
- super(SubprocessSdkWorkerHandler,
- self).__init__(state, provision_info, grpc_server)
+ super().__init__(state, provision_info, grpc_server)
self._worker_command_line = worker_command_line
def start_worker(self):
@@ -728,8 +727,7 @@
grpc_server # type: GrpcServer
):
# type: (...) -> None
- super(DockerSdkWorkerHandler,
- self).__init__(state, provision_info, grpc_server)
+ super().__init__(state, provision_info, grpc_server)
self._container_image = payload.container_image
self._container_id = None # type: Optional[bytes]
@@ -743,7 +741,7 @@
# 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(DockerSdkWorkerHandler, self).host_from_worker()
+ return super().host_from_worker()
def start_worker(self):
# type: () -> None
@@ -813,7 +811,8 @@
'SDK exited unexpectedly. '
'Final status is %s. Final log line is %s' % (
status.decode('utf-8'),
- logs.decode('utf-8').strip().split('\n')[-1])))
+ logs.decode('utf-8').strip().rsplit('\n',
+ maxsplit=1)[-1])))
time.sleep(5)
def stop_worker(self):
diff --git a/sdks/python/apache_beam/runners/portability/job_server.py b/sdks/python/apache_beam/runners/portability/job_server.py
index b1363c6..eda8755 100644
--- a/sdks/python/apache_beam/runners/portability/job_server.py
+++ b/sdks/python/apache_beam/runners/portability/job_server.py
@@ -121,7 +121,7 @@
class JavaJarJobServer(SubprocessJobServer):
def __init__(self, options):
- super(JavaJarJobServer, self).__init__()
+ super().__init__()
options = options.view_as(pipeline_options.JobServerOptions)
self._job_port = options.job_port
self._artifact_port = options.artifact_port
diff --git a/sdks/python/apache_beam/runners/portability/local_job_service.py b/sdks/python/apache_beam/runners/portability/local_job_service.py
index aedfc03..86049a9 100644
--- a/sdks/python/apache_beam/runners/portability/local_job_service.py
+++ b/sdks/python/apache_beam/runners/portability/local_job_service.py
@@ -74,7 +74,7 @@
subprocesses for the runner and worker(s).
"""
def __init__(self, staging_dir=None):
- super(LocalJobServicer, self).__init__()
+ super().__init__()
self._cleanup_staging_dir = staging_dir is None
self._staging_dir = staging_dir or tempfile.mkdtemp()
self._artifact_service = artifact_service.ArtifactStagingService(
@@ -125,11 +125,12 @@
return 'localhost'
def start_grpc_server(self, port=0):
- no_max_message_sizes = [("grpc.max_receive_message_length", -1),
- ("grpc.max_send_message_length", -1)]
+ 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=no_max_message_sizes)
+ 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)
@@ -183,8 +184,10 @@
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())
+ thread_pool_executor.shared_unbounded_instance(), options=options)
logging_port = logging_server.add_insecure_port('[::]:0')
logging_server.start()
logging_servicer = BeamFnLoggingServicer()
@@ -234,8 +237,7 @@
artifact_staging_endpoint, # type: Optional[endpoints_pb2.ApiServiceDescriptor]
artifact_service, # type: artifact_service.ArtifactStagingService
):
- super(BeamJob,
- self).__init__(job_id, provision_info.job_name, pipeline, options)
+ 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
@@ -246,7 +248,7 @@
def set_state(self, new_state):
"""Set the latest state as an int enum and notify consumers"""
- timestamp = super(BeamJob, self).set_state(new_state)
+ timestamp = super().set_state(new_state)
if timestamp is not None:
# Inform consumers of the new state.
for queue in self._state_queues:
@@ -389,7 +391,7 @@
}
def __init__(self, log_queues):
- super(JobLogHandler, self).__init__()
+ super().__init__()
self._last_id = 0
self._logged_thread = None
self._log_queues = log_queues
diff --git a/sdks/python/apache_beam/runners/portability/portable_runner.py b/sdks/python/apache_beam/runners/portability/portable_runner.py
index d040d2e..3d2ebcd 100644
--- a/sdks/python/apache_beam/runners/portability/portable_runner.py
+++ b/sdks/python/apache_beam/runners/portability/portable_runner.py
@@ -493,7 +493,7 @@
message_stream,
state_stream,
cleanup_callbacks=()):
- super(PipelineResult, self).__init__(beam_job_api_pb2.JobState.UNSPECIFIED)
+ super().__init__(beam_job_api_pb2.JobState.UNSPECIFIED)
self._job_service = job_service
self._job_id = job_id
self._messages = []
diff --git a/sdks/python/apache_beam/runners/portability/portable_runner_test.py b/sdks/python/apache_beam/runners/portability/portable_runner_test.py
index 5d4217f..b0404640 100644
--- a/sdks/python/apache_beam/runners/portability/portable_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/portable_runner_test.py
@@ -224,7 +224,7 @@
@unittest.skip("BEAM-7248")
class PortableRunnerOptimized(PortableRunnerTest):
def create_options(self):
- options = super(PortableRunnerOptimized, self).create_options()
+ 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(
@@ -236,7 +236,7 @@
# beam:runner:executable_stage:v1.
class PortableRunnerOptimizedWithoutFusion(PortableRunnerTest):
def create_options(self):
- options = super(PortableRunnerOptimizedWithoutFusion, self).create_options()
+ 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')
@@ -257,7 +257,7 @@
cls._worker_server.stop(1)
def create_options(self):
- options = super(PortableRunnerTestWithExternalEnv, self).create_options()
+ options = super().create_options()
options.view_as(PortableOptions).environment_type = 'EXTERNAL'
options.view_as(PortableOptions).environment_config = self._worker_address
return options
@@ -268,7 +268,7 @@
_use_subprocesses = True
def create_options(self):
- options = super(PortableRunnerTestWithSubprocesses, self).create_options()
+ options = super().create_options()
options.view_as(PortableOptions).environment_type = (
python_urns.SUBPROCESS_SDK)
options.view_as(PortableOptions).environment_config = (
@@ -297,7 +297,7 @@
_use_subprocesses = True
def create_options(self):
- options = super(PortableRunnerTestWithSubprocessesAndMultiWorkers, self) \
+ options = super() \
.create_options()
options.view_as(DirectOptions).direct_num_workers = 2
return options
@@ -396,7 +396,7 @@
"no docker image")
class PortableRunnerTestWithLocalDocker(PortableRunnerTest):
def create_options(self):
- options = super(PortableRunnerTestWithLocalDocker, self).create_options()
+ options = super().create_options()
options.view_as(PortableOptions).job_endpoint = 'embed'
return options
diff --git a/sdks/python/apache_beam/runners/portability/samza_runner_test.py b/sdks/python/apache_beam/runners/portability/samza_runner_test.py
index 2f60ad8..e946371 100644
--- a/sdks/python/apache_beam/runners/portability/samza_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/samza_runner_test.py
@@ -126,7 +126,7 @@
return 'localhost:%s' % cls.expansion_port
def create_options(self):
- options = super(SamzaRunnerTest, self).create_options()
+ options = super().create_options()
options.view_as(PortableOptions).environment_type = self.environment_type
options.view_as(
PortableOptions).environment_options = self.environment_options
@@ -140,8 +140,7 @@
def test_flattened_side_input(self):
# Blocked on support for transcoding
# https://issues.apache.org/jira/browse/BEAM-12681
- super(SamzaRunnerTest,
- self).test_flattened_side_input(with_transcoding=False)
+ super().test_flattened_side_input(with_transcoding=False)
def test_pack_combiners(self):
# Stages produced by translations.pack_combiners are fused
diff --git a/sdks/python/apache_beam/runners/portability/spark_runner.py b/sdks/python/apache_beam/runners/portability/spark_runner.py
index bb7b4c0..b1d754d 100644
--- a/sdks/python/apache_beam/runners/portability/spark_runner.py
+++ b/sdks/python/apache_beam/runners/portability/spark_runner.py
@@ -44,7 +44,7 @@
not portable_options.environment_type and
not portable_options.output_executable_path):
portable_options.environment_type = 'LOOPBACK'
- return super(SparkRunner, self).run_pipeline(pipeline, options)
+ return super().run_pipeline(pipeline, options)
def default_job_server(self, options):
spark_options = options.view_as(pipeline_options.SparkRunnerOptions)
@@ -73,7 +73,7 @@
class SparkJarJobServer(job_server.JavaJarJobServer):
def __init__(self, options):
- super(SparkJarJobServer, self).__init__(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
diff --git a/sdks/python/apache_beam/runners/portability/spark_runner_test.py b/sdks/python/apache_beam/runners/portability/spark_runner_test.py
index ba5f103..fb84c27 100644
--- a/sdks/python/apache_beam/runners/portability/spark_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/spark_runner_test.py
@@ -130,7 +130,7 @@
cls.spark_job_server_jar = spark_job_server_jar
def create_options(self):
- options = super(SparkRunnerTest, self).create_options()
+ options = super().create_options()
options.view_as(PortableOptions).environment_type = self.environment_type
options.view_as(
PortableOptions).environment_options = self.environment_options
@@ -175,8 +175,7 @@
def test_flattened_side_input(self):
# Blocked on support for transcoding
# https://jira.apache.org/jira/browse/BEAM-7236
- super(SparkRunnerTest,
- self).test_flattened_side_input(with_transcoding=False)
+ super().test_flattened_side_input(with_transcoding=False)
def test_custom_merging_window(self):
raise unittest.SkipTest("BEAM-11004")
diff --git a/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py b/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py
index 60b2e88..2f880d2 100644
--- a/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py
+++ b/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server.py
@@ -45,7 +45,7 @@
the pipeline artifacts.
"""
def __init__(self, rest_url, options):
- super(SparkUberJarJobServer, self).__init__()
+ super().__init__()
self._rest_url = rest_url
self._artifact_port = (
options.view_as(pipeline_options.JobServerOptions).artifact_port)
@@ -108,7 +108,7 @@
pipeline,
options,
artifact_port=0):
- super(SparkBeamJob, self).__init__(
+ super().__init__(
executable_jar,
job_id,
job_name,
@@ -218,7 +218,7 @@
timestamp = self.set_state(state)
if timestamp is None:
# State has not changed since last check. Use previous timestamp.
- return super(SparkBeamJob, self).get_state()
+ return super().get_state()
else:
return state, timestamp
diff --git a/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server_test.py b/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server_test.py
index d9bb7e6..6bb27b5 100644
--- a/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server_test.py
+++ b/sdks/python/apache_beam/runners/portability/spark_uber_jar_job_server_test.py
@@ -132,6 +132,7 @@
# 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,
diff --git a/sdks/python/apache_beam/runners/runner.py b/sdks/python/apache_beam/runners/runner.py
index 1030a53..b90eff5 100644
--- a/sdks/python/apache_beam/runners/runner.py
+++ b/sdks/python/apache_beam/runners/runner.py
@@ -41,7 +41,7 @@
__all__ = ['PipelineRunner', 'PipelineState', 'PipelineResult']
_RUNNER_MAP = {
- path.split('.')[-1].lower(): path
+ path.rsplit('.', maxsplit=1)[-1].lower(): path
for path in StandardOptions.ALL_KNOWN_RUNNERS
}
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py
index 497d613..532b7dd 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -129,8 +129,7 @@
data_channel # type: data_plane.DataChannel
):
# type: (...) -> None
- super(RunnerIOOperation,
- self).__init__(name_context, None, counter_factory, state_sampler)
+ 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
@@ -158,7 +157,7 @@
def finish(self):
# type: () -> None
self.output_stream.close()
- super(DataOutputOperation, self).finish()
+ super().finish()
class DataInputOperation(RunnerIOOperation):
@@ -175,7 +174,7 @@
data_channel # type: data_plane.GrpcClientDataChannel
):
# type: (...) -> None
- super(DataInputOperation, self).__init__(
+ super().__init__(
operation_name,
step_name,
consumers,
@@ -201,7 +200,7 @@
def start(self):
# type: () -> None
- super(DataInputOperation, self).start()
+ super().start()
with self.splitting_lock:
self.started = True
@@ -223,7 +222,7 @@
def monitoring_infos(self, transform_id, tag_to_pcollection_id):
# type: (str, Dict[str, str]) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo]
- all_monitoring_infos = super(DataInputOperation, self).monitoring_infos(
+ 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,
@@ -297,7 +296,7 @@
element_primaries, element_residuals = split
return index - 1, element_primaries, element_residuals, index + 1
# Otherwise, split at the closest element boundary.
- # pylint: disable=round-builtin
+ # 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.
@@ -330,7 +329,7 @@
with self.splitting_lock:
self.index = -1
self.stop = float('inf')
- super(DataInputOperation, self).reset()
+ super().reset()
class _StateBackedIterable(object):
diff --git a/sdks/python/apache_beam/runners/worker/channel_factory.py b/sdks/python/apache_beam/runners/worker/channel_factory.py
index cd859e6..6ad0f72 100644
--- a/sdks/python/apache_beam/runners/worker/channel_factory.py
+++ b/sdks/python/apache_beam/runners/worker/channel_factory.py
@@ -22,7 +22,10 @@
class GRPCChannelFactory(grpc.StreamStreamClientInterceptor):
- DEFAULT_OPTIONS = [("grpc.keepalive_time_ms", 20000)]
+ DEFAULT_OPTIONS = [
+ ("grpc.keepalive_time_ms", 20000),
+ ("grpc.keepalive_timeout_ms", 300000),
+ ]
def __init__(self):
pass
diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py
index d395cee..4baca68 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane.py
@@ -76,7 +76,7 @@
close_callback=None # type: Optional[Callable[[bytes], None]]
):
# type: (...) -> None
- super(ClosableOutputStream, self).__init__()
+ super().__init__()
self._close_callback = close_callback
def close(self):
@@ -117,7 +117,7 @@
flush_callback=None, # type: Optional[Callable[[bytes], None]]
size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD # type: int
):
- super(SizeBasedBufferingClosableOutputStream, self).__init__(close_callback)
+ super().__init__(close_callback)
self._flush_callback = flush_callback
self._size_flush_threshold = size_flush_threshold
@@ -147,8 +147,7 @@
time_flush_threshold_ms=_DEFAULT_TIME_FLUSH_THRESHOLD_MS # type: int
):
# type: (...) -> None
- super(TimeBasedBufferingClosableOutputStream,
- self).__init__(close_callback, flush_callback, size_flush_threshold)
+ 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()
@@ -159,7 +158,7 @@
def flush(self):
# type: () -> None
with self._flush_lock:
- super(TimeBasedBufferingClosableOutputStream, self).flush()
+ super().flush()
def close(self):
# type: () -> None
@@ -168,7 +167,7 @@
if self._periodic_flusher:
self._periodic_flusher.cancel()
self._periodic_flusher = None
- super(TimeBasedBufferingClosableOutputStream, self).close()
+ super().close()
def _schedule_periodic_flush(self):
# type: () -> None
@@ -663,7 +662,7 @@
data_buffer_time_limit_ms=0 # type: int
):
# type: (...) -> None
- super(GrpcClientDataChannel, self).__init__(data_buffer_time_limit_ms)
+ super().__init__(data_buffer_time_limit_ms)
self.set_inputs(data_stub.Data(self._write_outputs()))
diff --git a/sdks/python/apache_beam/runners/worker/log_handler.py b/sdks/python/apache_beam/runners/worker/log_handler.py
index 46157db..75cdcf5 100644
--- a/sdks/python/apache_beam/runners/worker/log_handler.py
+++ b/sdks/python/apache_beam/runners/worker/log_handler.py
@@ -72,7 +72,7 @@
def __init__(self, log_service_descriptor):
# type: (endpoints_pb2.ApiServiceDescriptor) -> None
- super(FnApiLogRecordHandler, self).__init__()
+ super().__init__()
self._alive = True
self._dropped_logs = 0
@@ -150,7 +150,7 @@
self._reader.join()
self.release()
# Unregister this handler.
- super(FnApiLogRecordHandler, self).close()
+ super().close()
except Exception:
# Log rather than raising exceptions, to avoid clobbering
# underlying errors that may have caused this to close
diff --git a/sdks/python/apache_beam/runners/worker/logger.py b/sdks/python/apache_beam/runners/worker/logger.py
index e171caf..6a86e00 100644
--- a/sdks/python/apache_beam/runners/worker/logger.py
+++ b/sdks/python/apache_beam/runners/worker/logger.py
@@ -43,7 +43,7 @@
class _PerThreadWorkerData(threading.local):
def __init__(self):
# type: () -> None
- super(_PerThreadWorkerData, self).__init__()
+ super().__init__()
# in the list, as going up and down all the way to zero incurs several
# reallocations.
self.stack = [] # type: List[Dict[str, Any]]
@@ -74,7 +74,7 @@
"""A JSON formatter class as expected by the logging standard module."""
def __init__(self, job_id, worker_id):
# type: (str, str) -> None
- super(JsonLogFormatter, self).__init__()
+ super().__init__()
self.job_id = job_id
self.worker_id = worker_id
diff --git a/sdks/python/apache_beam/runners/worker/opcounters.py b/sdks/python/apache_beam/runners/worker/opcounters.py
index bafbf9f..fad54aa 100644
--- a/sdks/python/apache_beam/runners/worker/opcounters.py
+++ b/sdks/python/apache_beam/runners/worker/opcounters.py
@@ -98,7 +98,7 @@
class NoOpTransformIOCounter(TransformIOCounter):
"""All operations for IO tracking are no-ops."""
def __init__(self):
- super(NoOpTransformIOCounter, self).__init__(None, None)
+ super().__init__(None, None)
def update_current_step(self):
pass
@@ -148,7 +148,7 @@
side input, and input_index is the index of the PCollectionView within
the list of inputs.
"""
- super(SideInputReadCounter, self).__init__(counter_factory, state_sampler)
+ super().__init__(counter_factory, state_sampler)
self.declaring_step = declaring_step
self.input_index = input_index
diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py
index 19546dd..1de42f9 100644
--- a/sdks/python/apache_beam/runners/worker/operations.py
+++ b/sdks/python/apache_beam/runners/worker/operations.py
@@ -21,6 +21,7 @@
"""Worker operations executor."""
# pytype: skip-file
+# pylint: disable=super-with-arguments
import collections
import logging
@@ -710,13 +711,12 @@
def process(self, o):
# type: (WindowedValue) -> None
with self.scoped_process_state:
- delayed_application = self.dofn_runner.process(o)
- if delayed_application:
+ delayed_applications = self.dofn_runner.process(o)
+ if delayed_applications:
assert self.execution_context is not None
- # TODO(BEAM-77746): there's disagreement between subclasses
- # of DoFnRunner over the return type annotations of process().
- self.execution_context.delayed_applications.append(
- (self, delayed_application)) # type: ignore[arg-type]
+ for delayed_application in delayed_applications:
+ self.execution_context.delayed_applications.append(
+ (self, delayed_application))
def finalize_bundle(self):
# type: () -> None
diff --git a/sdks/python/apache_beam/runners/worker/sideinputs_test.py b/sdks/python/apache_beam/runners/worker/sideinputs_test.py
index f609cf4..2e89b86 100644
--- a/sdks/python/apache_beam/runners/worker/sideinputs_test.py
+++ b/sdks/python/apache_beam/runners/worker/sideinputs_test.py
@@ -44,7 +44,7 @@
class FakeSourceReader(observable.ObservableMixin):
def __init__(self, items, notify_observers=False):
- super(FakeSourceReader, self).__init__()
+ super().__init__()
self.items = items
self.entered = False
self.exited = False
diff --git a/sdks/python/apache_beam/runners/worker/statesampler.py b/sdks/python/apache_beam/runners/worker/statesampler.py
index 7230b24..2d975ba 100644
--- a/sdks/python/apache_beam/runners/worker/statesampler.py
+++ b/sdks/python/apache_beam/runners/worker/statesampler.py
@@ -104,7 +104,7 @@
self.tracked_thread = None # type: Optional[threading.Thread]
self.finished = False
self.started = False
- super(StateSampler, self).__init__(sampling_period_ms)
+ super().__init__(sampling_period_ms)
@property
def stage_name(self):
@@ -114,7 +114,7 @@
def stop(self):
# type: () -> None
set_current_tracker(None)
- super(StateSampler, self).stop()
+ super().stop()
def stop_if_still_running(self):
# type: () -> None
@@ -125,7 +125,7 @@
# type: () -> None
self.tracked_thread = threading.current_thread()
set_current_tracker(self)
- super(StateSampler, self).start()
+ super().start()
self.started = True
def get_info(self):
@@ -171,12 +171,8 @@
else:
output_counter = self._counter_factory.get_counter(
counter_name, Counter.SUM)
- self._states_by_name[counter_name] = super(StateSampler,
- self)._scoped_state(
- counter_name,
- name_context,
- output_counter,
- metrics_container)
+ 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):
diff --git a/sdks/python/apache_beam/runners/worker/worker_pool_main.py b/sdks/python/apache_beam/runners/worker/worker_pool_main.py
index e5bfff8..eb5cdd9 100644
--- a/sdks/python/apache_beam/runners/worker/worker_pool_main.py
+++ b/sdks/python/apache_beam/runners/worker/worker_pool_main.py
@@ -76,8 +76,10 @@
container_executable=None # type: Optional[str]
):
# type: (...) -> 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())
+ thread_pool_executor.shared_unbounded_instance(), options=options)
worker_address = 'localhost:%s' % worker_server.add_insecure_port(
'[::]:%s' % port)
worker_pool = cls(
diff --git a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/winning_bids.py b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/winning_bids.py
index 94f84b6..52ffd48 100644
--- a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/winning_bids.py
+++ b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/winning_bids.py
@@ -45,7 +45,7 @@
class AuctionOrBidWindow(IntervalWindow):
"""Windows for open auctions and bids."""
def __init__(self, start, end, auction_id, is_auction_window):
- super(AuctionOrBidWindow, self).__init__(start, end)
+ super().__init__(start, end)
self.auction = auction_id
self.is_auction_window = is_auction_window
diff --git a/sdks/python/apache_beam/testing/load_tests/co_group_by_key_test.py b/sdks/python/apache_beam/testing/load_tests/co_group_by_key_test.py
index e7d212b..617e00d 100644
--- a/sdks/python/apache_beam/testing/load_tests/co_group_by_key_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/co_group_by_key_test.py
@@ -90,7 +90,7 @@
CO_INPUT_TAG = 'pc2'
def __init__(self):
- super(CoGroupByKeyTest, self).__init__()
+ super().__init__()
self.co_input_options = json.loads(
self.pipeline.get_option('co_input_options'))
self.iterations = self.get_option_or_default('iterations', 1)
diff --git a/sdks/python/apache_beam/testing/load_tests/combine_test.py b/sdks/python/apache_beam/testing/load_tests/combine_test.py
index d3a372f..9452730 100644
--- a/sdks/python/apache_beam/testing/load_tests/combine_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/combine_test.py
@@ -82,7 +82,7 @@
class CombineTest(LoadTest):
def __init__(self):
- super(CombineTest, self).__init__()
+ super().__init__()
self.fanout = self.get_option_or_default('fanout', 1)
try:
self.top_count = int(self.pipeline.get_option('top_count'))
diff --git a/sdks/python/apache_beam/testing/load_tests/group_by_key_test.py b/sdks/python/apache_beam/testing/load_tests/group_by_key_test.py
index 69ea74c..38724fc 100644
--- a/sdks/python/apache_beam/testing/load_tests/group_by_key_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/group_by_key_test.py
@@ -81,7 +81,7 @@
class GroupByKeyTest(LoadTest):
def __init__(self):
- super(GroupByKeyTest, self).__init__()
+ super().__init__()
self.fanout = self.get_option_or_default('fanout', 1)
self.iterations = self.get_option_or_default('iterations', 1)
diff --git a/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py b/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py
index f6d3340..7b975bb 100644
--- a/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py
+++ b/sdks/python/apache_beam/testing/load_tests/load_test_metrics_utils.py
@@ -48,7 +48,7 @@
from apache_beam.utils.timestamp import Timestamp
try:
- from google.cloud import bigquery # type: ignore
+ from google.cloud import bigquery # type: ignore[attr-defined]
from google.cloud.bigquery.schema import SchemaField
from google.cloud.exceptions import NotFound
except ImportError:
@@ -319,8 +319,7 @@
"""
def __init__(self, counter_metric, submit_timestamp, metric_id):
value = counter_metric.result
- super(CounterMetric,
- self).__init__(submit_timestamp, metric_id, value, counter_metric)
+ super().__init__(submit_timestamp, metric_id, value, counter_metric)
class DistributionMetric(Metric):
@@ -342,7 +341,7 @@
'not None.' % custom_label
_LOGGER.debug(msg)
raise ValueError(msg)
- super(DistributionMetric, self) \
+ super() \
.__init__(submit_timestamp, metric_id, value, dist_metric, custom_label)
@@ -361,8 +360,7 @@
# out of many steps
label = runtime_list[0].key.metric.namespace + \
'_' + RUNTIME_METRIC
- super(RuntimeMetric,
- self).__init__(submit_timestamp, metric_id, value, None, label)
+ super().__init__(submit_timestamp, metric_id, value, None, label)
def _prepare_runtime_metrics(self, distributions):
min_values = []
diff --git a/sdks/python/apache_beam/testing/load_tests/microbenchmarks_test.py b/sdks/python/apache_beam/testing/load_tests/microbenchmarks_test.py
index 2dff740..34d4080 100644
--- a/sdks/python/apache_beam/testing/load_tests/microbenchmarks_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/microbenchmarks_test.py
@@ -60,7 +60,7 @@
class MicroBenchmarksLoadTest(LoadTest):
def __init__(self):
- super(MicroBenchmarksLoadTest, self).__init__()
+ super().__init__()
def test(self):
self.extra_metrics.update(self._run_fn_api_runner_microbenchmark())
diff --git a/sdks/python/apache_beam/testing/load_tests/pardo_test.py b/sdks/python/apache_beam/testing/load_tests/pardo_test.py
index 6722fe3..989ed21 100644
--- a/sdks/python/apache_beam/testing/load_tests/pardo_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/pardo_test.py
@@ -88,7 +88,7 @@
class ParDoTest(LoadTest):
def __init__(self):
- super(ParDoTest, self).__init__()
+ super().__init__()
self.iterations = self.get_option_or_default('iterations')
self.number_of_counters = self.get_option_or_default(
'number_of_counters', 1)
diff --git a/sdks/python/apache_beam/testing/load_tests/sideinput_test.py b/sdks/python/apache_beam/testing/load_tests/sideinput_test.py
index f77e35c..745d961 100644
--- a/sdks/python/apache_beam/testing/load_tests/sideinput_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/sideinput_test.py
@@ -79,7 +79,7 @@
SDF_INITIAL_ELEMENTS = 1000
def __init__(self):
- super(SideInputTest, self).__init__()
+ super().__init__()
self.windows = self.get_option_or_default('window_count', default=1)
self.access_percentage = self.get_option_or_default(
diff --git a/sdks/python/apache_beam/testing/metric_result_matchers.py b/sdks/python/apache_beam/testing/metric_result_matchers.py
index a4a7f69..3c03865 100644
--- a/sdks/python/apache_beam/testing/metric_result_matchers.py
+++ b/sdks/python/apache_beam/testing/metric_result_matchers.py
@@ -80,8 +80,8 @@
self.step = _matcher_or_equal_to(step)
self.attempted = _matcher_or_equal_to(attempted)
self.committed = _matcher_or_equal_to(committed)
- labels = labels or dict()
- self.label_matchers = dict()
+ 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)
diff --git a/sdks/python/apache_beam/testing/metric_result_matchers_test.py b/sdks/python/apache_beam/testing/metric_result_matchers_test.py
index 9f4d408..3657356 100644
--- a/sdks/python/apache_beam/testing/metric_result_matchers_test.py
+++ b/sdks/python/apache_beam/testing/metric_result_matchers_test.py
@@ -79,7 +79,7 @@
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 dict()
+ labels = data_dict['labels'] if 'labels' in data_dict else {}
values = {}
for key in ['attempted', 'committed']:
if key in data_dict:
diff --git a/sdks/python/apache_beam/testing/test_pipeline.py b/sdks/python/apache_beam/testing/test_pipeline.py
index 910f149..9f70e4d 100644
--- a/sdks/python/apache_beam/testing/test_pipeline.py
+++ b/sdks/python/apache_beam/testing/test_pipeline.py
@@ -106,10 +106,10 @@
self.blocking = blocking
if options is None:
options = PipelineOptions(self.options_list)
- super(TestPipeline, self).__init__(runner, options)
+ super().__init__(runner, options)
def run(self, test_runner_api=True):
- result = super(TestPipeline, self).run(
+ result = super().run(
test_runner_api=(
False if self.not_use_test_runner_api else test_runner_api))
if self.blocking:
diff --git a/sdks/python/apache_beam/testing/test_stream.py b/sdks/python/apache_beam/testing/test_stream.py
index d655a90..734ca8e 100644
--- a/sdks/python/apache_beam/testing/test_stream.py
+++ b/sdks/python/apache_beam/testing/test_stream.py
@@ -286,7 +286,7 @@
endpoint: (str) a URL locating a TestStreamService.
"""
- super(TestStream, self).__init__()
+ super().__init__()
assert coder is not None
self.coder = coder
diff --git a/sdks/python/apache_beam/testing/test_utils.py b/sdks/python/apache_beam/testing/test_utils.py
index b58bfcd..46b82bb 100644
--- a/sdks/python/apache_beam/testing/test_utils.py
+++ b/sdks/python/apache_beam/testing/test_utils.py
@@ -23,7 +23,7 @@
# pytype: skip-file
import hashlib
-import imp
+import importlib
import os
import shutil
import tempfile
@@ -112,12 +112,12 @@
side_effect=patched_retry_with_exponential_backoff).start()
# Reload module after patching.
- imp.reload(module)
+ importlib.reload(module)
def remove_patches():
patch.stopall()
# Reload module again after removing patch.
- imp.reload(module)
+ importlib.reload(module)
testcase.addCleanup(remove_patches)
diff --git a/sdks/python/apache_beam/testing/util.py b/sdks/python/apache_beam/testing/util.py
index c42f90e..8c91812 100644
--- a/sdks/python/apache_beam/testing/util.py
+++ b/sdks/python/apache_beam/testing/util.py
@@ -96,11 +96,11 @@
actual = windowed_value.value
window_key = windowed_value.windows[0]
try:
- expected = _expected[window_key]
+ _expected[window_key]
except KeyError:
raise BeamAssertException(
'Failed assert: window {} not found in any expected ' \
- 'windows {}'.format(window_key, list(_expected.keys())))
+ '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
@@ -110,7 +110,7 @@
except ValueError:
raise BeamAssertException(
'Failed assert: element {} not found in window ' \
- '{}:{}'.format(actual, window_key, _expected[window_key]))
+ '{}:{}'.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.
diff --git a/sdks/python/apache_beam/tools/fn_api_runner_microbenchmark.py b/sdks/python/apache_beam/tools/fn_api_runner_microbenchmark.py
index a7bf7cb..a73b228 100644
--- a/sdks/python/apache_beam/tools/fn_api_runner_microbenchmark.py
+++ b/sdks/python/apache_beam/tools/fn_api_runner_microbenchmark.py
@@ -61,7 +61,6 @@
import random
import apache_beam as beam
-import apache_beam.typehints.typehints as typehints
from apache_beam.coders import VarIntCoder
from apache_beam.runners.portability.fn_api_runner import FnApiRunner
from apache_beam.tools import utils
@@ -69,6 +68,7 @@
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
diff --git a/sdks/python/apache_beam/tools/teststream_microbenchmark.py b/sdks/python/apache_beam/tools/teststream_microbenchmark.py
index 4b00de0..7c5bb61 100644
--- a/sdks/python/apache_beam/tools/teststream_microbenchmark.py
+++ b/sdks/python/apache_beam/tools/teststream_microbenchmark.py
@@ -45,12 +45,12 @@
import random
import apache_beam as beam
-import apache_beam.typehints.typehints as typehints
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
diff --git a/sdks/python/apache_beam/transforms/combinefn_lifecycle_pipeline.py b/sdks/python/apache_beam/transforms/combinefn_lifecycle_pipeline.py
index 1964082..51f66b3 100644
--- a/sdks/python/apache_beam/transforms/combinefn_lifecycle_pipeline.py
+++ b/sdks/python/apache_beam/transforms/combinefn_lifecycle_pipeline.py
@@ -38,7 +38,7 @@
instances = set() # type: Set[CallSequenceEnforcingCombineFn]
def __init__(self):
- super(CallSequenceEnforcingCombineFn, self).__init__()
+ super().__init__()
self._setup_called = False
self._teardown_called = False
diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py
index 41ad3df..65e8b04 100644
--- a/sdks/python/apache_beam/transforms/combiners.py
+++ b/sdks/python/apache_beam/transforms/combiners.py
@@ -58,7 +58,7 @@
class CombinerWithoutDefaults(ptransform.PTransform):
"""Super class to inherit without_defaults to built-in Combiners."""
def __init__(self, has_defaults=True):
- super(CombinerWithoutDefaults, self).__init__()
+ super().__init__()
self.has_defaults = has_defaults
def with_defaults(self, has_defaults=True):
@@ -191,7 +191,7 @@
reverse: (optional) whether to order things smallest to largest, rather
than largest to smallest
"""
- super(Top.Of, self).__init__()
+ super().__init__()
self._n = n
self._key = key
self._reverse = reverse
@@ -519,7 +519,7 @@
class Smallest(TopCombineFn):
def __init__(self, n):
- super(Smallest, self).__init__(n, reverse=True)
+ super().__init__(n, reverse=True)
def default_label(self):
return 'Smallest(%s)' % self._n
@@ -535,7 +535,7 @@
class FixedSizeGlobally(CombinerWithoutDefaults):
"""Sample n elements from the input PCollection without replacement."""
def __init__(self, n):
- super(Sample.FixedSizeGlobally, self).__init__()
+ super().__init__()
self._n = n
def expand(self, pcoll):
@@ -573,7 +573,7 @@
class SampleCombineFn(core.CombineFn):
"""CombineFn for all Sample transforms."""
def __init__(self, n):
- super(SampleCombineFn, self).__init__()
+ 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
@@ -659,10 +659,9 @@
]
def extract_output(self, accumulator, *args, **kwargs):
- return tuple([
+ return tuple(
c.extract_output(a, *args, **kwargs) for c,
- a in zip(self._combiners, accumulator)
- ])
+ a in zip(self._combiners, accumulator))
def teardown(self, *args, **kwargs):
for c in reversed(self._combiners):
@@ -753,7 +752,7 @@
class ToDictCombineFn(core.CombineFn):
"""CombineFn for to_dict."""
def create_accumulator(self):
- return dict()
+ return {}
def add_input(self, accumulator, element):
key, value = element
@@ -761,7 +760,7 @@
return accumulator
def merge_accumulators(self, accumulators):
- result = dict()
+ result = {}
for a in accumulators:
result.update(a)
return result
diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py
index 25b05df..2bd2d4f 100644
--- a/sdks/python/apache_beam/transforms/core.py
+++ b/sdks/python/apache_beam/transforms/core.py
@@ -19,10 +19,13 @@
# pytype: skip-file
+import concurrent.futures
import copy
import inspect
import logging
import random
+import sys
+import traceback
import types
import typing
@@ -731,7 +734,7 @@
# For cases such as set / list where fn is callable but not a function
self.process = lambda element: fn(element)
- super(CallableWrapperDoFn, self).__init__()
+ super().__init__()
def display_data(self):
# If the callable has a name, then it's likely a function, and
@@ -1009,7 +1012,7 @@
if not callable(fn):
raise TypeError('Expected a callable object instead of: %r' % fn)
- super(CallableWrapperCombineFn, self).__init__()
+ super().__init__()
self._fn = fn
self._buffer_size = buffer_size
@@ -1210,7 +1213,7 @@
exact positions where they appear in the argument lists.
"""
def __init__(self, fn, *args, **kwargs):
- super(ParDo, self).__init__(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]
@@ -1222,6 +1225,75 @@
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):
+ """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 records (those which threw exceptions during
+ processing) along with information about the errors raised.
+
+ For example, one would write::
+
+ good, bad = Map(maybe_error_raising_function).with_exception_handling()
+
+ and `good` will be a PCollection of mapped records and `bad` will contain
+ those that raised exceptions.
+
+
+ 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 records, 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 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).
+ threshold_windowing: Event-time windowing to use for threshold. Optional,
+ defaults to the windowing of the input.
+ """
+ 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)
+
def default_type_hints(self):
return self.fn.get_type_hints()
@@ -1267,7 +1339,7 @@
return pvalue.PCollection.from_(pcoll)
- def with_outputs(self, *tags, **main_kw):
+ def with_outputs(self, *tags, main=None, allow_unknown_tags=None):
"""Returns a tagged tuple allowing access to the outputs of a
:class:`ParDo`.
@@ -1297,13 +1369,11 @@
ValueError: if **main_kw** contains any key other than
``'main'``.
"""
- main_tag = main_kw.pop('main', None)
- if main_tag in tags:
+ if main in tags:
raise ValueError(
- 'Main output tag must be different from side output tags.')
- if main_kw:
- raise ValueError('Unexpected keyword arguments: %s' % list(main_kw))
- return _MultiParDo(self, tags, main_tag)
+ 'Main output tag %r must be different from side output tags %r.' %
+ (main, tags))
+ return _MultiParDo(self, tags, main, allow_unknown_tags)
def _do_fn_info(self):
return DoFnInfo.create(self.fn, self.args, self.kwargs)
@@ -1419,16 +1489,21 @@
class _MultiParDo(PTransform):
- def __init__(self, do_transform, tags, main_tag):
- super(_MultiParDo, self).__init__(do_transform.label)
+ 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)
+ pcoll.pipeline,
+ self._do_transform,
+ self._tags,
+ self._main_tag,
+ self._allow_unknown_tags)
class DoFnInfo(object):
@@ -1733,6 +1808,188 @@
return pardo
+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):
+ 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
+
+ def expand(self, pcoll):
+ result = pcoll | ParDo(
+ _ExceptionHandlingWrapperDoFn(
+ _SubprocessDoFn(self._fn) if self._use_subprocess else self._fn,
+ self._dead_letter_tag,
+ self._exc_class,
+ self._partial),
+ *self._args,
+ **self._kwargs).with_outputs(
+ self._dead_letter_tag, main=self._main_tag, allow_unknown_tags=True)
+
+ 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
+
+ input_count_view = pcoll | 'CountTotal' >> (
+ MaybeWindow() | Map(lambda _: 1)
+ | CombineGlobally(sum).as_singleton_view())
+ bad_count_pcoll = result[self._dead_letter_tag] | 'CountBad' >> (
+ MaybeWindow() | 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)
+
+ return result
+
+
+class _ExceptionHandlingWrapperDoFn(DoFn):
+ def __init__(self, fn, dead_letter_tag, exc_class, partial):
+ self._fn = fn
+ self._dead_letter_tag = dead_letter_tag
+ self._exc_class = exc_class
+ self._partial = partial
+
+ 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 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:
+ yield pvalue.TaggedOutput(
+ self._dead_letter_tag,
+ (
+ args[0], (
+ type(exn),
+ repr(exn),
+ traceback.format_exception(*sys.exc_info()))))
+
+
+class _SubprocessDoFn(DoFn):
+ """Process method run in a subprocess, turning hard crashes into exceptions.
+ """
+ def __init__(self, fn):
+ self._fn = fn
+ self._serialized_fn = pickler.dumps(fn)
+
+ 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._pool.shutdown()
+ self._pool = None
+
+ 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()
+ except concurrent.futures.process.BrokenProcessPool:
+ self._pool = None
+ raise
+
+ # 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
+
+
def Filter(fn, *args, **kwargs): # pylint: disable=invalid-name
""":func:`Filter` is a :func:`FlatMap` with its callable filtering out
elements.
@@ -1848,7 +2105,7 @@
'CombineGlobally can be used only with combineFn objects. '
'Received %r instead.' % (fn))
- super(CombineGlobally, self).__init__()
+ super().__init__()
self.fn = fn
self.args = args
self.kwargs = kwargs
@@ -2136,7 +2393,7 @@
combinefn, # type: CombineFn
runtime_type_check, # type: bool
):
- super(CombineValuesDoFn, self).__init__()
+ super().__init__()
self.combinefn = combinefn
self.runtime_type_check = runtime_type_check
@@ -2795,7 +3052,7 @@
accumulation_mode,
timestamp_combiner,
allowed_lateness)
- super(WindowInto, self).__init__(self.WindowIntoFn(self.windowing))
+ super().__init__(self.WindowIntoFn(self.windowing))
def get_windowing(self, unused_inputs):
# type: (typing.Any) -> Windowing
@@ -2811,7 +3068,7 @@
output_type = input_type
self.with_input_types(input_type)
self.with_output_types(output_type)
- return super(WindowInto, self).expand(pcoll)
+ 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]
@@ -2859,7 +3116,7 @@
provide pipeline information and should be considered mandatory.
"""
def __init__(self, **kwargs):
- super(Flatten, self).__init__()
+ super().__init__()
self.pipeline = kwargs.pop(
'pipeline', None) # type: typing.Optional[Pipeline]
if kwargs:
@@ -2905,7 +3162,7 @@
Args:
values: An object of values for the PCollection
"""
- super(Create, self).__init__()
+ super().__init__()
if isinstance(values, (str, bytes)):
raise TypeError(
'PTransform Create: Refusing to treat string as '
diff --git a/sdks/python/apache_beam/transforms/deduplicate_test.py b/sdks/python/apache_beam/transforms/deduplicate_test.py
index b6ec53d..392dac2 100644
--- a/sdks/python/apache_beam/transforms/deduplicate_test.py
+++ b/sdks/python/apache_beam/transforms/deduplicate_test.py
@@ -46,7 +46,7 @@
def __init__(self, *args, **kwargs):
self.runner = None
self.options = None
- super(DeduplicateTest, self).__init__(*args, **kwargs)
+ super().__init__(*args, **kwargs)
def set_runner(self, runner):
self.runner = runner
diff --git a/sdks/python/apache_beam/transforms/display.py b/sdks/python/apache_beam/transforms/display.py
index 449e45a..cd9de00 100644
--- a/sdks/python/apache_beam/transforms/display.py
+++ b/sdks/python/apache_beam/transforms/display.py
@@ -149,13 +149,28 @@
value = display_data_dict['value']
if isinstance(value, str):
return beam_runner_api_pb2.LabelledPayload(
- label=label, string_value=value)
+ 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)
- elif isinstance(value, (int, float, complex)):
+ 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, double_value=value)
+ 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,
+ key=display_data_dict['key'],
+ namespace=display_data_dict.get('namespace', ''))
else:
raise ValueError(
'Unsupported type %s for value of display data %s' %
diff --git a/sdks/python/apache_beam/transforms/external.py b/sdks/python/apache_beam/transforms/external.py
index 16840c4..005cb18 100644
--- a/sdks/python/apache_beam/transforms/external.py
+++ b/sdks/python/apache_beam/transforms/external.py
@@ -140,7 +140,7 @@
"""
:param tuple_instance: an instance of a typing.NamedTuple
"""
- super(NamedTupleBasedPayloadBuilder, self).__init__()
+ super().__init__()
self._tuple_instance = tuple_instance
def _get_named_tuple_instance(self):
@@ -236,8 +236,7 @@
:param args: parameter values of the constructor.
:param kwargs: parameter names and values of the constructor.
"""
- if (self._constructor_method or self._constructor_param_args or
- self._constructor_param_kwargs):
+ if self._has_constructor():
raise ValueError(
'Constructor or constructor method can only be specified once')
@@ -254,8 +253,7 @@
:param args: parameter values of the constructor method.
:param kwargs: parameter names and values of the constructor method.
"""
- if (self._constructor_method or self._constructor_param_args or
- self._constructor_param_kwargs):
+ if self._has_constructor():
raise ValueError(
'Constructor or constructor method can only be specified once')
@@ -276,6 +274,46 @@
"""
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)
+
+
+class JavaExternalTransform(ptransform.PTransform):
+ """A proxy for Java-implemented external transforms.
+
+ One builds these transforms just as one would in Java.
+ """
+ def __init__(self, class_name, expansion_service=None):
+ self._payload_builder = JavaClassLookupPayloadBuilder(class_name)
+ self._expansion_service = 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)
+
+ 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):
+ return pcolls | ExternalTransform(
+ common_urns.java_class_lookup,
+ self._payload_builder.build(),
+ self._expansion_service)
+
class AnnotationBasedPayloadBuilder(SchemaBasedPayloadBuilder):
"""
@@ -598,7 +636,7 @@
def __init__(self, channel, **kwargs):
self._channel = channel
self._kwargs = kwargs
- super(ExpansionAndArtifactRetrievalStub, self).__init__(channel, **kwargs)
+ super().__init__(channel, **kwargs)
def artifact_service(self):
return beam_artifact_api_pb2_grpc.ArtifactRetrievalServiceStub(
@@ -614,7 +652,7 @@
"""
def __init__(self, path_to_jar, extra_args=None):
if extra_args is None:
- extra_args = ['{{PORT}}']
+ extra_args = ['{{PORT}}', f'--filesToStage={path_to_jar}']
self._path_to_jar = path_to_jar
self._extra_args = extra_args
self._service_count = 0
@@ -648,7 +686,7 @@
def __init__(self, gradle_target, extra_args=None, gradle_appendix=None):
path_to_jar = subprocess_server.JavaJarServer.path_to_beam_jar(
gradle_target, gradle_appendix)
- super(BeamJarExpansionService, self).__init__(path_to_jar, extra_args)
+ super().__init__(path_to_jar, extra_args)
def memoize(func):
diff --git a/sdks/python/apache_beam/transforms/external_test.py b/sdks/python/apache_beam/transforms/external_test.py
index 8673ed5..83e72461 100644
--- a/sdks/python/apache_beam/transforms/external_test.py
+++ b/sdks/python/apache_beam/transforms/external_test.py
@@ -39,6 +39,7 @@
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 NamedTupleBasedPayloadBuilder
from apache_beam.typehints import typehints
from apache_beam.typehints.native_type_compatibility import convert_to_beam_type
@@ -334,7 +335,7 @@
mapping: typing.Mapping[str, float],
optional_integer: typing.Optional[int] = None,
expansion_service=None):
- super(AnnotatedTransform, self).__init__(
+ super().__init__(
self.URN,
AnnotationBasedPayloadBuilder(
self,
@@ -362,7 +363,7 @@
mapping: typehints.Dict[str, float],
optional_integer: typehints.Optional[int] = None,
expansion_service=None):
- super(AnnotatedTransform, self).__init__(
+ super().__init__(
self.URN,
AnnotationBasedPayloadBuilder(
self,
@@ -509,6 +510,42 @@
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, 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, 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, {})
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/transforms/periodicsequence.py b/sdks/python/apache_beam/transforms/periodicsequence.py
index 9fc9029..7976323 100644
--- a/sdks/python/apache_beam/transforms/periodicsequence.py
+++ b/sdks/python/apache_beam/transforms/periodicsequence.py
@@ -19,9 +19,9 @@
import time
import apache_beam as beam
-import apache_beam.runners.sdf_utils as sdf_utils
from apache_beam.io.restriction_trackers import OffsetRange
from apache_beam.io.restriction_trackers import OffsetRestrictionTracker
+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
@@ -124,6 +124,8 @@
PeriodicSequence guarantees that elements would not be output prior to given
runtime timestamp.
'''
+
+ # pylint: disable=unused-private-member
def __init_(self):
pass
diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py
index b314abd..40626aa 100644
--- a/sdks/python/apache_beam/transforms/ptransform.py
+++ b/sdks/python/apache_beam/transforms/ptransform.py
@@ -202,8 +202,7 @@
class _MaterializedDoOutputsTuple(pvalue.DoOutputsTuple):
def __init__(self, deferred, results_by_tag):
- super(_MaterializedDoOutputsTuple,
- self).__init__(None, None, deferred._tags, deferred._main_tag)
+ super().__init__(None, None, deferred._tags, deferred._main_tag)
self._deferred = deferred
self._results_by_tag = results_by_tag
@@ -351,7 +350,7 @@
def __init__(self, label=None):
# type: (Optional[str]) -> None
- super(PTransform, self).__init__()
+ super().__init__()
self.label = label # type: ignore # https://github.com/python/mypy/issues/3004
@property
@@ -402,7 +401,7 @@
input_type_hint)
validate_composite_type_param(
input_type_hint, 'Type hints for a PTransform')
- return super(PTransform, self).with_input_types(input_type_hint)
+ return super().with_input_types(input_type_hint)
def with_output_types(self, type_hint):
"""Annotates the output type of a :class:`PTransform` with a type-hint.
@@ -423,7 +422,7 @@
"""
type_hint = native_type_compatibility.convert_to_beam_type(type_hint)
validate_composite_type_param(type_hint, 'Type hints for a PTransform')
- return super(PTransform, self).with_output_types(type_hint)
+ return super().with_output_types(type_hint)
def with_resource_hints(self, **kwargs): # type: (...) -> PTransform
"""Adds resource hints to the :class:`PTransform`.
@@ -806,7 +805,7 @@
class _ChainedPTransform(PTransform):
def __init__(self, *parts):
# type: (*PTransform) -> None
- super(_ChainedPTransform, self).__init__(label=self._chain_label(parts))
+ super().__init__(label=self._chain_label(parts))
self._parts = parts
def _chain_label(self, parts):
@@ -842,10 +841,10 @@
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(PTransformWithSideInputs, self).__init__()
+ super().__init__()
- if (any([isinstance(v, pvalue.PCollection) for v in args]) or
- any([isinstance(v, pvalue.PCollection) for v in kwargs.values()])):
+ 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 '
@@ -898,7 +897,7 @@
:class:`PTransform` object. This allows chaining type-hinting related
methods.
"""
- super(PTransformWithSideInputs, self).with_input_types(input_type_hint)
+ super().with_input_types(input_type_hint)
side_inputs_arg_hints = native_type_compatibility.convert_to_beam_types(
side_inputs_arg_hints)
@@ -964,7 +963,7 @@
class _PTransformFnPTransform(PTransform):
"""A class wrapper for a function-based transform."""
def __init__(self, fn, *args, **kwargs):
- super(_PTransformFnPTransform, self).__init__()
+ super().__init__()
self._fn = fn
self._args = args
self._kwargs = kwargs
@@ -1031,7 +1030,7 @@
class CustomMapper(PTransform):
def __init__(self, mapfn):
- super(CustomMapper, self).__init__()
+ super().__init__()
self.mapfn = mapfn
def expand(self, pcoll):
@@ -1084,7 +1083,7 @@
class _NamedPTransform(PTransform):
def __init__(self, transform, label):
- super(_NamedPTransform, self).__init__(label)
+ super().__init__(label)
self.transform = transform
def __ror__(self, pvalueish, _unused=None):
diff --git a/sdks/python/apache_beam/transforms/ptransform_test.py b/sdks/python/apache_beam/transforms/ptransform_test.py
index ccf0a55..13d533c 100644
--- a/sdks/python/apache_beam/transforms/ptransform_test.py
+++ b/sdks/python/apache_beam/transforms/ptransform_test.py
@@ -21,6 +21,7 @@
import collections
import operator
+import os
import pickle
import random
import re
@@ -32,11 +33,12 @@
import hamcrest as hc
import pytest
+from parameterized import parameterized_class
import apache_beam as beam
-import apache_beam.pvalue as pvalue
import apache_beam.transforms.combiners as combine
-import apache_beam.typehints as typehints
+from apache_beam import pvalue
+from apache_beam import typehints
from apache_beam.io.iobase import Read
from apache_beam.metrics import Metrics
from apache_beam.metrics.metric import MetricsFilter
@@ -2524,6 +2526,235 @@
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_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 TestPTransformFn(TypeHintTestCase):
def test_type_checking_fail(self):
@beam.ptransform_fn
diff --git a/sdks/python/apache_beam/transforms/sql.py b/sdks/python/apache_beam/transforms/sql.py
index 30d5464..4102c73 100644
--- a/sdks/python/apache_beam/transforms/sql.py
+++ b/sdks/python/apache_beam/transforms/sql.py
@@ -85,7 +85,7 @@
"""
expansion_service = expansion_service or BeamJarExpansionService(
':sdks:java:extensions:sql:expansion-service:shadowJar')
- super(SqlTransform, self).__init__(
+ super().__init__(
self.URN,
NamedTupleBasedPayloadBuilder(
SqlTransformSchema(query=query, dialect=dialect)),
diff --git a/sdks/python/apache_beam/transforms/stats_test.py b/sdks/python/apache_beam/transforms/stats_test.py
index 7394380..bf634c0 100644
--- a/sdks/python/apache_beam/transforms/stats_test.py
+++ b/sdks/python/apache_beam/transforms/stats_test.py
@@ -621,7 +621,7 @@
[391, 977, 1221, 1526,
954], [782, 977, 1221, 1526,
1908], [3125, 3907, 9766, 12208, 15259]]
- test_data = list()
+ test_data = []
i = 0
for epsilon in epsilons:
j = 0
diff --git a/sdks/python/apache_beam/transforms/trigger.py b/sdks/python/apache_beam/transforms/trigger.py
index e5f7c24..d7a17ac 100644
--- a/sdks/python/apache_beam/transforms/trigger.py
+++ b/sdks/python/apache_beam/transforms/trigger.py
@@ -111,7 +111,7 @@
# TODO(robertwb): Also store the coder (perhaps extracted from the combine_fn)
def __init__(self, tag, combine_fn):
- super(_CombiningValueStateTag, self).__init__(tag)
+ super().__init__(tag)
if not combine_fn:
raise ValueError('combine_fn must be specified.')
if not isinstance(combine_fn, core.CombineFn):
@@ -148,7 +148,7 @@
class _WatermarkHoldStateTag(_StateTag):
def __init__(self, tag, timestamp_combiner_impl):
- super(_WatermarkHoldStateTag, self).__init__(tag)
+ super().__init__(tag)
self.timestamp_combiner_impl = timestamp_combiner_impl
def __repr__(self):
@@ -724,15 +724,8 @@
self.underlying.reset(window, context)
def may_lose_data(self, windowing):
- """Repeatedly may only lose data if the underlying trigger may not have
- its condition met.
-
- For underlying triggers that may finish, Repeatedly overrides that
- behavior.
- """
- return (
- self.underlying.may_lose_data(windowing)
- & DataLossReason.CONDITION_NOT_GUARANTEED)
+ """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):
@@ -1243,7 +1236,7 @@
class _UnwindowedValues(observable.ObservableMixin):
"""Exposes iterable of windowed values as iterable of unwindowed values."""
def __init__(self, windowed_values):
- super(_UnwindowedValues, self).__init__()
+ super().__init__()
self._windowed_values = windowed_values
def __iter__(self):
diff --git a/sdks/python/apache_beam/transforms/trigger_test.py b/sdks/python/apache_beam/transforms/trigger_test.py
index ed43094..f0f4902 100644
--- a/sdks/python/apache_beam/transforms/trigger_test.py
+++ b/sdks/python/apache_beam/transforms/trigger_test.py
@@ -496,7 +496,7 @@
self._test(
AfterWatermark(late=AfterCount(5)),
60,
- DataLossReason.CONDITION_NOT_GUARANTEED)
+ DataLossReason.NO_POTENTIAL_LOSS)
def test_after_count_one(self):
self._test(AfterCount(1), 0, DataLossReason.MAY_FINISH)
@@ -515,8 +515,7 @@
self._test(Repeatedly(AfterCount(1)), 0, DataLossReason.NO_POTENTIAL_LOSS)
def test_repeatedly_condition_underlying(self):
- self._test(
- Repeatedly(AfterCount(2)), 0, DataLossReason.CONDITION_NOT_GUARANTEED)
+ self._test(Repeatedly(AfterCount(2)), 0, DataLossReason.NO_POTENTIAL_LOSS)
def test_after_any_some_unsafe(self):
self._test(
@@ -532,7 +531,7 @@
def test_after_any_different_reasons(self):
self._test(
- AfterAny(Repeatedly(AfterCount(2)), AfterProcessingTime()),
+ AfterAny(AfterCount(2), AfterProcessingTime()),
0,
DataLossReason.MAY_FINISH | DataLossReason.CONDITION_NOT_GUARANTEED)
diff --git a/sdks/python/apache_beam/transforms/userstate.py b/sdks/python/apache_beam/transforms/userstate.py
index 84184d4..1276118 100644
--- a/sdks/python/apache_beam/transforms/userstate.py
+++ b/sdks/python/apache_beam/transforms/userstate.py
@@ -134,7 +134,7 @@
if coder is None:
coder = self.combine_fn.get_accumulator_coder()
- super(CombiningValueStateSpec, self).__init__(name, coder)
+ super().__init__(name, coder)
def to_runner_api(self, context):
# type: (PipelineContext) -> beam_runner_api_pb2.StateSpec
diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py
index 63bbece..77c7f85 100644
--- a/sdks/python/apache_beam/transforms/util.py
+++ b/sdks/python/apache_beam/transforms/util.py
@@ -632,7 +632,7 @@
Arguments:
window_coder: coders.Coder object to be used on windows.
"""
- super(_IdentityWindowFn, self).__init__()
+ super().__init__()
if window_coder is None:
raise ValueError('window_coder should not be None')
self._window_coder = window_coder
@@ -768,9 +768,9 @@
"""
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()]):
+ if all(isinstance(arg, AsSideInput)
+ for arg in args) and all(isinstance(kwarg, AsSideInput)
+ for kwarg in kwargs.values()):
return pcoll | Map(
lambda v,
*args,
@@ -963,6 +963,7 @@
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)
diff --git a/sdks/python/apache_beam/transforms/window.py b/sdks/python/apache_beam/transforms/window.py
index de16c73..522f661 100644
--- a/sdks/python/apache_beam/transforms/window.py
+++ b/sdks/python/apache_beam/transforms/window.py
@@ -318,7 +318,7 @@
def __init__(self):
# type: () -> None
- super(GlobalWindow, self).__init__(GlobalWindow._getTimestampFromProto())
+ super().__init__(GlobalWindow._getTimestampFromProto())
def __repr__(self):
return 'GlobalWindow'
diff --git a/sdks/python/apache_beam/transforms/window_test.py b/sdks/python/apache_beam/transforms/window_test.py
index 7369090..9b2254f 100644
--- a/sdks/python/apache_beam/transforms/window_test.py
+++ b/sdks/python/apache_beam/transforms/window_test.py
@@ -149,7 +149,7 @@
class TestMergeContext(WindowFn.MergeContext):
def __init__(self):
- super(TestMergeContext, self).__init__(running)
+ super().__init__(running)
def merge(self, to_be_merged, merge_result):
for w in to_be_merged:
diff --git a/sdks/python/apache_beam/typehints/schemas.py b/sdks/python/apache_beam/typehints/schemas.py
index d62ba54..9be2f8c 100644
--- a/sdks/python/apache_beam/typehints/schemas.py
+++ b/sdks/python/apache_beam/typehints/schemas.py
@@ -67,6 +67,7 @@
from uuid import uuid4
import numpy as np
+from google.protobuf import text_format
from apache_beam.portability.api import schema_pb2
from apache_beam.typehints import row_type
@@ -239,10 +240,19 @@
from apache_beam import coders
type_name = 'BeamSchema_{}'.format(schema.id.replace('-', '_'))
- user_type = NamedTuple(
- type_name,
- [(field.name, typing_from_runner_api(field.type))
- for field in schema.fields])
+
+ subfields = []
+ for field in schema.fields:
+ try:
+ field_py_type = typing_from_runner_api(field.type)
+ except ValueError as e:
+ raise ValueError(
+ "Failed to decode schema due to an issue with Field proto:\n\n" +
+ text_format.MessageToString(field)) from e
+
+ subfields.append((field.name, field_py_type))
+
+ user_type = NamedTuple(type_name, subfields)
setattr(user_type, _BEAM_SCHEMA_ID, schema.id)
@@ -266,6 +276,9 @@
return LogicalType.from_runner_api(
fieldtype_proto.logical_type).language_type()
+ else:
+ raise ValueError(f"Unrecognized type_info: {type_info!r}")
+
def _hydrate_namedtuple_instance(encoded_schema, values):
return named_tuple_from_schema(
diff --git a/sdks/python/apache_beam/typehints/schemas_test.py b/sdks/python/apache_beam/typehints/schemas_test.py
index c662bac..952ac54 100644
--- a/sdks/python/apache_beam/typehints/schemas_test.py
+++ b/sdks/python/apache_beam/typehints/schemas_test.py
@@ -275,6 +275,19 @@
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())
+ ])
+
+ # Should raise an exception referencing the problem field
+ self.assertRaisesRegex(
+ ValueError,
+ "type_with_no_typeinfo",
+ lambda: named_tuple_from_schema(schema_proto))
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/typehints/trivial_inference.py b/sdks/python/apache_beam/typehints/trivial_inference.py
index cc6534d..0d760db 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference.py
@@ -53,7 +53,7 @@
(name, instance_to_type(value)) for name, value in o.as_dict().items()
])
elif t not in typehints.DISALLOWED_PRIMITIVE_TYPES:
- # pylint: disable=deprecated-types-field
+ # pylint: disable=bad-option-value
if t == BoundMethod:
return types.MethodType
return t
diff --git a/sdks/python/apache_beam/typehints/typecheck.py b/sdks/python/apache_beam/typehints/typecheck.py
index 6c4ba25..9fe1804 100644
--- a/sdks/python/apache_beam/typehints/typecheck.py
+++ b/sdks/python/apache_beam/typehints/typecheck.py
@@ -44,7 +44,7 @@
class AbstractDoFnWrapper(DoFn):
"""An abstract class to create wrapper around DoFn"""
def __init__(self, dofn):
- super(AbstractDoFnWrapper, self).__init__()
+ super().__init__()
self.dofn = dofn
def _inspect_start_bundle(self):
@@ -78,7 +78,7 @@
class OutputCheckWrapperDoFn(AbstractDoFnWrapper):
"""A DoFn that verifies against common errors in the output type."""
def __init__(self, dofn, full_label):
- super(OutputCheckWrapperDoFn, self).__init__(dofn)
+ super().__init__(dofn)
self.full_label = full_label
def wrapper(self, method, args, kwargs):
@@ -116,7 +116,7 @@
"""A wrapper around a DoFn which performs type-checking of input and output.
"""
def __init__(self, dofn, type_hints, label=None):
- super(TypeCheckWrapperDoFn, self).__init__(dofn)
+ super().__init__(dofn)
self._process_fn = self.dofn._process_argspec_fn()
if type_hints.input_types:
input_args, input_kwargs = type_hints.input_types
diff --git a/sdks/python/apache_beam/typehints/typehints.py b/sdks/python/apache_beam/typehints/typehints.py
index 136ad01..7c1a813 100644
--- a/sdks/python/apache_beam/typehints/typehints.py
+++ b/sdks/python/apache_beam/typehints/typehints.py
@@ -251,7 +251,7 @@
for index, elem in enumerate(sequence_instance):
try:
check_constraint(self.inner_type, elem)
- except SimpleTypeHintError as e:
+ 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, '
@@ -599,7 +599,7 @@
"""
class TupleSequenceConstraint(SequenceTypeConstraint):
def __init__(self, type_param):
- super(TupleHint.TupleSequenceConstraint, self).__init__(type_param, tuple)
+ super().__init__(type_param, tuple)
def __repr__(self):
return 'Tuple[%s, ...]' % _unified_repr(self.inner_type)
@@ -610,7 +610,7 @@
return all(
is_consistent_with(elem, self.inner_type)
for elem in sub.tuple_types)
- return super(TupleSequenceConstraint, self)._consistent_with_check_(sub)
+ return super()._consistent_with_check_(sub)
class TupleConstraint(IndexableTypeConstraint):
def __init__(self, type_params):
@@ -731,7 +731,7 @@
"""
class ListConstraint(SequenceTypeConstraint):
def __init__(self, list_type):
- super(ListHint.ListConstraint, self).__init__(list_type, list)
+ super().__init__(list_type, list)
def __repr__(self):
return 'List[%s]' % _unified_repr(self.inner_type)
@@ -912,7 +912,7 @@
"""
class SetTypeConstraint(SequenceTypeConstraint):
def __init__(self, type_param):
- super(SetHint.SetTypeConstraint, self).__init__(type_param, set)
+ super().__init__(type_param, set)
def __repr__(self):
return 'Set[%s]' % _unified_repr(self.inner_type)
diff --git a/sdks/python/apache_beam/typehints/typehints_test.py b/sdks/python/apache_beam/typehints/typehints_test.py
index 2e0658c..2b39b60b 100644
--- a/sdks/python/apache_beam/typehints/typehints_test.py
+++ b/sdks/python/apache_beam/typehints/typehints_test.py
@@ -24,7 +24,6 @@
import typing
import unittest
-import apache_beam.typehints.typehints as typehints
from apache_beam import Map
from apache_beam import PTransform
from apache_beam.pvalue import PBegin
@@ -38,6 +37,7 @@
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
diff --git a/sdks/python/apache_beam/utils/counters.py b/sdks/python/apache_beam/utils/counters.py
index 320a9cd..bcc883e 100644
--- a/sdks/python/apache_beam/utils/counters.py
+++ b/sdks/python/apache_beam/utils/counters.py
@@ -114,7 +114,7 @@
output_index=None,
io_target=None):
origin = origin or CounterName.SYSTEM
- return super(CounterName, cls).__new__(
+ return super().__new__(
cls,
name,
stage_name,
@@ -206,7 +206,7 @@
def __init__(self, name, combine_fn):
# type: (CounterName, cy_combiners.AccumulatorCombineFn) -> None
assert isinstance(combine_fn, cy_combiners.AccumulatorCombineFn)
- super(AccumulatorCombineFnCounter, self).__init__(name, combine_fn)
+ super().__init__(name, combine_fn)
self.reset()
def update(self, value):
@@ -268,4 +268,4 @@
this method returns hence the returned iterable may be stale.
"""
with self._lock:
- return self.counters.values() # pylint: disable=dict-values-not-iterating
+ return self.counters.values() # pylint: disable=bad-option-value
diff --git a/sdks/python/apache_beam/utils/histogram.py b/sdks/python/apache_beam/utils/histogram.py
index 13bb5c2..83533c5 100644
--- a/sdks/python/apache_beam/utils/histogram.py
+++ b/sdks/python/apache_beam/utils/histogram.py
@@ -107,7 +107,7 @@
elif f == float('inf'):
return '>=%s' % self._bucket_type.range_to()
else:
- return str(int(round(f))) # pylint: disable=round-builtin
+ return str(int(round(f))) # pylint: disable=bad-option-value
with self._lock:
return (
diff --git a/sdks/python/apache_beam/utils/profiler.py b/sdks/python/apache_beam/utils/profiler.py
index de9f943..d10703c 100644
--- a/sdks/python/apache_beam/utils/profiler.py
+++ b/sdks/python/apache_beam/utils/profiler.py
@@ -23,7 +23,7 @@
# pytype: skip-file
# mypy: check-untyped-defs
-import cProfile # pylint: disable=bad-python3-import
+import cProfile
import io
import logging
import os
diff --git a/sdks/python/apache_beam/utils/shared.py b/sdks/python/apache_beam/utils/shared.py
index 79d7027..23622ef 100644
--- a/sdks/python/apache_beam/utils/shared.py
+++ b/sdks/python/apache_beam/utils/shared.py
@@ -72,7 +72,7 @@
def construct_table():
# Construct the rainbow table from the table elements.
# The table contains lines in the form "string::hash"
- result = dict()
+ result = {}
for key, value in table_elements:
result[value] = key
return result
@@ -209,7 +209,7 @@
self._lock = threading.Lock()
# Dictionary of references to shared control blocks
- self._cache_map = dict()
+ self._cache_map = {}
# Tuple of (key, obj), where obj is an object we explicitly hold a reference
# to keep it alive
diff --git a/sdks/python/apache_beam/utils/subprocess_server.py b/sdks/python/apache_beam/utils/subprocess_server.py
index d1466a2..7035cad 100644
--- a/sdks/python/apache_beam/utils/subprocess_server.py
+++ b/sdks/python/apache_beam/utils/subprocess_server.py
@@ -159,7 +159,7 @@
dict(__init__=lambda self: setattr(self, 'replacements', {})))()
def __init__(self, stub_class, path_to_jar, java_arguments):
- super(JavaJarServer, self).__init__(
+ super().__init__(
stub_class, ['java', '-jar', path_to_jar] + list(java_arguments))
self._existing_service = path_to_jar if _is_service_endpoint(
path_to_jar) else None
@@ -172,13 +172,13 @@
raise RuntimeError(
'Java must be installed on this system to use this '
'transform/runner.')
- return super(JavaJarServer, self).start_process()
+ return super().start_process()
def stop_process(self):
if self._existing_service:
pass
else:
- return super(JavaJarServer, self).stop_process()
+ return super().stop_process()
@classmethod
def jar_name(cls, artifact_id, version, classifier=None, appendix=None):
diff --git a/sdks/python/apache_beam/utils/thread_pool_executor.py b/sdks/python/apache_beam/utils/thread_pool_executor.py
index afe172a..e1e8ad5 100644
--- a/sdks/python/apache_beam/utils/thread_pool_executor.py
+++ b/sdks/python/apache_beam/utils/thread_pool_executor.py
@@ -41,7 +41,7 @@
class _Worker(threading.Thread):
def __init__(self, idle_worker_queue, work_item):
- super(_Worker, self).__init__()
+ super().__init__()
self._idle_worker_queue = idle_worker_queue
self._work_item = work_item
self._wake_semaphore = threading.Semaphore(0)
diff --git a/sdks/python/apache_beam/version.py b/sdks/python/apache_beam/version.py
index 264029d..435125c 100644
--- a/sdks/python/apache_beam/version.py
+++ b/sdks/python/apache_beam/version.py
@@ -17,4 +17,4 @@
"""Apache Beam SDK version information and utilities."""
-__version__ = '2.34.0.dev'
+__version__ = '2.35.0.dev'
diff --git a/sdks/python/container/base_image_requirements.txt b/sdks/python/container/base_image_requirements.txt
index 6c158cd..b98d530 100644
--- a/sdks/python/container/base_image_requirements.txt
+++ b/sdks/python/container/base_image_requirements.txt
@@ -29,11 +29,11 @@
crcmod==1.7
dill==0.3.1.1
future==0.18.2
-grpcio==1.34.0
+grpcio==1.40.0
hdfs==2.5.8
httplib2==0.19.1
oauth2client==4.1.3
-protobuf==3.12.2
+protobuf==3.17.3
pyarrow==3.0.0
pydot==1.4.1
pymongo==3.10.1
@@ -78,10 +78,10 @@
scipy==1.4.1
scikit-learn==0.24.1
pandas==1.1.5 ; python_version<"3.7"
-pandas==1.2.4 ; python_version>="3.7"
+pandas==1.3.3 ; python_version>="3.7"
protorpc==0.12.0
python-gflags==3.1.2
-tensorflow==2.5.1
+tensorflow==2.6.0
nltk==3.5.0
# Packages needed for testing.
diff --git a/sdks/python/container/license_scripts/dep_urls_py.yaml b/sdks/python/container/license_scripts/dep_urls_py.yaml
index cdc45a8..ce2c7f6 100644
--- a/sdks/python/container/license_scripts/dep_urls_py.yaml
+++ b/sdks/python/container/license_scripts/dep_urls_py.yaml
@@ -48,6 +48,8 @@
license: "https://raw.githubusercontent.com/chardet/chardet/master/LICENSE"
certifi:
license: "https://raw.githubusercontent.com/certifi/python-certifi/master/LICENSE"
+ clang:
+ license: "https://raw.githubusercontent.com/llvm/llvm-project/main/clang/LICENSE.TXT"
cython:
license: "https://raw.githubusercontent.com/cython/cython/master/LICENSE.txt"
dataclasses:
diff --git a/sdks/python/gen_protos.py b/sdks/python/gen_protos.py
index 2fa3cac..4e63078 100644
--- a/sdks/python/gen_protos.py
+++ b/sdks/python/gen_protos.py
@@ -60,8 +60,8 @@
This is executed at build time rather than dynamically on import to ensure
that it is compatible with static type checkers like mypy.
"""
- import google.protobuf.message as message
import google.protobuf.pyext._message as pyext_message
+ from google.protobuf import message
class Context(object):
INDENT = ' '
@@ -307,6 +307,7 @@
p.join()
if p.exitcode:
raise ValueError("Proto generation failed (see log for details).")
+
else:
log.info('Regenerating Python proto definitions (%s).' % regenerate)
builtin_protos = pkg_resources.resource_filename('grpc_tools', '_proto')
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 514f4e7..7ce7056 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -148,7 +148,7 @@
'pymongo>=3.8.0,<4.0.0',
'oauth2client>=2.0.1,<5',
'protobuf>=3.12.2,<4',
- 'pyarrow>=0.15.1,<5.0.0',
+ 'pyarrow>=0.15.1,<6.0.0',
'pydot>=1.2.0,<2',
'python-dateutil>=2.8.0,<3',
'pytz>=2018.3',
@@ -165,7 +165,7 @@
REQUIRED_TEST_PACKAGES = [
'freezegun>=0.3.12',
'mock>=1.0.1,<3.0.0',
- 'pandas>=1.0,<1.3.0',
+ 'pandas<2.0.0',
'parameterized>=0.7.1,<0.8.0',
'pyhamcrest>=1.9,!=1.10.0,<2.0.0',
'pyyaml>=3.12,<6.0.0',
@@ -190,6 +190,7 @@
'google-cloud-pubsub>=0.39.0,<2',
# GCP packages required by tests
'google-cloud-bigquery>=1.6.0,<3',
+ 'google-cloud-bigquery-storage>=2.6.3',
'google-cloud-core>=0.28.1,<2',
'google-cloud-bigtable>=0.31.1,<2',
'google-cloud-spanner>=1.13.0,<2',
@@ -206,6 +207,7 @@
'facets-overview>=1.0.0,<2',
'ipython>=7,<8',
'ipykernel>=5.2.0,<6',
+ 'ipywidgets>=7.6.5,<8',
# Skip version 6.1.13 due to
# https://github.com/jupyter/jupyter_client/issues/637
'jupyter-client>=6.1.11,<6.1.13',
@@ -215,7 +217,7 @@
INTERACTIVE_BEAM_TEST = [
# notebok utils
'nbformat>=5.0.5,<6',
- 'nbconvert>=5.6.1,<6',
+ 'nbconvert>=6.2.0,<7',
# headless chrome based integration tests
'selenium>=3.141.0,<4',
'needle>=0.5.0,<1',
@@ -242,7 +244,7 @@
class cmd(original_cmd, object):
def run(self):
gen_protos.generate_proto_files()
- super(cmd, self).run()
+ super().run()
return cmd
except ImportError:
@@ -258,76 +260,79 @@
'Python %s.%s. You may encounter bugs or missing features.' %
(sys.version_info.major, sys.version_info.minor))
-setuptools.setup(
- name=PACKAGE_NAME,
- version=PACKAGE_VERSION,
- description=PACKAGE_DESCRIPTION,
- long_description=PACKAGE_LONG_DESCRIPTION,
- 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',
- 'portability/api/*.yaml'
- ]
- },
- ext_modules=cythonize([
- # Make sure to use language_level=3 cython directive in files below.
- '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/counters.py',
- 'apache_beam/utils/windowed_value.py',
- ]),
- install_requires=REQUIRED_PACKAGES,
- python_requires=python_requires,
- # BEAM-8840: Do NOT use tests_require or setup_requires.
- extras_require={
- 'docs': ['Sphinx>=1.5.2,<2.0'],
- 'test': REQUIRED_TEST_PACKAGES,
- 'gcp': GCP_REQUIREMENTS,
- 'interactive': INTERACTIVE_BEAM,
- 'interactive_test': INTERACTIVE_BEAM_TEST,
- 'aws': AWS_REQUIREMENTS,
- 'azure': AZURE_REQUIREMENTS
- },
- 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.6',
- 'Programming Language :: Python :: 3.7',
- 'Programming Language :: Python :: 3.8',
- # When updating vesion 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={
- 'build_py': generate_protos_first(build_py),
- 'develop': generate_protos_first(develop),
- 'egg_info': generate_protos_first(egg_info),
- 'test': generate_protos_first(test),
- 'mypy': generate_protos_first(mypy),
- },
-)
+
+if __name__ == '__main__':
+ setuptools.setup(
+ name=PACKAGE_NAME,
+ version=PACKAGE_VERSION,
+ description=PACKAGE_DESCRIPTION,
+ long_description=PACKAGE_LONG_DESCRIPTION,
+ 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',
+ 'portability/api/*.yaml'
+ ]
+ },
+ ext_modules=cythonize([
+ # Make sure to use language_level=3 cython directive in files below.
+ '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/counters.py',
+ 'apache_beam/utils/windowed_value.py',
+ ]),
+ install_requires=REQUIRED_PACKAGES,
+ python_requires=python_requires,
+ # BEAM-8840: Do NOT use tests_require or setup_requires.
+ extras_require={
+ 'docs': ['Sphinx>=1.5.2,<2.0'],
+ 'test': REQUIRED_TEST_PACKAGES,
+ 'gcp': GCP_REQUIREMENTS,
+ 'interactive': INTERACTIVE_BEAM,
+ 'interactive_test': INTERACTIVE_BEAM_TEST,
+ 'aws': AWS_REQUIREMENTS,
+ 'azure': AZURE_REQUIREMENTS,
+ 'dataframe': ['pandas>=1.0,<1.4']
+ },
+ 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.6',
+ 'Programming Language :: Python :: 3.7',
+ 'Programming Language :: Python :: 3.8',
+ # When updating vesion 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={
+ 'build_py': generate_protos_first(build_py),
+ 'develop': generate_protos_first(develop),
+ 'egg_info': generate_protos_first(egg_info),
+ 'test': generate_protos_first(test),
+ 'mypy': generate_protos_first(mypy),
+ },
+ )
diff --git a/sdks/python/test-suites/tox/common.gradle b/sdks/python/test-suites/tox/common.gradle
index af74725..285a8a7 100644
--- a/sdks/python/test-suites/tox/common.gradle
+++ b/sdks/python/test-suites/tox/common.gradle
@@ -33,9 +33,7 @@
project.task("preCommitPy${pythonVersionSuffix}") {
// Generates coverage reports only once, in Py38, to remove duplicated work
if (pythonVersionSuffix.equals('38')) {
- dependsOn = ["testPy38CloudCoverage", "testPy38Cython",
- "testPy38pyarrow-0", "testPy38pyarrow-1", "testPy38pyarrow-2",
- "testPy38pyarrow-3", "testPy38pyarrow-4"]
+ dependsOn = ["testPy38CloudCoverage", "testPy38Cython"]
} else {
dependsOn = ["testPy${pythonVersionSuffix}Cloud", "testPy${pythonVersionSuffix}Cython"]
}
diff --git a/sdks/python/test-suites/tox/py38/build.gradle b/sdks/python/test-suites/tox/py38/build.gradle
index 8497d57..6e54a86 100644
--- a/sdks/python/test-suites/tox/py38/build.gradle
+++ b/sdks/python/test-suites/tox/py38/build.gradle
@@ -34,16 +34,12 @@
// TODO(BEAM-8954): Remove this once tox uses isolated builds.
testPy38Cython.mustRunAfter testPython38, testPy38CloudCoverage
-toxTask "testPy38pyarrow-0", "py38-pyarrow-0"
-toxTask "testPy38pyarrow-1", "py38-pyarrow-1"
-toxTask "testPy38pyarrow-2", "py38-pyarrow-2"
-toxTask "testPy38pyarrow-3", "py38-pyarrow-3"
-toxTask "testPy38pyarrow-4", "py38-pyarrow-4"
-test.dependsOn "testPy38pyarrow-0"
-test.dependsOn "testPy38pyarrow-1"
-test.dependsOn "testPy38pyarrow-2"
-test.dependsOn "testPy38pyarrow-3"
-test.dependsOn "testPy38pyarrow-4"
+(0..5).each {version ->
+ // Create a test task for each major version of pyarrow
+ toxTask "testPy38pyarrow-$version", "py38-pyarrow-$version"
+ test.dependsOn "testPy38pyarrow-$version"
+ preCommitPy38.dependsOn "testPy38pyarrow-$version"
+}
toxTask "whitespacelint", "whitespacelint"
diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini
index a5ed2db..ca3e0b8 100644
--- a/sdks/python/tox.ini
+++ b/sdks/python/tox.ini
@@ -30,7 +30,7 @@
# allow apps that support color to use it.
passenv=TERM
# Set [] options for pip installation of apache-beam tarball.
-extras = test
+extras = test,dataframe
# Don't warn that these commands aren't installed.
whitelist_externals =
false
@@ -88,7 +88,7 @@
{toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
[testenv:py{36,37,38}-cloud]
-extras = test,gcp,interactive,aws,azure
+extras = test,gcp,interactive,dataframe,aws,azure
commands =
{toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
@@ -98,7 +98,7 @@
codecov
pytest-cov==2.9.0
passenv = GIT_* BUILD_* ghprb* CHANGE_ID BRANCH_NAME JENKINS_* CODECOV_*
-extras = test,gcp,interactive,aws
+extras = test,gcp,interactive,dataframe,aws
commands =
-rm .coverage
{toxinidir}/scripts/run_pytest.sh {envname} "{posargs}" "--cov-report=xml --cov=. --cov-append"
@@ -110,9 +110,9 @@
# keep the version of pylint in sync with the 'rev' in .pre-commit-config.yaml
deps =
-r build-requirements.txt
- astroid<2.4,>=2.3.0
+ astroid<2.9,>=2.8.0
pycodestyle==2.3.1
- pylint==2.4.3
+ pylint==2.11.1
isort==4.2.15
flake8==3.5.0
commands =
@@ -138,7 +138,7 @@
python setup.py mypy
[testenv:py38-docs]
-extras = test,gcp,docs,interactive
+extras = test,gcp,docs,interactive,dataframe
deps =
Sphinx==1.8.5
sphinx_rtd_theme==0.4.3
@@ -197,7 +197,7 @@
# pulls in the latest docutils. Uncomment this line once botocore does not
# conflict with Sphinx:
# extras = docs,test,gcp,aws,interactive,interactive_test
-extras = test,gcp,aws,interactive,interactive_test
+extras = test,gcp,aws,dataframe,interactive,interactive_test
passenv = WORKSPACE
commands =
time {toxinidir}/scripts/run_dependency_check.sh
@@ -233,7 +233,7 @@
commands =
{toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/spark_runner_test.py {posargs}
-[testenv:py{36,37,38}-pyarrow-{0,1,2,3,4}]
+[testenv:py{36,37,38}-pyarrow-{0,1,2,3,4,5}]
deps =
0: pyarrow>=0.15.1,<0.18.0
1: pyarrow>=1,<2
@@ -243,7 +243,9 @@
{0,1,2}: numpy<1.20.0
3: pyarrow>=3,<4
4: pyarrow>=4,<5
+ 5: pyarrow>=5,<6
commands =
# Log pyarrow and numpy version for debugging
/bin/sh -c "pip freeze | grep -E '(pyarrow|numpy)'"
- {toxinidir}/scripts/run_pytest.sh {envname} '-m uses_pyarrow'
+ # TODO(BEAM-12985): Running run_pytest.sh with -m causes us to run all the tests twice
+ {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}" '-m uses_pyarrow'
diff --git a/website/www/site/content/en/community/contact-us.md b/website/www/site/content/en/community/contact-us.md
index c071818..3af9984 100644
--- a/website/www/site/content/en/community/contact-us.md
+++ b/website/www/site/content/en/community/contact-us.md
@@ -32,13 +32,13 @@
## Mailing list, what are they and how they work
-The official communication channels for Apache projects are their mailing lists, and Apache Beam has two main lists: [user@beam.apache.org](user@beam.apache.org) and [dev@beam.apache.org](dev@beam.apache.org). The topics for each of them can be seen in the section above.
+The official communication channels for Apache projects are their mailing lists, and Apache Beam has two main lists: [user@beam.apache.org](https://lists.apache.org/list.html?user@beam.apache.org) and [dev@beam.apache.org](https://lists.apache.org/list.html?dev@beam.apache.org). The topics for each of them can be seen in the section above.
-### Subsribe and Unsubscribe:
+### Subscribe and Unsubscribe:
-Prior to sending emails to these lists, you need to subscribe. To subscribe, send a blank email to [user-subscribe@beam.apache.org](user-subscribe@beam.apache.org) or [dev-subscribe@beam.apache.org](dev-subscribe@beam.apache.org) depending on the list you want to write to.
+Prior to sending emails to these lists, you need to subscribe. To subscribe, send a blank email to [user-subscribe@beam.apache.org](mailto:user-subscribe@beam.apache.org) or [dev-subscribe@beam.apache.org](mailto:dev-subscribe@beam.apache.org) depending on the list you want to write to.
-To unsubscribe, send a blank email to [user-unsubscribe@beam.apache.org](user-unsubscribe@beam.apache.org) or [dev-unsubscribe@beam.apache.org](dev-unsubscribe@beam.apache.org) depending on the list you want to unsubscribe.
+To unsubscribe, send a blank email to [user-unsubscribe@beam.apache.org](mailto:user-unsubscribe@beam.apache.org) or [dev-unsubscribe@beam.apache.org](mailto:dev-unsubscribe@beam.apache.org) depending on the list you want to unsubscribe.
### Useful Tips for Sending Emails
diff --git a/website/www/site/content/en/documentation/dsls/dataframes/overview.md b/website/www/site/content/en/documentation/dsls/dataframes/overview.md
index 8d2c0a8..0620da4 100644
--- a/website/www/site/content/en/documentation/dsls/dataframes/overview.md
+++ b/website/www/site/content/en/documentation/dsls/dataframes/overview.md
@@ -30,9 +30,18 @@
If you’re new to pandas DataFrames, you can get started by reading [10 minutes to pandas](https://pandas.pydata.org/pandas-docs/stable/user_guide/10min.html), which shows you how to import and work with the `pandas` package. pandas is an open-source Python library for data manipulation and analysis. It provides data structures that simplify working with relational or labeled data. One of these data structures is the [DataFrame](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html), which contains two-dimensional tabular data and provides labeled rows and columns for the data.
-## Using DataFrames
+## Pre-requisites
-To use Beam DataFrames, you need to install Apache Beam version 2.26.0 or higher (for complete setup instructions, see the [Apache Beam Python SDK Quickstart](https://beam.apache.org/get-started/quickstart-py/)) and pandas version 1.0 or higher. You can use DataFrames as shown in the following example, which reads New York City taxi data from a CSV file, performs a grouped aggregation, and writes the output back to CSV:
+To use Beam DataFrames, you need to install Beam python version 2.26.0 or higher (for complete setup instructions, see the [Apache Beam Python SDK Quickstart](https://beam.apache.org/get-started/quickstart-py/)) and a supported `pandas` version. In Beam 2.34.0 and newer the easiest way to do this is with the "dataframe" extra:
+
+```
+pip install apache_beam[dataframe]
+```
+
+Note that the _same_ `pandas` version should be installed on workers when executing DataFrame API pipelines on distributed runners. Reference [`base_image_requirements.txt`](https://github.com/apache/beam/blob/master/sdks/python/container/base_image_requirements.txt) for the Beam release you are using to see what version of `pandas` will be used by default on workers.
+
+## Using DataFrames
+You can use DataFrames as shown in the following example, which reads New York City taxi data from a CSV file, performs a grouped aggregation, and writes the output back to CSV:
{{< highlight py >}}
from apache_beam.dataframe.io import read_csv
diff --git a/website/www/site/content/en/documentation/glossary.md b/website/www/site/content/en/documentation/glossary.md
index acca8a3..a4f6d0c 100644
--- a/website/www/site/content/en/documentation/glossary.md
+++ b/website/www/site/content/en/documentation/glossary.md
@@ -19,16 +19,16 @@
## Aggregation
-A transform pattern for computing a value from multiple input elements. Aggregation is similar to the reduce operation in the [MapReduce](https://en.wikipedia.org/wiki/MapReduce) model. Aggregation transforms include Count (computes the count of all elements in the aggregation), Max (computes the maximum element in the aggregation), and Sum (computes the sum of all elements in the aggregation).
+A transform pattern for computing a value from multiple input elements. Aggregation is similar to the reduce operation in the [MapReduce](https://en.wikipedia.org/wiki/MapReduce) model. Aggregation transforms include Combine (applies a user-defined function to all elements in the aggregation), Count (computes the count of all elements in the aggregation), Max (computes the maximum element in the aggregation), and Sum (computes the sum of all elements in the aggregation).
-For a complete list of aggregation transforms, see:
+For a list of built-in aggregation transforms, see:
* [Java Transform catalog](/documentation/transforms/java/overview/#aggregation)
* [Python Transform catalog](/documentation/transforms/python/overview/#aggregation)
## Apply
-A method for invoking a transform on a PCollection. Each transform in the Beam SDKs has a generic `apply` method (or pipe operator `|`). Invoking multiple Beam transforms is similar to method chaining, but with a difference: You apply the transform to the input PCollection, passing the transform itself as an argument, and the operation returns the output PCollection. Because of Beam’s deferred execution model, applying a transform does not immediately execute that transform.
+A method for invoking a transform on an input PCollection (or set of PCollections) to produce one or more output PCollections. The `apply` method is attached to the PCollection (or value). Invoking multiple Beam transforms is similar to method chaining, but with a difference: You apply the transform to the input PCollection, passing the transform itself as an argument, and the operation returns the output PCollection. Because of Beam’s deferred execution model, applying a transform does not immediately execute that transform.
To learn more, see:
@@ -44,7 +44,7 @@
## Bounded data
-A dataset of a known, fixed size. A PCollection can be bounded or unbounded, depending on the source of the data that it represents. Reading from a batch data source, such as a file or a database, creates a bounded PCollection. Beam also supports reading a bounded amount of data from an unbounded source.
+A dataset of a known, fixed size (alternatively, a dataset that is not growing over time). A PCollection can be bounded or unbounded, depending on the source of the data that it represents. Reading from a batch data source, such as a file or a database, creates a bounded PCollection. Beam also supports reading a bounded amount of data from an unbounded source.
To learn more, see:
@@ -52,7 +52,7 @@
## Bundle
-The processing unit for elements in a PCollection. Instead of processing all elements in a PCollection simultaneously, Beam processes the elements in bundles. The runner handles the division of the collection into bundles, and in doing so it may optimize the bundle size for the use case. For example, a streaming runner might process smaller bundles than a batch runner.
+The processing and commit/retry unit for elements in a PCollection. Instead of processing all elements in a PCollection simultaneously, Beam processes the elements in bundles. The runner handles the division of the collection into bundles, and in doing so it may optimize the bundle size for the use case. For example, a streaming runner might process smaller bundles than a batch runner.
To learn more, see:
@@ -94,7 +94,7 @@
## Composite transform
-A PTransform that expands into many PTransforms. Composite transforms have a nested structure, in which a complex transform applies one or more simpler transforms. These simpler transforms could be existing Beam operations like ParDo, Combine, or GroupByKey, or they could be other composite transforms. Nesting multiple transforms inside a single composite transform can make your pipeline more modular and easier to understand.
+A PTransform that expands into many PTransforms. Composite transforms have a nested structure, in which a complex transform applies one or more simpler transforms. These simpler transforms could be existing Beam operations like ParDo, Combine, or GroupByKey, or they could be other composite transforms. Nesting multiple transforms inside a single composite transform can make your pipeline more modular and easier to understand. Many of the built-in transforms are composite transforms.
To learn more, see:
@@ -118,7 +118,7 @@
## Deferred execution
-A feature of the Beam execution model. Beam operations are deferred, meaning that the result of a given operation may not be available for control flow. Deferred execution allows the Beam API to support parallel processing of data.
+A feature of the Beam execution model. Beam operations are deferred, meaning that the result of a given operation may not be available for control flow. Deferred execution allows the Beam API to support parallel processing of data and perform pipeline-level optimizations.
## Distribution (metric)
@@ -130,7 +130,7 @@
## DoFn
-A function object used by ParDo (or some other transform) to process the elements of a PCollection. A DoFn is a user-defined function, meaning that it contains custom code that defines a data processing task in your pipeline. The Beam system invokes a DoFn one or more times to process some arbitrary bundle of elements, but Beam doesn’t guarantee an exact number of invocations.
+A function object used by ParDo (or some other transform) to process the elements of a PCollection, often producing elements for an output PCollection. A DoFn is a user-defined function, meaning that it contains custom code that defines a data processing task in your pipeline. The Beam system invokes a DoFn one or more times to process some arbitrary bundle of elements, but Beam doesn’t guarantee an exact number of invocations.
To learn more, see:
@@ -167,7 +167,7 @@
## Event time
-The time a data event occurs, determined by a timestamp on an element. This is in contrast to processing time, which is when an element is processed in a pipeline. An event could be, for example, a user interaction or a write to an error log. There’s no guarantee that events will appear in a pipeline in order of event time.
+The time a data event occurs, determined by a timestamp on an element. This is in contrast to processing time, which is when an element is processed in a pipeline. An event could be, for example, a user interaction or a write to an error log. There’s no guarantee that events will appear in a pipeline in order of event time, but windowing and timers let you reason correctly about event time.
To learn more, see:
@@ -176,7 +176,7 @@
## Expansion Service
-A service that enables a pipeline to apply (expand) cross-language transforms defined in other SDKs. For example, by connecting to a Java expansion service, the Python SDK can apply transforms implemented in Java. Currently SDKs define expansion services as local processes, but in the future Beam may support long-running expansion services. The development of expansion services is part of the ongoing effort to support multi-language pipelines.
+A service that enables a pipeline to apply (expand) cross-language transforms defined in other SDKs. For example, by connecting to a Java expansion service, the Python SDK can apply transforms implemented in Java. Currently, SDKs typically start up expansion services as local processes, but in the future Beam may support long-running expansion services. The development of expansion services is part of the ongoing effort to support multi-language pipelines.
## Flatten
One of the core PTransforms. Flatten merges multiple PCollections into a single logical PCollection.
@@ -187,9 +187,13 @@
* [Flatten (Java)](/documentation/transforms/java/other/flatten/)
* [Flatten (Python)](/documentation/transforms/python/other/flatten/)
+## Fn API
+
+An interface that lets a runner invoke SDK-specific user-defined functions. The Fn API, together with the Runner API, supports the ability to mix and match SDKs and runners. Used together, the Fn and Runner APIs let new SDKs run on every runner, and let new runners run pipelines from every SDK.
+
## Fusion
-An optimization that Beam runners can apply before running a pipeline. When one transform outputs a PCollection that’s consumed by another transform, or when two or more transforms take the same PCollection as input, a runner may be able to fuse the transforms together into a single processing unit (a *stage* in Dataflow). Fusion can make pipeline execution more efficient by preventing I/O operations.
+An optimization that Beam runners can apply before running a pipeline. When one transform outputs a PCollection that’s consumed by another transform, or when two or more transforms take the same PCollection as input, a runner may be able to fuse the transforms together into a single processing unit (a *stage* in Dataflow). The consuming DoFn processes elements as they are emitted by the producing DoFn, rather than waiting for the entire intermediate PCollection to be computed. Fusion can make pipeline execution more efficient by preventing I/O operations.
## Gauge (metric)
@@ -220,7 +224,7 @@
## Map
-An element-wise PTransform that applies a user-defined function (UDF) to each element in a PCollection. Using Map, you can transform each individual element, but you can't change the number of elements.
+An element-wise PTransform that applies a user-defined function (UDF) to each element in a PCollection. Using Map, you can transform each individual element into a new element, but you can't change the number of elements.
To learn more, see:
@@ -245,7 +249,7 @@
## ParDo
-The lowest-level element-wise PTransform. For each element in an input PCollection, ParDo applies a function and emits zero, one, or multiple elements to an output PCollection. “ParDo” is short for “Parallel Do.” It’s similar to the map operation in a [MapReduce](https://en.wikipedia.org/wiki/MapReduce) algorithm, the `apply` method from a DataFrame, or the `UPDATE` keyword from SQL.
+The lowest-level element-wise PTransform. For each element in an input PCollection, ParDo applies a function and emits zero, one, or multiple elements to an output PCollection. “ParDo” is short for “Parallel Do.” It’s similar to the map operation in a [MapReduce](https://en.wikipedia.org/wiki/MapReduce) algorithm and the reduce operation when following a GroupByKey. ParDo is also comparable to the `apply` method from a DataFrame, or the `UPDATE` keyword from SQL.
To learn more, see:
@@ -255,7 +259,7 @@
## Partition
-An element-wise PTransform that splits a single PCollection into a fixed number of smaller PCollections. Partition requires a user-defined function (UDF) to determine how to split up the elements of the input collection into the resulting output collections. The number of partitions must be determined at graph construction time, meaning that you can’t determine the number of partitions using data calculated by the running pipeline.
+An element-wise PTransform that splits a single PCollection into a fixed number of smaller, disjoint PCollections. Partition requires a user-defined function (UDF) to determine how to split up the elements of the input collection into the resulting output collections. The number of partitions must be determined at graph construction time, meaning that you can’t determine the number of partitions using data calculated by the running pipeline.
To learn more, see:
@@ -273,7 +277,7 @@
## Pipe operator (`|`)
-Delimits a step in a Python pipeline. For example: `[Final Output PCollection] = ([Initial Input PCollection] | [First Transform] | [Second Transform] | [Third Transform])`. The output of each transform is passed from left to right as input to the next transform. The pipe operator in Python is equivalent to the `apply` method in Java (in other words, the pipe applies a transform to a PCollection).
+Delimits a step in a Python pipeline. For example: `[Final Output PCollection] = ([Initial Input PCollection] | [First Transform] | [Second Transform] | [Third Transform])`. The output of each transform is passed from left to right as input to the next transform. The pipe operator in Python is equivalent to the `apply` method in Java (in other words, the pipe applies a transform to a PCollection), and usage is similar to the pipe operator in shell scripts, which lets you pass the output of one program into the input of another.
To learn more, see:
@@ -281,7 +285,7 @@
## Pipeline
-An encapsulation of your entire data processing task, including reading input data from a source, transforming that data, and writing output data to a sink. You can think of a pipeline as a Beam program that uses PTransforms to process PCollections. The transforms in a pipeline can be represented as a directed acyclic graph (DAG). All Beam driver programs must create a pipeline.
+An encapsulation of your entire data processing task, including reading input data from a source, transforming that data, and writing output data to a sink. You can think of a pipeline as a Beam program that uses PTransforms to process PCollections. (Alternatively, you can think of it as a single, executable composite PTransform with no inputs or outputs.) The transforms in a pipeline can be represented as a directed acyclic graph (DAG). All Beam driver programs must create a pipeline.
To learn more, see:
@@ -292,7 +296,7 @@
## Processing time
-The time at which an element is processed at some stage in a pipeline. Processing time is not the same as event time, which is the time at which a data event occurs. Processing time is determined by the clock on the system processing the element. There’s no guarantee that elements will be processed in order of event time.
+The real-world time at which an element is processed at some stage in a pipeline. Processing time is not the same as event time, which is the time at which a data event occurs. Processing time is determined by the clock on the system processing the element. There’s no guarantee that elements will be processed in order of event time.
To learn more, see:
@@ -327,7 +331,7 @@
## Schema
-A language-independent type definition for a PCollection. The schema for a PCollection defines elements of that PCollection as an ordered list of named fields. Each field has a name, a type, and possibly a set of user options. Schemas provide a way to reason about types across different programming-language APIs.
+A language-independent type definition for the elements of a PCollection. The schema for a PCollection defines elements of that PCollection as an ordered list of named fields. Each field has a name, a type, and possibly a set of user options. Schemas provide a way to reason about types across different programming-language APIs. They also let you describe data transformations more succinctly and at a higher level.
To learn more, see:
@@ -336,7 +340,7 @@
## Session
-A time interval for grouping data events. A session is defined by some minimum gap duration between events. For example, a data stream representing user mouse activity may have periods with high concentrations of clicks followed by periods of inactivity. A session can represent such a pattern of activity followed by inactivity.
+A time interval for grouping data events. A session is defined by some minimum gap duration between events. For example, a data stream representing user mouse activity may have periods with high concentrations of clicks followed by periods of inactivity. A session can represent such a pattern of activity delimited by inactivity.
To learn more, see:
@@ -345,7 +349,7 @@
## Side input
-Additional input to a PTransform. Side input is input that you provide in addition to the main input PCollection. A DoFn can access side input each time it processes an element in the PCollection. Side inputs are useful if your transform needs to inject additional data at runtime.
+Additional input to a PTransform that is provided in its entirety, rather than element-by-element. Side input is input that you provide in addition to the main input PCollection. A DoFn can access side input each time it processes an element in the PCollection.
To learn more, see:
@@ -380,9 +384,13 @@
* [Splittable DoFns](/documentation/programming-guide/#splittable-dofns)
* [Splittable DoFn in Apache Beam is Ready to Use](/blog/splittable-do-fn-is-available/)
+## Stage
+
+The unit of fused transforms in a pipeline. Runners can perform fusion optimization to make pipeline execution more efficient. In Dataflow, the pipeline is conceptualized as a graph of fused stages.
+
## State
-Persistent values that a PTransform can access. The state API lets you augment element-wise operations (for example, ParDo or Map) with mutable state. Using the state API, you can read from, and write to, state as you process each element of a PCollection. You can use the state API together with the timer API to create processing tasks that give you fine-grained control over the workflow.
+Persistent values that a PTransform can access. The state API lets you augment element-wise operations (for example, ParDo or Map) with mutable state. Using the state API, you can read from, and write to, state as you process each element of a PCollection. You can use the state API together with the timer API to create processing tasks that give you fine-grained control over the workflow. State is always local to a key and window.
To learn more, see:
@@ -410,7 +418,7 @@
## Timestamp
-A point in time associated with an element in a PCollection and used to assign a window to the element. The source that creates the PCollection assigns each element an initial timestamp, often corresponding to when the element was read or added. But you can also manually assign timestamps. This can be useful if elements have an inherent timestamp, but the timestamp is somewhere in the structure of the element itself (for example, a time field in a server log entry).
+A point in event time associated with an element in a PCollection and used to assign a window to the element. The source that creates the PCollection assigns each element an initial timestamp, often corresponding to when the element was read or added. But you can also manually assign timestamps. This can be useful if elements have an inherent timestamp, but the timestamp is somewhere in the structure of the element itself (for example, a time field in a server log entry).
To learn more, see:
@@ -431,7 +439,7 @@
## Unbounded data
-A dataset of unlimited size. A PCollection can be bounded or unbounded, depending on the source of the data that it represents. Reading from a streaming or continuously-updating data source, such as Pub/Sub or Kafka, typically creates an unbounded PCollection.
+A dataset that grows over time, with elements processed as they arrive. A PCollection can be bounded or unbounded, depending on the source of the data that it represents. Reading from a streaming or continuously-updating data source, such as Pub/Sub or Kafka, typically creates an unbounded PCollection.
To learn more, see:
@@ -449,7 +457,7 @@
## Watermark
-The point in event time when all data in a window can be expected to have arrived in the pipeline. Watermarks provide a way to estimate the completeness of input data. Every PCollection has an associated watermark. Once the watermark progresses past the end of a window, any element that arrives with a timestamp in that window is considered late data.
+An estimate on the lower bound of the timestamps that will be seen (in the future) at this point of the pipeline. Watermarks provide a way to estimate the completeness of input data. Every PCollection has an associated watermark. Once the watermark progresses past the end of a window, any element that arrives with a timestamp in that window is considered late data.
To learn more, see:
@@ -465,7 +473,7 @@
## Worker
-A container, process, or virtual machine (VM) that handles some part of the parallel processing of a pipeline. The Beam model doesn’t support synchronizing shared state across worker machines. Instead, each worker node has its own independent copy of state. A Beam runner might serialize elements between machines for communication purposes and for other reasons such as persistence.
+A container, process, or virtual machine (VM) that handles some part of the parallel processing of a pipeline. Each worker node has its own independent copy of state. A Beam runner might serialize elements between machines for communication purposes and for other reasons such as persistence.
To learn more, see:
diff --git a/website/www/site/content/en/documentation/patterns/cross-language.md b/website/www/site/content/en/documentation/patterns/cross-language.md
index 8cdd6a3..b67db17 100644
--- a/website/www/site/content/en/documentation/patterns/cross-language.md
+++ b/website/www/site/content/en/documentation/patterns/cross-language.md
@@ -95,7 +95,7 @@
@ptransform.PTransform.register_urn(URN, None)
class PythonTransform(ptransform.PTransform):
def __init__(self):
- super(PythonTransform, self).__init__()
+ super().__init__()
def expand(self, pcoll):
return (pcoll
diff --git a/website/www/site/content/en/documentation/programming-guide.md b/website/www/site/content/en/documentation/programming-guide.md
index 61b6b97e..229b215 100644
--- a/website/www/site/content/en/documentation/programming-guide.md
+++ b/website/www/site/content/en/documentation/programming-guide.md
@@ -1462,7 +1462,8 @@
{{< paragraph class="language-go" >}}
If your `PCollection` uses any non-global windowing function, the Beam Go SDK
-behaves the same way as with global windowing.
+behaves the same way as with global windowing. Windows that are empty in the input
+ `PCollection` will likewise be empty in the output collection.
{{< /paragraph >}}
##### 4.2.4.6. Combining values in a keyed PCollection {#combining-values-in-a-keyed-pcollection}
@@ -1651,7 +1652,7 @@
<span class="language-java language-py">
-> **Note:** These requirements apply to subclasses of `DoFn`</span> (a function object
+> **Note:** These requirements apply to subclasses of `DoFn`(a function object
> used with the [ParDo](#pardo) transform), `CombineFn` (a function object used
> with the [Combine](#combine) transform), and `WindowFn` (a function object
> used with the [Window](#windowing) transform).
@@ -1660,7 +1661,7 @@
<span class="language-go">
-> **Note:** These requirements apply to `DoFn`s</span> (a function object
+> **Note:** These requirements apply to `DoFn`s (a function object
> used with the [ParDo](#pardo) transform), `CombineFn`s (a function object used
> with the [Combine](#combine) transform), and `WindowFn`s (a function object
> used with the [Window](#windowing) transform).
@@ -1834,7 +1835,7 @@
### 4.5. Additional outputs {#additional-outputs}
-{{< paragraph class="language-java language-python" >}}
+{{< paragraph class="language-java language-py" >}}
While `ParDo` always produces a main output `PCollection` (as the return value
from `apply`), you can also have your `ParDo` produce any number of additional
output `PCollection`s. If you choose to have multiple outputs, your `ParDo`
@@ -2666,6 +2667,11 @@
In Python you can use the following set of classes to represent the purchase schema. Beam will automatically infer the correct schema based on the members of the class.
{{< /paragraph >}}
+{{< paragraph class="language-go" >}}
+In Go, schema encoding is used by default for struct types, with Exported fields becoming part of the schema.
+Beam will automatically infer the schema based on the fields and field tags of the struct, and their order.
+{{< /paragraph >}}
+
{{< highlight java >}}
@DefaultSchema(JavaBeanSchema.class)
public class Purchase {
@@ -2731,6 +2737,10 @@
purchase_amount: float
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/06schemas.go" schema_define >}}
+{{< /highlight >}}
+
{{< paragraph class="language-java" >}}
Using JavaBean classes as above is one way to map a schema to Java classes. However multiple Java classes might have
the same schema, in which case the different Java types can often be used interchangeably. Beam will add implicit
@@ -2887,6 +2897,13 @@
In Java, a logical type is specified as a subclass of the `LogicalType` class. A custom Java class can be specified to represent the logical type and conversion functions must be supplied to convert back and forth between this Java class and the underlying Schema type representation. For example, the logical type representing nanosecond timestamp might be implemented as follows
{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+In Go, a logical type is specified with a custom implementation of the `beam.SchemaProvider` interface.
+For example, the logical type provider representing nanosecond timestamps
+might be implemented as follows
+{{< /paragraph >}}
+
{{< highlight java >}}
// A Logical type using java.time.Instant to represent the logical type.
public class TimestampNanos implements LogicalType<Instant, Row> {
@@ -2909,11 +2926,39 @@
}
{{< /highlight >}}
+{{< highlight go >}}
+// Define a logical provider like so:
+{{< code_sample "sdks/go/examples/snippets/06schemas.go" schema_logical_provider >}}
+
+// Register it like so:
+{{< code_sample "sdks/go/examples/snippets/06schemas.go" schema_logical_register >}}
+{{< /highlight >}}
+
#### 6.4.2. Useful logical types {#built-in-logical-types}
+{{< paragraph class="language-py" >}}
+Currently the Python SDK provides minimal convenience logical types,
+other than to handle `MicrosInstant`.
+{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+Currently the Go SDK provides minimal convenience logical types,
+other than to handle additional integer primitives, and `time.Time`.
+{{< /paragraph >}}
+
##### **EnumerationType**
+{{< paragraph class="language-py" >}}
+This convenience builder doesn't yet exist for the Python SDK.
+{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+This convenience builder doesn't yet exist for the Go SDK.
+{{< /paragraph >}}
+
+{{< paragraph class="language-java" >}}
This logical type allows creating an enumeration type consisting of a set of named constants.
+{{< /paragraph >}}
{{< highlight java >}}
Schema schema = Schema.builder()
@@ -2922,8 +2967,10 @@
.build();
{{< /highlight >}}
+{{< paragraph class="language-java" >}}
The value of this field is stored in the row as an INT32 type, however the logical type defines a value type that lets
you access the enumeration either as a string or a value. For example:
+{{< /paragraph >}}
{{< highlight java >}}
EnumerationType.Value enumValue = enumType.valueOf("RED");
@@ -2931,18 +2978,32 @@
enumValue.toString(); // Returns "RED", the string value of the constant
{{< /highlight >}}
+{{< paragraph class="language-java" >}}
Given a row object with an enumeration field, you can also extract the field as the enumeration value.
+{{< /paragraph >}}
{{< highlight java >}}
EnumerationType.Value enumValue = row.getLogicalTypeValue("color", EnumerationType.Value.class);
{{< /highlight >}}
+{{< paragraph class="language-java" >}}
Automatic schema inference from Java POJOs and JavaBeans automatically converts Java enums to EnumerationType logical
types.
+{{< /paragraph >}}
##### **OneOfType**
+{{< paragraph class="language-py" >}}
+This convenience builder doesn't yet exist for the Python SDK.
+{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+This convenience builder doesn't yet exist for the Go SDK.
+{{< /paragraph >}}
+
+{{< paragraph class="language-java" >}}
OneOfType allows creating a disjoint union type over a set of schema fields. For example:
+{{< /paragraph >}}
{{< highlight java >}}
Schema schema = Schema.builder()
@@ -2954,9 +3015,11 @@
.build();
{{< /highlight >}}
+{{< paragraph class="language-java" >}}
The value of this field is stored in the row as another Row type, where all the fields are marked as nullable. The
logical type however defines a Value object that contains an enumeration value indicating which field was set and allows
getting just that field:
+{{< /paragraph >}}
{{< highlight java >}}
// Returns an enumeration indicating all possible case values for the enum.
@@ -2978,17 +3041,28 @@
}
{{< /highlight >}}
+{{< paragraph class="language-java" >}}
In the above example we used the field names in the switch statement for clarity, however the enum integer values could
also be used.
+{{< /paragraph >}}
### 6.5. Creating Schemas {#creating-schemas}
-In order to take advantage of schemas, your `PCollection`s must have a schema attached to it. Often, the source itself will attach a schema to the PCollection. For example, when using `AvroIO` to read Avro files, the source can automatically infer a Beam schema from the Avro schema and attach that to the Beam `PCollection`. However not all sources produce schemas. In addition, often Beam pipelines have intermediate stages and types, and those also can benefit from the expressiveness of schemas.
+In order to take advantage of schemas, your `PCollection`s must have a schema attached to it.
+Often, the source itself will attach a schema to the PCollection.
+For example, when using `AvroIO` to read Avro files, the source can automatically infer a Beam schema from the Avro schema and attach that to the Beam `PCollection`.
+However not all sources produce schemas.
+In addition, often Beam pipelines have intermediate stages and types, and those also can benefit from the expressiveness of schemas.
#### 6.5.1. Inferring schemas {#inferring-schemas}
+{{< language-switcher java py go >}}
+
{{< paragraph class="language-java" >}}
-Beam is able to infer schemas from a variety of common Java types. The `@DefaultSchema` annotation can be used to tell Beam to infer schemas from a specific type. The annotation takes a `SchemaProvider` as an argument, and `SchemaProvider` classes are already built in for common Java types. The `SchemaRegistry` can also be invoked programmatically for cases where it is not practical to annotate the Java type itself.
+Beam is able to infer schemas from a variety of common Java types.
+The `@DefaultSchema` annotation can be used to tell Beam to infer schemas from a specific type.
+The annotation takes a `SchemaProvider` as an argument, and `SchemaProvider` classes are already built in for common Java types.
+The `SchemaRegistry` can also be invoked programmatically for cases where it is not practical to annotate the Java type itself.
{{< /paragraph >}}
{{< paragraph class="language-java" >}}
@@ -3191,6 +3265,53 @@
purchase_amount=float(item["purchase_amount"])))
{{< /highlight >}}
+{{< paragraph class="language-go" >}}
+Beam currently only infers schemas for exported fields in Go structs.
+{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+**Structs**
+{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+Beam will automatically infer schemas for all Go structs used
+as PCollection elements, and default to encoding them using
+schema encoding.
+{{< /paragraph >}}
+
+{{< highlight go >}}
+type Transaction struct{
+ Bank string
+ PurchaseAmount float64
+
+ checksum []byte // ignored
+}
+{{< /highlight >}}
+
+{{< paragraph class="language-go" >}}
+Unexported fields are ignored, and cannot be automatically infered as part of the schema.
+Fields of type func, channel, unsafe.Pointer, or uintptr will be ignored by inference.
+Fields of interface types are ignored, unless a schema provider
+is registered for them.
+{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+By default, schema field names will match the exported struct field names.
+In the above example, "Bank" and "PurchaseAmount" are the schema field names.
+A schema field name can be overridden with a struct tag for the field.
+{{< /paragraph >}}
+
+{{< highlight go >}}
+type Transaction struct{
+ Bank string `beam:"bank"`
+ PurchaseAmount float64 `beam:"purchase_amount"`
+}
+{{< /highlight >}}
+
+{{< paragraph class="language-go" >}}
+Overriding schema field names is useful for compatibility cross language transforms,
+as schema fields may have different requirements or restrictions from Go exported fields.
+{{< /paragraph >}}
### 6.6. Using Schema Transforms {#using-schemas}
@@ -3198,6 +3319,10 @@
named fields allows for simple and readable aggregations that reference fields by name, similar to the aggregations in
a SQL expression.
+{{< paragraph class="language-go" >}}
+Beam does not yet support Schema transforms natively in Go. However, it will be implemented with the following behavior.
+{{< /paragraph >}}
+
#### 6.6.1. Field selection syntax
The advantage of schemas is that they allow referencing of element fields by name. Beam provides a selection syntax for
@@ -3723,6 +3848,10 @@
##### **Input conversion**
+{{< paragraph class="language-go" >}}
+Beam does not yet support input conversion in Go.
+{{< /paragraph >}}
+
Since Beam knows the schema of the source `PCollection`, it can automatically convert the elements to any Java type for
which a matching schema is known. For example, using the above-mentioned Transaction schema, say we have the following
`PCollection`:
@@ -3789,6 +3918,8 @@
## 7. Data encoding and type safety {#data-encoding-and-type-safety}
+{{< language-switcher java py go >}}
+
When Beam runners execute your pipeline, they often need to materialize the
intermediate data in your `PCollection`s, which requires converting elements to
and from byte strings. The Beam SDKs use objects called `Coder`s to describe how
@@ -3817,6 +3948,15 @@
package.
{{< /paragraph >}}
+{{< paragraph class="language-go" >}}
+Standard Go types like `int`, `int64` `float64`, `[]byte`, and `string` and more are coded using builtin coders.
+Structs and pointers to structs default using Beam Schema Row encoding.
+However, users can build and register custom coders with `beam.RegisterCoder`.
+You can find available Coder functions in the
+[coder](https://pkg.go.dev/github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coders)
+package.
+{{< /paragraph >}}
+
> Note that coders do not necessarily have a 1:1 relationship with types. For
> example, the Integer type can have multiple valid coders, and input and output
> data can use different Integer coders. A transform might have Integer-typed
@@ -3858,6 +3998,11 @@
type.
{{< /paragraph >}}
+{{< paragraph class="language-go" >}}
+The Beam SDK for Go allows users to register default coder
+implementations with `beam.RegisterCoder`.
+{{< /paragraph >}}
+
{{< paragraph class="language-java" >}}
By default, the Beam SDK for Java automatically infers the `Coder` for the
elements of a `PCollection` produced by a `PTransform` using the type parameter
@@ -3880,12 +4025,24 @@
(in the default pipeline `CoderRegistry`, this is `BytesCoder`).
{{< /paragraph >}}
-> NOTE: If you create your `PCollection` from in-memory data by using the
+{{< paragraph class="language-go" >}}
+By default, the Beam SDK for Go automatically infers the `Coder` for the elements of an output `PCollection` by the output of the transform's function object, such as a `DoFn`.
+ In the case of `ParDo`, for example a `DoFn`
+with the parameters of `v int, emit func(string)` accepts an input element of type `int`
+and produces an output element of type `string`.
+In such a case, the Beam SDK for Go will automatically infer the default `Coder` for the output `PCollection` to be the `string_utf8` coder.
+{{< /paragraph >}}
+
+<span class="language-java">
+
+> **Note:** If you create your `PCollection` from in-memory data by using the
> `Create` transform, you cannot rely on coder inference and default coders.
> `Create` does not have access to any typing information for its arguments, and
> may not be able to infer a coder if the argument list contains a value whose
> exact run-time class doesn't have a default coder registered.
+</span>
+
{{< paragraph class="language-java" >}}
When using `Create`, the simplest way to ensure that you have the correct coder
is by invoking `withCoder` when you apply the `Create` transform.
@@ -4019,8 +4176,13 @@
This allows you to determine (or set) the default Coder for a Python type.
{{< /paragraph >}}
+{{< paragraph class="language-go" >}}
+You can use the `beam.NewCoder` function to determine the default Coder for a Go type.
+{{< /paragraph >}}
+
#### 7.2.2. Setting the default coder for a type {#setting-default-coder}
+{{< paragraph class="language-java language-py" >}}
To set the default Coder for a
<span class="language-java">Java</span><span class="language-py">Python</span>
type for a particular pipeline, you obtain and modify the pipeline's
@@ -4031,11 +4193,23 @@
<span class="language-java">`CoderRegistry.registerCoder`</span>
<span class="language-py">`CoderRegistry.register_coder`</span>
to register a new `Coder` for the target type.
+{{< /paragraph >}}
+{{< paragraph class="language-go" >}}
+To set the default Coder for a Go type you use the function `beam.RegisterCoder` to register a encoder and decoder functions for the target type.
+However, built in types like `int`, `string`, `float64`, etc cannot have their coders overridde.
+{{< /paragraph >}}
+
+{{< paragraph class="language-java language-py" >}}
The following example code demonstrates how to set a default Coder, in this case
`BigEndianIntegerCoder`, for
<span class="language-java">Integer</span><span class="language-py">int</span>
values for a pipeline.
+{{< /paragraph >}}
+
+{{< paragraph class="language-go" >}}
+The following example code demonstrates how to set a custom Coder for `MyCustomType` elements.
+{{< /paragraph >}}
{{< highlight java >}}
PipelineOptions options = PipelineOptionsFactory.create();
@@ -4049,9 +4223,26 @@
apache_beam.coders.registry.register_coder(int, BigEndianIntegerCoder)
{{< /highlight >}}
+{{< highlight go >}}
+type MyCustomType struct{
+ ...
+}
+
+// See documentation on beam.RegisterCoder for other supported coder forms.
+
+func encode(MyCustomType) []byte { ... }
+
+func decode(b []byte) MyCustomType { ... }
+
+func init() {
+ beam.RegisterCoder(reflect.TypeOf((*MyCustomType)(nil)).Elem(), encode, decode)
+}
+{{< /highlight >}}
+
#### 7.2.3. Annotating a custom data type with a default coder {#annotating-custom-type-default-coder}
-{{< paragraph class="language-java" >}}
+<span class="language-java">
+
If your pipeline program defines a custom data type, you can use the
`@DefaultCoder` annotation to specify the coder to use with that type.
By default, Beam will use `SerializableCoder` which uses Java serialization,
@@ -4064,10 +4255,11 @@
For key/value pairs, the correctness of key-based operations
(GroupByKey, Combine) and per-key State depends on having a deterministic
- coder for the key.
+ coder for the key
You can use the `@DefaultCoder` annotation to set a new default as follows:
-{{< /paragraph >}}
+
+</span>
{{< highlight java >}}
@DefaultCoder(AvroCoder.class)
@@ -4094,9 +4286,10 @@
}
{{< /highlight >}}
-{{< paragraph class="language-py" >}}
-The Beam SDK for Python does not support annotating data types with a default
-coder. If you would like to set a default coder, use the method described in the
+{{< paragraph class="language-py language-go" >}}
+The Beam SDK for <span class="language-py">Python</span><span class="language-go">Go</span>
+does not support annotating data types with a default coder.
+If you would like to set a default coder, use the method described in the
previous section, *Setting the default coder for a type*.
{{< /paragraph >}}
@@ -4222,7 +4415,7 @@
* Sliding Time Windows
* Per-Session Windows
* Single Global Window
-* Calendar-based Windows (not supported by the Beam SDK for Python)
+* Calendar-based Windows (not supported by the Beam SDK for Python or Go)
You can also define your own `WindowFn` if you have a more complex need.
@@ -4327,6 +4520,10 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test.py" setting_fixed_windows >}}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/08windowing.go" setting_fixed_windows >}}
+{{< /highlight >}}
+
#### 8.3.2. Sliding time windows {#using-sliding-time-windows}
The following example code shows how to apply `Window` to divide a `PCollection`
@@ -4343,6 +4540,10 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test.py" setting_sliding_windows >}}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/08windowing.go" setting_sliding_windows >}}
+{{< /highlight >}}
+
#### 8.3.3. Session windows {#using-session-windows}
The following example code shows how to apply `Window` to divide a `PCollection`
@@ -4359,6 +4560,10 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test.py" setting_session_windows >}}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/08windowing.go" setting_session_windows >}}
+{{< /highlight >}}
+
Note that the sessions are per-key — each key in the collection will have its
own session groupings depending on the data distribution.
@@ -4378,6 +4583,10 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test.py" setting_global_window >}}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/08windowing.go" setting_global_window >}}
+{{< /highlight >}}
+
### 8.4. Watermarks and late data {#watermarks-and-late-data}
In any data processing system, there is a certain amount of lag between the time
@@ -4445,6 +4654,10 @@
allowed_lateness=Duration(seconds=2*24*60*60)) # 2 days
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/08windowing.go" setting_allowed_lateness >}}
+{{< /highlight >}}
+
When you set `.withAllowedLateness` on a `PCollection`, that allowed lateness
propagates forward to any subsequent `PCollection` derived from the first
`PCollection` you applied allowed lateness to. If you want to change the allowed
@@ -4489,8 +4702,21 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test.py" setting_timestamp >}}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/08windowing.go" setting_timestamp >}}
+
+// Use the DoFn with ParDo as normal.
+{{< code_sample "sdks/go/examples/snippets/08windowing.go" setting_timestamp_pipeline >}}
+{{< /highlight >}}
+
## 9. Triggers {#triggers}
+<span class="language-go">
+
+> **Note:** The Trigger API in the Beam SDK for Go is currently experimental and subject to change.
+
+</span>
+
When collecting and grouping data into windows, Beam uses **triggers** to
determine when to emit the aggregated results of each window (referred to as a
*pane*). If you use Beam's default windowing configuration and [default
@@ -4553,7 +4779,8 @@
timestamps attached to the data elements. The watermark is a global progress
metric, and is Beam's notion of input completeness within your pipeline at any
given point. <span class="language-java">`AfterWatermark.pastEndOfWindow()`</span>
-<span class="language-py">`AfterWatermark`</span> *only* fires when the
+<span class="language-py">`AfterWatermark`</span>
+<span class="language-go">`window.AfterEndOfWindow`</span> *only* fires when the
watermark passes the end of the window.
In addition, you can configure triggers that fire if your pipeline receives data
@@ -4578,6 +4805,10 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test.py" model_early_late_triggers >}}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/09triggers.go" after_window_trigger >}}
+{{< /highlight >}}
+
#### 9.1.1. Default trigger {#default-trigger}
The default trigger for a `PCollection` is based on event time, and emits the
@@ -4594,7 +4825,8 @@
The `AfterProcessingTime` trigger operates on *processing time*. For example,
the <span class="language-java">`AfterProcessingTime.pastFirstElementInPane()`</span>
-<span class="language-py">`AfterProcessingTime`</span> trigger emits a window
+<span class="language-py">`AfterProcessingTime`</span>
+<span class="language-go">`window.TriggerAfterProcessingTime()`</span> trigger emits a window
after a certain amount of processing time has passed since data was received.
The processing time is determined by the system clock, rather than the data
element's timestamp.
@@ -4607,14 +4839,16 @@
Beam provides one data-driven trigger,
<span class="language-java">`AfterPane.elementCountAtLeast()`</span>
-<span class="language-py">`AfterCount`</span>. This trigger works on an element
+<span class="language-py">`AfterCount`</span>
+<span class="language-go">`window.TriggerAfterCount()`</span>. This trigger works on an element
count; it fires after the current pane has collected at least *N* elements. This
allows a window to emit early results (before all the data has accumulated),
which can be particularly useful if you are using a single global window.
It is important to note that if, for example, you specify
<span class="language-java">`.elementCountAtLeast(50)`</span>
-<span class="language-py">AfterCount(50)</span> and only 32 elements arrive,
+<span class="language-py">AfterCount(50)</span>
+<span class="language-go">`window.TriggerAfterCount(50)`</span> and only 32 elements arrive,
those 32 elements sit around forever. If the 32 elements are important to you,
consider using [composite triggers](#composite-triggers) to combine multiple
conditions. This allows you to specify multiple firing conditions such as "fire
@@ -4623,7 +4857,7 @@
### 9.4. Setting a trigger {#setting-a-trigger}
When you set a windowing function for a `PCollection` by using the
-<span class="language-java">`Window`</span><span class="language-py">`WindowInto`</span>
+<span class="language-java">`Window`</span><span class="language-py">`WindowInto`</span><span class="language-go">`beam.WindowInto`</span>
transform, you can also specify a trigger.
{{< paragraph class="language-java" >}}
@@ -4643,6 +4877,14 @@
sets the window's **accumulation mode**.
{{< /paragraph >}}
+{{< paragraph class="language-go" >}}
+You set the trigger(s) for a `PCollection` by passing in the `beam.Trigger` parameter
+when you use the `beam.WindowInto` transform. This code sample sets a time-based
+trigger for a `PCollection`, which emits results one minute after the first
+element in that window has been processed.
+ The `beam.AccumulationMode` parameter sets the window's **accumulation mode**.
+{{< /paragraph >}}
+
{{< highlight java >}}
PCollection<String> pc = ...;
pc.apply(Window.<String>into(FixedWindows.of(1, TimeUnit.MINUTES))
@@ -4658,6 +4900,10 @@
accumulation_mode=AccumulationMode.DISCARDING)
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/09triggers.go" setting_a_trigger >}}
+{{< /highlight >}}
+
#### 9.4.1. Window accumulation modes {#window-accumulation-modes}
When you specify a trigger, you must also set the the window's **accumulation
@@ -4679,6 +4925,13 @@
`DISCARDING`.
{{< /paragraph >}}
+{{< paragraph class="language-go" >}}
+To set a window to accumulate the panes that are produced when the trigger
+fires, set the `beam.AccumulationMode` parameter to `beam.PanesAccumulate()` when you set the
+trigger. To set a window to discard fired panes, set `beam.AccumulationMode` to
+`beam.PanesDiscard()`.
+{{< /paragraph >}}
+
Let's look an example that uses a `PCollection` with fixed-time windowing and a
data-based trigger. This is something you might do if, for example, each window
represented a ten-minute running average, but you wanted to display the current
@@ -4728,8 +4981,10 @@
to the late data. If allowed lateness is set, the default trigger will emit new
results immediately whenever late data arrives.
-You set the allowed lateness by using `.withAllowedLateness()` when you set your
-windowing function:
+You set the allowed lateness by using <span class="language-java">`.withAllowedLateness()`</span>
+<span class="language-py">`allowed_lateness`</span>
+<span class="language-go">`beam.AllowedLateness()`</span>
+when you set your windowing function:
{{< highlight java >}}
PCollection<String> pc = ...;
@@ -4749,10 +5004,17 @@
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/09triggers.go" setting_allowed_lateness >}}
+{{< /highlight >}}
+
This allowed lateness propagates to all `PCollection`s derived as a result of
applying transforms to the original `PCollection`. If you want to change the
allowed lateness later in your pipeline, you can apply
-`Window.configure().withAllowedLateness()` again, explicitly.
+<span class="language-java">`Window.configure().withAllowedLateness()`</span>
+<span class="language-py">`allowed_lateness`</span>
+<span class="language-go">`beam.AllowedLateness()`</span>
+again, explicitly.
### 9.5. Composite triggers {#composite-triggers}
@@ -4825,6 +5087,10 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test.py" model_composite_triggers >}}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/09triggers.go" model_composite_triggers >}}
+{{< /highlight >}}
+
#### 9.5.3. Other composite triggers {#other-composite-triggers}
You can also build other sorts of composite triggers. The following example code
@@ -4873,6 +5139,12 @@
There are three types of metrics that are supported for the moment: `Counter`, `Distribution` and
`Gauge`.
+{{< paragraph class="language-go" >}}
+In the Beam SDK for Go, a `context.Context` provided by the framework must be passed to the metric
+or the metric value will not be recorded. The framework will automatically provide a valid
+`context.Context` to `ProcessElement` and similar methods when it's the first parameter.
+{{< /paragraph >}}
+
**Counter**: A metric that reports a single long value and can be incremented or decremented.
{{< highlight java >}}
@@ -4886,6 +5158,16 @@
}
{{< /highlight >}}
+{{< highlight go >}}
+var counter = beam.NewCounter("namespace", "counter1")
+
+func (fn *MyDoFn) ProcessElement(ctx context.Context, ...) {
+ // count the elements
+ counter.Inc(ctx, 1)
+ ...
+}
+{{< /highlight >}}
+
**Distribution**: A metric that reports information about the distribution of reported values.
{{< highlight java >}}
@@ -4900,6 +5182,16 @@
}
{{< /highlight >}}
+{{< highlight go >}}
+var distribution = beam.NewDistribution("namespace", "distribution1")
+
+func (fn *MyDoFn) ProcessElement(ctx context.Context, v int64, ...) {
+ // create a distribution (histogram) of the values
+ distribution.Inc(ctx, v)
+ ...
+}
+{{< /highlight >}}
+
**Gauge**: A metric that reports the latest value out of reported values. Since metrics are
collected from many workers the value may not be the absolute last, but one of the latest values.
@@ -4915,10 +5207,29 @@
}
{{< /highlight >}}
+{{< highlight go >}}
+var gauge = beam.NewGauge("namespace", "gauge1")
+
+func (fn *MyDoFn) ProcessElement(ctx context.Context, v int64, ...) {
+ // create a gauge (latest value received) of the values
+ gauge.Set(ctx, v)
+ ...
+}
+{{< /highlight >}}
+
### 10.3. Querying metrics {#querying-metrics}
+{{< paragraph class="language-java language-python">}}
`PipelineResult` has a method `metrics()` which returns a `MetricResults` object that allows
accessing metrics. The main method available in `MetricResults` allows querying for all metrics
matching a given filter.
+{{< /paragraph >}}
+
+{{< paragraph class="language-go">}}
+`beam.PipelineResult` has a method `Metrics()` which returns a `metrics.Results` object that allows
+accessing metrics. The main method available in `metrics.Results` allows querying for all metrics
+matching a given filter. It takes in a predicate with a `SingleResult` parameter type, which can
+be used for custom filters.
+{{< /paragraph >}}
{{< highlight java >}}
public interface PipelineResult {
@@ -4943,6 +5254,10 @@
}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/10metrics.go" metrics_query >}}
+{{< /highlight >}}
+
### 10.4. Using metrics in pipeline {#using-metrics}
Below, there is a simple example of how to use a `Counter` metric in a user pipeline.
@@ -4981,6 +5296,10 @@
}
{{< /highlight >}}
+{{< highlight go >}}
+{{< code_sample "sdks/go/examples/snippets/10metrics.go" metrics_pipeline >}}
+{{< /highlight >}}
+
### 10.5. Export metrics {#export-metrics}
Beam metrics can be exported to external sinks. If a metrics sink is set up in the configuration, the runner will push metrics to it at a default 5s period.
@@ -6350,6 +6669,11 @@
Currently Python external transforms are limited to dependencies available in core Beam SDK Harness.
+#### 13.1.3. Creating cross-language Go transforms
+
+Go currently does not support creating cross-language transforms, only using cross-language
+transforms from other languages; see more at [BEAM-9923](https://issues.apache.org/jira/browse/BEAM-9923).
+
### 13.2. Using cross-language transforms {#use-x-lang-transforms}
Depending on the SDK language of the pipeline, you can use a high-level SDK-wrapper class, or a low-level transform class to access a cross-language transform.
@@ -6369,7 +6693,7 @@
3. Include [External.of(...)](https://github.com/apache/beam/blob/master/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/External.java) when instantiating your pipeline. Reference the URN, payload, and expansion service. For examples, see the [cross-language transform test suite](https://github.com/apache/beam/blob/master/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ValidateRunnerXlangTest.java).
4. After the job has been submitted to the Beam runner, shutdown the expansion service by terminating the expansion service process.
-#### 13.2.2 Using cross-language transforms in a Python pipeline
+#### 13.2.2. Using cross-language transforms in a Python pipeline
If a Python-specific wrapper for a cross-language transform is available, use that; otherwise, you have to use the lower-level [ExternalTransform](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/transforms/external.py) class to access the transform.
@@ -6416,6 +6740,68 @@
4. After the job has been submitted to the Beam runner, shutdown the expansion service by terminating the expansion service process.
+#### 13.2.3. Using cross-language transforms in a Go pipeline
+
+If a Go-specific wrapper for a cross-language is available, use that; otherwise, you have to use the
+lower-level [CrossLanguage](https://pkg.go.dev/github.com/apache/beam/sdks/go/pkg/beam#CrossLanguage)
+function to access the transform.
+
+**Expansion Services**
+
+The Go SDK does not yet support automatically starting an expansion service. In order to use
+cross-language transforms, you must manually start any necessary expansion services on your local
+machine and ensure they are accessible to your code during pipeline construction; see more at
+[BEAM-12862](https://issues.apache.org/jira/browse/BEAM-12862).
+
+**Using an SDK wrapper**
+
+To use a cross-language transform through an SDK wrapper, import the package for the SDK wrapper
+and call it from your pipeline as shown in the example:
+
+{{< highlight >}}
+import (
+ "github.com/apache/beam/sdks/v2/go/pkg/beam/io/xlang/kafkaio"
+)
+
+// Kafka Read using previously defined values.
+kafkaRecords := kafkaio.Read(
+ s,
+ expansionAddr, // Address of expansion service.
+ bootstrapAddr,
+ []string{topicName},
+ kafkaio.MaxNumRecords(numRecords),
+ kafkaio.ConsumerConfigs(map[string]string{"auto.offset.reset": "earliest"}))
+{{< /highlight >}}
+
+**Using the CrossLanguage function**
+
+When an SDK-specific wrapper isn't available, you will have to access the cross-language transform through the `beam.CrossLanguage` function.
+
+1. Make sure you have the appropriate expansion service running. See the expansion service section for details.
+2. Make sure the transform you're trying to use is available and can be used by the expansion service.
+ Refer to [Creating cross-language transforms](#create-x-lang-transforms) for details.
+3. Use the `beam.CrossLanguage` function in your pipeline as appropriate. Reference the URN, Payload,
+ expansion service address, and define inputs and outputs. You can use the
+ [beam.CrossLanguagePayload](https://pkg.go.dev/github.com/apache/beam/sdks/go/pkg/beam#CrossLanguagePayload)
+ function as a helper for encoding a payload. You can use the
+ [beam.UnnamedInput](https://pkg.go.dev/github.com/apache/beam/sdks/go/pkg/beam#UnnamedInput) and
+ [beam.UnnamedOutput](https://pkg.go.dev/github.com/apache/beam/sdks/go/pkg/beam#UnnamedOutput)
+ functions as shortcuts for single, unnamed inputs/outputs or define a map for named ones.
+
+ {{< highlight >}}
+type prefixPayload struct {
+ Data string `beam:"data"`
+}
+urn := "beam:transforms:xlang:test:prefix"
+payload := beam.CrossLanguagePayload(prefixPayload{Data: prefix})
+expansionAddr := "localhost:8097"
+outT := beam.UnnamedOutput(typex.New(reflectx.String))
+res := beam.CrossLanguage(s, urn, payload, expansionAddr, beam.UnnamedInput(inputPCol), outT)
+ {{< /highlight >}}
+
+4. After the job has been submitted to the Beam runner, shutdown the expansion service by
+ terminating the expansion service process.
+
### 13.3. Runner Support {#x-lang-transform-runner-support}
Currently, portable runners such as Flink, Spark, and the Direct runner can be used with multi-language pipelines.
diff --git a/website/www/site/static/js/copy-to-clipboard.js b/website/www/site/static/js/copy-to-clipboard.js
index fad14d02..4555f7e 100644
--- a/website/www/site/static/js/copy-to-clipboard.js
+++ b/website/www/site/static/js/copy-to-clipboard.js
@@ -13,20 +13,18 @@
$(document).ready(function() {
function copy() {
$(".copy").click(function(){
- var text=$(this).siblings()[$(this).siblings().length-1].childNodes[0].innerHTML;
+ var text=$(this).siblings()[$(this).siblings().length-1].childNodes[0].innerText;
const el=document.createElement('textarea');
el.value=text;document.body.appendChild(el);
el.select();document.execCommand('copy');
document.body.removeChild(el);
- alert('copied to clipboard');
})
$(".just-copy").click(function(){
- var text=$(this).parent().siblings()[0].innerHTML;
+ var text=$(this).parent().siblings()[0].innerText;
const el=document.createElement('textarea');
el.value=text;document.body.appendChild(el);
el.select();document.execCommand('copy');
document.body.removeChild(el);
- alert('copied to clipboard');
})
}
let code = document.querySelectorAll('pre'),
