Merge pull request #10683 from EDjur/BEAM-7810/valueprovider-namespace
[BEAM-7810] Added ValueProvider support for Datastore query namespaces
diff --git a/.test-infra/dataproc/flink_cluster.sh b/.test-infra/dataproc/flink_cluster.sh
index 8b6f939..ad88578 100755
--- a/.test-infra/dataproc/flink_cluster.sh
+++ b/.test-infra/dataproc/flink_cluster.sh
@@ -106,7 +106,7 @@
local job_server_config=`gcloud compute ssh --quiet --zone=$GCLOUD_ZONE yarn@$MASTER_NAME --command="curl -s \"http://$YARN_APPLICATION_MASTER/jobmanager/config\""`
local key="jobmanager.rpc.port"
local yarn_application_master_host=`echo $YARN_APPLICATION_MASTER | cut -d ":" -f1`
- local jobmanager_rpc_port=`echo $job_server_config | python -c "import sys, json; print [ e['value'] for e in json.load(sys.stdin) if e['key'] == u'$key'][0]"`
+ local jobmanager_rpc_port=`echo $job_server_config | python -c "import sys, json; print([e['value'] for e in json.load(sys.stdin) if e['key'] == u'$key'][0])"`
local detached_mode_params=$([[ $DETACHED_MODE == "true" ]] && echo " -Nf >& /dev/null" || echo "")
diff --git a/.test-infra/jenkins/CommonTestProperties.groovy b/.test-infra/jenkins/CommonTestProperties.groovy
index 7bf585e..74a8bd0 100644
--- a/.test-infra/jenkins/CommonTestProperties.groovy
+++ b/.test-infra/jenkins/CommonTestProperties.groovy
@@ -28,7 +28,8 @@
enum Runner {
DATAFLOW("DataflowRunner"),
SPARK("SparkRunner"),
- FLINK("TestFlinkRunner"),
+ SPARK_STRUCTURED_STREAMING("SparkStructuredStreamingRunner"),
+ FLINK("FlinkRunner"),
DIRECT("DirectRunner"),
PORTABLE("PortableRunner")
@@ -36,6 +37,7 @@
JAVA: [
DATAFLOW: ":runners:google-cloud-dataflow-java",
SPARK: ":runners:spark",
+ SPARK_STRUCTURED_STREAMING: ":runners:spark",
FLINK: ":runners:flink:1.9",
DIRECT: ":runners:direct-java"
],
diff --git a/.test-infra/jenkins/LoadTestsBuilder.groovy b/.test-infra/jenkins/LoadTestsBuilder.groovy
index d7ed2d5..8b5e825 100644
--- a/.test-infra/jenkins/LoadTestsBuilder.groovy
+++ b/.test-infra/jenkins/LoadTestsBuilder.groovy
@@ -48,6 +48,12 @@
}
}
+ static String parseOptions(Map<String, ?> options) {
+ options.collect {
+ "--${it.key}=$it.value".replace('\"', '\\\"').replace('\'', '\\\'')
+ }.join(' ')
+ }
+
static String getBigQueryDataset(String baseName, TriggeringContext triggeringContext) {
if (triggeringContext == TriggeringContext.PR) {
return baseName + '_PRs'
@@ -76,12 +82,6 @@
throw new RuntimeException("No task name defined for SDK: $SDK")
}
}
-
- private static String parseOptions(Map<String, ?> options) {
- options.collect {
- "--${it.key}=$it.value".replace('\"', '\\\"').replace('\'', '\\\'')
- }.join(' ')
- }
}
diff --git a/.test-infra/jenkins/NexmarkBuilder.groovy b/.test-infra/jenkins/NexmarkBuilder.groovy
index 9cdef21..942b964 100644
--- a/.test-infra/jenkins/NexmarkBuilder.groovy
+++ b/.test-infra/jenkins/NexmarkBuilder.groovy
@@ -52,12 +52,10 @@
suite(context, "NEXMARK IN SQL STREAMING MODE USING ${runner} RUNNER", runner, sdk, options)
}
- static void batchOnlyJob(context, Map<String, Object> jobSpecificOptions, TriggeringContext triggeringContext) {
- Runner runner = Runner.SPARK
- SDK sdk = SDK.JAVA
+ static void batchOnlyJob(context, Runner runner, SDK sdk, Map<String, Object> jobSpecificOptions, TriggeringContext triggeringContext) {
Map<String, Object> options = getFullOptions(jobSpecificOptions, runner, triggeringContext)
- options.put('streaming', false)
+ options.put('streaming', false)
suite(context, "NEXMARK IN BATCH MODE USING ${runner} RUNNER", runner, sdk, options)
options.put('queryLanguage', 'sql')
diff --git a/.test-infra/jenkins/README.md b/.test-infra/jenkins/README.md
index 471d914..ef1cc3e 100644
--- a/.test-infra/jenkins/README.md
+++ b/.test-infra/jenkins/README.md
@@ -59,11 +59,11 @@
| beam_PostCommit_Java_PVR_Flink_Streaming | [cron](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Streaming/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Streaming_PR/) | `Run Java Flink PortableValidatesRunner Streaming` | [](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Flink_Streaming) |
| beam_PostCommit_Java_PVR_Spark_Batch | [cron](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Spark_Batch/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Spark_Batch_PR/) | `Run Java Spark PortableValidatesRunner Batch` | [](https://builds.apache.org/job/beam_PostCommit_Java_PVR_Spark_Batch) |
| beam_PostCommit_Java_PortabilityApi | [cron](https://builds.apache.org/job/beam_PostCommit_Java_PortabilityApi/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_PortabilityApi_PR/) | `Run Java PortabilityApi PostCommit` | [](https://builds.apache.org/job/beam_PostCommit_Java_PortabilityApi) |
-| beam_PostCommit_Java11_Dataflow_Examples | [cron](https://builds.apache.org/job/beam_PostCommit_Java11_Examples_Dataflow/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java11_Examples_Dataflow_PR/) | `Run Java examples on Dataflow with Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java11_Examples_Dataflow) |
-| beam_PostCommit_Java11_Dataflow_Portability_Examples | [cron](https://builds.apache.org/job/beam_PostCommit_Java11_Examples_Dataflow_Portability/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java11_Examples_Dataflow_Portability_PR/) | `Run Java Portability examples on Dataflow with Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java11_Examples_Dataflow_Portability) |
-| beam_PostCommit_Java11_ValidatesRunner_Dataflow | [cron](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_Dataflow/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_Dataflow_PR/) | `Run Dataflow ValidatesRunner Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_Dataflow) |
-| beam_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow | [cron](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow_PR/) | `Run Dataflow PortabilityApi ValidatesRunner with Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow) |
-| beam_PostCommit_Java11_ValidatesRunner_Direct | [cron](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_Direct), [phrase](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_Direct_PR) | `Run Direct ValidatesRunner in Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java11_ValidatesRunner_Direct) |
+| beam_PostCommit_Java_Dataflow_Examples_Java11 | [cron](https://builds.apache.org/job/beam_PostCommit_Java_Examples_Dataflow_Java11/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_Examples_Dataflow_Java11_PR/) | `Run Java examples on Dataflow Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java_Examples_Dataflow_Java11) |
+| beam_PostCommit_Java_Dataflow_Portability_Examples_Java11 | [cron](https://builds.apache.org/job/beam_PostCommit_Java_Examples_Dataflow_Portability_Java11/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_Examples_Dataflow_Portability_Java11_PR/) | `Run Java Portability examples on Dataflow Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java_Examples_Dataflow_Portability_Java11) |
+| beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11 | [cron](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11_PR/) | `Run Dataflow ValidatesRunner Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11) |
+| beam_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11 | [cron](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11_PR/) | `Run Dataflow PortabilityApi ValidatesRunner Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11) |
+| beam_PostCommit_Java_ValidatesRunner_Direct_Java11 | [cron](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Direct_Java11), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Direct_Java11_PR) | `Run Direct ValidatesRunner Java 11` | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Direct_Java11) |
| beam_PostCommit_Java_ValidatesRunner_Apex | [cron](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Apex/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Apex_PR/) | `Run Apex ValidatesRunner` | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Apex) |
| beam_PostCommit_Java_ValidatesRunner_Dataflow | [cron](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_PR/) | `Run Dataflow ValidatesRunner` | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow) |
| beam_PostCommit_Java_ValidatesRunner_Flink | [cron](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink/), [phrase](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink_PR/) | `Run Flink ValidatesRunner` | [](https://builds.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink) |
@@ -75,6 +75,8 @@
| beam_PostCommit_Py_VR_Dataflow | [cron](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow/), [phrase](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow_PR/) | `Run Python Dataflow ValidatesRunner` | [](https://builds.apache.org/job/beam_PostCommit_Py_VR_Dataflow) |
| beam_PostCommit_Py_ValCont | [cron](https://builds.apache.org/job/beam_PostCommit_Py_ValCont/), [phrase](https://builds.apache.org/job/beam_PostCommit_Py_ValCont_PR/) | `Run Python Dataflow ValidatesContainer` | [](https://builds.apache.org/job/beam_PostCommit_Py_ValCont) |
| beam_PostCommit_Python35_VR_Flink | [cron](https://builds.apache.org/job/beam_PostCommit_Python35_VR_Flink/), [phrase](https://builds.apache.org/job/beam_PostCommit_Python35_VR_Flink/) | `Run Python 3.5 Flink ValidatesRunner` | [](https://builds.apache.org/job/beam_PostCommit_Python35_VR_Flink) |
+| beam_PostCommit_Python_Chicago_Taxi_Example_Dataflow | [cron](https://builds.apache.org/job/beam_PostCommit_Python_Chicago_Taxi_Dataflow/), [phrase](https://builds.apache.org/job/beam_PostCommit_Python_Chicago_Taxi_Dataflow_PR/) | `Run Chicago Taxi on Dataflow` | [](https://builds.apache.org/job/beam_PostCommit_Python_Chicago_Taxi_Dataflow) |
+| beam_PostCommit_Python_Chicago_Taxi_Example_Flink | [cron](https://builds.apache.org/job/beam_PostCommit_Python_Chicago_Taxi_Flink/), [phrase](https://builds.apache.org/job/beam_PostCommit_Python_Chicago_Taxi_Flink_PR/) | `Run Chicago Taxi on Flink` | [](https://builds.apache.org/job/beam_PostCommit_Python_Chicago_Taxi_Flink) |
| beam_PostCommit_Python_MongoDBIO_IT | [cron](https://builds.apache.org/job/beam_PostCommit_Python_MongoDBIO_IT), [phrase](https://builds.apache.org/job/beam_PostCommit_Python_MongoDBIO_IT_PR/) | `Run Python MongoDBIO_IT` | [](https://builds.apache.org/job/beam_PostCommit_Python_MongoDBIO_IT) |
| beam_PostCommit_Python_VR_Spark | [cron](https://builds.apache.org/job/beam_PostCommit_Python_VR_Spark/), [phrase](https://builds.apache.org/job/beam_PostCommit_Python_VR_Spark/) | `Run Python Spark ValidatesRunner` | [](https://builds.apache.org/job/beam_PostCommit_Python_VR_Spark) |
| beam_PostCommit_Python2 | [cron](https://builds.apache.org/job/beam_PostCommit_Python2), [phrase](https://builds.apache.org/job/beam_PostCommit_Python2_PR/) | `Run Python 2 PostCommit` | [](https://builds.apache.org/job/beam_PostCommit_Python2) |
diff --git a/.test-infra/jenkins/job_CancelStaleDataflowJobs.groovy b/.test-infra/jenkins/job_CancelStaleDataflowJobs.groovy
index a03a1d0..e32d14a 100644
--- a/.test-infra/jenkins/job_CancelStaleDataflowJobs.groovy
+++ b/.test-infra/jenkins/job_CancelStaleDataflowJobs.groovy
@@ -37,7 +37,6 @@
steps {
gradle {
rootBuildScriptDir(commonJobProperties.checkoutDir)
- tasks(':beam-test-tools:check')
tasks(':beam-test-tools:cancelStaleDataflowJobs')
commonJobProperties.setGradleSwitches(delegate)
}
diff --git a/.test-infra/jenkins/job_PerformanceTests_SQLIO_Java.groovy b/.test-infra/jenkins/job_PerformanceTests_SQLIO_Java.groovy
new file mode 100644
index 0000000..3de488a
--- /dev/null
+++ b/.test-infra/jenkins/job_PerformanceTests_SQLIO_Java.groovy
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 CommonJobProperties as common
+
+def jobConfigs = [
+ [
+ title : 'SQL BigQueryIO with push-down Batch Performance Test Java',
+ triggerPhrase: 'Run SQLBigQueryIO Batch Performance Test Java',
+ name : 'beam_SQLBigQueryIO_Batch_Performance_Test_Java',
+ itClass : 'org.apache.beam.sdk.extensions.sql.meta.provider.bigquery.BigQueryIOPushDownIT',
+ properties: [
+ project : 'apache-beam-testing',
+ tempLocation : 'gs://temp-storage-for-perf-tests/loadtests',
+ tempRoot : 'gs://temp-storage-for-perf-tests/loadtests',
+ metricsBigQueryDataset: 'beam_performance',
+ metricsBigQueryTable : 'sql_bqio_read_java_batch',
+ runner : "DataflowRunner",
+ maxNumWorkers : '5',
+ numWorkers : '5',
+ autoscalingAlgorithm : 'NONE',
+ ]
+ ]
+]
+
+jobConfigs.forEach { jobConfig -> createPostCommitJob(jobConfig)}
+
+private void createPostCommitJob(jobConfig) {
+ job(jobConfig.name) {
+ description(jobConfig.description)
+ common.setTopLevelMainJobProperties(delegate)
+ common.enablePhraseTriggeringFromPullRequest(delegate, jobConfig.title, jobConfig.triggerPhrase)
+ common.setAutoJob(delegate, 'H */6 * * *')
+ publishers {
+ archiveJunit('**/build/test-results/**/*.xml')
+ }
+
+ steps {
+ gradle {
+ rootBuildScriptDir(common.checkoutDir)
+ common.setGradleSwitches(delegate)
+ switches("--info")
+ switches("-DintegrationTestPipelineOptions=\'${common.joinOptionsWithNestedJsonValues(jobConfig.properties)}\'")
+ switches("-DintegrationTestRunner=dataflow")
+ tasks(":sdks:java:extensions:sql:perf-tests:integrationTest --tests ${jobConfig.itClass}")
+ }
+ }
+ }
+}
diff --git a/.test-infra/jenkins/job_PostCommit_Java11_Dataflow_Examples.groovy b/.test-infra/jenkins/job_PostCommit_Java_Dataflow_Examples_Java11.groovy
similarity index 90%
rename from .test-infra/jenkins/job_PostCommit_Java11_Dataflow_Examples.groovy
rename to .test-infra/jenkins/job_PostCommit_Java_Dataflow_Examples_Java11.groovy
index 9e1004c..1e88537 100644
--- a/.test-infra/jenkins/job_PostCommit_Java11_Dataflow_Examples.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_Dataflow_Examples_Java11.groovy
@@ -19,8 +19,8 @@
import PostcommitJobBuilder
import CommonJobProperties as commonJobProperties
-PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java11_Examples_Dataflow',
- 'Run Java examples on Dataflow with Java 11', 'Google Cloud Dataflow Runner Examples Java 11', this) {
+PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_Java11',
+ 'Run Java examples on Dataflow Java 11', 'Google Cloud Dataflow Runner Examples Java 11', this) {
description('Runs the Java Examples suite on the Java 11 enabled Dataflow runner.')
diff --git a/.test-infra/jenkins/job_PostCommit_Java11_Dataflow_Portability_Examples.groovy b/.test-infra/jenkins/job_PostCommit_Java_Dataflow_Portability_Examples_Java11.groovy
similarity index 89%
rename from .test-infra/jenkins/job_PostCommit_Java11_Dataflow_Portability_Examples.groovy
rename to .test-infra/jenkins/job_PostCommit_Java_Dataflow_Portability_Examples_Java11.groovy
index eadf255..1222c06 100644
--- a/.test-infra/jenkins/job_PostCommit_Java11_Dataflow_Portability_Examples.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_Dataflow_Portability_Examples_Java11.groovy
@@ -19,8 +19,8 @@
import PostcommitJobBuilder
import CommonJobProperties as commonJobProperties
-PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java11_Examples_Dataflow_Portability',
- 'Run Java Portability examples on Dataflow with Java 11', 'Google Cloud Dataflow Portability Runner Examples Java 11', this) {
+PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_Portability_Java11',
+ 'Run Java Portability examples on Dataflow Java 11', 'Google Cloud Dataflow Portability Runner Examples Java 11', this) {
description('Runs the Java Examples suite on the Java 11 enabled Dataflow runner with Portability API.')
diff --git a/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Flink.groovy b/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Flink.groovy
index cbcd0ba..30cf5ee 100644
--- a/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Flink.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Flink.groovy
@@ -43,6 +43,8 @@
switches('-Pnexmark.runner=":runners:flink:1.9"' +
' -Pnexmark.args="' +
[NexmarkBigqueryProperties.nexmarkBigQueryArgs,
+ '--runner=FlinkRunner',
+ '--shutdownSourcesOnFinalWatermark=true',
'--streaming=false',
'--suite=SMOKE',
'--streamTimeout=60' ,
@@ -58,6 +60,8 @@
switches('-Pnexmark.runner=":runners:flink:1.9"' +
' -Pnexmark.args="' +
[NexmarkBigqueryProperties.nexmarkBigQueryArgs,
+ '--runner=FlinkRunner',
+ '--shutdownSourcesOnFinalWatermark=true',
'--streaming=true',
'--suite=SMOKE',
'--streamTimeout=60' ,
@@ -73,6 +77,8 @@
switches('-Pnexmark.runner=":runners:flink:1.9"' +
' -Pnexmark.args="' +
[NexmarkBigqueryProperties.nexmarkBigQueryArgs,
+ '--runner=FlinkRunner',
+ '--shutdownSourcesOnFinalWatermark=true',
'--queryLanguage=sql',
'--streaming=false',
'--suite=SMOKE',
@@ -88,6 +94,8 @@
switches('-Pnexmark.runner=":runners:flink:1.9"' +
' -Pnexmark.args="' +
[NexmarkBigqueryProperties.nexmarkBigQueryArgs,
+ '--runner=FlinkRunner',
+ '--shutdownSourcesOnFinalWatermark=true',
'--queryLanguage=sql',
'--streaming=true',
'--suite=SMOKE',
@@ -107,8 +115,9 @@
commonJobProperties.setTopLevelMainJobProperties(delegate, 'master', 240)
def final JOB_SPECIFIC_OPTIONS = [
- 'suite' : 'SMOKE',
- 'streamTimeout': 60,
+ 'suite' : 'SMOKE',
+ 'streamTimeout' : 60,
+ 'shutdownSourcesOnFinalWatermark' : true,
]
Nexmark.standardJob(delegate, Runner.FLINK, SDK.JAVA, JOB_SPECIFIC_OPTIONS, TriggeringContext.PR)
diff --git a/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy b/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy
index 5fb7d24..838fac0 100644
--- a/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_Nexmark_Spark.groovy
@@ -17,6 +17,8 @@
*/
import CommonJobProperties as commonJobProperties
+import CommonTestProperties.Runner
+import CommonTestProperties.SDK
import CommonTestProperties.TriggeringContext
import NexmarkBigqueryProperties
import NexmarkBuilder as Nexmark
@@ -64,6 +66,39 @@
'--manageResources=false',
'--monitorJobs=true'].join(' '))
}
+ shell('echo *** RUN NEXMARK IN BATCH MODE USING SPARK STRUCTURED STREAMING RUNNER ***')
+ gradle {
+ rootBuildScriptDir(commonJobProperties.checkoutDir)
+ tasks(':sdks:java:testing:nexmark:run')
+ commonJobProperties.setGradleSwitches(delegate)
+ switches('-Pnexmark.runner=":runners:spark"' +
+ ' -Pnexmark.args="' +
+ [NexmarkBigqueryProperties.nexmarkBigQueryArgs,
+ '--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 STRUCTURED STREAMING RUNNER ***')
+ gradle {
+ rootBuildScriptDir(commonJobProperties.checkoutDir)
+ tasks(':sdks:java:testing:nexmark:run')
+ commonJobProperties.setGradleSwitches(delegate)
+ switches('-Pnexmark.runner=":runners:spark"' +
+ ' -Pnexmark.args="' +
+ [NexmarkBigqueryProperties.nexmarkBigQueryArgs,
+ '--runner=SparkStructuredStreamingRunner',
+ '--queryLanguage=sql',
+ '--streaming=false',
+ '--suite=SMOKE',
+ '--streamTimeout=60' ,
+ '--manageResources=false',
+ '--monitorJobs=true'].join(' '))
+ }
}
}
@@ -78,7 +113,14 @@
'suite' : 'SMOKE',
'streamTimeout': 60
]
-
// Spark doesn't run streaming jobs, therefore run only batch variants.
- Nexmark.batchOnlyJob(delegate, JOB_SPECIFIC_OPTIONS, TriggeringContext.PR)
+ Nexmark.batchOnlyJob(delegate, Runner.SPARK, SDK.JAVA, JOB_SPECIFIC_OPTIONS, TriggeringContext.PR)
+
+ def final SPARK_STRUCTURED_STREAMING_JOB_SPECIFIC_OPTIONS = [
+ 'suite' : 'SMOKE',
+ 'streamTimeout': 60,
+ // Skip query 3 (SparkStructuredStreamingRunner does not support State/Timers yet)
+ 'skipQueries' : 3,
+ ]
+ Nexmark.batchOnlyJob(delegate, Runner.SPARK_STRUCTURED_STREAMING, SDK.JAVA, SPARK_STRUCTURED_STREAMING_JOB_SPECIFIC_OPTIONS, TriggeringContext.PR)
}
diff --git a/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Dataflow_Java11.groovy b/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Dataflow_Java11.groovy
index a1e7fc9..431c922 100644
--- a/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Dataflow_Java11.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Dataflow_Java11.groovy
@@ -20,7 +20,7 @@
import PostcommitJobBuilder
-PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java11_ValidatesRunner_Dataflow',
+PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11',
'Run Dataflow ValidatesRunner Java 11', 'Google Cloud Dataflow Runner ValidatesRunner Tests On Java 11', this) {
description('Runs the ValidatesRunner suite on the Dataflow runner with Java 11 worker harness.')
diff --git a/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Direct_Java11.groovy b/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Direct_Java11.groovy
index 5e3c8cf..f38058a 100644
--- a/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Direct_Java11.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_Direct_Java11.groovy
@@ -22,8 +22,8 @@
// This job runs the suite of ValidatesRunner tests using Java 11 against the Direct
// runner compiled with Java 8.
-PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java11_ValidatesRunner_Direct',
- 'Run Direct ValidatesRunner in Java 11', 'Direct Runner ValidatesRunner Tests for Java 11', this) {
+PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_ValidatesRunner_Direct_Java11',
+ 'Run Direct ValidatesRunner Java 11', 'Direct Runner ValidatesRunner Tests for Java 11', this) {
description('Builds the Direct Runner with Java 8 and runs ValidatesRunner test suite in Java 11.')
diff --git a/.test-infra/jenkins/job_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow.groovy b/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11.groovy
similarity index 87%
rename from .test-infra/jenkins/job_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow.groovy
rename to .test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11.groovy
index 6ac6ff8..7eb1533 100644
--- a/.test-infra/jenkins/job_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11.groovy
@@ -20,8 +20,8 @@
import PostcommitJobBuilder
-PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java11_ValidatesRunner_PortabilityApi_Dataflow',
- 'Run Dataflow PortabilityApi ValidatesRunner with Java 11', 'Google Cloud Dataflow Runner PortabilityApi ValidatesRunner Java 11 Tests', this) {
+PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_ValidatesRunner_PortabilityApi_Dataflow_Java11',
+ 'Run Dataflow PortabilityApi ValidatesRunner Java 11', 'Google Cloud Dataflow Runner PortabilityApi ValidatesRunner Java 11 Tests', this) {
description('Runs the ValidatesRunner suite on the Java 11 enabled Dataflow PortabilityApi runner.')
diff --git a/.test-infra/jenkins/job_PostCommit_Python_Chicago_Taxi_Example_Dataflow.groovy b/.test-infra/jenkins/job_PostCommit_Python_Chicago_Taxi_Example_Dataflow.groovy
index 56d46be..859abfe 100644
--- a/.test-infra/jenkins/job_PostCommit_Python_Chicago_Taxi_Example_Dataflow.groovy
+++ b/.test-infra/jenkins/job_PostCommit_Python_Chicago_Taxi_Example_Dataflow.groovy
@@ -37,7 +37,7 @@
steps {
gradle {
rootBuildScriptDir(commonJobProperties.checkoutDir)
- tasks(':sdks:python:test-suites:dataflow:py2:dataflowChicagoTaxiExample')
+ tasks(':sdks:python:test-suites:dataflow:py2:chicagoTaxiExample')
switches('-PgcsRoot=gs://temp-storage-for-perf-tests/chicago-taxi')
}
}
@@ -57,7 +57,7 @@
steps {
gradle {
rootBuildScriptDir(commonJobProperties.checkoutDir)
- tasks(':sdks:python:test-suites:dataflow:py2:dataflowChicagoTaxiExample')
+ tasks(':sdks:python:test-suites:dataflow:py2:chicagoTaxiExample')
switches('-PgcsRoot=gs://temp-storage-for-perf-tests/chicago-taxi')
}
}
diff --git a/.test-infra/jenkins/job_PostCommit_Python_Chicago_Taxi_Example_Flink.groovy b/.test-infra/jenkins/job_PostCommit_Python_Chicago_Taxi_Example_Flink.groovy
new file mode 100644
index 0000000..4291b42
--- /dev/null
+++ b/.test-infra/jenkins/job_PostCommit_Python_Chicago_Taxi_Example_Flink.groovy
@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 CommonJobProperties as commonJobProperties
+import CronJobBuilder
+import Docker
+import Flink
+import LoadTestsBuilder
+import PostcommitJobBuilder
+
+def chicagoTaxiJob = { scope ->
+ scope.description('Runs the Chicago Taxi Example on the Flink runner.')
+ commonJobProperties.setTopLevelMainJobProperties(scope, 'master', 120)
+
+ def numberOfWorkers = 5
+
+ Docker publisher = new Docker(scope, LoadTestsBuilder.DOCKER_CONTAINER_REGISTRY)
+ publisher.publish(':sdks:python:container:py2:docker', 'python2.7_sdk')
+ publisher.publish(':runners:flink:1.9:job-server-container:docker', 'flink1.9_job_server')
+ String pythonHarnessImageTag = publisher.getFullImageName('python2.7_sdk')
+ Flink flink = new Flink(scope, 'beam_PostCommit_Python_Chicago_Taxi_Flink')
+ flink.setUp([pythonHarnessImageTag], numberOfWorkers, publisher.getFullImageName('flink1.9_job_server'))
+
+ def pipelineOptions = [
+ parallelism : numberOfWorkers,
+ job_endpoint : 'localhost:8099',
+ environment_config : pythonHarnessImageTag,
+ environment_type : 'DOCKER',
+ execution_mode_for_batch: 'BATCH_FORCED',
+ ]
+
+ scope.steps {
+ gradle {
+ rootBuildScriptDir(commonJobProperties.checkoutDir)
+ tasks(':sdks:python:test-suites:portable:py2:chicagoTaxiExample')
+ switches('-PgcsRoot=gs://temp-storage-for-perf-tests/chicago-taxi')
+ switches("-PpipelineOptions=\"${LoadTestsBuilder.parseOptions(pipelineOptions)}\"")
+ }
+ }
+}
+
+PostcommitJobBuilder.postCommitJob(
+ 'beam_PostCommit_Python_Chicago_Taxi_Flink',
+ 'Run Chicago Taxi on Flink',
+ 'Google Cloud Flink Runner Chicago Taxi Example',
+ this
+) {
+ chicagoTaxiJob(delegate)
+}
+
+CronJobBuilder.cronJob(
+ 'beam_PostCommit_Python_Chicago_Taxi_Flink',
+ 'H 14 * * *',
+ this
+) {
+ chicagoTaxiJob(delegate)
+}
diff --git a/.test-infra/junitxml_report.py b/.test-infra/junitxml_report.py
index 27c3818..ec1a899 100644
--- a/.test-infra/junitxml_report.py
+++ b/.test-infra/junitxml_report.py
@@ -41,6 +41,9 @@
if child.tag == 'skipped':
assert status == ''
status = 'S'
+ elif child.tag == 'failure':
+ assert status == ''
+ status = 'F'
elif child.tag in ['system-err', 'system-out']:
pass
else:
diff --git a/.test-infra/tools/build.gradle b/.test-infra/tools/build.gradle
index aabeca0..53445b6 100644
--- a/.test-infra/tools/build.gradle
+++ b/.test-infra/tools/build.gradle
@@ -16,27 +16,6 @@
* limitations under the License.
*/
-plugins {
- id 'org.apache.beam.module'
-}
-
-applyGoNature()
-
-repositories { mavenCentral() }
-
-clean {
- delete '.gogradle'
-}
-
-golang {
- packagePath = 'github.com/apache/beam/.test-infra/tools'
-}
-
-check.dependsOn goTest
-
-task cancelStaleDataflowJobs(type: com.github.blindpirate.gogradle.Go) {
- dependsOn goVendor
- go('get golang.org/x/oauth2/google')
- go('get google.golang.org/api/dataflow/v1b3')
- go('run stale_dataflow_jobs_cleaner.go')
+task cancelStaleDataflowJobs(type: Exec) {
+ commandLine './stale_dataflow_jobs_cleaner.sh'
}
diff --git a/.test-infra/tools/stale_dataflow_jobs_cleaner.go b/.test-infra/tools/stale_dataflow_jobs_cleaner.go
deleted file mode 100644
index 6361e27..0000000
--- a/.test-infra/tools/stale_dataflow_jobs_cleaner.go
+++ /dev/null
@@ -1,117 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * License); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an AS IS BASIS,
- * WITHOUT 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 (
- "context"
- "log"
- "strings"
- "time"
-
- "golang.org/x/oauth2/google"
- df "google.golang.org/api/dataflow/v1b3"
-)
-
-const (
- longRunningPrefix = "long-running-"
-)
-
-// client contains methods for listing and cancelling jobs, extracted to allow easier testing.
-type client interface {
- CurrentTime() time.Time
- ListJobs(projectId string) ([]*df.Job, error)
- CancelJob(job *df.Job) error
-}
-
-// dataflowClient implements the client interface for Google Cloud Dataflow.
-type dataflowClient struct {
- s *df.ProjectsJobsService
-}
-
-// newDataflowClient creates a new Dataflow ProjectsJobsService.
-func newDataflowClient() (*dataflowClient, error) {
- ctx := context.Background()
- cl, err := google.DefaultClient(ctx, df.CloudPlatformScope)
- if err != nil {
- return nil, err
- }
- service, err := df.New(cl)
- if err != nil {
- return nil, err
- }
- return &dataflowClient{s: df.NewProjectsJobsService(service)}, nil
-}
-
-// CurrentTime gets the time Now.
-func (c dataflowClient) CurrentTime() time.Time {
- return time.Now()
-}
-
-// ListJobs lists the active Dataflow jobs for a project.
-func (c dataflowClient) ListJobs(projectId string) ([]*df.Job, error) {
- resp, err := c.s.Aggregated(projectId).Filter("ACTIVE").Fields("jobs(id,name,projectId,createTime)").Do()
- if err != nil {
- return nil, err
- }
- return resp.Jobs, nil
-}
-
-// CancelJob requests the cancellation od a Dataflow job.
-func (c dataflowClient) CancelJob(job *df.Job) error {
- jobDone := df.Job{
- RequestedState: "JOB_STATE_DONE",
- }
- _, err := c.s.Update(job.ProjectId, job.Id, &jobDone).Do()
- return err
-}
-
-// cleanDataflowJobs cancels stale Dataflow jobs, excluding the longRunningPrefix prefixed jobs.
-func cleanDataflowJobs(c client, projectId string, hoursStale float64) error {
- now := c.CurrentTime()
- jobs, err := c.ListJobs(projectId)
- if err != nil {
- return err
- }
- for _, j := range jobs {
- t, err := time.Parse(time.RFC3339, j.CreateTime)
- if err != nil {
- return err
- }
- hoursSinceCreate := now.Sub(t).Hours()
- log.Printf("Job %v %v %v %v %.2f\n", j.ProjectId, j.Id, j.Name, j.CreateTime, hoursSinceCreate)
- if hoursSinceCreate > hoursStale && !strings.HasPrefix(j.Name, longRunningPrefix) {
- log.Printf("Attempting to cancel %v\n", j.Id)
- c.CancelJob(j)
- }
- }
- return nil
-}
-
-func main() {
- client, err := newDataflowClient()
- if err != nil {
- log.Fatalf("Error creating dataflow client, %v", err)
- }
- // Cancel any jobs older than 3 hours.
- err = cleanDataflowJobs(client, "apache-beam-testing", 3.0)
- if err != nil {
- log.Fatalf("Error cleaning dataflow jobs, %v", err)
- }
- log.Printf("Done")
-}
diff --git a/.test-infra/tools/stale_dataflow_jobs_cleaner.sh b/.test-infra/tools/stale_dataflow_jobs_cleaner.sh
new file mode 100755
index 0000000..e563048
--- /dev/null
+++ b/.test-infra/tools/stale_dataflow_jobs_cleaner.sh
@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Cancels active Dataflow jobs older than 3 hours.
+#
+set -euo pipefail
+
+STALE_JOBS=$(gcloud dataflow jobs list --created-before=-P3H \
+--format='value(JOB_ID)' --status=active --region=us-central1)
+
+if [[ ${STALE_JOBS} ]]; then
+ gcloud dataflow jobs cancel --region=us-central1 ${STALE_JOBS}
+else
+ echo "No stale jobs found."
+fi
diff --git a/.test-infra/tools/stale_dataflow_jobs_cleaner_test.go b/.test-infra/tools/stale_dataflow_jobs_cleaner_test.go
deleted file mode 100644
index 342052a..0000000
--- a/.test-infra/tools/stale_dataflow_jobs_cleaner_test.go
+++ /dev/null
@@ -1,74 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * License); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an AS IS BASIS,
- * WITHOUT 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 (
- df "google.golang.org/api/dataflow/v1b3"
- "reflect"
- "testing"
- "time"
-)
-
-var (
- currentTime time.Time = time.Now()
- jobsReturned = []*df.Job{}
- cancelledJobs = []*df.Job{}
-)
-
-type fakeClient struct{}
-
-func (c fakeClient) ListJobs(projectId string) ([]*df.Job, error) {
- return jobsReturned, nil
-}
-
-func (c fakeClient) CancelJob(job *df.Job) error {
- cancelledJobs = append(cancelledJobs, job)
- return nil
-}
-
-func (c fakeClient) CurrentTime() time.Time {
- return currentTime
-}
-
-func helperForJobCancel(t *testing.T, hoursStale float64, jobList []*df.Job, expectedCancelled []*df.Job) {
- var c fakeClient
- jobsReturned = jobList
- cancelledJobs = []*df.Job{}
- cleanDataflowJobs(c, "some-project-id", 2.0)
- if !reflect.DeepEqual(cancelledJobs, expectedCancelled) {
- t.Errorf("Cancelled arrays not as expected actual=%v, expected=%v", cancelledJobs, expectedCancelled)
- }
-}
-
-func TestEmptyJobList(t *testing.T) {
- helperForJobCancel(t, 2.0, []*df.Job{}, []*df.Job{})
-}
-
-func TestNotExpiredJob(t *testing.T) {
- // Just under 2 hours.
- createTime := currentTime.Add(-(2*time.Hour - time.Second))
- helperForJobCancel(t, 2.0, []*df.Job{&df.Job{CreateTime: createTime.Format(time.RFC3339)}}, []*df.Job{})
-}
-
-func TestExpiredJob(t *testing.T) {
- // Just over 2 hours.
- createTime := currentTime.Add(-(2*time.Hour + time.Second))
- job := &df.Job{CreateTime: createTime.Format(time.RFC3339)}
- helperForJobCancel(t, 2.0, []*df.Job{job}, []*df.Job{job})
-}
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
deleted file mode 100644
index 2784be4..0000000
--- a/CONTRIBUTING.md
+++ /dev/null
@@ -1,29 +0,0 @@
-<!--
- Licensed to the Apache Software Foundation (ASF) under one
- or more contributor license agreements. See the NOTICE file
- distributed with this work for additional information
- regarding copyright ownership. The ASF licenses this file
- to you under the Apache License, Version 2.0 (the
- "License"); you may not use this file except in compliance
- with the License. You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing,
- software distributed under the License is distributed on an
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- KIND, either express or implied. See the License for the
- specific language governing permissions and limitations
- under the License.
--->
-
-# How to Contribute
-
-The Apache Beam community welcomes contributions from anyone!
-Please see our [contribution guide](https://beam.apache.org/contribute/contribution-guide/)
-for details, such as:
-
-* Sharing your intent with the community
-* Development setup and testing your changes
-* Submitting a pull request and finding a reviewer
-
diff --git a/README.md b/README.md
index 8d7b9ee..2f30c9a 100644
--- a/README.md
+++ b/README.md
@@ -83,20 +83,8 @@
## Getting Started
-Please refer to the Quickstart[[Java](https://beam.apache.org/get-started/quickstart-java), [Python](https://beam.apache.org/get-started/quickstart-py), [Go](https://beam.apache.org/get-started/quickstart-go)] available on our website.
-
-If you'd like to build and install the whole project from the source distribution, you may need some additional tools installed
-in your system. In a Debian-based distribution:
-
-```
-sudo apt-get install \
- openjdk-8-jdk \
- python-setuptools \
- python-pip \
- virtualenv
-```
-
-Then please use the standard `./gradlew build` command.
+To learn how to write Beam pipelines, read the Quickstart for [[Java](https://beam.apache.org/get-started/quickstart-java), [Python](https://beam.apache.org/get-started/quickstart-py), or
+[Go](https://beam.apache.org/get-started/quickstart-go)] available on our website.
## Contact Us
@@ -106,11 +94,12 @@
* [Subscribe](mailto:dev-subscribe@beam.apache.org) or [mail](mailto:dev@beam.apache.org) the [dev@beam.apache.org](http://mail-archives.apache.org/mod_mbox/beam-dev/) list.
* Report issues on [JIRA](https://issues.apache.org/jira/browse/BEAM).
-We also have a [contributor's guide](https://beam.apache.org/contribute/contribution-guide/).
+Instructions for building and testing Beam itself
+are in the [contribution guide](https://beam.apache.org/contribute/).
## More Information
-* [Apache Beam](http://beam.apache.org)
-* [Overview](http://beam.apache.org/use/beam-overview/)
+* [Apache Beam](https://beam.apache.org)
+* [Overview](https://beam.apache.org/use/beam-overview/)
* Quickstart: [Java](https://beam.apache.org/get-started/quickstart-java), [Python](https://beam.apache.org/get-started/quickstart-py), [Go](https://beam.apache.org/get-started/quickstart-go)
* [Community metrics](https://s.apache.org/beam-community-metrics)
diff --git a/build.gradle b/build.gradle
index cc9152d..b63b08e 100644
--- a/build.gradle
+++ b/build.gradle
@@ -303,7 +303,7 @@
configurations { linkageCheckerJava }
dependencies {
- linkageCheckerJava "com.google.cloud.tools:dependencies:1.0.1"
+ linkageCheckerJava "com.google.cloud.tools:dependencies:1.1.2"
}
// We need to evaluate all the projects first so that we can find depend on all the
@@ -318,7 +318,14 @@
dependsOn project.getTasksByName('publishMavenJavaPublicationToMavenLocal', true /* recursively */)
classpath = project.configurations.linkageCheckerJava
main = 'com.google.cloud.tools.opensource.classpath.LinkageCheckerMain'
- args '-a', project.javaLinkageArtifactIds.split(',').collect({"${project.ext.mavenGroupId}:${it}:${project.version}"}).join(',')
+ args '-a', project.javaLinkageArtifactIds.split(',').collect({
+ if (it.contains(':')) {
+ "${project.ext.mavenGroupId}:${it}"
+ } else {
+ // specify the version if not provided
+ "${project.ext.mavenGroupId}:${it}:${project.version}"
+ }
+ }).join(',')
doLast {
println "NOTE: This task published artifacts into your local Maven repository. You may want to remove them manually."
}
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 1236f06..2c2aea9 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -297,7 +297,7 @@
// Automatically use the official release version if we are performing a release
// otherwise append '-SNAPSHOT'
- project.version = '2.19.0'
+ project.version = '2.20.0'
if (!isRelease(project)) {
project.version += '-SNAPSHOT'
}
@@ -362,17 +362,19 @@
def apex_malhar_version = "3.4.0"
def aws_java_sdk_version = "1.11.519"
def aws_java_sdk2_version = "2.5.71"
- def cassandra_driver_version = "3.6.0"
+ def cassandra_driver_version = "3.8.0"
def classgraph_version = "4.8.56"
+ def gax_version = "1.52.0"
def generated_grpc_beta_version = "0.44.0"
- def generated_grpc_ga_version = "1.43.0"
+ def generated_grpc_ga_version = "1.83.0"
def generated_grpc_dc_beta_version = "0.27.0-alpha"
- def google_auth_version = "0.12.0"
+ def google_auth_version = "0.19.0"
def google_clients_version = "1.28.0"
def google_cloud_bigdataoss_version = "1.9.16"
def google_cloud_core_version = "1.61.0"
def google_cloud_spanner_version = "1.6.0"
- def grpc_version = "1.17.1"
+ def google_http_clients_version = "1.34.0"
+ def grpc_version = "1.25.0"
def guava_version = "25.1-jre"
def hadoop_version = "2.8.5"
def hamcrest_version = "2.1"
@@ -387,7 +389,6 @@
def protobuf_version = "3.11.1"
def quickcheck_version = "0.8"
def spark_version = "2.4.4"
- def spark_structured_streaming_version = "2.4.0"
// A map of maps containing common libraries used per language. To use:
// dependencies {
@@ -430,13 +431,13 @@
classgraph : "io.github.classgraph:classgraph:$classgraph_version",
commons_codec : "commons-codec:commons-codec:1.14",
commons_compress : "org.apache.commons:commons-compress:1.19",
- commons_csv : "org.apache.commons:commons-csv:1.4",
- commons_io_1x : "commons-io:commons-io:1.3.2",
- commons_io_2x : "commons-io:commons-io:2.5",
- commons_lang3 : "org.apache.commons:commons-lang3:3.6",
+ commons_csv : "org.apache.commons:commons-csv:1.7",
+ commons_io : "commons-io:commons-io:2.6",
+ commons_lang3 : "org.apache.commons:commons-lang3:3.9",
commons_math3 : "org.apache.commons:commons-math3:3.6.1",
error_prone_annotations : "com.google.errorprone:error_prone_annotations:2.0.15",
- gax_grpc : "com.google.api:gax-grpc:1.38.0",
+ gax : "com.google.api:gax:$gax_version",
+ gax_grpc : "com.google.api:gax-grpc:$gax_version",
google_api_client : "com.google.api-client:google-api-client:$google_clients_version",
google_api_client_jackson2 : "com.google.api-client:google-api-client-jackson2:$google_clients_version",
google_api_client_java6 : "com.google.api-client:google-api-client-java6:$google_clients_version",
@@ -445,7 +446,7 @@
google_api_services_clouddebugger : "com.google.apis:google-api-services-clouddebugger:v2-rev20181114-$google_clients_version",
google_api_services_cloudresourcemanager : "com.google.apis:google-api-services-cloudresourcemanager:v1-rev20181015-$google_clients_version",
google_api_services_dataflow : "com.google.apis:google-api-services-dataflow:v1b3-rev20190927-$google_clients_version",
- google_api_services_pubsub : "com.google.apis:google-api-services-pubsub:v1-rev20181213-$google_clients_version",
+ google_api_services_pubsub : "com.google.apis:google-api-services-pubsub:v1-rev20191111-$google_clients_version",
google_api_services_storage : "com.google.apis:google-api-services-storage:v1-rev20181109-$google_clients_version",
google_auth_library_credentials : "com.google.auth:google-auth-library-credentials:$google_auth_version",
google_auth_library_oauth2_http : "com.google.auth:google-auth-library-oauth2-http:$google_auth_version",
@@ -457,15 +458,16 @@
google_cloud_dataflow_java_proto_library_all: "com.google.cloud.dataflow:google-cloud-dataflow-java-proto-library-all:0.5.160304",
google_cloud_datastore_v1_proto_client : "com.google.cloud.datastore:datastore-v1-proto-client:1.6.3",
google_cloud_spanner : "com.google.cloud:google-cloud-spanner:$google_cloud_spanner_version",
- google_http_client : "com.google.http-client:google-http-client:$google_clients_version",
- google_http_client_jackson : "com.google.http-client:google-http-client-jackson:$google_clients_version",
- google_http_client_jackson2 : "com.google.http-client:google-http-client-jackson2:$google_clients_version",
- google_http_client_protobuf : "com.google.http-client:google-http-client-protobuf:$google_clients_version",
+ google_http_client : "com.google.http-client:google-http-client:$google_http_clients_version",
+ google_http_client_jackson : "com.google.http-client:google-http-client-jackson:1.29.2",
+ google_http_client_jackson2 : "com.google.http-client:google-http-client-jackson2:$google_http_clients_version",
+ google_http_client_protobuf : "com.google.http-client:google-http-client-protobuf:$google_http_clients_version",
google_oauth_client : "com.google.oauth-client:google-oauth-client:$google_clients_version",
google_oauth_client_java6 : "com.google.oauth-client:google-oauth-client-java6:$google_clients_version",
grpc_all : "io.grpc:grpc-all:$grpc_version",
grpc_auth : "io.grpc:grpc-auth:$grpc_version",
grpc_core : "io.grpc:grpc-core:$grpc_version",
+ grpc_context : "io.grpc:grpc-context:$grpc_version",
grpc_google_cloud_datacatalog_v1beta1 : "com.google.api.grpc:grpc-google-cloud-datacatalog-v1beta1:$generated_grpc_dc_beta_version",
grpc_google_cloud_pubsub_v1 : "com.google.api.grpc:grpc-google-cloud-pubsub-v1:$generated_grpc_ga_version",
grpc_protobuf : "io.grpc:grpc-protobuf:$grpc_version",
@@ -526,7 +528,7 @@
spark_streaming : "org.apache.spark:spark-streaming_2.11:$spark_version",
stax2_api : "org.codehaus.woodstox:stax2-api:3.1.4",
vendored_bytebuddy_1_9_3 : "org.apache.beam:beam-vendor-bytebuddy-1_9_3:0.1",
- vendored_grpc_1_21_0 : "org.apache.beam:beam-vendor-grpc-1_21_0:0.1",
+ vendored_grpc_1_26_0 : "org.apache.beam:beam-vendor-grpc-1_26_0:0.1",
vendored_guava_26_0_jre : "org.apache.beam:beam-vendor-guava-26_0-jre:0.1",
vendored_calcite_1_20_0 : "org.apache.beam:beam-vendor-calcite-1_20_0:0.1",
woodstox_core_asl : "org.codehaus.woodstox:woodstox-core-asl:4.4.1",
@@ -733,7 +735,7 @@
// configurations because they are never required to be shaded or become a
// dependency of the output.
def compileOnlyAnnotationDeps = [
- "com.google.auto.value:auto-value-annotations:1.6.3",
+ "com.google.auto.value:auto-value-annotations:1.7",
"com.google.auto.service:auto-service-annotations:1.0-rc6",
"com.google.j2objc:j2objc-annotations:1.3",
// These dependencies are needed to avoid error-prone warnings on package-info.java files,
@@ -760,7 +762,7 @@
// Add common annotation processors to all Java projects
def annotationProcessorDeps = [
- "com.google.auto.value:auto-value:1.6.3",
+ "com.google.auto.value:auto-value:1.7",
"com.google.auto.service:auto-service:1.0-rc6",
]
@@ -829,7 +831,7 @@
project.apply plugin: 'com.github.spotbugs'
project.dependencies {
spotbugs "com.github.spotbugs:spotbugs:3.1.12"
- spotbugs "com.google.auto.value:auto-value:1.6.3"
+ spotbugs "com.google.auto.value:auto-value:1.7"
compileOnlyAnnotationDeps.each { dep -> spotbugs dep }
}
project.spotbugs {
@@ -956,7 +958,7 @@
FileTree exposedClasses = project.zipTree(it).matching {
include "**/*.class"
// BEAM-5919: Exclude paths for Java 9 multi-release jars.
- exclude "META-INF/versions/*/module-info.class"
+ exclude "**/module-info.class"
configuration.shadowJarValidationExcludes.each {
exclude "$it"
exclude "META-INF/versions/*/$it"
@@ -1434,19 +1436,21 @@
/** ***********************************************************************************************/
+ // applyGrpcNature should only be applied to projects who wish to use
+ // unvendored gRPC / protobuf dependencies.
project.ext.applyGrpcNature = {
project.apply plugin: "com.google.protobuf"
project.protobuf {
protoc {
// The artifact spec for the Protobuf Compiler
- artifact = "com.google.protobuf:protoc:3.6.0" }
+ artifact = "com.google.protobuf:protoc:$protobuf_version" }
// Configure the codegen plugins
plugins {
// An artifact spec for a protoc plugin, with "grpc" as
// the identifier, which can be referred to in the "plugins"
// container of the "generateProtoTasks" closure.
- grpc { artifact = "io.grpc:protoc-gen-grpc-java:1.13.1" }
+ grpc { artifact = "io.grpc:protoc-gen-grpc-java:$grpc_version" }
}
generateProtoTasks {
@@ -1483,6 +1487,8 @@
/** ***********************************************************************************************/
+ // applyPortabilityNature should only be applied to projects that want to use
+ // vendored gRPC / protobuf dependencies.
project.ext.applyPortabilityNature = {
PortabilityNatureConfiguration configuration = it ? it as PortabilityNatureConfiguration : new PortabilityNatureConfiguration()
@@ -1497,10 +1503,10 @@
archivesBaseName: configuration.archivesBaseName,
automaticModuleName: configuration.automaticModuleName,
shadowJarValidationExcludes: it.shadowJarValidationExcludes,
- shadowClosure: GrpcVendoring.shadowClosure() << {
+ shadowClosure: GrpcVendoring_1_26_0.shadowClosure() << {
// We perform all the code relocations but don't include
// any of the actual dependencies since they will be supplied
- // by org.apache.beam:beam-vendor-grpc-v1p21p0:0.1
+ // by org.apache.beam:beam-vendor-grpc-v1p26p0:0.1
dependencies {
include(dependency { return false })
}
@@ -1517,14 +1523,14 @@
project.protobuf {
protoc {
// The artifact spec for the Protobuf Compiler
- artifact = "com.google.protobuf:protoc:3.7.1" }
+ artifact = "com.google.protobuf:protoc:${GrpcVendoring_1_26_0.protobuf_version}" }
// Configure the codegen plugins
plugins {
// An artifact spec for a protoc plugin, with "grpc" as
// the identifier, which can be referred to in the "plugins"
// container of the "generateProtoTasks" closure.
- grpc { artifact = "io.grpc:protoc-gen-grpc-java:1.21.0" }
+ grpc { artifact = "io.grpc:protoc-gen-grpc-java:${GrpcVendoring_1_26_0.grpc_version}" }
}
generateProtoTasks {
@@ -1538,7 +1544,7 @@
}
}
- project.dependencies GrpcVendoring.dependenciesClosure() << { shadow project.ext.library.java.vendored_grpc_1_21_0 }
+ project.dependencies GrpcVendoring_1_26_0.dependenciesClosure() << { shadow project.ext.library.java.vendored_grpc_1_26_0 }
}
/** ***********************************************************************************************/
@@ -1808,15 +1814,17 @@
// distribution tarball generated by :sdks:python:sdist.
project.configurations { distTarBall }
- project.task('installGcpTest', dependsOn: 'setupVirtualenv') {
+ project.task('installGcpTest') {
+ dependsOn 'setupVirtualenv'
+ dependsOn ':sdks:python:sdist'
doLast {
+ def distTarBall = "${pythonRootDir}/build/apache-beam.tar.gz"
project.exec {
executable 'sh'
- args '-c', ". ${project.ext.envdir}/bin/activate && pip install --retries 10 -e ${pythonRootDir}/[gcp,test]"
+ args '-c', ". ${project.ext.envdir}/bin/activate && pip install --retries 10 ${distTarBall}[gcp,test]"
}
}
}
- project.installGcpTest.mustRunAfter project.configurations.distTarBall
project.task('cleanPython') {
doLast {
@@ -1939,6 +1947,8 @@
"--parallelism=2",
"--shutdown_sources_on_final_watermark",
"--sdk_worker_parallelism=1",
+ "--flink_job_server_jar=${project.project(':runners:flink:1.9:job-server').shadowJar.archivePath}",
+ "--spark_job_server_jar=${project.project(':runners:spark:job-server').shadowJar.archivePath}",
]
if (isStreaming)
options += [
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/GrpcVendoring.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/GrpcVendoring.groovy
deleted file mode 100644
index 96c6bf8..0000000
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/GrpcVendoring.groovy
+++ /dev/null
@@ -1,150 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * License); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an AS IS BASIS,
- * WITHOUT 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.gradle
-
-import org.gradle.api.Project
-
-/**
- * Utilities for working with our vendored version of gRPC.
- */
-class GrpcVendoring {
- /** Returns the list of compile time dependencies. */
- static List<String> dependencies() {
- return [
- 'com.google.guava:guava:26.0-jre',
- 'com.google.protobuf:protobuf-java:3.7.1',
- 'com.google.protobuf:protobuf-java-util:3.7.1',
- 'com.google.code.gson:gson:2.7',
- 'io.grpc:grpc-auth:1.21.0',
- 'io.grpc:grpc-core:1.21.0',
- 'io.grpc:grpc-context:1.21.0',
- 'io.grpc:grpc-netty:1.21.0',
- 'io.grpc:grpc-protobuf:1.21.0',
- 'io.grpc:grpc-stub:1.21.0',
- 'io.netty:netty-transport-native-epoll:4.1.34.Final',
- // tcnative version from https://github.com/grpc/grpc-java/blob/master/SECURITY.md#netty
- 'io.netty:netty-tcnative-boringssl-static:2.0.22.Final',
- 'com.google.auth:google-auth-library-credentials:0.13.0',
- 'io.grpc:grpc-testing:1.21.0',
- 'com.google.api.grpc:proto-google-common-protos:1.12.0',
- 'io.opencensus:opencensus-api:0.21.0',
- 'io.opencensus:opencensus-contrib-grpc-metrics:0.21.0',
- ]
- }
-
- /**
- * Returns the list of runtime time dependencies that should be exported as runtime
- * dependencies within the vendored jar.
- */
- static List<String> runtimeDependencies() {
- return [
- 'com.google.errorprone:error_prone_annotations:2.3.2',
- ]
- }
-
- static Map<String, String> relocations() {
- // The relocation paths below specifically use gRPC and the full version string as
- // the code relocation prefix. See https://lists.apache.org/thread.html/4c12db35b40a6d56e170cd6fc8bb0ac4c43a99aa3cb7dbae54176815@%3Cdev.beam.apache.org%3E
- // for further details.
-
- // To produce the list of necessary relocations, one needs to start with a set of target
- // packages that one wants to vendor, find all necessary transitive dependencies of that
- // set and provide relocations for each such that all necessary packages and their
- // dependencies are relocated. Any optional dependency that doesn't need relocation
- // must be excluded via an 'exclude' rule. There is additional complexity of libraries that use
- // JNI or reflection and have to be handled on case by case basis by learning whether
- // they support relocation and how would one go about doing it by reading any documentation
- // those libraries may provide. The 'validateShadedJarDoesntLeakNonOrgApacheBeamClasses'
- // ensures that there are no classes outside of the 'org.apache.beam' namespace.
-
- String version = "v1p21p0";
- String prefix = "org.apache.beam.vendor.grpc.${version}";
- List<String> packagesToRelocate = [
- // guava uses the com.google.common and com.google.thirdparty package namespaces
- "com.google.common",
- "com.google.thirdparty",
- "com.google.protobuf",
- "com.google.gson",
- "io.grpc",
- "com.google.auth",
- "com.google.api",
- "com.google.cloud",
- "com.google.logging",
- "com.google.longrunning",
- "com.google.rpc",
- "com.google.type",
- "io.opencensus",
- "io.netty"
- ]
-
- return packagesToRelocate.collectEntries {
- [ (it): "${prefix}.${it}" ]
- } + [
- // Adapted from https://github.com/grpc/grpc-java/blob/e283f70ad91f99c7fee8b31b605ef12a4f9b1690/netty/shaded/build.gradle#L41
- // We "io.netty": "${prefix}.io.netty",have to be careful with these replacements as they must not match any
- // string in NativeLibraryLoader, else they cause corruption. Note that
- // this includes concatenation of string literals and constants.
- 'META-INF/native/libnetty': "META-INF/native/liborg_apache_beam_vendor_grpc_${version}_netty",
- 'META-INF/native/netty': "META-INF/native/org_apache_beam_vendor_grpc_${version}_netty",
- ]
- }
-
- /** Returns the list of shading exclusions. */
- static List<String> exclusions() {
- return [
- // Don't include android annotations, errorprone, checkerframework, JDK8 annotations, objenesis, junit, and mockito in the vendored jar
- "android/annotation/**/",
- "com/google/errorprone/**",
- "com/google/instrumentation/**",
- "com/google/j2objc/annotations/**",
- "javax/annotation/**",
- "junit/**",
- "org/checkerframework/**",
- "org/codehaus/mojo/animal_sniffer/**",
- "org/hamcrest/**",
- "org/junit/**",
- "org/mockito/**",
- "org/objenesis/**",
- ]
- }
-
- /**
- * Returns a closure contaning the dependencies map used for shading gRPC within the main
- * Apache Beam project.
- */
- static Object dependenciesClosure() {
- return {
- dependencies().each { compile it }
- runtimeDependencies().each { shadow it }
- }
- }
-
- /**
- * Returns a closure with the code relocation configuration for shading gRPC within the main
- * Apache Beam project.
- */
- static Object shadowClosure() {
- return {
- relocations().each { srcNamespace, destNamespace ->
- relocate srcNamespace, destNamespace
- }
- exclusions().each { exclude it }
- }
- }
-}
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/GrpcVendoring_1_26_0.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/GrpcVendoring_1_26_0.groovy
new file mode 100644
index 0000000..de87bdf
--- /dev/null
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/GrpcVendoring_1_26_0.groovy
@@ -0,0 +1,217 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT 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.gradle
+
+import org.gradle.api.Project
+
+/**
+ * Utilities for working with our vendored version of gRPC.
+ */
+class GrpcVendoring_1_26_0 {
+
+ static def guava_version = "26.0-jre"
+ static def protobuf_version = "3.11.0"
+ static def grpc_version = "1.26.0"
+ static def gson_version = "2.8.6"
+ static def netty_version = "4.1.42.Final"
+ static def google_auth_version = "0.18.0"
+ static def proto_google_common_protos_version = "1.12.0"
+ static def opencensus_version = "0.24.0"
+ static def perfmark_version = "0.19.0"
+ static def lzma_java_version = "1.3"
+ static def protobuf_javanano_version = "3.0.0-alpha-5"
+ static def jzlib_version = "1.1.3"
+ static def compress_lzf_version = "1.0.3"
+ static def lz4_version = "1.3.0"
+ static def bouncycastle_version = "1.54"
+ static def conscrypt_version = "1.3.0"
+ static def alpn_api_version = "1.1.2.v20150522"
+ static def npn_api_version = "1.1.1.v20141010"
+ static def jboss_marshalling_version = "1.4.11.Final"
+ static def jboss_modules_version = "1.1.0.Beta1"
+
+ /** Returns the list of compile time dependencies. */
+ static List<String> dependencies() {
+ return [
+ "com.google.guava:guava:$guava_version",
+ "com.google.protobuf:protobuf-java:$protobuf_version",
+ "com.google.protobuf:protobuf-java-util:$protobuf_version",
+ "com.google.code.gson:gson:$gson_version",
+ "io.grpc:grpc-auth:$grpc_version",
+ "io.grpc:grpc-core:$grpc_version",
+ "io.grpc:grpc-context:$grpc_version",
+ "io.grpc:grpc-netty:$grpc_version",
+ "io.grpc:grpc-protobuf:$grpc_version",
+ "io.grpc:grpc-stub:$grpc_version",
+ "io.netty:netty-transport-native-epoll:$netty_version",
+ // tcnative version from https://github.com/grpc/grpc-java/blob/master/SECURITY.md#netty
+ "io.netty:netty-tcnative-boringssl-static:2.0.26.Final",
+ "com.google.auth:google-auth-library-credentials:$google_auth_version",
+ "io.grpc:grpc-testing:$grpc_version",
+ "com.google.api.grpc:proto-google-common-protos:$proto_google_common_protos_version",
+ "io.opencensus:opencensus-api:$opencensus_version",
+ "io.opencensus:opencensus-contrib-grpc-metrics:$opencensus_version",
+ "io.perfmark:perfmark-api:$perfmark_version",
+ "com.github.jponge:lzma-java:$lzma_java_version",
+ "com.google.protobuf.nano:protobuf-javanano:$protobuf_javanano_version",
+ "com.jcraft:jzlib:$jzlib_version",
+ "com.ning:compress-lzf:$compress_lzf_version",
+ "net.jpountz.lz4:lz4:$lz4_version",
+ "org.bouncycastle:bcpkix-jdk15on:$bouncycastle_version",
+ "org.bouncycastle:bcprov-jdk15on:$bouncycastle_version",
+ "org.conscrypt:conscrypt-openjdk-uber:$conscrypt_version",
+ "org.eclipse.jetty.alpn:alpn-api:$alpn_api_version",
+ "org.eclipse.jetty.npn:npn-api:$npn_api_version",
+ "org.jboss.marshalling:jboss-marshalling:$jboss_marshalling_version",
+ "org.jboss.modules:jboss-modules:$jboss_modules_version"
+ ]
+ }
+
+ /**
+ * Returns the list of runtime time dependencies that should be exported as runtime
+ * dependencies within the vendored jar.
+ */
+ static List<String> runtimeDependencies() {
+ return [
+ 'com.google.errorprone:error_prone_annotations:2.3.3',
+ 'commons-logging:commons-logging:1.2',
+ 'org.apache.logging.log4j:log4j-api:2.6.2',
+ 'org.slf4j:slf4j-api:1.7.21'
+ ]
+ }
+
+ /**
+ * Returns the list of test dependencies.
+ */
+ static List<String> testDependencies() {
+ return [
+ 'junit:junit:4.12',
+ ]
+ }
+
+ static Map<String, String> relocations() {
+ // The relocation paths below specifically use gRPC and the full version string as
+ // the code relocation prefix. See https://lists.apache.org/thread.html/4c12db35b40a6d56e170cd6fc8bb0ac4c43a99aa3cb7dbae54176815@%3Cdev.beam.apache.org%3E
+ // for further details.
+
+ // To produce the list of necessary relocations, one needs to start with a set of target
+ // packages that one wants to vendor, find all necessary transitive dependencies of that
+ // set and provide relocations for each such that all necessary packages and their
+ // dependencies are relocated. Any optional dependency that doesn't need relocation
+ // must be excluded via an 'exclude' rule. There is additional complexity of libraries that use
+ // JNI or reflection and have to be handled on case by case basis by learning whether
+ // they support relocation and how would one go about doing it by reading any documentation
+ // those libraries may provide. The 'validateShadedJarDoesntLeakNonOrgApacheBeamClasses'
+ // ensures that there are no classes outside of the 'org.apache.beam' namespace.
+
+ String version = "v1p26p0";
+ String prefix = "org.apache.beam.vendor.grpc.${version}";
+ List<String> packagesToRelocate = [
+ // guava uses the com.google.common and com.google.thirdparty package namespaces
+ "com.google.common",
+ "com.google.thirdparty",
+ "com.google.protobuf",
+ "com.google.gson",
+ "com.google.auth",
+ "com.google.api",
+ "com.google.cloud",
+ "com.google.logging",
+ "com.google.longrunning",
+ "com.google.rpc",
+ "com.google.type",
+ "io.grpc",
+ "io.netty",
+ "io.opencensus",
+ "io.perfmark",
+ "com.google.protobuf.nano",
+ "com.jcraft",
+ "com.ning",
+ "com.sun",
+ "lzma",
+ "net.jpountz",
+ "org.bouncycastle",
+ "org.cservenak.streams",
+ "org.conscrypt",
+ "org.eclipse.jetty.alpn",
+ "org.eclipse.jetty.npn",
+ "org.jboss.marshalling",
+ "org.jboss.modules"
+ ]
+
+ return packagesToRelocate.collectEntries {
+ [ (it): "${prefix}.${it}" ]
+ } + [
+ // Adapted from https://github.com/grpc/grpc-java/blob/e283f70ad91f99c7fee8b31b605ef12a4f9b1690/netty/shaded/build.gradle#L41
+ // We "io.netty": "${prefix}.io.netty",have to be careful with these replacements as they must not match any
+ // string in NativeLibraryLoader, else they cause corruption. Note that
+ // this includes concatenation of string literals and constants.
+ 'META-INF/native/libnetty': "META-INF/native/liborg_apache_beam_vendor_grpc_${version}_netty",
+ 'META-INF/native/netty': "META-INF/native/org_apache_beam_vendor_grpc_${version}_netty",
+ ]
+ }
+
+ /** Returns the list of shading exclusions. */
+ static List<String> exclusions() {
+ return [
+ // Don't include android annotations, errorprone, checkerframework, JDK8 annotations, objenesis, junit,
+ // commons-logging, log4j, slf4j and mockito in the vendored jar
+ "android/annotation/**/",
+ "com/google/errorprone/**",
+ "com/google/instrumentation/**",
+ "com/google/j2objc/annotations/**",
+ "javax/annotation/**",
+ "junit/**",
+ "org/apache/commons/logging/**",
+ "org/apache/log/**",
+ "org/apache/log4j/**",
+ "org/apache/logging/log4j/**",
+ "org/checkerframework/**",
+ "org/codehaus/mojo/animal_sniffer/**",
+ "org/hamcrest/**",
+ "org/junit/**",
+ "org/mockito/**",
+ "org/objenesis/**",
+ "org/slf4j/**",
+ ]
+ }
+
+ /**
+ * Returns a closure contaning the dependencies map used for shading gRPC within the main
+ * Apache Beam project.
+ */
+ static Object dependenciesClosure() {
+ return {
+ dependencies().each { compile it }
+ runtimeDependencies().each { shadow it }
+ }
+ }
+
+ /**
+ * Returns a closure with the code relocation configuration for shading gRPC within the main
+ * Apache Beam project.
+ */
+ static Object shadowClosure() {
+ return {
+ relocations().each { srcNamespace, destNamespace ->
+ relocate srcNamespace, destNamespace
+ }
+ exclusions().each { exclude it }
+ }
+ }
+}
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/VendorJavaPlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/VendorJavaPlugin.groovy
index 24ca6e1..77c3019 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/VendorJavaPlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/VendorJavaPlugin.groovy
@@ -21,6 +21,7 @@
import org.gradle.api.GradleException
import org.gradle.api.Plugin
import org.gradle.api.Project
+import org.gradle.api.artifacts.ProjectDependency
import org.gradle.api.file.FileTree
import org.gradle.api.publish.maven.MavenPublication
@@ -51,6 +52,7 @@
static class VendorJavaPluginConfig {
List<String> dependencies
List<String> runtimeDependencies
+ List<String> testDependencies
Map<String, String> relocations
List<String> exclusions
String groupId
@@ -96,7 +98,8 @@
project.dependencies {
config.dependencies.each { compile it }
- config.runtimeDependencies.each { runtime it }
+ config.runtimeDependencies.each { runtimeOnly it }
+ config.testDependencies.each { compileOnly it}
}
// Create a task which emulates the maven-archiver plugin in generating a
@@ -132,7 +135,7 @@
include "**/*.class"
exclude "org/apache/beam/vendor/**"
// BEAM-5919: Exclude paths for Java 9 multi-release jars.
- exclude "META-INF/versions/*/module-info.class"
+ exclude "**/module-info.class"
exclude "META-INF/versions/*/org/apache/beam/vendor/**"
}
if (exposedClasses.files) {
@@ -281,6 +284,55 @@
}
pom.withXml {
+ def root = asNode()
+ def dependenciesNode = root.appendNode('dependencies')
+ def generateDependenciesFromConfiguration = { param ->
+ project.configurations."${param.configuration}".allDependencies.each {
+ def dependencyNode = dependenciesNode.appendNode('dependency')
+ def appendClassifier = { dep ->
+ dep.artifacts.each { art ->
+ if (art.hasProperty('classifier')) {
+ dependencyNode.appendNode('classifier', art.classifier)
+ }
+ }
+ }
+
+ if (it instanceof ProjectDependency) {
+ dependencyNode.appendNode('groupId', it.getDependencyProject().mavenGroupId)
+ dependencyNode.appendNode('artifactId', it.getDependencyProject().archivesBaseName)
+ dependencyNode.appendNode('version', it.version)
+ dependencyNode.appendNode('scope', param.scope)
+ appendClassifier(it)
+ } else {
+ dependencyNode.appendNode('groupId', it.group)
+ dependencyNode.appendNode('artifactId', it.name)
+ dependencyNode.appendNode('version', it.version)
+ dependencyNode.appendNode('scope', param.scope)
+ appendClassifier(it)
+ }
+
+ // Start with any exclusions that were added via configuration exclude rules.
+ // Then add all the exclusions that are specific to the dependency (if any
+ // were declared). Finally build the node that represents all exclusions.
+ def exclusions = []
+ exclusions += project.configurations."${param.configuration}".excludeRules
+ if (it.hasProperty('excludeRules')) {
+ exclusions += it.excludeRules
+ }
+ if (!exclusions.empty) {
+ def exclusionsNode = dependencyNode.appendNode('exclusions')
+ exclusions.each { exclude ->
+ def exclusionNode = exclusionsNode.appendNode('exclusion')
+ exclusionNode.appendNode('groupId', exclude.group)
+ exclusionNode.appendNode('artifactId', exclude.module)
+ }
+ }
+ }
+ }
+
+ generateDependenciesFromConfiguration(configuration: 'runtimeOnly', scope: 'runtime')
+ generateDependenciesFromConfiguration(configuration: 'compileOnly', scope: 'provided')
+
// NB: This must come after asNode() logic, as it seems asNode()
// removes XML comments.
// TODO: Load this from file?
diff --git a/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java b/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java
index a758417..bfb11ae 100644
--- a/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java
+++ b/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.examples;
+import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import java.util.ArrayList;
@@ -28,7 +29,6 @@
import java.util.concurrent.ThreadLocalRandom;
import org.apache.beam.examples.common.ExampleUtils;
import org.apache.beam.examples.common.WriteOneFilePerWindow.PerWindowFiles;
-import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
@@ -74,6 +74,10 @@
.withInitialBackoff(DEFAULT_SLEEP_DURATION)
.withMaxRetries(MAX_READ_RETRIES);
+ public static WordCountsMatcher containsWordCounts(SortedMap<String, Long> expectedWordCounts) {
+ return new WordCountsMatcher(expectedWordCounts);
+ }
+
/** Options for the {@link WindowedWordCount} Integration Test. */
public interface WindowedWordCountITOptions
extends WindowedWordCount.Options, TestPipelineOptions, StreamingOptions {}
@@ -183,30 +187,27 @@
}
}
- options.setOnSuccessMatcher(new WordCountsMatcher(expectedWordCounts, expectedOutputFiles));
-
WindowedWordCount.runWindowedWordCount(options);
+
+ assertThat(expectedOutputFiles, containsWordCounts(expectedWordCounts));
}
/**
* A matcher that bakes in expected word counts, so they can be read directly via some other
* mechanism, and compares a sharded output file with the result.
*/
- private static class WordCountsMatcher extends TypeSafeMatcher<PipelineResult>
- implements SerializableMatcher<PipelineResult> {
+ private static class WordCountsMatcher extends TypeSafeMatcher<List<ShardedFile>>
+ implements SerializableMatcher<List<ShardedFile>> {
private final SortedMap<String, Long> expectedWordCounts;
- private final List<ShardedFile> outputFiles;
private SortedMap<String, Long> actualCounts;
- public WordCountsMatcher(
- SortedMap<String, Long> expectedWordCounts, List<ShardedFile> outputFiles) {
+ private WordCountsMatcher(SortedMap<String, Long> expectedWordCounts) {
this.expectedWordCounts = expectedWordCounts;
- this.outputFiles = outputFiles;
}
@Override
- public boolean matchesSafely(PipelineResult pipelineResult) {
+ public boolean matchesSafely(List<ShardedFile> outputFiles) {
try {
// Load output data
List<String> outputLines = new ArrayList<>();
@@ -238,7 +239,7 @@
}
@Override
- public void describeMismatchSafely(PipelineResult pResult, Description description) {
+ public void describeMismatchSafely(List<ShardedFile> shardedFiles, Description description) {
equalTo(expectedWordCounts).describeMismatch(actualCounts, description);
}
}
diff --git a/examples/java/src/test/java/org/apache/beam/examples/WordCountIT.java b/examples/java/src/test/java/org/apache/beam/examples/WordCountIT.java
index 836ea08d..11d3d78 100644
--- a/examples/java/src/test/java/org/apache/beam/examples/WordCountIT.java
+++ b/examples/java/src/test/java/org/apache/beam/examples/WordCountIT.java
@@ -17,14 +17,17 @@
*/
package org.apache.beam.examples;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
+import static org.hamcrest.MatcherAssert.assertThat;
+
import java.util.Date;
import org.apache.beam.examples.WordCount.WordCountOptions;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.FileChecksumMatcher;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.apache.beam.sdk.util.NumberedShardedFile;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -63,9 +66,9 @@
.resolve("output", StandardResolveOptions.RESOLVE_DIRECTORY)
.resolve("results", StandardResolveOptions.RESOLVE_FILE)
.toString());
- options.setOnSuccessMatcher(
- new FileChecksumMatcher(DEFAULT_OUTPUT_CHECKSUM, options.getOutput() + "*-of-*"));
-
WordCount.runWordCount(options);
+ assertThat(
+ new NumberedShardedFile(options.getOutput() + "*-of-*"),
+ fileContentsHaveChecksum(DEFAULT_OUTPUT_CHECKSUM));
}
}
diff --git a/examples/java/src/test/java/org/apache/beam/examples/complete/TfIdfIT.java b/examples/java/src/test/java/org/apache/beam/examples/complete/TfIdfIT.java
index e6701c8..64fdd79 100644
--- a/examples/java/src/test/java/org/apache/beam/examples/complete/TfIdfIT.java
+++ b/examples/java/src/test/java/org/apache/beam/examples/complete/TfIdfIT.java
@@ -17,15 +17,18 @@
*/
package org.apache.beam.examples.complete;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
+import static org.hamcrest.MatcherAssert.assertThat;
+
import java.util.Date;
import java.util.regex.Pattern;
import org.apache.beam.examples.complete.TfIdf.Options;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.FileChecksumMatcher;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.apache.beam.sdk.util.NumberedShardedFile;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -64,10 +67,10 @@
.resolve("output", StandardResolveOptions.RESOLVE_DIRECTORY)
.resolve("results", StandardResolveOptions.RESOLVE_FILE)
.toString());
- options.setOnSuccessMatcher(
- new FileChecksumMatcher(
- EXPECTED_OUTPUT_CHECKSUM, options.getOutput() + "*-of-*.csv", DEFAULT_SHARD_TEMPLATE));
-
TfIdf.runTfIdf(options);
+
+ assertThat(
+ new NumberedShardedFile(options.getOutput() + "*-of-*.csv", DEFAULT_SHARD_TEMPLATE),
+ fileContentsHaveChecksum(EXPECTED_OUTPUT_CHECKSUM));
}
}
diff --git a/examples/java/src/test/java/org/apache/beam/examples/complete/TopWikipediaSessionsIT.java b/examples/java/src/test/java/org/apache/beam/examples/complete/TopWikipediaSessionsIT.java
index 9eb9a86..765fed5 100644
--- a/examples/java/src/test/java/org/apache/beam/examples/complete/TopWikipediaSessionsIT.java
+++ b/examples/java/src/test/java/org/apache/beam/examples/complete/TopWikipediaSessionsIT.java
@@ -17,13 +17,16 @@
*/
package org.apache.beam.examples.complete;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
+import static org.hamcrest.MatcherAssert.assertThat;
+
import java.util.Date;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.FileChecksumMatcher;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.apache.beam.sdk.util.NumberedShardedFile;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -60,9 +63,11 @@
.resolve("output", StandardResolveOptions.RESOLVE_DIRECTORY)
.resolve("results", StandardResolveOptions.RESOLVE_FILE)
.toString());
- options.setOnSuccessMatcher(
- new FileChecksumMatcher(DEFAULT_OUTPUT_CHECKSUM, options.getOutput() + "*-of-*"));
TopWikipediaSessions.run(options);
+
+ assertThat(
+ new NumberedShardedFile(options.getOutput() + "*-of-*"),
+ fileContentsHaveChecksum(DEFAULT_OUTPUT_CHECKSUM));
}
}
diff --git a/examples/java/src/test/java/org/apache/beam/examples/cookbook/BigQueryTornadoesIT.java b/examples/java/src/test/java/org/apache/beam/examples/cookbook/BigQueryTornadoesIT.java
index c53f49b..0825478 100644
--- a/examples/java/src/test/java/org/apache/beam/examples/cookbook/BigQueryTornadoesIT.java
+++ b/examples/java/src/test/java/org/apache/beam/examples/cookbook/BigQueryTornadoesIT.java
@@ -17,6 +17,8 @@
*/
package org.apache.beam.examples.cookbook;
+import static org.hamcrest.MatcherAssert.assertThat;
+
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryOptions;
import org.apache.beam.sdk.io.gcp.testing.BigqueryMatcher;
@@ -63,11 +65,11 @@
private void runE2EBigQueryTornadoesTest(BigQueryTornadoesITOptions options) throws Exception {
String query = String.format("SELECT month, tornado_count FROM [%s]", options.getOutput());
- options.setOnSuccessMatcher(
- new BigqueryMatcher(
- options.getAppName(), options.getProject(), query, DEFAULT_OUTPUT_CHECKSUM));
-
BigQueryTornadoes.runBigQueryTornadoes(options);
+
+ assertThat(
+ BigqueryMatcher.createQuery(options.getAppName(), options.getProject(), query),
+ BigqueryMatcher.queryResultHasChecksum(DEFAULT_OUTPUT_CHECKSUM));
}
@Test
diff --git a/examples/kotlin/src/test/java/org/apache/beam/examples/WindowedWordCountITKotlin.kt b/examples/kotlin/src/test/java/org/apache/beam/examples/WindowedWordCountITKotlin.kt
index ee2ec23..9298bfe 100644
--- a/examples/kotlin/src/test/java/org/apache/beam/examples/WindowedWordCountITKotlin.kt
+++ b/examples/kotlin/src/test/java/org/apache/beam/examples/WindowedWordCountITKotlin.kt
@@ -35,6 +35,7 @@
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.collect.ImmutableList
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists
+import org.hamcrest.MatcherAssert.assertThat
import org.hamcrest.Description
import org.hamcrest.Matchers.equalTo
import org.hamcrest.TypeSafeMatcher
diff --git a/examples/kotlin/src/test/java/org/apache/beam/examples/WordCountITKotlin.java b/examples/kotlin/src/test/java/org/apache/beam/examples/WordCountITKotlin.java
index ef016b9..dba629c 100644
--- a/examples/kotlin/src/test/java/org/apache/beam/examples/WordCountITKotlin.java
+++ b/examples/kotlin/src/test/java/org/apache/beam/examples/WordCountITKotlin.java
@@ -17,15 +17,18 @@
*/
package org.apache.beam.examples;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
+import static org.hamcrest.MatcherAssert.assertThat;
+
import java.util.Date;
import org.apache.beam.examples.kotlin.WordCount;
import org.apache.beam.examples.kotlin.WordCount.WordCountOptions;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.FileChecksumMatcher;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.apache.beam.sdk.util.NumberedShardedFile;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -64,9 +67,11 @@
.resolve("output", StandardResolveOptions.RESOLVE_DIRECTORY)
.resolve("results", StandardResolveOptions.RESOLVE_FILE)
.toString());
- options.setOnSuccessMatcher(
- new FileChecksumMatcher(DEFAULT_OUTPUT_CHECKSUM, options.getOutput() + "*-of-*"));
WordCount.runWordCount(options);
+
+ assertThat(
+ new NumberedShardedFile(options.getOutput() + "*-of-*"),
+ fileContentsHaveChecksum(DEFAULT_OUTPUT_CHECKSUM));
}
}
diff --git a/examples/notebooks/get-started/try-apache-beam-java.ipynb b/examples/notebooks/get-started/try-apache-beam-java.ipynb
index 40d648a..101df82 100644
--- a/examples/notebooks/get-started/try-apache-beam-java.ipynb
+++ b/examples/notebooks/get-started/try-apache-beam-java.ipynb
@@ -593,8 +593,8 @@
"\n",
"> Task :runShadow\n",
"WARNING: An illegal reflective access operation has occurred\n",
- "WARNING: Illegal reflective access by org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.UnsafeUtil (file:/content/build/install/content-shadow/lib/WordCount.jar) to field java.nio.Buffer.address\n",
- "WARNING: Please consider reporting this to the maintainers of org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.UnsafeUtil\n",
+ "WARNING: Illegal reflective access by org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.UnsafeUtil (file:/content/build/install/content-shadow/lib/WordCount.jar) to field java.nio.Buffer.address\n",
+ "WARNING: Please consider reporting this to the maintainers of org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.UnsafeUtil\n",
"WARNING: Use --illegal-access=warn to enable warnings of further illegal reflective access operations\n",
"WARNING: All illegal access operations will be denied in a future release\n",
"Mar 04, 2019 11:00:24 PM org.apache.beam.sdk.io.FileBasedSource getEstimatedSizeBytes\n",
@@ -735,8 +735,8 @@
"\n",
">> java -jar WordCount.jar\n",
"WARNING: An illegal reflective access operation has occurred\n",
- "WARNING: Illegal reflective access by org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.UnsafeUtil (file:/content/WordCount.jar) to field java.nio.Buffer.address\n",
- "WARNING: Please consider reporting this to the maintainers of org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.UnsafeUtil\n",
+ "WARNING: Illegal reflective access by org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.UnsafeUtil (file:/content/WordCount.jar) to field java.nio.Buffer.address\n",
+ "WARNING: Please consider reporting this to the maintainers of org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.UnsafeUtil\n",
"WARNING: Use --illegal-access=warn to enable warnings of further illegal reflective access operations\n",
"WARNING: All illegal access operations will be denied in a future release\n",
"Mar 04, 2019 11:00:49 PM org.apache.beam.sdk.io.FileBasedSource getEstimatedSizeBytes\n",
@@ -981,8 +981,8 @@
"\n",
"> Task :runShadow\n",
"WARNING: An illegal reflective access operation has occurred\n",
- "WARNING: Illegal reflective access by org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.UnsafeUtil (file:/content/build/install/content-shadow/lib/WordCount.jar) to field java.nio.Buffer.address\n",
- "WARNING: Please consider reporting this to the maintainers of org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.UnsafeUtil\n",
+ "WARNING: Illegal reflective access by org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.UnsafeUtil (file:/content/build/install/content-shadow/lib/WordCount.jar) to field java.nio.Buffer.address\n",
+ "WARNING: Please consider reporting this to the maintainers of org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.UnsafeUtil\n",
"WARNING: Use --illegal-access=warn to enable warnings of further illegal reflective access operations\n",
"WARNING: All illegal access operations will be denied in a future release\n",
"Mar 04, 2019 11:01:26 PM org.apache.beam.sdk.io.FileBasedSource getEstimatedSizeBytes\n",
@@ -1096,4 +1096,4 @@
]
}
]
-}
\ No newline at end of file
+}
diff --git a/gradle.properties b/gradle.properties
index d758ae4..2b9eb79 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -23,7 +23,7 @@
signing.gnupg.executable=gpg
signing.gnupg.useLegacyGpg=true
-version=2.19.0-SNAPSHOT
-python_sdk_version=2.19.0.dev
+version=2.20.0-SNAPSHOT
+sdk_version=2.20.0.dev
javaVersion=1.8
diff --git a/learning/katas/python/Common Transforms/Aggregation/Count/tests.py b/learning/katas/python/Common Transforms/Aggregation/Count/tests.py
index a63ecd4..ccc96cb 100644
--- a/learning/katas/python/Common Transforms/Aggregation/Count/tests.py
+++ b/learning/katas/python/Common Transforms/Aggregation/Count/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -29,5 +30,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Common Transforms/Aggregation/Largest/tests.py b/learning/katas/python/Common Transforms/Aggregation/Largest/tests.py
index 1a89b9d..0fced6e 100644
--- a/learning/katas/python/Common Transforms/Aggregation/Largest/tests.py
+++ b/learning/katas/python/Common Transforms/Aggregation/Largest/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -29,5 +30,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Common Transforms/Aggregation/Mean/tests.py b/learning/katas/python/Common Transforms/Aggregation/Mean/tests.py
index cba3205..9e99d39 100644
--- a/learning/katas/python/Common Transforms/Aggregation/Mean/tests.py
+++ b/learning/katas/python/Common Transforms/Aggregation/Mean/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -29,5 +30,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Common Transforms/Aggregation/Smallest/tests.py b/learning/katas/python/Common Transforms/Aggregation/Smallest/tests.py
index 90d1223..5bfa80c 100644
--- a/learning/katas/python/Common Transforms/Aggregation/Smallest/tests.py
+++ b/learning/katas/python/Common Transforms/Aggregation/Smallest/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_answer_placeholders_text_deleted, test_is_not_empty
def test_output():
@@ -29,5 +30,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Common Transforms/Aggregation/Sum/tests.py b/learning/katas/python/Common Transforms/Aggregation/Sum/tests.py
index 9756b6d..e761cdf 100644
--- a/learning/katas/python/Common Transforms/Aggregation/Sum/tests.py
+++ b/learning/katas/python/Common Transforms/Aggregation/Sum/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -29,5 +30,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Common Transforms/Filter/Filter/tests.py b/learning/katas/python/Common Transforms/Filter/Filter/tests.py
index 40e7837..03487d2 100644
--- a/learning/katas/python/Common Transforms/Filter/Filter/tests.py
+++ b/learning/katas/python/Common Transforms/Filter/Filter/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_filter():
@@ -39,6 +41,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_filter()
test_output()
diff --git a/learning/katas/python/Common Transforms/Filter/ParDo/tests.py b/learning/katas/python/Common Transforms/Filter/ParDo/tests.py
index a91e9bc..b1d475b 100644
--- a/learning/katas/python/Common Transforms/Filter/ParDo/tests.py
+++ b/learning/katas/python/Common Transforms/Filter/ParDo/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -29,5 +30,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Core Transforms/Branching/Branching/tests.py b/learning/katas/python/Core Transforms/Branching/Branching/tests.py
index 8278966..de1fea6 100644
--- a/learning/katas/python/Core Transforms/Branching/Branching/tests.py
+++ b/learning/katas/python/Core Transforms/Branching/Branching/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -39,5 +40,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Core Transforms/CoGroupByKey/CoGroupByKey/tests.py b/learning/katas/python/Core Transforms/CoGroupByKey/CoGroupByKey/tests.py
index 59c8515..da12782 100644
--- a/learning/katas/python/Core Transforms/CoGroupByKey/CoGroupByKey/tests.py
+++ b/learning/katas/python/Core Transforms/CoGroupByKey/CoGroupByKey/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -33,5 +34,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Core Transforms/Combine/Combine PerKey/tests.py b/learning/katas/python/Core Transforms/Combine/Combine PerKey/tests.py
index 7890e0a..e804283 100644
--- a/learning/katas/python/Core Transforms/Combine/Combine PerKey/tests.py
+++ b/learning/katas/python/Core Transforms/Combine/Combine PerKey/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_combine_placeholders():
@@ -35,7 +37,7 @@
PLAYER_3 = 'Player 3'
answers = [str((PLAYER_1, 115)), str((PLAYER_2, 85)), str((PLAYER_3, 25))]
- print answers
+ print(answers)
if all(num in output for num in answers):
passed()
@@ -44,6 +46,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_combine_placeholders()
test_output()
diff --git a/learning/katas/python/Core Transforms/Combine/CombineFn/tests.py b/learning/katas/python/Core Transforms/Combine/CombineFn/tests.py
index 6e8ff93..656e5b3 100644
--- a/learning/katas/python/Core Transforms/Combine/CombineFn/tests.py
+++ b/learning/katas/python/Core Transforms/Combine/CombineFn/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_combine_placeholders():
@@ -39,6 +41,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_combine_placeholders()
test_output()
diff --git a/learning/katas/python/Core Transforms/Combine/Simple Function/tests.py b/learning/katas/python/Core Transforms/Combine/Simple Function/tests.py
index 70cb7f6..2d740d8 100644
--- a/learning/katas/python/Core Transforms/Combine/Simple Function/tests.py
+++ b/learning/katas/python/Core Transforms/Combine/Simple Function/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_combine_placeholders():
@@ -39,6 +41,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_combine_placeholders()
test_output()
diff --git a/learning/katas/python/Core Transforms/Composite Transform/Composite Transform/tests.py b/learning/katas/python/Core Transforms/Composite Transform/Composite Transform/tests.py
index 24477d1..cc7db80 100644
--- a/learning/katas/python/Core Transforms/Composite Transform/Composite Transform/tests.py
+++ b/learning/katas/python/Core Transforms/Composite Transform/Composite Transform/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_composite_expand_method():
@@ -39,6 +41,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_composite_expand_method()
test_output()
diff --git a/learning/katas/python/Core Transforms/Flatten/Flatten/tests.py b/learning/katas/python/Core Transforms/Flatten/Flatten/tests.py
index 285dc48..c2caa2e 100644
--- a/learning/katas/python/Core Transforms/Flatten/Flatten/tests.py
+++ b/learning/katas/python/Core Transforms/Flatten/Flatten/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_flatten():
@@ -39,6 +41,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_flatten()
test_output()
diff --git a/learning/katas/python/Core Transforms/GroupByKey/GroupByKey/tests.py b/learning/katas/python/Core Transforms/GroupByKey/GroupByKey/tests.py
index afc8e36..8f9ffd5 100644
--- a/learning/katas/python/Core Transforms/GroupByKey/GroupByKey/tests.py
+++ b/learning/katas/python/Core Transforms/GroupByKey/GroupByKey/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -31,5 +32,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/Core Transforms/Map/FlatMap/tests.py b/learning/katas/python/Core Transforms/Map/FlatMap/tests.py
index 55bd104..6eaaa64 100644
--- a/learning/katas/python/Core Transforms/Map/FlatMap/tests.py
+++ b/learning/katas/python/Core Transforms/Map/FlatMap/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_flatmap():
@@ -39,6 +41,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_flatmap()
test_output()
diff --git a/learning/katas/python/Core Transforms/Map/Map/tests.py b/learning/katas/python/Core Transforms/Map/Map/tests.py
index df1544f..52789ea 100644
--- a/learning/katas/python/Core Transforms/Map/Map/tests.py
+++ b/learning/katas/python/Core Transforms/Map/Map/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_map():
@@ -39,6 +41,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_map()
test_output()
diff --git a/learning/katas/python/Core Transforms/Map/ParDo OneToMany/tests.py b/learning/katas/python/Core Transforms/Map/ParDo OneToMany/tests.py
index 8f9fe86..b934821 100644
--- a/learning/katas/python/Core Transforms/Map/ParDo OneToMany/tests.py
+++ b/learning/katas/python/Core Transforms/Map/ParDo OneToMany/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_dofn_process_method():
@@ -49,7 +51,8 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_dofn_process_method()
test_pardo()
test_output()
diff --git a/learning/katas/python/Core Transforms/Map/ParDo/tests.py b/learning/katas/python/Core Transforms/Map/ParDo/tests.py
index f6dd859..5591318 100644
--- a/learning/katas/python/Core Transforms/Map/ParDo/tests.py
+++ b/learning/katas/python/Core Transforms/Map/ParDo/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_dofn_process_method():
@@ -49,7 +51,8 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_dofn_process_method()
test_pardo()
test_output()
diff --git a/learning/katas/python/Core Transforms/Partition/Partition/tests.py b/learning/katas/python/Core Transforms/Partition/Partition/tests.py
index 6c5568a..bbeeaf7 100644
--- a/learning/katas/python/Core Transforms/Partition/Partition/tests.py
+++ b/learning/katas/python/Core Transforms/Partition/Partition/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_partition():
@@ -48,6 +50,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_partition()
test_output()
diff --git a/learning/katas/python/Core Transforms/Side Input/Side Input/tests.py b/learning/katas/python/Core Transforms/Side Input/Side Input/tests.py
index 0d7146d..8fdd7da 100644
--- a/learning/katas/python/Core Transforms/Side Input/Side Input/tests.py
+++ b/learning/katas/python/Core Transforms/Side Input/Side Input/tests.py
@@ -22,8 +22,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, \
- get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_dofn_process_method():
@@ -64,7 +65,8 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_dofn_process_method()
test_pardo()
test_output()
diff --git a/learning/katas/python/Core Transforms/Side Output/Side Output/tests.py b/learning/katas/python/Core Transforms/Side Output/Side Output/tests.py
index b525cae..1af8439 100644
--- a/learning/katas/python/Core Transforms/Side Output/Side Output/tests.py
+++ b/learning/katas/python/Core Transforms/Side Output/Side Output/tests.py
@@ -14,8 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, \
- get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_dofn_process_method():
@@ -59,7 +60,8 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_dofn_process_method()
test_pardo()
test_output()
diff --git a/learning/katas/python/Examples/Word Count/Word Count/task-info.yaml b/learning/katas/python/Examples/Word Count/Word Count/task-info.yaml
index 0eef4b3..435527d 100644
--- a/learning/katas/python/Examples/Word Count/Word Count/task-info.yaml
+++ b/learning/katas/python/Examples/Word Count/Word Count/task-info.yaml
@@ -23,7 +23,7 @@
visible: true
placeholders:
- offset: 1021
- length: 133
+ length: 136
placeholder_text: TODO()
- name: tests.py
visible: false
diff --git a/learning/katas/python/Examples/Word Count/Word Count/task.html b/learning/katas/python/Examples/Word Count/Word Count/task.html
index a963aab..82ce81c 100644
--- a/learning/katas/python/Examples/Word Count/Word Count/task.html
+++ b/learning/katas/python/Examples/Word Count/Word Count/task.html
@@ -33,4 +33,8 @@
<div class="hint">
Refer to your katas above.
</div>
+<div class="hint">
+ Use <a href="https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.core.html#apache_beam.transforms.core.MapTuple">
+ MapTuple</a> to unpack key-value pair into different function arguments.
+</div>
</html>
diff --git a/learning/katas/python/Examples/Word Count/Word Count/task.py b/learning/katas/python/Examples/Word Count/Word Count/task.py
index dd1daeb..10b7cf8 100644
--- a/learning/katas/python/Examples/Word Count/Word Count/task.py
+++ b/learning/katas/python/Examples/Word Count/Word Count/task.py
@@ -29,7 +29,7 @@
| beam.FlatMap(lambda sentence: sentence.split())
| beam.combiners.Count.PerElement()
- | beam.Map(lambda (k, v): k + ":" + str(v))
+ | beam.MapTuple(lambda k, v: k + ":" + str(v))
| LogElements())
diff --git a/learning/katas/python/Examples/Word Count/Word Count/tests.py b/learning/katas/python/Examples/Word Count/Word Count/tests.py
index 2f0a085..16b7bf5 100644
--- a/learning/katas/python/Examples/Word Count/Word Count/tests.py
+++ b/learning/katas/python/Examples/Word Count/Word Count/tests.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_file_output
+from test_helper import failed, passed, get_file_output, \
+ test_is_not_empty, test_answer_placeholders_text_deleted
def test_output():
@@ -35,5 +36,6 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_output()
diff --git a/learning/katas/python/IO/Built-in IOs/Built-in IOs/task-info.yaml b/learning/katas/python/IO/Built-in IOs/Built-in IOs/task-info.yaml
index 45ce4ef..c164ed0 100644
--- a/learning/katas/python/IO/Built-in IOs/Built-in IOs/task-info.yaml
+++ b/learning/katas/python/IO/Built-in IOs/Built-in IOs/task-info.yaml
@@ -17,7 +17,7 @@
# under the License.
#
-type: edu
+type: theory
files:
- name: task.py
visible: true
diff --git a/learning/katas/python/IO/Built-in IOs/Built-in IOs/task.html b/learning/katas/python/IO/Built-in IOs/Built-in IOs/task.html
index 7d6cc8d..ef1b208 100644
--- a/learning/katas/python/IO/Built-in IOs/Built-in IOs/task.html
+++ b/learning/katas/python/IO/Built-in IOs/Built-in IOs/task.html
@@ -27,7 +27,6 @@
Transforms</a> page for a list of the currently available I/O transforms.
</p>
<p>
- <b>Note:</b> There is no kata for this task. Please click the "Check" button and
- proceed to the next task.
+ <b>Note:</b> There is no kata for this task. Please proceed to the next task.
</p>
</html>
diff --git a/learning/katas/python/IO/TextIO/ReadFromText/tests.py b/learning/katas/python/IO/TextIO/ReadFromText/tests.py
index 413bda7..273aada 100644
--- a/learning/katas/python/IO/TextIO/ReadFromText/tests.py
+++ b/learning/katas/python/IO/TextIO/ReadFromText/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_readfromtext_method():
@@ -50,6 +52,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_readfromtext_method()
test_output()
diff --git a/learning/katas/python/Introduction/Hello Beam/Hello Beam/tests.py b/learning/katas/python/Introduction/Hello Beam/Hello Beam/tests.py
index 3ee6516..33d45a6 100644
--- a/learning/katas/python/Introduction/Hello Beam/Hello Beam/tests.py
+++ b/learning/katas/python/Introduction/Hello Beam/Hello Beam/tests.py
@@ -14,7 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from test_helper import run_common_tests, failed, passed, get_answer_placeholders, get_file_output
+from test_helper import failed, passed, \
+ get_answer_placeholders, get_file_output, test_is_not_empty, \
+ test_answer_placeholders_text_deleted
def test_answer_placeholders():
@@ -36,6 +38,7 @@
if __name__ == '__main__':
- run_common_tests()
+ test_is_not_empty()
+ test_answer_placeholders_text_deleted()
test_answer_placeholders()
test_output()
diff --git a/learning/katas/python/course-info.yaml b/learning/katas/python/course-info.yaml
index b14f13a..ede92f2 100644
--- a/learning/katas/python/course-info.yaml
+++ b/learning/katas/python/course-info.yaml
@@ -22,7 +22,7 @@
summary: "This course provides a series of katas to get familiar with Apache Beam.\
\ \n\nApache Beam website – https://beam.apache.org/"
programming_language: Python
-programming_language_version: 2.7
+programming_language_version: 3.7
content:
- Introduction
- Core Transforms
diff --git a/learning/katas/python/course-remote-info.yaml b/learning/katas/python/course-remote-info.yaml
index ed9c8a7..c26dd2e 100644
--- a/learning/katas/python/course-remote-info.yaml
+++ b/learning/katas/python/course-remote-info.yaml
@@ -1,2 +1,2 @@
id: 54532
-update_date: Wed, 19 Jun 2019 10:36:17 UTC
+update_date: Mon, 20 Jan 2020 15:20:43 UTC
diff --git a/learning/katas/python/log_elements.py b/learning/katas/python/log_elements.py
index 471996ec30..ccff89b 100644
--- a/learning/katas/python/log_elements.py
+++ b/learning/katas/python/log_elements.py
@@ -26,7 +26,7 @@
self.prefix = prefix
def process(self, element, **kwargs):
- print self.prefix + str(element)
+ print(self.prefix + str(element))
yield element
def __init__(self, label=None, prefix=''):
diff --git a/learning/katas/python/requirements.txt b/learning/katas/python/requirements.txt
index c53a288..3a79b42 100644
--- a/learning/katas/python/requirements.txt
+++ b/learning/katas/python/requirements.txt
@@ -14,5 +14,5 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-apache-beam==2.13.0
-apache-beam[test]==2.13.0
+apache-beam==2.17.0
+apache-beam[test]==2.17.0
diff --git a/model/fn-execution/src/main/proto/beam_fn_api.proto b/model/fn-execution/src/main/proto/beam_fn_api.proto
index 3824e90..17440e6 100644
--- a/model/fn-execution/src/main/proto/beam_fn_api.proto
+++ b/model/fn-execution/src/main/proto/beam_fn_api.proto
@@ -889,3 +889,32 @@
// Stop the SDK worker.
rpc StopWorker (StopWorkerRequest) returns (StopWorkerResponse) {}
}
+
+// Request from runner to SDK Harness asking for its status. For more details see
+// https://s.apache.org/beam-fn-api-harness-status
+message WorkerStatusRequest {
+ // (Required) Unique ID identifying this request.
+ string id = 1;
+}
+
+// Response from SDK Harness to runner containing the debug related status info.
+message WorkerStatusResponse {
+ // (Required) Unique ID from the original request.
+ string id = 1;
+
+ // (Optional) Error message if exception encountered generating the status response.
+ string error = 2;
+
+ // (Optional) Status debugging info reported by SDK harness worker. Content and
+ // format is not strongly enforced but should be print-friendly and
+ // appropriate as an HTTP response body for end user. For details of the preferred
+ // info to include in the message see
+ // https://s.apache.org/beam-fn-api-harness-status
+ string status_info = 3;
+}
+
+// API for SDKs to report debug-related statuses to runner during pipeline execution.
+service BeamFnWorkerStatus {
+ rpc WorkerStatus (stream WorkerStatusResponse)
+ returns (stream WorkerStatusRequest) {}
+}
diff --git a/model/pipeline/src/main/proto/beam_runner_api.proto b/model/pipeline/src/main/proto/beam_runner_api.proto
index df5d59d..57c5295 100644
--- a/model/pipeline/src/main/proto/beam_runner_api.proto
+++ b/model/pipeline/src/main/proto/beam_runner_api.proto
@@ -395,6 +395,9 @@
// (Optional) Only set when this ParDo can request bundle finalization.
bool requests_finalization = 8;
+
+ // (Optional) A mapping of local timer family names to timer specifications.
+ map<string, TimerFamilySpec> timer_family_specs = 9;
}
// Parameters that a UDF might require.
@@ -461,6 +464,11 @@
string timer_coder_id = 2;
}
+message TimerFamilySpec {
+ TimeDomain.Enum time_domain = 1;
+ string timer_family_coder_id = 2;
+}
+
message IsBounded {
enum Enum {
UNSPECIFIED = 0;
@@ -1079,21 +1087,6 @@
FunctionSpec window_mapping_fn = 3;
}
-// Settings that decide the coder type of wire coder.
-message WireCoderSetting {
- // (Required) The URN of the wire coder.
- // Note that only windowed value coder or parameterized windowed value coder are supported.
- string urn = 1;
-
- // (Optional) The data specifying any parameters to the URN. If
- // the URN is beam:coder:windowed_value:v1, this may be omitted. If the URN is
- // beam:coder:param_windowed_value:v1, the payload is an encoded windowed
- // value using the beam:coder:windowed_value:v1 coder parameterized by
- // a beam:coder:bytes:v1 element coder and the window coder that this
- // param_windowed_value coder uses.
- bytes payload = 2;
-}
-
// An environment for executing UDFs. By default, an SDK container URL, but
// can also be a process forked by a command, or an externally managed process.
message Environment {
@@ -1302,8 +1295,8 @@
// because ExecutableStages use environments directly. This may change in the future.
Environment environment = 1;
- // set the wire coder of this executable stage
- WireCoderSetting wire_coder_setting = 9;
+ // The wire coder settings of this executable stage
+ repeated WireCoderSetting wire_coder_settings = 9;
// (Required) Input PCollection id. This must be present as a value in the inputs of any
// PTransform the ExecutableStagePayload is the payload of.
@@ -1333,6 +1326,10 @@
// this ExecutableStagePayload must be represented within this field.
repeated TimerId timers = 8;
+ // The timerfamilies required for this executable stage. Each timer familyof each PTransform within
+ // this ExecutableStagePayload must be represented within this field.
+ repeated TimerFamilyId timerFamilies = 10;
+
// A reference to a side input. Side inputs are uniquely identified by PTransform id and
// local name.
message SideInputId {
@@ -1362,4 +1359,36 @@
// (Required) The local name of this timer for the PTransform that references it.
string local_name = 2;
}
+
+ // A reference to a timer. Timers are uniquely identified by PTransform id and
+ // local name.
+ message TimerFamilyId {
+ // (Required) The id of the PTransform that references this timer family.
+ string transform_id = 1;
+
+ // (Required) The local name of this timer family for the PTransform that references it.
+ string local_name = 2;
+ }
+ // Settings that decide the coder type of wire coder.
+ message WireCoderSetting {
+ // (Required) The URN of the wire coder.
+ // Note that only windowed value coder or parameterized windowed value coder are supported.
+ string urn = 1;
+
+ // (Optional) The data specifying any parameters to the URN. If
+ // the URN is beam:coder:windowed_value:v1, this may be omitted. If the URN is
+ // beam:coder:param_windowed_value:v1, the payload is an encoded windowed
+ // value using the beam:coder:windowed_value:v1 coder parameterized by
+ // a beam:coder:bytes:v1 element coder and the window coder that this
+ // param_windowed_value coder uses.
+ bytes payload = 2;
+
+ // (Required) The target(PCollection or Timer) this setting applies to.
+ oneof target {
+ // The input or output PCollection id this setting applies to.
+ string input_or_output_id = 3;
+ // The timer id this setting applies to.
+ TimerId timer = 4;
+ }
+ }
}
diff --git a/release/src/main/scripts/build_release_candidate.sh b/release/src/main/scripts/build_release_candidate.sh
index 1f54d96..a367b60 100755
--- a/release/src/main/scripts/build_release_candidate.sh
+++ b/release/src/main/scripts/build_release_candidate.sh
@@ -46,10 +46,7 @@
WEBSITE_ROOT_DIR=beam-site
PYTHON_VER=("python2.7" "python3.5" "python3.6" "python3.7")
-FLINK_VER=("$(ls -1 runners/flink | awk '/^[0-9]+\.[0-9]+$/{print}')")
-if [[ "${#FLINK_VER[@]}" = 0 ]]; then
- echo "WARNING: Failed to list Flink versions. Are you in the root directory?"
-fi
+FLINK_VER=("1.7" "1.8" "1.9")
echo "================Setting Up Environment Variables==========="
echo "Which release version are you working on: "
@@ -224,7 +221,7 @@
echo '-------------------Generating and Pushing Python images-----------------'
./gradlew :sdks:python:container:buildAll -Pdocker-tag=${RELEASE}_rc${RC_NUM}
for ver in "${PYTHON_VER[@]}"; do
- docker push apachebeam/${ver}_sdk:${RELEASE}_rc${RC_NUM} &
+ docker push apachebeam/${ver}_sdk:${RELEASE}_rc${RC_NUM} &
done
echo '-------------------Generating and Pushing Java images-----------------'
@@ -236,7 +233,7 @@
echo '-------------Generating and Pushing Flink job server images-------------'
echo "Building containers for the following Flink versions:" "${FLINK_VER[@]}"
for ver in "${FLINK_VER[@]}"; do
- ./gradlew ":runners:flink:${ver}:job-server-container:dockerPush" -Pdocker-tag="${RELEASE}_rc${RC_NUM}"
+ ./gradlew ":runners:flink:${ver}:job-server-container:dockerPush" -Pdocker-tag="${RELEASE}_rc${RC_NUM}"
done
rm -rf ~/${PYTHON_ARTIFACTS_DIR}
diff --git a/release/src/main/scripts/publish_docker_images.sh b/release/src/main/scripts/publish_docker_images.sh
old mode 100644
new mode 100755
index 44a133e..e80c446
--- a/release/src/main/scripts/publish_docker_images.sh
+++ b/release/src/main/scripts/publish_docker_images.sh
@@ -24,69 +24,98 @@
set -e
-source release/src/main/scripts/build_release_candidate.sh
+PYTHON_VER=("python2.7" "python3.5" "python3.6" "python3.7")
+FLINK_VER=("1.7" "1.8" "1.9")
echo "Publish SDK docker images to Docker Hub."
+
+echo "================Setting Up Environment Variables==========="
+echo "Which release version are you working on: "
+read RELEASE
+
+echo "================Setting Up RC candidate Variables==========="
+echo "From which RC candidate do you create publish docker image? (ex: rc0, rc1) "
+read RC_VERSION
+
+echo "================Confirmimg Release and RC version==========="
+echo "We are using ${RC_VERSION} to create docker images for ${RELEASE}."
echo "Do you want to proceed? [y|N]"
read confirmation
if [[ $confirmation = "y" ]]; then
- echo "============Publishing SDK docker images on docker hub========="
- cd ~
- if [[ -d ${LOCAL_PYTHON_STAGING_DIR} ]]; then
- rm -rf ${LOCAL_PYTHON_STAGING_DIR}
- fi
- mkdir -p ${LOCAL_PYTHON_STAGING_DIR}
- cd ${LOCAL_PYTHON_STAGING_DIR}
-
- echo '-------------------Cloning Beam Release Branch-----------------'
- git clone ${GIT_REPO_URL}
- cd ${BEAM_ROOT_DIR}
- git checkout ${RELEASE_BRANCH}
echo '-------------------Generating and Pushing Python images-----------------'
- ./gradlew :sdks:python:container:buildAll -Pdocker-tag=${RELEASE}
for ver in "${PYTHON_VER[@]}"; do
- docker push apachebeam/${ver}_sdk:${RELEASE}
- docker tag apachebeam/${ver}_sdk:${RELEASE} apachebeam/${ver}_sdk:latest
- docker push apachebeam/${ver}_sdk:latest
- done
+ # Pull varified RC from dockerhub.
+ docker pull apachebeam/${ver}_sdk:${RELEASE}_${RC_VERSION}
- echo '-------------------Generating and Pushing Java images-----------------'
- ./gradlew :sdks:java:container:dockerPush -Pdocker-tag=${RELEASE}
- docker tag apachebeam/java_sdk:${RELEASE} apachebeam/java_sdk:latest
- docker push apachebeam/java_sdk:latest
+ # Tag with ${RELEASE} and push to dockerhub.
+ docker tag apachebeam/${ver}_sdk:${RELEASE}_${RC_VERSION} apachebeam/${ver}_sdk:${RELEASE}
+ docker push apachebeam/${ver}_sdk:${RELEASE}
- echo '-------------------Generating and Pushing Go images-----------------'
- ./gradlew :sdks:go:container:dockerPush -Pdocker-tag=${RELEASE}
- docker tag apachebeam/go_sdk:${RELEASE} apachebeam/go_sdk:latest
- docker push apachebeam/go_sdk:latest
+ # Tag with latest and push to dockerhub.
+ docker tag apachebeam/${ver}_sdk:${RELEASE}_${RC_VERSION} apachebeam/${ver}_sdk:latest
+ docker push apachebeam/${ver}_sdk:latest
- echo '-------------Generating and Pushing Flink job server images-------------'
- echo "Building containers for the following Flink versions:" "${FLINK_VER[@]}"
- for ver in "${FLINK_VER[@]}"; do
- ./gradlew ":runners:flink:${ver}:job-server-container:docker" -Pdocker-tag="${RELEASE}"
- FLINK_IMAGE_NAME=apachebeam/flink${ver}_job_server
- docker push "${FLINK_IMAGE_NAME}:${RELEASE}"
- docker tag "${FLINK_IMAGE_NAME}:${RELEASE}" "${FLINK_IMAGE_NAME}:latest"
- docker push "${FLINK_IMAGE_NAME}:latest"
- done
-
- rm -rf ~/${PYTHON_ARTIFACTS_DIR}
-
- echo "-------------------Clean up SDK docker images at local-------------------"
- for ver in "${PYTHON_VER[@]}"; do
+ # Cleanup images from local
+ docker rmi -f apachebeam/${ver}_sdk:${RELEASE}_${RC_VERSION}
docker rmi -f apachebeam/${ver}_sdk:${RELEASE}
docker rmi -f apachebeam/${ver}_sdk:latest
done
+ echo '-------------------Generating and Pushing Java images-----------------'
+ # Pull varified RC from dockerhub.
+ docker pull apachebeam/java_sdk:${RELEASE}_${RC_VERSION}
+
+ # Tag with ${RELEASE} and push to dockerhub.
+ docker tag apachebeam/java_sdk:${RELEASE}_${RC_VERSION} apachebeam/java_sdk:${RELEASE}
+ docker push apachebeam/java_sdk:${RELEASE}
+
+ # Tag with latest and push to dockerhub.
+ docker tag apachebeam/java_sdk:${RELEASE}_${RC_VERSION} apachebeam/java_sdk:latest
+ docker push apachebeam/java_sdk:latest
+
+ # Cleanup images from local
+ docker rmi -f apachebeam/java_sdk:${RELEASE}_${RC_VERSION}
docker rmi -f apachebeam/java_sdk:${RELEASE}
docker rmi -f apachebeam/java_sdk:latest
+ echo '-------------------Generating and Pushing Go images-----------------'
+ # Pull varified RC from dockerhub.
+ docker pull apachebeam/go_sdk:${RELEASE}_${RC_VERSION}
+
+ # Tag with ${RELEASE} and push to dockerhub.
+ docker tag apachebeam/go_sdk:${RELEASE}_${RC_VERSION} apachebeam/go_sdk:${RELEASE}
+ docker push apachebeam/go_sdk:${RELEASE}
+
+ # Tag with latest and push to dockerhub.
+ docker tag apachebeam/go_sdk:${RELEASE}_${RC_VERSION} apachebeam/go_sdk:latest
+ docker push apachebeam/go_sdk:latest
+
+ # Cleanup images from local
+ docker rmi -f apachebeam/go_sdk:${RELEASE}_${RC_VERSION}
docker rmi -f apachebeam/go_sdk:${RELEASE}
docker rmi -f apachebeam/go_sdk:latest
+ echo '-------------Generating and Pushing Flink job server images-------------'
+ echo "Publishing images for the following Flink versions:" "${FLINK_VER[@]}"
+ echo "Make sure the versions are correct, then press any key to proceed."
+ read
for ver in "${FLINK_VER[@]}"; do
FLINK_IMAGE_NAME=apachebeam/flink${ver}_job_server
+
+ # Pull verified RC from dockerhub.
+ docker pull "${FLINK_IMAGE_NAME}:${RELEASE}_${RC_VERSION}"
+
+ # Tag with ${RELEASE} and push to dockerhub.
+ docker tag "${FLINK_IMAGE_NAME}:${RELEASE}_${RC_VERSION}" "${FLINK_IMAGE_NAME}:${RELEASE}"
+ docker push "${FLINK_IMAGE_NAME}:${RELEASE}"
+
+ # Tag with latest and push to dockerhub.
+ docker tag "${FLINK_IMAGE_NAME}:${RELEASE}_${RC_VERSION}" "${FLINK_IMAGE_NAME}:latest"
+ docker push "${FLINK_IMAGE_NAME}:latest"
+
+ # Cleanup images from local
+ docker rmi -f "${FLINK_IMAGE_NAME}:${RELEASE}_${RC_VERSION}"
docker rmi -f "${FLINK_IMAGE_NAME}:${RELEASE}"
docker rmi -f "${FLINK_IMAGE_NAME}:latest"
done
diff --git a/release/src/main/scripts/set_version.sh b/release/src/main/scripts/set_version.sh
index 5844b73..b52dfc9 100755
--- a/release/src/main/scripts/set_version.sh
+++ b/release/src/main/scripts/set_version.sh
@@ -67,7 +67,7 @@
sed -i -e "s/version=.*/version=$TARGET_VERSION/" gradle.properties
sed -i -e "s/project.version = .*/project.version = '$TARGET_VERSION'/" buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
sed -i -e "s/^__version__ = .*/__version__ = '${TARGET_VERSION}'/" sdks/python/apache_beam/version.py
- sed -i -e "s/python_sdk_version=.*/python_sdk_version=$TARGET_VERSION/" gradle.properties
+ sed -i -e "s/sdk_version=.*/sdk_version=$TARGET_VERSION/" gradle.properties
# TODO: [BEAM-4767]
sed -i -e "s/'dataflow.container_version' : .*/'dataflow.container_version' : 'beam-${RELEASE}'/" runners/google-cloud-dataflow-java/build.gradle
else
@@ -79,7 +79,7 @@
sed -i -e "s/version=.*/version=$TARGET_VERSION-SNAPSHOT/" gradle.properties
sed -i -e "s/project.version = .*/project.version = '$TARGET_VERSION'/" buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
sed -i -e "s/^__version__ = .*/__version__ = '${TARGET_VERSION}.dev'/" sdks/python/apache_beam/version.py
- sed -i -e "s/python_sdk_version=.*/python_sdk_version=$TARGET_VERSION.dev/" gradle.properties
+ sed -i -e "s/sdk_version=.*/sdk_version=$TARGET_VERSION.dev/" gradle.properties
sed -i -e "s/'dataflow.container_version' : .*/'dataflow.container_version' : 'beam-master-.*'/" runners/google-cloud-dataflow-java/build.gradle
fi
diff --git a/release/src/main/scripts/verify_release_build.sh b/release/src/main/scripts/verify_release_build.sh
index 8442e9f..f9b0480 100755
--- a/release/src/main/scripts/verify_release_build.sh
+++ b/release/src/main/scripts/verify_release_build.sh
@@ -46,7 +46,6 @@
# To run all PostCommit jobs
"Run Go PostCommit"
"Run Java PostCommit"
- "Run Java PostCommit"
"Run Java PortabilityApi PostCommit"
"Run Java Flink PortableValidatesRunner Batch"
"Run Java Flink PortableValidatesRunner Streaming"
@@ -152,7 +151,7 @@
if [[ ! -z `which hub` ]]; then
git checkout -b ${WORKING_BRANCH} origin/${RELEASE_BRANCH} --quiet
- touch empty_file.txt
+ touch empty_file.json
git add .
git commit -m "Add empty file in order to create a test PR" --quiet
git push -f ${GITHUB_USERNAME} --quiet
diff --git a/runners/apex/build.gradle b/runners/apex/build.gradle
index 739fd9d..d204a0b 100644
--- a/runners/apex/build.gradle
+++ b/runners/apex/build.gradle
@@ -42,7 +42,6 @@
compile library.java.apex_common
compile library.java.malhar_library
compile library.java.apex_engine
- compile library.java.commons_lang3
compile library.java.apex_engine
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
// ApexStateInternalsTest extends abstract StateInternalsTest
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexYarnLauncher.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexYarnLauncher.java
index d1ae4ec..15e3968 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexYarnLauncher.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexYarnLauncher.java
@@ -62,11 +62,11 @@
import org.apache.apex.api.Launcher.LauncherException;
import org.apache.apex.api.Launcher.ShutdownMode;
import org.apache.apex.api.YarnAppLauncher;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.SerializationUtils;
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.Splitter;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
import org.apache.commons.io.IOUtils;
-import org.apache.commons.lang3.SerializationUtils;
import org.apache.hadoop.conf.Configuration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java
index 6d1f4b0..4a19f14 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java
@@ -28,6 +28,8 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.ImmutablePair;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair;
import org.apache.beam.runners.apex.ApexPipelineOptions;
import org.apache.beam.runners.apex.translation.utils.ApexStateInternals;
import org.apache.beam.runners.apex.translation.utils.ApexStateInternals.ApexStateBackend;
@@ -46,8 +48,6 @@
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-import org.apache.commons.lang3.tuple.ImmutablePair;
-import org.apache.commons.lang3.tuple.Pair;
/** Maintains context data for {@link TransformTranslator}s. */
@SuppressWarnings({"rawtypes", "unchecked", "TypeParameterUnusedInFormals"})
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java
index 79bb6ef..4841c6a 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java
@@ -384,7 +384,12 @@
checkArgument(namespace instanceof WindowNamespace);
BoundedWindow window = ((WindowNamespace<?>) namespace).getWindow();
pushbackDoFnRunner.onTimer(
- timerData.getTimerId(), window, timerData.getTimestamp(), timerData.getDomain());
+ timerData.getTimerId(),
+ timerData.getTimerFamilyId(),
+ window,
+ timerData.getTimestamp(),
+ timerData.getOutputTimestamp(),
+ timerData.getDomain());
}
pushbackDoFnRunner.finishBundle();
}
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternals.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternals.java
index 682cbed..886f153 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternals.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternals.java
@@ -99,7 +99,7 @@
}
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
this.eventTimeTimeTimers.deleteTimer(getKeyBytes(this.currentKey), namespace, timerId);
this.processingTimeTimers.deleteTimer(getKeyBytes(this.currentKey), namespace, timerId);
}
diff --git a/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternalsTest.java b/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternalsTest.java
index 1d7e3f8..180d2fe 100644
--- a/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternalsTest.java
+++ b/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/operators/ApexTimerInternalsTest.java
@@ -56,11 +56,13 @@
timerInternals.setContext(key1, StringUtf8Coder.of(), Instant.now(), null);
TimerData timerData0 =
- TimerData.of("timerData0", StateNamespaces.global(), instant0, TimeDomain.EVENT_TIME);
+ TimerData.of(
+ "timerData0", StateNamespaces.global(), instant0, instant0, TimeDomain.EVENT_TIME);
timerInternals.setTimer(timerData0);
TimerData timerData1 =
- TimerData.of("timerData1", StateNamespaces.global(), instant1, TimeDomain.EVENT_TIME);
+ TimerData.of(
+ "timerData1", StateNamespaces.global(), instant1, instant1, TimeDomain.EVENT_TIME);
timerInternals.setTimer(timerData1);
timerInternals.fireReadyTimers(instant0.getMillis(), timerProcessor, TimeDomain.EVENT_TIME);
@@ -94,18 +96,20 @@
timerInternals.setContext(key1, StringUtf8Coder.of(), Instant.now(), null);
TimerData timerData0 =
- TimerData.of("timerData0", StateNamespaces.global(), instant0, TimeDomain.EVENT_TIME);
+ TimerData.of(
+ "timerData0", StateNamespaces.global(), instant0, instant0, TimeDomain.EVENT_TIME);
timerInternals.setTimer(timerData0);
TimerData timerData1 =
- TimerData.of("timerData1", StateNamespaces.global(), instant1, TimeDomain.EVENT_TIME);
+ TimerData.of(
+ "timerData1", StateNamespaces.global(), instant1, instant1, TimeDomain.EVENT_TIME);
timerInternals.setTimer(timerData1);
Map<?, Set<Slice>> timerMap = timerInternals.getTimerSet(TimeDomain.EVENT_TIME).getMap();
assertEquals(1, timerMap.size());
assertEquals(2, timerMap.values().iterator().next().size());
- timerInternals.deleteTimer(timerData0.getNamespace(), timerData0.getTimerId());
+ timerInternals.deleteTimer(timerData0.getNamespace(), timerData0.getTimerId(), "");
assertEquals(1, timerMap.size());
assertEquals(1, timerMap.values().iterator().next().size());
diff --git a/runners/core-construction-java/build.gradle b/runners/core-construction-java/build.gradle
index b9c842f..3de798a 100644
--- a/runners/core-construction-java/build.gradle
+++ b/runners/core-construction-java/build.gradle
@@ -36,7 +36,7 @@
compile project(path: ":model:pipeline", configuration: "shadow")
compile project(path: ":model:job-management", configuration: "shadow")
compile project(path: ":sdks:java:core", configuration: "shadow")
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
compile library.java.vendored_guava_26_0_jre
compile library.java.classgraph
compile library.java.jackson_core
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ArtifactServiceStager.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ArtifactServiceStager.java
index 4212916..29e47b8 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ArtifactServiceStager.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ArtifactServiceStager.java
@@ -50,9 +50,9 @@
import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc.ArtifactStagingServiceStub;
import org.apache.beam.sdk.util.MoreFutures;
import org.apache.beam.sdk.util.ThrowingSupplier;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Channel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Channel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hasher;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ListeningExecutorService;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/BeamUrns.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/BeamUrns.java
index f1f30dc..e4fc6d7 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/BeamUrns.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/BeamUrns.java
@@ -18,7 +18,7 @@
package org.apache.beam.runners.core.construction;
import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ProtocolMessageEnum;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ProtocolMessageEnum;
/** Returns the standard URN of a given enum annotated with [(standard_urn)]. */
public class BeamUrns {
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslation.java
index 8e1021d..86f0178 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslation.java
@@ -28,7 +28,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.util.SerializableUtils;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.collect.BiMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableBiMap;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslators.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslators.java
index 6f156dd..81b7922 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslators.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CoderTranslators.java
@@ -33,7 +33,7 @@
import org.apache.beam.sdk.util.InstanceBuilder;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
/** {@link CoderTranslator} implementations for known coder types. */
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CombineTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CombineTranslation.java
index e5edc35..5ea9ef7 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CombineTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CombineTranslation.java
@@ -39,7 +39,7 @@
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java
index b89c5b6..5027bb4 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java
@@ -34,7 +34,7 @@
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
/**
* Utility methods for translating a {@link View} transforms to and from {@link RunnerApi}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DefaultExpansionServiceClientFactory.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DefaultExpansionServiceClientFactory.java
index 1586be8..a25007e 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DefaultExpansionServiceClientFactory.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DefaultExpansionServiceClientFactory.java
@@ -23,7 +23,7 @@
import org.apache.beam.model.expansion.v1.ExpansionApi;
import org.apache.beam.model.expansion.v1.ExpansionServiceGrpc;
import org.apache.beam.model.pipeline.v1.Endpoints;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
/** Default factory for ExpansionServiceClient used by External transform. */
public class DefaultExpansionServiceClientFactory implements ExpansionServiceClientFactory {
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 c7cd235..57836b9 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
@@ -19,8 +19,8 @@
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Any;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.BoolValue;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Any;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.BoolValue;
/** Utilities for going to/from DisplayData protos. */
public class DisplayDataTranslation {
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Environments.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Environments.java
index ed94642..73bf534 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Environments.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Environments.java
@@ -32,7 +32,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.StandardEnvironments;
import org.apache.beam.sdk.util.ReleaseInfo;
import org.apache.beam.sdk.util.common.ReflectHelpers;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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;
@@ -55,7 +55,7 @@
* container.
*/
private static final String JAVA_SDK_HARNESS_CONTAINER_URL =
- "apachebeam/java_sdk:" + ReleaseInfo.getReleaseInfo().getVersion();
+ "apachebeam/java_sdk:" + ReleaseInfo.getReleaseInfo().getSdkVersion();
public static final Environment JAVA_SDK_HARNESS_ENVIRONMENT =
createDockerEnvironment(JAVA_SDK_HARNESS_CONTAINER_URL);
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/External.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/External.java
index d58346b..45665fd 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/External.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/External.java
@@ -38,8 +38,8 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannelBuilder;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ModelCoders.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ModelCoders.java
index 3d6d4dd..929d7a8 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ModelCoders.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ModelCoders.java
@@ -25,7 +25,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.Coder;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.StandardCoders;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
/** Utilities and constants ot interact with coders that are part of the Beam Model. */
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionViewTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionViewTranslation.java
index d9ddb93..402e7d6 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionViewTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionViewTranslation.java
@@ -30,7 +30,7 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
/** Utilities for interacting with PCollection view protos. */
public class PCollectionViewTranslation {
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
index c21cd38..fbdc259 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
@@ -19,6 +19,7 @@
import static org.apache.beam.runners.core.construction.BeamUrns.getUrn;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
import java.io.IOException;
import java.util.Collection;
@@ -60,68 +61,135 @@
* buffers}.
*/
public class PTransformTranslation {
+ // We specifically copy the values here so that they can be used in switch case statements
+ // and we validate that the value matches the actual URN in the static block below.
- public static final String PAR_DO_TRANSFORM_URN = getUrn(StandardPTransforms.Primitives.PAR_DO);
- public static final String FLATTEN_TRANSFORM_URN = getUrn(StandardPTransforms.Primitives.FLATTEN);
- public static final String GROUP_BY_KEY_TRANSFORM_URN =
- getUrn(StandardPTransforms.Primitives.GROUP_BY_KEY);
- public static final String IMPULSE_TRANSFORM_URN = getUrn(StandardPTransforms.Primitives.IMPULSE);
- public static final String ASSIGN_WINDOWS_TRANSFORM_URN =
- getUrn(StandardPTransforms.Primitives.ASSIGN_WINDOWS);
- public static final String TEST_STREAM_TRANSFORM_URN =
- getUrn(StandardPTransforms.Primitives.TEST_STREAM);
- public static final String MAP_WINDOWS_TRANSFORM_URN =
- getUrn(StandardPTransforms.Primitives.MAP_WINDOWS);
+ // Primitives
+ public static final String PAR_DO_TRANSFORM_URN = "beam:transform:pardo:v1";
+ public static final String FLATTEN_TRANSFORM_URN = "beam:transform:flatten:v1";
+ public static final String GROUP_BY_KEY_TRANSFORM_URN = "beam:transform:group_by_key:v1";
+ public static final String IMPULSE_TRANSFORM_URN = "beam:transform:impulse:v1";
+ public static final String ASSIGN_WINDOWS_TRANSFORM_URN = "beam:transform:window_into:v1";
+ public static final String TEST_STREAM_TRANSFORM_URN = "beam:transform:teststream:v1";
+ public static final String MAP_WINDOWS_TRANSFORM_URN = "beam:transform:map_windows:v1";
+ public static final String MERGE_WINDOWS_TRANSFORM_URN = "beam:transform:merge_windows:v1";
+ // DeprecatedPrimitives
/**
* @deprecated SDKs should move away from creating `Read` transforms and migrate to using Impulse
* + SplittableDoFns.
*/
- @Deprecated
- public static final String READ_TRANSFORM_URN =
- getUrn(StandardPTransforms.DeprecatedPrimitives.READ);
+ @Deprecated public static final String READ_TRANSFORM_URN = "beam:transform:read:v1";
+
/**
* @deprecated runners should move away from translating `CreatePCollectionView` and treat this as
* part of the translation for a `ParDo` side input.
*/
@Deprecated
- public static final String CREATE_VIEW_TRANSFORM_URN =
- getUrn(StandardPTransforms.DeprecatedPrimitives.CREATE_VIEW);
+ public static final String CREATE_VIEW_TRANSFORM_URN = "beam:transform:create_view:v1";
- public static final String COMBINE_PER_KEY_TRANSFORM_URN =
- getUrn(StandardPTransforms.Composites.COMBINE_PER_KEY);
- public static final String COMBINE_GLOBALLY_TRANSFORM_URN =
- getUrn(StandardPTransforms.Composites.COMBINE_GLOBALLY);
- public static final String COMBINE_GROUPED_VALUES_TRANSFORM_URN =
- getUrn(CombineComponents.COMBINE_GROUPED_VALUES);
+ // Composites
+ public static final String COMBINE_PER_KEY_TRANSFORM_URN = "beam:transform:combine_per_key:v1";
+ public static final String COMBINE_GLOBALLY_TRANSFORM_URN = "beam:transform:combine_globally:v1";
+ public static final String RESHUFFLE_URN = "beam:transform:reshuffle:v1";
+ public static final String WRITE_FILES_TRANSFORM_URN = "beam:transform:write_files:v1";
+
+ // CombineComponents
public static final String COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN =
- getUrn(CombineComponents.COMBINE_PER_KEY_PRECOMBINE);
+ "beam:transform:combine_per_key_precombine:v1";
public static final String COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN =
- getUrn(CombineComponents.COMBINE_PER_KEY_MERGE_ACCUMULATORS);
+ "beam:transform:combine_per_key_merge_accumulators:v1";
public static final String COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN =
- getUrn(CombineComponents.COMBINE_PER_KEY_EXTRACT_OUTPUTS);
- public static final String RESHUFFLE_URN = getUrn(StandardPTransforms.Composites.RESHUFFLE);
- public static final String WRITE_FILES_TRANSFORM_URN =
- getUrn(StandardPTransforms.Composites.WRITE_FILES);
+ "beam:transform:combine_per_key_extract_outputs:v1";
+ public static final String COMBINE_GROUPED_VALUES_TRANSFORM_URN =
+ "beam:transform:combine_grouped_values:v1";
// SplittableParDoComponents
public static final String SPLITTABLE_PAIR_WITH_RESTRICTION_URN =
- getUrn(SplittableParDoComponents.PAIR_WITH_RESTRICTION);
+ "beam:transform:sdf_pair_with_restriction:v1";
public static final String SPLITTABLE_SPLIT_RESTRICTION_URN =
- getUrn(SplittableParDoComponents.SPLIT_RESTRICTION);
+ "beam:transform:sdf_split_restriction:v1";
+ /**
+ * @deprecated runners should move away from using `SplittableProcessKeyedElements` and prefer to
+ * internalize any necessary SplittableDoFn expansion.
+ */
+ @Deprecated
public static final String SPLITTABLE_PROCESS_KEYED_URN =
- getUrn(SplittableParDoComponents.PROCESS_KEYED_ELEMENTS);
- public static final String SPLITTABLE_PROCESS_ELEMENTS_URN =
- getUrn(SplittableParDoComponents.PROCESS_ELEMENTS);
- public static final String SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN =
- getUrn(SplittableParDoComponents.SPLIT_AND_SIZE_RESTRICTIONS);
- public static final String SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN =
- getUrn(SplittableParDoComponents.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS);
+ "beam:transform:sdf_process_keyed_elements:v1";
- public static final String ITERABLE_SIDE_INPUT =
- getUrn(RunnerApi.StandardSideInputTypes.Enum.ITERABLE);
- public static final String MULTIMAP_SIDE_INPUT =
- getUrn(RunnerApi.StandardSideInputTypes.Enum.MULTIMAP);
+ public static final String SPLITTABLE_PROCESS_ELEMENTS_URN =
+ "beam:transform:sdf_process_elements:v1";
+ public static final String SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN =
+ "beam:transform:sdf_split_and_size_restrictions:v1";
+ public static final String SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN =
+ "beam:transform:sdf_process_sized_element_and_restrictions:v1";
+
+ static {
+ // Primitives
+ checkState(PAR_DO_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.PAR_DO)));
+ checkState(FLATTEN_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.FLATTEN)));
+ checkState(
+ GROUP_BY_KEY_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.GROUP_BY_KEY)));
+ checkState(IMPULSE_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.IMPULSE)));
+ checkState(
+ ASSIGN_WINDOWS_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.ASSIGN_WINDOWS)));
+ checkState(
+ TEST_STREAM_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.TEST_STREAM)));
+ checkState(
+ MAP_WINDOWS_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.MAP_WINDOWS)));
+ checkState(
+ MERGE_WINDOWS_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Primitives.MERGE_WINDOWS)));
+
+ // DeprecatedPrimitives
+ checkState(READ_TRANSFORM_URN.equals(getUrn(StandardPTransforms.DeprecatedPrimitives.READ)));
+ checkState(
+ CREATE_VIEW_TRANSFORM_URN.equals(
+ getUrn(StandardPTransforms.DeprecatedPrimitives.CREATE_VIEW)));
+
+ // Composites
+ checkState(
+ COMBINE_PER_KEY_TRANSFORM_URN.equals(
+ getUrn(StandardPTransforms.Composites.COMBINE_PER_KEY)));
+ checkState(
+ COMBINE_GLOBALLY_TRANSFORM_URN.equals(
+ getUrn(StandardPTransforms.Composites.COMBINE_GLOBALLY)));
+ checkState(RESHUFFLE_URN.equals(getUrn(StandardPTransforms.Composites.RESHUFFLE)));
+ checkState(
+ WRITE_FILES_TRANSFORM_URN.equals(getUrn(StandardPTransforms.Composites.WRITE_FILES)));
+
+ // CombineComponents
+ checkState(
+ COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN.equals(
+ getUrn(CombineComponents.COMBINE_PER_KEY_PRECOMBINE)));
+ checkState(
+ COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN.equals(
+ getUrn(CombineComponents.COMBINE_PER_KEY_MERGE_ACCUMULATORS)));
+ checkState(
+ COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN.equals(
+ getUrn(CombineComponents.COMBINE_PER_KEY_EXTRACT_OUTPUTS)));
+ checkState(
+ COMBINE_GROUPED_VALUES_TRANSFORM_URN.equals(
+ getUrn(CombineComponents.COMBINE_GROUPED_VALUES)));
+
+ // SplittableParDoComponents
+ checkState(
+ SPLITTABLE_PAIR_WITH_RESTRICTION_URN.equals(
+ getUrn(SplittableParDoComponents.PAIR_WITH_RESTRICTION)));
+ checkState(
+ SPLITTABLE_SPLIT_RESTRICTION_URN.equals(
+ getUrn(SplittableParDoComponents.SPLIT_RESTRICTION)));
+ checkState(
+ SPLITTABLE_PROCESS_KEYED_URN.equals(
+ getUrn(SplittableParDoComponents.PROCESS_KEYED_ELEMENTS)));
+ checkState(
+ SPLITTABLE_PROCESS_ELEMENTS_URN.equals(getUrn(SplittableParDoComponents.PROCESS_ELEMENTS)));
+ checkState(
+ SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN.equals(
+ getUrn(SplittableParDoComponents.SPLIT_AND_SIZE_RESTRICTIONS)));
+ checkState(
+ SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN.equals(
+ getUrn(SplittableParDoComponents.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS)));
+ }
private static final Collection<TransformTranslator<?>> KNOWN_TRANSLATORS =
loadKnownTranslators();
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
index f26f49ba..7e6ba7b 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
@@ -18,6 +18,11 @@
package org.apache.beam.runners.core.construction;
import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
+import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN;
+import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN;
+import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN;
+import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN;
+import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN;
import static org.apache.beam.sdk.transforms.reflect.DoFnSignatures.getStateSpecOrThrow;
import static org.apache.beam.sdk.transforms.reflect.DoFnSignatures.getTimerSpecOrThrow;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
@@ -74,8 +79,8 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
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.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
@@ -254,10 +259,26 @@
translateTimerSpec(getTimerSpecOrThrow(timer.getValue(), doFn), newComponents);
timerSpecs.put(timer.getKey(), spec);
}
+
return timerSpecs;
}
@Override
+ public Map<String, RunnerApi.TimerFamilySpec> translateTimerFamilySpecs(
+ SdkComponents newComponents) {
+ Map<String, RunnerApi.TimerFamilySpec> timerFamilySpecs = new HashMap<>();
+ for (Map.Entry<String, DoFnSignature.TimerFamilyDeclaration> timerFamily :
+ signature.timerFamilyDeclarations().entrySet()) {
+ RunnerApi.TimerFamilySpec spec =
+ translateTimerFamilySpec(
+ DoFnSignatures.getTimerFamilySpecOrThrow(timerFamily.getValue(), doFn),
+ newComponents);
+ timerFamilySpecs.put(timerFamily.getKey(), spec);
+ }
+ return timerFamilySpecs;
+ }
+
+ @Override
public boolean isSplittable() {
return signature.processElement().isSplittable();
}
@@ -414,7 +435,13 @@
public static RunnerApi.PCollection getMainInput(
RunnerApi.PTransform ptransform, Components components) throws IOException {
checkArgument(
- ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN),
+ PAR_DO_TRANSFORM_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_PAIR_WITH_RESTRICTION_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_SPLIT_RESTRICTION_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_PROCESS_ELEMENTS_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN.equals(
+ ptransform.getSpec().getUrn()),
"Unexpected payload type %s",
ptransform.getSpec().getUrn());
return components.getPcollectionsOrThrow(
@@ -425,7 +452,13 @@
public static String getMainInputName(RunnerApi.PTransformOrBuilder ptransform)
throws IOException {
checkArgument(
- ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN),
+ PAR_DO_TRANSFORM_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_PAIR_WITH_RESTRICTION_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_SPLIT_RESTRICTION_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_PROCESS_ELEMENTS_URN.equals(ptransform.getSpec().getUrn())
+ || SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN.equals(
+ ptransform.getSpec().getUrn()),
"Unexpected payload type %s",
ptransform.getSpec().getUrn());
ParDoPayload payload = ParDoPayload.parseFrom(ptransform.getSpec().getPayload());
@@ -563,6 +596,14 @@
.build();
}
+ public static RunnerApi.TimerFamilySpec translateTimerFamilySpec(
+ TimerSpec timer, SdkComponents components) {
+ return RunnerApi.TimerFamilySpec.newBuilder()
+ .setTimeDomain(translateTimeDomain(timer.getTimeDomain()))
+ .setTimerFamilyCoderId(registerCoderOrThrow(components, Timer.Coder.of(VoidCoder.of())))
+ .build();
+ }
+
private static RunnerApi.TimeDomain.Enum translateTimeDomain(TimeDomain timeDomain) {
switch (timeDomain) {
case EVENT_TIME:
@@ -680,7 +721,9 @@
public static boolean usesStateOrTimers(AppliedPTransform<?, ?, ?> transform) throws IOException {
ParDoPayload payload = getParDoPayload(transform);
- return payload.getStateSpecsCount() > 0 || payload.getTimerSpecsCount() > 0;
+ return payload.getStateSpecsCount() > 0
+ || payload.getTimerSpecsCount() > 0
+ || payload.getTimerFamilySpecsCount() > 0;
}
public static boolean isSplittable(AppliedPTransform<?, ?, ?> transform) throws IOException {
@@ -709,6 +752,8 @@
Map<String, RunnerApi.TimerSpec> translateTimerSpecs(SdkComponents newComponents);
+ Map<String, RunnerApi.TimerFamilySpec> translateTimerFamilySpecs(SdkComponents newComponents);
+
boolean isSplittable();
String translateRestrictionCoderId(SdkComponents newComponents);
@@ -722,6 +767,7 @@
.addAllParameters(parDo.translateParameters())
.putAllStateSpecs(parDo.translateStateSpecs(components))
.putAllTimerSpecs(parDo.translateTimerSpecs(components))
+ .putAllTimerFamilySpecs(parDo.translateTimerFamilySpecs(components))
.putAllSideInputs(parDo.translateSideInputs(components))
.setSplittable(parDo.isSplittable())
.setRestrictionCoderId(parDo.translateRestrictionCoderId(components))
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslation.java
index baf7c36..56e5d06 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslation.java
@@ -27,9 +27,9 @@
import java.util.Map;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.common.ReflectHelpers;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.JsonFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.JsonFormat;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.CaseFormat;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReadTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReadTranslation.java
index 81bafab..94288e5 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReadTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReadTranslation.java
@@ -38,8 +38,8 @@
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
/**
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
index f89874e..2b700d5 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
@@ -32,6 +32,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.Parameter;
import org.apache.beam.model.pipeline.v1.RunnerApi.SideInput;
import org.apache.beam.model.pipeline.v1.RunnerApi.StateSpec;
+import org.apache.beam.model.pipeline.v1.RunnerApi.TimerFamilySpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.TimerSpec;
import org.apache.beam.runners.core.construction.PTransformTranslation.TransformPayloadTranslator;
import org.apache.beam.runners.core.construction.ParDoTranslation.ParDoLike;
@@ -399,6 +400,13 @@
}
@Override
+ public Map<String, TimerFamilySpec> translateTimerFamilySpecs(
+ SdkComponents newComponents) {
+ // SDFs don't have timers.
+ return ImmutableMap.of();
+ }
+
+ @Override
public boolean isSplittable() {
return true;
}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDoNaiveBounded.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDoNaiveBounded.java
index 8a41d7e..ce888c9 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDoNaiveBounded.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDoNaiveBounded.java
@@ -29,6 +29,7 @@
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -243,6 +244,11 @@
}
@Override
+ public TimerMap timerFamily(String tagId) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
public Object schemaElement(int index) {
throw new UnsupportedOperationException();
}
@@ -253,6 +259,11 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
return new OutputReceiver<OutputT>() {
@Override
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java
index 1b747c1..abbc328 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java
@@ -36,7 +36,7 @@
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TimestampedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.joda.time.Duration;
import org.joda.time.Instant;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Timer.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Timer.java
index d994762..072dbc7 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Timer.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Timer.java
@@ -47,13 +47,9 @@
/** Returns a timer for the given timestamp with a user specified payload. */
public static <T> Timer<T> of(Instant timestamp, @Nullable T payload) {
- return new AutoValue_Timer(timestamp, timestamp, payload);
+ return new AutoValue_Timer(timestamp, payload);
}
- /** Returns a timer for the given timestamp with a user specified payload and outputTimestamp. */
- public static <T> Timer<T> of(Instant timestamp, Instant outputTimestamp, @Nullable T payload) {
- return new AutoValue_Timer(timestamp, outputTimestamp, payload);
- }
/**
* Returns the timestamp of when the timer is scheduled to fire.
*
@@ -62,9 +58,6 @@
*/
public abstract Instant getTimestamp();
- /* Returns the outputTimestamps */
- public abstract Instant getOutputTimestamp();
-
/** A user supplied payload. */
@Nullable
public abstract T getPayload();
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowIntoTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowIntoTranslation.java
index 0d72861..d7fafcd 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowIntoTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowIntoTranslation.java
@@ -33,7 +33,7 @@
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.Window.Assign;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
/**
* Utility methods for translating a {@link Window.Assign} to and from {@link RunnerApi}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java
index 63f662f..bbb31a7 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java
@@ -41,10 +41,10 @@
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.Durations;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.Timestamps;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.Durations;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.Timestamps;
import org.joda.time.Duration;
/** Utilities for working with {@link WindowingStrategy WindowingStrategies}. */
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java
index 0df16a2..986a585 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java
@@ -45,7 +45,7 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.collect.Lists;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionServer.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionServer.java
index 12b52f4..5859d69 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionServer.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionServer.java
@@ -20,8 +20,8 @@
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.concurrent.TimeUnit;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.netty.NettyServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.netty.NettyServerBuilder;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
/** A {@link Server gRPC Server} for an ExpansionService. */
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionService.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionService.java
index aab393b..45559bb 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionService.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/expansion/ExpansionService.java
@@ -54,9 +54,9 @@
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
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.CaseFormat;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Converter;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java
index dd2d374..21d252d 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java
@@ -17,9 +17,8 @@
*/
package org.apache.beam.runners.core.construction.graph;
-import static org.apache.beam.runners.core.construction.BeamUrns.getUrn;
-
import java.util.Collection;
+import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -29,11 +28,11 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.SideInputId;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.TimerId;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.UserStateId;
+import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
-import org.apache.beam.model.pipeline.v1.RunnerApi.WireCoderSetting;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
@@ -73,12 +72,12 @@
Environment getEnvironment();
/**
- * Returns the {@link WireCoderSetting} this stage executes in.
+ * Returns a set of {@link WireCoderSetting}s this stage executes in.
*
* <p>A {@link WireCoderSetting} consists of settings which is used to configure the type of the
- * wire coder.
+ * wire coder for a dedicated PCollection.
*/
- WireCoderSetting getWireCoderSetting();
+ Collection<WireCoderSetting> getWireCoderSettings();
/**
* Returns the root {@link PCollectionNode} of this {@link ExecutableStage}. This {@link
@@ -145,7 +144,7 @@
ExecutableStagePayload.Builder payload = ExecutableStagePayload.newBuilder();
payload.setEnvironment(getEnvironment());
- payload.setWireCoderSetting(getWireCoderSetting());
+ payload.addAllWireCoderSettings(getWireCoderSettings());
// Populate inputs and outputs of the stage payload and outer PTransform simultaneously.
PCollectionNode input = getInputPCollection();
@@ -220,7 +219,7 @@
static ExecutableStage fromPayload(ExecutableStagePayload payload) {
Components components = payload.getComponents();
Environment environment = payload.getEnvironment();
- WireCoderSetting wireCoderSetting = payload.getWireCoderSetting();
+ Collection<WireCoderSetting> wireCoderSettings = payload.getWireCoderSettingsList();
PCollectionNode input =
PipelineNode.pCollection(
@@ -254,12 +253,12 @@
timers,
transforms,
outputs,
- wireCoderSetting);
+ wireCoderSettings);
}
- /** The default wire coder, i.e., WINDOWED_VALUE coder. */
- WireCoderSetting DEFAULT_WIRE_CODER_SETTING =
- WireCoderSetting.newBuilder()
- .setUrn(getUrn(RunnerApi.StandardCoders.Enum.WINDOWED_VALUE))
- .build();
+ /**
+ * The default wire coder settings which returns an empty list, i.e., the WireCoder for each
+ * PCollection and timer will be a WINDOWED_VALUE coder.
+ */
+ Collection<WireCoderSetting> DEFAULT_WIRE_CODER_SETTINGS = Collections.emptyList();
}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java
index cecbee9..3d7d414 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java
@@ -31,7 +31,7 @@
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
import org.apache.beam.sdk.transforms.Flatten;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
import org.slf4j.Logger;
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java
index 00cca32..4eaad5a 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java
@@ -416,7 +416,7 @@
stage.getTimers(),
pTransformNodes,
stage.getOutputPCollections(),
- stage.getWireCoderSetting());
+ stage.getWireCoderSettings());
}
/**
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java
index 87f2076..6ec3da2 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.core.construction.graph;
-import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTING;
+import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTINGS;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import java.util.ArrayDeque;
@@ -140,7 +140,7 @@
timers,
fusedTransforms.build(),
materializedPCollections,
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
}
private static Environment getStageEnvironment(
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java
index 0092056..a66958f 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java
@@ -22,7 +22,7 @@
import java.util.stream.Collectors;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.Environment;
-import org.apache.beam.model.pipeline.v1.RunnerApi.WireCoderSetting;
+import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
@@ -39,7 +39,7 @@
Collection<TimerReference> timers,
Collection<PTransformNode> transforms,
Collection<PCollectionNode> outputs,
- WireCoderSetting wireCoderSetting) {
+ Collection<WireCoderSetting> wireCoderSettings) {
Components prunedComponents =
components
.toBuilder()
@@ -57,7 +57,7 @@
timers,
transforms,
outputs,
- wireCoderSetting);
+ wireCoderSettings);
}
public static ImmutableExecutableStage of(
@@ -69,7 +69,7 @@
Collection<TimerReference> timers,
Collection<PTransformNode> transforms,
Collection<PCollectionNode> outputs,
- WireCoderSetting wireCoderSetting) {
+ Collection<WireCoderSetting> wireCoderSettings) {
return new AutoValue_ImmutableExecutableStage(
components,
environment,
@@ -79,7 +79,7 @@
ImmutableSet.copyOf(timers),
ImmutableSet.copyOf(transforms),
ImmutableSet.copyOf(outputs),
- wireCoderSetting);
+ wireCoderSettings);
}
@Override
@@ -108,5 +108,5 @@
public abstract Collection<PCollectionNode> getOutputPCollections();
@Override
- public abstract WireCoderSetting getWireCoderSetting();
+ public abstract Collection<WireCoderSetting> getWireCoderSettings();
}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java
index def7de8..cd5d8ba 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java
@@ -309,7 +309,7 @@
stage.getTimers(),
updatedTransforms,
updatedOutputs,
- stage.getWireCoderSetting());
+ stage.getWireCoderSettings());
}
/**
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/PipelineValidator.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/PipelineValidator.java
index e46612c..c0ffb5c 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/PipelineValidator.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/PipelineValidator.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.core.construction.graph;
-import static org.apache.beam.runners.core.construction.BeamUrns.getUrn;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import java.util.Map;
@@ -28,13 +27,10 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms.CombineComponents;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms.Composites;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms.Primitives;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms.SplittableParDoComponents;
import org.apache.beam.model.pipeline.v1.RunnerApi.TestStreamPayload;
import org.apache.beam.model.pipeline.v1.RunnerApi.WindowIntoPayload;
import org.apache.beam.model.pipeline.v1.RunnerApi.WindowingStrategy;
+import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
@@ -51,32 +47,41 @@
private static final ImmutableMap<String, TransformValidator> VALIDATORS =
ImmutableMap.<String, TransformValidator>builder()
- .put(getUrn(Primitives.PAR_DO), PipelineValidator::validateParDo)
+ .put(PTransformTranslation.PAR_DO_TRANSFORM_URN, PipelineValidator::validateParDo)
// Nothing to validate for FLATTEN, GROUP_BY_KEY, IMPULSE
- .put(getUrn(Primitives.ASSIGN_WINDOWS), PipelineValidator::validateAssignWindows)
- .put(getUrn(Primitives.TEST_STREAM), PipelineValidator::validateTestStream)
+ .put(
+ PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN,
+ PipelineValidator::validateAssignWindows)
+ .put(
+ PTransformTranslation.TEST_STREAM_TRANSFORM_URN,
+ PipelineValidator::validateTestStream)
// Nothing to validate for MAP_WINDOWS, READ, CREATE_VIEW.
- .put(getUrn(Composites.COMBINE_PER_KEY), PipelineValidator::validateCombine)
- .put(getUrn(Composites.COMBINE_GLOBALLY), PipelineValidator::validateCombine)
+ .put(
+ PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN,
+ PipelineValidator::validateCombine)
+ .put(
+ PTransformTranslation.COMBINE_GLOBALLY_TRANSFORM_URN,
+ PipelineValidator::validateCombine)
// Nothing to validate for RESHUFFLE and WRITE_FILES
.put(
- getUrn(CombineComponents.COMBINE_PER_KEY_PRECOMBINE),
+ PTransformTranslation.COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN,
PipelineValidator::validateCombine)
.put(
- getUrn(CombineComponents.COMBINE_PER_KEY_MERGE_ACCUMULATORS),
+ PTransformTranslation.COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN,
PipelineValidator::validateCombine)
.put(
- getUrn(CombineComponents.COMBINE_PER_KEY_EXTRACT_OUTPUTS),
+ PTransformTranslation.COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN,
PipelineValidator::validateCombine)
- .put(getUrn(CombineComponents.COMBINE_GROUPED_VALUES), PipelineValidator::validateCombine)
.put(
- getUrn(SplittableParDoComponents.PAIR_WITH_RESTRICTION),
+ PTransformTranslation.COMBINE_GROUPED_VALUES_TRANSFORM_URN,
+ PipelineValidator::validateCombine)
+ .put(
+ PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN,
PipelineValidator::validateParDo)
.put(
- getUrn(SplittableParDoComponents.SPLIT_RESTRICTION), PipelineValidator::validateParDo)
- .put(
- getUrn(SplittableParDoComponents.PROCESS_KEYED_ELEMENTS),
+ PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN,
PipelineValidator::validateParDo)
+ .put(PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN, PipelineValidator::validateParDo)
.put(ExecutableStage.URN, PipelineValidator::validateExecutableStage)
.build();
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java
index 4ed19da..099294d 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java
@@ -62,7 +62,7 @@
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
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.collect.ImmutableSet;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ArtifactServiceStagerTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ArtifactServiceStagerTest.java
index 99c14e5..4cc2b68 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ArtifactServiceStagerTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ArtifactServiceStagerTest.java
@@ -33,9 +33,9 @@
import java.util.Set;
import org.apache.beam.model.jobmanagement.v1.ArtifactApi.ArtifactMetadata;
import org.apache.beam.runners.core.construction.ArtifactServiceStager.StagedFile;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing;
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/CommonCoderTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/CommonCoderTest.java
index 56fae2f..4ec1a7d 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/CommonCoderTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/CommonCoderTest.java
@@ -71,7 +71,7 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.Row;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
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.Splitter;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ExternalTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ExternalTest.java
index b399472..1cedfc5 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ExternalTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ExternalTest.java
@@ -38,11 +38,11 @@
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.TypeDescriptors;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ConnectivityState;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ConnectivityState;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerBuilder;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.junit.AfterClass;
import org.junit.BeforeClass;
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/InMemoryArtifactStagerService.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/InMemoryArtifactStagerService.java
index cb850bd..34431b6 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/InMemoryArtifactStagerService.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/InMemoryArtifactStagerService.java
@@ -35,7 +35,7 @@
import org.apache.beam.model.jobmanagement.v1.ArtifactApi.PutArtifactRequest.ContentCase;
import org.apache.beam.model.jobmanagement.v1.ArtifactApi.PutArtifactResponse;
import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc.ArtifactStagingServiceImplBase;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing;
/**
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslationTest.java
index a482d02..14d8c1c 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PipelineOptionsTranslationTest.java
@@ -30,9 +30,9 @@
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.common.ReflectHelpers;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.NullValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Value;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.NullValue;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Value;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowIntoTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowIntoTranslationTest.java
index 2db4d70..4b3f7aa 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowIntoTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowIntoTranslationTest.java
@@ -38,7 +38,7 @@
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.Window.Assign;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.joda.time.Duration;
import org.joda.time.Instant;
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/expansion/ExpansionServiceTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/expansion/ExpansionServiceTest.java
index 6024c10..b78f1da 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/expansion/ExpansionServiceTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/expansion/ExpansionServiceTest.java
@@ -46,7 +46,7 @@
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Impulse;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java
index 863165c..9a69ed8 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.core.construction.graph;
-import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTING;
+import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTINGS;
import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -106,7 +106,7 @@
Collections.singleton(timerRef),
Collections.singleton(PipelineNode.pTransform("pt", pt)),
Collections.singleton(PipelineNode.pCollection("output.out", output)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
PTransform stagePTransform = stage.toPTransform("foo");
assertThat(stagePTransform.getOutputsMap(), hasValue("output.out"));
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java
index 578d506..69336b7 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.core.construction.graph;
-import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTING;
+import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTINGS;
import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.equalTo;
@@ -100,7 +100,7 @@
Collections.singleton(timerRef),
Collections.singleton(PipelineNode.pTransform("pt", pt)),
Collections.singleton(PipelineNode.pCollection("output.out", output)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
assertThat(stage.getComponents().containsTransforms("pt"), is(true));
assertThat(stage.getComponents().containsTransforms("other_pt"), is(false));
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java
index 8e83eef..8d45782 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.core.construction.graph;
-import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTING;
+import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTINGS;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables.getOnlyElement;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.empty;
@@ -121,7 +121,7 @@
ImmutableList.of(),
ImmutableList.of(PipelineNode.pTransform("one", one)),
ImmutableList.of(PipelineNode.pCollection(oneOut.getUniqueName(), oneOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
ExecutableStage twoStage =
ImmutableExecutableStage.of(
components,
@@ -132,7 +132,7 @@
ImmutableList.of(),
ImmutableList.of(PipelineNode.pTransform("two", two)),
ImmutableList.of(PipelineNode.pCollection(twoOut.getUniqueName(), twoOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
PTransformNode redTransform = PipelineNode.pTransform("red", red);
PTransformNode blueTransform = PipelineNode.pTransform("blue", blue);
QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components);
@@ -241,7 +241,7 @@
ImmutableList.of(
PipelineNode.pTransform("one", one), PipelineNode.pTransform("shared", shared)),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
ExecutableStage twoStage =
ImmutableExecutableStage.of(
components,
@@ -253,7 +253,7 @@
ImmutableList.of(
PipelineNode.pTransform("two", two), PipelineNode.pTransform("shared", shared)),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
PTransformNode redTransform = PipelineNode.pTransform("red", red);
PTransformNode blueTransform = PipelineNode.pTransform("blue", blue);
QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components);
@@ -373,7 +373,7 @@
ImmutableList.of(),
ImmutableList.of(PipelineNode.pTransform("one", one), sharedTransform),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
PTransformNode redTransform = PipelineNode.pTransform("red", red);
PTransformNode blueTransform = PipelineNode.pTransform("blue", blue);
QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components);
@@ -547,7 +547,7 @@
ImmutableList.of(
PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut),
PipelineNode.pCollection(otherSharedOut.getUniqueName(), otherSharedOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
ExecutableStage oneStage =
ImmutableExecutableStage.of(
components,
@@ -559,7 +559,7 @@
ImmutableList.of(
PipelineNode.pTransform("one", one), PipelineNode.pTransform("shared", shared)),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
ExecutableStage twoStage =
ImmutableExecutableStage.of(
components,
@@ -573,7 +573,7 @@
PipelineNode.pTransform("otherShared", otherShared)),
ImmutableList.of(
PipelineNode.pCollection(otherSharedOut.getUniqueName(), otherSharedOut)),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
PTransformNode redTransform = PipelineNode.pTransform("red", red);
PTransformNode blueTransform = PipelineNode.pTransform("blue", blue);
QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components);
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ProtoOverridesTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ProtoOverridesTest.java
index 5215d72..f27f38d 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ProtoOverridesTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ProtoOverridesTest.java
@@ -36,7 +36,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
import org.apache.beam.model.pipeline.v1.RunnerApi.WindowingStrategy;
import org.apache.beam.runners.core.construction.graph.ProtoOverrides.TransformReplacement;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java
index cf28436..e2fc262 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java
@@ -38,7 +38,13 @@
* Calls a {@link DoFn DoFn's} {@link DoFn.OnTimer @OnTimer} method for the given timer in the
* given window.
*/
- void onTimer(String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain);
+ void onTimer(
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain);
/**
* Calls a {@link DoFn DoFn's} {@link DoFn.FinishBundle @FinishBundle} method and performs
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java
index 286e60b..7fbfaf0 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java
@@ -122,9 +122,14 @@
WindowTracing.trace("{}.setTimer: {}", getClass().getSimpleName(), timerData);
@Nullable
- TimerData existing = existingTimers.get(timerData.getNamespace(), timerData.getTimerId());
+ TimerData existing =
+ existingTimers.get(
+ timerData.getNamespace(), timerData.getTimerId() + '+' + timerData.getTimerFamilyId());
if (existing == null) {
- existingTimers.put(timerData.getNamespace(), timerData.getTimerId(), timerData);
+ existingTimers.put(
+ timerData.getNamespace(),
+ timerData.getTimerId() + '+' + timerData.getTimerFamilyId(),
+ timerData);
timersForDomain(timerData.getDomain()).add(timerData);
} else {
checkArgument(
@@ -138,7 +143,10 @@
NavigableSet<TimerData> timers = timersForDomain(timerData.getDomain());
timers.remove(existing);
timers.add(timerData);
- existingTimers.put(timerData.getNamespace(), timerData.getTimerId(), timerData);
+ existingTimers.put(
+ timerData.getNamespace(),
+ timerData.getTimerId() + '+' + timerData.getTimerFamilyId(),
+ timerData);
}
}
}
@@ -151,8 +159,8 @@
/** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */
@Deprecated
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
- TimerData existing = existingTimers.get(namespace, timerId);
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
+ TimerData existing = existingTimers.get(namespace, timerId + '+' + timerFamilyId);
if (existing != null) {
deleteTimer(existing);
}
@@ -163,7 +171,8 @@
@Override
public void deleteTimer(TimerData timer) {
WindowTracing.trace("{}.deleteTimer: {}", getClass().getSimpleName(), timer);
- existingTimers.remove(timer.getNamespace(), timer.getTimerId());
+ existingTimers.remove(
+ timer.getNamespace(), timer.getTimerId() + '+' + timer.getTimerFamilyId());
timersForDomain(timer.getDomain()).remove(timer);
}
@@ -321,7 +330,8 @@
if (!timers.isEmpty() && currentTime.isAfter(timers.first().getTimestamp())) {
TimerData timer = timers.pollFirst();
- existingTimers.remove(timer.getNamespace(), timer.getTimerId());
+ existingTimers.remove(
+ timer.getNamespace(), timer.getTimerId() + '+' + timer.getTimerFamilyId());
return timer;
} else {
return null;
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
index 4865e82..8f19b5f 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
@@ -82,8 +82,13 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
- doFnRunner.onTimer(timerId, window, timestamp, timeDomain);
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
+ doFnRunner.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
}
@Override
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
index 9edc558..8f09c5d 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
@@ -30,6 +30,7 @@
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.FinishBundleContext;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
@@ -137,6 +138,12 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException(
+ "Cannot access timerId as parameter outside of @OnTimer method.");
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Access to time domain not supported in ProcessElement");
@@ -213,6 +220,12 @@
throw new UnsupportedOperationException(
"Access to timers not supported in Splittable DoFn");
}
+
+ @Override
+ public TimerMap timerFamily(String tagId) {
+ throw new UnsupportedOperationException(
+ "Access to timerFamily not supported in Splittable DoFn");
+ }
});
processContext.cancelScheduledCheckpoint();
@Nullable KV<RestrictionT, Instant> residual = processContext.getTakenCheckpoint();
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java
index d65b5f4..c310c49 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java
@@ -83,7 +83,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
throw new UnsupportedOperationException("User timers unsupported in ProcessFn");
}
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/PushbackSideInputDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/PushbackSideInputDoFnRunner.java
index cc2e86a..32a61af 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/PushbackSideInputDoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/PushbackSideInputDoFnRunner.java
@@ -43,7 +43,13 @@
Iterable<WindowedValue<InputT>> processElementInReadyWindows(WindowedValue<InputT> elem);
/** Calls the underlying {@link DoFn.OnTimer} method. */
- void onTimer(String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain);
+ void onTimer(
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain);
/** Calls the underlying {@link DoFn.FinishBundle} method. */
void finishBundle();
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
index 472a9d2..0b41602 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
@@ -21,6 +21,7 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
import java.util.Collection;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -33,6 +34,7 @@
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.state.TimerSpec;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
@@ -186,18 +188,22 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
// The effective timestamp is when derived elements will have their timestamp set, if not
- // otherwise specified. If this is an event time timer, then they have the timestamp of the
- // timer itself. Otherwise, they are set to the input timestamp, which is by definition
+ // otherwise specified. If this is an event time timer, then they have the timer's output
+ // timestamp. Otherwise, they are set to the input timestamp, which is by definition
// non-late.
Instant effectiveTimestamp;
switch (timeDomain) {
case EVENT_TIME:
- effectiveTimestamp = timestamp;
+ effectiveTimestamp = outputTimestamp;
break;
-
case PROCESSING_TIME:
case SYNCHRONIZED_PROCESSING_TIME:
effectiveTimestamp = stepContext.timerInternals().currentInputWatermarkTime();
@@ -208,8 +214,8 @@
}
OnTimerArgumentProvider argumentProvider =
- new OnTimerArgumentProvider(window, effectiveTimestamp, timeDomain);
- invoker.invokeOnTimer(timerId, argumentProvider);
+ new OnTimerArgumentProvider(timerId, window, timestamp, effectiveTimestamp, timeDomain);
+ invoker.invokeOnTimer(timerId, timerFamilyId, argumentProvider);
}
private void invokeProcessElement(WindowedValue<InputT> elem) {
@@ -324,6 +330,11 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException("Cannot access timerId outside of @OnTimer method.");
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access time domain outside of @ProcessTimer method.");
@@ -370,6 +381,12 @@
throw new UnsupportedOperationException(
"Cannot access timers outside of @ProcessElement and @OnTimer methods.");
}
+
+ @Override
+ public TimerMap timerFamily(String tagId) {
+ throw new UnsupportedOperationException(
+ "Cannot access timer family outside of @ProcessElement and @OnTimer methods");
+ }
}
/** B A concrete implementation of {@link DoFn.FinishBundleContext}. */
@@ -444,6 +461,12 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException(
+ "Cannot access timerId as parameter outside of @OnTimer method.");
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access time domain outside of @ProcessTimer method.");
@@ -492,6 +515,12 @@
}
@Override
+ public TimerMap timerFamily(String tagId) {
+ throw new UnsupportedOperationException(
+ "Cannot access timerFamily outside of @ProcessElement and @OnTimer methods.");
+ }
+
+ @Override
public void output(OutputT output, Instant timestamp, BoundedWindow window) {
output(mainOutputTag, output, timestamp, window);
}
@@ -664,6 +693,12 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException(
+ "Cannot access timerId as parameter outside of @OnTimer method.");
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access time domain outside of @ProcessTimer method.");
@@ -718,6 +753,18 @@
throw new RuntimeException(e);
}
}
+
+ @Override
+ public TimerMap timerFamily(String timerFamilyId) {
+ try {
+ TimerSpec spec =
+ (TimerSpec) signature.timerFamilyDeclarations().get(timerFamilyId).field().get(fn);
+ return new TimerInternalsTimerMap(
+ timerFamilyId, window(), getNamespace(), spec, stepContext.timerInternals());
+ } catch (IllegalAccessException e) {
+ throw new RuntimeException(e);
+ }
+ }
}
/**
@@ -727,8 +774,10 @@
private class OnTimerArgumentProvider extends DoFn<InputT, OutputT>.OnTimerContext
implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {
private final BoundedWindow window;
+ private final Instant fireTimestamp;
private final Instant timestamp;
private final TimeDomain timeDomain;
+ private final String timerId;
/** Lazily initialized; should only be accessed via {@link #getNamespace()}. */
private @Nullable StateNamespace namespace;
@@ -748,9 +797,15 @@
}
private OnTimerArgumentProvider(
- BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ BoundedWindow window,
+ Instant fireTimestamp,
+ Instant timestamp,
+ TimeDomain timeDomain) {
fn.super();
+ this.timerId = timerId;
this.window = window;
+ this.fireTimestamp = fireTimestamp;
this.timestamp = timestamp;
this.timeDomain = timeDomain;
}
@@ -761,6 +816,11 @@
}
@Override
+ public Instant fireTimestamp() {
+ return fireTimestamp;
+ }
+
+ @Override
public BoundedWindow window() {
return window;
}
@@ -818,6 +878,11 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ return timerId;
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
return timeDomain();
}
@@ -872,6 +937,18 @@
}
@Override
+ public TimerMap timerFamily(String timerFamilyId) {
+ try {
+ TimerSpec spec =
+ (TimerSpec) signature.timerFamilyDeclarations().get(timerFamilyId).field().get(fn);
+ return new TimerInternalsTimerMap(
+ timerFamilyId, window(), getNamespace(), spec, stepContext.timerInternals());
+ } catch (IllegalAccessException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
public PipelineOptions getPipelineOptions() {
return options;
}
@@ -906,7 +983,10 @@
private final BoundedWindow window;
private final StateNamespace namespace;
private final String timerId;
+ private final String timerFamilyId;
private final TimerSpec spec;
+ private Instant target;
+ private Instant outputTimestamp;
private Duration period = Duration.ZERO;
private Duration offset = Duration.ZERO;
@@ -919,36 +999,36 @@
this.window = window;
this.namespace = namespace;
this.timerId = timerId;
+ this.timerFamilyId = "";
+ this.spec = spec;
+ this.timerInternals = timerInternals;
+ }
+
+ public TimerInternalsTimer(
+ BoundedWindow window,
+ StateNamespace namespace,
+ String timerId,
+ String timerFamilyId,
+ TimerSpec spec,
+ TimerInternals timerInternals) {
+ this.window = window;
+ this.namespace = namespace;
+ this.timerId = timerId;
+ this.timerFamilyId = timerFamilyId;
this.spec = spec;
this.timerInternals = timerInternals;
}
@Override
public void set(Instant target) {
- // Verifies that the time domain of this timer is acceptable for absolute timers.
- if (!TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
- throw new IllegalStateException(
- "Can only set relative timers in processing time domain. Use #setRelative()");
- }
-
- // Ensures that the target time is reasonable. For event time timers this means that the time
- // should be prior to window GC time.
- if (TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
- Instant windowExpiry = window.maxTimestamp().plus(allowedLateness);
- checkArgument(
- !target.isAfter(windowExpiry),
- "Attempted to set event time timer for %s but that is after"
- + " the expiration of window %s",
- target,
- windowExpiry);
- }
-
- setUnderlyingTimer(target);
+ this.target = target;
+ verifyAbsoluteTimeDomain();
+ setAndVerifyOutputTimestamp();
+ setUnderlyingTimer();
}
@Override
public void setRelative() {
- Instant target;
Instant now = getCurrentTime();
if (period.equals(Duration.ZERO)) {
target = now.plus(offset);
@@ -957,7 +1037,9 @@
target = millisSinceStart == 0 ? now : now.plus(period).minus(millisSinceStart);
}
target = minTargetAndGcTime(target);
- setUnderlyingTimer(target);
+
+ setAndVerifyOutputTimestamp();
+ setUnderlyingTimer();
}
@Override
@@ -986,13 +1068,58 @@
return target;
}
+ @Override
+ public Timer withOutputTimestamp(Instant outputTimestamp) {
+ this.outputTimestamp = outputTimestamp;
+ return this;
+ }
+
+ /** Verifies that the time domain of this timer is acceptable for absolute timers. */
+ private void verifyAbsoluteTimeDomain() {
+ if (!TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
+ throw new IllegalStateException(
+ "Cannot only set relative timers in processing time domain." + " Use #setRelative()");
+ }
+ }
+
+ /**
+ *
+ *
+ * <ul>
+ * Ensures that:
+ * <li>Users can't set {@code outputTimestamp} for processing time timers.
+ * <li>Event time timers' {@code outputTimestamp} is set before window expiration.
+ * </ul>
+ */
+ private void setAndVerifyOutputTimestamp() {
+ // Output timestamp is currently not supported in processing time timers.
+ if (outputTimestamp != null && !TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
+ throw new IllegalStateException("Cannot set outputTimestamp in processing time domain.");
+ }
+ // Output timestamp is set to the delivery time if not initialized by an user.
+ if (outputTimestamp == null) {
+ outputTimestamp = target;
+ }
+
+ if (TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
+ Instant windowExpiry = window.maxTimestamp().plus(allowedLateness);
+ checkArgument(
+ !target.isAfter(windowExpiry),
+ "Attempted to set event time timer that outputs for %s but that is"
+ + " after the expiration of window %s",
+ target,
+ windowExpiry);
+ }
+ }
+
/**
* Sets the timer for the target time without checking anything about whether it is a reasonable
* thing to do. For example, absolute processing time timers are not really sensible since the
* user has no way to compute a good choice of time.
*/
- private void setUnderlyingTimer(Instant target) {
- timerInternals.setTimer(namespace, timerId, "", target, target, spec.getTimeDomain());
+ private void setUnderlyingTimer() {
+ timerInternals.setTimer(
+ namespace, timerId, timerFamilyId, target, outputTimestamp, spec.getTimeDomain());
}
private Instant getCurrentTime() {
@@ -1009,4 +1136,46 @@
}
}
}
+
+ private class TimerInternalsTimerMap implements TimerMap {
+
+ Map<String, Timer> timers = new HashMap<>();
+ private final TimerInternals timerInternals;
+ private final BoundedWindow window;
+ private final StateNamespace namespace;
+ private final TimerSpec spec;
+ private final String timerFamilyId;
+
+ public TimerInternalsTimerMap(
+ String timerFamilyId,
+ BoundedWindow window,
+ StateNamespace namespace,
+ TimerSpec spec,
+ TimerInternals timerInternals) {
+ this.window = window;
+ this.namespace = namespace;
+ this.spec = spec;
+ this.timerInternals = timerInternals;
+ this.timerFamilyId = timerFamilyId;
+ }
+
+ @Override
+ public void set(String timerId, Instant absoluteTime) {
+ Timer timer =
+ new TimerInternalsTimer(window, namespace, timerId, timerFamilyId, spec, timerInternals);
+ timer.set(absoluteTime);
+ timers.put(timerId, timer);
+ }
+
+ @Override
+ public Timer get(String timerId) {
+ if (timers.get(timerId) == null) {
+ Timer timer =
+ new TimerInternalsTimer(
+ window, namespace, timerId, timerFamilyId, spec, timerInternals);
+ timers.put(timerId, timer);
+ }
+ return timers.get(timerId);
+ }
+ }
}
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java
index 36a89fe..b27e046 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java
@@ -108,8 +108,13 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
- underlying.onTimer(timerId, window, timestamp, timeDomain);
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
+ underlying.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
}
@Override
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java
index f69c74a..f6170ce 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java
@@ -38,7 +38,7 @@
/**
* A customized {@link DoFnRunner} that handles late data dropping and garbage collection for
* stateful {@link DoFn DoFns}. It registers a GC timer in {@link #processElement(WindowedValue)}
- * and does cleanup in {@link #onTimer(String, BoundedWindow, Instant, TimeDomain)}
+ * and does cleanup in {@link #onTimer(String, BoundedWindow, Instant, Instant, TimeDomain)}
*
* @param <InputT> the type of the {@link DoFn} (main) input elements
* @param <OutputT> the type of the {@link DoFn} (main) output elements
@@ -117,7 +117,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
if (cleanupTimer.isForWindow(timerId, window, timestamp, timeDomain)) {
stateCleaner.clearForWindow(window);
// There should invoke the onWindowExpiration of DoFn
@@ -134,7 +139,7 @@
window,
cleanupTimer.currentInputWatermarkTime());
} else {
- doFnRunner.onTimer(timerId, window, timestamp, timeDomain);
+ doFnRunner.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
}
}
}
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java
index f9f23ca..5ea3fb7 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java
@@ -80,7 +80,7 @@
/** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */
@Deprecated
- void deleteTimer(StateNamespace namespace, String timerId);
+ void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId);
/** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */
@Deprecated
@@ -187,9 +187,20 @@
// When adding a new field, make sure to add it to the compareTo() method.
+ /** Construct a {@link TimerData} for the given parameters. */
+ public static TimerData of(
+ String timerId,
+ StateNamespace namespace,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain domain) {
+ return new AutoValue_TimerInternals_TimerData(
+ timerId, "", namespace, timestamp, outputTimestamp, domain);
+ }
+
/**
- * Construct a {@link TimerData} for the given parameters, where the timer ID is automatically
- * generated.
+ * Construct a {@link TimerData} for the given parameters except for {@code outputTimestamp}.
+ * {@code outputTimestamp} is set to timer {@code timestamp}.
*/
public static TimerData of(
String timerId,
@@ -210,21 +221,41 @@
public static TimerData of(
String timerId, StateNamespace namespace, Instant timestamp, TimeDomain domain) {
return new AutoValue_TimerInternals_TimerData(
- timerId, timerId, namespace, timestamp, timestamp, domain);
+ timerId, "", namespace, timestamp, timestamp, domain);
+ }
+
+ public static TimerData of(
+ String timerId,
+ String timerFamilyId,
+ StateNamespace namespace,
+ Instant timestamp,
+ TimeDomain domain) {
+ return new AutoValue_TimerInternals_TimerData(
+ timerId, timerFamilyId, namespace, timestamp, timestamp, domain);
}
/**
- * Construct a {@link TimerData} for the given parameters, where the timer ID is
+ * Construct a {@link TimerData} for the given parameters except for timer ID. Timer ID is
* deterministically generated from the {@code timestamp} and {@code domain}.
*/
- public static TimerData of(StateNamespace namespace, Instant timestamp, TimeDomain domain) {
+ public static TimerData of(
+ StateNamespace namespace, Instant timestamp, Instant outputTimestamp, TimeDomain domain) {
String timerId =
new StringBuilder()
.append(domain.ordinal())
.append(':')
.append(timestamp.getMillis())
.toString();
- return of(timerId, namespace, timestamp, domain);
+ return of(timerId, namespace, timestamp, outputTimestamp, domain);
+ }
+
+ /**
+ * Construct a {@link TimerData} for the given parameters, where the timer ID is
+ * deterministically generated from the {@code timestamp} and {@code domain}. Also, output
+ * timestamp is set to the timer timestamp by default.
+ */
+ public static TimerData of(StateNamespace namespace, Instant timestamp, TimeDomain domain) {
+ return of(namespace, timestamp, timestamp, domain);
}
/**
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleExecutionState.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleExecutionState.java
index 07eeeba..881261c 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleExecutionState.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/SimpleExecutionState.java
@@ -97,7 +97,7 @@
String userStepName =
this.labelsMetadata.getOrDefault(MonitoringInfoConstants.Labels.PTRANSFORM, null);
StringBuilder message = new StringBuilder();
- message.append("Processing stuck");
+ message.append("Operation ongoing");
if (userStepName != null) {
message.append(" in step ").append(userStepName);
}
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java
index dd106329..3ab7932 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java
@@ -71,15 +71,15 @@
Instant laterTimestamp = new Instant(42);
underTest.advanceInputWatermark(new Instant(0));
- underTest.setTimer(NS1, ID1, ID1, earlyTimestamp, earlyTimestamp, TimeDomain.EVENT_TIME);
- underTest.setTimer(NS1, ID1, ID1, laterTimestamp, laterTimestamp, TimeDomain.EVENT_TIME);
+ underTest.setTimer(NS1, ID1, "", earlyTimestamp, earlyTimestamp, TimeDomain.EVENT_TIME);
+ underTest.setTimer(NS1, ID1, "", laterTimestamp, laterTimestamp, TimeDomain.EVENT_TIME);
underTest.advanceInputWatermark(earlyTimestamp.plus(1L));
assertThat(underTest.removeNextEventTimer(), nullValue());
underTest.advanceInputWatermark(laterTimestamp.plus(1L));
assertThat(
underTest.removeNextEventTimer(),
- equalTo(TimerData.of(ID1, NS1, laterTimestamp, TimeDomain.EVENT_TIME)));
+ equalTo(TimerData.of(ID1, "", NS1, laterTimestamp, TimeDomain.EVENT_TIME)));
}
@Test
@@ -87,8 +87,8 @@
InMemoryTimerInternals underTest = new InMemoryTimerInternals();
Instant timestamp = new Instant(42);
underTest.setTimer(NS1, ID1, ID1, timestamp, timestamp, TimeDomain.EVENT_TIME);
- underTest.deleteTimer(NS1, ID1);
- underTest.deleteTimer(NS1, ID1);
+ underTest.deleteTimer(NS1, ID1, ID1);
+ underTest.deleteTimer(NS1, ID1, ID1);
}
@Test
@@ -98,7 +98,7 @@
underTest.advanceInputWatermark(new Instant(0));
underTest.setTimer(NS1, ID1, ID1, timestamp, timestamp, TimeDomain.EVENT_TIME);
- underTest.deleteTimer(NS1, ID1);
+ underTest.deleteTimer(NS1, ID1, ID1);
underTest.advanceInputWatermark(new Instant(43));
assertThat(underTest.removeNextEventTimer(), nullValue());
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java
index 10972d6..f3f1628 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java
@@ -119,7 +119,12 @@
thrown.expectCause(is(fn.exceptionToThrow));
runner.onTimer(
- ThrowingDoFn.TIMER_ID, GlobalWindow.INSTANCE, new Instant(0), TimeDomain.EVENT_TIME);
+ ThrowingDoFn.TIMER_ID,
+ ThrowingDoFn.TIMER_ID,
+ GlobalWindow.INSTANCE,
+ new Instant(0),
+ new Instant(0),
+ TimeDomain.EVENT_TIME);
}
/**
@@ -239,8 +244,10 @@
// the method call.
runner.onTimer(
DoFnWithTimers.TIMER_ID,
+ "",
GlobalWindow.INSTANCE,
currentTime.plus(offset),
+ currentTime.plus(offset),
TimeDomain.EVENT_TIME);
assertThat(
@@ -248,8 +255,10 @@
contains(
TimerData.of(
DoFnWithTimers.TIMER_ID,
+ "",
StateNamespaces.window(windowFn.windowCoder(), GlobalWindow.INSTANCE),
currentTime.plus(offset),
+ currentTime.plus(offset),
TimeDomain.EVENT_TIME)));
}
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java
index 28b387e..8f88d41 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java
@@ -283,7 +283,8 @@
// Mocking is not easily compatible with annotation analysis, so we manually record
// the method call.
- runner.onTimer(timerId, window, new Instant(timestamp), TimeDomain.EVENT_TIME);
+ runner.onTimer(
+ timerId, "", window, new Instant(timestamp), new Instant(timestamp), TimeDomain.EVENT_TIME);
assertThat(
underlying.firedTimers,
@@ -320,12 +321,19 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
firedTimers.add(
TimerData.of(
timerId,
+ timerFamilyId,
StateNamespaces.window(IntervalWindow.getCoder(), (IntervalWindow) window),
timestamp,
+ outputTimestamp,
timeDomain));
}
@@ -458,7 +466,13 @@
StateNamespace namespace = timer.getNamespace();
checkArgument(namespace instanceof StateNamespaces.WindowNamespace);
BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
- toTrigger.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ toTrigger.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
}
}
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java
index 85b3c0b..be4e321 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java
@@ -220,7 +220,13 @@
StateNamespace namespace = timer.getNamespace();
checkArgument(namespace instanceof StateNamespaces.WindowNamespace);
BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
- toTrigger.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ toTrigger.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
}
}
diff --git a/runners/direct-java/build.gradle b/runners/direct-java/build.gradle
index b8836a8..6d46528 100644
--- a/runners/direct-java/build.gradle
+++ b/runners/direct-java/build.gradle
@@ -65,7 +65,7 @@
compile project(it)
}
shadow project(path: ":sdks:java:core", configuration: "shadow")
- shadow library.java.vendored_grpc_1_21_0
+ shadow library.java.vendored_grpc_1_26_0
shadow library.java.joda_time
shadow library.java.slf4j_api
shadow library.java.args4j
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java
index 0261bf6..7b6dbf1 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java
@@ -75,7 +75,7 @@
/** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */
@Deprecated
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
throw new UnsupportedOperationException("Canceling of timer by ID is not yet supported.");
}
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java
index 31eb80b..5f41175 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluator.java
@@ -222,7 +222,13 @@
public void onTimer(TimerData timer, BoundedWindow window) {
try {
- fnRunner.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ fnRunner.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
} catch (Exception e) {
throw UserCodeException.wrap(e);
}
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
index f99d07f..f30ee41 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
@@ -103,7 +103,8 @@
if (signature.processElement().isSplittable()) {
return SplittableParDo.forAppliedParDo((AppliedPTransform) application);
} else if (signature.stateDeclarations().size() > 0
- || signature.timerDeclarations().size() > 0) {
+ || signature.timerDeclarations().size() > 0
+ || signature.timerFamilyDeclarations().size() > 0) {
return new GbkThenStatefulParDo(
fn,
ParDoTranslation.getMainOutputTag(application),
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java
index e1080e5..0dd826d 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java
@@ -44,12 +44,14 @@
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.WatermarkHoldState;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.StateDeclaration;
import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
@@ -71,6 +73,8 @@
private final ParDoEvaluatorFactory<KV<K, InputT>, OutputT> delegateFactory;
+ private final EvaluationContext evaluationContext;
+
StatefulParDoEvaluatorFactory(EvaluationContext evaluationContext, PipelineOptions options) {
this.delegateFactory =
new ParDoEvaluatorFactory<>(
@@ -92,6 +96,8 @@
CacheBuilder.newBuilder()
.weakValues()
.build(new CleanupSchedulingLoader(evaluationContext));
+
+ this.evaluationContext = evaluationContext;
}
@Override
@@ -146,7 +152,13 @@
application.getTransform().getSchemaInformation(),
application.getTransform().getSideInputMapping());
- return new StatefulParDoEvaluator<>(delegateEvaluator);
+ DirectStepContext stepContext =
+ evaluationContext
+ .getExecutionContext(application, inputBundle.getKey())
+ .getStepContext(evaluationContext.getStepName(application));
+
+ stepContext.stateInternals().commit();
+ return new StatefulParDoEvaluator<>(delegateEvaluator, stepContext);
}
private class CleanupSchedulingLoader
@@ -241,10 +253,14 @@
private final List<TimerData> pushedBackTimers = new ArrayList<>();
private final DirectTimerInternals timerInternals;
+ DirectStepContext stepContext;
+
public StatefulParDoEvaluator(
- DoFnLifecycleManagerRemovingTransformEvaluator<KV<K, InputT>> delegateEvaluator) {
+ DoFnLifecycleManagerRemovingTransformEvaluator<KV<K, InputT>> delegateEvaluator,
+ DirectStepContext stepContext) {
this.delegateEvaluator = delegateEvaluator;
this.timerInternals = delegateEvaluator.getParDoEvaluator().getStepContext().timerInternals();
+ this.stepContext = stepContext;
}
@Override
@@ -269,6 +285,12 @@
WindowNamespace<?> windowNamespace = (WindowNamespace) timer.getNamespace();
BoundedWindow timerWindow = windowNamespace.getWindow();
delegateEvaluator.onTimer(timer, timerWindow);
+
+ StateTag<WatermarkHoldState> timerWatermarkHoldTag = setTimerTag(timer);
+
+ stepContext.stateInternals().state(timer.getNamespace(), timerWatermarkHoldTag).clear();
+ stepContext.stateInternals().commit();
+
if (timerInternals.containsUpdateForTimeBefore(currentInputWatermark)) {
break;
}
@@ -278,15 +300,41 @@
@Override
public TransformResult<KeyedWorkItem<K, KV<K, InputT>>> finishBundle() throws Exception {
+
TransformResult<KV<K, InputT>> delegateResult = delegateEvaluator.finishBundle();
+ boolean isTimerDeclared = false;
+ for (TimerData timerData : delegateResult.getTimerUpdate().getSetTimers()) {
+ StateTag<WatermarkHoldState> timerWatermarkHoldTag = setTimerTag(timerData);
+
+ stepContext
+ .stateInternals()
+ .state(timerData.getNamespace(), timerWatermarkHoldTag)
+ .add(timerData.getOutputTimestamp());
+ isTimerDeclared = true;
+ }
+
+ CopyOnAccessInMemoryStateInternals state;
+ Instant watermarkHold;
+
+ if (isTimerDeclared && delegateResult.getState() != null) { // For both State and Timer Holds
+ state = delegateResult.getState();
+ watermarkHold = stepContext.commitState().getEarliestWatermarkHold();
+ } else if (isTimerDeclared) { // For only Timer holds
+ state = stepContext.commitState();
+ watermarkHold = state.getEarliestWatermarkHold();
+ } else { // For only State ( non Timer ) holds
+ state = delegateResult.getState();
+ watermarkHold = delegateResult.getWatermarkHold();
+ }
+
TimerUpdate timerUpdate =
delegateResult.getTimerUpdate().withPushedBackTimers(pushedBackTimers);
pushedBackTimers.clear();
StepTransformResult.Builder<KeyedWorkItem<K, KV<K, InputT>>> regroupedResult =
StepTransformResult.<KeyedWorkItem<K, KV<K, InputT>>>withHold(
- delegateResult.getTransform(), delegateResult.getWatermarkHold())
+ delegateResult.getTransform(), watermarkHold)
.withTimerUpdate(timerUpdate)
- .withState(delegateResult.getState())
+ .withState(state)
.withMetricUpdates(delegateResult.getLogicalMetricUpdates())
.addOutput(Lists.newArrayList(delegateResult.getOutputBundles()));
@@ -306,4 +354,11 @@
return regroupedResult.build();
}
}
+
+ private static StateTag<WatermarkHoldState> setTimerTag(TimerData timerData) {
+ return StateTags.makeSystemTagInternal(
+ StateTags.watermarkStateInternal(
+ "timer-" + timerData.getTimerId() + "+" + timerData.getTimerFamilyId(),
+ TimestampCombiner.EARLIEST));
+ }
}
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
index d9e7ac2..265ebb1 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
@@ -327,7 +327,7 @@
if (pendingTimers.isEmpty()) {
return BoundedWindow.TIMESTAMP_MAX_VALUE;
} else {
- return pendingTimers.firstEntry().getElement().getTimestamp();
+ return pendingTimers.firstEntry().getElement().getOutputTimestamp();
}
}
@@ -342,7 +342,8 @@
if (TimeDomain.EVENT_TIME.equals(timer.getDomain())) {
@Nullable
TimerData existingTimer =
- existingTimersForKey.get(timer.getNamespace(), timer.getTimerId());
+ existingTimersForKey.get(
+ timer.getNamespace(), timer.getTimerId() + '+' + timer.getTimerFamilyId());
if (existingTimer == null) {
pendingTimers.add(timer);
@@ -357,7 +358,8 @@
keyTimers.add(timer);
}
- existingTimersForKey.put(timer.getNamespace(), timer.getTimerId(), timer);
+ existingTimersForKey.put(
+ timer.getNamespace(), timer.getTimerId() + '+' + timer.getTimerFamilyId(), timer);
}
}
@@ -365,12 +367,15 @@
if (TimeDomain.EVENT_TIME.equals(timer.getDomain())) {
@Nullable
TimerData existingTimer =
- existingTimersForKey.get(timer.getNamespace(), timer.getTimerId());
+ existingTimersForKey.get(
+ timer.getNamespace(), timer.getTimerId() + '+' + timer.getTimerFamilyId());
if (existingTimer != null) {
pendingTimers.remove(existingTimer);
keyTimers.remove(existingTimer);
- existingTimersForKey.remove(existingTimer.getNamespace(), existingTimer.getTimerId());
+ existingTimersForKey.remove(
+ existingTimer.getNamespace(),
+ existingTimer.getTimerId() + '+' + existingTimer.getTimerFamilyId());
}
}
}
@@ -465,7 +470,8 @@
Instant oldWatermark = currentWatermark.get();
Instant newWatermark =
INSTANT_ORDERING.min(
- inputWatermark.get(), inputWatermark.getEarliestTimerTimestamp(), holds.getMinHold());
+ inputWatermark.get(), holds.getMinHold(), inputWatermark.getEarliestTimerTimestamp());
+
newWatermark = INSTANT_ORDERING.max(oldWatermark, newWatermark);
currentWatermark.set(newWatermark);
return updateAndTrace(getName(), oldWatermark, newWatermark);
@@ -618,7 +624,9 @@
@Nullable
TimerData existingTimer =
- existingTimersForKey.get(addedTimer.getNamespace(), addedTimer.getTimerId());
+ existingTimersForKey.get(
+ addedTimer.getNamespace(),
+ addedTimer.getTimerId() + '+' + addedTimer.getTimerFamilyId());
if (existingTimer == null) {
timerQueue.add(addedTimer);
} else if (!existingTimer.equals(addedTimer)) {
@@ -626,7 +634,10 @@
timerQueue.add(addedTimer);
} // else the timer is already set identically, so noop.
- existingTimersForKey.put(addedTimer.getNamespace(), addedTimer.getTimerId(), addedTimer);
+ existingTimersForKey.put(
+ addedTimer.getNamespace(),
+ addedTimer.getTimerId() + '+' + addedTimer.getTimerFamilyId(),
+ addedTimer);
}
for (TimerData deletedTimer : update.deletedTimers) {
@@ -637,12 +648,16 @@
@Nullable
TimerData existingTimer =
- existingTimersForKey.get(deletedTimer.getNamespace(), deletedTimer.getTimerId());
+ existingTimersForKey.get(
+ deletedTimer.getNamespace(),
+ deletedTimer.getTimerId() + '+' + deletedTimer.getTimerFamilyId());
if (existingTimer != null) {
pendingTimers.remove(deletedTimer);
timerQueue.remove(deletedTimer);
- existingTimersForKey.remove(existingTimer.getNamespace(), existingTimer.getTimerId());
+ existingTimersForKey.remove(
+ existingTimer.getNamespace(),
+ existingTimer.getTimerId() + '+' + existingTimer.getTimerFamilyId());
}
}
@@ -956,7 +971,7 @@
Map<ExecutableT, Set<String>> transformsWithAlreadyExtractedTimers, ExecutableT executable) {
return update -> {
- String timerIdWithNs = TimerUpdate.getTimerIdWithNamespace(update);
+ String timerIdWithNs = TimerUpdate.getTimerIdAndTimerFamilyIdWithNamespace(update);
transformsWithAlreadyExtractedTimers.compute(
executable,
(k, v) -> {
@@ -1228,7 +1243,8 @@
v = new HashSet<>();
}
final Set<String> toUpdate = v;
- newTimers.forEach(td -> toUpdate.add(TimerUpdate.getTimerIdWithNamespace(td)));
+ newTimers.forEach(
+ td -> toUpdate.add(TimerUpdate.getTimerIdAndTimerFamilyIdWithNamespace(td)));
return v;
});
allTimers.addAll(firedTimers);
@@ -1583,11 +1599,13 @@
private static Map<String, TimerData> indexTimerData(Iterable<? extends TimerData> timerData) {
return StreamSupport.stream(timerData.spliterator(), false)
- .collect(Collectors.toMap(TimerUpdate::getTimerIdWithNamespace, e -> e, (a, b) -> b));
+ .collect(
+ Collectors.toMap(
+ TimerUpdate::getTimerIdAndTimerFamilyIdWithNamespace, e -> e, (a, b) -> b));
}
- private static String getTimerIdWithNamespace(TimerData td) {
- return td.getNamespace() + td.getTimerId();
+ private static String getTimerIdAndTimerFamilyIdWithNamespace(TimerData td) {
+ return td.getNamespace() + td.getTimerId() + td.getTimerFamilyId();
}
private TimerUpdate(
@@ -1644,7 +1662,7 @@
Set<TimerData> pushedBack = Sets.newHashSet(pushedBackTimers);
Map<String, TimerData> newSetTimers = indexTimerData(setTimers);
for (TimerData td : completedTimers) {
- String timerIdWithNs = getTimerIdWithNamespace(td);
+ String timerIdWithNs = getTimerIdAndTimerFamilyIdWithNamespace(td);
if (!pushedBack.contains(td)) {
timersToComplete.add(td);
} else if (!newSetTimers.containsKey(timerIdWithNs)) {
diff --git a/runners/flink/flink_runner.gradle b/runners/flink/flink_runner.gradle
index 1c617d8..9817358 100644
--- a/runners/flink/flink_runner.gradle
+++ b/runners/flink/flink_runner.gradle
@@ -139,7 +139,7 @@
compile project(":runners:core-construction-java")
compile project(":runners:java-fn-execution")
compile project(":sdks:java:extensions:google-cloud-platform-core")
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
compile library.java.jackson_annotations
compile library.java.slf4j_api
compile library.java.joda_time
@@ -154,7 +154,6 @@
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
// FlinkStateInternalsTest extends abstract StateInternalsTest
testCompile project(path: ":runners:core-java", configuration: "testRuntime")
- testCompile library.java.commons_lang3
testCompile library.java.hamcrest_core
testCompile library.java.junit
testCompile library.java.mockito_core
@@ -192,11 +191,15 @@
])
systemProperty "beamTestPipelineOptions", pipelineOptions
classpath = configurations.validatesRunner
- testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs, project(":runners:core-java").sourceSets.test.output.classesDirs)
+ testClassesDirs = files(
+ project(":sdks:java:core").sourceSets.test.output.classesDirs,
+ project(":runners:core-java").sourceSets.test.output.classesDirs,
+ )
// maxParallelForks decreased from 4 in order to avoid OOM errors
maxParallelForks 2
useJUnit {
includeCategories 'org.apache.beam.sdk.testing.ValidatesRunner'
+ excludeCategories 'org.apache.beam.sdk.testing.UsesTimerMap'
excludeCategories 'org.apache.beam.sdk.testing.FlattenWithHeterogeneousCoders'
excludeCategories 'org.apache.beam.sdk.testing.LargeKeys$Above100MB'
excludeCategories 'org.apache.beam.sdk.testing.UsesCommittedMetrics'
diff --git a/runners/flink/job-server/flink_job_server.gradle b/runners/flink/job-server/flink_job_server.gradle
index 27a116f..12b4d65 100644
--- a/runners/flink/job-server/flink_job_server.gradle
+++ b/runners/flink/job-server/flink_job_server.gradle
@@ -151,6 +151,7 @@
if (streaming) {
excludeCategories 'org.apache.beam.sdk.testing.UsesTestStreamWithProcessingTime'
excludeCategories 'org.apache.beam.sdk.testing.UsesTestStreamWithMultipleStages'
+ excludeCategories 'org.apache.beam.sdk.testing.UsesTestStreamWithOutputTimestamp'
} else {
excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream'
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java
index ee40fb6..84a2e05 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java
@@ -75,7 +75,7 @@
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.BiMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkJobInvoker.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkJobInvoker.java
index 6e65637..67f664e 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkJobInvoker.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkJobInvoker.java
@@ -19,18 +19,18 @@
import static org.apache.beam.runners.core.construction.resources.PipelineResources.detectClassPathResourcesToStage;
-import java.io.IOException;
import java.util.UUID;
import javax.annotation.Nullable;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.core.construction.PipelineOptionsTranslation;
+import org.apache.beam.runners.flink.translation.utils.Workarounds;
import org.apache.beam.runners.fnexecution.jobsubmission.JobInvocation;
import org.apache.beam.runners.fnexecution.jobsubmission.JobInvoker;
import org.apache.beam.runners.fnexecution.jobsubmission.PortablePipelineJarCreator;
import org.apache.beam.runners.fnexecution.jobsubmission.PortablePipelineRunner;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.sdk.options.PortablePipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ListeningExecutorService;
import org.slf4j.Logger;
@@ -56,8 +56,9 @@
RunnerApi.Pipeline pipeline,
Struct options,
@Nullable String retrievalToken,
- ListeningExecutorService executorService)
- throws IOException {
+ ListeningExecutorService executorService) {
+ Workarounds.restoreOriginalStdOutAndStdErrIfApplicable();
+
// TODO: How to make Java/Python agree on names of keys and their values?
LOG.trace("Parsing pipeline options");
FlinkPipelineOptions flinkOptions =
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java
index 137aa00..f46bf7b 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java
@@ -45,7 +45,7 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.PortablePipelineOptions;
import org.apache.beam.sdk.options.PortablePipelineOptions.RetrievalServiceType;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.flink.api.common.JobExecutionResult;
import org.apache.flink.client.program.DetachedEnvironment;
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkRunner.java
index 3c53a42..9454ba2 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkRunner.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkRunner.java
@@ -26,6 +26,7 @@
import java.util.SortedSet;
import java.util.TreeSet;
import org.apache.beam.runners.core.metrics.MetricsPusher;
+import org.apache.beam.runners.flink.translation.utils.Workarounds;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.PipelineRunner;
@@ -89,9 +90,10 @@
return new FlinkRunner(flinkOptions);
}
- private FlinkRunner(FlinkPipelineOptions options) {
+ protected FlinkRunner(FlinkPipelineOptions options) {
this.options = options;
this.ptransformViewsWithNonDeterministicKeyCoders = new HashSet<>();
+ Workarounds.restoreOriginalStdOutAndStdErrIfApplicable();
}
@Override
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
index 92b07a4..cbc437b 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPortablePipelineTranslator.java
@@ -90,7 +90,7 @@
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.beam.sdk.values.ValueWithRecordId;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.BiMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.HashMultiset;
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
index cdb3060..d98a601 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
@@ -1310,46 +1310,12 @@
new CreateStreamingFlinkViewPayloadTranslator())
.put(
SplittableParDoViaKeyedWorkItems.ProcessElements.class,
- new SplittableParDoProcessElementsTranslator())
- .put(
- SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class,
- new SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator())
+ PTransformTranslation.TransformPayloadTranslator.NotSerializable.forUrn(
+ SPLITTABLE_PROCESS_URN))
.build();
}
}
- /**
- * A translator just to vend the URN. This will need to be moved to runners-core-construction-java
- * once SDF is reorganized appropriately.
- */
- private static class SplittableParDoProcessElementsPayloadTranslator
- extends PTransformTranslation.TransformPayloadTranslator.NotSerializable<
- SplittableParDoViaKeyedWorkItems.ProcessElements<?, ?, ?, ?>> {
-
- private SplittableParDoProcessElementsPayloadTranslator() {}
-
- @Override
- public String getUrn(SplittableParDoViaKeyedWorkItems.ProcessElements<?, ?, ?, ?> transform) {
- return SPLITTABLE_PROCESS_URN;
- }
- }
-
- /**
- * A translator just to vend the URN. This will need to be moved to runners-core-construction-java
- * once SDF is reorganized appropriately.
- */
- private static class SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator
- extends PTransformTranslation.TransformPayloadTranslator.NotSerializable<
- SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems<?, ?>> {
-
- private SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator() {}
-
- @Override
- public String getUrn(SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems<?, ?> transform) {
- return SplittableParDo.SPLITTABLE_GBKIKWI_URN;
- }
- }
-
/** A translator just to vend the URN. */
private static class CreateStreamingFlinkViewPayloadTranslator
extends PTransformTranslation.TransformPayloadTranslator.NotSerializable<
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java
index 9d853a2..ce54d5b 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java
@@ -68,12 +68,14 @@
@Override
public void onTimer(
final String timerId,
+ final String timerFamilyId,
final BoundedWindow window,
final Instant timestamp,
+ final Instant outputTimestamp,
final TimeDomain timeDomain) {
try (Closeable ignored =
MetricsEnvironment.scopedMetricsContainer(container.getMetricsContainer(stepName))) {
- delegate.onTimer(timerId, window, timestamp, timeDomain);
+ delegate.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
} catch (IOException e) {
throw new RuntimeException(e);
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkDoFnFunction.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkDoFnFunction.java
index fa6867f..a34d840 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkDoFnFunction.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkDoFnFunction.java
@@ -26,7 +26,7 @@
import org.apache.beam.runners.flink.FlinkPipelineOptions;
import org.apache.beam.runners.flink.metrics.DoFnRunnerWithMetricsUpdate;
import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
-import org.apache.beam.runners.flink.translation.utils.FlinkClassloading;
+import org.apache.beam.runners.flink.translation.utils.Workarounds;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.options.PipelineOptions;
@@ -162,7 +162,7 @@
metricContainer.registerMetricsForPipelineResult();
Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
} finally {
- FlinkClassloading.deleteStaticCaches();
+ Workarounds.deleteStaticCaches();
}
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java
index 6aee09f..d7fc2de 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java
@@ -35,7 +35,7 @@
import org.apache.beam.runners.flink.FlinkPipelineOptions;
import org.apache.beam.runners.flink.metrics.DoFnRunnerWithMetricsUpdate;
import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
-import org.apache.beam.runners.flink.translation.utils.FlinkClassloading;
+import org.apache.beam.runners.flink.translation.utils.Workarounds;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.options.PipelineOptions;
@@ -215,7 +215,13 @@
StateNamespace namespace = timer.getNamespace();
checkArgument(namespace instanceof StateNamespaces.WindowNamespace);
BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
- doFnRunner.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ doFnRunner.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
}
@Override
@@ -234,7 +240,7 @@
metricContainer.registerMetricsForPipelineResult();
Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
} finally {
- FlinkClassloading.deleteStaticCaches();
+ Workarounds.deleteStaticCaches();
}
}
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/Workarounds.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/Workarounds.java
new file mode 100644
index 0000000..77baba3
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/Workarounds.java
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.flink.translation.utils;
+
+import com.fasterxml.jackson.databind.type.TypeFactory;
+import java.io.FileDescriptor;
+import java.io.FileOutputStream;
+import java.io.PrintStream;
+import org.apache.flink.api.java.ExecutionEnvironment;
+import org.apache.flink.client.program.OptimizerPlanEnvironment;
+
+/** Workarounds for dealing with limitations of Flink or its libraries. */
+public class Workarounds {
+
+ public static void deleteStaticCaches() {
+ // Clear cache to get rid of any references to the Flink Classloader
+ // See https://jira.apache.org/jira/browse/BEAM-6460
+ TypeFactory.defaultInstance().clearCache();
+ }
+
+ /**
+ * Flink uses the {@link org.apache.flink.client.program.OptimizerPlanEnvironment} which replaces
+ * stdout/stderr during job graph creation. This was intended only for previewing the plan, but
+ * other parts of Flink, e.g. the Rest API have started to use this code as well. To be able to
+ * inspect the output before execution, we use this method to restore the original stdout/stderr.
+ */
+ public static void restoreOriginalStdOutAndStdErrIfApplicable() {
+ if (ExecutionEnvironment.getExecutionEnvironment() instanceof OptimizerPlanEnvironment) {
+ System.setOut(new PrintStream(new FileOutputStream(FileDescriptor.out)));
+ System.setErr(new PrintStream(new FileOutputStream(FileDescriptor.err)));
+ }
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java
index 8e63679..dbefa68 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java
@@ -58,8 +58,8 @@
import org.apache.beam.runners.flink.metrics.DoFnRunnerWithMetricsUpdate;
import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
import org.apache.beam.runners.flink.translation.types.CoderTypeSerializer;
-import org.apache.beam.runners.flink.translation.utils.FlinkClassloading;
import org.apache.beam.runners.flink.translation.utils.NoopLock;
+import org.apache.beam.runners.flink.translation.utils.Workarounds;
import org.apache.beam.runners.flink.translation.wrappers.streaming.stableinput.BufferingDoFnRunner;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkBroadcastStateInternals;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkStateInternals;
@@ -467,7 +467,7 @@
public void dispose() throws Exception {
try {
Optional.ofNullable(checkFinishBundleTimer).ifPresent(timer -> timer.cancel(true));
- FlinkClassloading.deleteStaticCaches();
+ Workarounds.deleteStaticCaches();
Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
} finally {
// This releases all task's resources. We need to call this last
@@ -653,7 +653,10 @@
Instant watermarkHold = keyedStateInternals.watermarkHold();
long combinedWatermarkHold = Math.min(watermarkHold.getMillis(), getPushbackWatermarkHold());
-
+ if (timerInternals.getWatermarkHoldMs() < Long.MAX_VALUE) {
+ combinedWatermarkHold =
+ Math.min(combinedWatermarkHold, timerInternals.getWatermarkHoldMs());
+ }
long potentialOutputWatermark = Math.min(pushedBackInputWatermark, combinedWatermarkHold);
if (potentialOutputWatermark > currentOutputWatermark) {
@@ -811,9 +814,14 @@
// This is a user timer, so namespace must be WindowNamespace
checkArgument(namespace instanceof WindowNamespace);
BoundedWindow window = ((WindowNamespace) namespace).getWindow();
- timerInternals.cleanupPendingTimer(timer.getNamespace());
+ timerInternals.cleanupPendingTimer(timer.getNamespace(), true);
pushbackDoFnRunner.onTimer(
- timerData.getTimerId(), window, timerData.getTimestamp(), timerData.getDomain());
+ timerData.getTimerId(),
+ timerData.getTimerFamilyId(),
+ window,
+ timerData.getTimestamp(),
+ timerData.getOutputTimestamp(),
+ timerData.getDomain());
}
private void setCurrentInputWatermark(long currentInputWatermark) {
@@ -1079,6 +1087,8 @@
*/
final MapState<String, TimerData> pendingTimersById;
+ long watermarkHoldMs = Long.MAX_VALUE;
+
private FlinkTimerInternals() {
MapStateDescriptor<String, TimerData> pendingTimersByIdStateDescriptor =
new MapStateDescriptor<>(
@@ -1086,6 +1096,22 @@
this.pendingTimersById = getKeyedStateStore().getMapState(pendingTimersByIdStateDescriptor);
}
+ long getWatermarkHoldMs() {
+ return watermarkHoldMs;
+ }
+
+ void updateWatermarkHold() {
+ this.watermarkHoldMs = Long.MAX_VALUE;
+ try {
+ for (TimerData timerData : pendingTimersById.values()) {
+ this.watermarkHoldMs =
+ Math.min(timerData.getOutputTimestamp().getMillis(), this.watermarkHoldMs);
+ }
+ } catch (Exception e) {
+ throw new RuntimeException("Exception while reading set of timers", e);
+ }
+ }
+
@Override
public void setTimer(
StateNamespace namespace,
@@ -1113,6 +1139,7 @@
// before we set the new one.
cancelPendingTimerById(contextTimerId);
registerTimer(timer, contextTimerId);
+ updateWatermarkHold();
} catch (Exception e) {
throw new RuntimeException("Failed to set timer", e);
}
@@ -1137,13 +1164,16 @@
private void cancelPendingTimerById(String contextTimerId) throws Exception {
TimerData oldTimer = pendingTimersById.get(contextTimerId);
if (oldTimer != null) {
- deleteTimer(oldTimer);
+ deleteTimerInternal(oldTimer, false);
}
}
- void cleanupPendingTimer(TimerData timer) {
+ void cleanupPendingTimer(TimerData timer, boolean updateWatermark) {
try {
pendingTimersById.remove(getContextTimerId(timer.getTimerId(), timer.getNamespace()));
+ if (updateWatermark) {
+ updateWatermarkHold();
+ }
} catch (Exception e) {
throw new RuntimeException("Failed to cleanup state with pending timers", e);
}
@@ -1157,7 +1187,7 @@
/** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */
@Deprecated
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
throw new UnsupportedOperationException("Canceling of a timer by ID is not yet supported.");
}
@@ -1165,16 +1195,21 @@
public void deleteTimer(StateNamespace namespace, String timerId, TimeDomain timeDomain) {
try {
cancelPendingTimerById(getContextTimerId(timerId, namespace));
+ updateWatermarkHold();
} catch (Exception e) {
throw new RuntimeException("Failed to cancel timer", e);
}
}
/** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */
- @Deprecated
@Override
+ @Deprecated
public void deleteTimer(TimerData timerKey) {
- cleanupPendingTimer(timerKey);
+ deleteTimerInternal(timerKey, true);
+ }
+
+ void deleteTimerInternal(TimerData timerKey, boolean updateWatermark) {
+ cleanupPendingTimer(timerKey, true);
long time = timerKey.getTimestamp().getMillis();
switch (timerKey.getDomain()) {
case EVENT_TIME:
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
index 891cef8..3af14ca 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
@@ -87,7 +87,7 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.sdk.v2.sdk.extensions.protobuf.ByteStringCoder;
@@ -661,7 +661,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
Object timerKey = keyForTimer.get();
Preconditions.checkNotNull(timerKey, "Key for timer needs to be set before calling onTimer");
Preconditions.checkNotNull(remoteBundle, "Call to onTimer outside of a bundle");
@@ -674,7 +679,7 @@
WindowedValue<KV<Object, Timer>> timerValue =
WindowedValue.of(
KV.of(timerKey, Timer.of(timestamp, new byte[0])),
- timestamp,
+ outputTimestamp,
Collections.singleton(window),
PaneInfo.NO_FIRING);
try {
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtils.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtils.java
index ccd10d4..3e27bb1 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtils.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtils.java
@@ -32,7 +32,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StructuredCoder;
import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
/**
* Utility functions for dealing with key encoding. Beam requires keys to be compared in binary
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java
index 4f8e0d7..e955c19 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java
@@ -159,7 +159,7 @@
@Override
protected void fireTimer(InternalTimer<ByteBuffer, TimerInternals.TimerData> timer) {
- timerInternals.cleanupPendingTimer(timer.getNamespace());
+ timerInternals.cleanupPendingTimer(timer.getNamespace(), true);
if (timer.getNamespace().getDomain().equals(TimeDomain.EVENT_TIME)) {
// ignore this, it can only be a state cleanup timers from StatefulDoFnRunner and ProcessFn
// does its own state cleanup and should never set event-time timers.
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java
index e2058bb..8b4cb24 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/WindowDoFnOperator.java
@@ -126,7 +126,7 @@
@Override
protected void fireTimer(InternalTimer<ByteBuffer, TimerData> timer) {
- timerInternals.cleanupPendingTimer(timer.getNamespace());
+ timerInternals.cleanupPendingTimer(timer.getNamespace(), true);
doFnRunner.processElement(
WindowedValue.valueInGlobalWindow(
KeyedWorkItems.timersWorkItem(
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
index 2539687..a3b40d0 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/UnboundedSourceWrapper.java
@@ -27,7 +27,7 @@
import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
import org.apache.beam.runners.flink.metrics.ReaderInvocationUtil;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
-import org.apache.beam.runners.flink.translation.utils.FlinkClassloading;
+import org.apache.beam.runners.flink.translation.utils.Workarounds;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.SerializableCoder;
@@ -358,7 +358,7 @@
}
}
} finally {
- FlinkClassloading.deleteStaticCaches();
+ Workarounds.deleteStaticCaches();
}
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java
index b0f9304..5c5ca6a 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java
@@ -67,20 +67,30 @@
static final class Timer implements BufferedElement {
private final String timerId;
+ private final String timerFamilyId;
private final BoundedWindow window;
private final Instant timestamp;
+ private final Instant outputTimestamp;
private final TimeDomain timeDomain;
- Timer(String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ Timer(
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
this.timerId = timerId;
this.window = window;
this.timestamp = timestamp;
this.timeDomain = timeDomain;
+ this.outputTimestamp = outputTimestamp;
+ this.timerFamilyId = timerFamilyId;
}
@Override
public void processWith(DoFnRunner doFnRunner) {
- doFnRunner.onTimer(timerId, window, timestamp, timeDomain);
+ doFnRunner.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
}
@Override
@@ -130,8 +140,10 @@
outStream.write(TIMER_MAGIC_BYTE);
Timer timer = (Timer) value;
STRING_CODER.encode(timer.timerId, outStream);
+ STRING_CODER.encode(timer.timerFamilyId, outStream);
windowCoder.encode(timer.window, outStream);
INSTANT_CODER.encode(timer.timestamp, outStream);
+ INSTANT_CODER.encode(timer.outputTimestamp, outStream);
outStream.write(timer.timeDomain.ordinal());
} else {
throw new IllegalStateException("Unexpected element " + value);
@@ -147,8 +159,10 @@
case TIMER_MAGIC_BYTE:
return new Timer(
STRING_CODER.decode(inStream),
+ STRING_CODER.decode(inStream),
windowCoder.decode(inStream),
INSTANT_CODER.decode(inStream),
+ INSTANT_CODER.decode(inStream),
TimeDomain.values()[inStream.read()]);
default:
throw new IllegalStateException(
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java
index 80aabc2..367ed32 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java
@@ -117,9 +117,15 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
currentBufferingElementsHandler.buffer(
- new BufferedElements.Timer(timerId, window, timestamp, timeDomain));
+ new BufferedElements.Timer(
+ timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain));
}
@Override
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkJobInvokerTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkJobInvokerTest.java
new file mode 100644
index 0000000..6397d4a
--- /dev/null
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkJobInvokerTest.java
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.flink;
+
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.core.construction.PipelineOptionsTranslation;
+import org.apache.beam.runners.core.construction.PipelineTranslation;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.io.GenerateSequence;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
+import org.apache.flink.client.program.OptimizerPlanEnvironment;
+import org.apache.flink.client.program.PackagedProgram;
+import org.apache.flink.client.program.ProgramInvocationException;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.optimizer.Optimizer;
+import org.hamcrest.CoreMatchers;
+import org.hamcrest.MatcherAssert;
+import org.hamcrest.core.StringContains;
+import org.junit.Assert;
+import org.junit.Test;
+
+/** Tests for {@link FlinkJobInvoker}. */
+public class FlinkJobInvokerTest {
+
+ @Test
+ public void testEnsureStdoutStdErrIsRestored() throws Exception {
+ PackagedProgram packagedProgram = new PackagedProgram(getClass());
+ OptimizerPlanEnvironment env = new OptimizerPlanEnvironment(new Optimizer(new Configuration()));
+ try {
+ // Flink will throw an error because no job graph will be generated by the main method
+ env.getOptimizedPlan(packagedProgram);
+ Assert.fail("This should have failed to create the Flink Plan.");
+ } catch (ProgramInvocationException e) {
+ // Test that Flink wasn't able to intercept the stdout/stderr and we printed to the regular
+ // output instead
+ MatcherAssert.assertThat(
+ e.getMessage(),
+ CoreMatchers.allOf(
+ StringContains.containsString("System.out: (none)"),
+ StringContains.containsString("System.err: (none)")));
+ }
+ }
+
+ /** Main method for {@code testEnsureStdoutStdErrIsRestored()}. */
+ public static void main(String[] args) {
+ Pipeline p = Pipeline.create();
+ p.apply(GenerateSequence.from(0));
+
+ RunnerApi.Pipeline pipeline = PipelineTranslation.toProto(p);
+ Struct options = PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create());
+
+ FlinkJobInvoker flinkJobInvoker =
+ FlinkJobInvoker.create(new FlinkJobServerDriver.FlinkServerConfiguration());
+ // This will call Workarounds.restoreOriginalStdOutAndStdErr() which we want to test
+ flinkJobInvoker.invokeWithExecutor(pipeline, options, "retrievalToken", null);
+ }
+}
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java
index c0d6ceb..48c4ed5 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkPipelineOptionsTest.java
@@ -23,6 +23,7 @@
import java.util.Collections;
import java.util.HashMap;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.SerializationUtils;
import org.apache.beam.runners.flink.translation.wrappers.streaming.DoFnOperator;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -37,7 +38,6 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.commons.lang3.SerializationUtils;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.ExecutionMode;
import org.apache.flink.api.common.typeinfo.TypeHint;
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRequiresStableInputTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRequiresStableInputTest.java
index d5be223..6908563 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRequiresStableInputTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRequiresStableInputTest.java
@@ -17,9 +17,9 @@
*/
package org.apache.beam.runners.flink;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
import static org.hamcrest.MatcherAssert.assertThat;
-import java.util.Collections;
import java.util.Date;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
@@ -31,8 +31,6 @@
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.FileChecksumMatcher;
-import org.apache.beam.sdk.testing.SerializableMatchers;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.MapElements;
@@ -167,12 +165,11 @@
waitUntilJobIsDone();
assertThat(
- new FlinkRunnerResult(Collections.emptyMap(), 1L),
- SerializableMatchers.allOf(
- new FileChecksumMatcher(
- VALUE_CHECKSUM, new FilePatternMatchingShardedFile(singleOutputPrefix + "*")),
- new FileChecksumMatcher(
- VALUE_CHECKSUM, new FilePatternMatchingShardedFile(multiOutputPrefix + "*"))));
+ new FilePatternMatchingShardedFile(singleOutputPrefix + "*"),
+ fileContentsHaveChecksum(VALUE_CHECKSUM));
+ assertThat(
+ new FilePatternMatchingShardedFile(multiOutputPrefix + "*"),
+ fileContentsHaveChecksum(VALUE_CHECKSUM));
}
private JobGraph getJobGraph(Pipeline pipeline) {
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerTest.java
new file mode 100644
index 0000000..182f65e
--- /dev/null
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkRunnerTest.java
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.flink;
+
+import static org.hamcrest.CoreMatchers.allOf;
+
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.io.GenerateSequence;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.flink.client.program.OptimizerPlanEnvironment;
+import org.apache.flink.client.program.PackagedProgram;
+import org.apache.flink.client.program.ProgramInvocationException;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.optimizer.Optimizer;
+import org.hamcrest.MatcherAssert;
+import org.hamcrest.core.StringContains;
+import org.junit.Assert;
+import org.junit.Test;
+
+/** Test for {@link FlinkRunner}. */
+public class FlinkRunnerTest {
+
+ @Test
+ public void testEnsureStdoutStdErrIsRestored() throws Exception {
+ PackagedProgram packagedProgram = new PackagedProgram(getClass());
+ OptimizerPlanEnvironment env = new OptimizerPlanEnvironment(new Optimizer(new Configuration()));
+ try {
+ // Flink will throw an error because no job graph will be generated by the main method
+ env.getOptimizedPlan(packagedProgram);
+ Assert.fail("This should have failed to create the Flink Plan.");
+ } catch (ProgramInvocationException e) {
+ // Test that Flink wasn't able to intercept the stdout/stderr and we printed to the regular
+ // output instead
+ MatcherAssert.assertThat(
+ e.getMessage(),
+ allOf(
+ StringContains.containsString("System.out: (none)"),
+ StringContains.containsString("System.err: (none)")));
+ }
+ }
+
+ /** Main method for {@code testEnsureStdoutStdErrIsRestored()}. */
+ public static void main(String[] args) {
+ FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
+ options.setRunner(NotExecutingFlinkRunner.class);
+ Pipeline p = Pipeline.create(options);
+ p.apply(GenerateSequence.from(0));
+
+ // This will call Workarounds.restoreOriginalStdOutAndStdErr() through the constructor of
+ // FlinkRunner
+ p.run();
+ }
+
+ private static class NotExecutingFlinkRunner extends FlinkRunner {
+
+ protected NotExecutingFlinkRunner(FlinkPipelineOptions options) {
+ // Stdout/Stderr is restored here
+ super(options);
+ }
+
+ @SuppressWarnings("unused")
+ public static NotExecutingFlinkRunner fromOptions(PipelineOptions options) {
+ return new NotExecutingFlinkRunner(options.as(FlinkPipelineOptions.class));
+ }
+
+ @Override
+ public PipelineResult run(Pipeline pipeline) {
+ // Do not execute to test the stdout printing
+ return null;
+ }
+ }
+}
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java
index 61d8906..fb79841 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java
@@ -45,7 +45,7 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.join.RawUnionValue;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.flink.api.common.cache.DistributedCache;
import org.apache.flink.api.common.functions.RuntimeContext;
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java
index 7fdec35..3c5f44b 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java
@@ -53,6 +53,8 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.SerializationUtils;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.mutable.MutableObject;
import org.apache.beam.runners.core.InMemoryStateInternals;
import org.apache.beam.runners.core.InMemoryTimerInternals;
import org.apache.beam.runners.core.StateNamespace;
@@ -94,15 +96,13 @@
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.sdk.v2.sdk.extensions.protobuf.ByteStringCoder;
-import org.apache.commons.lang3.SerializationUtils;
-import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.flink.api.common.cache.DistributedCache;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.typeinfo.TypeInformation;
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtilsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtilsTest.java
index 274b2bf..817b5e8 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtilsTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/FlinkKeyUtilsTest.java
@@ -26,7 +26,7 @@
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
import org.junit.Test;
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElementsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElementsTest.java
index 9ebdefc..0828a22 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElementsTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElementsTest.java
@@ -52,7 +52,12 @@
WindowedValue.of("test", new Instant(2), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING));
BufferedElement timerElement =
new BufferedElements.Timer(
- "timerId", GlobalWindow.INSTANCE, new Instant(1), TimeDomain.EVENT_TIME);
+ "timerId",
+ "timerId",
+ GlobalWindow.INSTANCE,
+ new Instant(1),
+ new Instant(1),
+ TimeDomain.EVENT_TIME);
testRoundTrip(ImmutableList.of(element), coder);
testRoundTrip(ImmutableList.of(timerElement), coder);
diff --git a/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle
index d8ddffa..990ede6 100644
--- a/runners/google-cloud-dataflow-java/build.gradle
+++ b/runners/google-cloud-dataflow-java/build.gradle
@@ -39,7 +39,7 @@
filter org.apache.tools.ant.filters.ReplaceTokens, tokens: [
'dataflow.legacy_environment_major_version' : '7',
'dataflow.fnapi_environment_major_version' : '7',
- 'dataflow.container_version' : 'beam-master-20191226'
+ 'dataflow.container_version' : 'beam-master-20200123'
]
}
@@ -81,8 +81,9 @@
compile library.java.jackson_databind
compile library.java.joda_time
compile library.java.slf4j_api
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
testCompile library.java.hamcrest_core
+ testCompile library.java.guava_testlib
testCompile library.java.junit
testCompile project(path: ":sdks:java:io:google-cloud-platform", configuration: "testRuntime")
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
@@ -217,7 +218,7 @@
def defaultDockerImageName = containerImageName(
name: "java_sdk",
root: "apachebeam",
- tag: project.version)
+ tag: project.sdk_version)
doLast {
exec {
commandLine "docker", "tag", "${defaultDockerImageName}", "${dockerImageName}"
diff --git a/runners/google-cloud-dataflow-java/examples/build.gradle b/runners/google-cloud-dataflow-java/examples/build.gradle
index 8054488..aabb0d2 100644
--- a/runners/google-cloud-dataflow-java/examples/build.gradle
+++ b/runners/google-cloud-dataflow-java/examples/build.gradle
@@ -56,7 +56,7 @@
"--tempRoot=${gcsTempRoot}",
"--runner=TestDataflowRunner",
"--dataflowWorkerJar=${dataflowWorkerJar}",
- workerHarnessContainerImage.isEmpty() ?'':"--workerHarnessContainerImage=${workerHarnessContainerImage}"
+ "--workerHarnessContainerImage=${workerHarnessContainerImage}"
] + additionalOptions
systemProperty "beamTestPipelineOptions", JsonOutput.toJson(preCommitBeamTestPipelineOptions)
}
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
index fa45b7f..0d100ad5 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
@@ -108,8 +108,8 @@
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
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.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/PrimitiveParDoSingleFactory.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/PrimitiveParDoSingleFactory.java
index 4db0e3d..eb42be2 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/PrimitiveParDoSingleFactory.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/PrimitiveParDoSingleFactory.java
@@ -18,9 +18,11 @@
package org.apache.beam.runners.dataflow;
import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
+import static org.apache.beam.runners.core.construction.ParDoTranslation.translateTimerFamilySpec;
import static org.apache.beam.runners.core.construction.ParDoTranslation.translateTimerSpec;
import static org.apache.beam.sdk.options.ExperimentalOptions.hasExperiment;
import static org.apache.beam.sdk.transforms.reflect.DoFnSignatures.getStateSpecOrThrow;
+import static org.apache.beam.sdk.transforms.reflect.DoFnSignatures.getTimerFamilySpecOrThrow;
import static org.apache.beam.sdk.transforms.reflect.DoFnSignatures.getTimerSpecOrThrow;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
@@ -242,6 +244,20 @@
}
@Override
+ public Map<String, RunnerApi.TimerFamilySpec> translateTimerFamilySpecs(
+ SdkComponents newComponents) {
+ Map<String, RunnerApi.TimerFamilySpec> timerFamilySpecs = new HashMap<>();
+ for (Map.Entry<String, DoFnSignature.TimerFamilyDeclaration> timerFamily :
+ signature.timerFamilyDeclarations().entrySet()) {
+ RunnerApi.TimerFamilySpec spec =
+ translateTimerFamilySpec(
+ getTimerFamilySpecOrThrow(timerFamily.getValue(), doFn), newComponents);
+ timerFamilySpecs.put(timerFamily.getKey(), spec);
+ }
+ return timerFamilySpecs;
+ }
+
+ @Override
public boolean isSplittable() {
return signature.processElement().isSplittable();
}
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/CloudObject.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/CloudObject.java
index e341004..8e8589d 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/CloudObject.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/CloudObject.java
@@ -22,6 +22,7 @@
import com.google.api.client.json.GenericJson;
import com.google.api.client.util.Key;
import java.util.Map;
+import java.util.Objects;
import javax.annotation.Nullable;
/**
@@ -182,4 +183,18 @@
public CloudObject clone() {
return (CloudObject) super.clone();
}
+
+ @Override
+ public boolean equals(Object otherObject) {
+ if (!(otherObject instanceof CloudObject)) {
+ return false;
+ }
+ CloudObject other = (CloudObject) otherObject;
+ return Objects.equals(className, other.className) && super.equals(otherObject);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(className, super.hashCode());
+ }
}
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/OutputReference.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/OutputReference.java
index f8b7784..7c1b9e4f 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/OutputReference.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/OutputReference.java
@@ -21,6 +21,7 @@
import com.google.api.client.json.GenericJson;
import com.google.api.client.util.Key;
+import java.util.Objects;
/**
* A representation used by {@link com.google.api.services.dataflow.model.Step}s to reference the
@@ -40,4 +41,21 @@
this.stepName = checkNotNull(stepName);
this.outputName = checkNotNull(outputName);
}
+
+ @Override
+ public boolean equals(Object otherObject) {
+ if (!(otherObject instanceof OutputReference)) {
+ return false;
+ }
+ OutputReference other = (OutputReference) otherObject;
+ return Objects.equals(type, other.type)
+ && Objects.equals(stepName, other.stepName)
+ && Objects.equals(outputName, other.outputName)
+ && super.equals(other);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(type, stepName, outputName, super.hashCode());
+ }
}
diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/CloudObjectTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/CloudObjectTest.java
new file mode 100644
index 0000000..118bab8
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/CloudObjectTest.java
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.dataflow.util;
+
+import com.google.common.testing.EqualsTester;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class CloudObjectTest {
+
+ @Test
+ public void testEquality() {
+ new EqualsTester()
+ .addEqualityGroup(CloudObject.forFloat(1.0), CloudObject.forFloat(1.0))
+ .addEqualityGroup(CloudObject.forInteger(3L), CloudObject.forInteger(3L))
+ .addEqualityGroup(CloudObject.forFloat(3.0))
+ .addEqualityGroup(CloudObject.forString("foo"), CloudObject.forString("foo"))
+ .addEqualityGroup(CloudObject.forClassName("foo.Bar"), CloudObject.forClassName("foo.Bar"))
+ .addEqualityGroup(
+ CloudObject.fromSpec(ImmutableMap.of(PropertyNames.OBJECT_TYPE_NAME, "ValuesDoFn")),
+ CloudObject.fromSpec(ImmutableMap.of(PropertyNames.OBJECT_TYPE_NAME, "ValuesDoFn")))
+ .testEquals();
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/OutputReferenceTest.java
similarity index 61%
copy from runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
copy to runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/OutputReferenceTest.java
index a114f40..f0817b5 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
+++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/OutputReferenceTest.java
@@ -15,16 +15,21 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.runners.flink.translation.utils;
+package org.apache.beam.runners.dataflow.util;
-import com.fasterxml.jackson.databind.type.TypeFactory;
+import com.google.common.testing.EqualsTester;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
-/** Utilities for dealing with classloading. */
-public class FlinkClassloading {
+@RunWith(JUnit4.class)
+public class OutputReferenceTest {
- public static void deleteStaticCaches() {
- // Clear cache to get rid of any references to the Flink Classloader
- // See https://jira.apache.org/jira/browse/BEAM-6460
- TypeFactory.defaultInstance().clearCache();
+ @Test
+ public void testEquality() {
+ new EqualsTester()
+ .addEqualityGroup(new OutputReference("sA", "oA"), new OutputReference("sA", "oA"))
+ .addEqualityGroup(new OutputReference("sB", "oB"), new OutputReference("sB", "oB"))
+ .testEquals();
}
}
diff --git a/runners/google-cloud-dataflow-java/worker/build.gradle b/runners/google-cloud-dataflow-java/worker/build.gradle
index 6b866d7..a0e3356 100644
--- a/runners/google-cloud-dataflow-java/worker/build.gradle
+++ b/runners/google-cloud-dataflow-java/worker/build.gradle
@@ -76,7 +76,7 @@
compile project(":runners:java-fn-execution")
compile project(":sdks:java:fn-execution")
compile project(path: ":runners:google-cloud-dataflow-java:worker:windmill", configuration: "shadow")
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
compile google_api_services_dataflow
compile library.java.avro
compile library.java.google_api_client
diff --git a/runners/google-cloud-dataflow-java/worker/legacy-worker/build.gradle b/runners/google-cloud-dataflow-java/worker/legacy-worker/build.gradle
index 5ee2d63..d2912e4 100644
--- a/runners/google-cloud-dataflow-java/worker/legacy-worker/build.gradle
+++ b/runners/google-cloud-dataflow-java/worker/legacy-worker/build.gradle
@@ -53,7 +53,7 @@
library.java.jackson_databind,
library.java.joda_time,
library.java.slf4j_api,
- library.java.vendored_grpc_1_21_0,
+ library.java.vendored_grpc_1_26_0,
]
def sdk_provided_shaded_project_dependencies = [
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java
index a5fb660..5cc0db7 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java
@@ -21,7 +21,7 @@
import com.google.api.services.dataflow.model.WorkItem;
import java.io.Closeable;
import java.io.IOException;
-import java.util.List;
+import java.util.Collection;
import java.util.function.Function;
import javax.annotation.Nullable;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort;
@@ -46,6 +46,7 @@
import org.apache.beam.runners.dataflow.worker.graph.RegisterNodeFunction;
import org.apache.beam.runners.dataflow.worker.graph.ReplacePgbkWithPrecombineFunction;
import org.apache.beam.runners.dataflow.worker.status.DebugCapture;
+import org.apache.beam.runners.dataflow.worker.status.DebugCapture.Capturable;
import org.apache.beam.runners.dataflow.worker.status.WorkerStatusPages;
import org.apache.beam.runners.dataflow.worker.util.MemoryMonitor;
import org.apache.beam.sdk.fn.IdGenerator;
@@ -265,7 +266,7 @@
}
private static DebugCapture.Manager initializeAndStartDebugCaptureManager(
- DataflowWorkerHarnessOptions options, List<DebugCapture.Capturable> debugCapturePages) {
+ DataflowWorkerHarnessOptions options, Collection<Capturable> debugCapturePages) {
DebugCapture.Manager result = new DebugCapture.Manager(options, debugCapturePages);
result.start();
return result;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ByteStringCoder.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ByteStringCoder.java
index 3d9a2c4..eab920f 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ByteStringCoder.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ByteStringCoder.java
@@ -23,7 +23,7 @@
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.util.VarInt;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.ByteStreams;
/**
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContext.java
index 18ce7c5..166f5f0 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContext.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContext.java
@@ -228,7 +228,7 @@
protected String getLullMessage(Thread trackedThread, Duration millis) {
StringBuilder message = new StringBuilder();
- message.append("Processing stuck");
+ message.append("Operation ongoing");
if (getStepName() != null) {
message.append(" in step ").append(getStepName().userName());
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java
index 2903cc0..3cb92c9 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java
@@ -109,7 +109,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
throw new UnsupportedOperationException("Unsupported for ProcessFn");
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowRunnerHarness.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowRunnerHarness.java
index eda6b03..e3de7cb 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowRunnerHarness.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowRunnerHarness.java
@@ -33,12 +33,14 @@
import org.apache.beam.runners.dataflow.worker.fn.logging.BeamFnLoggingService;
import org.apache.beam.runners.dataflow.worker.fn.stream.ServerStreamObserverFactory;
import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingInitializer;
+import org.apache.beam.runners.dataflow.worker.status.SdkWorkerStatusServlet;
import org.apache.beam.runners.fnexecution.GrpcContextHeaderAccessorProvider;
import org.apache.beam.runners.fnexecution.ServerFactory;
import org.apache.beam.runners.fnexecution.control.FnApiControlClient;
import org.apache.beam.runners.fnexecution.state.GrpcStateService;
+import org.apache.beam.runners.fnexecution.status.BeamWorkerStatusGrpcService;
import org.apache.beam.sdk.io.FileSystems;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -62,6 +64,7 @@
// critical traffic protected from best effort traffic.
ApiServiceDescriptor controlApiService = DataflowWorkerHarnessHelper.getControlDescriptor();
ApiServiceDescriptor loggingApiService = DataflowWorkerHarnessHelper.getLoggingDescriptor();
+ ApiServiceDescriptor statusApiService = DataflowWorkerHarnessHelper.getStatusDescriptor();
LOG.info(
"{} started, using port {} for control, {} for logging.",
@@ -93,6 +96,7 @@
Server servicesServer = null;
Server loggingServer = null;
+ Server statusServer = null;
try (BeamFnLoggingService beamFnLoggingService =
new BeamFnLoggingService(
loggingApiService,
@@ -110,6 +114,11 @@
controlApiService,
streamObserverFactory::from,
GrpcContextHeaderAccessorProvider.getHeaderAccessor());
+ BeamWorkerStatusGrpcService beamWorkerStatusGrpcService =
+ statusApiService == null
+ ? null
+ : BeamWorkerStatusGrpcService.create(
+ statusApiService, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
GrpcStateService beamFnStateService = GrpcStateService.create()) {
servicesServer =
@@ -120,22 +129,41 @@
loggingServer =
serverFactory.create(ImmutableList.of(beamFnLoggingService), loggingApiService);
+ // gRPC server for obtaining SDK harness runtime status information.
+ if (beamWorkerStatusGrpcService != null) {
+ statusServer =
+ serverFactory.create(ImmutableList.of(beamWorkerStatusGrpcService), statusApiService);
+ }
+
start(
pipeline,
pipelineOptions,
beamFnControlService,
beamFnDataService,
controlApiService,
- beamFnStateService);
+ beamFnStateService,
+ beamWorkerStatusGrpcService);
+
+ if (statusServer != null) {
+ statusServer.shutdown();
+ }
servicesServer.shutdown();
loggingServer.shutdown();
+
+ // wait 30 secs for outstanding requests to finish.
+ if (statusServer != null) {
+ statusServer.awaitTermination(30, TimeUnit.SECONDS);
+ }
+ servicesServer.awaitTermination(30, TimeUnit.SECONDS);
+ loggingServer.awaitTermination(30, TimeUnit.SECONDS);
} finally {
- if (servicesServer != null) {
- servicesServer.awaitTermination(30, TimeUnit.SECONDS);
+ if (statusServer != null && !statusServer.isTerminated()) {
+ statusServer.shutdownNow();
+ }
+ if (servicesServer != null && !servicesServer.isTerminated()) {
servicesServer.shutdownNow();
}
- if (loggingServer != null) {
- loggingServer.awaitTermination(30, TimeUnit.SECONDS);
+ if (loggingServer != null && !loggingServer.isTerminated()) {
loggingServer.shutdownNow();
}
}
@@ -148,7 +176,8 @@
BeamFnControlService beamFnControlService,
BeamFnDataGrpcService beamFnDataService,
ApiServiceDescriptor stateApiServiceDescriptor,
- GrpcStateService beamFnStateService)
+ GrpcStateService beamFnStateService,
+ BeamWorkerStatusGrpcService beamWorkerStatusGrpcService)
throws Exception {
SdkHarnessRegistry sdkHarnessRegistry =
@@ -161,6 +190,12 @@
StreamingDataflowWorker worker =
StreamingDataflowWorker.forStreamingFnWorkerHarness(
Collections.emptyList(), client, pipelineOptions, pipeline, sdkHarnessRegistry);
+ // Add SDK status servlet and capture page only if Fn worker status server is started.
+ if (beamWorkerStatusGrpcService != null) {
+ SdkWorkerStatusServlet sdkWorkerStatusServlet =
+ new SdkWorkerStatusServlet(beamWorkerStatusGrpcService);
+ worker.addWorkerStatusPage(sdkWorkerStatusServlet);
+ }
worker.startStatusPages();
worker.start();
ExecutorService executor = Executors.newSingleThreadExecutor();
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelper.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelper.java
index 5449462..36454db 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelper.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelper.java
@@ -31,7 +31,7 @@
import org.apache.beam.runners.dataflow.worker.ExperimentContext.Experiment;
import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingInitializer;
import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
import org.conscrypt.OpenSSLProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java
index e50f1bd..6296461 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java
@@ -82,7 +82,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
throw new UnsupportedOperationException(
String.format("Timers are not supported by %s", GroupAlsoByWindowFn.class.getSimpleName()));
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowParDoFnFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowParDoFnFactory.java
index b12f889..bef1466 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowParDoFnFactory.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowParDoFnFactory.java
@@ -58,7 +58,7 @@
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/MetricTrackingWindmillServerStub.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/MetricTrackingWindmillServerStub.java
index 734f49e..f08bc80 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/MetricTrackingWindmillServerStub.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/MetricTrackingWindmillServerStub.java
@@ -30,7 +30,7 @@
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetDataStream;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.SettableFuture;
import org.joda.time.Duration;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java
index 147fd76..4cc73a1 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java
@@ -36,7 +36,7 @@
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
/**
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReaderCache.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReaderCache.java
index 6b00560..8b4439d 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReaderCache.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReaderCache.java
@@ -22,7 +22,7 @@
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java
index 712a017..4c011ed 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java
@@ -357,7 +357,13 @@
private void processUserTimer(TimerData timer) throws Exception {
if (fnSignature.timerDeclarations().containsKey(timer.getTimerId())) {
BoundedWindow window = ((WindowNamespace) timer.getNamespace()).getWindow();
- fnRunner.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ fnRunner.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
}
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StateFetcher.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StateFetcher.java
index 3c804db..1dbbd6c 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StateFetcher.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StateFetcher.java
@@ -38,7 +38,7 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
index d6d017a..f2b0b27 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
@@ -36,6 +36,7 @@
import java.io.PrintWriter;
import java.util.ArrayDeque;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.Collections;
import java.util.Deque;
import java.util.HashMap;
@@ -128,8 +129,8 @@
import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
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.Optional;
@@ -818,7 +819,7 @@
new TimerTask() {
@Override
public void run() {
- List<Capturable> pages = statusPages.getDebugCapturePages();
+ Collection<Capturable> pages = statusPages.getDebugCapturePages();
if (pages.isEmpty()) {
LOG.warn("No captured status pages.");
}
@@ -871,6 +872,13 @@
statusPages.start();
}
+ public void addWorkerStatusPage(BaseStatusServlet page) {
+ statusPages.addServlet(page);
+ if (page instanceof Capturable) {
+ statusPages.addCapturePage((Capturable) page);
+ }
+ }
+
public void stop() {
try {
if (globalConfigRefreshTimer != null) {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java
index bf5efee..41da4b3 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java
@@ -133,7 +133,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
throw new UnsupportedOperationException(
"Attempt to deliver a timer to a DoFn, but timers are not supported in Dataflow.");
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
index 76aa8b0..a761ae4 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
@@ -54,7 +54,7 @@
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java
index 02a1292..f7e86c7 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java
@@ -77,7 +77,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
throw new UnsupportedOperationException(
"Attempt to deliver a timer to a DoFn, but timers are not supported in Dataflow.");
}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java
index 2c00c99..eb95afb 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcher.java
@@ -48,8 +48,8 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Parser;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Parser;
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.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillNamespacePrefix.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillNamespacePrefix.java
index eba5c5d..2dd7006 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillNamespacePrefix.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillNamespacePrefix.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.dataflow.worker;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
/**
* A prefix for a Windmill state or timer tag to separate user state and timers from system state
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java
index 60ddce5..cd35038 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java
@@ -40,7 +40,7 @@
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.ValueWithRecordId;
import org.apache.beam.sdk.values.ValueWithRecordId.ValueWithRecordIdCoder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
class WindmillSink<T> extends Sink<WindowedValue<T>> {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateCache.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateCache.java
index b419a38..eb18ef8 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateCache.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateCache.java
@@ -32,7 +32,7 @@
import org.apache.beam.runners.dataflow.worker.status.StatusDataProvider;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.util.Weighted;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Equivalence;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
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 9c9779e..a3619ab 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
@@ -55,7 +55,7 @@
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.util.CombineFnUtil;
import org.apache.beam.sdk.util.Weighted;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateReader.java
index 0050602..75ee1cb 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateReader.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillStateReader.java
@@ -40,7 +40,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.Weighted;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.Function;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Objects;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java
index b2ba62e..5d92d2a 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java
@@ -29,7 +29,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.collect.HashBasedTable;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Table;
@@ -88,8 +88,14 @@
@Override
public void setTimer(TimerData timerKey) {
- timers.put(timerKey.getTimerId(), timerKey.getNamespace(), timerKey);
- timerStillPresent.put(timerKey.getTimerId(), timerKey.getNamespace(), true);
+ timers.put(
+ getTimerDataKey(timerKey.getTimerId(), timerKey.getTimerFamilyId()),
+ timerKey.getNamespace(),
+ timerKey);
+ timerStillPresent.put(
+ getTimerDataKey(timerKey.getTimerId(), timerKey.getTimerFamilyId()),
+ timerKey.getNamespace(),
+ true);
}
@Override
@@ -101,20 +107,31 @@
Instant outputTimestamp,
TimeDomain timeDomain) {
timers.put(
- timerId,
+ getTimerDataKey(timerId, timerFamilyId),
namespace,
TimerData.of(timerId, timerFamilyId, namespace, timestamp, outputTimestamp, timeDomain));
- timerStillPresent.put(timerId, namespace, true);
+ timerStillPresent.put(getTimerDataKey(timerId, timerFamilyId), namespace, true);
+ }
+
+ private String getTimerDataKey(String timerId, String timerFamilyId) {
+ // Identifies timer uniquely with timerFamilyId
+ return timerId + '+' + timerFamilyId;
}
@Override
public void deleteTimer(TimerData timerKey) {
- timers.put(timerKey.getTimerId(), timerKey.getNamespace(), timerKey);
- timerStillPresent.put(timerKey.getTimerId(), timerKey.getNamespace(), false);
+ timers.put(
+ getTimerDataKey(timerKey.getTimerId(), timerKey.getTimerFamilyId()),
+ timerKey.getNamespace(),
+ timerKey);
+ timerStillPresent.put(
+ getTimerDataKey(timerKey.getTimerId(), timerKey.getTimerFamilyId()),
+ timerKey.getNamespace(),
+ false);
}
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
throw new UnsupportedOperationException("Canceling a timer by ID is not yet supported.");
}
@@ -184,7 +201,7 @@
.setStateFamily(stateFamily)
.setReset(true)
.addTimestamps(
- WindmillTimeUtils.harnessToWindmillTimestamp(timerData.getTimestamp()));
+ WindmillTimeUtils.harnessToWindmillTimestamp(timerData.getOutputTimestamp()));
}
} else {
// Deleting a timer. If it is a user timer, clear the hold
@@ -246,7 +263,7 @@
// The tag is a path-structure string but cheaper to parse than a proper URI. It follows
// this pattern, where no component but the ID can contain a slash
//
- // prefix namespace '+' id
+ // prefix namespace '+' id '+' familyId
//
// prefix ::= '/' prefix_char
// namespace ::= '/' | '/' window '/'
@@ -269,13 +286,19 @@
prefix.byteString());
int namespaceStart = prefix.byteString().size(); // drop the prefix, leave the begin slash
int namespaceEnd = tag.indexOf('+', namespaceStart); // keep the end slash, drop the +
-
String namespaceString = tag.substring(namespaceStart, namespaceEnd);
- String id = tag.substring(namespaceEnd + 1);
+ String timerIdPlusTimerFamilyId = tag.substring(namespaceEnd + 1); // timerId+timerFamilyId
+ int timerIdEnd = timerIdPlusTimerFamilyId.indexOf('+'); // end of timerId
+ // if no '+' found then timerFamilyId is empty string else they have a '+' separator
+ String familyId = timerIdEnd == -1 ? "" : timerIdPlusTimerFamilyId.substring(timerIdEnd + 1);
+ String id =
+ timerIdEnd == -1
+ ? timerIdPlusTimerFamilyId
+ : timerIdPlusTimerFamilyId.substring(0, timerIdEnd);
StateNamespace namespace = StateNamespaces.fromString(namespaceString, windowCoder);
Instant timestamp = WindmillTimeUtils.windmillToHarnessTimestamp(timer.getTimestamp());
- return TimerData.of(id, namespace, timestamp, timerTypeToTimeDomain(timer.getType()));
+ return TimerData.of(id, familyId, namespace, timestamp, timerTypeToTimeDomain(timer.getType()));
}
/**
@@ -285,13 +308,27 @@
* <p>This is necessary because Windmill will deduplicate based only on this tag.
*/
public static ByteString timerTag(WindmillNamespacePrefix prefix, TimerData timerData) {
- String tagString =
- new StringBuilder()
- .append(prefix.byteString().toStringUtf8()) // this never ends with a slash
- .append(timerData.getNamespace().stringKey()) // this must begin and end with a slash
- .append('+')
- .append(timerData.getTimerId()) // this is arbitrary; currently unescaped
- .toString();
+ String tagString;
+ // Timers without timerFamily would have timerFamily would be an empty string
+ if ("".equals(timerData.getTimerFamilyId())) {
+ tagString =
+ new StringBuilder()
+ .append(prefix.byteString().toStringUtf8()) // this never ends with a slash
+ .append(timerData.getNamespace().stringKey()) // this must begin and end with a slash
+ .append('+')
+ .append(timerData.getTimerId()) // this is arbitrary; currently unescaped
+ .toString();
+ } else {
+ tagString =
+ new StringBuilder()
+ .append(prefix.byteString().toStringUtf8()) // this never ends with a slash
+ .append(timerData.getNamespace().stringKey()) // this must begin and end with a slash
+ .append('+')
+ .append(timerData.getTimerId()) // this is arbitrary; currently unescaped
+ .append('+')
+ .append(timerData.getTimerFamilyId())
+ .toString();
+ }
return ByteString.copyFromUtf8(tagString);
}
@@ -300,14 +337,30 @@
* hold that is only freed after the timer fires.
*/
public static ByteString timerHoldTag(WindmillNamespacePrefix prefix, TimerData timerData) {
- String tagString =
- new StringBuilder()
- .append(prefix.byteString().toStringUtf8()) // this never ends with a slash
- .append(TIMER_HOLD_PREFIX) // this never ends with a slash
- .append(timerData.getNamespace().stringKey()) // this must begin and end with a slash
- .append('+')
- .append(timerData.getTimerId()) // this is arbitrary; currently unescaped
- .toString();
+ String tagString;
+ if ("".equals(timerData.getTimerFamilyId())) {
+ tagString =
+ new StringBuilder()
+ .append(prefix.byteString().toStringUtf8()) // this never ends with a slash
+ .append(TIMER_HOLD_PREFIX) // this never ends with a slash
+ .append(timerData.getNamespace().stringKey()) // this must begin and end with a slash
+ .append('+')
+ .append(timerData.getTimerId()) // this is arbitrary; currently unescaped
+ .toString();
+ } else {
+ tagString =
+ new StringBuilder()
+ .append(prefix.byteString().toStringUtf8()) // this never ends with a slash
+ .append(TIMER_HOLD_PREFIX) // this never ends with a slash
+ .append(timerData.getNamespace().stringKey()) // this must begin and end with a slash
+ .append('+')
+ .append(timerData.getTimerId()) // this is arbitrary; currently unescaped
+ .append('+')
+ .append(
+ timerData.getTimerFamilyId()) // use to differentiate same timerId in different
+ // timerMap
+ .toString();
+ }
return ByteString.copyFromUtf8(tagString);
}
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 4964a89..344798e 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
@@ -61,7 +61,7 @@
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.ValueWithRecordId;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlService.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlService.java
index d701083..865d033 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlService.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlService.java
@@ -26,7 +26,7 @@
import org.apache.beam.runners.dataflow.worker.fn.grpc.BeamFnService;
import org.apache.beam.runners.fnexecution.HeaderAccessor;
import org.apache.beam.runners.fnexecution.control.FnApiControlClient;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperation.java
index bf42c4d..2879c18 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperation.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperation.java
@@ -73,8 +73,8 @@
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.MoreFutures;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
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.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiver.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiver.java
index e3d2277..89b2b7d 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiver.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiver.java
@@ -120,7 +120,7 @@
TimerInternals timerInternals = stepContext.namespacedToUser().timerInternals();
timerInternals.setTimer(
- namespace, timerId, "", timer.getTimestamp(), timer.getOutputTimestamp(), timeDomain);
+ namespace, timerId, "", timer.getTimestamp(), windowedValue.getTimestamp(), timeDomain);
timerIdToKey.put(timerId, windowedValue.getValue().getKey());
timerIdToPayload.put(timerId, timer.getPayload());
@@ -144,7 +144,7 @@
KV.of(
timerIdToKey.get(timerData.getTimerId()),
Timer.of(timerData.getTimestamp(), timerIdToPayload.get(timerData.getTimerId()))),
- timerData.getTimestamp(),
+ timerData.getOutputTimestamp(),
Collections.singleton(window),
PaneInfo.NO_FIRING);
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcService.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcService.java
index dcde104..cad6a8e 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcService.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcService.java
@@ -40,7 +40,7 @@
import org.apache.beam.sdk.fn.data.LogicalEndpoint;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingService.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingService.java
index d1b62d1..2045317 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingService.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingService.java
@@ -28,8 +28,8 @@
import org.apache.beam.runners.dataflow.worker.fn.grpc.BeamFnService;
import org.apache.beam.runners.dataflow.worker.logging.DataflowWorkerLoggingMDC;
import org.apache.beam.runners.fnexecution.HeaderAccessor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactory.java
index 1fabd6d..e51c0ee 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactory.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactory.java
@@ -27,9 +27,9 @@
import org.apache.beam.sdk.fn.stream.ForwardingClientResponseObserver;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ServerCallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ServerCallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* A {@link StreamObserver} factory that wraps provided {@link CallStreamObserver}s making them flow
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/CreateExecutableStageNodeFunction.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/CreateExecutableStageNodeFunction.java
index 9b75433..1b888c6 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/CreateExecutableStageNodeFunction.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/CreateExecutableStageNodeFunction.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.dataflow.worker.graph;
-import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTING;
+import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTINGS;
import static org.apache.beam.runners.dataflow.util.Structs.getBytes;
import static org.apache.beam.runners.dataflow.util.Structs.getString;
import static org.apache.beam.runners.dataflow.worker.graph.LengthPrefixUnknownCoders.forSideInputInfos;
@@ -45,7 +45,6 @@
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Environment;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.UserStateId;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms;
import org.apache.beam.runners.core.construction.*;
import org.apache.beam.runners.core.construction.graph.ExecutableStage;
import org.apache.beam.runners.core.construction.graph.ImmutableExecutableStage;
@@ -79,8 +78,8 @@
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
@@ -101,15 +100,15 @@
private static final String JAVA_SOURCE_URN = "beam:source:java:0.1";
public static final String COMBINE_PER_KEY_URN =
- BeamUrns.getUrn(StandardPTransforms.Composites.COMBINE_PER_KEY);
+ PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN;
public static final String COMBINE_PRECOMBINE_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_PRECOMBINE);
+ PTransformTranslation.COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN;
public static final String COMBINE_MERGE_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_MERGE_ACCUMULATORS);
+ PTransformTranslation.COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN;
public static final String COMBINE_EXTRACT_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_EXTRACT_OUTPUTS);
+ PTransformTranslation.COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN;
public static final String COMBINE_GROUPED_VALUES_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_GROUPED_VALUES);
+ PTransformTranslation.COMBINE_GROUPED_VALUES_TRANSFORM_URN;
private static final String SERIALIZED_SOURCE = "serialized_source";
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
@@ -487,7 +486,7 @@
executableStageTimers,
executableStageTransforms,
executableStageOutputs,
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
return ExecutableStageNode.create(
executableStage,
ptransformIdToNameContexts.build(),
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodes.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodes.java
index 0d96981..83ed105 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodes.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodes.java
@@ -37,7 +37,7 @@
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/RegisterNodeFunction.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/RegisterNodeFunction.java
index b44e42c..e1c3614 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/RegisterNodeFunction.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/graph/RegisterNodeFunction.java
@@ -43,9 +43,7 @@
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.SideInput;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms;
import org.apache.beam.runners.core.SideInputReader;
-import org.apache.beam.runners.core.construction.BeamUrns;
import org.apache.beam.runners.core.construction.CoderTranslation;
import org.apache.beam.runners.core.construction.Environments;
import org.apache.beam.runners.core.construction.PTransformTranslation;
@@ -80,8 +78,8 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
@@ -102,15 +100,15 @@
private static final String JAVA_SOURCE_URN = "beam:source:java:0.1";
public static final String COMBINE_PER_KEY_URN =
- BeamUrns.getUrn(StandardPTransforms.Composites.COMBINE_PER_KEY);
+ PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN;
public static final String COMBINE_PRECOMBINE_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_PRECOMBINE);
+ PTransformTranslation.COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN;
public static final String COMBINE_MERGE_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_MERGE_ACCUMULATORS);
+ PTransformTranslation.COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN;
public static final String COMBINE_EXTRACT_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_EXTRACT_OUTPUTS);
+ PTransformTranslation.COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN;
public static final String COMBINE_GROUPED_VALUES_URN =
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_GROUPED_VALUES);
+ PTransformTranslation.COMBINE_GROUPED_VALUES_TRANSFORM_URN;
private static final String SERIALIZED_SOURCE = "serialized_source";
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/DebugCapture.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/DebugCapture.java
index f02ddda..c6f873f 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/DebugCapture.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/DebugCapture.java
@@ -30,6 +30,7 @@
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
@@ -99,13 +100,13 @@
private String project, job, host, region;
private Dataflow client = null;
private ScheduledExecutorService executor = null;
- private List<Capturable> capturables;
+ private Collection<Capturable> capturables;
private boolean enabled;
private long lastCaptureUsec = 0;
@VisibleForTesting Config captureConfig = new Config();
- public Manager(DataflowWorkerHarnessOptions options, List<Capturable> capturables) {
+ public Manager(DataflowWorkerHarnessOptions options, Collection<Capturable> capturables) {
try {
client = options.getDataflowClient();
} catch (Exception e) {
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/SdkWorkerStatusServlet.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/SdkWorkerStatusServlet.java
new file mode 100644
index 0000000..b304952
--- /dev/null
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/SdkWorkerStatusServlet.java
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.dataflow.worker.status;
+
+import com.google.common.base.Strings;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import javax.servlet.ServletException;
+import javax.servlet.ServletOutputStream;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import org.apache.beam.runners.dataflow.worker.status.DebugCapture.Capturable;
+import org.apache.beam.runners.fnexecution.status.BeamWorkerStatusGrpcService;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Servlet dedicated to provide live status info retrieved from SDK Harness. Note this is different
+ * from {@link WorkerStatusPages} which incorporates all info for Dataflow runner including this
+ * SDKWorkerStatus page.
+ */
+public class SdkWorkerStatusServlet extends BaseStatusServlet implements Capturable {
+
+ private static final Logger LOG = LoggerFactory.getLogger(SdkWorkerStatusServlet.class);
+ private final transient BeamWorkerStatusGrpcService statusGrpcService;
+
+ public SdkWorkerStatusServlet(BeamWorkerStatusGrpcService statusGrpcService) {
+ super("sdk_status");
+ this.statusGrpcService = statusGrpcService;
+ }
+
+ @Override
+ protected void doGet(HttpServletRequest request, HttpServletResponse response)
+ throws IOException, ServletException {
+ String id = request.getParameter("id");
+ if (Strings.isNullOrEmpty(id)) {
+ // return all connected sdk statuses if no id provided.
+ response.setContentType("text/html;charset=utf-8");
+ ServletOutputStream writer = response.getOutputStream();
+ try (PrintWriter out =
+ new PrintWriter(new OutputStreamWriter(writer, StandardCharsets.UTF_8))) {
+ captureData(out);
+ }
+ } else {
+ response.setContentType("text/plain;charset=utf-8");
+ ServletOutputStream writer = response.getOutputStream();
+ writer.println(statusGrpcService.getSingleWorkerStatus(id, 10, TimeUnit.SECONDS));
+ }
+ response.setStatus(HttpServletResponse.SC_OK);
+ response.flushBuffer();
+ }
+
+ @Override
+ public String pageName() {
+ return "/sdk_status";
+ }
+
+ @Override
+ public void captureData(PrintWriter writer) {
+ Map<String, String> allStatuses = statusGrpcService.getAllWorkerStatuses(10, TimeUnit.SECONDS);
+
+ writer.println("<html>");
+ writer.println("<h1>SDK harness</h1>");
+ // add links to each sdk section for easier navigation.
+ for (String sdkId : allStatuses.keySet()) {
+ writer.print(String.format("<a href=\"#%s\">%s</a> ", sdkId, sdkId));
+ }
+ writer.println();
+
+ for (Map.Entry<String, String> entry : allStatuses.entrySet()) {
+ writer.println(String.format("<h2 id=\"%s\">%s</h2>", entry.getKey(), entry.getKey()));
+ writer.println("<a href=\"#top\">return to top</a>");
+ writer.println("<div style=\"white-space:pre-wrap\">");
+ writer.println(entry.getValue());
+ writer.println("</div>");
+ writer.println("");
+ }
+ writer.println("</html>");
+ }
+}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java
index cd3b9de..764a607 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPages.java
@@ -18,7 +18,8 @@
package org.apache.beam.runners.dataflow.worker.status;
import java.io.IOException;
-import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
import java.util.List;
import java.util.function.BooleanSupplier;
import javax.servlet.ServletException;
@@ -39,6 +40,7 @@
private static final Logger LOG = LoggerFactory.getLogger(WorkerStatusPages.class);
private final Server statusServer;
+ private final List<Capturable> capturePages;
private final StatuszServlet statuszServlet = new StatuszServlet();
private final ThreadzServlet threadzServlet = new ThreadzServlet();
private final ServletHandler servletHandler = new ServletHandler();
@@ -46,6 +48,7 @@
@VisibleForTesting
WorkerStatusPages(Server server, MemoryMonitor memoryMonitor, BooleanSupplier healthyIndicator) {
this.statusServer = server;
+ this.capturePages = new ArrayList<>();
this.statusServer.setHandler(servletHandler);
// Install the default servlets (threadz, healthz, heapz, statusz)
@@ -54,6 +57,9 @@
addServlet(new HeapzServlet(memoryMonitor));
addServlet(statuszServlet);
+ // Add default capture pages (threadz, statusz)
+ this.capturePages.add(threadzServlet);
+ this.capturePages.add(statuszServlet);
// Add some status pages
addStatusDataProvider("resources", "Resources", memoryMonitor);
}
@@ -107,8 +113,12 @@
}
/** Returns the set of pages than should be captured by DebugCapture. */
- public List<Capturable> getDebugCapturePages() {
- return Arrays.asList(threadzServlet, statuszServlet);
+ public Collection<Capturable> getDebugCapturePages() {
+ return this.capturePages;
+ }
+
+ public void addCapturePage(Capturable page) {
+ this.capturePages.add(page);
}
/** Redirect all invalid pages to /statusz. */
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/DirectStreamObserver.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/DirectStreamObserver.java
index a54733d..7565ba2 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/DirectStreamObserver.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/DirectStreamObserver.java
@@ -19,8 +19,8 @@
import java.util.concurrent.Phaser;
import javax.annotation.concurrent.ThreadSafe;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* A {@link StreamObserver} which uses synchronization on the underlying {@link CallStreamObserver}
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java
index 74d8e4d..d7eba1f 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/ForwardingClientResponseObserver.java
@@ -17,9 +17,9 @@
*/
package org.apache.beam.runners.dataflow.worker.windmill;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientCallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientResponseObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientCallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientResponseObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* A {@link ClientResponseObserver} which delegates all {@link StreamObserver} calls.
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServer.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServer.java
index c64803d..632ab07 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServer.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServer.java
@@ -84,17 +84,17 @@
import org.apache.beam.sdk.util.BackOffUtils;
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.util.Sleeper;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.CallCredentials;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Channel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.auth.MoreCallCredentials;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.netty.GrpcSslContexts;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.netty.NegotiationType;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.netty.NettyChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.CallCredentials;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Channel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.auth.MoreCallCredentials;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.netty.GrpcSslContexts;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.netty.NegotiationType;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.netty.NettyChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
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.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Splitter;
@@ -232,11 +232,11 @@
*/
private static class VendoredRequestMetadataCallbackAdapter
implements com.google.auth.RequestMetadataCallback {
- private final org.apache.beam.vendor.grpc.v1p21p0.com.google.auth.RequestMetadataCallback
+ private final org.apache.beam.vendor.grpc.v1p26p0.com.google.auth.RequestMetadataCallback
callback;
private VendoredRequestMetadataCallbackAdapter(
- org.apache.beam.vendor.grpc.v1p21p0.com.google.auth.RequestMetadataCallback callback) {
+ org.apache.beam.vendor.grpc.v1p26p0.com.google.auth.RequestMetadataCallback callback) {
this.callback = callback;
}
@@ -260,7 +260,7 @@
* delegate to reduce maintenance burden.
*/
private static class VendoredCredentialsAdapter
- extends org.apache.beam.vendor.grpc.v1p21p0.com.google.auth.Credentials {
+ extends org.apache.beam.vendor.grpc.v1p26p0.com.google.auth.Credentials {
private final com.google.auth.Credentials credentials;
private VendoredCredentialsAdapter(com.google.auth.Credentials credentials) {
@@ -281,7 +281,7 @@
public void getRequestMetadata(
final URI uri,
Executor executor,
- final org.apache.beam.vendor.grpc.v1p21p0.com.google.auth.RequestMetadataCallback
+ final org.apache.beam.vendor.grpc.v1p26p0.com.google.auth.RequestMetadataCallback
callback) {
credentials.getRequestMetadata(
uri, executor, new VendoredRequestMetadataCallbackAdapter(callback));
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/StreamObserverFactory.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/StreamObserverFactory.java
index 6731951..0216766 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/StreamObserverFactory.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/StreamObserverFactory.java
@@ -20,8 +20,8 @@
import java.util.function.Function;
import org.apache.beam.sdk.fn.stream.AdvancingPhaser;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* Uses {@link PipelineOptions} to configure which underlying {@link StreamObserver} implementation
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/harness/test/TestStreams.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/harness/test/TestStreams.java
index 07ccdb1..e88054c 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/harness/test/TestStreams.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/harness/test/TestStreams.java
@@ -19,8 +19,8 @@
import java.util.function.Consumer;
import java.util.function.Supplier;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/** Utility methods which enable testing of {@link StreamObserver}s. */
public class TestStreams {
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowMatchers.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowMatchers.java
index 469add3..4b69f07 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowMatchers.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowMatchers.java
@@ -18,7 +18,7 @@
package org.apache.beam.runners.dataflow.worker;
import java.io.Serializable;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.hamcrest.Description;
import org.hamcrest.TypeSafeMatcher;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContextTest.java
index 3412de3..a207f49 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContextTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowOperationContextTest.java
@@ -244,7 +244,8 @@
assertThat(
contents,
Matchers.allOf(
- Matchers.containsString("Processing stuck in step " + NameContextsForTests.USER_NAME),
+ Matchers.containsString(
+ "Operation ongoing in step " + NameContextsForTests.USER_NAME),
Matchers.containsString(" without outputting or completing in state somestate"),
Matchers.containsString("userpackage.SomeUserDoFn.helperMethod"),
Matchers.not(Matchers.containsString(SimpleDoFnRunner.class.getName()))));
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelperTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelperTest.java
index 66be62d..453dbea 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelperTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkerHarnessHelperTest.java
@@ -32,7 +32,7 @@
import org.apache.beam.runners.dataflow.worker.testing.RestoreDataflowLoggingMDC;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.RestoreSystemProperties;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubReaderTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubReaderTest.java
index fddfdc5..9cae865 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubReaderTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubReaderTest.java
@@ -31,7 +31,7 @@
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Test;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java
index 9f45286..a01356e 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java
@@ -29,7 +29,7 @@
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Test;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/ReaderCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/ReaderCacheTest.java
index 51f20a4..1820ef1 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/ReaderCacheTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/ReaderCacheTest.java
@@ -26,7 +26,7 @@
import java.io.IOException;
import java.util.concurrent.TimeUnit;
import org.apache.beam.sdk.io.UnboundedSource;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Stopwatch;
import org.joda.time.Duration;
import org.junit.Before;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StateFetcherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StateFetcherTest.java
index 5d90a85..6dc4f7b 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StateFetcherTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StateFetcherTest.java
@@ -43,7 +43,7 @@
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
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 c954bdb..7ce23a4 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
@@ -137,9 +137,9 @@
import org.apache.beam.sdk.values.ValueWithRecordId;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString.Output;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString.Output;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java
index 0acb1dc..bf9e875 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java
@@ -75,7 +75,7 @@
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Before;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java
index 11322a9..81a2bf5 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java
@@ -48,7 +48,7 @@
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.joda.time.Instant;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
index 03d4376..b0577e4 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
@@ -65,7 +65,7 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.hamcrest.Matchers;
import org.joda.time.Instant;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunnerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunnerTest.java
index 24d17ff..8d87a10 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunnerTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunnerTest.java
@@ -63,7 +63,7 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.hamcrest.Matchers;
import org.joda.time.Duration;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java
index d70d7f6..c1a9945 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java
@@ -49,7 +49,7 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
import org.hamcrest.Matchers;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java
index 4441f35..de62567 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java
@@ -40,7 +40,7 @@
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.hamcrest.Matchers;
import org.joda.time.Instant;
import org.junit.Before;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java
index bff116f..0b89c99 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillReaderIteratorBaseTest.java
@@ -26,7 +26,7 @@
import java.util.List;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateCacheTest.java
index 64f425a..23fe22f 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateCacheTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateCacheTest.java
@@ -28,7 +28,7 @@
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Test;
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 f708500..e693672 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
@@ -50,7 +50,7 @@
import org.apache.beam.sdk.transforms.Sum;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Futures;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateReaderTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateReaderTest.java
index ef529df..f2628ff 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateReaderTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillStateReaderTest.java
@@ -29,8 +29,8 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString.Output;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString.Output;
import org.hamcrest.Matchers;
import org.joda.time.Instant;
import org.junit.Before;
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 181183d..15ff67f 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
@@ -102,7 +102,7 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.ValueWithRecordId;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
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.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlServiceTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlServiceTest.java
index 0cac04e..f986783 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlServiceTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/BeamFnControlServiceTest.java
@@ -35,9 +35,9 @@
import org.apache.beam.runners.fnexecution.control.FnApiControlClient;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.net.HostAndPort;
import org.junit.Before;
@@ -88,7 +88,7 @@
Server server = ServerFactory.createDefault().create(ImmutableList.of(service), descriptor);
String url = service.getApiServiceDescriptor().getUrl();
BeamFnControlGrpc.BeamFnControlStub clientStub =
- BeamFnControlGrpc.newStub(ManagedChannelBuilder.forTarget(url).usePlaintext(true).build());
+ BeamFnControlGrpc.newStub(ManagedChannelBuilder.forTarget(url).usePlaintext().build());
// Connect from the client.
clientStub.control(requestObserver);
@@ -134,9 +134,9 @@
String url = service.getApiServiceDescriptor().getUrl();
BeamFnControlGrpc.BeamFnControlStub clientStub =
- BeamFnControlGrpc.newStub(ManagedChannelBuilder.forTarget(url).usePlaintext(true).build());
+ BeamFnControlGrpc.newStub(ManagedChannelBuilder.forTarget(url).usePlaintext().build());
BeamFnControlGrpc.BeamFnControlStub anotherClientStub =
- BeamFnControlGrpc.newStub(ManagedChannelBuilder.forTarget(url).usePlaintext(true).build());
+ BeamFnControlGrpc.newStub(ManagedChannelBuilder.forTarget(url).usePlaintext().build());
// Connect from the client.
clientStub.control(requestObserver);
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperationTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperationTest.java
index eb3d21d..a89dac6 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperationTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/RegisterAndProcessBundleOperationTest.java
@@ -80,7 +80,7 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.ValueInSingleWindow.Coder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableTable;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiverTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiverTest.java
index 2067a3a..a8c0586 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiverTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/control/TimerReceiverTest.java
@@ -70,7 +70,7 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
@@ -234,7 +234,9 @@
long testTimerOffset = 123456;
// Arbitrary key.
Object timer = timerBytes("X", testTimerOffset);
- Object windowedTimer = WindowedValue.valueInGlobalWindow(timer);
+ Object windowedTimer =
+ WindowedValue.timestampedValueInGlobalWindow(
+ timer, BoundedWindow.TIMESTAMP_MIN_VALUE.plus(testTimerOffset));
// Simulate the SDK Harness sending a timer element to the Runner Harness.
assertTrue(timerReceiver.receive(timerOutputPCollection, windowedTimer));
@@ -349,10 +351,14 @@
long testTimerOffset = 123456;
// Arbitrary key.
Object timer1 = timerBytes("X", testTimerOffset);
- Object windowedTimer1 = WindowedValue.valueInGlobalWindow(timer1);
+ Object windowedTimer1 =
+ WindowedValue.timestampedValueInGlobalWindow(
+ timer1, BoundedWindow.TIMESTAMP_MIN_VALUE.plus(testTimerOffset));
Object timer2 = timerBytes("Y", testTimerOffset);
- Object windowedTimer2 = WindowedValue.valueInGlobalWindow(timer2);
+ Object windowedTimer2 =
+ WindowedValue.timestampedValueInGlobalWindow(
+ timer2, BoundedWindow.TIMESTAMP_MIN_VALUE.plus(testTimerOffset));
// Simulate the SDK Harness sending a timer element to the Runner Harness.
assertTrue(
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcServiceTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcServiceTest.java
index 9c2b57a..ad514a8 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcServiceTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/data/BeamFnDataGrpcServiceTest.java
@@ -51,22 +51,22 @@
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.BindableService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.CallOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Channel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ClientCall;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ClientInterceptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Metadata;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Metadata.Key;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.MethodDescriptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerInterceptors;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.BindableService;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.CallOptions;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Channel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ClientCall;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ClientInterceptor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Metadata;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Metadata.Key;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.MethodDescriptor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerInterceptors;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingServiceTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingServiceTest.java
index 55b81e0..114ded8 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingServiceTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingServiceTest.java
@@ -38,12 +38,12 @@
import org.apache.beam.runners.dataflow.worker.fn.stream.ServerStreamObserverFactory;
import org.apache.beam.runners.fnexecution.GrpcContextHeaderAccessorProvider;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.BindableService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.BindableService;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.net.HostAndPort;
import org.junit.After;
import org.junit.Test;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactoryTest.java
index 43d6975..e41fd69 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactoryTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/fn/stream/ServerStreamObserverFactoryTest.java
@@ -24,8 +24,8 @@
import org.apache.beam.sdk.fn.stream.BufferingStreamObserver;
import org.apache.beam.sdk.fn.stream.DirectStreamObserver;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodesTest.java
index d1115c9..0e36473 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodesTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/graph/InsertFetchAndFilterStreamingSideInputNodesTest.java
@@ -54,7 +54,7 @@
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Equivalence;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Equivalence.Wrapper;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/logging/DataflowWorkerLoggingHandlerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/logging/DataflowWorkerLoggingHandlerTest.java
index 568fcff..84adfcd 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/logging/DataflowWorkerLoggingHandlerTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/logging/DataflowWorkerLoggingHandlerTest.java
@@ -35,7 +35,7 @@
import org.apache.beam.runners.dataflow.worker.NameContextsForTests;
import org.apache.beam.runners.dataflow.worker.TestOperationContext.TestDataflowExecutionState;
import org.apache.beam.runners.dataflow.worker.testing.RestoreDataflowLoggingMDC;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Timestamp;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Timestamp;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.junit.After;
import org.junit.Before;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/testing/GenericJsonMatcherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/testing/GenericJsonMatcherTest.java
index cac9fe3..146fb17 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/testing/GenericJsonMatcherTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/testing/GenericJsonMatcherTest.java
@@ -53,7 +53,9 @@
assertThat(actual, is(jsonOf(expected)));
} catch (AssertionError ex) {
assertEquals(
- "\nExpected: is {\"foo\":\"expected\"}\n but: was <{foo=actual}>", ex.getMessage());
+ "\nExpected: is {\"foo\":\"expected\"}\n"
+ + " but: was <GenericData{classInfo=[], {foo=actual}}>",
+ ex.getMessage());
// pass
return;
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServerTest.java
index 9adce9f..b889138 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServerTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServerTest.java
@@ -59,13 +59,13 @@
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.CommitWorkStream;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetDataStream;
import org.apache.beam.runners.dataflow.worker.windmill.WindmillServerStub.GetWorkStream;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.util.MutableHandlerRegistry;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.util.MutableHandlerRegistry;
import org.hamcrest.Matchers;
import org.joda.time.Instant;
import org.junit.After;
diff --git a/runners/java-fn-execution/build.gradle b/runners/java-fn-execution/build.gradle
index f032d8f..434ac02 100644
--- a/runners/java-fn-execution/build.gradle
+++ b/runners/java-fn-execution/build.gradle
@@ -30,7 +30,7 @@
compile project(":sdks:java:fn-execution")
compile project(":runners:core-construction-java")
compile project(path: ":vendor:sdks-java-extensions-protobuf", configuration: "shadow")
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
compile library.java.slf4j_api
compile library.java.args4j
testCompile project(":sdks:java:harness")
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/FnService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/FnService.java
index 3055b0b..634657a 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/FnService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/FnService.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.fnexecution;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.BindableService;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.BindableService;
/** An interface sharing common behavior with services used during execution of user Fns. */
public interface FnService extends AutoCloseable, BindableService {
@@ -26,8 +26,8 @@
*
* <p>There should be no more calls to any service method by the time a call to {@link #close()}
* begins. Specifically, this means that a {@link
- * org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server} that this service is bound to should have
- * completed a call to the {@link org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server#shutdown()}
+ * org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server} that this service is bound to should have
+ * completed a call to the {@link org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server#shutdown()}
* method, and all future incoming calls will be rejected.
*/
@Override
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProvider.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProvider.java
index 5d758a2..4c7899c 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProvider.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProvider.java
@@ -17,14 +17,14 @@
*/
package org.apache.beam.runners.fnexecution;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Context;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Contexts;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Metadata;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Metadata.Key;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerCall;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerCall.Listener;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerCallHandler;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerInterceptor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Context;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Contexts;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Metadata;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Metadata.Key;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerCall;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerCall.Listener;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerCallHandler;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerInterceptor;
/**
* A HeaderAccessorProvider which intercept the header in a GRPC request and expose the relevant
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcFnServer.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcFnServer.java
index f7a4a4b..fbb12af 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcFnServer.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/GrpcFnServer.java
@@ -20,7 +20,7 @@
import java.io.IOException;
import java.util.concurrent.TimeUnit;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
/**
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/InProcessServerFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/InProcessServerFactory.java
index a899cb2..e72b0bc 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/InProcessServerFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/InProcessServerFactory.java
@@ -21,10 +21,10 @@
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.BindableService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerInterceptors;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.BindableService;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerInterceptors;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
/**
* A {@link ServerFactory} which creates {@link Server servers} with the {@link
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/ServerFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/ServerFactory.java
index ff0d5b4..30f6b8b 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/ServerFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/ServerFactory.java
@@ -29,16 +29,16 @@
import java.util.function.Supplier;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.sdk.fn.channel.SocketAddressFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.BindableService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerInterceptors;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.netty.NettyServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.EpollEventLoopGroup;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.EpollServerDomainSocketChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.EpollServerSocketChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.unix.DomainSocketAddress;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.util.internal.ThreadLocalRandom;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.BindableService;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerInterceptors;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.netty.NettyServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.EpollEventLoopGroup;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.EpollServerDomainSocketChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.EpollServerSocketChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.unix.DomainSocketAddress;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.util.internal.ThreadLocalRandom;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.net.HostAndPort;
/** A {@link Server gRPC server} factory. */
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactRetrievalService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactRetrievalService.java
index 72af9e81..a9b04b3 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactRetrievalService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactRetrievalService.java
@@ -29,11 +29,11 @@
import org.apache.beam.model.jobmanagement.v1.ArtifactApi.ArtifactMetadata;
import org.apache.beam.model.jobmanagement.v1.ArtifactApi.ProxyManifest;
import org.apache.beam.model.jobmanagement.v1.ArtifactRetrievalServiceGrpc;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.JsonFormat;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.JsonFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.Cache;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.cache.CacheBuilder;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactStagingService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactStagingService.java
index 0b0fadf..ae951d9 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactStagingService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/AbstractArtifactStagingService.java
@@ -35,11 +35,11 @@
import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc.ArtifactStagingServiceImplBase;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.fnexecution.FnService;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.JsonFormat;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.JsonFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hasher;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing;
import org.slf4j.Logger;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactStagingService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactStagingService.java
index c9baa17..74bce71 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactStagingService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactStagingService.java
@@ -31,8 +31,8 @@
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.util.MimeTypes;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClient.java
index 9051051..e0e5d50 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClient.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClient.java
@@ -29,9 +29,9 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest;
import org.apache.beam.sdk.fn.stream.SynchronizedStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolService.java
index 598e1db..82409ae 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolService.java
@@ -26,7 +26,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnControlGrpc;
import org.apache.beam.runners.fnexecution.FnService;
import org.apache.beam.runners.fnexecution.HeaderAccessor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java
index cd81c0a..80b3f61 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java
@@ -33,6 +33,7 @@
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
+import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.runners.core.construction.ModelCoders;
@@ -56,7 +57,7 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableTable;
@@ -119,15 +120,19 @@
ImmutableMap.builder();
ImmutableMap.Builder<String, Coder> remoteOutputCodersBuilder = ImmutableMap.builder();
+ WireCoderSetting wireCoderSetting =
+ stage.getWireCoderSettings().stream()
+ .filter(ws -> ws.getInputOrOutputId().equals(stage.getInputPCollection().getId()))
+ .findAny()
+ .orElse(WireCoderSetting.getDefaultInstance());
// The order of these does not matter.
inputDestinationsBuilder.put(
stage.getInputPCollection().getId(),
- addStageInput(
- dataEndpoint, stage.getInputPCollection(), components, stage.getWireCoderSetting()));
+ addStageInput(dataEndpoint, stage.getInputPCollection(), components, wireCoderSetting));
remoteOutputCodersBuilder.putAll(
addStageOutputs(
- dataEndpoint, stage.getOutputPCollections(), components, stage.getWireCoderSetting()));
+ dataEndpoint, stage.getOutputPCollections(), components, stage.getWireCoderSettings()));
Map<String, Map<String, SideInputSpec>> sideInputSpecs = addSideInputs(stage, components);
@@ -192,10 +197,15 @@
ApiServiceDescriptor dataEndpoint,
Collection<PCollectionNode> outputPCollections,
Components.Builder components,
- RunnerApi.WireCoderSetting wireCoderSetting)
+ Collection<WireCoderSetting> wireCoderSettings)
throws IOException {
Map<String, Coder<WindowedValue<?>>> remoteOutputCoders = new LinkedHashMap<>();
for (PCollectionNode outputPCollection : outputPCollections) {
+ WireCoderSetting wireCoderSetting =
+ wireCoderSettings.stream()
+ .filter(ws -> ws.getInputOrOutputId().equals(outputPCollection.getId()))
+ .findAny()
+ .orElse(WireCoderSetting.getDefaultInstance());
OutputEncoding outputEncoding =
addStageOutput(dataEndpoint, components, outputPCollection, wireCoderSetting);
remoteOutputCoders.put(outputEncoding.getPTransformId(), outputEncoding.getCoder());
@@ -207,7 +217,7 @@
ApiServiceDescriptor dataEndpoint,
PCollectionNode inputPCollection,
Components.Builder components,
- RunnerApi.WireCoderSetting wireCoderSetting)
+ WireCoderSetting wireCoderSetting)
throws IOException {
String inputWireCoderId =
WireCoders.addSdkWireCoder(inputPCollection, components, wireCoderSetting);
@@ -235,7 +245,7 @@
ApiServiceDescriptor dataEndpoint,
Components.Builder components,
PCollectionNode outputPCollection,
- RunnerApi.WireCoderSetting wireCoderSetting)
+ WireCoderSetting wireCoderSetting)
throws IOException {
String outputWireCoderId =
WireCoders.addSdkWireCoder(outputPCollection, components, wireCoderSetting);
@@ -385,6 +395,17 @@
.setCoderId(timerCoderId)
.build();
+ // The wire coder setting for both input and output of the timer. We haven't provided
+ // different settings for input and output now, because the timerCollectionSpec is same for
+ // both input and output of the timer.
+ WireCoderSetting wireCoderSetting =
+ stage.getWireCoderSettings().stream()
+ .filter(
+ ws ->
+ ws.getTimer().getTransformId().equals(timerReference.transform().getId())
+ && ws.getTimer().getLocalName().equals(timerReference.localName()))
+ .findAny()
+ .orElse(WireCoderSetting.getDefaultInstance());
// "Unroll" the timers into PCollections.
String inputTimerPCollectionId =
SyntheticComponents.uniqueId(
@@ -398,7 +419,7 @@
dataEndpoint,
PipelineNode.pCollection(inputTimerPCollectionId, timerCollectionSpec),
components,
- stage.getWireCoderSetting()));
+ wireCoderSetting));
String outputTimerPCollectionId =
SyntheticComponents.uniqueId(
String.format(
@@ -411,7 +432,7 @@
dataEndpoint,
components,
PipelineNode.pCollection(outputTimerPCollectionId, timerCollectionSpec),
- stage.getWireCoderSetting());
+ wireCoderSetting);
outputTransformCodersBuilder.put(outputEncoding.getPTransformId(), outputEncoding.getCoder());
components.putTransforms(
timerReference.transform().getId(),
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java
index 5d58c5c..1c283bf 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/data/GrpcDataService.java
@@ -36,7 +36,7 @@
import org.apache.beam.sdk.fn.data.LogicalEndpoint;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.SettableFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobService.java
index 2f4df48..abda968 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobService.java
@@ -21,7 +21,7 @@
import java.util.List;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.function.Consumer;
import java.util.function.Function;
import org.apache.beam.model.jobmanagement.v1.JobApi;
@@ -52,11 +52,12 @@
import org.apache.beam.sdk.fn.stream.SynchronizedStreamObserver;
import org.apache.beam.sdk.function.ThrowingConsumer;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -72,12 +73,16 @@
public class InMemoryJobService extends JobServiceGrpc.JobServiceImplBase implements FnService {
private static final Logger LOG = LoggerFactory.getLogger(InMemoryJobService.class);
+ /** The default maximum number of completed invocations to keep. */
+ public static final int DEFAULT_MAX_INVOCATION_HISTORY = 10;
+
/**
* Creates an InMemoryJobService.
*
* @param stagingServiceDescriptor Endpoint for the staging service.
* @param stagingServiceTokenProvider Function mapping a preparationId to a staging service token.
- * @param invoker A JobInvoker that will actually create the jobs.
+ * @param cleanupJobFn A cleanup function to run, parameterized with the staging token of a job.
+ * @param invoker A JobInvoker which creates the jobs.
* @return A new InMemoryJobService.
*/
public static InMemoryJobService create(
@@ -86,22 +91,60 @@
ThrowingConsumer<Exception, String> cleanupJobFn,
JobInvoker invoker) {
return new InMemoryJobService(
- stagingServiceDescriptor, stagingServiceTokenProvider, cleanupJobFn, invoker);
+ stagingServiceDescriptor,
+ stagingServiceTokenProvider,
+ cleanupJobFn,
+ invoker,
+ DEFAULT_MAX_INVOCATION_HISTORY);
}
- private final ConcurrentMap<String, JobPreparation> preparations;
- private final ConcurrentMap<String, JobInvocation> invocations;
- private final ConcurrentMap<String, String> stagingSessionTokens;
+ /**
+ * Creates an InMemoryJobService.
+ *
+ * @param stagingServiceDescriptor The endpoint for the staging service.
+ * @param stagingServiceTokenProvider Function mapping a preparationId to a staging service token.
+ * @param cleanupJobFn A cleanup function to run, parameterized with the staging token of a job.
+ * @param invoker A JobInvoker which creates the jobs.
+ * @param maxInvocationHistory The maximum number of completed invocations to keep.
+ * @return A new InMemoryJobService.
+ */
+ public static InMemoryJobService create(
+ Endpoints.ApiServiceDescriptor stagingServiceDescriptor,
+ Function<String, String> stagingServiceTokenProvider,
+ ThrowingConsumer<Exception, String> cleanupJobFn,
+ JobInvoker invoker,
+ int maxInvocationHistory) {
+ return new InMemoryJobService(
+ stagingServiceDescriptor,
+ stagingServiceTokenProvider,
+ cleanupJobFn,
+ invoker,
+ maxInvocationHistory);
+ }
+
+ /** Map of preparationId to preparation. */
+ private final ConcurrentHashMap<String, JobPreparation> preparations;
+ /** Map of preparationId to staging token. */
+ private final ConcurrentHashMap<String, String> stagingSessionTokens;
+ /** Map of invocationId to invocation. */
+ private final ConcurrentHashMap<String, JobInvocation> invocations;
+ /** InvocationIds of completed invocations in least-recently-completed order. */
+ private final ConcurrentLinkedDeque<String> completedInvocationsIds;
+
private final Endpoints.ApiServiceDescriptor stagingServiceDescriptor;
private final Function<String, String> stagingServiceTokenProvider;
private final ThrowingConsumer<Exception, String> cleanupJobFn;
private final JobInvoker invoker;
+ /** The maximum number of past invocations to keep. */
+ private final int maxInvocationHistory;
+
private InMemoryJobService(
Endpoints.ApiServiceDescriptor stagingServiceDescriptor,
Function<String, String> stagingServiceTokenProvider,
ThrowingConsumer<Exception, String> cleanupJobFn,
- JobInvoker invoker) {
+ JobInvoker invoker,
+ int maxInvocationHistory) {
this.stagingServiceDescriptor = stagingServiceDescriptor;
this.stagingServiceTokenProvider = stagingServiceTokenProvider;
this.cleanupJobFn = cleanupJobFn;
@@ -109,9 +152,10 @@
this.preparations = new ConcurrentHashMap<>();
this.invocations = new ConcurrentHashMap<>();
-
- // Map "preparation ID" to staging token
this.stagingSessionTokens = new ConcurrentHashMap<>();
+ this.completedInvocationsIds = new ConcurrentLinkedDeque<>();
+ Preconditions.checkArgument(maxInvocationHistory >= 0);
+ this.maxInvocationHistory = maxInvocationHistory;
}
@Override
@@ -196,20 +240,25 @@
}
String stagingSessionToken = stagingSessionTokens.get(preparationId);
stagingSessionTokens.remove(preparationId);
- if (cleanupJobFn != null) {
- try {
+ try {
+ if (cleanupJobFn != null) {
cleanupJobFn.accept(stagingSessionToken);
- } catch (Exception e) {
- LOG.warn(
- "Failed to remove job staging directory for token {}: {}",
- stagingSessionToken,
- e);
}
+ } catch (Exception e) {
+ LOG.warn(
+ "Failed to remove job staging directory for token {}: {}",
+ stagingSessionToken,
+ e);
+ } finally {
+ onFinishedInvocationCleanup(invocationId);
}
});
invocation.start();
invocations.put(invocationId, invocation);
+ // Cleanup this preparation because we are running it now.
+ // If we fail, we need to prepare again.
+ preparations.remove(preparationId);
RunJobResponse response = RunJobResponse.newBuilder().setJobId(invocationId).build();
responseObserver.onNext(response);
responseObserver.onCompleted();
@@ -344,7 +393,9 @@
event -> {
syncResponseObserver.onNext(
JobMessagesResponse.newBuilder().setStateResponse(event).build());
- if (JobInvocation.isTerminated(event.getState())) {
+ // The terminal state is always updated after the last message, that's
+ // why we can end the stream here.
+ if (JobInvocation.isTerminated(invocation.getStateEvent().getState())) {
responseObserver.onCompleted();
}
};
@@ -353,8 +404,11 @@
syncResponseObserver.onNext(
JobMessagesResponse.newBuilder().setMessageResponse(message).build());
- invocation.addStateListener(stateListener);
invocation.addMessageListener(messageListener);
+ // The order matters here. Make sure to send all the message first because the stream
+ // will be ended by the terminal state request.
+ invocation.addStateListener(stateListener);
+
} catch (StatusRuntimeException | StatusException e) {
responseObserver.onError(e);
} catch (Exception e) {
@@ -421,4 +475,14 @@
}
return invocation;
}
+
+ private void onFinishedInvocationCleanup(String invocationId) {
+ completedInvocationsIds.addLast(invocationId);
+ while (completedInvocationsIds.size() > maxInvocationHistory) {
+ // Clean up invocations
+ // "preparations" is cleaned up when adding to "invocations"
+ // "stagingTokens" is cleaned up when the invocation finishes
+ invocations.remove(completedInvocationsIds.removeFirst());
+ }
+ }
}
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocation.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocation.java
index 2da0592..6776493 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocation.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocation.java
@@ -36,7 +36,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.sdk.PipelineResult;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.Timestamps;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.Timestamps;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.FutureCallback;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Futures;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ListenableFuture;
@@ -53,12 +53,13 @@
private final PortablePipelineRunner pipelineRunner;
private final JobInfo jobInfo;
private final ListeningExecutorService executorService;
- private List<Consumer<JobStateEvent>> stateObservers;
- private List<Consumer<JobMessage>> messageObservers;
+ private final List<JobStateEvent> stateHistory;
+ private final List<JobMessage> messageHistory;
+ private final List<Consumer<JobStateEvent>> stateObservers;
+ private final List<Consumer<JobMessage>> messageObservers;
private JobApi.MetricResults metrics;
private PortablePipelineResult resultHandle;
@Nullable private ListenableFuture<PortablePipelineResult> invocationFuture;
- private List<JobStateEvent> stateHistory;
public JobInvocation(
JobInfo jobInfo,
@@ -73,6 +74,7 @@
this.messageObservers = new ArrayList<>();
this.invocationFuture = null;
this.stateHistory = new ArrayList<>();
+ this.messageHistory = new ArrayList<>();
this.metrics = JobApi.MetricResults.newBuilder().build();
this.setState(JobState.Enum.STOPPED);
}
@@ -217,6 +219,9 @@
/** Listen for job messages with a {@link Consumer}. */
public synchronized void addMessageListener(Consumer<JobMessage> messageStreamObserver) {
+ for (JobMessage msg : messageHistory) {
+ messageStreamObserver.accept(msg);
+ }
messageObservers.add(messageStreamObserver);
}
@@ -243,6 +248,7 @@
}
private synchronized void sendMessage(JobMessage message) {
+ messageHistory.add(message);
for (Consumer<JobMessage> observer : messageObservers) {
observer.accept(message);
}
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvoker.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvoker.java
index 7612d8b..0f66c38 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvoker.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvoker.java
@@ -22,7 +22,7 @@
import java.util.concurrent.ThreadFactory;
import javax.annotation.Nullable;
import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ListeningExecutorService;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.MoreExecutors;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobPreparation.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobPreparation.java
index a304093..30b0774 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobPreparation.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobPreparation.java
@@ -19,7 +19,7 @@
import com.google.auto.value.AutoValue;
import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
/** A job that has been prepared, but not invoked. */
@AutoValue
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobServerDriver.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobServerDriver.java
index d0061d2..d61162f 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobServerDriver.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/JobServerDriver.java
@@ -63,7 +63,8 @@
artifactStagingServer.getService().removeArtifacts(stagingSessionToken);
}
},
- invoker);
+ invoker,
+ configuration.getMaxInvocationHistory());
}
/** Configuration for the jobServer. */
@@ -97,6 +98,9 @@
handler = ExplicitBooleanOptionHandler.class)
private boolean cleanArtifactsPerJob = true;
+ @Option(name = "--history-size", usage = "The maximum number of completed jobs to keep.")
+ private int maxInvocationHistory = 10;
+
public String getHost() {
return host;
}
@@ -120,6 +124,10 @@
public boolean isCleanArtifactsPerJob() {
return cleanArtifactsPerJob;
}
+
+ public int getMaxInvocationHistory() {
+ return maxInvocationHistory;
+ }
}
protected static ServerFactory createJobServerFactory(ServerConfiguration configuration) {
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarCreator.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarCreator.java
index c14098a..bb70158 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarCreator.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarCreator.java
@@ -58,9 +58,9 @@
import org.apache.beam.sdk.metrics.MetricResults;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PortablePipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.MessageOrBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.JsonFormat;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.MessageOrBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.JsonFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
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.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarUtils.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarUtils.java
index 291605a..d32f1e1 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarUtils.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/jobsubmission/PortablePipelineJarUtils.java
@@ -27,9 +27,9 @@
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Message.Builder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.JsonFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Message.Builder;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.JsonFormat;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.ByteStreams;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingService.java
index a37a2f3..aa7117d 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingService.java
@@ -24,7 +24,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.LogControl;
import org.apache.beam.model.fnexecution.v1.BeamFnLoggingGrpc;
import org.apache.beam.runners.fnexecution.FnService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/JobInfo.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/JobInfo.java
index aea6bb3..f9d9c66 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/JobInfo.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/JobInfo.java
@@ -20,7 +20,7 @@
import com.google.auto.value.AutoValue;
import java.io.Serializable;
import org.apache.beam.model.fnexecution.v1.ProvisionApi;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
/**
* A subset of {@link org.apache.beam.model.fnexecution.v1.ProvisionApi.ProvisionInfo} that
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionService.java
index 4fec80c..aeece77 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionService.java
@@ -24,7 +24,7 @@
import org.apache.beam.model.fnexecution.v1.ProvisionServiceGrpc;
import org.apache.beam.model.fnexecution.v1.ProvisionServiceGrpc.ProvisionServiceImplBase;
import org.apache.beam.runners.fnexecution.FnService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* A {@link ProvisionServiceImplBase provision service} that returns a static response to all calls.
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/splittabledofn/SDFFeederViaStateAndTimers.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/splittabledofn/SDFFeederViaStateAndTimers.java
index 920dae6..d4edbd2 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/splittabledofn/SDFFeederViaStateAndTimers.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/splittabledofn/SDFFeederViaStateAndTimers.java
@@ -43,8 +43,8 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.Durations;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.Durations;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.joda.time.Instant;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/GrpcStateService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/GrpcStateService.java
index 9c72d81..df10910 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/GrpcStateService.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/GrpcStateService.java
@@ -28,8 +28,8 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnStateGrpc;
import org.apache.beam.runners.fnexecution.FnService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ServerCallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ServerCallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/** An implementation of the Beam Fn State service. */
public class GrpcStateService extends BeamFnStateGrpc.BeamFnStateImplBase
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/InMemoryBagUserStateFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/InMemoryBagUserStateFactory.java
index f840864..988f63db 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/InMemoryBagUserStateFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/InMemoryBagUserStateFactory.java
@@ -31,7 +31,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
/**
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/StateRequestHandlers.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/StateRequestHandlers.java
index 26dd6ac..b05776e 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/StateRequestHandlers.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/state/StateRequestHandlers.java
@@ -51,7 +51,7 @@
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.common.Reiterable;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.sdk.v2.sdk.extensions.protobuf.ByteStringCoder;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/BeamWorkerStatusGrpcService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/BeamWorkerStatusGrpcService.java
new file mode 100644
index 0000000..c3f9bba
--- /dev/null
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/BeamWorkerStatusGrpcService.java
@@ -0,0 +1,223 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.fnexecution.status;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusResponse;
+import org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc.BeamFnWorkerStatusImplBase;
+import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
+import org.apache.beam.runners.fnexecution.FnService;
+import org.apache.beam.runners.fnexecution.HeaderAccessor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
+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.Strings;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A Fn Status service which can collect run-time status information from SDK harnesses for
+ * debugging purpose.
+ */
+public class BeamWorkerStatusGrpcService extends BeamFnWorkerStatusImplBase implements FnService {
+
+ private static final Logger LOG = LoggerFactory.getLogger(BeamWorkerStatusGrpcService.class);
+ private static final String DEFAULT_EXCEPTION_RESPONSE =
+ "Error: exception encountered getting status from SDK harness";
+
+ private final HeaderAccessor headerAccessor;
+ private final Map<String, CompletableFuture<WorkerStatusClient>> connectedClient =
+ Collections.synchronizedMap(new HashMap<>());
+ private final AtomicBoolean isClosed = new AtomicBoolean();
+
+ private BeamWorkerStatusGrpcService(
+ ApiServiceDescriptor apiServiceDescriptor, HeaderAccessor headerAccessor) {
+ this.headerAccessor = headerAccessor;
+ LOG.info("Launched Beam Fn Status service at {}", apiServiceDescriptor);
+ }
+
+ /**
+ * Create new instance of {@link BeamWorkerStatusGrpcService}.
+ *
+ * @param apiServiceDescriptor describes the configuration for the endpoint the server will
+ * expose.
+ * @param headerAccessor headerAccessor gRPC header accessor used to obtain SDK harness worker id.
+ * @return {@link BeamWorkerStatusGrpcService}
+ */
+ public static BeamWorkerStatusGrpcService create(
+ ApiServiceDescriptor apiServiceDescriptor, HeaderAccessor headerAccessor) {
+ return new BeamWorkerStatusGrpcService(apiServiceDescriptor, headerAccessor);
+ }
+
+ @Override
+ public void close() throws Exception {
+ if (isClosed.getAndSet(true)) {
+ return;
+ }
+ synchronized (connectedClient) {
+ for (CompletableFuture<WorkerStatusClient> clientFuture : connectedClient.values()) {
+ if (clientFuture.isDone()) {
+ clientFuture.get().close();
+ }
+ }
+ connectedClient.clear();
+ }
+ }
+
+ @Override
+ public StreamObserver<WorkerStatusResponse> workerStatus(
+ StreamObserver<WorkerStatusRequest> requestObserver) {
+ if (isClosed.get()) {
+ throw new IllegalStateException("BeamWorkerStatusGrpcService already closed.");
+ }
+ String workerId = headerAccessor.getSdkWorkerId();
+ LOG.info("Beam Fn Status client connected with id {}", workerId);
+
+ WorkerStatusClient fnApiStatusClient =
+ WorkerStatusClient.forRequestObserver(workerId, requestObserver);
+ connectedClient.compute(
+ workerId,
+ (k, existingClientFuture) -> {
+ if (existingClientFuture != null) {
+ try {
+ if (existingClientFuture.isDone()) {
+ LOG.info(
+ "SDK Worker {} was connected to status server previously, disconnecting old client",
+ workerId);
+ existingClientFuture.get().close();
+ } else {
+ existingClientFuture.complete(fnApiStatusClient);
+ return existingClientFuture;
+ }
+ } catch (IOException | InterruptedException | ExecutionException e) {
+ LOG.warn("Error closing worker status client", e);
+ }
+ }
+ return CompletableFuture.completedFuture(fnApiStatusClient);
+ });
+ return fnApiStatusClient.getResponseObserver();
+ }
+
+ /**
+ * Get the latest SDK worker status from the client's corresponding SDK harness.
+ *
+ * @param workerId worker id of the SDK harness.
+ * @return {@link CompletableFuture} of WorkerStatusResponse from SDK harness.
+ */
+ public String getSingleWorkerStatus(String workerId, long timeout, TimeUnit timeUnit) {
+ if (isClosed.get()) {
+ throw new IllegalStateException("BeamWorkerStatusGrpcService already closed.");
+ }
+ try {
+ return getWorkerStatus(workerId).get(timeout, timeUnit);
+ } catch (InterruptedException | ExecutionException | TimeoutException e) {
+ return handleAndReturnExceptionResponse(e);
+ }
+ }
+
+ /**
+ * Get all the statuses from all connected SDK harnesses within specified timeout. Any errors
+ * getting status from the SDK harnesses will be returned in the map.
+ *
+ * @param timeout max time waiting for the response from each SDK harness.
+ * @param timeUnit timeout time unit.
+ * @return All the statuses in a map keyed by the SDK harness id.
+ */
+ public Map<String, String> getAllWorkerStatuses(long timeout, TimeUnit timeUnit) {
+ if (isClosed.get()) {
+ throw new IllegalStateException("BeamWorkerStatusGrpcService already closed.");
+ }
+ // return result in worker id sorted map.
+ Map<String, String> allStatuses = new ConcurrentSkipListMap<>(Comparator.naturalOrder());
+ Set<String> connectedClientIdsCopy;
+ synchronized (connectedClient) {
+ connectedClientIdsCopy = ImmutableSet.copyOf(connectedClient.keySet());
+ }
+ connectedClientIdsCopy
+ .parallelStream()
+ .forEach(
+ workerId ->
+ allStatuses.put(workerId, getSingleWorkerStatus(workerId, timeout, timeUnit)));
+
+ return allStatuses;
+ }
+
+ @VisibleForTesting
+ CompletableFuture<String> getWorkerStatus(String workerId) {
+ CompletableFuture<WorkerStatusClient> statusClient;
+ try {
+ statusClient = getStatusClient(workerId);
+ if (!statusClient.isDone()) {
+ return CompletableFuture.completedFuture("Error: Not connected.");
+ }
+ CompletableFuture<WorkerStatusResponse> future = statusClient.get().getWorkerStatus();
+ return future.thenApply(this::getStatusErrorOrInfo);
+ } catch (ExecutionException | InterruptedException e) {
+ return CompletableFuture.completedFuture(handleAndReturnExceptionResponse(e));
+ }
+ }
+
+ /**
+ * Get the status api client connected to the SDK harness with specified workerId.
+ *
+ * @param workerId worker id of the SDK harness.
+ * @return CompletableFuture of {@link WorkerStatusClient}.
+ */
+ @VisibleForTesting
+ CompletableFuture<WorkerStatusClient> getStatusClient(String workerId) {
+ return connectedClient.computeIfAbsent(workerId, k -> new CompletableFuture<>());
+ }
+
+ /**
+ * Return Error field from WorkerStatusResponse if not empty, otherwise return the StatusInfo
+ * field.
+ */
+ private String getStatusErrorOrInfo(WorkerStatusResponse response) {
+ return !Strings.isNullOrEmpty(response.getError())
+ ? response.getError()
+ : response.getStatusInfo();
+ }
+
+ private String handleAndReturnExceptionResponse(Exception e) {
+ LOG.warn(DEFAULT_EXCEPTION_RESPONSE, e);
+ if (e instanceof InterruptedException) {
+ Thread.currentThread().interrupt();
+ }
+ StringBuilder response = new StringBuilder();
+ response
+ .append(DEFAULT_EXCEPTION_RESPONSE)
+ .append(": ")
+ .append(e.getClass().getCanonicalName());
+ if (e.getMessage() != null) {
+ response.append(": ").append(e.getMessage());
+ }
+ return response.toString();
+ }
+}
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/WorkerStatusClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/WorkerStatusClient.java
new file mode 100644
index 0000000..eacb3fc
--- /dev/null
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/WorkerStatusClient.java
@@ -0,0 +1,165 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.fnexecution.status;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusResponse;
+import org.apache.beam.sdk.fn.IdGenerator;
+import org.apache.beam.sdk.fn.IdGenerators;
+import org.apache.beam.sdk.fn.stream.SynchronizedStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Client for handling requests and responses over Fn Worker Status Api between runner and SDK
+ * harness.
+ */
+class WorkerStatusClient implements Closeable {
+
+ public static final Logger LOG = LoggerFactory.getLogger(WorkerStatusClient.class);
+ private final IdGenerator idGenerator = IdGenerators.incrementingLongs();
+ private final StreamObserver<WorkerStatusRequest> requestReceiver;
+ private final Map<String, CompletableFuture<WorkerStatusResponse>> pendingResponses =
+ Collections.synchronizedMap(new HashMap<>());
+ private final String workerId;
+ private AtomicBoolean isClosed = new AtomicBoolean(false);
+
+ private WorkerStatusClient(String workerId, StreamObserver<WorkerStatusRequest> requestReceiver) {
+ this.requestReceiver = SynchronizedStreamObserver.wrapping(requestReceiver);
+ this.workerId = workerId;
+ }
+
+ /**
+ * Create new status api client with SDK harness worker id and request observer.
+ *
+ * @param workerId SDK harness worker id.
+ * @param requestObserver The outbound request observer this client uses to send new status
+ * requests to its corresponding SDK harness.
+ * @return {@link WorkerStatusClient}
+ */
+ public static WorkerStatusClient forRequestObserver(
+ String workerId, StreamObserver<WorkerStatusRequest> requestObserver) {
+ return new WorkerStatusClient(workerId, requestObserver);
+ }
+
+ /**
+ * Get the latest sdk worker status from the client's corresponding SDK harness. A random id will
+ * be used to specify the request_id field.
+ *
+ * @return {@link CompletableFuture} of the SDK harness status response.
+ */
+ public CompletableFuture<WorkerStatusResponse> getWorkerStatus() {
+ WorkerStatusRequest request =
+ WorkerStatusRequest.newBuilder().setId(idGenerator.getId()).build();
+ return getWorkerStatus(request);
+ }
+
+ /**
+ * Get the latest sdk worker status from the client's corresponding SDK harness with request.
+ *
+ * @param request WorkerStatusRequest to be sent to SDK harness.
+ * @return {@link CompletableFuture} of the SDK harness status response.
+ */
+ CompletableFuture<WorkerStatusResponse> getWorkerStatus(WorkerStatusRequest request) {
+ CompletableFuture<WorkerStatusResponse> future = new CompletableFuture<>();
+ if (isClosed.get()) {
+ future.completeExceptionally(new RuntimeException("Worker status client already closed."));
+ return future;
+ }
+ this.pendingResponses.put(request.getId(), future);
+ this.requestReceiver.onNext(request);
+ return future;
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (isClosed.getAndSet(true)) {
+ return;
+ }
+ synchronized (pendingResponses) {
+ for (CompletableFuture<WorkerStatusResponse> pendingResponse : pendingResponses.values()) {
+ pendingResponse.completeExceptionally(
+ new RuntimeException("Fn Status Api client shut down while waiting for the request"));
+ }
+ pendingResponses.clear();
+ }
+ requestReceiver.onCompleted();
+ }
+
+ /** Check if the client connection has already been closed. */
+ public boolean isClosed() {
+ return isClosed.get();
+ }
+
+ /** Get the worker id for the client's corresponding SDK harness. */
+ public String getWorkerId() {
+ return this.workerId;
+ }
+
+ /** Get the response observer of this client for retrieving inbound worker status responses. */
+ public StreamObserver<WorkerStatusResponse> getResponseObserver() {
+ return new ResponseStreamObserver();
+ }
+
+ /**
+ * ResponseObserver for handling status responses. Each request will be cached with it's
+ * request_id. Upon receiving response from SDK harness with this StreamObserver, the future
+ * mapped to same request_id will be finished accordingly.
+ */
+ private class ResponseStreamObserver implements StreamObserver<WorkerStatusResponse> {
+
+ @Override
+ public void onNext(WorkerStatusResponse response) {
+ if (isClosed.get()) {
+ return;
+ }
+ CompletableFuture<WorkerStatusResponse> future = pendingResponses.remove(response.getId());
+ if (future != null) {
+ future.complete(response);
+ } else {
+ LOG.warn(
+ String.format(
+ "Received response for status with unknown response id %s and status %s",
+ response.getId(), response.getStatusInfo()));
+ }
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ LOG.error("{} received error {}", WorkerStatusClient.class.getSimpleName(), throwable);
+ onCompleted();
+ }
+
+ @Override
+ public void onCompleted() {
+ try {
+ close();
+ } catch (IOException e) {
+ LOG.warn("Error closing Fn status api client", e);
+ }
+ }
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/package-info.java
similarity index 65%
copy from runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
copy to runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/package-info.java
index a114f40..65edce9 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/status/package-info.java
@@ -15,16 +15,6 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.runners.flink.translation.utils;
-import com.fasterxml.jackson.databind.type.TypeFactory;
-
-/** Utilities for dealing with classloading. */
-public class FlinkClassloading {
-
- public static void deleteStaticCaches() {
- // Clear cache to get rid of any references to the Flink Classloader
- // See https://jira.apache.org/jira/browse/BEAM-6460
- TypeFactory.defaultInstance().clearCache();
- }
-}
+/** Worker Status API services. */
+package org.apache.beam.runners.fnexecution.status;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java
index 6e6c80a..56b8179 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java
@@ -40,7 +40,7 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.BiMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableBiMap;
@@ -139,10 +139,11 @@
Preconditions.checkArgument(namespace instanceof StateNamespaces.WindowNamespace);
BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
Instant timestamp = timer.getTimestamp();
+ Instant outputTimestamp = timer.getOutputTimestamp();
WindowedValue<KV<Object, Timer>> timerValue =
WindowedValue.of(
KV.of(currentTimerKey, Timer.of(timestamp, new byte[0])),
- timestamp,
+ outputTimestamp,
Collections.singleton(window),
PaneInfo.NO_FIRING);
timerConsumer.accept(timer.getTimerId(), timerValue);
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/wire/WireCoders.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/wire/WireCoders.java
index f37901b..8cf7cd3 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/wire/WireCoders.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/wire/WireCoders.java
@@ -18,11 +18,11 @@
package org.apache.beam.runners.fnexecution.wire;
import static org.apache.beam.runners.core.construction.BeamUrns.getUrn;
-import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTING;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import java.io.IOException;
import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting;
import org.apache.beam.runners.core.construction.ModelCoders;
import org.apache.beam.runners.core.construction.RehydratedComponents;
import org.apache.beam.runners.core.construction.SyntheticComponents;
@@ -45,7 +45,7 @@
public static String addSdkWireCoder(
PCollectionNode pCollectionNode,
RunnerApi.Components.Builder components,
- RunnerApi.WireCoderSetting wireCoderSetting) {
+ WireCoderSetting wireCoderSetting) {
return addWireCoder(pCollectionNode, components, false, wireCoderSetting);
}
@@ -60,7 +60,7 @@
public static String addRunnerWireCoder(
PCollectionNode pCollectionNode,
RunnerApi.Components.Builder components,
- RunnerApi.WireCoderSetting wireCoderSetting) {
+ WireCoderSetting wireCoderSetting) {
return addWireCoder(pCollectionNode, components, true, wireCoderSetting);
}
@@ -72,7 +72,8 @@
*/
public static <T> Coder<WindowedValue<T>> instantiateRunnerWireCoder(
PCollectionNode pCollectionNode, RunnerApi.Components components) throws IOException {
- return instantiateRunnerWireCoder(pCollectionNode, components, DEFAULT_WIRE_CODER_SETTING);
+ return instantiateRunnerWireCoder(
+ pCollectionNode, components, WireCoderSetting.getDefaultInstance());
}
/**
@@ -84,7 +85,7 @@
public static <T> Coder<WindowedValue<T>> instantiateRunnerWireCoder(
PCollectionNode pCollectionNode,
RunnerApi.Components components,
- RunnerApi.WireCoderSetting wireCoderSetting)
+ WireCoderSetting wireCoderSetting)
throws IOException {
// NOTE: We discard the new set of components so we don't bother to ensure it's consistent with
// the caller's view.
@@ -104,7 +105,7 @@
PCollectionNode pCollectionNode,
RunnerApi.Components.Builder components,
boolean useByteArrayCoder,
- RunnerApi.WireCoderSetting wireCoderSetting) {
+ WireCoderSetting wireCoderSetting) {
String elementCoderId = pCollectionNode.getPCollection().getCoderId();
String windowingStrategyId = pCollectionNode.getPCollection().getWindowingStrategyId();
String windowCoderId =
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProviderTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProviderTest.java
index 532e904..85763a7 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProviderTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/GrpcContextHeaderAccessorProviderTest.java
@@ -23,16 +23,16 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements;
import org.apache.beam.model.fnexecution.v1.BeamFnDataGrpc;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.CallOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Channel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ClientCall;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ClientInterceptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Metadata;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.MethodDescriptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.CallOptions;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Channel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ClientCall;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ClientInterceptor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Metadata;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.MethodDescriptor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.junit.Assert;
import org.junit.Test;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/ServerFactoryTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/ServerFactoryTest.java
index 0972d7b..2dbed76 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/ServerFactoryTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/ServerFactoryTest.java
@@ -42,11 +42,11 @@
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.sdk.fn.channel.ManagedChannelFactory;
import org.apache.beam.sdk.fn.test.TestStreams;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.Epoll;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.Epoll;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.net.HostAndPort;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactServicesTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactServicesTest.java
index 9585530..2479da0 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactServicesTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/BeamFileSystemArtifactServicesTest.java
@@ -58,10 +58,10 @@
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.InProcessServerFactory;
import org.apache.beam.sdk.io.FileSystems;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ClassLoaderArtifactServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ClassLoaderArtifactServiceTest.java
index 65d54a9..849da70 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ClassLoaderArtifactServiceTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ClassLoaderArtifactServiceTest.java
@@ -42,10 +42,10 @@
import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc;
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.InProcessServerFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.junit.Assert;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactoryTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactoryTest.java
index b5ac3c6..1d53d35 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactoryTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DefaultJobBundleFactoryTest.java
@@ -62,8 +62,8 @@
import org.apache.beam.sdk.options.ExperimentalOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.PortablePipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.junit.Assert;
import org.junit.Before;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolServiceTest.java
index cb65a0e..0af2602 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolServiceTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientPoolServiceTest.java
@@ -35,8 +35,8 @@
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.InProcessServerFactory;
import org.apache.beam.sdk.util.MoreFutures;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientTest.java
index 341b53c..63fbe55 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/FnApiControlClientTest.java
@@ -31,7 +31,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse;
import org.apache.beam.sdk.util.MoreFutures;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
index d9d51d4..365f679 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
@@ -111,7 +111,7 @@
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Collections2;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
@@ -1016,7 +1016,9 @@
@TimerId("event") Timer eventTimeTimer,
@TimerId("processing") Timer processingTimeTimer) {
context.output(KV.of("main" + context.element().getKey(), ""));
- eventTimeTimer.set(context.timestamp().plus(1L));
+ eventTimeTimer
+ .withOutputTimestamp(context.timestamp())
+ .set(context.timestamp().plus(1L));
processingTimeTimer.offset(Duration.millis(2L));
processingTimeTimer.setRelative();
}
@@ -1027,7 +1029,9 @@
@TimerId("event") Timer eventTimeTimer,
@TimerId("processing") Timer processingTimeTimer) {
context.output(KV.of("event", ""));
- eventTimeTimer.set(context.timestamp().plus(11L));
+ eventTimeTimer
+ .withOutputTimestamp(context.timestamp())
+ .set(context.timestamp().plus(11L));
processingTimeTimer.offset(Duration.millis(12L));
processingTimeTimer.setRelative();
}
@@ -1038,7 +1042,9 @@
@TimerId("event") Timer eventTimeTimer,
@TimerId("processing") Timer processingTimeTimer) {
context.output(KV.of("processing", ""));
- eventTimeTimer.set(context.timestamp().plus(21L));
+ eventTimeTimer
+ .withOutputTimestamp(context.timestamp())
+ .set(context.timestamp().plus(21L));
processingTimeTimer.offset(Duration.millis(22L));
processingTimeTimer.setRelative();
}
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java
index be08b58..adf843d 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/data/GrpcDataServiceTest.java
@@ -47,10 +47,10 @@
import org.apache.beam.sdk.fn.test.TestStreams;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobServiceTest.java
index e7b01af..2a9b366 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobServiceTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/InMemoryJobServiceTest.java
@@ -17,34 +17,39 @@
*/
package org.apache.beam.runners.fnexecution.jobsubmission;
+import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.collection.IsCollectionWithSize.hasSize;
import static org.hamcrest.core.Is.is;
import static org.hamcrest.core.Is.isA;
import static org.hamcrest.core.IsNull.notNullValue;
-import static org.junit.Assert.assertThat;
import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import java.util.ArrayList;
import java.util.List;
+import java.util.UUID;
+import java.util.function.Consumer;
import org.apache.beam.model.jobmanagement.v1.JobApi;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
/** Tests for {@link InMemoryJobService}. */
@RunWith(JUnit4.class)
public class InMemoryJobServiceTest {
+
private static final String TEST_JOB_NAME = "test-job";
private static final String TEST_JOB_ID = "test-job-id";
private static final String TEST_RETRIEVAL_TOKEN = "test-staging-token";
@@ -57,6 +62,8 @@
.setPipelineOptions(TEST_OPTIONS)
.build();
+ private final int maxInvocationHistory = 3;
+
Endpoints.ApiServiceDescriptor stagingServiceDescriptor;
@Mock JobInvoker invoker;
@Mock JobInvocation invocation;
@@ -186,7 +193,79 @@
verify(invocation, times(1)).start();
}
+ @Test
+ public void testInvocationCleanup() {
+ final int maxInvocationHistory = 3;
+ service =
+ InMemoryJobService.create(
+ stagingServiceDescriptor, session -> "token", null, invoker, maxInvocationHistory);
+
+ assertThat(getNumberOfInvocations(), is(0));
+
+ Job job1 = runJob();
+ assertThat(getNumberOfInvocations(), is(1));
+ Job job2 = runJob();
+ assertThat(getNumberOfInvocations(), is(2));
+ Job job3 = runJob();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory));
+
+ // All running invocations must be available and never be discarded
+ // even if they exceed the max history size
+ Job job4 = runJob();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory + 1));
+
+ // We need to have more than maxInvocationHistory completed jobs for the cleanup to trigger
+ job1.finish();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory + 1));
+ job2.finish();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory + 1));
+ job3.finish();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory + 1));
+
+ // The fourth finished job exceeds maxInvocationHistory and triggers the cleanup
+ job4.finish();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory));
+
+ // Run a new job after the cleanup
+ Job job5 = runJob();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory + 1));
+ job5.finish();
+ assertThat(getNumberOfInvocations(), is(maxInvocationHistory));
+ }
+
+ private Job runJob() {
+ when(invocation.getId()).thenReturn(UUID.randomUUID().toString());
+ prepareAndRunJob();
+ // Retrieve the state listener for this invocation
+ ArgumentCaptor<Consumer<JobApi.JobStateEvent>> stateListener =
+ ArgumentCaptor.forClass(Consumer.class);
+ verify(invocation, atLeastOnce()).addStateListener(stateListener.capture());
+ return new Job(stateListener.getValue());
+ }
+
+ private int getNumberOfInvocations() {
+ RecordingObserver<JobApi.GetJobsResponse> recorder = new RecordingObserver<>();
+ final JobApi.GetJobsRequest getJobsRequest = JobApi.GetJobsRequest.newBuilder().build();
+ service.getJobs(getJobsRequest, recorder);
+ return recorder.getValue().getJobInfoCount();
+ }
+
+ private static class Job {
+ private Consumer<JobApi.JobStateEvent> stateListener;
+
+ private Job(Consumer<JobApi.JobStateEvent> stateListener) {
+ this.stateListener = stateListener;
+ }
+
+ void finish() {
+ JobApi.JobStateEvent terminalEvent =
+ JobApi.JobStateEvent.newBuilder().setState(JobApi.JobState.Enum.DONE).build();
+ stateListener.accept(terminalEvent);
+ }
+ }
+
private static class RecordingObserver<T> implements StreamObserver<T> {
+
ArrayList<T> values = new ArrayList<>();
Throwable error = null;
boolean isCompleted = false;
@@ -206,6 +285,11 @@
isCompleted = true;
}
+ T getValue() {
+ assert values.size() == 1;
+ return values.get(0);
+ }
+
boolean isSuccessful() {
return isCompleted && error == null;
}
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocationTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocationTest.java
index 30e34d4..f9c256f 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocationTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/jobsubmission/JobInvocationTest.java
@@ -33,7 +33,7 @@
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.metrics.MetricResults;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ListeningExecutorService;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.MoreExecutors;
import org.joda.time.Duration;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingServiceTest.java
index 3bfda79..39a5e55 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingServiceTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingServiceTest.java
@@ -37,9 +37,9 @@
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.InProcessServerFactory;
import org.apache.beam.sdk.fn.test.TestStreams;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionServiceTest.java
index 850a070..a0fd1b7 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionServiceTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/provisioning/StaticGrpcProvisionServiceTest.java
@@ -31,11 +31,11 @@
import org.apache.beam.model.fnexecution.v1.ProvisionServiceGrpc.ProvisionServiceBlockingStub;
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.InProcessServerFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ListValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.NullValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Value;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ListValue;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.NullValue;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Value;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/state/GrpcStateServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/state/GrpcStateServiceTest.java
index f8b3f29..aa986c9 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/state/GrpcStateServiceTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/state/GrpcStateServiceTest.java
@@ -31,8 +31,8 @@
import java.util.concurrent.TimeUnit;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.sdk.fn.test.TestStreams;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/status/BeamWorkerStatusGrpcServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/status/BeamWorkerStatusGrpcServiceTest.java
new file mode 100644
index 0000000..9b04914
--- /dev/null
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/status/BeamWorkerStatusGrpcServiceTest.java
@@ -0,0 +1,192 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.fnexecution.status;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import org.apache.beam.fn.harness.control.AddHarnessIdInterceptor;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusResponse;
+import org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc;
+import org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc.BeamFnWorkerStatusStub;
+import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
+import org.apache.beam.runners.fnexecution.GrpcContextHeaderAccessorProvider;
+import org.apache.beam.runners.fnexecution.GrpcFnServer;
+import org.apache.beam.runners.fnexecution.InProcessServerFactory;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.testing.GrpcCleanupRule;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(JUnit4.class)
+public class BeamWorkerStatusGrpcServiceTest {
+
+ @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+ private static final String ID = "id";
+ private BeamWorkerStatusGrpcService service;
+ private GrpcFnServer<BeamWorkerStatusGrpcService> server;
+ private ManagedChannel channel;
+ private BeamFnWorkerStatusStub stub;
+ @Mock private StreamObserver<WorkerStatusRequest> mockObserver;
+
+ @Before
+ public void setUp() throws Exception {
+ MockitoAnnotations.initMocks(this);
+ service =
+ BeamWorkerStatusGrpcService.create(
+ ApiServiceDescriptor.newBuilder().setUrl(UUID.randomUUID().toString()).build(),
+ GrpcContextHeaderAccessorProvider.getHeaderAccessor());
+ server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create());
+ channel = InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).build();
+ stub =
+ BeamFnWorkerStatusGrpc.newStub(channel)
+ .withInterceptors(AddHarnessIdInterceptor.create(ID));
+ grpcCleanup.register(server.getServer());
+ grpcCleanup.register(channel);
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ if (service != null) {
+ service.close();
+ }
+ }
+
+ @Test
+ public void testClientConnected() throws Exception {
+ StreamObserver<WorkerStatusResponse> workerStatusResponseStreamObserver =
+ stub.workerStatus(mockObserver);
+ WorkerStatusClient client = waitAndGetStatusClient(ID);
+ assertNotNull(client);
+ }
+
+ @Test
+ public void testGetWorkerStatusNoResponse() throws Exception {
+ StreamObserver<WorkerStatusResponse> unused = stub.workerStatus(mockObserver);
+ waitAndGetStatusClient(ID);
+ String response = service.getSingleWorkerStatus("id", 1, TimeUnit.MILLISECONDS);
+ assertEquals(
+ "Error: exception encountered getting status from SDK harness: java.util.concurrent.TimeoutException",
+ response);
+ }
+
+ @Test
+ public void testGetWorkerStatusSuccess() throws Exception {
+ StreamObserver<WorkerStatusResponse> observer = stub.workerStatus(mockObserver);
+ waitAndGetStatusClient(ID);
+ doAnswer(
+ (invocation) -> {
+ WorkerStatusRequest request = (WorkerStatusRequest) invocation.getArguments()[0];
+ observer.onNext(
+ WorkerStatusResponse.newBuilder()
+ .setId(request.getId())
+ .setStatusInfo("status")
+ .build());
+ return null;
+ })
+ .when(mockObserver)
+ .onNext(any());
+
+ CompletableFuture<String> future = service.getWorkerStatus(ID);
+ String response = future.get(5, TimeUnit.SECONDS);
+ assertEquals("status", response);
+ }
+
+ @Test
+ public void testGetWorkerStatusReturnError() throws Exception {
+ StreamObserver<WorkerStatusResponse> observer = stub.workerStatus(mockObserver);
+ waitAndGetStatusClient(ID);
+ doAnswer(
+ (invocation) -> {
+ WorkerStatusRequest request = (WorkerStatusRequest) invocation.getArguments()[0];
+ observer.onNext(
+ WorkerStatusResponse.newBuilder()
+ .setId(request.getId())
+ .setError("error")
+ .build());
+ return null;
+ })
+ .when(mockObserver)
+ .onNext(any());
+
+ CompletableFuture<String> future = service.getWorkerStatus(ID);
+ String response = future.get(5, TimeUnit.SECONDS);
+ assertEquals("error", response);
+ }
+
+ @Test
+ public void testGetAllWorkerStatuses() throws Exception {
+ Set<String> ids = Sets.newHashSet("id0", "id3", "id11", "id12", "id21");
+ for (String id : ids) {
+ StreamObserver<WorkerStatusRequest> requestObserverMock = mock(StreamObserver.class);
+ BeamFnWorkerStatusStub workerStatusStub =
+ BeamFnWorkerStatusGrpc.newStub(channel)
+ .withInterceptors(AddHarnessIdInterceptor.create(id));
+ StreamObserver<WorkerStatusResponse> observer =
+ workerStatusStub.workerStatus(requestObserverMock);
+ // wait for connection before proceeding to avoid race condition.
+ waitAndGetStatusClient(id);
+ doAnswer(
+ (invocation) -> {
+ WorkerStatusRequest request = (WorkerStatusRequest) invocation.getArguments()[0];
+ observer.onNext(
+ WorkerStatusResponse.newBuilder()
+ .setId(request.getId())
+ .setStatusInfo("status")
+ .build());
+ return null;
+ })
+ .when(requestObserverMock)
+ .onNext(any());
+ }
+ Map<String, String> allWorkerStatuses = service.getAllWorkerStatuses(5, TimeUnit.SECONDS);
+
+ assertEquals(ids, allWorkerStatuses.keySet());
+
+ for (String id : ids) {
+ assertEquals("status", allWorkerStatuses.get(id));
+ }
+ }
+
+ private WorkerStatusClient waitAndGetStatusClient(String id)
+ throws InterruptedException, ExecutionException, TimeoutException {
+ // wait for status client connection, and get the corresponding client.
+ CompletableFuture<WorkerStatusClient> clientFuture = service.getStatusClient(id);
+ return clientFuture.get(1, TimeUnit.SECONDS);
+ }
+}
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/status/WorkerStatusClientTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/status/WorkerStatusClientTest.java
new file mode 100644
index 0000000..5fa143e
--- /dev/null
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/status/WorkerStatusClientTest.java
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.fnexecution.status;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.verify;
+
+import java.io.IOException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusResponse;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(JUnit4.class)
+public class WorkerStatusClientTest {
+
+ @Mock public StreamObserver<BeamFnApi.WorkerStatusRequest> mockObserver;
+ private WorkerStatusClient client;
+
+ @Before
+ public void setup() {
+ MockitoAnnotations.initMocks(this);
+ client = WorkerStatusClient.forRequestObserver("ID", mockObserver);
+ }
+
+ @Test
+ public void testGetWorkerStatusSuccess() throws Exception {
+ CompletableFuture<WorkerStatusResponse> workerStatus =
+ client.getWorkerStatus(WorkerStatusRequest.newBuilder().setId("123").build());
+ client
+ .getResponseObserver()
+ .onNext(WorkerStatusResponse.newBuilder().setId("123").setStatusInfo("status").build());
+ Assert.assertEquals("status", workerStatus.get().getStatusInfo());
+ }
+
+ @Test
+ public void testGetWorkerStatusError() throws Exception {
+ CompletableFuture<WorkerStatusResponse> workerStatus =
+ client.getWorkerStatus(WorkerStatusRequest.newBuilder().setId("123").build());
+ client
+ .getResponseObserver()
+ .onNext(WorkerStatusResponse.newBuilder().setId("123").setError("error").build());
+ Assert.assertEquals("error", workerStatus.get().getError());
+ }
+
+ @Test
+ public void testGetWorkerStatusRequestSent() {
+ CompletableFuture<WorkerStatusResponse> workerStatus = client.getWorkerStatus();
+ verify(mockObserver).onNext(any(WorkerStatusRequest.class));
+ }
+
+ @Test
+ public void testUnknownRequestIdResponseIgnored() {
+ CompletableFuture<WorkerStatusResponse> workerStatus = client.getWorkerStatus();
+ client
+ .getResponseObserver()
+ .onNext(WorkerStatusResponse.newBuilder().setId("unknown").setStatusInfo("status").build());
+ Assert.assertFalse(workerStatus.isDone());
+ }
+
+ @Test
+ public void testCloseOutstandingRequest() throws IOException {
+ CompletableFuture<WorkerStatusResponse> workerStatus = client.getWorkerStatus();
+ client.close();
+ Assert.assertThrows(ExecutionException.class, workerStatus::get);
+ }
+}
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/translation/BatchSideInputHandlerFactoryTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/translation/BatchSideInputHandlerFactoryTest.java
index e1ccad6..1f45945 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/translation/BatchSideInputHandlerFactoryTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/translation/BatchSideInputHandlerFactoryTest.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.runners.fnexecution.translation;
-import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTING;
+import static org.apache.beam.runners.core.construction.graph.ExecutableStage.DEFAULT_WIRE_CODER_SETTINGS;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -236,6 +236,6 @@
Collections.emptyList(),
Collections.emptyList(),
Collections.emptyList(),
- DEFAULT_WIRE_CODER_SETTING);
+ DEFAULT_WIRE_CODER_SETTINGS);
}
}
diff --git a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java
index e291117..76e8375 100644
--- a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java
+++ b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java
@@ -92,7 +92,13 @@
TimerInternals.TimerData timer, DoFnRunner<KV<?, ?>, ?> doFnRunner) {
StateNamespace namespace = timer.getNamespace();
BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
- doFnRunner.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ doFnRunner.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
}
@Override
diff --git a/runners/portability/java/build.gradle b/runners/portability/java/build.gradle
index 5425b8f..9c529a6 100644
--- a/runners/portability/java/build.gradle
+++ b/runners/portability/java/build.gradle
@@ -33,7 +33,7 @@
compile library.java.hamcrest_library
compile project(":runners:java-fn-execution")
compile project(path: ":sdks:java:harness", configuration: "shadow")
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
compile library.java.slf4j_api
testCompile project(path: ":runners:core-construction-java", configuration: "testRuntime")
testCompile library.java.hamcrest_core
diff --git a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/ExternalWorkerService.java b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/ExternalWorkerService.java
index 028a934..363b013 100644
--- a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/ExternalWorkerService.java
+++ b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/ExternalWorkerService.java
@@ -25,7 +25,7 @@
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.ServerFactory;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java
index bcfa321..820e93b 100644
--- a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java
+++ b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/JobServicePipelineResult.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.runners.portability;
+import java.util.Iterator;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
@@ -26,24 +27,27 @@
import org.apache.beam.model.jobmanagement.v1.JobApi.CancelJobRequest;
import org.apache.beam.model.jobmanagement.v1.JobApi.CancelJobResponse;
import org.apache.beam.model.jobmanagement.v1.JobApi.GetJobStateRequest;
+import org.apache.beam.model.jobmanagement.v1.JobApi.JobMessage;
+import org.apache.beam.model.jobmanagement.v1.JobApi.JobMessagesRequest;
+import org.apache.beam.model.jobmanagement.v1.JobApi.JobMessagesResponse;
import org.apache.beam.model.jobmanagement.v1.JobApi.JobStateEvent;
import org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.metrics.MetricResults;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
class JobServicePipelineResult implements PipelineResult, AutoCloseable {
- private static final long POLL_INTERVAL_MS = 10 * 1000;
+ private static final long POLL_INTERVAL_MS = 3_000;
private static final Logger LOG = LoggerFactory.getLogger(JobServicePipelineResult.class);
private final ByteString jobId;
private final CloseableResource<JobServiceBlockingStub> jobService;
- @Nullable private State terminationState;
+ @Nullable private State terminalState;
@Nullable private final Runnable cleanup;
private org.apache.beam.model.jobmanagement.v1.JobApi.MetricResults jobMetrics;
@@ -51,14 +55,14 @@
ByteString jobId, CloseableResource<JobServiceBlockingStub> jobService, Runnable cleanup) {
this.jobId = jobId;
this.jobService = jobService;
- this.terminationState = null;
+ this.terminalState = null;
this.cleanup = cleanup;
}
@Override
public State getState() {
- if (terminationState != null) {
- return terminationState;
+ if (terminalState != null) {
+ return terminalState;
}
JobServiceBlockingStub stub = jobService.get();
JobStateEvent response =
@@ -98,26 +102,16 @@
@Override
public State waitUntilFinish() {
- if (terminationState != null) {
- return terminationState;
+ if (terminalState != null) {
+ return terminalState;
}
- JobServiceBlockingStub stub = jobService.get();
- GetJobStateRequest request = GetJobStateRequest.newBuilder().setJobIdBytes(jobId).build();
- JobStateEvent response = stub.getState(request);
- State lastState = getJavaState(response.getState());
- while (!lastState.isTerminal()) {
- try {
- Thread.sleep(POLL_INTERVAL_MS);
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- throw new RuntimeException(e);
- }
- response = stub.getState(request);
- lastState = getJavaState(response.getState());
+ try {
+ waitForTerminalState();
+ propagateErrors();
+ return terminalState;
+ } finally {
+ close();
}
- close();
- terminationState = lastState;
- return lastState;
}
@Override
@@ -139,6 +133,41 @@
}
}
+ private void waitForTerminalState() {
+ JobServiceBlockingStub stub = jobService.get();
+ GetJobStateRequest request = GetJobStateRequest.newBuilder().setJobIdBytes(jobId).build();
+ JobStateEvent response = stub.getState(request);
+ State lastState = getJavaState(response.getState());
+ while (!lastState.isTerminal()) {
+ try {
+ Thread.sleep(POLL_INTERVAL_MS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ response = stub.getState(request);
+ lastState = getJavaState(response.getState());
+ }
+ terminalState = lastState;
+ }
+
+ private void propagateErrors() {
+ if (terminalState != State.DONE) {
+ JobMessagesRequest messageStreamRequest =
+ JobMessagesRequest.newBuilder().setJobIdBytes(jobId).build();
+ Iterator<JobMessagesResponse> messageStreamIterator =
+ jobService.get().getMessageStream(messageStreamRequest);
+ while (messageStreamIterator.hasNext()) {
+ JobMessage messageResponse = messageStreamIterator.next().getMessageResponse();
+ if (messageResponse.getImportance() == JobMessage.MessageImportance.JOB_MESSAGE_ERROR) {
+ throw new RuntimeException(
+ "The Runner experienced the following error during execution:\n"
+ + messageResponse.getMessageText());
+ }
+ }
+ }
+ }
+
private static State getJavaState(JobApi.JobState.Enum protoState) {
switch (protoState) {
case UNSPECIFIED:
diff --git a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/PortableRunner.java b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/PortableRunner.java
index 72c3cbc..31c7be0 100644
--- a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/PortableRunner.java
+++ b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/PortableRunner.java
@@ -53,8 +53,8 @@
import org.apache.beam.sdk.options.PipelineOptionsValidator;
import org.apache.beam.sdk.options.PortablePipelineOptions;
import org.apache.beam.sdk.util.ZipFiles;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
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.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
diff --git a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/testing/TestJobService.java b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/testing/TestJobService.java
index 76b8974..53578b7 100644
--- a/runners/portability/java/src/main/java/org/apache/beam/runners/portability/testing/TestJobService.java
+++ b/runners/portability/java/src/main/java/org/apache/beam/runners/portability/testing/TestJobService.java
@@ -27,7 +27,7 @@
import org.apache.beam.model.jobmanagement.v1.JobApi.RunJobResponse;
import org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceImplBase;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* A JobService for tests.
diff --git a/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java b/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java
index 40f7900..9f011f3 100644
--- a/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java
+++ b/runners/portability/java/src/test/java/org/apache/beam/runners/portability/PortableRunnerTest.java
@@ -38,9 +38,9 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.PortablePipelineOptions;
import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Timestamp;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Timestamp;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/runners/samza/build.gradle b/runners/samza/build.gradle
index 18dbea4..469c961 100644
--- a/runners/samza/build.gradle
+++ b/runners/samza/build.gradle
@@ -46,7 +46,7 @@
compile library.java.slf4j_api
compile library.java.joda_time
compile library.java.commons_compress
- compile library.java.commons_io_2x
+ compile library.java.commons_io
compile library.java.args4j
compile "org.apache.samza:samza-api:$samza_version"
compile "org.apache.samza:samza-core_2.11:$samza_version"
@@ -58,7 +58,6 @@
compile "org.apache.kafka:kafka-clients:0.11.0.2"
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
testCompile project(path: ":runners:core-java", configuration: "testRuntime")
- testCompile library.java.commons_lang3
testCompile library.java.hamcrest_core
testCompile library.java.junit
testCompile library.java.mockito_core
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaJobServerDriver.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaJobServerDriver.java
index f21d666..b820687 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaJobServerDriver.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/SamzaJobServerDriver.java
@@ -32,7 +32,7 @@
import org.apache.beam.runners.fnexecution.jobsubmission.JobInvoker;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ListeningExecutorService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -101,7 +101,8 @@
}
},
stagingSessionToken -> {},
- jobInvoker);
+ jobInvoker,
+ InMemoryJobService.DEFAULT_MAX_INVOCATION_HISTORY);
}
public void run() throws Exception {
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/UnboundedSourceSystem.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/UnboundedSourceSystem.java
index b118a3c..c77ddb1 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/UnboundedSourceSystem.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/adapter/UnboundedSourceSystem.java
@@ -34,6 +34,7 @@
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.core.serialization.Base64Serializer;
import org.apache.beam.runners.samza.SamzaPipelineOptions;
@@ -48,7 +49,6 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
-import org.apache.commons.lang3.StringUtils;
import org.apache.samza.Partition;
import org.apache.samza.SamzaException;
import org.apache.samza.config.Config;
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java
index 101ee80..aefcf6d 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java
@@ -57,8 +57,16 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
- withMetrics(() -> underlying.onTimer(timerId, window, timestamp, timeDomain));
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
+ withMetrics(
+ () ->
+ underlying.onTimer(
+ timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain));
}
@Override
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java
index 795b8c0..cd140c7 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java
@@ -452,7 +452,13 @@
// Need to pass in the keyed TimerData here
((DoFnRunnerWithKeyedInternals) fnRunner).onTimer(keyedTimerData, window);
} else {
- pushbackFnRunner.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ pushbackFnRunner.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
}
}
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java
index 6fb2bd3..3b2d1cb 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java
@@ -62,7 +62,13 @@
try {
final TimerInternals.TimerData timer = keyedTimerData.getTimerData();
- onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
} finally {
clearKeyedInternals();
}
@@ -70,10 +76,15 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {
checkState(keyedInternals.getKey() != null, "Key is not set for timer");
- underlying.onTimer(timerId, window, timestamp, timeDomain);
+ underlying.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
}
@Override
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/KeyedInternals.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/KeyedInternals.java
index 330fb24..7673bd3 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/KeyedInternals.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/KeyedInternals.java
@@ -139,8 +139,8 @@
}
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
- getInternals().deleteTimer(namespace, timerId);
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
+ getInternals().deleteTimer(namespace, timerId, timerFamilyId);
}
@Override
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java
index 3b1b938..99439a2 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java
@@ -262,7 +262,12 @@
@Override
public void onTimer(
- String timerId, BoundedWindow window, Instant timestamp, TimeDomain timeDomain) {}
+ String timerId,
+ String timerFamilyId,
+ BoundedWindow window,
+ Instant timestamp,
+ Instant outputTimestamp,
+ TimeDomain timeDomain) {}
@Override
public void finishBundle() {
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaTimerInternalsFactory.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaTimerInternalsFactory.java
index 9ac082b..22c2ac9 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaTimerInternalsFactory.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaTimerInternalsFactory.java
@@ -248,7 +248,7 @@
}
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
deleteTimer(TimerData.of(timerId, namespace, Instant.now(), TimeDomain.EVENT_TIME));
}
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 5ddcf09..b3ee221 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
@@ -28,6 +28,7 @@
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.core.serialization.Base64Serializer;
import org.apache.beam.runners.samza.SamzaExecutionEnvironment;
@@ -36,7 +37,6 @@
import org.apache.beam.runners.samza.runtime.SamzaStoreStateInternals;
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.collect.ImmutableMap;
-import org.apache.commons.lang3.StringUtils;
import org.apache.samza.config.ApplicationConfig;
import org.apache.samza.config.Config;
import org.apache.samza.config.ConfigFactory;
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/WindowAssignTranslator.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/WindowAssignTranslator.java
index 95c7328..114a256 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/WindowAssignTranslator.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/translation/WindowAssignTranslator.java
@@ -28,7 +28,7 @@
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.samza.operators.MessageStream;
/**
diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/util/SamzaPipelineTranslatorUtils.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/util/SamzaPipelineTranslatorUtils.java
index 758515a4..b1bfc83 100644
--- a/runners/samza/src/main/java/org/apache/beam/runners/samza/util/SamzaPipelineTranslatorUtils.java
+++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/util/SamzaPipelineTranslatorUtils.java
@@ -28,7 +28,7 @@
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
/** Utilities for pipeline translation. */
diff --git a/runners/spark/build.gradle b/runners/spark/build.gradle
index 3de3bc1..495b1d9 100644
--- a/runners/spark/build.gradle
+++ b/runners/spark/build.gradle
@@ -70,14 +70,12 @@
provided library.java.spark_streaming
provided library.java.spark_network_common
provided library.java.hadoop_common
- provided library.java.commons_lang3
- provided library.java.commons_io_2x
+ provided library.java.commons_io
provided library.java.hamcrest_core
provided library.java.hamcrest_library
provided "com.esotericsoftware.kryo:kryo:2.21"
runtimeOnly library.java.jackson_module_scala
runtimeOnly "org.scala-lang:scala-library:2.11.8"
- compile "org.scala-lang.modules:scala-java8-compat_2.11:0.9.0"
testCompile project(":sdks:java:io:kafka")
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
// SparkStateInternalsTest extends abstract StateInternalsTest
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobInvoker.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobInvoker.java
index 2ea261f..956e9bc 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobInvoker.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkJobInvoker.java
@@ -27,7 +27,7 @@
import org.apache.beam.runners.fnexecution.jobsubmission.PortablePipelineRunner;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.sdk.options.PortablePipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.ListeningExecutorService;
import org.slf4j.Logger;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java
index b85a71e..0b91c4f 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java
@@ -21,6 +21,7 @@
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.List;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils;
import org.apache.beam.runners.spark.translation.EvaluationContext;
import org.apache.beam.runners.spark.translation.SparkPipelineTranslator;
import org.apache.beam.runners.spark.translation.TransformEvaluator;
@@ -32,7 +33,6 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Joiner;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
-import org.apache.commons.lang3.StringUtils;
/**
* Pipeline visitor for translating a Beam pipeline into equivalent Spark operations. Used for
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java
index d0c1c0f..e5166a3 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkPipelineRunner.java
@@ -49,7 +49,7 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.PortablePipelineOptions;
import org.apache.beam.sdk.options.PortablePipelineOptions.RetrievalServiceType;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.spark.api.java.JavaSparkContext;
import org.kohsuke.args4j.CmdLineException;
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/MetricsAccumulator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/MetricsAccumulator.java
index a18b92d..8313d73 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/MetricsAccumulator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/metrics/MetricsAccumulator.java
@@ -37,7 +37,7 @@
* For resilience, {@link AccumulatorV2 Accumulators} are required to be wrapped in a Singleton.
*
* @see <a
- * href="https://spark.apache.org/docs/latest/api/java/org/apache/spark/util/AccumulatorV2.html">accumulatorsV2</a>
+ * href="https://spark.apache.org/docs/2.4.4/streaming-programming-guide.html#accumulators-broadcast-variables-and-checkpoints">accumulatorsV2</a>
*/
public class MetricsAccumulator {
private static final Logger LOG = LoggerFactory.getLogger(MetricsAccumulator.class);
@@ -74,6 +74,8 @@
}
}
LOG.info("Instantiated metrics accumulator: " + instance.value());
+ } else {
+ instance.reset();
}
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
index 1d5b36b..58fbee9 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
@@ -165,7 +165,7 @@
}
@Override
- public void deleteTimer(StateNamespace namespace, String timerId) {
+ public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) {
throw new UnsupportedOperationException("Deleting a timer by ID is not yet supported.");
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/MetricsAccumulator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/MetricsAccumulator.java
index ac44bf8..e182cf0 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/MetricsAccumulator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/metrics/MetricsAccumulator.java
@@ -28,7 +28,7 @@
* For resilience, {@link AccumulatorV2 Accumulators} are required to be wrapped in a Singleton.
*
* @see <a
- * href="https://spark.apache.org/docs/latest/api/java/org/apache/spark/util/AccumulatorV2.html">accumulatorsV2</a>
+ * href="https://spark.apache.org/docs/2.4.4/streaming-programming-guide.html#accumulators-broadcast-variables-and-checkpoints">accumulatorsV2</a>
*/
public class MetricsAccumulator {
private static final Logger LOG = LoggerFactory.getLogger(MetricsAccumulator.class);
@@ -51,6 +51,8 @@
}
}
LOG.info("Instantiated metrics accumulator: " + instance.value());
+ } else {
+ instance.reset();
}
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SchemaHelpers.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SchemaHelpers.java
index 13fbfb8..b778c46 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SchemaHelpers.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/SchemaHelpers.java
@@ -24,14 +24,16 @@
/** A {@link SchemaHelpers} for the Spark Batch Runner. */
public class SchemaHelpers {
+ private static final StructType BINARY_SCHEMA =
+ new StructType(
+ new StructField[] {
+ StructField.apply("binaryStructField", DataTypes.BinaryType, true, Metadata.empty())
+ });
+
public static StructType binarySchema() {
// we use a binary schema for now because:
// using a empty schema raises a indexOutOfBoundsException
// using a NullType schema stores null in the elements
- StructField[] array = new StructField[1];
- StructField binaryStructField =
- StructField.apply("binaryStructField", DataTypes.BinaryType, true, Metadata.empty());
- array[0] = binaryStructField;
- return new StructType(array);
+ return BINARY_SCHEMA;
}
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java
index 46dc282..55d97ba 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java
@@ -71,11 +71,13 @@
@Override
public void onTimer(
final String timerId,
+ final String timerFamilyId,
final BoundedWindow window,
final Instant timestamp,
+ final Instant outputTimestamp,
final TimeDomain timeDomain) {
try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metricsContainer())) {
- delegate.onTimer(timerId, window, timestamp, timeDomain);
+ delegate.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
} catch (IOException e) {
throw new RuntimeException(e);
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
index eca87bd..b77f94c 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpers.java
@@ -21,8 +21,10 @@
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
import java.util.Objects;
+import javax.annotation.Nullable;
import org.apache.beam.runners.spark.structuredstreaming.translation.SchemaHelpers;
import org.apache.beam.sdk.coders.Coder;
import org.apache.spark.sql.Encoder;
@@ -47,67 +49,26 @@
/** {@link Encoders} utility class. */
public class EncoderHelpers {
-
- /*
- --------- Bridges from Beam Coders to Spark Encoders
- */
-
/**
* Wrap a Beam coder into a Spark Encoder using Catalyst Expression Encoders (which uses java code
* generation).
*/
- public static <T> Encoder<T> fromBeamCoder(Coder<T> beamCoder) {
+ public static <T> Encoder<T> fromBeamCoder(Coder<T> coder) {
+ Class<? super T> clazz = coder.getEncodedTypeDescriptor().getRawType();
+ ClassTag<T> classTag = ClassTag$.MODULE$.apply(clazz);
+ List<Expression> serializers =
+ Collections.singletonList(
+ new EncodeUsingBeamCoder<>(new BoundReference(0, new ObjectType(clazz), true), coder));
- List<Expression> serialiserList = new ArrayList<>();
- Class<? super T> claz = beamCoder.getEncodedTypeDescriptor().getRawType();
- BeamCoderWrapper<T> beamCoderWrapper = new BeamCoderWrapper<>(beamCoder);
-
- serialiserList.add(
- new EncodeUsingBeamCoder<>(
- new BoundReference(0, new ObjectType(claz), true), beamCoderWrapper));
- ClassTag<T> classTag = ClassTag$.MODULE$.apply(claz);
return new ExpressionEncoder<>(
SchemaHelpers.binarySchema(),
false,
- JavaConversions.collectionAsScalaIterable(serialiserList).toSeq(),
+ JavaConversions.collectionAsScalaIterable(serializers).toSeq(),
new DecodeUsingBeamCoder<>(
- new Cast(new GetColumnByOrdinal(0, BinaryType), BinaryType),
- classTag,
- beamCoderWrapper),
+ new Cast(new GetColumnByOrdinal(0, BinaryType), BinaryType), classTag, coder),
classTag);
}
- public static class BeamCoderWrapper<T> implements Serializable {
-
- private Coder<T> beamCoder;
-
- public BeamCoderWrapper(Coder<T> beamCoder) {
- this.beamCoder = beamCoder;
- }
-
- public byte[] encode(boolean isInputNull, T inputValue) {
- if (isInputNull) {
- return null;
- } else {
- java.io.ByteArrayOutputStream baos = new java.io.ByteArrayOutputStream();
- try {
- beamCoder.encode(inputValue, baos);
- } catch (Exception e) {
- throw org.apache.beam.sdk.util.UserCodeException.wrap(e);
- }
- return baos.toByteArray();
- }
- }
-
- public T decode(boolean isInputNull, byte[] inputValue) {
- try {
- return isInputNull ? null : beamCoder.decode(new java.io.ByteArrayInputStream(inputValue));
- } catch (Exception e) {
- throw org.apache.beam.sdk.util.UserCodeException.wrap(e);
- }
- }
- }
-
/**
* Catalyst Expression that serializes elements using Beam {@link Coder}.
*
@@ -116,12 +77,12 @@
public static class EncodeUsingBeamCoder<T> extends UnaryExpression
implements NonSQLExpression, Serializable {
- private Expression child;
- private BeamCoderWrapper<T> beamCoderWrapper;
+ private final Expression child;
+ private final Coder<T> coder;
- public EncodeUsingBeamCoder(Expression child, BeamCoderWrapper<T> beamCoderWrapper) {
+ public EncodeUsingBeamCoder(Expression child, Coder<T> coder) {
this.child = child;
- this.beamCoderWrapper = beamCoderWrapper;
+ this.coder = coder;
}
@Override
@@ -131,28 +92,27 @@
@Override
public ExprCode doGenCode(CodegenContext ctx, ExprCode ev) {
- String accessCode =
- ctx.addReferenceObj(
- "beamCoderWrapper", beamCoderWrapper, BeamCoderWrapper.class.getName());
+ String accessCode = ctx.addReferenceObj("coder", coder, coder.getClass().getName());
ExprCode input = child.genCode(ctx);
String javaType = CodeGenerator.javaType(dataType());
List<String> parts = new ArrayList<>();
List<Object> args = new ArrayList<>();
/*
- CODE GENERATED
- final $javaType ${ev.value} = $beamCoderWrapper.encode(${input.isNull}, ${input.value});
+ CODE GENERATED
+ final ${javaType} ${ev.value} = org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.EncodeUsingBeamCoder.encode(${input.value}, ${coder});
*/
parts.add("final ");
args.add(javaType);
parts.add(" ");
args.add(ev.value());
- parts.add(" = ");
- args.add(accessCode);
- parts.add(".encode(");
+ parts.add(
+ " = org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.EncodeUsingBeamCoder.encode(");
args.add(input.isNull());
parts.add(", ");
args.add(input.value());
+ parts.add(", ");
+ args.add(accessCode);
parts.add(");");
StringContext sc =
@@ -174,7 +134,7 @@
case 0:
return child;
case 1:
- return beamCoderWrapper;
+ return coder;
default:
throw new ArrayIndexOutOfBoundsException("productElement out of bounds");
}
@@ -199,12 +159,20 @@
return false;
}
EncodeUsingBeamCoder<?> that = (EncodeUsingBeamCoder<?>) o;
- return beamCoderWrapper.equals(that.beamCoderWrapper) && child.equals(that.child);
+ return child.equals(that.child) && coder.equals(that.coder);
}
@Override
public int hashCode() {
- return Objects.hash(super.hashCode(), child, beamCoderWrapper);
+ return Objects.hash(super.hashCode(), child, coder);
+ }
+
+ /**
+ * Convert value to byte array (invoked by generated code in {@link #doGenCode(CodegenContext,
+ * ExprCode)}).
+ */
+ public static <T> byte[] encode(boolean isNull, @Nullable T value, Coder<T> coder) {
+ return isNull ? null : CoderHelpers.toByteArray(value, coder);
}
}
@@ -216,15 +184,14 @@
public static class DecodeUsingBeamCoder<T> extends UnaryExpression
implements NonSQLExpression, Serializable {
- private Expression child;
- private ClassTag<T> classTag;
- private BeamCoderWrapper<T> beamCoderWrapper;
+ private final Expression child;
+ private final ClassTag<T> classTag;
+ private final Coder<T> coder;
- public DecodeUsingBeamCoder(
- Expression child, ClassTag<T> classTag, BeamCoderWrapper<T> beamCoderWrapper) {
+ public DecodeUsingBeamCoder(Expression child, ClassTag<T> classTag, Coder<T> coder) {
this.child = child;
this.classTag = classTag;
- this.beamCoderWrapper = beamCoderWrapper;
+ this.coder = coder;
}
@Override
@@ -234,16 +201,15 @@
@Override
public ExprCode doGenCode(CodegenContext ctx, ExprCode ev) {
- String accessCode =
- ctx.addReferenceObj(
- "beamCoderWrapper", beamCoderWrapper, BeamCoderWrapper.class.getName());
+ String accessCode = ctx.addReferenceObj("coder", coder, coder.getClass().getName());
ExprCode input = child.genCode(ctx);
String javaType = CodeGenerator.javaType(dataType());
+
List<String> parts = new ArrayList<>();
List<Object> args = new ArrayList<>();
/*
- CODE GENERATED:
- final $javaType ${ev.value} = ($javaType) $beamCoderWrapper.decode(${input.isNull}, ${input.value});
+ CODE GENERATED:
+ final ${javaType} ${ev.value} = (${javaType}) org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.DecodeUsingBeamCoder.decode(${input.value}, ${coder});
*/
parts.add("final ");
args.add(javaType);
@@ -251,13 +217,15 @@
args.add(ev.value());
parts.add(" = (");
args.add(javaType);
- parts.add(") ");
- args.add(accessCode);
- parts.add(".decode(");
+ parts.add(
+ ") org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.DecodeUsingBeamCoder.decode(");
args.add(input.isNull());
parts.add(", ");
args.add(input.value());
+ parts.add(", ");
+ args.add(accessCode);
parts.add(");");
+
StringContext sc =
new StringContext(JavaConversions.collectionAsScalaIterable(parts).toSeq());
Block code =
@@ -278,7 +246,7 @@
case 1:
return classTag;
case 2:
- return beamCoderWrapper;
+ return coder;
default:
throw new ArrayIndexOutOfBoundsException("productElement out of bounds");
}
@@ -303,14 +271,20 @@
return false;
}
DecodeUsingBeamCoder<?> that = (DecodeUsingBeamCoder<?>) o;
- return child.equals(that.child)
- && classTag.equals(that.classTag)
- && beamCoderWrapper.equals(that.beamCoderWrapper);
+ return child.equals(that.child) && classTag.equals(that.classTag) && coder.equals(that.coder);
}
@Override
public int hashCode() {
- return Objects.hash(super.hashCode(), child, classTag, beamCoderWrapper);
+ return Objects.hash(super.hashCode(), child, classTag, coder);
+ }
+
+ /**
+ * Convert value from byte array (invoked by generated code in {@link #doGenCode(CodegenContext,
+ * ExprCode)}).
+ */
+ public static <T> T decode(boolean isNull, @Nullable byte[] serialized, Coder<T> coder) {
+ return isNull ? null : CoderHelpers.fromByteArray(serialized, coder);
}
}
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java
index 845dc63..013f860 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java
@@ -71,11 +71,13 @@
@Override
public void onTimer(
final String timerId,
+ final String timerFamilyId,
final BoundedWindow window,
final Instant timestamp,
+ final Instant outputTimestamp,
final TimeDomain timeDomain) {
try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metricsContainer())) {
- delegate.onTimer(timerId, window, timestamp, timeDomain);
+ delegate.onTimer(timerId, timerFamilyId, window, timestamp, outputTimestamp, timeDomain);
} catch (IOException e) {
throw new RuntimeException(e);
}
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkProcessContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkProcessContext.java
index e978f46..9cbbeda 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkProcessContext.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/SparkProcessContext.java
@@ -161,7 +161,13 @@
StateNamespace namespace = timer.getNamespace();
checkArgument(namespace instanceof StateNamespaces.WindowNamespace);
BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
- doFnRunner.onTimer(timer.getTimerId(), window, timer.getTimestamp(), timer.getDomain());
+ doFnRunner.onTimer(
+ timer.getTimerId(),
+ timer.getTimerFamilyId(),
+ window,
+ timer.getTimestamp(),
+ timer.getOutputTimestamp(),
+ timer.getDomain());
}
}
}
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/utils/EncodersTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpersTest.java
similarity index 79%
rename from runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/utils/EncodersTest.java
rename to runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpersTest.java
index 8327fd8..54db4fa 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/utils/EncodersTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/helpers/EncoderHelpersTest.java
@@ -15,13 +15,12 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.runners.spark.structuredstreaming.utils;
+package org.apache.beam.runners.spark.structuredstreaming.translation.helpers;
import static org.junit.Assert.assertEquals;
-import java.util.ArrayList;
+import java.util.Arrays;
import java.util.List;
-import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.SparkSession;
@@ -31,7 +30,7 @@
/** Test of the wrapping of Beam Coders as Spark ExpressionEncoders. */
@RunWith(JUnit4.class)
-public class EncodersTest {
+public class EncoderHelpersTest {
@Test
public void beamCoderToSparkEncoderTest() {
@@ -40,13 +39,9 @@
.appName("beamCoderToSparkEncoderTest")
.master("local[4]")
.getOrCreate();
- List<Integer> data = new ArrayList<>();
- data.add(1);
- data.add(2);
- data.add(3);
+ List<Integer> data = Arrays.asList(1, 2, 3);
Dataset<Integer> dataset =
sparkSession.createDataset(data, EncoderHelpers.fromBeamCoder(VarIntCoder.of()));
- List<Integer> results = dataset.collectAsList();
- assertEquals(data, results);
+ assertEquals(data, dataset.collectAsList());
}
}
diff --git a/sdks/go/README.md b/sdks/go/README.md
index ddbbe9c..f2d56d4 100644
--- a/sdks/go/README.md
+++ b/sdks/go/README.md
@@ -108,8 +108,7 @@
See [BUILD.md](./BUILD.md) for how to build Go code in general. See
-[CONTAINERS.md](../CONTAINERS.md) for how to build and push the Go
-SDK harness container image.
+[container documentation](https://beam.apache.org/documentation/runtime/environments/#building-container-images) for how to build and push the Go SDK harness container image.
## Issues
diff --git a/sdks/go/examples/build.gradle b/sdks/go/examples/build.gradle
index f7299f1..c040dd5 100644
--- a/sdks/go/examples/build.gradle
+++ b/sdks/go/examples/build.gradle
@@ -69,4 +69,9 @@
go 'build -o ./build/bin/${GOOS}_${GOARCH}/wordcount github.com/apache/beam/sdks/go/examples/wordcount'
go 'build -o ./build/bin/${GOOS}_${GOARCH}/yatzy github.com/apache/beam/sdks/go/examples/yatzy'
}
+
+ // Ignore spurious vet errors during check for [BEAM-8992].
+ goVet {
+ continueOnFailure = true
+ }
}
diff --git a/sdks/go/pkg/beam/core/metrics/metrics.go b/sdks/go/pkg/beam/core/metrics/metrics.go
index a29446b..a21f109 100644
--- a/sdks/go/pkg/beam/core/metrics/metrics.go
+++ b/sdks/go/pkg/beam/core/metrics/metrics.go
@@ -49,10 +49,13 @@
import (
"context"
"fmt"
+ "hash/fnv"
"sort"
"sync"
+ "sync/atomic"
"time"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/util/ioutilx"
"github.com/apache/beam/sdks/go/pkg/beam/log"
"github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
"github.com/golang/protobuf/ptypes"
@@ -64,22 +67,57 @@
// using metrics requires the PTransform have a context.Context
// argument.
+// perBundle is a struct that retains per transform countersets.
+// TODO(lostluck): Migrate the exec package to use these to extract
+// metric data for export to runner, rather than the global store.
+type perBundle struct {
+ mu sync.Mutex
+ css []*ptCounterSet
+}
+
+type nameHash uint64
+
+// ptCounterSet is the internal tracking struct for a single ptransform
+// in a single bundle for all counter types.
+type ptCounterSet struct {
+ // We store the user path access to the cells in metric type segregated
+ // maps. At present, caching the name hash,
+ counters map[nameHash]*counter
+ distributions map[nameHash]*distribution
+ gauges map[nameHash]*gauge
+}
+
type ctxKey string
-const bundleKey ctxKey = "beam:bundle"
-const ptransformKey ctxKey = "beam:ptransform"
+const (
+ bundleKey ctxKey = "beam:bundle"
+ ptransformKey ctxKey = "beam:ptransform"
+ counterSetKey ctxKey = "beam:counterset"
+)
// beamCtx is a caching context for IDs necessary to place metric updates.
-// Allocating contexts and searching for PTransformIDs for every element
+// Allocating contexts and searching for PTransformIDs for every element
// is expensive, so we avoid it if possible.
type beamCtx struct {
context.Context
bundleID, ptransformID string
+ bs *perBundle
+ cs *ptCounterSet
}
-// Value lifts the beam contLift the keys value for faster lookups when not available.
+// Value implements context.Value for beamCtx, and lifts the
+// values for metrics keys for value for faster lookups.
func (ctx *beamCtx) Value(key interface{}) interface{} {
switch key {
+ case counterSetKey:
+ if ctx.cs == nil {
+ if cs := ctx.Context.Value(key); cs != nil {
+ ctx.cs = cs.(*ptCounterSet)
+ } else {
+ return nil
+ }
+ }
+ return ctx.cs
case bundleKey:
if ctx.bundleID == "" {
if id := ctx.Context.Value(key); id != nil {
@@ -102,24 +140,30 @@
return ctx.Context.Value(key)
}
+func (ctx *beamCtx) String() string {
+ return fmt.Sprintf("beamCtx[%s;%s]", ctx.bundleID, ctx.ptransformID)
+}
+
// SetBundleID sets the id of the current Bundle.
func SetBundleID(ctx context.Context, id string) context.Context {
// Checking for *beamCtx is an optimization, so we don't dig deeply
// for ids if not necessary.
if bctx, ok := ctx.(*beamCtx); ok {
- return &beamCtx{Context: bctx.Context, bundleID: id, ptransformID: bctx.ptransformID}
+ return &beamCtx{Context: bctx.Context, bundleID: id, bs: &perBundle{}, ptransformID: bctx.ptransformID}
}
- return &beamCtx{Context: ctx, bundleID: id}
+ return &beamCtx{Context: ctx, bundleID: id, bs: &perBundle{}}
}
// SetPTransformID sets the id of the current PTransform.
+// Must only be called on a context returened by SetBundleID.
func SetPTransformID(ctx context.Context, id string) context.Context {
// Checking for *beamCtx is an optimization, so we don't dig deeply
// for ids if not necessary.
if bctx, ok := ctx.(*beamCtx); ok {
- return &beamCtx{Context: bctx.Context, bundleID: bctx.bundleID, ptransformID: id}
+ return &beamCtx{Context: bctx.Context, bundleID: bctx.bundleID, bs: bctx.bs, ptransformID: id}
}
- return &beamCtx{Context: ctx, ptransformID: id}
+ panic(fmt.Sprintf("SetPTransformID called before SetBundleID for %v", id))
+ return nil // never runs.
}
const (
@@ -138,9 +182,55 @@
return key
}
+func getCounterSet(ctx context.Context) *ptCounterSet {
+ if id := ctx.Value(counterSetKey); id != nil {
+ return id.(*ptCounterSet)
+ }
+ // It's not set anywhere and wasn't hoisted, so create it.
+ if bctx, ok := ctx.(*beamCtx); ok {
+ bctx.bs.mu.Lock()
+ cs := &ptCounterSet{
+ counters: make(map[nameHash]*counter),
+ distributions: make(map[nameHash]*distribution),
+ gauges: make(map[nameHash]*gauge),
+ }
+ bctx.bs.css = append(bctx.bs.css, cs)
+ bctx.cs = cs
+ bctx.bs.mu.Unlock()
+ return cs
+ }
+ panic("counterSet missing, beam isn't set up properly.")
+ return nil // never runs.
+}
+
+type kind uint8
+
+const (
+ kindUnknown kind = iota
+ kindSumCounter
+ kindDistribution
+ kindGauge
+)
+
+func (t kind) String() string {
+ switch t {
+ case kindSumCounter:
+ return "Counter"
+ case kindDistribution:
+ return "Distribution"
+ case kindGauge:
+ return "Gauge"
+ default:
+ panic(fmt.Sprintf("Unknown metric type value: %v", uint8(t)))
+ }
+}
+
// userMetric knows how to convert it's value to a Metrics_User proto.
+// TODO(lostluck): Move proto translation to the harness package to
+// avoid the proto dependency outside of the harness.
type userMetric interface {
toProto() *fnexecution_v1.Metrics_User
+ kind() kind
}
// name is a pair of strings identifying a specific metric.
@@ -159,43 +249,87 @@
return name{namespace: ns, name: n}
}
+// We hash the name to a uint64 so we avoid using go's native string hashing for
+// every use of a metrics. uint64s have faster lookup than strings as a result.
+// Collisions are possible, but statistically unlikely as namespaces and names
+// are usually short enough to avoid this.
+var (
+ hasherMu sync.Mutex
+ hasher = fnv.New64a()
+)
+
+func hashName(ns, n string) nameHash {
+ hasherMu.Lock()
+ hasher.Reset()
+ var buf [64]byte
+ b := buf[:]
+ hashString(ns, b)
+ hashString(n, b)
+ h := hasher.Sum64()
+ hasherMu.Unlock()
+ return nameHash(h)
+}
+
+// hashString hashes a string with the package level hasher
+// and requires posession of the hasherMu lock. The byte
+// slice is assumed to be backed by a [64]byte.
+func hashString(s string, b []byte) {
+ l := len(s)
+ i := 0
+ for len(s)-i > 64 {
+ n := i + 64
+ copy(b, s[i:n])
+ ioutilx.WriteUnsafe(hasher, b)
+ i = n
+ }
+ n := l - i
+ copy(b, s[i:])
+ ioutilx.WriteUnsafe(hasher, b[:n])
+}
+
type key struct {
name name
bundle, ptransform string
}
var (
- // mu protects access to store
+ // mu protects access to the global store
mu sync.RWMutex
// store is a map of BundleIDs to PtransformIDs to userMetrics.
// it permits us to extract metric protos for runners per data Bundle, and
// per PTransform.
+ // TODO(lostluck): Migrate exec/plan.go to manage it's own perBundle counter stores.
+ // Note: This is safe to use to read the metrics concurrently with user modification
+ // in bundles, as only initial use modifies this map, and only contains the resulting
+ // metrics.
store = make(map[string]map[string]map[name]userMetric)
-
- // We store the user path access to the cells in metric type segregated
- // sync.Maps. Using sync.Maps lets metrics with disjoint keys have concurrent
- // access to the cells, and using separate sync.Map per metric type
- // simplifies code understanding, since each only contains a single type of
- // cell.
-
- countersMu sync.RWMutex
- counters = make(map[key]*counter)
- distributionsMu sync.RWMutex
- distributions = make(map[key]*distribution)
- gaugesMu sync.RWMutex
- gauges = make(map[key]*gauge)
)
-// TODO(lostluck): 2018/03/05 Use a common internal beam now() instead, once that exists.
-var now = time.Now
+// storeMetric stores a metric away on its first use so it may be retrieved later on.
+// In the event of a name collision, storeMetric can panic, so it's prudent to release
+// locks if they are no longer required.
+func storeMetric(key key, m userMetric) {
+ mu.Lock()
+ defer mu.Unlock()
+ if _, ok := store[key.bundle]; !ok {
+ store[key.bundle] = make(map[string]map[name]userMetric)
+ }
+ if _, ok := store[key.bundle][key.ptransform]; !ok {
+ store[key.bundle][key.ptransform] = make(map[name]userMetric)
+ }
+ if ms, ok := store[key.bundle][key.ptransform][key.name]; ok {
+ if ms.kind() != m.kind() {
+ panic(fmt.Sprintf("metric name %s being reused for a second metric in a single PTransform", key.name))
+ }
+ return
+ }
+ store[key.bundle][key.ptransform][key.name] = m
+}
// Counter is a simple counter for incrementing and decrementing a value.
type Counter struct {
name name
- // The following are a fast cache of the key and storage
- mu sync.Mutex
- k key
- c *counter
+ hash nameHash
}
func (m *Counter) String() string {
@@ -204,38 +338,26 @@
// NewCounter returns the Counter with the given namespace and name.
func NewCounter(ns, n string) *Counter {
- mn := newName(ns, n)
return &Counter{
- name: mn,
+ name: newName(ns, n),
+ hash: hashName(ns, n),
}
}
// Inc increments the counter within the given PTransform context by v.
func (m *Counter) Inc(ctx context.Context, v int64) {
- key := getContextKey(ctx, m.name)
- cs := &counter{
- value: v,
- }
- m.mu.Lock()
- if m.k == key {
- m.c.inc(v)
- m.mu.Unlock()
- return
- }
- m.k = key
- countersMu.Lock()
- if c, loaded := counters[key]; loaded {
- m.c = c
- countersMu.Unlock()
- m.mu.Unlock()
+ cs := getCounterSet(ctx)
+ if c, ok := cs.counters[m.hash]; ok {
c.inc(v)
return
}
- m.c = cs
- counters[key] = cs
- countersMu.Unlock()
- m.mu.Unlock()
- storeMetric(key, cs)
+ // We're the first to create this metric!
+ c := &counter{
+ value: v,
+ }
+ cs.counters[m.hash] = c
+ key := getContextKey(ctx, m.name)
+ storeMetric(key, c)
}
// Dec decrements the counter within the given PTransform context by v.
@@ -246,13 +368,10 @@
// counter is a metric cell for counter values.
type counter struct {
value int64
- mu sync.Mutex
}
func (m *counter) inc(v int64) {
- m.mu.Lock()
- m.value += v
- m.mu.Unlock()
+ atomic.AddInt64(&m.value, v)
}
func (m *counter) String() string {
@@ -262,25 +381,23 @@
// toProto returns a Metrics_User populated with the Data messages, but not the name. The
// caller needs to populate with the metric's name.
func (m *counter) toProto() *fnexecution_v1.Metrics_User {
- m.mu.Lock()
- defer m.mu.Unlock()
return &fnexecution_v1.Metrics_User{
Data: &fnexecution_v1.Metrics_User_CounterData_{
CounterData: &fnexecution_v1.Metrics_User_CounterData{
- Value: m.value,
+ Value: atomic.LoadInt64(&m.value),
},
},
}
}
+func (m *counter) kind() kind {
+ return kindSumCounter
+}
+
// Distribution is a simple distribution of values.
type Distribution struct {
name name
-
- // The following are a fast cache of the key and storage
- mu sync.Mutex
- k key
- d *distribution
+ hash nameHash
}
func (m *Distribution) String() string {
@@ -289,41 +406,29 @@
// NewDistribution returns the Distribution with the given namespace and name.
func NewDistribution(ns, n string) *Distribution {
- mn := newName(ns, n)
return &Distribution{
- name: mn,
+ name: newName(ns, n),
+ hash: hashName(ns, n),
}
}
// Update updates the distribution within the given PTransform context with v.
func (m *Distribution) Update(ctx context.Context, v int64) {
- key := getContextKey(ctx, m.name)
- ds := &distribution{
+ cs := getCounterSet(ctx)
+ if d, ok := cs.distributions[m.hash]; ok {
+ d.update(v)
+ return
+ }
+ // We're the first to create this metric!
+ d := &distribution{
count: 1,
sum: v,
min: v,
max: v,
}
- m.mu.Lock()
- if m.k == key {
- m.d.update(v)
- m.mu.Unlock()
- return
- }
- m.k = key
- distributionsMu.Lock()
- if d, loaded := distributions[key]; loaded {
- m.d = d
- distributionsMu.Unlock()
- m.mu.Unlock()
- d.update(v)
- return
- }
- m.d = ds
- distributions[key] = ds
- distributionsMu.Unlock()
- m.mu.Unlock()
- storeMetric(key, ds)
+ cs.distributions[m.hash] = d
+ key := getContextKey(ctx, m.name)
+ storeMetric(key, d)
}
// distribution is a metric cell for distribution values.
@@ -366,14 +471,14 @@
}
}
+func (m *distribution) kind() kind {
+ return kindDistribution
+}
+
// Gauge is a time, value pair metric.
type Gauge struct {
name name
-
- // The following are a fast cache of the key and storage
- mu sync.Mutex
- k key
- g *gauge
+ hash nameHash
}
func (m *Gauge) String() string {
@@ -382,57 +487,30 @@
// NewGauge returns the Gauge with the given namespace and name.
func NewGauge(ns, n string) *Gauge {
- mn := newName(ns, n)
return &Gauge{
- name: mn,
+ name: newName(ns, n),
+ hash: hashName(ns, n),
}
}
+// TODO(lostluck): 2018/03/05 Use a common internal beam now() instead, once that exists.
+var now = time.Now
+
// Set sets the gauge to the given value, and associates it with the current time on the clock.
func (m *Gauge) Set(ctx context.Context, v int64) {
- key := getContextKey(ctx, m.name)
- gs := &gauge{
- t: now(),
- v: v,
- }
- m.mu.Lock()
- if m.k == key {
- m.g.set(v)
- m.mu.Unlock()
- return
- }
- m.k = key
- gaugesMu.Lock()
- if g, loaded := gauges[key]; loaded {
- m.g = g
- gaugesMu.Unlock()
- m.mu.Unlock()
+ cs := getCounterSet(ctx)
+ if g, ok := cs.gauges[m.hash]; ok {
g.set(v)
return
}
- m.g = gs
- gauges[key] = gs
- gaugesMu.Unlock()
- m.mu.Unlock()
- storeMetric(key, gs)
-}
-
-// storeMetric stores a metric away on its first use so it may be retrieved later on.
-// In the event of a name collision, storeMetric can panic, so it's prudent to release
-// locks if they are no longer required.
-func storeMetric(key key, m userMetric) {
- mu.Lock()
- defer mu.Unlock()
- if _, ok := store[key.bundle]; !ok {
- store[key.bundle] = make(map[string]map[name]userMetric)
+ // We're the first to create this metric!
+ g := &gauge{
+ t: now(),
+ v: v,
}
- if _, ok := store[key.bundle][key.ptransform]; !ok {
- store[key.bundle][key.ptransform] = make(map[name]userMetric)
- }
- if _, ok := store[key.bundle][key.ptransform][key.name]; ok {
- panic(fmt.Sprintf("metric name %s being reused for a second metric in a single PTransform", key.name))
- }
- store[key.bundle][key.ptransform][key.name] = m
+ cs.gauges[m.hash] = g
+ key := getContextKey(ctx, m.name)
+ storeMetric(key, g)
}
// gauge is a metric cell for gauge values.
@@ -466,8 +544,12 @@
}
}
+func (m *gauge) kind() kind {
+ return kindGauge
+}
+
func (m *gauge) String() string {
- return fmt.Sprintf("time: %s value: %d", m.t, m.v)
+ return fmt.Sprintf("%v time: %s value: %d", m.kind(), m.t, m.v)
}
// ToProto exports all collected metrics for the given BundleID and PTransform ID pair.
@@ -505,23 +587,19 @@
func dumpTo(p func(format string, args ...interface{})) {
mu.RLock()
defer mu.RUnlock()
- countersMu.RLock()
- defer countersMu.RUnlock()
- distributionsMu.RLock()
- defer distributionsMu.RUnlock()
- gaugesMu.RLock()
- defer gaugesMu.RUnlock()
var bs []string
for b := range store {
bs = append(bs, b)
}
sort.Strings(bs)
+
for _, b := range bs {
var pts []string
for pt := range store[b] {
pts = append(pts, pt)
}
sort.Strings(pts)
+
for _, pt := range pts {
var ns []name
for n := range store[b][pt] {
@@ -537,17 +615,10 @@
return false
})
p("Bundle: %q - PTransformID: %q", b, pt)
+
for _, n := range ns {
- key := key{name: n, bundle: b, ptransform: pt}
- if m, ok := counters[key]; ok {
- p("\t%s - %s", key.name, m)
- }
- if m, ok := distributions[key]; ok {
- p("\t%s - %s", key.name, m)
- }
- if m, ok := gauges[key]; ok {
- p("\t%s - %s", key.name, m)
- }
+ m := store[b][pt][n]
+ p("\t%s - %s", n, m)
}
}
}
@@ -558,9 +629,6 @@
func Clear() {
mu.Lock()
store = make(map[string]map[string]map[name]userMetric)
- counters = make(map[key]*counter)
- distributions = make(map[key]*distribution)
- gauges = make(map[key]*gauge)
mu.Unlock()
}
@@ -570,21 +638,6 @@
// No concurrency races since mu guards all access to store,
// and the metric cell sync.Maps are goroutine safe.
mu.Lock()
- pts := store[b]
- countersMu.Lock()
- distributionsMu.Lock()
- gaugesMu.Lock()
- for pt, m := range pts {
- for n := range m {
- key := key{name: n, bundle: b, ptransform: pt}
- delete(counters, key)
- delete(distributions, key)
- delete(gauges, key)
- }
- }
- countersMu.Unlock()
- distributionsMu.Unlock()
- gaugesMu.Unlock()
delete(store, b)
mu.Unlock()
}
diff --git a/sdks/go/pkg/beam/core/metrics/metrics_test.go b/sdks/go/pkg/beam/core/metrics/metrics_test.go
index 3ef2de9..c2d62c0 100644
--- a/sdks/go/pkg/beam/core/metrics/metrics_test.go
+++ b/sdks/go/pkg/beam/core/metrics/metrics_test.go
@@ -25,194 +25,150 @@
// bID is a bundleId to use in the tests, if nothing more specific is needed.
const bID = "bID"
-// TestRobustness validates metrics not panicking if the context doesn't
-// have the bundle or transform ID.
-func TestRobustness(t *testing.T) {
- m := NewCounter("Test", "myCount")
- m.Inc(context.Background(), 3)
- ptCtx := SetPTransformID(context.Background(), "MY_TRANSFORM")
- m.Inc(ptCtx, 3)
- bCtx := SetBundleID(context.Background(), bID)
- m.Inc(bCtx, 3)
-}
-
-func TestBeamContext(t *testing.T) {
- t.Run("ptransformID", func(t *testing.T) {
- ptID := "MY_TRANSFORM"
- ctx := SetPTransformID(context.Background(), ptID)
- key := getContextKey(ctx, name{})
- if key.bundle != bundleIDUnset {
- t.Errorf("bundleId = %q, want %q", key.bundle, bundleIDUnset)
- }
- if key.ptransform != ptID {
- t.Errorf("ptransformId = %q, want %q", key.ptransform, ptID)
- }
-
- })
-
- t.Run("bundleID", func(t *testing.T) {
- ctx := SetBundleID(context.Background(), bID)
- key := getContextKey(ctx, name{})
- if key.bundle != bID {
- t.Errorf("bundleId = %q, want %q", key.bundle, bID)
- }
- if key.ptransform != ptransformIDUnset {
- t.Errorf("ptransformId = %q, want %q", key.ptransform, ptransformIDUnset)
- }
- })
-}
-
func ctxWith(b, pt string) context.Context {
ctx := context.Background()
- ctx = SetPTransformID(ctx, pt)
ctx = SetBundleID(ctx, b)
+ ctx = SetPTransformID(ctx, pt)
return ctx
}
func TestCounter_Inc(t *testing.T) {
+ ctxA := ctxWith(bID, "A")
+ ctxB := ctxWith(bID, "B")
+
tests := []struct {
- ns, n, key string // Counter name and PTransform context
- inc int64
- value int64 // Internal variable to check
+ ns, n string // Counter name
+ ctx context.Context
+ inc int64
+ value int64 // Internal variable to check
}{
- {ns: "inc1", n: "count", key: "A", inc: 1, value: 1},
- {ns: "inc1", n: "count", key: "A", inc: 1, value: 2},
- {ns: "inc1", n: "ticker", key: "A", inc: 1, value: 1},
- {ns: "inc1", n: "ticker", key: "A", inc: 2, value: 3},
- {ns: "inc1", n: "count", key: "B", inc: 1, value: 1},
- {ns: "inc1", n: "count", key: "B", inc: 1, value: 2},
- {ns: "inc1", n: "ticker", key: "B", inc: 1, value: 1},
- {ns: "inc1", n: "ticker", key: "B", inc: 2, value: 3},
- {ns: "inc2", n: "count", key: "A", inc: 1, value: 1},
- {ns: "inc2", n: "count", key: "A", inc: 1, value: 2},
- {ns: "inc2", n: "ticker", key: "A", inc: 1, value: 1},
- {ns: "inc2", n: "ticker", key: "A", inc: 2, value: 3},
- {ns: "inc2", n: "count", key: "B", inc: 1, value: 1},
- {ns: "inc2", n: "count", key: "B", inc: 1, value: 2},
- {ns: "inc2", n: "ticker", key: "B", inc: 1, value: 1},
- {ns: "inc2", n: "ticker", key: "B", inc: 2, value: 3},
+ {ns: "inc1", n: "count", ctx: ctxA, inc: 1, value: 1},
+ {ns: "inc1", n: "count", ctx: ctxA, inc: 1, value: 2},
+ {ns: "inc1", n: "ticker", ctx: ctxA, inc: 1, value: 1},
+ {ns: "inc1", n: "ticker", ctx: ctxA, inc: 2, value: 3},
+ {ns: "inc1", n: "count", ctx: ctxB, inc: 1, value: 1},
+ {ns: "inc1", n: "count", ctx: ctxB, inc: 1, value: 2},
+ {ns: "inc1", n: "ticker", ctx: ctxB, inc: 1, value: 1},
+ {ns: "inc1", n: "ticker", ctx: ctxB, inc: 2, value: 3},
+ {ns: "inc2", n: "count", ctx: ctxA, inc: 1, value: 1},
+ {ns: "inc2", n: "count", ctx: ctxA, inc: 1, value: 2},
+ {ns: "inc2", n: "ticker", ctx: ctxA, inc: 1, value: 1},
+ {ns: "inc2", n: "ticker", ctx: ctxA, inc: 2, value: 3},
+ {ns: "inc2", n: "count", ctx: ctxB, inc: 1, value: 1},
+ {ns: "inc2", n: "count", ctx: ctxB, inc: 1, value: 2},
+ {ns: "inc2", n: "ticker", ctx: ctxB, inc: 1, value: 1},
+ {ns: "inc2", n: "ticker", ctx: ctxB, inc: 2, value: 3},
}
for _, test := range tests {
- t.Run(fmt.Sprintf("add %d to %s.%s[%q] value: %d", test.inc, test.ns, test.n, test.key, test.value),
+ t.Run(fmt.Sprintf("add %d to %s.%s[%v] value: %d", test.inc, test.ns, test.n, test.ctx, test.value),
func(t *testing.T) {
m := NewCounter(test.ns, test.n)
- ctx := ctxWith(bID, test.key)
- m.Inc(ctx, test.inc)
+ m.Inc(test.ctx, test.inc)
- key := key{name: name{namespace: test.ns, name: test.n}, bundle: bID, ptransform: test.key}
- countersMu.Lock()
- c, ok := counters[key]
- countersMu.Unlock()
- if !ok {
- t.Fatalf("Unable to find Counter for key %v", key)
- }
+ cs := getCounterSet(test.ctx)
+ c := cs.counters[m.hash]
if got, want := c.value, test.value; got != want {
- t.Errorf("GetCounter(%q,%q).Inc(%s, %d) c.value got %v, want %v", test.ns, test.n, test.key, test.inc, got, want)
+ t.Errorf("GetCounter(%q,%q).Inc(%v, %d) c.value got %v, want %v", test.ns, test.n, test.ctx, test.inc, got, want)
}
})
}
}
func TestCounter_Dec(t *testing.T) {
+ ctxA := ctxWith(bID, "A")
+ ctxB := ctxWith(bID, "B")
+
tests := []struct {
- ns, n, key string // Counter name and PTransform context
- dec int64
- value int64 // Internal variable to check
+ ns, n string // Counter name
+ ctx context.Context
+ dec int64
+ value int64 // Internal variable to check
}{
- {ns: "dec1", n: "count", key: "A", dec: 1, value: -1},
- {ns: "dec1", n: "count", key: "A", dec: 1, value: -2},
- {ns: "dec1", n: "ticker", key: "A", dec: 1, value: -1},
- {ns: "dec1", n: "ticker", key: "A", dec: 2, value: -3},
- {ns: "dec1", n: "count", key: "B", dec: 1, value: -1},
- {ns: "dec1", n: "count", key: "B", dec: 1, value: -2},
- {ns: "dec1", n: "ticker", key: "B", dec: 1, value: -1},
- {ns: "dec1", n: "ticker", key: "B", dec: 2, value: -3},
- {ns: "dec2", n: "count", key: "A", dec: 1, value: -1},
- {ns: "dec2", n: "count", key: "A", dec: 1, value: -2},
- {ns: "dec2", n: "ticker", key: "A", dec: 1, value: -1},
- {ns: "dec2", n: "ticker", key: "A", dec: 2, value: -3},
- {ns: "dec2", n: "count", key: "B", dec: 1, value: -1},
- {ns: "dec2", n: "count", key: "B", dec: 1, value: -2},
- {ns: "dec2", n: "ticker", key: "B", dec: 1, value: -1},
- {ns: "dec2", n: "ticker", key: "B", dec: 2, value: -3},
+ {ns: "dec1", n: "count", ctx: ctxA, dec: 1, value: -1},
+ {ns: "dec1", n: "count", ctx: ctxA, dec: 1, value: -2},
+ {ns: "dec1", n: "ticker", ctx: ctxA, dec: 1, value: -1},
+ {ns: "dec1", n: "ticker", ctx: ctxA, dec: 2, value: -3},
+ {ns: "dec1", n: "count", ctx: ctxB, dec: 1, value: -1},
+ {ns: "dec1", n: "count", ctx: ctxB, dec: 1, value: -2},
+ {ns: "dec1", n: "ticker", ctx: ctxB, dec: 1, value: -1},
+ {ns: "dec1", n: "ticker", ctx: ctxB, dec: 2, value: -3},
+ {ns: "dec2", n: "count", ctx: ctxA, dec: 1, value: -1},
+ {ns: "dec2", n: "count", ctx: ctxA, dec: 1, value: -2},
+ {ns: "dec2", n: "ticker", ctx: ctxA, dec: 1, value: -1},
+ {ns: "dec2", n: "ticker", ctx: ctxA, dec: 2, value: -3},
+ {ns: "dec2", n: "count", ctx: ctxB, dec: 1, value: -1},
+ {ns: "dec2", n: "count", ctx: ctxB, dec: 1, value: -2},
+ {ns: "dec2", n: "ticker", ctx: ctxB, dec: 1, value: -1},
+ {ns: "dec2", n: "ticker", ctx: ctxB, dec: 2, value: -3},
}
for _, test := range tests {
- t.Run(fmt.Sprintf("subtract %d to %s.%s[%q] value: %d", test.dec, test.ns, test.n, test.key, test.value),
+ t.Run(fmt.Sprintf("subtract %d to %s.%s[%v] value: %d", test.dec, test.ns, test.n, test.ctx, test.value),
func(t *testing.T) {
m := NewCounter(test.ns, test.n)
- ctx := ctxWith(bID, test.key)
- m.Dec(ctx, test.dec)
+ m.Dec(test.ctx, test.dec)
- key := key{name: name{namespace: test.ns, name: test.n}, bundle: bID, ptransform: test.key}
- countersMu.Lock()
- c, ok := counters[key]
- countersMu.Unlock()
- if !ok {
- t.Fatalf("Unable to find Counter for key %v", key)
- }
+ cs := getCounterSet(test.ctx)
+ c := cs.counters[m.hash]
if got, want := c.value, test.value; got != want {
- t.Errorf("GetCounter(%q,%q).Dec(%s, %d) c.value got %v, want %v", test.ns, test.n, test.key, test.dec, got, want)
+ t.Errorf("GetCounter(%q,%q).Dec(%v, %d) c.value got %v, want %v", test.ns, test.n, test.ctx, test.dec, got, want)
}
})
}
}
func TestDistribution_Update(t *testing.T) {
+ ctxA := ctxWith(bID, "A")
+ ctxB := ctxWith(bID, "B")
tests := []struct {
- ns, n, key string // Gauge name and PTransform context
+ ns, n string // Counter name
+ ctx context.Context
v int64
count, sum, min, max int64 // Internal variables to check
}{
- {ns: "update1", n: "latency", key: "A", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update1", n: "latency", key: "A", v: 1, count: 2, sum: 2, min: 1, max: 1},
- {ns: "update1", n: "latency", key: "A", v: 1, count: 3, sum: 3, min: 1, max: 1},
- {ns: "update1", n: "size", key: "A", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update1", n: "size", key: "A", v: 2, count: 2, sum: 3, min: 1, max: 2},
- {ns: "update1", n: "size", key: "A", v: 3, count: 3, sum: 6, min: 1, max: 3},
- {ns: "update1", n: "size", key: "A", v: -4, count: 4, sum: 2, min: -4, max: 3},
- {ns: "update1", n: "size", key: "A", v: 1, count: 5, sum: 3, min: -4, max: 3},
- {ns: "update1", n: "latency", key: "B", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update1", n: "latency", key: "B", v: 1, count: 2, sum: 2, min: 1, max: 1},
- {ns: "update1", n: "size", key: "B", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update1", n: "size", key: "B", v: 2, count: 2, sum: 3, min: 1, max: 2},
- {ns: "update2", n: "latency", key: "A", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update2", n: "latency", key: "A", v: 1, count: 2, sum: 2, min: 1, max: 1},
- {ns: "update2", n: "size", key: "A", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update2", n: "size", key: "A", v: 2, count: 2, sum: 3, min: 1, max: 2},
- {ns: "update2", n: "latency", key: "B", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update2", n: "latency", key: "B", v: 1, count: 2, sum: 2, min: 1, max: 1},
- {ns: "update2", n: "size", key: "B", v: 1, count: 1, sum: 1, min: 1, max: 1},
- {ns: "update2", n: "size", key: "B", v: 2, count: 2, sum: 3, min: 1, max: 2},
- {ns: "update1", n: "size", key: "A", v: 1, count: 6, sum: 4, min: -4, max: 3},
+ {ns: "update1", n: "latency", ctx: ctxA, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update1", n: "latency", ctx: ctxA, v: 1, count: 2, sum: 2, min: 1, max: 1},
+ {ns: "update1", n: "latency", ctx: ctxA, v: 1, count: 3, sum: 3, min: 1, max: 1},
+ {ns: "update1", n: "size", ctx: ctxA, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update1", n: "size", ctx: ctxA, v: 2, count: 2, sum: 3, min: 1, max: 2},
+ {ns: "update1", n: "size", ctx: ctxA, v: 3, count: 3, sum: 6, min: 1, max: 3},
+ {ns: "update1", n: "size", ctx: ctxA, v: -4, count: 4, sum: 2, min: -4, max: 3},
+ {ns: "update1", n: "size", ctx: ctxA, v: 1, count: 5, sum: 3, min: -4, max: 3},
+ {ns: "update1", n: "latency", ctx: ctxB, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update1", n: "latency", ctx: ctxB, v: 1, count: 2, sum: 2, min: 1, max: 1},
+ {ns: "update1", n: "size", ctx: ctxB, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update1", n: "size", ctx: ctxB, v: 2, count: 2, sum: 3, min: 1, max: 2},
+ {ns: "update2", n: "latency", ctx: ctxA, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update2", n: "latency", ctx: ctxA, v: 1, count: 2, sum: 2, min: 1, max: 1},
+ {ns: "update2", n: "size", ctx: ctxA, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update2", n: "size", ctx: ctxA, v: 2, count: 2, sum: 3, min: 1, max: 2},
+ {ns: "update2", n: "latency", ctx: ctxB, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update2", n: "latency", ctx: ctxB, v: 1, count: 2, sum: 2, min: 1, max: 1},
+ {ns: "update2", n: "size", ctx: ctxB, v: 1, count: 1, sum: 1, min: 1, max: 1},
+ {ns: "update2", n: "size", ctx: ctxB, v: 2, count: 2, sum: 3, min: 1, max: 2},
+ {ns: "update1", n: "size", ctx: ctxA, v: 1, count: 6, sum: 4, min: -4, max: 3},
}
for _, test := range tests {
- t.Run(fmt.Sprintf("add %d to %s.%s[%q] count: %d sum: %d", test.v, test.ns, test.n, test.key, test.count, test.sum),
+ t.Run(fmt.Sprintf("add %d to %s.%s[%q] count: %d sum: %d", test.v, test.ns, test.n, test.ctx, test.count, test.sum),
func(t *testing.T) {
m := NewDistribution(test.ns, test.n)
- ctx := ctxWith(bID, test.key)
- m.Update(ctx, test.v)
+ m.Update(test.ctx, test.v)
- key := key{name: name{namespace: test.ns, name: test.n}, bundle: bID, ptransform: test.key}
- distributionsMu.Lock()
- d, ok := distributions[key]
- distributionsMu.Unlock()
- if !ok {
- t.Fatalf("Unable to find Distribution for key %v", key)
- }
+ cs := getCounterSet(test.ctx)
+ d := cs.distributions[m.hash]
if got, want := d.count, test.count; got != want {
- t.Errorf("GetDistribution(%q,%q).Update(%s, %d) d.count got %v, want %v", test.ns, test.n, test.key, test.v, got, want)
+ t.Errorf("GetDistribution(%q,%q).Update(%v, %d) d.count got %v, want %v", test.ns, test.n, test.ctx, test.v, got, want)
}
if got, want := d.sum, test.sum; got != want {
- t.Errorf("GetDistribution(%q,%q).Update(%s, %d) d.sum got %v, want %v", test.ns, test.n, test.key, test.v, got, want)
+ t.Errorf("GetDistribution(%q,%q).Update(%v, %d) d.sum got %v, want %v", test.ns, test.n, test.ctx, test.v, got, want)
}
if got, want := d.min, test.min; got != want {
- t.Errorf("GetDistribution(%q,%q).Update(%s, %d) d.min got %v, want %v", test.ns, test.n, test.key, test.v, got, want)
+ t.Errorf("GetDistribution(%q,%q).Update(%v, %d) d.min got %v, want %v", test.ns, test.n, test.ctx, test.v, got, want)
}
if got, want := d.max, test.max; got != want {
- t.Errorf("GetDistribution(%q,%q).Update(%s, %d) d.max got %v, want %v", test.ns, test.n, test.key, test.v, got, want)
+ t.Errorf("GetDistribution(%q,%q).Update(%v, %d) d.max got %v, want %v", test.ns, test.n, test.ctx, test.v, got, want)
}
})
}
@@ -223,74 +179,51 @@
}
func TestGauge_Set(t *testing.T) {
+ ctxA := ctxWith(bID, "A")
+ ctxB := ctxWith(bID, "B")
tests := []struct {
- ns, n, key string // Gauge name and PTransform context
- v int64
- t time.Time
+ ns, n string // Counter name
+ ctx context.Context
+ v int64
+ t time.Time
}{
- {ns: "set1", n: "load", key: "A", v: 1, t: time.Unix(0, 0)},
- {ns: "set1", n: "load", key: "A", v: 1, t: time.Unix(0, 0)},
- {ns: "set1", n: "speed", key: "A", v: 1, t: time.Unix(0, 0)},
- {ns: "set1", n: "speed", key: "A", v: 2, t: time.Unix(0, 0)},
- {ns: "set1", n: "load", key: "B", v: 1, t: time.Unix(0, 0)},
- {ns: "set1", n: "load", key: "B", v: 1, t: time.Unix(0, 0)},
- {ns: "set1", n: "speed", key: "B", v: 1, t: time.Unix(0, 0)},
- {ns: "set1", n: "speed", key: "B", v: 2, t: time.Unix(0, 0)},
- {ns: "set2", n: "load", key: "A", v: 1, t: time.Unix(0, 0)},
- {ns: "set2", n: "load", key: "A", v: 1, t: time.Unix(0, 0)},
- {ns: "set2", n: "speed", key: "A", v: 1, t: time.Unix(0, 0)},
- {ns: "set2", n: "speed", key: "A", v: 2, t: time.Unix(0, 0)},
- {ns: "set2", n: "load", key: "B", v: 1, t: time.Unix(0, 0)},
- {ns: "set2", n: "load", key: "B", v: 1, t: time.Unix(0, 0)},
- {ns: "set2", n: "speed", key: "B", v: 1, t: time.Unix(0, 0)},
- {ns: "set2", n: "speed", key: "B", v: 2, t: time.Unix(0, 0)},
+ {ns: "set1", n: "load", ctx: ctxA, v: 1, t: time.Unix(0, 0)},
+ {ns: "set1", n: "load", ctx: ctxA, v: 1, t: time.Unix(0, 0)},
+ {ns: "set1", n: "speed", ctx: ctxA, v: 1, t: time.Unix(0, 0)},
+ {ns: "set1", n: "speed", ctx: ctxA, v: 2, t: time.Unix(0, 0)},
+ {ns: "set1", n: "load", ctx: ctxB, v: 1, t: time.Unix(0, 0)},
+ {ns: "set1", n: "load", ctx: ctxB, v: 1, t: time.Unix(0, 0)},
+ {ns: "set1", n: "speed", ctx: ctxB, v: 1, t: time.Unix(0, 0)},
+ {ns: "set1", n: "speed", ctx: ctxB, v: 2, t: time.Unix(0, 0)},
+ {ns: "set2", n: "load", ctx: ctxA, v: 1, t: time.Unix(0, 0)},
+ {ns: "set2", n: "load", ctx: ctxA, v: 1, t: time.Unix(0, 0)},
+ {ns: "set2", n: "speed", ctx: ctxA, v: 1, t: time.Unix(0, 0)},
+ {ns: "set2", n: "speed", ctx: ctxA, v: 2, t: time.Unix(0, 0)},
+ {ns: "set2", n: "load", ctx: ctxB, v: 1, t: time.Unix(0, 0)},
+ {ns: "set2", n: "load", ctx: ctxB, v: 1, t: time.Unix(0, 0)},
+ {ns: "set2", n: "speed", ctx: ctxB, v: 1, t: time.Unix(0, 0)},
+ {ns: "set2", n: "speed", ctx: ctxB, v: 2, t: time.Unix(0, 0)},
}
for _, test := range tests {
- t.Run(fmt.Sprintf("set %s.%s[%q] to %d at %v", test.ns, test.n, test.key, test.v, test.t),
+ t.Run(fmt.Sprintf("set %s.%s[%v] to %d at %v", test.ns, test.n, test.ctx, test.v, test.t),
func(t *testing.T) {
m := NewGauge(test.ns, test.n)
- ctx := ctxWith(bID, test.key)
now = testclock(test.t)
- m.Set(ctx, test.v)
+ m.Set(test.ctx, test.v)
- key := key{name: name{namespace: test.ns, name: test.n}, bundle: bID, ptransform: test.key}
- gaugesMu.Lock()
- g, ok := gauges[key]
- gaugesMu.Unlock()
- if !ok {
- t.Fatalf("Unable to find Gauge for key %v", key)
- }
+ cs := getCounterSet(test.ctx)
+ g := cs.gauges[m.hash]
if got, want := g.v, test.v; got != want {
- t.Errorf("GetGauge(%q,%q).Set(%s, %d) g.v got %v, want %v", test.ns, test.n, test.key, test.v, got, want)
+ t.Errorf("GetGauge(%q,%q).Set(%v, %d) g.v got %v, want %v", test.ns, test.n, test.ctx, test.v, got, want)
}
if got, want := g.t, test.t; got != want {
- t.Errorf("GetGauge(%q,%q).Set(%s, %d) t.t got %v, want %v", test.ns, test.n, test.key, test.v, got, want)
+ t.Errorf("GetGauge(%q,%q).Set(%v, %d) t.t got %v, want %v", test.ns, test.n, test.ctx, test.v, got, want)
}
})
}
}
-type metricType uint8
-
-const (
- counterType metricType = iota
- distributionType
- gaugeType
-)
-
-func (t metricType) String() string {
- switch t {
- case counterType:
- return "Counter"
- case distributionType:
- return "Distribution"
- case gaugeType:
- return "Gauge"
- default:
- panic(fmt.Sprintf("Unknown metric type value: %v", uint8(t)))
- }
-}
func TestNameCollisions(t *testing.T) {
ns, c, d, g := "collisions", "counter", "distribution", "gauge"
// Checks that user code panics if a counter attempts to be defined in the same PTransform
@@ -302,17 +235,17 @@
NewDistribution(ns, d).Update(ctxWith(bID, d), 1)
NewGauge(ns, g).Set(ctxWith(bID, g), 1)
tests := []struct {
- existing, new metricType
+ existing, new kind
}{
- {existing: counterType, new: counterType},
- {existing: counterType, new: distributionType},
- {existing: counterType, new: gaugeType},
- {existing: distributionType, new: counterType},
- {existing: distributionType, new: distributionType},
- {existing: distributionType, new: gaugeType},
- {existing: gaugeType, new: counterType},
- {existing: gaugeType, new: distributionType},
- {existing: gaugeType, new: gaugeType},
+ {existing: kindSumCounter, new: kindSumCounter},
+ {existing: kindSumCounter, new: kindDistribution},
+ {existing: kindSumCounter, new: kindGauge},
+ {existing: kindDistribution, new: kindSumCounter},
+ {existing: kindDistribution, new: kindDistribution},
+ {existing: kindDistribution, new: kindGauge},
+ {existing: kindGauge, new: kindSumCounter},
+ {existing: kindGauge, new: kindDistribution},
+ {existing: kindGauge, new: kindGauge},
}
for _, test := range tests {
t.Run(fmt.Sprintf("%s name collides with %s", test.existing, test.new),
@@ -325,26 +258,26 @@
}
t.Error("panic expected")
} else {
- t.Log("reusing names is fine when the metrics the same type:", test.existing, test.new)
+ t.Log("reusing names is fine when the metric is the same type:", test.existing, test.new)
}
}()
var name string
switch test.existing {
- case counterType:
+ case kindSumCounter:
name = c
- case distributionType:
+ case kindDistribution:
name = d
- case gaugeType:
+ case kindGauge:
name = g
default:
t.Fatalf("unknown existing metricType with value: %v", int(test.existing))
}
switch test.new {
- case counterType:
+ case kindSumCounter:
NewCounter(ns, name).Inc(ctxWith(bID, name), 1)
- case distributionType:
+ case kindDistribution:
NewDistribution(ns, name).Update(ctxWith(bID, name), 1)
- case gaugeType:
+ case kindGauge:
NewGauge(ns, name).Set(ctxWith(bID, name), 1)
default:
t.Fatalf("unknown new metricType with value: %v", int(test.new))
@@ -398,14 +331,16 @@
}
}
-// Run on @lostluck's desktop:
+// Run on @lostluck's desktop (2020/01/21) go1.13.4
//
-// BenchmarkMetrics/counter_inplace-12 5000000 243 ns/op 128 B/op 2 allocs/op
-// BenchmarkMetrics/distribution_inplace-12 5000000 252 ns/op 160 B/op 2 allocs/op
-// BenchmarkMetrics/gauge_inplace-12 5000000 266 ns/op 160 B/op 2 allocs/op
-// BenchmarkMetrics/counter_predeclared-12 20000000 74.3 ns/op 16 B/op 1 allocs/op
-// BenchmarkMetrics/distribution_predeclared-12 20000000 79.6 ns/op 48 B/op 1 allocs/op
-// BenchmarkMetrics/gauge_predeclared-12 20000000 92.9 ns/op 48 B/op 1 allocs/op
+// Allocs & bytes should be consistetn within go versions, but ns/op is relative to the running machine.
+//
+// BenchmarkMetrics/counter_inplace-12 4814373 243 ns/op 48 B/op 1 allocs/op
+// BenchmarkMetrics/distribution_inplace-12 4455957 273 ns/op 48 B/op 1 allocs/op
+// BenchmarkMetrics/gauge_inplace-12 4605908 265 ns/op 48 B/op 1 allocs/op
+// BenchmarkMetrics/counter_predeclared-12 75339600 15.5 ns/op 0 B/op 0 allocs/op
+// BenchmarkMetrics/distribution_predeclared-12 49202775 24.4 ns/op 0 B/op 0 allocs/op
+// BenchmarkMetrics/gauge_predeclared-12 46614810 28.3 ns/op 0 B/op 0 allocs/op
func BenchmarkMetrics(b *testing.B) {
Clear()
pt, c, d, g := "bench.bundle.data", "counter", "distribution", "gauge"
diff --git a/sdks/go/pkg/beam/core/runtime/exec/datasource.go b/sdks/go/pkg/beam/core/runtime/exec/datasource.go
index 06131b7..5d50061 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/datasource.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/datasource.go
@@ -137,7 +137,7 @@
switch {
case size >= 0:
// Single chunk streams are fully read in and buffered in memory.
- var buf []FullValue
+ buf := make([]FullValue, 0, size)
buf, err = readStreamToBuffer(cv, r, int64(size), buf)
if err != nil {
return nil, err
@@ -156,10 +156,12 @@
case chunk == 0: // End of stream, return buffer.
return &FixedReStream{Buf: buf}, nil
case chunk > 0: // Non-zero chunk, read that many elements from the stream, and buffer them.
- buf, err = readStreamToBuffer(cv, r, chunk, buf)
+ chunkBuf := make([]FullValue, 0, chunk)
+ chunkBuf, err = readStreamToBuffer(cv, r, chunk, chunkBuf)
if err != nil {
return nil, err
}
+ buf = append(buf, chunkBuf...)
case chunk == -1: // State backed iterable!
chunk, err := coder.DecodeVarInt(r)
if err != nil {
diff --git a/sdks/go/pkg/beam/metrics.go b/sdks/go/pkg/beam/metrics.go
index 67da5b1..0d4c0bb 100644
--- a/sdks/go/pkg/beam/metrics.go
+++ b/sdks/go/pkg/beam/metrics.go
@@ -21,8 +21,15 @@
"github.com/apache/beam/sdks/go/pkg/beam/core/metrics"
)
+// Implementation Note: The wrapping of the embedded methods
+// is to allow better GoDocs for the methods on the proxy types.
+
// Counter is a metric that can be incremented and decremented,
// and is aggregated by the sum.
+//
+// Counters are safe to use in multiple bundles simultaneously, but
+// not generally threadsafe. Your DoFn needs to manage the thread
+// safety of Beam metrics for any additional concurrency it uses.
type Counter struct {
*metrics.Counter
}
@@ -44,6 +51,10 @@
// Distribution is a metric that records various statistics about the distribution
// of reported values.
+//
+// Distributions are safe to use in multiple bundles simultaneously, but
+// not generally threadsafe. Your DoFn needs to manage the thread
+// safety of Beam metrics for any additional concurrency it uses.
type Distribution struct {
*metrics.Distribution
}
@@ -60,6 +71,10 @@
// Gauge is a metric that can have its new value set, and is aggregated by taking
// the last reported value.
+//
+// Gauge are safe to use in multiple bundles simultaneously, but
+// not generally threadsafe. Your DoFn needs to manage the thread
+// safety of Beam metrics for any additional concurrency it uses.
type Gauge struct {
*metrics.Gauge
}
diff --git a/sdks/go/pkg/beam/metrics_test.go b/sdks/go/pkg/beam/metrics_test.go
index 11b44101..abd19e0 100644
--- a/sdks/go/pkg/beam/metrics_test.go
+++ b/sdks/go/pkg/beam/metrics_test.go
@@ -29,8 +29,8 @@
func ctxWithPtransformID(id string) context.Context {
ctx := context.Background()
- ctx = metrics.SetPTransformID(ctx, id)
ctx = metrics.SetBundleID(ctx, "exampleBundle")
+ ctx = metrics.SetPTransformID(ctx, id)
return ctx
}
@@ -69,7 +69,7 @@
func Example_metricsReusable() {
- // Metrics can be used in multiple DoFns
+ // Metric proxies can be used in multiple DoFns
c := beam.NewCounter("example.reusable", "count")
extractWordsDofn := func(ctx context.Context, line string, emit func(string)) {
diff --git a/sdks/go/pkg/beam/partition.go b/sdks/go/pkg/beam/partition.go
index 0d13afa..071fb48 100644
--- a/sdks/go/pkg/beam/partition.go
+++ b/sdks/go/pkg/beam/partition.go
@@ -23,38 +23,53 @@
"github.com/apache/beam/sdks/go/pkg/beam/core/funcx"
"github.com/apache/beam/sdks/go/pkg/beam/core/graph"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/typex"
"github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
)
var (
- sig = &funcx.Signature{Args: []reflect.Type{TType}, Return: []reflect.Type{reflectx.Int}} // T -> int
+ sig = &funcx.Signature{Args: []reflect.Type{TType}, Return: []reflect.Type{reflectx.Int}} // T -> int
+ sigKV = &funcx.Signature{Args: []reflect.Type{TType, UType}, Return: []reflect.Type{reflectx.Int}} // KV<T, U> -> int
)
// Partition takes a PCollection<T> and a PartitionFn, uses the PartitionFn to
// split the elements of the input PCollection into N partitions, and returns
// a []PCollection<T> that bundles N PCollection<T>s containing the split elements.
+//
+// A PartitionFn has the signature `func(T) int.`
+//
+// T is permitted to be a KV.
func Partition(s Scope, n int, fn interface{}, col PCollection) []PCollection {
s = s.Scope(fmt.Sprintf("Partition(%v)", n))
if n < 1 {
panic(fmt.Sprintf("n must be > 0"))
}
- t := col.Type().Type()
- funcx.MustSatisfy(fn, funcx.Replace(sig, TType, t))
+ var emit reflect.Type
+ var in []reflect.Type
+ if typex.IsKV(col.Type()) {
+ comps := col.Type().Components()
+ k, v := comps[0].Type(), comps[1].Type()
+ funcx.MustSatisfy(fn, funcx.Replace(funcx.Replace(sigKV, TType, k), UType, v))
+ emit = reflect.FuncOf([]reflect.Type{EventTimeType, k, v}, nil, false)
+ in = []reflect.Type{EventTimeType, k, v}
+ } else {
+ t := col.Type().Type()
+ funcx.MustSatisfy(fn, funcx.Replace(sig, TType, t))
+ emit = reflect.FuncOf([]reflect.Type{EventTimeType, t}, nil, false)
+ in = []reflect.Type{EventTimeType, t}
+ }
// The partitionFn is a DoFn with a signature that is dependent on the input, so
// neither reflection nor type-specialization is adequate. Instead, it uses a
// dynamic function.
-
- emit := reflect.FuncOf([]reflect.Type{EventTimeType, t}, nil, false)
- in := []reflect.Type{EventTimeType, t}
for i := 0; i < n; i++ {
in = append(in, emit)
}
fnT := reflect.FuncOf(in, []reflect.Type{reflectx.Error}, false)
- data, err := json.Marshal(partitionData{N: n, Fn: EncodedFunc{Fn: reflectx.MakeFunc(fn)}})
+ data, err := json.Marshal(partitionData{KV: typex.IsKV(col.Type()), N: n, Fn: EncodedFunc{Fn: reflectx.MakeFunc(fn)}})
if err != nil {
panic(errors.WithContext(err, "encoding partition function"))
}
@@ -64,6 +79,7 @@
// partitionData contains the data needed for the partition DoFn generator.
type partitionData struct {
+ KV bool `json:"kv"`
N int `json:"n"`
Fn EncodedFunc `json:"fn"`
}
@@ -106,11 +122,58 @@
return []interface{}{err}
}
+// partitionFnKV is a Func with the following underlying type:
+//
+// fn : (EventTime, K, V, emit_1, emit_2, ..., emit_N) -> error
+//
+// where emit_i : (EventTime, K, V) -> () and N is given by the encoded
+// partitionData value. For any input element, it invokes to the
+// given partition function to determine which emitter to use.
+type partitionFnKV struct {
+ name string
+ t reflect.Type
+ n int
+ fnKV reflectx.Func2x1
+}
+
+func (f *partitionFnKV) Name() string {
+ return f.name
+}
+
+func (f *partitionFnKV) Type() reflect.Type {
+ return f.t
+}
+
+func (f *partitionFnKV) Call(args []interface{}) []interface{} {
+ timestamp := args[0]
+ key := args[1]
+ value := args[2]
+
+ n := f.fnKV.Call2x1(key, value).(int)
+ if n < 0 || n >= f.n {
+ return []interface{}{errors.Errorf("partitionFn(%v) = %v, want [0,%v)", value, n, f.n)}
+ }
+
+ emit := args[n+3]
+ reflectx.MakeFunc3x0(emit).Call3x0(timestamp, key, value)
+
+ var err error
+ return []interface{}{err}
+}
+
func makePartitionFn(name string, t reflect.Type, enc []byte) reflectx.Func {
var data partitionData
if err := json.Unmarshal(enc, &data); err != nil {
panic(errors.WithContext(err, "unmarshalling partitionFn data"))
}
+ if data.KV {
+ return &partitionFnKV{
+ name: name,
+ t: t,
+ n: data.N,
+ fnKV: reflectx.ToFunc2x1(data.Fn.Fn),
+ }
+ }
return &partitionFn{
name: name,
t: t,
diff --git a/sdks/go/pkg/beam/partition_test.go b/sdks/go/pkg/beam/partition_test.go
index 74a5c2a..7d285cd 100644
--- a/sdks/go/pkg/beam/partition_test.go
+++ b/sdks/go/pkg/beam/partition_test.go
@@ -28,6 +28,9 @@
beam.RegisterFunction(identityMinus2)
beam.RegisterFunction(mod2)
beam.RegisterFunction(less3)
+ beam.RegisterFunction(extractKV)
+ beam.RegisterFunction(combineKV)
+ beam.RegisterFunction(mod2Keys)
}
func identity(n int) int { return n }
@@ -87,6 +90,56 @@
}
}
+type kvIntInt struct {
+ K, V int
+}
+
+func extractKV(v kvIntInt) (int, int) {
+ return v.K, v.V
+}
+
+func combineKV(k, v int) kvIntInt {
+ return kvIntInt{k, v}
+}
+
+func mod2Keys(k, v int) int {
+ return mod2(k)
+}
+
+func TestPartitionKV(t *testing.T) {
+ tests := []struct {
+ in []kvIntInt
+ n int
+ fn interface{}
+ out0 []kvIntInt
+ }{
+ {
+ []kvIntInt{{1, 1}, {1, 2}, {2, 3}, {3, 4}},
+ 2,
+ mod2Keys,
+ []kvIntInt{{2, 3}},
+ },
+ {
+ []kvIntInt{{1, 1}, {1, 2}, {2, 3}, {3, 4}},
+ 11,
+ mod2Keys,
+ []kvIntInt{{2, 3}},
+ },
+ }
+
+ for _, test := range tests {
+ p, s, in, exp := ptest.CreateList2(test.in, test.out0)
+ kvs := beam.ParDo(s, extractKV, in)
+ parts := beam.Partition(s, test.n, test.fn, kvs)
+ out := beam.ParDo(s, combineKV, parts[0])
+ passert.Equals(s, out, exp)
+
+ if err := ptest.Run(p); err != nil {
+ t.Errorf("Partition(%v)[0] != %v: %v", test.in, test.out0, err)
+ }
+ }
+}
+
func TestPartitionFailures(t *testing.T) {
tests := []struct {
in []int
diff --git a/sdks/java/container/build.gradle b/sdks/java/container/build.gradle
index ca5b3cf..5c182c9 100644
--- a/sdks/java/container/build.gradle
+++ b/sdks/java/container/build.gradle
@@ -74,7 +74,7 @@
root: project.rootProject.hasProperty(["docker-repository-root"]) ?
project.rootProject["docker-repository-root"] : "apachebeam",
tag: project.rootProject.hasProperty(["docker-tag"]) ?
- project.rootProject["docker-tag"] : project.version)
+ project.rootProject["docker-tag"] : project.sdk_version)
dockerfile project.file("./${dockerfileName}")
files "./build/"
}
diff --git a/sdks/java/core/build.gradle b/sdks/java/core/build.gradle
index a14305d..ea0f5c9 100644
--- a/sdks/java/core/build.gradle
+++ b/sdks/java/core/build.gradle
@@ -46,6 +46,7 @@
processResources {
filter org.apache.tools.ant.filters.ReplaceTokens, tokens: [
'pom.version': version,
+ 'pom.sdk_version': sdk_version,
'timestamp': new Date().format("yyyy-MM-dd HH:mm")
]
}
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 27863ce..a125444 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
@@ -33,8 +33,9 @@
import java.util.SortedSet;
import javax.annotation.Nullable;
import org.apache.avro.AvroRuntimeException;
+import org.apache.avro.Conversion;
+import org.apache.avro.LogicalType;
import org.apache.avro.Schema;
-import org.apache.avro.data.TimeConversions.TimestampConversion;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
@@ -59,6 +60,8 @@
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Supplier;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
/**
* A {@link Coder} using Avro binary format.
@@ -239,7 +242,7 @@
@Override
public ReflectData get() {
ReflectData reflectData = new ReflectData(clazz.getClassLoader());
- reflectData.addLogicalTypeConversion(new TimestampConversion());
+ reflectData.addLogicalTypeConversion(new JodaTimestampConversion());
return reflectData;
}
}
@@ -715,4 +718,35 @@
public int hashCode() {
return Objects.hash(schemaSupplier.get(), typeDescriptor);
}
+
+ /**
+ * Conversion for DateTime.
+ *
+ * <p>This is a copy from Avro 1.8's TimestampConversion, which is renamed in Avro 1.9. Defining
+ * own copy gives flexibility for Beam Java SDK to work with Avro 1.8 and 1.9 at runtime.
+ *
+ * @see <a href="https://issues.apache.org/jira/browse/BEAM-9144">BEAM-9144: Beam's own Avro
+ * TimeConversion class in beam-sdk-java-core</a>
+ */
+ public static class JodaTimestampConversion extends Conversion<DateTime> {
+ @Override
+ public Class<DateTime> getConvertedType() {
+ return DateTime.class;
+ }
+
+ @Override
+ public String getLogicalTypeName() {
+ return "timestamp-millis";
+ }
+
+ @Override
+ public DateTime fromLong(Long millisFromEpoch, Schema schema, LogicalType type) {
+ return new DateTime(millisFromEpoch, DateTimeZone.UTC);
+ }
+
+ @Override
+ public Long toLong(DateTime timestamp, Schema schema, LogicalType type) {
+ return timestamp.getMillis();
+ }
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java
index 9359c75..889925a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoder.java
@@ -102,7 +102,8 @@
/**
* Returns a {@link SchemaCoder} for the specified class. If no schema is registered for this
- * class, then throws {@link NoSuchSchemaException}.
+ * class, then throws {@link NoSuchSchemaException}. The parameter functions to convert from and
+ * to Rows <b>must</b> implement the equals contract.
*/
public static <T> SchemaCoder<T> of(
Schema schema,
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 6784712..380e324 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
@@ -31,7 +31,7 @@
import org.apache.beam.sdk.schemas.Schema.TypeName;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.Row;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/AvroUtils.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/AvroUtils.java
index 80c337e..b9b9119 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/AvroUtils.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/utils/AvroUtils.java
@@ -20,6 +20,9 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
import java.lang.reflect.Method;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
@@ -27,6 +30,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
@@ -34,7 +38,6 @@
import org.apache.avro.LogicalType;
import org.apache.avro.LogicalTypes;
import org.apache.avro.Schema.Type;
-import org.apache.avro.data.TimeConversions;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericFixed;
import org.apache.avro.generic.GenericRecord;
@@ -47,6 +50,7 @@
import org.apache.avro.util.Utf8;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AvroCoder;
+import org.apache.beam.sdk.coders.AvroCoder.JodaTimestampConversion;
import org.apache.beam.sdk.schemas.AvroRecordSchema;
import org.apache.beam.sdk.schemas.FieldValueGetter;
import org.apache.beam.sdk.schemas.FieldValueTypeInformation;
@@ -90,8 +94,8 @@
static {
// This works around a bug in the Avro library (AVRO-1891) around SpecificRecord's handling
// of DateTime types.
- SpecificData.get().addLogicalTypeConversion(new TimeConversions.TimestampConversion());
- GenericData.get().addLogicalTypeConversion(new TimeConversions.TimestampConversion());
+ SpecificData.get().addLogicalTypeConversion(new JodaTimestampConversion());
+ GenericData.get().addLogicalTypeConversion(new JodaTimestampConversion());
}
// Unwrap an AVRO schema into the base type an whether it is nullable.
@@ -422,7 +426,37 @@
*/
public static SerializableFunction<GenericRecord, Row> getGenericRecordToRowFunction(
@Nullable Schema schema) {
- return g -> toBeamRowStrict(g, schema);
+ return new GenericRecordToRowFn(schema);
+ }
+
+ private static class GenericRecordToRowFn implements SerializableFunction<GenericRecord, Row> {
+ private final Schema schema;
+
+ GenericRecordToRowFn(Schema schema) {
+ this.schema = schema;
+ }
+
+ @Override
+ public Row apply(GenericRecord input) {
+ return toBeamRowStrict(input, schema);
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ GenericRecordToRowFn that = (GenericRecordToRowFn) other;
+ return schema.equals(that.schema);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(schema);
+ }
}
/**
@@ -431,7 +465,50 @@
*/
public static SerializableFunction<Row, GenericRecord> getRowToGenericRecordFunction(
@Nullable org.apache.avro.Schema avroSchema) {
- return g -> toGenericRecord(g, avroSchema);
+ return new RowToGenericRecordFn(avroSchema);
+ }
+
+ private static class RowToGenericRecordFn implements SerializableFunction<Row, GenericRecord> {
+ private transient org.apache.avro.Schema avroSchema;
+
+ RowToGenericRecordFn(@Nullable org.apache.avro.Schema avroSchema) {
+ this.avroSchema = avroSchema;
+ }
+
+ @Override
+ public GenericRecord apply(Row input) {
+ return toGenericRecord(input, avroSchema);
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ RowToGenericRecordFn that = (RowToGenericRecordFn) other;
+ return avroSchema.equals(that.avroSchema);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(avroSchema);
+ }
+
+ private void writeObject(ObjectOutputStream out) throws IOException {
+ final String avroSchemaAsString = (avroSchema == null) ? null : avroSchema.toString();
+ out.writeObject(avroSchemaAsString);
+ }
+
+ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+ final String avroSchemaAsString = (String) in.readObject();
+ avroSchema =
+ (avroSchemaAsString == null)
+ ? null
+ : new org.apache.avro.Schema.Parser().parse(avroSchemaAsString);
+ }
}
/**
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/Timer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/Timer.java
index 4da1278..ba996b7 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/Timer.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/Timer.java
@@ -66,4 +66,10 @@
* period}.
*/
Timer align(Duration period);
+
+ /**
+ * Sets event time timer's output timestamp. Output watermark will be held at this timestamp until
+ * the timer fires.
+ */
+ Timer withOutputTimestamp(Instant outputTime);
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/TimerMap.java
similarity index 65%
copy from runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
copy to sdks/java/core/src/main/java/org/apache/beam/sdk/state/TimerMap.java
index a114f40..92c6d59 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/TimerMap.java
@@ -15,16 +15,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.runners.flink.translation.utils;
+package org.apache.beam.sdk.state;
-import com.fasterxml.jackson.databind.type.TypeFactory;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.joda.time.Instant;
-/** Utilities for dealing with classloading. */
-public class FlinkClassloading {
+@Experimental(Experimental.Kind.TIMERS)
+public interface TimerMap {
- public static void deleteStaticCaches() {
- // Clear cache to get rid of any references to the Flink Classloader
- // See https://jira.apache.org/jira/browse/BEAM-6460
- TypeFactory.defaultInstance().clearCache();
- }
+ void set(String timerId, Instant absoluteTime);
+
+ Timer get(String timerId);
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/TimerSpecs.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/TimerSpecs.java
index 32bc54c..d5a2159 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/TimerSpecs.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/TimerSpecs.java
@@ -29,6 +29,10 @@
return new AutoValue_TimerSpecs_SimpleTimerSpec(timeDomain);
}
+ public static TimerSpec timerMap(TimeDomain timeDomain) {
+ return new AutoValue_TimerSpecs_SimpleTimerSpec(timeDomain);
+ }
+
/** A straightforward POJO {@link TimerSpec}. Package-level access for AutoValue. */
@AutoValue
abstract static class SimpleTimerSpec implements TimerSpec {
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/FileChecksumMatcher.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/FileChecksumMatcher.java
index 28349ab..497d69b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/FileChecksumMatcher.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/FileChecksumMatcher.java
@@ -18,17 +18,12 @@
package org.apache.beam.sdk.testing;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
-import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
-import java.util.regex.Pattern;
import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.util.FluentBackoff;
-import org.apache.beam.sdk.util.NumberedShardedFile;
import org.apache.beam.sdk.util.ShardedFile;
import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
@@ -41,25 +36,25 @@
import org.slf4j.LoggerFactory;
/**
- * Matcher to verify file checksum in E2E test.
+ * Matcher to verify checksum of the contents of an {@link ShardedFile} in E2E test.
*
* <p>For example:
*
* <pre>{@code
- * assertThat(job, new FileChecksumMatcher(checksumString, filePath));
+ * assertThat(new NumberedShardedFile(filePath), fileContentsHaveChecksum(checksumString));
* }</pre>
*
* or
*
* <pre>{@code
- * assertThat(job, new FileChecksumMatcher(checksumString, filePath, shardTemplate));
+ * assertThat(new NumberedShardedFile(filePath, shardTemplate), fileContentsHaveChecksum(checksumString));
* }</pre>
*
* <p>Checksum of outputs is generated based on SHA-1 algorithm. If output file is empty, SHA-1 hash
* of empty string (da39a3ee5e6b4b0d3255bfef95601890afd80709) is used as expected.
*/
-public class FileChecksumMatcher extends TypeSafeMatcher<PipelineResult>
- implements SerializableMatcher<PipelineResult> {
+public class FileChecksumMatcher extends TypeSafeMatcher<ShardedFile>
+ implements SerializableMatcher<ShardedFile> {
private static final Logger LOG = LoggerFactory.getLogger(FileChecksumMatcher.class);
@@ -70,83 +65,40 @@
.withInitialBackoff(DEFAULT_SLEEP_DURATION)
.withMaxRetries(MAX_READ_RETRIES);
- private static final Pattern DEFAULT_SHARD_TEMPLATE =
- Pattern.compile("(?x) \\S* (?<shardnum> \\d+) -of- (?<numshards> \\d+)");
-
private final String expectedChecksum;
- private final ShardedFile shardedFile;
+ private String actualChecksum;
- /** Access via {@link #getActualChecksum()}. */
- @Nullable private String actualChecksum;
-
- /**
- * Constructor that uses default shard template.
- *
- * @param checksum expected checksum string used to verify file content.
- * @param filePath path of files that's to be verified.
- */
- public FileChecksumMatcher(String checksum, String filePath) {
- this(checksum, filePath, DEFAULT_SHARD_TEMPLATE);
- }
-
- /**
- * Constructor using a custom shard template.
- *
- * @param checksum expected checksum string used to verify file content.
- * @param filePath path of files that's to be verified.
- * @param shardTemplate template of shard name to parse out the total number of shards which is
- * used in I/O retry to avoid inconsistency of filesystem. Customized template should assign
- * name "numshards" to capturing group - total shard number.
- */
- public FileChecksumMatcher(String checksum, String filePath, Pattern shardTemplate) {
+ private FileChecksumMatcher(String checksum) {
checkArgument(
!Strings.isNullOrEmpty(checksum), "Expected valid checksum, but received %s", checksum);
- checkArgument(
- !Strings.isNullOrEmpty(filePath), "Expected valid file path, but received %s", filePath);
- checkNotNull(
- shardTemplate,
- "Expected non-null shard pattern. "
- + "Please call the other constructor to use default pattern: %s",
- DEFAULT_SHARD_TEMPLATE);
-
this.expectedChecksum = checksum;
- this.shardedFile = new NumberedShardedFile(filePath, shardTemplate);
}
- /**
- * Constructor using an entirely custom {@link ShardedFile} implementation.
- *
- * <p>For internal use only.
- */
- public FileChecksumMatcher(String expectedChecksum, ShardedFile shardedFile) {
- this.expectedChecksum = expectedChecksum;
- this.shardedFile = shardedFile;
+ public static FileChecksumMatcher fileContentsHaveChecksum(String checksum) {
+ return new FileChecksumMatcher(checksum);
}
@Override
- public boolean matchesSafely(PipelineResult pipelineResult) {
- return getActualChecksum().equals(expectedChecksum);
+ public boolean matchesSafely(ShardedFile shardedFile) {
+ return getActualChecksum(shardedFile).equals(expectedChecksum);
}
/**
- * Computes a checksum of the sharded file specified in the constructor. Not safe to call until
- * the writing is complete.
+ * Computes a checksum of the given sharded file. Not safe to call until the writing is complete.
*/
- private String getActualChecksum() {
- if (actualChecksum == null) {
- // Load output data
- List<String> outputs;
- try {
- outputs = shardedFile.readFilesWithRetries(Sleeper.DEFAULT, BACK_OFF_FACTORY.backoff());
- } catch (Exception e) {
- throw new RuntimeException(String.format("Failed to read from: %s", shardedFile), e);
- }
-
- // Verify outputs. Checksum is computed using SHA-1 algorithm
- actualChecksum = computeHash(outputs);
- LOG.debug("Generated checksum: {}", actualChecksum);
+ private String getActualChecksum(ShardedFile shardedFile) {
+ // Load output data
+ List<String> outputs;
+ try {
+ outputs = shardedFile.readFilesWithRetries(Sleeper.DEFAULT, BACK_OFF_FACTORY.backoff());
+ } catch (Exception e) {
+ throw new RuntimeException(String.format("Failed to read from: %s", shardedFile), e);
}
+ // Verify outputs. Checksum is computed using SHA-1 algorithm
+ actualChecksum = computeHash(outputs);
+ LOG.debug("Generated checksum: {}", actualChecksum);
+
return actualChecksum;
}
@@ -168,7 +120,7 @@
}
@Override
- public void describeMismatchSafely(PipelineResult pResult, Description description) {
- description.appendText("was (").appendText(getActualChecksum()).appendText(")");
+ public void describeMismatchSafely(ShardedFile shardedFile, Description description) {
+ description.appendText("was (").appendText(actualChecksum).appendText(")");
}
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesTestStreamWithOutputTimestamp.java
similarity index 64%
copy from runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
copy to sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesTestStreamWithOutputTimestamp.java
index a114f40..e498ac8 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesTestStreamWithOutputTimestamp.java
@@ -15,16 +15,10 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.runners.flink.translation.utils;
+package org.apache.beam.sdk.testing;
-import com.fasterxml.jackson.databind.type.TypeFactory;
-
-/** Utilities for dealing with classloading. */
-public class FlinkClassloading {
-
- public static void deleteStaticCaches() {
- // Clear cache to get rid of any references to the Flink Classloader
- // See https://jira.apache.org/jira/browse/BEAM-6460
- TypeFactory.defaultInstance().clearCache();
- }
-}
+/**
+ * Category tag for validation tests which use outputTimestamp. Tests tagged with {@link
+ * UsesTestStreamWithOutputTimestamp} should be run for runners which support outputTimestamp.
+ */
+public interface UsesTestStreamWithOutputTimestamp {}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesTimerMap.java
similarity index 64%
copy from runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
copy to sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesTimerMap.java
index a114f40..2b7bb67 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/UsesTimerMap.java
@@ -15,16 +15,9 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.runners.flink.translation.utils;
-
-import com.fasterxml.jackson.databind.type.TypeFactory;
-
-/** Utilities for dealing with classloading. */
-public class FlinkClassloading {
-
- public static void deleteStaticCaches() {
- // Clear cache to get rid of any references to the Flink Classloader
- // See https://jira.apache.org/jira/browse/BEAM-6460
- TypeFactory.defaultInstance().clearCache();
- }
-}
+package org.apache.beam.sdk.testing;
+/**
+ * Category tag for validation tests which use timerMap. Tests tagged with {@link UsesTimerMap}
+ * should be run for runners which support timerMap.
+ */
+public interface UsesTimerMap {}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
index 27e8517..61fb365 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
@@ -279,6 +279,9 @@
/** Returns the timestamp of the current timer. */
public abstract Instant timestamp();
+ /** Returns the output timestamp of the current timer. */
+ public abstract Instant fireTimestamp();
+
/** Returns the window in which the timer is firing. */
public abstract BoundedWindow window();
@@ -433,6 +436,16 @@
@Experimental(Kind.TIMERS)
public @interface TimerId {
/** The timer ID. */
+ String value() default "";
+ }
+
+ /** Parameter annotation for the TimerMap for a {@link ProcessElement} method. */
+ @Documented
+ @Retention(RetentionPolicy.RUNTIME)
+ @Target({ElementType.FIELD, ElementType.PARAMETER})
+ @Experimental(Kind.TIMERS)
+ public @interface TimerFamily {
+ /** The TimerMap tag ID. */
String value();
}
@@ -464,6 +477,25 @@
}
/**
+ * Annotation for registering a callback for a timerFamily.
+ *
+ * <p>See the javadoc for {@link TimerFamily} for use in a full example.
+ *
+ * <p>The method annotated with {@code @OnTimerFamily} may have parameters according to the same
+ * logic as {@link ProcessElement}, but limited to the {@link BoundedWindow}, {@link State}
+ * subclasses, and {@link org.apache.beam.sdk.state.TimerMap}. State and timer parameters must be
+ * annotated with their {@link StateId} and {@link TimerId} respectively.
+ */
+ @Documented
+ @Retention(RetentionPolicy.RUNTIME)
+ @Target(ElementType.METHOD)
+ @Experimental(Kind.TIMERS)
+ public @interface OnTimerFamily {
+ /** The timer ID. */
+ String value();
+ }
+
+ /**
* Annotation for the method to use for performing actions on window expiration. For example,
* users can use this annotation to write a method that extracts a value saved in a state before
* it gets garbage collected on window expiration.
@@ -618,7 +650,10 @@
@Target(ElementType.PARAMETER)
public @interface Element {}
- /** Parameter annotation for the input element timestamp for a {@link ProcessElement} method. */
+ /**
+ * Parameter annotation for the input element timestamp for {@link ProcessElement}, {@link
+ * GetInitialRestriction}, {@link SplitRestriction}, and {@link NewTracker} methods.
+ */
@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.PARAMETER)
@@ -723,7 +758,23 @@
* Annotation for the method that maps an element to an initial restriction for a <a
* href="https://s.apache.org/splittable-do-fn">splittable</a> {@link DoFn}.
*
- * <p>Signature: {@code RestrictionT getInitialRestriction(InputT element);}
+ * <p>Signature: {@code RestrictionT getInitialRestriction(InputT element, <optional arguments>);}
+ *
+ * <p>The optional arguments are allowed to be:
+ *
+ * <ul>
+ * <li>If one of its arguments is tagged with the {@link Timestamp} annotation, then it will be
+ * passed the timestamp of the current element being processed; the argument must be of type
+ * {@link Instant}.
+ * <li>If one of its arguments is a subtype of {@link BoundedWindow}, then it will be passed the
+ * window of the current element. When applied by {@link ParDo} the subtype of {@link
+ * BoundedWindow} must match the type of windows on the input {@link PCollection}. If the
+ * window is not accessed a runner may perform additional optimizations.
+ * <li>If one of its arguments is of type {@link PaneInfo}, then it will be passed information
+ * about the current triggering pane.
+ * <li>If one of the parameters is of type {@link PipelineOptions}, then it will be passed the
+ * options for the current pipeline.
+ * </ul>
*/
// TODO: Make the InputT parameter optional.
@Documented
@@ -788,7 +839,23 @@
* be processed in parallel.
*
* <p>Signature: {@code void splitRestriction(InputT element, RestrictionT restriction,
- * OutputReceiver<RestrictionT> receiver);}
+ * OutputReceiver<RestrictionT> receiver, <optional arguments>);}
+ *
+ * <p>The optional arguments are allowed to be:
+ *
+ * <ul>
+ * <li>If one of its arguments is tagged with the {@link Timestamp} annotation, then it will be
+ * passed the timestamp of the current element being processed; the argument must be of type
+ * {@link Instant}.
+ * <li>If one of its arguments is a subtype of {@link BoundedWindow}, then it will be passed the
+ * window of the current element. When applied by {@link ParDo} the subtype of {@link
+ * BoundedWindow} must match the type of windows on the input {@link PCollection}. If the
+ * window is not accessed a runner may perform additional optimizations.
+ * <li>If one of its arguments is of type {@link PaneInfo}, then it will be passed information
+ * about the current triggering pane.
+ * <li>If one of the parameters is of type {@link PipelineOptions}, then it will be passed the
+ * options for the current pipeline.
+ * </ul>
*
* <p>Optional: if this method is omitted, the restriction will not be split (equivalent to
* defining the method and outputting the {@code restriction} unchanged).
@@ -804,8 +871,25 @@
* Annotation for the method that creates a new {@link RestrictionTracker} for the restriction of
* a <a href="https://s.apache.org/splittable-do-fn">splittable</a> {@link DoFn}.
*
- * <p>Signature: {@code MyRestrictionTracker newTracker(RestrictionT restriction);} where {@code
- * MyRestrictionTracker} must be a subtype of {@code RestrictionTracker<RestrictionT>}.
+ * <p>Signature: {@code MyRestrictionTracker newTracker(RestrictionT restriction, <optional
+ * arguments>);} where {@code MyRestrictionTracker} must be a subtype of {@code
+ * RestrictionTracker<RestrictionT>}.
+ *
+ * <p>The optional arguments are allowed to be:
+ *
+ * <ul>
+ * <li>If one of its arguments is tagged with the {@link Timestamp} annotation, then it will be
+ * passed the timestamp of the current element being processed; the argument must be of type
+ * {@link Instant}.
+ * <li>If one of its arguments is a subtype of {@link BoundedWindow}, then it will be passed the
+ * window of the current element. When applied by {@link ParDo} the subtype of {@link
+ * BoundedWindow} must match the type of windows on the input {@link PCollection}. If the
+ * window is not accessed a runner may perform additional optimizations.
+ * <li>If one of its arguments is of type {@link PaneInfo}, then it will be passed information
+ * about the current triggering pane.
+ * <li>If one of the parameters is of type {@link PipelineOptions}, then it will be passed the
+ * options for the current pipeline.
+ * </ul>
*/
@Documented
@Retention(RetentionPolicy.RUNTIME)
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
index 3bec152..5fb574d 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
@@ -33,6 +33,7 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
import org.apache.beam.sdk.transforms.DoFn.OnTimerContext;
@@ -267,6 +268,12 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException(
+ "Cannot access timerId as parameter outside of @OnTimer method.");
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Not expected to access TimeDomain from @ProcessElement");
@@ -307,6 +314,11 @@
public Timer timer(String timerId) {
throw new UnsupportedOperationException("DoFnTester doesn't support timers yet");
}
+
+ @Override
+ public TimerMap timerFamily(String tagId) {
+ throw new UnsupportedOperationException("DoFnTester doesn't support timerFamily yet");
+ }
});
} catch (UserCodeException e) {
unwrapUserCodeException(e);
@@ -688,6 +700,11 @@
}
@Override
+ public Void dispatch(DoFnSignature.Parameter.TimerIdParameter p) {
+ return null;
+ }
+
+ @Override
protected Void dispatchDefault(DoFnSignature.Parameter p) {
throw new UnsupportedOperationException(
"Parameter " + p + " not supported by DoFnTester");
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
index aec80c4..bca2eb4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
@@ -548,6 +548,9 @@
for (OnTimerMethod method : signature.onTimerMethods().values()) {
validateWindowTypeForMethod(actualWindowT, method);
}
+ for (DoFnSignature.OnTimerFamilyMethod method : signature.onTimerFamilyMethods().values()) {
+ validateWindowTypeForMethod(actualWindowT, method);
+ }
}
private static void validateWindowTypeForMethod(
@@ -586,6 +589,15 @@
"%s is splittable and uses timers, but these are not compatible",
fn.getClass().getName()));
}
+
+ // TimerFamily is semantically incompatible with splitting
+ if (!signature.timerFamilyDeclarations().isEmpty()
+ && signature.processElement().isSplittable()) {
+ throw new UnsupportedOperationException(
+ String.format(
+ "%s is splittable and uses timer family, but these are not compatible",
+ fn.getClass().getName()));
+ }
}
/**
* Extract information on how the DoFn uses schemas. In particular, if the schema of an element
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
index 2f54c27..33013fe 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
@@ -33,6 +33,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.ProcessElement;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.OnTimerMethod;
@@ -50,6 +51,7 @@
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.StateParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TaggedOutputReceiverParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimeDomainParameter;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimerFamilyParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimerParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimestampParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.WindowParameter;
@@ -114,6 +116,8 @@
public static final String STATE_PARAMETER_METHOD = "state";
public static final String TIMER_PARAMETER_METHOD = "timer";
public static final String SIDE_INPUT_PARAMETER_METHOD = "sideInput";
+ public static final String TIMER_FAMILY_PARAMETER_METHOD = "timerFamily";
+ public static final String TIMER_ID_PARAMETER_METHOD = "timerId";
/**
* Returns a {@link ByteBuddyDoFnInvokerFactory} shared with all other invocations, so that its
@@ -185,7 +189,9 @@
@Override
public void invokeOnTimer(
- String timerId, DoFnInvoker.ArgumentProvider<InputT, OutputT> arguments) {
+ String timerId,
+ String timerFamilyId,
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> arguments) {
@Nullable OnTimerInvoker onTimerInvoker = onTimerInvokers.get(timerId);
if (onTimerInvoker != null) {
@@ -206,6 +212,59 @@
}
}
+ /**
+ * Internal base class for generated {@link DoFnInvoker} instances.
+ *
+ * <p>This class should <i>not</i> be extended directly, or by Beam users. It must be public for
+ * generated instances to have adequate access, as they are generated "inside" the invoked {@link
+ * DoFn} class.
+ */
+ public abstract static class DoFnInvokerTimerFamily<
+ InputT, OutputT, DoFnT extends DoFn<InputT, OutputT>>
+ implements DoFnInvoker<InputT, OutputT> {
+ protected DoFnT delegate;
+
+ private Map<String, OnTimerInvoker> onTimerInvokers = new HashMap<>();
+
+ public DoFnInvokerTimerFamily(DoFnT delegate) {
+ this.delegate = delegate;
+ }
+
+ /**
+ * Associates the given timerFamily ID with the given {@link OnTimerInvoker}.
+ *
+ * <p>ByteBuddy does not like to generate conditional code, so we use a map + lookup of the
+ * timer ID rather than a generated conditional branch to choose which OnTimerInvoker to invoke.
+ */
+ void addOnTimerFamilyInvoker(String timerFamilyId, OnTimerInvoker onTimerInvoker) {
+ this.onTimerInvokers.put(timerFamilyId, onTimerInvoker);
+ }
+
+ @Override
+ public void invokeOnTimer(
+ String timerId,
+ String timerFamilyId,
+ DoFnInvoker.ArgumentProvider<InputT, OutputT> arguments) {
+ @Nullable OnTimerInvoker onTimerInvoker = onTimerInvokers.get(timerFamilyId);
+
+ if (onTimerInvoker != null) {
+ onTimerInvoker.invokeOnTimer(arguments);
+ } else {
+ throw new IllegalArgumentException(
+ String.format(
+ "Attempted to invoke timerFamily %s on %s, but that timerFamily is not registered."
+ + " This is the responsibility of the runner, which must only deliver"
+ + " registered timers.",
+ timerFamilyId, delegate.getClass().getName()));
+ }
+ }
+
+ @Override
+ public DoFn<InputT, OutputT> getFn() {
+ return delegate;
+ }
+ }
+
/** @return the {@link DoFnInvoker} for the given {@link DoFn}. */
public <InputT, OutputT> DoFnInvoker<InputT, OutputT> newByteBuddyInvoker(
DoFnSignature signature, DoFn<InputT, OutputT> fn) {
@@ -216,17 +275,33 @@
fn.getClass());
try {
- @SuppressWarnings("unchecked")
- DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>> invoker =
- (DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>>)
- getByteBuddyInvokerConstructor(signature).newInstance(fn);
+ if (signature.timerFamilyDeclarations().size() > 0) {
+ @SuppressWarnings("unchecked")
+ DoFnInvokerTimerFamily<InputT, OutputT, DoFn<InputT, OutputT>> invoker =
+ (DoFnInvokerTimerFamily<InputT, OutputT, DoFn<InputT, OutputT>>)
+ getByteBuddyInvokerConstructor(signature).newInstance(fn);
- for (OnTimerMethod onTimerMethod : signature.onTimerMethods().values()) {
- invoker.addOnTimerInvoker(
- onTimerMethod.id(), OnTimerInvokers.forTimer(fn, onTimerMethod.id()));
+ for (DoFnSignature.OnTimerFamilyMethod onTimerFamilyMethod :
+ signature.onTimerFamilyMethods().values()) {
+ invoker.addOnTimerFamilyInvoker(
+ onTimerFamilyMethod.id(),
+ OnTimerInvokers.forTimerFamily(fn, onTimerFamilyMethod.id()));
+ }
+ return invoker;
+ } else {
+
+ @SuppressWarnings("unchecked")
+ DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>> invoker =
+ (DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>>)
+ getByteBuddyInvokerConstructor(signature).newInstance(fn);
+
+ for (OnTimerMethod onTimerMethod : signature.onTimerMethods().values()) {
+ invoker.addOnTimerInvoker(
+ onTimerMethod.id(), OnTimerInvokers.forTimer(fn, onTimerMethod.id()));
+ }
+ return invoker;
}
- return invoker;
} catch (InstantiationException
| IllegalAccessException
| IllegalArgumentException
@@ -246,7 +321,12 @@
Class<? extends DoFn<?, ?>> fnClass = signature.fnClass();
Constructor<?> constructor = byteBuddyInvokerConstructorCache.get(fnClass);
if (constructor == null) {
- Class<? extends DoFnInvoker<?, ?>> invokerClass = generateInvokerClass(signature);
+ Class<? extends DoFnInvoker<?, ?>> invokerClass =
+ generateInvokerClass(
+ signature,
+ signature.timerFamilyDeclarations().size() > 0
+ ? DoFnInvokerTimerFamily.class
+ : DoFnInvokerBase.class);
try {
constructor = invokerClass.getConstructor(fnClass);
} catch (IllegalArgumentException | NoSuchMethodException | SecurityException e) {
@@ -293,7 +373,8 @@
}
/** Generates a {@link DoFnInvoker} class for the given {@link DoFnSignature}. */
- private static Class<? extends DoFnInvoker<?, ?>> generateInvokerClass(DoFnSignature signature) {
+ private static Class<? extends DoFnInvoker<?, ?>> generateInvokerClass(
+ DoFnSignature signature, Class<? extends DoFnInvoker> clazz) {
Class<? extends DoFn<?, ?>> fnClass = signature.fnClass();
final TypeDescription clazzDescription = new TypeDescription.ForLoadedType(fnClass);
@@ -307,12 +388,12 @@
.withSuffix(DoFnInvoker.class.getSimpleName()))
// class <invoker class> extends DoFnInvokerBase {
- .subclass(DoFnInvokerBase.class, ConstructorStrategy.Default.NO_CONSTRUCTORS)
+ .subclass(clazz, ConstructorStrategy.Default.NO_CONSTRUCTORS)
// public <invoker class>(<fn class> delegate) { this.delegate = delegate; }
.defineConstructor(Visibility.PUBLIC)
.withParameter(fnClass)
- .intercept(new InvokerConstructor())
+ .intercept(new InvokerConstructor(clazz))
// public invokeProcessElement(ProcessContext, ExtraContextFactory) {
// delegate.<@ProcessElement>(... pass just the right args ...);
@@ -786,6 +867,16 @@
}
@Override
+ public StackManipulation dispatch(TimerFamilyParameter p) {
+ return new StackManipulation.Compound(
+ new TextConstant(p.referent().id()),
+ MethodInvocation.invoke(
+ getExtraContextFactoryMethodDescription(
+ TIMER_FAMILY_PARAMETER_METHOD, String.class)),
+ TypeCasting.to(new TypeDescription.ForLoadedType(TimerMap.class)));
+ }
+
+ @Override
public StackManipulation dispatch(DoFnSignature.Parameter.PipelineOptionsParameter p) {
return simpleExtraContextParameter(PIPELINE_OPTIONS_PARAMETER_METHOD);
}
@@ -799,6 +890,15 @@
SIDE_INPUT_PARAMETER_METHOD, String.class)),
TypeCasting.to(new TypeDescription.ForLoadedType(p.elementT().getRawType())));
}
+
+ @Override
+ public StackManipulation dispatch(DoFnSignature.Parameter.TimerIdParameter p) {
+ return new StackManipulation.Compound(
+ pushDelegate,
+ MethodInvocation.invoke(
+ getExtraContextFactoryMethodDescription(
+ TIMER_ID_PARAMETER_METHOD, DoFn.class)));
+ }
});
}
@@ -984,6 +1084,12 @@
* for a constructor that takes a single argument and assigns it to the delegate field.
*/
private static final class InvokerConstructor implements Implementation {
+ Class<? extends DoFnInvoker> clazz;
+
+ InvokerConstructor(Class<? extends DoFnInvoker> clazz) {
+ this.clazz = clazz;
+ }
+
@Override
public InstrumentedType prepare(InstrumentedType instrumentedType) {
return instrumentedType;
@@ -1000,7 +1106,7 @@
MethodVariableAccess.REFERENCE.loadFrom(1),
// Invoke the super constructor (default constructor of Object)
MethodInvocation.invoke(
- new TypeDescription.ForLoadedType(DoFnInvokerBase.class)
+ new TypeDescription.ForLoadedType(clazz)
.getDeclaredMethods()
.filter(
ElementMatchers.isConstructor()
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java
index b997653..b305d78 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java
@@ -81,6 +81,31 @@
}
}
+ public <InputT, OutputT> OnTimerInvoker<InputT, OutputT> forTimerFamily(
+ DoFn<InputT, OutputT> fn, String timerId) {
+
+ @SuppressWarnings("unchecked")
+ Class<? extends DoFn<?, ?>> fnClass = (Class<? extends DoFn<?, ?>>) fn.getClass();
+ try {
+ OnTimerMethodSpecifier onTimerMethodSpecifier =
+ OnTimerMethodSpecifier.forClassAndTimerId(fnClass, timerId);
+ Constructor<?> constructor = constructorTimerFamilyCache.get(onTimerMethodSpecifier);
+
+ return (OnTimerInvoker<InputT, OutputT>) constructor.newInstance(fn);
+ } catch (InstantiationException
+ | IllegalAccessException
+ | IllegalArgumentException
+ | InvocationTargetException
+ | SecurityException
+ | ExecutionException e) {
+ throw new RuntimeException(
+ String.format(
+ "Unable to construct @%s invoker for %s",
+ DoFn.OnTimerFamily.class.getSimpleName(), fn.getClass().getName()),
+ e);
+ }
+ }
+
public static ByteBuddyOnTimerInvokerFactory only() {
return INSTANCE;
}
@@ -120,6 +145,34 @@
}
}
});
+
+ /**
+ * A cache of constructors of generated {@link OnTimerInvoker} classes, keyed by {@link
+ * OnTimerMethodSpecifier}.
+ *
+ * <p>Needed because generating an invoker class is expensive, and to avoid generating an
+ * excessive number of classes consuming PermGen memory in Java's that still have PermGen.
+ */
+ private final LoadingCache<OnTimerMethodSpecifier, Constructor<?>> constructorTimerFamilyCache =
+ CacheBuilder.newBuilder()
+ .build(
+ new CacheLoader<OnTimerMethodSpecifier, Constructor<?>>() {
+ @Override
+ public Constructor<?> load(final OnTimerMethodSpecifier onTimerMethodSpecifier)
+ throws Exception {
+ DoFnSignature signature =
+ DoFnSignatures.getSignature(onTimerMethodSpecifier.fnClass());
+ Class<? extends OnTimerInvoker<?, ?>> invokerClass =
+ generateOnTimerFamilyInvokerClass(
+ signature, onTimerMethodSpecifier.timerId());
+ try {
+ return invokerClass.getConstructor(signature.fnClass());
+ } catch (IllegalArgumentException | NoSuchMethodException | SecurityException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ });
+
/**
* Generates a {@link OnTimerInvoker} class for the given {@link DoFnSignature} and {@link
* TimerId}.
@@ -176,6 +229,89 @@
return res;
}
+ private static Class<? extends OnTimerInvoker<?, ?>> generateOnTimerFamilyInvokerClass(
+ DoFnSignature signature, String timerId) {
+ Class<? extends DoFn<?, ?>> fnClass = signature.fnClass();
+
+ final TypeDescription clazzDescription = new TypeDescription.ForLoadedType(fnClass);
+
+ final String suffix =
+ String.format(
+ "%s$%s$%s",
+ OnTimerInvoker.class.getSimpleName(),
+ CharMatcher.javaLetterOrDigit().retainFrom(timerId),
+ BaseEncoding.base64().omitPadding().encode(timerId.getBytes(Charsets.UTF_8)));
+
+ DynamicType.Builder<?> builder =
+ new ByteBuddy()
+ // Create subclasses inside the target class, to have access to
+ // private and package-private bits
+ .with(StableInvokerNamingStrategy.forDoFnClass(fnClass).withSuffix(suffix))
+
+ // class <invoker class> implements OnTimerInvoker {
+ .subclass(OnTimerInvoker.class, ConstructorStrategy.Default.NO_CONSTRUCTORS)
+
+ // private final <fn class> delegate;
+ .defineField(
+ FN_DELEGATE_FIELD_NAME, fnClass, Visibility.PRIVATE, FieldManifestation.FINAL)
+
+ // <invoker class>(<fn class> delegate) { this.delegate = delegate; }
+ .defineConstructor(Visibility.PUBLIC)
+ .withParameter(fnClass)
+ .intercept(new InvokerConstructor())
+
+ // public invokeOnTimer(DoFn.ArgumentProvider) {
+ // this.delegate.<@OnTimer method>(... pass the right args ...)
+ // }
+ .method(ElementMatchers.named("invokeOnTimer"))
+ .intercept(
+ new InvokeOnTimerFamilyDelegation(
+ clazzDescription, signature.onTimerFamilyMethods().get(timerId)));
+
+ DynamicType.Unloaded<?> unloaded = builder.make();
+
+ @SuppressWarnings("unchecked")
+ Class<? extends OnTimerInvoker<?, ?>> res =
+ (Class<? extends OnTimerInvoker<?, ?>>)
+ unloaded
+ .load(
+ findClassLoader(fnClass.getClassLoader()),
+ ClassLoadingStrategy.Default.INJECTION)
+ .getLoaded();
+ return res;
+ }
+
+ /**
+ * An "invokeOnTimer" method implementation akin to @ProcessElement, but simpler because no
+ * splitting-related parameters need to be handled.
+ */
+ private static class InvokeOnTimerFamilyDelegation
+ extends DoFnMethodWithExtraParametersDelegation {
+
+ private final DoFnSignature.OnTimerFamilyMethod signature;
+
+ public InvokeOnTimerFamilyDelegation(
+ TypeDescription clazzDescription, DoFnSignature.OnTimerFamilyMethod signature) {
+ super(clazzDescription, signature);
+ this.signature = signature;
+ }
+
+ @Override
+ public InstrumentedType prepare(InstrumentedType instrumentedType) {
+ // Remember the field description of the instrumented type.
+ // Kind of a hack to set the protected value, because the instrumentedType
+ // is only available to prepare, while we need this information in
+ // beforeDelegation
+ delegateField =
+ instrumentedType
+ .getDeclaredFields() // the delegate is declared on the OnTimerInvoker
+ .filter(ElementMatchers.named(FN_DELEGATE_FIELD_NAME))
+ .getOnly();
+ // Delegating the method call doesn't require any changes to the instrumented type.
+ return instrumentedType;
+ }
+ }
+
/**
* An "invokeOnTimer" method implementation akin to @ProcessElement, but simpler because no
* splitting-related parameters need to be handled.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
index fe31c64..5aaea28 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
@@ -23,6 +23,7 @@
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.FinishBundle;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
@@ -69,7 +70,8 @@
DoFn.ProcessContinuation invokeProcessElement(ArgumentProvider<InputT, OutputT> extra);
/** Invoke the appropriate {@link DoFn.OnTimer} method on the bound {@link DoFn}. */
- void invokeOnTimer(String timerId, ArgumentProvider<InputT, OutputT> arguments);
+ void invokeOnTimer(
+ String timerId, String timerFamilyId, ArgumentProvider<InputT, OutputT> arguments);
/** Invoke the {@link DoFn.GetInitialRestriction} method on the bound {@link DoFn}. */
@SuppressWarnings("TypeParameterUnusedInFormals")
@@ -170,6 +172,13 @@
/** Returns the timer for the given {@link TimerId}. */
Timer timer(String timerId);
+
+ /**
+ * Returns the timerMap for the given {@link org.apache.beam.sdk.transforms.DoFn.TimerFamily}.
+ */
+ TimerMap timerFamily(String tagId);
+
+ String timerId(DoFn<InputT, OutputT> doFn);
}
/**
@@ -202,6 +211,14 @@
}
@Override
+ public TimerMap timerFamily(String tagId) {
+ throw new UnsupportedOperationException(
+ String.format(
+ "Should never call non-overridden methods of %s",
+ FakeArgumentProvider.class.getSimpleName()));
+ }
+
+ @Override
public InputT schemaElement(int index) {
throw new UnsupportedOperationException(
String.format(
@@ -218,6 +235,14 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException(
+ String.format(
+ "Should never call non-overridden methods of %s",
+ FakeArgumentProvider.class.getSimpleName()));
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
String.format(
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
index 5737ac9..1ea3547 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
@@ -93,6 +93,9 @@
/** Timer declarations present on the {@link DoFn} class. Immutable. */
public abstract Map<String, TimerDeclaration> timerDeclarations();
+ /** TimerMap declarations present on the {@link DoFn} class. Immutable. */
+ public abstract Map<String, TimerFamilyDeclaration> timerFamilyDeclarations();
+
/** Field access declaration. */
@Nullable
public abstract Map<String, FieldAccessDeclaration> fieldAccessDeclarations();
@@ -117,6 +120,10 @@
@Nullable
public abstract Map<String, OnTimerMethod> onTimerMethods();
+ /** Details about this {@link DoFn}'s {@link DoFn.OnTimerFamily} methods. */
+ @Nullable
+ public abstract Map<String, OnTimerFamilyMethod> onTimerFamilyMethods();
+
/** @deprecated use {@link #usesState()}, it's cleaner */
@Deprecated
public boolean isStateful() {
@@ -130,7 +137,7 @@
/** Whether the {@link DoFn} described by this signature uses timers. */
public boolean usesTimers() {
- return timerDeclarations().size() > 0;
+ return timerDeclarations().size() > 0 || timerFamilyDeclarations().size() > 0;
}
static Builder builder() {
@@ -167,11 +174,16 @@
abstract Builder setTimerDeclarations(Map<String, TimerDeclaration> timerDeclarations);
+ abstract Builder setTimerFamilyDeclarations(
+ Map<String, TimerFamilyDeclaration> timerFamilyDeclarations);
+
abstract Builder setFieldAccessDeclarations(
Map<String, FieldAccessDeclaration> fieldAccessDeclaration);
abstract Builder setOnTimerMethods(Map<String, OnTimerMethod> onTimerMethods);
+ abstract Builder setOnTimerFamilyMethods(Map<String, OnTimerFamilyMethod> onTimerFamilyMethods);
+
abstract DoFnSignature build();
}
@@ -245,6 +257,10 @@
return cases.dispatch((TimeDomainParameter) this);
} else if (this instanceof SideInputParameter) {
return cases.dispatch((SideInputParameter) this);
+ } else if (this instanceof TimerFamilyParameter) {
+ return cases.dispatch((TimerFamilyParameter) this);
+ } else if (this instanceof TimerIdParameter) {
+ return cases.dispatch((TimerIdParameter) this);
} else {
throw new IllegalStateException(
String.format(
@@ -289,6 +305,10 @@
ResultT dispatch(SideInputParameter p);
+ ResultT dispatch(TimerFamilyParameter p);
+
+ ResultT dispatch(TimerIdParameter p);
+
/** A base class for a visitor with a default method for cases it is not interested in. */
abstract class WithDefault<ResultT> implements Cases<ResultT> {
@@ -335,6 +355,11 @@
}
@Override
+ public ResultT dispatch(TimerIdParameter p) {
+ return dispatchDefault(p);
+ }
+
+ @Override
public ResultT dispatch(TimeDomainParameter p) {
return dispatchDefault(p);
}
@@ -378,6 +403,11 @@
public ResultT dispatch(SideInputParameter p) {
return dispatchDefault(p);
}
+
+ @Override
+ public ResultT dispatch(TimerFamilyParameter p) {
+ return dispatchDefault(p);
+ }
}
}
@@ -392,6 +422,8 @@
new AutoValue_DoFnSignature_Parameter_OnTimerContextParameter();
private static final TimestampParameter TIMESTAMP_PARAMETER =
new AutoValue_DoFnSignature_Parameter_TimestampParameter();
+ private static final TimerIdParameter TIMER_ID_PARAMETER =
+ new AutoValue_DoFnSignature_Parameter_TimerIdParameter();
private static final PaneInfoParameter PANE_INFO_PARAMETER =
new AutoValue_DoFnSignature_Parameter_PaneInfoParameter();
private static final TimeDomainParameter TIME_DOMAIN_PARAMETER =
@@ -423,6 +455,10 @@
return TIMESTAMP_PARAMETER;
}
+ public static TimerIdParameter timerIdParameter() {
+ return TIMER_ID_PARAMETER;
+ }
+
public static SideInputParameter sideInputParameter(
TypeDescriptor<?> elementT, String sideInputId) {
return new AutoValue_DoFnSignature_Parameter_SideInputParameter.Builder()
@@ -476,6 +512,10 @@
return new AutoValue_DoFnSignature_Parameter_TimerParameter(decl);
}
+ public static TimerFamilyParameter timerFamilyParameter(TimerFamilyDeclaration decl) {
+ return new AutoValue_DoFnSignature_Parameter_TimerFamilyParameter(decl);
+ }
+
/** Descriptor for a {@link Parameter} of a subtype of {@link PipelineOptions}. */
@AutoValue
public abstract static class PipelineOptionsParameter extends Parameter {
@@ -564,6 +604,11 @@
TimestampParameter() {}
}
+ @AutoValue
+ public abstract static class TimerIdParameter extends Parameter {
+ TimerIdParameter() {}
+ }
+
/**
* Descriptor for a {@link Parameter} representing the time domain of a timer.
*
@@ -689,6 +734,15 @@
public abstract TimerDeclaration referent();
}
+
+ /** Descriptor for a {@link Parameter} of type {@link DoFn.TimerFamily}. */
+ @AutoValue
+ public abstract static class TimerFamilyParameter extends Parameter {
+ // Package visible for AutoValue
+ TimerFamilyParameter() {}
+
+ public abstract TimerFamilyDeclaration referent();
+ }
}
/** Describes a {@link DoFn.ProcessElement} method. */
@@ -830,6 +884,48 @@
}
}
+ /** Describes a {@link DoFn.OnTimerFamily} method. */
+ @AutoValue
+ public abstract static class OnTimerFamilyMethod implements MethodWithExtraParameters {
+
+ /** The id on the method's {@link DoFn.TimerId} annotation. */
+ public abstract String id();
+
+ /** The annotated method itself. */
+ @Override
+ public abstract Method targetMethod();
+
+ /**
+ * Whether this method requires stable input, expressed via {@link
+ * org.apache.beam.sdk.transforms.DoFn.RequiresStableInput}. For timers, this means that any
+ * state must be stably persisted prior to calling it.
+ */
+ public abstract boolean requiresStableInput();
+
+ /** The window type used by this method, if any. */
+ @Nullable
+ @Override
+ public abstract TypeDescriptor<? extends BoundedWindow> windowT();
+
+ /** Types of optional parameters of the annotated method, in the order they appear. */
+ @Override
+ public abstract List<Parameter> extraParameters();
+
+ static OnTimerFamilyMethod create(
+ Method targetMethod,
+ String id,
+ boolean requiresStableInput,
+ TypeDescriptor<? extends BoundedWindow> windowT,
+ List<Parameter> extraParameters) {
+ return new AutoValue_DoFnSignature_OnTimerFamilyMethod(
+ id,
+ targetMethod,
+ requiresStableInput,
+ windowT,
+ Collections.unmodifiableList(extraParameters));
+ }
+ }
+
/** Describes a {@link DoFn.OnWindowExpiration} method. */
@AutoValue
public abstract static class OnWindowExpirationMethod implements MethodWithExtraParameters {
@@ -883,6 +979,21 @@
}
}
+ /**
+ * Describes a timer family declaration; a field of type {@link TimerSpec} annotated with {@link
+ * DoFn.TimerFamily}.
+ */
+ @AutoValue
+ public abstract static class TimerFamilyDeclaration {
+ public abstract String id();
+
+ public abstract Field field();
+
+ static TimerFamilyDeclaration create(String id, Field field) {
+ return new AutoValue_DoFnSignature_TimerFamilyDeclaration(id, field);
+ }
+ }
+
/** Describes a {@link DoFn.StartBundle} or {@link DoFn.FinishBundle} method. */
@AutoValue
public abstract static class BundleMethod implements DoFnMethod {
@@ -945,7 +1056,7 @@
/** Describes a {@link DoFn.GetInitialRestriction} method. */
@AutoValue
- public abstract static class GetInitialRestrictionMethod implements DoFnMethod {
+ public abstract static class GetInitialRestrictionMethod implements MethodWithExtraParameters {
/** The annotated method itself. */
@Override
public abstract Method targetMethod();
@@ -953,14 +1064,28 @@
/** Type of the returned restriction. */
public abstract TypeDescriptor<?> restrictionT();
- static GetInitialRestrictionMethod create(Method targetMethod, TypeDescriptor<?> restrictionT) {
- return new AutoValue_DoFnSignature_GetInitialRestrictionMethod(targetMethod, restrictionT);
+ /** The window type used by this method, if any. */
+ @Nullable
+ @Override
+ public abstract TypeDescriptor<? extends BoundedWindow> windowT();
+
+ /** Types of optional parameters of the annotated method, in the order they appear. */
+ @Override
+ public abstract List<Parameter> extraParameters();
+
+ static GetInitialRestrictionMethod create(
+ Method targetMethod,
+ TypeDescriptor<?> restrictionT,
+ TypeDescriptor<? extends BoundedWindow> windowT,
+ List<Parameter> extraParameters) {
+ return new AutoValue_DoFnSignature_GetInitialRestrictionMethod(
+ targetMethod, restrictionT, windowT, extraParameters);
}
}
/** Describes a {@link DoFn.SplitRestriction} method. */
@AutoValue
- public abstract static class SplitRestrictionMethod implements DoFnMethod {
+ public abstract static class SplitRestrictionMethod implements MethodWithExtraParameters {
/** The annotated method itself. */
@Override
public abstract Method targetMethod();
@@ -968,14 +1093,28 @@
/** Type of the restriction taken and returned. */
public abstract TypeDescriptor<?> restrictionT();
- static SplitRestrictionMethod create(Method targetMethod, TypeDescriptor<?> restrictionT) {
- return new AutoValue_DoFnSignature_SplitRestrictionMethod(targetMethod, restrictionT);
+ /** The window type used by this method, if any. */
+ @Nullable
+ @Override
+ public abstract TypeDescriptor<? extends BoundedWindow> windowT();
+
+ /** Types of optional parameters of the annotated method, in the order they appear. */
+ @Override
+ public abstract List<Parameter> extraParameters();
+
+ static SplitRestrictionMethod create(
+ Method targetMethod,
+ TypeDescriptor<?> restrictionT,
+ TypeDescriptor<? extends BoundedWindow> windowT,
+ List<Parameter> extraParameters) {
+ return new AutoValue_DoFnSignature_SplitRestrictionMethod(
+ targetMethod, restrictionT, windowT, extraParameters);
}
}
/** Describes a {@link DoFn.NewTracker} method. */
@AutoValue
- public abstract static class NewTrackerMethod implements DoFnMethod {
+ public abstract static class NewTrackerMethod implements MethodWithExtraParameters {
/** The annotated method itself. */
@Override
public abstract Method targetMethod();
@@ -986,9 +1125,23 @@
/** Type of the returned {@link RestrictionTracker}. */
public abstract TypeDescriptor<?> trackerT();
+ /** The window type used by this method, if any. */
+ @Nullable
+ @Override
+ public abstract TypeDescriptor<? extends BoundedWindow> windowT();
+
+ /** Types of optional parameters of the annotated method, in the order they appear. */
+ @Override
+ public abstract List<Parameter> extraParameters();
+
static NewTrackerMethod create(
- Method targetMethod, TypeDescriptor<?> restrictionT, TypeDescriptor<?> trackerT) {
- return new AutoValue_DoFnSignature_NewTrackerMethod(targetMethod, restrictionT, trackerT);
+ Method targetMethod,
+ TypeDescriptor<?> restrictionT,
+ TypeDescriptor<?> trackerT,
+ TypeDescriptor<? extends BoundedWindow> windowT,
+ List<Parameter> extraParameters) {
+ return new AutoValue_DoFnSignature_NewTrackerMethod(
+ targetMethod, restrictionT, trackerT, windowT, extraParameters);
}
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
index b3fde4f..3c5faf4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
@@ -45,6 +45,7 @@
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.state.TimerSpec;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
@@ -57,10 +58,12 @@
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.RestrictionTrackerParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.StateParameter;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimerFamilyParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimerParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.WindowParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.StateDeclaration;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerDeclaration;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerFamilyDeclaration;
import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
@@ -98,11 +101,14 @@
Parameter.PipelineOptionsParameter.class,
Parameter.TimerParameter.class,
Parameter.StateParameter.class,
- Parameter.SideInputParameter.class);
+ Parameter.SideInputParameter.class,
+ Parameter.TimerFamilyParameter.class);
private static final ImmutableList<Class<? extends Parameter>>
ALLOWED_SPLITTABLE_PROCESS_ELEMENT_PARAMETERS =
ImmutableList.of(
+ Parameter.WindowParameter.class,
+ Parameter.PaneInfoParameter.class,
Parameter.PipelineOptionsParameter.class,
Parameter.ElementParameter.class,
Parameter.TimestampParameter.class,
@@ -122,7 +128,24 @@
Parameter.OutputReceiverParameter.class,
Parameter.TaggedOutputReceiverParameter.class,
Parameter.TimerParameter.class,
- Parameter.StateParameter.class);
+ Parameter.StateParameter.class,
+ Parameter.TimerFamilyParameter.class,
+ Parameter.TimerIdParameter.class);
+
+ private static final ImmutableList<Class<? extends Parameter>>
+ ALLOWED_ON_TIMER_FAMILY_PARAMETERS =
+ ImmutableList.of(
+ Parameter.OnTimerContextParameter.class,
+ Parameter.TimestampParameter.class,
+ Parameter.TimeDomainParameter.class,
+ Parameter.WindowParameter.class,
+ Parameter.PipelineOptionsParameter.class,
+ Parameter.OutputReceiverParameter.class,
+ Parameter.TaggedOutputReceiverParameter.class,
+ Parameter.TimerParameter.class,
+ Parameter.StateParameter.class,
+ Parameter.TimerFamilyParameter.class,
+ Parameter.TimerIdParameter.class);
private static final Collection<Class<? extends Parameter>>
ALLOWED_ON_WINDOW_EXPIRATION_PARAMETERS =
@@ -133,6 +156,28 @@
Parameter.TaggedOutputReceiverParameter.class,
Parameter.StateParameter.class);
+ private static final Collection<Class<? extends Parameter>>
+ ALLOWED_GET_INITIAL_RESTRICTION_PARAMETERS =
+ ImmutableList.of(
+ Parameter.WindowParameter.class,
+ Parameter.TimestampParameter.class,
+ Parameter.PaneInfoParameter.class,
+ Parameter.PipelineOptionsParameter.class);
+
+ private static final Collection<Class<? extends Parameter>> ALLOWED_SPLIT_RESTRICTION_PARAMETERS =
+ ImmutableList.of(
+ Parameter.WindowParameter.class,
+ Parameter.TimestampParameter.class,
+ Parameter.PaneInfoParameter.class,
+ Parameter.PipelineOptionsParameter.class);
+
+ private static final Collection<Class<? extends Parameter>> ALLOWED_NEW_TRACKER_PARAMETERS =
+ ImmutableList.of(
+ Parameter.WindowParameter.class,
+ Parameter.TimestampParameter.class,
+ Parameter.PaneInfoParameter.class,
+ Parameter.PipelineOptionsParameter.class);
+
/** @return the {@link DoFnSignature} for the given {@link DoFn} instance. */
public static <FnT extends DoFn<?, ?>> DoFnSignature signatureForDoFn(FnT fn) {
return getSignature(fn.getClass());
@@ -154,6 +199,7 @@
private final Map<String, StateDeclaration> stateDeclarations = new HashMap<>();
private final Map<String, TimerDeclaration> timerDeclarations = new HashMap<>();
+ private final Map<String, TimerFamilyDeclaration> timerFamilyDeclarations = new HashMap<>();
private final Map<String, FieldAccessDeclaration> fieldAccessDeclarations = new HashMap<>();
private FnAnalysisContext() {}
@@ -173,6 +219,14 @@
return Collections.unmodifiableMap(timerDeclarations);
}
+ /**
+ * TimerMap parameters declared in this context, keyed by {@link
+ * org.apache.beam.sdk.transforms.DoFn.TimerFamily}. Unmodifiable.
+ */
+ public Map<String, TimerFamilyDeclaration> getTimerFamilyDeclarations() {
+ return Collections.unmodifiableMap(timerFamilyDeclarations);
+ }
+
/** Field access declaration declared in this context. */
@Nullable
public Map<String, FieldAccessDeclaration> getFieldAccessDeclarations() {
@@ -193,12 +247,22 @@
timerDeclarations.put(decl.id(), decl);
}
+ public void addTimerFamilyDeclaration(TimerFamilyDeclaration decl) {
+ timerFamilyDeclarations.put(decl.id(), decl);
+ }
+
public void addTimerDeclarations(Iterable<TimerDeclaration> decls) {
for (TimerDeclaration decl : decls) {
addTimerDeclaration(decl);
}
}
+ public void addTimerFamilyDeclarations(Iterable<TimerFamilyDeclaration> decls) {
+ for (TimerFamilyDeclaration decl : decls) {
+ addTimerFamilyDeclaration(decl);
+ }
+ }
+
public void addFieldAccessDeclaration(FieldAccessDeclaration decl) {
fieldAccessDeclarations.put(decl.id(), decl);
}
@@ -220,6 +284,7 @@
private final Map<String, StateParameter> stateParameters = new HashMap<>();
private final Map<String, TimerParameter> timerParameters = new HashMap<>();
+ private final Map<String, TimerFamilyParameter> timerFamilyParameters = new HashMap<>();
private final List<Parameter> extraParameters = new ArrayList<>();
@Nullable private TypeDescriptor<? extends BoundedWindow> windowT;
@@ -258,6 +323,13 @@
public Map<String, TimerParameter> getTimerParameters() {
return Collections.unmodifiableMap(timerParameters);
}
+ /**
+ * TimerMap parameters declared in this context, keyed by {@link
+ * org.apache.beam.sdk.transforms.DoFn.TimerFamily}.
+ */
+ public Map<String, TimerFamilyParameter> getTimerFamilyParameters() {
+ return Collections.unmodifiableMap(timerFamilyParameters);
+ }
/** Extra parameters in their entirety. Unmodifiable. */
public List<Parameter> getExtraParameters() {
return Collections.unmodifiableList(extraParameters);
@@ -282,6 +354,10 @@
TimerParameter timerParameter = (TimerParameter) param;
timerParameters.put(timerParameter.referent().id(), timerParameter);
}
+ if (param instanceof TimerFamilyParameter) {
+ TimerFamilyParameter timerFamilyParameter = (TimerFamilyParameter) param;
+ timerFamilyParameters.put(timerFamilyParameter.referent().id(), timerFamilyParameter);
+ }
}
/** Create an empty context, with no declarations. */
@@ -340,6 +416,7 @@
FnAnalysisContext fnContext = FnAnalysisContext.create();
fnContext.addStateDeclarations(analyzeStateDeclarations(errors, fnClass).values());
fnContext.addTimerDeclarations(analyzeTimerDeclarations(errors, fnClass).values());
+ fnContext.addTimerFamilyDeclarations(analyzeTimerFamilyDeclarations(errors, fnClass).values());
fnContext.addFieldAccessDeclarations(analyzeFieldAccessDeclaration(errors, fnClass).values());
Method processElementMethod =
@@ -385,6 +462,36 @@
}
signatureBuilder.setOnTimerMethods(onTimerMethodMap);
+ // Check for TimerFamily
+ Collection<Method> onTimerFamilyMethods =
+ declaredMethodsWithAnnotation(DoFn.OnTimerFamily.class, fnClass, DoFn.class);
+ HashMap<String, DoFnSignature.OnTimerFamilyMethod> onTimerFamilyMethodMap =
+ Maps.newHashMapWithExpectedSize(onTimerFamilyMethods.size());
+
+ for (Method onTimerFamilyMethod : onTimerFamilyMethods) {
+ String id = onTimerFamilyMethod.getAnnotation(DoFn.OnTimerFamily.class).value();
+ errors.checkArgument(
+ fnContext.getTimerFamilyDeclarations().containsKey(id),
+ "Callback %s is for undeclared timerFamily %s",
+ onTimerFamilyMethod,
+ id);
+
+ TimerFamilyDeclaration timerDecl = fnContext.getTimerFamilyDeclarations().get(id);
+ errors.checkArgument(
+ timerDecl.field().getDeclaringClass().equals(getDeclaringClass(onTimerFamilyMethod)),
+ "Callback %s is for timerFamily %s declared in a different class %s."
+ + " TimerFamily callbacks must be declared in the same lexical scope as their timer",
+ onTimerFamilyMethod,
+ id,
+ timerDecl.field().getDeclaringClass().getCanonicalName());
+
+ onTimerFamilyMethodMap.put(
+ id,
+ analyzeOnTimerFamilyMethod(
+ errors, fnT, onTimerFamilyMethod, id, inputT, outputT, fnContext));
+ }
+ signatureBuilder.setOnTimerFamilyMethods(onTimerFamilyMethodMap);
+
// Check the converse - that all timers have a callback. This could be relaxed to only
// those timers used in methods, once method parameter lists support timers.
for (TimerDeclaration decl : fnContext.getTimerDeclarations().values()) {
@@ -395,6 +502,16 @@
decl.id());
}
+ // Check the converse - that all timer family have a callback.
+
+ for (TimerFamilyDeclaration decl : fnContext.getTimerFamilyDeclarations().values()) {
+ errors.checkArgument(
+ onTimerFamilyMethodMap.containsKey(decl.id()),
+ "No callback registered via %s for timerFamily %s",
+ DoFn.OnTimerFamily.class.getSimpleName(),
+ decl.id());
+ }
+
ErrorReporter processElementErrors =
errors.forMethod(DoFn.ProcessElement.class, processElementMethod);
DoFnSignature.ProcessElementMethod processElement =
@@ -438,7 +555,12 @@
errors.forMethod(DoFn.GetInitialRestriction.class, getInitialRestrictionMethod);
signatureBuilder.setGetInitialRestriction(
analyzeGetInitialRestrictionMethod(
- getInitialRestrictionErrors, fnT, getInitialRestrictionMethod, inputT));
+ getInitialRestrictionErrors,
+ fnT,
+ getInitialRestrictionMethod,
+ inputT,
+ outputT,
+ fnContext));
}
if (splitRestrictionMethod != null) {
@@ -446,7 +568,7 @@
errors.forMethod(DoFn.SplitRestriction.class, splitRestrictionMethod);
signatureBuilder.setSplitRestriction(
analyzeSplitRestrictionMethod(
- splitRestrictionErrors, fnT, splitRestrictionMethod, inputT));
+ splitRestrictionErrors, fnT, splitRestrictionMethod, inputT, outputT, fnContext));
}
if (getRestrictionCoderMethod != null) {
@@ -460,13 +582,15 @@
if (newTrackerMethod != null) {
ErrorReporter newTrackerErrors = errors.forMethod(DoFn.NewTracker.class, newTrackerMethod);
signatureBuilder.setNewTracker(
- analyzeNewTrackerMethod(newTrackerErrors, fnT, newTrackerMethod));
+ analyzeNewTrackerMethod(
+ newTrackerErrors, fnT, newTrackerMethod, inputT, outputT, fnContext));
}
signatureBuilder.setIsBoundedPerElement(inferBoundedness(fnT, processElement, errors));
signatureBuilder.setStateDeclarations(fnContext.getStateDeclarations());
signatureBuilder.setTimerDeclarations(fnContext.getTimerDeclarations());
+ signatureBuilder.setTimerFamilyDeclarations(fnContext.getTimerFamilyDeclarations());
signatureBuilder.setFieldAccessDeclarations(fnContext.getFieldAccessDeclarations());
DoFnSignature signature = signatureBuilder.build();
@@ -747,6 +871,51 @@
}
@VisibleForTesting
+ static DoFnSignature.OnTimerFamilyMethod analyzeOnTimerFamilyMethod(
+ ErrorReporter errors,
+ TypeDescriptor<? extends DoFn<?, ?>> fnClass,
+ Method m,
+ String timerFamilyId,
+ TypeDescriptor<?> inputT,
+ TypeDescriptor<?> outputT,
+ FnAnalysisContext fnContext) {
+ errors.checkArgument(void.class.equals(m.getReturnType()), "Must return void");
+
+ Type[] params = m.getGenericParameterTypes();
+
+ MethodAnalysisContext methodContext = MethodAnalysisContext.create();
+
+ boolean requiresStableInput = m.isAnnotationPresent(DoFn.RequiresStableInput.class);
+
+ @Nullable TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnClass, m);
+
+ List<DoFnSignature.Parameter> extraParameters = new ArrayList<>();
+ ErrorReporter onTimerErrors = errors.forMethod(DoFn.OnTimerFamily.class, m);
+ for (int i = 0; i < params.length; ++i) {
+ Parameter parameter =
+ analyzeExtraParameter(
+ onTimerErrors,
+ fnContext,
+ methodContext,
+ fnClass,
+ ParameterDescription.of(
+ m,
+ i,
+ fnClass.resolveType(params[i]),
+ Arrays.asList(m.getParameterAnnotations()[i])),
+ inputT,
+ outputT);
+
+ checkParameterOneOf(errors, parameter, ALLOWED_ON_TIMER_FAMILY_PARAMETERS);
+
+ extraParameters.add(parameter);
+ }
+
+ return DoFnSignature.OnTimerFamilyMethod.create(
+ m, timerFamilyId, requiresStableInput, windowT, extraParameters);
+ }
+
+ @VisibleForTesting
static DoFnSignature.OnWindowExpirationMethod analyzeOnWindowExpirationMethod(
ErrorReporter errors,
TypeDescriptor<? extends DoFn<?, ?>> fnClass,
@@ -812,6 +981,7 @@
TypeDescriptor<?> trackerT = getTrackerType(fnClass, m);
TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnClass, m);
+
for (int i = 0; i < params.length; ++i) {
Parameter extraParam =
analyzeExtraParameter(
@@ -993,6 +1163,43 @@
return Parameter.timerParameter(timerDecl);
+ } else if (hasTimerIdAnnotation(param.getAnnotations())) {
+ boolean isValidTimerIdForTimerFamily =
+ fnContext.getTimerFamilyDeclarations().size() > 0 && rawType.equals(String.class);
+ paramErrors.checkArgument(
+ isValidTimerIdForTimerFamily, "%s not allowed here", DoFn.TimerId.class.getSimpleName());
+ return Parameter.timerIdParameter();
+ } else if (rawType.equals(TimerMap.class)) {
+ String id = getTimerFamilyId(param.getAnnotations());
+
+ paramErrors.checkArgument(
+ id != null,
+ "%s missing %s annotation",
+ TimerMap.class.getSimpleName(),
+ DoFn.TimerFamily.class.getSimpleName());
+
+ paramErrors.checkArgument(
+ !methodContext.getTimerFamilyParameters().containsKey(id),
+ "duplicate %s: \"%s\"",
+ DoFn.TimerFamily.class.getSimpleName(),
+ id);
+
+ TimerFamilyDeclaration timerDecl = fnContext.getTimerFamilyDeclarations().get(id);
+ paramErrors.checkArgument(
+ timerDecl != null,
+ "reference to undeclared %s: \"%s\"",
+ DoFn.TimerFamily.class.getSimpleName(),
+ id);
+
+ paramErrors.checkArgument(
+ timerDecl.field().getDeclaringClass().equals(getDeclaringClass(param.getMethod())),
+ "%s %s declared in a different class %s."
+ + " Timers may be referenced only in the lexical scope where they are declared.",
+ DoFn.TimerFamily.class.getSimpleName(),
+ id,
+ timerDecl.field().getDeclaringClass().getName());
+
+ return Parameter.timerFamilyParameter(timerDecl);
} else if (State.class.isAssignableFrom(rawType)) {
// m.getParameters() is not available until Java 8
String id = getStateId(param.getAnnotations());
@@ -1032,13 +1239,7 @@
return Parameter.stateParameter(stateDecl);
} else {
- List<String> allowedParamTypes =
- Arrays.asList(
- formatType(new TypeDescriptor<BoundedWindow>() {}),
- formatType(new TypeDescriptor<RestrictionTracker<?, ?>>() {}));
- paramErrors.throwIllegalArgument(
- "%s is not a valid context parameter. Should be one of %s",
- formatType(paramT), allowedParamTypes);
+ paramErrors.throwIllegalArgument("%s is not a valid context parameter.", formatType(paramT));
// Unreachable
return null;
}
@@ -1051,6 +1252,12 @@
}
@Nullable
+ private static String getTimerFamilyId(List<Annotation> annotations) {
+ DoFn.TimerFamily timerFamilyId = findFirstOfType(annotations, DoFn.TimerFamily.class);
+ return timerFamilyId != null ? timerFamilyId.value() : null;
+ }
+
+ @Nullable
private static String getStateId(List<Annotation> annotations) {
DoFn.StateId stateId = findFirstOfType(annotations, DoFn.StateId.class);
return stateId != null ? stateId.value() : null;
@@ -1087,6 +1294,10 @@
return annotations.stream().anyMatch(a -> a.annotationType().equals(DoFn.SideInput.class));
}
+ private static boolean hasTimerIdAnnotation(List<Annotation> annotations) {
+ return annotations.stream().anyMatch(a -> a.annotationType().equals(DoFn.TimerId.class));
+ }
+
@Nullable
private static TypeDescriptor<?> getTrackerType(TypeDescriptor<?> fnClass, Method method) {
Type[] params = method.getGenericParameterTypes();
@@ -1158,19 +1369,44 @@
@VisibleForTesting
static DoFnSignature.GetInitialRestrictionMethod analyzeGetInitialRestrictionMethod(
ErrorReporter errors,
- TypeDescriptor<? extends DoFn> fnT,
+ TypeDescriptor<? extends DoFn<?, ?>> fnT,
Method m,
- TypeDescriptor<?> inputT) {
+ TypeDescriptor<?> inputT,
+ TypeDescriptor<?> outputT,
+ FnAnalysisContext fnContext) {
// Method is of the form:
// @GetInitialRestriction
- // RestrictionT getInitialRestriction(InputT element);
+ // RestrictionT getInitialRestriction(InputT element, ... additional optional parameters ...);
+
Type[] params = m.getGenericParameterTypes();
errors.checkArgument(
- params.length == 1 && fnT.resolveType(params[0]).equals(inputT),
- "Must take a single argument of type %s",
+ params.length >= 1 && fnT.resolveType(params[0]).equals(inputT),
+ "First argument must be of type %s",
formatType(inputT));
+
+ MethodAnalysisContext methodContext = MethodAnalysisContext.create();
+ TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnT, m);
+ for (int i = 1; i < params.length; ++i) {
+ Parameter extraParam =
+ analyzeExtraParameter(
+ errors,
+ fnContext,
+ methodContext,
+ fnT,
+ ParameterDescription.of(
+ m, i, fnT.resolveType(params[i]), Arrays.asList(m.getParameterAnnotations()[i])),
+ inputT,
+ outputT);
+
+ methodContext.addParameter(extraParam);
+ }
+
+ for (Parameter parameter : methodContext.getExtraParameters()) {
+ checkParameterOneOf(errors, parameter, ALLOWED_GET_INITIAL_RESTRICTION_PARAMETERS);
+ }
+
return DoFnSignature.GetInitialRestrictionMethod.create(
- m, fnT.resolveType(m.getGenericReturnType()));
+ m, fnT.resolveType(m.getGenericReturnType()), windowT, methodContext.extraParameters);
}
/**
@@ -1186,16 +1422,19 @@
@VisibleForTesting
static DoFnSignature.SplitRestrictionMethod analyzeSplitRestrictionMethod(
ErrorReporter errors,
- TypeDescriptor<? extends DoFn> fnT,
+ TypeDescriptor<? extends DoFn<?, ?>> fnT,
Method m,
- TypeDescriptor<?> inputT) {
+ TypeDescriptor<?> inputT,
+ TypeDescriptor<?> outputT,
+ FnAnalysisContext fnContext) {
// Method is of the form:
// @SplitRestriction
- // void splitRestriction(InputT element, RestrictionT restriction);
+ // void splitRestriction(InputT element, RestrictionT restriction, ... additional optional
+ // parameters ...);
errors.checkArgument(void.class.equals(m.getReturnType()), "Must return void");
Type[] params = m.getGenericParameterTypes();
- errors.checkArgument(params.length == 3, "Must have exactly 3 arguments");
+ errors.checkArgument(params.length >= 3, "Must have at least 3 arguments");
errors.checkArgument(
fnT.resolveType(params[0]).equals(inputT),
"First argument must be the element type %s",
@@ -1210,7 +1449,43 @@
formatType(expectedReceiverT),
formatType(receiverT));
- return DoFnSignature.SplitRestrictionMethod.create(m, restrictionT);
+ MethodAnalysisContext methodContext = MethodAnalysisContext.create();
+ TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnT, m);
+ for (int i = 3; i < params.length; ++i) {
+ Parameter extraParam =
+ analyzeExtraParameter(
+ errors,
+ fnContext,
+ methodContext,
+ fnT,
+ ParameterDescription.of(
+ m, i, fnT.resolveType(params[i]), Arrays.asList(m.getParameterAnnotations()[i])),
+ inputT,
+ outputT);
+
+ methodContext.addParameter(extraParam);
+ }
+
+ for (Parameter parameter : methodContext.getExtraParameters()) {
+ checkParameterOneOf(errors, parameter, ALLOWED_SPLIT_RESTRICTION_PARAMETERS);
+ }
+
+ return DoFnSignature.SplitRestrictionMethod.create(
+ m, restrictionT, windowT, methodContext.getExtraParameters());
+ }
+
+ private static ImmutableMap<String, TimerFamilyDeclaration> analyzeTimerFamilyDeclarations(
+ ErrorReporter errors, Class<?> fnClazz) {
+ Map<String, TimerFamilyDeclaration> declarations = new HashMap<>();
+ for (Field field : declaredFieldsWithAnnotation(DoFn.TimerFamily.class, fnClazz, DoFn.class)) {
+ // TimerSpec fields may generally be private, but will be accessed via the signature
+ field.setAccessible(true);
+ String id = field.getAnnotation(DoFn.TimerFamily.class).value();
+ validateTimerFamilyField(errors, declarations, id, field);
+ declarations.put(id, TimerFamilyDeclaration.create(id, field));
+ }
+
+ return ImmutableMap.copyOf(declarations);
}
private static ImmutableMap<String, TimerDeclaration> analyzeTimerDeclarations(
@@ -1256,6 +1531,41 @@
}
}
+ /**
+ * Returns successfully if the field is valid, otherwise throws an exception via its {@link
+ * ErrorReporter} parameter describing validation failures for the timer family declaration.
+ */
+ private static void validateTimerFamilyField(
+ ErrorReporter errors,
+ Map<String, TimerFamilyDeclaration> declarations,
+ String id,
+ Field field) {
+
+ if (declarations.containsKey(id)) {
+ errors.throwIllegalArgument(
+ "Duplicate %s \"%s\", used on both of [%s] and [%s]",
+ DoFn.TimerFamily.class.getSimpleName(),
+ id,
+ field.toString(),
+ declarations.get(id).field().toString());
+ }
+
+ Class<?> timerSpecRawType = field.getType();
+ if (!(timerSpecRawType.equals(TimerSpec.class))) {
+ errors.throwIllegalArgument(
+ "%s annotation on non-%s field [%s]",
+ DoFn.TimerFamily.class.getSimpleName(),
+ TimerSpec.class.getSimpleName(),
+ field.toString());
+ }
+
+ if (!Modifier.isFinal(field.getModifiers())) {
+ errors.throwIllegalArgument(
+ "Non-final field %s annotated with %s. TimerMap declarations must be final.",
+ field.toString(), DoFn.TimerFamily.class.getSimpleName());
+ }
+ }
+
/** Generates a {@link TypeDescriptor} for {@code Coder<T>} given {@code T}. */
private static <T> TypeDescriptor<Coder<T>> coderTypeOf(TypeDescriptor<T> elementT) {
return new TypeDescriptor<Coder<T>>() {}.where(new TypeParameter<T>() {}, elementT);
@@ -1286,12 +1596,17 @@
@VisibleForTesting
static DoFnSignature.NewTrackerMethod analyzeNewTrackerMethod(
- ErrorReporter errors, TypeDescriptor<? extends DoFn> fnT, Method m) {
+ ErrorReporter errors,
+ TypeDescriptor<? extends DoFn<?, ?>> fnT,
+ Method m,
+ TypeDescriptor<?> inputT,
+ TypeDescriptor<?> outputT,
+ FnAnalysisContext fnContext) {
// Method is of the form:
// @NewTracker
- // TrackerT newTracker(RestrictionT restriction);
+ // TrackerT newTracker(RestrictionT restriction, ... additional optional parameters ...);
Type[] params = m.getGenericParameterTypes();
- errors.checkArgument(params.length == 1, "Must have a single argument");
+ errors.checkArgument(params.length >= 1, "Must have at least one argument");
TypeDescriptor<?> restrictionT = fnT.resolveType(params[0]);
TypeDescriptor<?> trackerT = fnT.resolveType(m.getGenericReturnType());
@@ -1301,7 +1616,30 @@
"Returns %s, but must return a subtype of %s",
formatType(trackerT),
formatType(expectedTrackerT));
- return DoFnSignature.NewTrackerMethod.create(m, restrictionT, trackerT);
+
+ MethodAnalysisContext methodContext = MethodAnalysisContext.create();
+ TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnT, m);
+ for (int i = 1; i < params.length; ++i) {
+ Parameter extraParam =
+ analyzeExtraParameter(
+ errors,
+ fnContext,
+ methodContext,
+ fnT,
+ ParameterDescription.of(
+ m, i, fnT.resolveType(params[i]), Arrays.asList(m.getParameterAnnotations()[i])),
+ inputT,
+ outputT);
+
+ methodContext.addParameter(extraParam);
+ }
+
+ for (Parameter parameter : methodContext.getExtraParameters()) {
+ checkParameterOneOf(errors, parameter, ALLOWED_NEW_TRACKER_PARAMETERS);
+ }
+
+ return DoFnSignature.NewTrackerMethod.create(
+ m, restrictionT, trackerT, windowT, methodContext.getExtraParameters());
}
private static Collection<Method> declaredMethodsWithAnnotation(
@@ -1578,4 +1916,27 @@
timerDeclaration.field().getName()));
}
}
+
+ public static TimerSpec getTimerFamilySpecOrThrow(
+ TimerFamilyDeclaration timerFamilyDeclaration, DoFn<?, ?> target) {
+ try {
+ Object fieldValue = timerFamilyDeclaration.field().get(target);
+ checkState(
+ fieldValue instanceof TimerSpec,
+ "Malformed %s class %s: timer declaration field %s does not have type %s.",
+ DoFn.class.getSimpleName(),
+ target.getClass().getName(),
+ timerFamilyDeclaration.field().getName(),
+ TimerSpec.class);
+
+ return (TimerSpec) timerFamilyDeclaration.field().get(target);
+ } catch (IllegalAccessException exc) {
+ throw new RuntimeException(
+ String.format(
+ "Malformed %s class %s: timer declaration field %s is not accessible.",
+ DoFn.class.getSimpleName(),
+ target.getClass().getName(),
+ timerFamilyDeclaration.field().getName()));
+ }
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerInvokers.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerInvokers.java
index 287828a..366e904 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerInvokers.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerInvokers.java
@@ -39,4 +39,9 @@
DoFn<InputT, OutputT> fn, String timerId) {
return ByteBuddyOnTimerInvokerFactory.only().forTimer(fn, timerId);
}
+
+ public static <InputT, OutputT> OnTimerInvoker<InputT, OutputT> forTimerFamily(
+ DoFn<InputT, OutputT> fn, String timerFamilyId) {
+ return ByteBuddyOnTimerInvokerFactory.only().forTimerFamily(fn, timerFamilyId);
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ReleaseInfo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ReleaseInfo.java
index 741f373..08a962a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ReleaseInfo.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ReleaseInfo.java
@@ -49,11 +49,16 @@
return getProperties().get("name");
}
- /** Provides the SDK version. */
+ /** Provides the BEAM version. ie: 2.18.0-SNAPSHOT */
public String getVersion() {
return getProperties().get("version");
}
+ /** Provides the SDK version. ie: 2.18.0 or 2.18.0.dev */
+ public String getSdkVersion() {
+ return getProperties().get("sdk_version");
+ }
+
/////////////////////////////////////////////////////////////////////////
private static final Logger LOG = LoggerFactory.getLogger(ReleaseInfo.class);
private static final String DEFAULT_NAME = "Apache Beam SDK for Java";
@@ -79,6 +84,9 @@
if (!properties.containsKey("version")) {
properties.setProperty("version", DEFAULT_VERSION);
}
+ if (!properties.containsKey("sdk_version")) {
+ properties.setProperty("sdk_version", DEFAULT_VERSION);
+ }
INSTANCE = new AutoValue_ReleaseInfo(ImmutableMap.copyOf((Map) properties));
}
}
diff --git a/sdks/java/core/src/main/resources/org/apache/beam/sdk/sdk.properties b/sdks/java/core/src/main/resources/org/apache/beam/sdk/sdk.properties
index 38181c4..3320a4c 100644
--- a/sdks/java/core/src/main/resources/org/apache/beam/sdk/sdk.properties
+++ b/sdks/java/core/src/main/resources/org/apache/beam/sdk/sdk.properties
@@ -17,6 +17,7 @@
# SDK source version
version=@pom.version@
+sdk_version=@pom.sdk_version@
build.date=@timestamp@
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/RequiresStableInputIT.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/RequiresStableInputIT.java
index 1403753..6d06d33 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/RequiresStableInputIT.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/RequiresStableInputIT.java
@@ -17,6 +17,9 @@
*/
package org.apache.beam.sdk;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
+import static org.hamcrest.MatcherAssert.assertThat;
+
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
@@ -28,8 +31,6 @@
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.FileChecksumMatcher;
-import org.apache.beam.sdk.testing.SerializableMatchers;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
import org.apache.beam.sdk.transforms.Create;
@@ -137,13 +138,6 @@
.resolve("key-", StandardResolveOptions.RESOLVE_FILE)
.toString();
- options.setOnSuccessMatcher(
- SerializableMatchers.allOf(
- new FileChecksumMatcher(
- VALUE_CHECKSUM, new FilePatternMatchingShardedFile(singleOutputPrefix + "*")),
- new FileChecksumMatcher(
- VALUE_CHECKSUM, new FilePatternMatchingShardedFile(multiOutputPrefix + "*"))));
-
Pipeline p = Pipeline.create(options);
SerializableFunction<Void, Void> firstTime =
@@ -168,5 +162,12 @@
.withOutputTags(new TupleTag<>(), TupleTagList.empty()));
p.run().waitUntilFinish();
+
+ assertThat(
+ new FilePatternMatchingShardedFile(singleOutputPrefix + "*"),
+ fileContentsHaveChecksum(VALUE_CHECKSUM));
+ assertThat(
+ new FilePatternMatchingShardedFile(multiOutputPrefix + "*"),
+ fileContentsHaveChecksum(VALUE_CHECKSUM));
}
}
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/utils/AvroUtilsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/utils/AvroUtilsTest.java
index 3679e21..cd27cb1 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/utils/AvroUtilsTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/utils/AvroUtilsTest.java
@@ -48,7 +48,9 @@
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.utils.AvroGenerators.RecordSchemaGenerator;
import org.apache.beam.sdk.schemas.utils.AvroUtils.TypeWithNullability;
+import org.apache.beam.sdk.testing.CoderProperties;
import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
@@ -501,6 +503,12 @@
}
@Test
+ public void testRowToGenericRecordFunction() {
+ SerializableUtils.ensureSerializable(AvroUtils.getRowToGenericRecordFunction(NULL_SCHEMA));
+ SerializableUtils.ensureSerializable(AvroUtils.getRowToGenericRecordFunction(null));
+ }
+
+ @Test
public void testGenericRecordToBeamRow() {
GenericRecord genericRecord = getGenericRecord();
Row row = AvroUtils.toBeamRowStrict(getGenericRecord(), null);
@@ -513,6 +521,12 @@
}
@Test
+ public void testGenericRecordToRowFunction() {
+ SerializableUtils.ensureSerializable(AvroUtils.getGenericRecordToRowFunction(Schema.of()));
+ SerializableUtils.ensureSerializable(AvroUtils.getGenericRecordToRowFunction(null));
+ }
+
+ @Test
public void testAvroSchemaCoders() {
Pipeline pipeline = Pipeline.create();
org.apache.avro.Schema schema =
@@ -533,6 +547,7 @@
assertFalse(records.hasSchema());
records.setCoder(AvroUtils.schemaCoder(schema));
assertTrue(records.hasSchema());
+ CoderProperties.coderSerializable(records.getCoder());
AvroGeneratedUser user = new AvroGeneratedUser("foo", 42, "green");
PCollection<AvroGeneratedUser> users =
@@ -540,6 +555,7 @@
assertFalse(users.hasSchema());
users.setCoder(AvroUtils.schemaCoder((AvroCoder<AvroGeneratedUser>) users.getCoder()));
assertTrue(users.hasSchema());
+ CoderProperties.coderSerializable(users.getCoder());
}
public static ContainsField containsField(Function<org.apache.avro.Schema, Boolean> predicate) {
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/FileChecksumMatcherTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/FileChecksumMatcherTest.java
index fd3927f..b2c1eb1 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/FileChecksumMatcherTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/FileChecksumMatcherTest.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.testing;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsString;
@@ -24,7 +25,7 @@
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.regex.Pattern;
-import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.util.NumberedShardedFile;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.Files;
import org.junit.Rule;
import org.junit.Test;
@@ -32,8 +33,6 @@
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
-import org.mockito.Mock;
-import org.mockito.Mockito;
/** Tests for {@link FileChecksumMatcher}. */
@RunWith(JUnit4.class)
@@ -41,53 +40,28 @@
@Rule public TemporaryFolder tmpFolder = new TemporaryFolder();
@Rule public ExpectedException thrown = ExpectedException.none();
- @Mock private PipelineResult pResult = Mockito.mock(PipelineResult.class);
-
@Test
public void testPreconditionChecksumIsNull() throws IOException {
- String tmpPath = tmpFolder.newFile().getPath();
-
thrown.expect(IllegalArgumentException.class);
thrown.expectMessage(containsString("Expected valid checksum, but received"));
- new FileChecksumMatcher(null, tmpPath);
+ fileContentsHaveChecksum(null);
}
@Test
public void testPreconditionChecksumIsEmpty() throws IOException {
- String tmpPath = tmpFolder.newFile().getPath();
-
thrown.expect(IllegalArgumentException.class);
thrown.expectMessage(containsString("Expected valid checksum, but received"));
- new FileChecksumMatcher("", tmpPath);
- }
-
- @Test
- public void testPreconditionFilePathIsEmpty() {
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage(containsString("Expected valid file path, but received"));
- new FileChecksumMatcher("checksumString", "");
- }
-
- @Test
- public void testPreconditionShardTemplateIsNull() throws IOException {
- String tmpPath = tmpFolder.newFile().getPath();
-
- thrown.expect(NullPointerException.class);
- thrown.expectMessage(
- containsString(
- "Expected non-null shard pattern. "
- + "Please call the other constructor to use default pattern:"));
- new FileChecksumMatcher("checksumString", tmpPath, null);
+ fileContentsHaveChecksum("");
}
@Test
public void testMatcherThatVerifiesSingleFile() throws IOException {
File tmpFile = tmpFolder.newFile("result-000-of-001");
Files.write("Test for file checksum verifier.", tmpFile, StandardCharsets.UTF_8);
- FileChecksumMatcher matcher =
- new FileChecksumMatcher("a8772322f5d7b851777f820fc79d050f9d302915", tmpFile.getPath());
- assertThat(pResult, matcher);
+ assertThat(
+ new NumberedShardedFile(tmpFile.getPath()),
+ fileContentsHaveChecksum("a8772322f5d7b851777f820fc79d050f9d302915"));
}
@Test
@@ -99,24 +73,19 @@
Files.write("it is not a question.", tmpFile2, StandardCharsets.UTF_8);
Files.write("tmp", tmpFile3, StandardCharsets.UTF_8);
- FileChecksumMatcher matcher =
- new FileChecksumMatcher(
- "90552392c28396935fe4f123bd0b5c2d0f6260c8",
- tmpFolder.getRoot().toPath().resolve("result-*").toString());
-
- assertThat(pResult, matcher);
+ assertThat(
+ new NumberedShardedFile(tmpFolder.getRoot().toPath().resolve("result-*").toString()),
+ fileContentsHaveChecksum("90552392c28396935fe4f123bd0b5c2d0f6260c8"));
}
@Test
public void testMatcherThatVerifiesFileWithEmptyContent() throws IOException {
File emptyFile = tmpFolder.newFile("result-000-of-001");
Files.write("", emptyFile, StandardCharsets.UTF_8);
- FileChecksumMatcher matcher =
- new FileChecksumMatcher(
- "da39a3ee5e6b4b0d3255bfef95601890afd80709",
- tmpFolder.getRoot().toPath().resolve("*").toString());
- assertThat(pResult, matcher);
+ assertThat(
+ new NumberedShardedFile(tmpFolder.getRoot().toPath().resolve("*").toString()),
+ fileContentsHaveChecksum("da39a3ee5e6b4b0d3255bfef95601890afd80709"));
}
@Test
@@ -129,12 +98,10 @@
Pattern customizedTemplate =
Pattern.compile("(?x) result (?<shardnum>\\d+) - total (?<numshards>\\d+)");
- FileChecksumMatcher matcher =
- new FileChecksumMatcher(
- "90552392c28396935fe4f123bd0b5c2d0f6260c8",
- tmpFolder.getRoot().toPath().resolve("*").toString(),
- customizedTemplate);
- assertThat(pResult, matcher);
+ assertThat(
+ new NumberedShardedFile(
+ tmpFolder.getRoot().toPath().resolve("*").toString(), customizedTemplate),
+ fileContentsHaveChecksum("90552392c28396935fe4f123bd0b5c2d0f6260c8"));
}
}
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestStreamTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestStreamTest.java
index e48b6b2..841a75c 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestStreamTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/TestStreamTest.java
@@ -191,7 +191,7 @@
TimestampedValue.of("firstPane", new Instant(100)),
TimestampedValue.of("alsoFirstPane", new Instant(200)))
.addElements(TimestampedValue.of("onTimePane", new Instant(500)))
- .advanceWatermarkTo(new Instant(1001L))
+ .advanceWatermarkTo(new Instant(1000L))
.addElements(
TimestampedValue.of("finalLatePane", new Instant(750)),
TimestampedValue.of("alsoFinalLatePane", new Instant(250)))
@@ -469,8 +469,7 @@
}
@Test
- @Category(UsesTestStreamWithProcessingTime.class)
- public void testCoder() throws Exception {
+ public void testTestStreamCoder() throws Exception {
TestStream<String> testStream =
TestStream.create(StringUtf8Coder.of())
.addElements("hey")
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
index db57335..455761f 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.transforms;
+import static junit.framework.TestCase.assertTrue;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasKey;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasType;
@@ -64,6 +65,7 @@
import org.apache.beam.sdk.coders.SetCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.GenerateSequence;
import org.apache.beam.sdk.options.Default;
@@ -78,6 +80,7 @@
import org.apache.beam.sdk.state.StateSpecs;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.state.TimerSpec;
import org.apache.beam.sdk.state.TimerSpecs;
import org.apache.beam.sdk.state.ValueState;
@@ -94,6 +97,7 @@
import org.apache.beam.sdk.testing.UsesStrictTimerOrdering;
import org.apache.beam.sdk.testing.UsesTestStream;
import org.apache.beam.sdk.testing.UsesTestStreamWithProcessingTime;
+import org.apache.beam.sdk.testing.UsesTimerMap;
import org.apache.beam.sdk.testing.UsesTimersInParDo;
import org.apache.beam.sdk.testing.ValidatesRunner;
import org.apache.beam.sdk.transforms.DoFn.OnTimer;
@@ -2863,29 +2867,28 @@
public void testEventTimeTimerBounded() throws Exception {
final String timerId = "foo";
- DoFn<KV<String, Integer>, Integer> fn =
- new DoFn<KV<String, Integer>, Integer>() {
+ DoFn<KV<String, Long>, Long> fn =
+ new DoFn<KV<String, Long>, Long>() {
@TimerId(timerId)
private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@ProcessElement
- public void processElement(@TimerId(timerId) Timer timer, OutputReceiver<Integer> r) {
+ public void processElement(@TimerId(timerId) Timer timer, OutputReceiver<Long> r) {
timer.offset(Duration.standardSeconds(1)).setRelative();
- r.output(3);
+ r.output(3L);
}
@OnTimer(timerId)
- public void onTimer(TimeDomain timeDomain, OutputReceiver<Integer> r) {
+ public void onTimer(TimeDomain timeDomain, OutputReceiver<Long> r) {
if (timeDomain.equals(TimeDomain.EVENT_TIME)) {
- r.output(42);
+ r.output(42L);
}
}
};
- PCollection<Integer> output =
- pipeline.apply(Create.of(KV.of("hello", 37))).apply(ParDo.of(fn));
- PAssert.that(output).containsInAnyOrder(3, 42);
+ PCollection<Long> output = pipeline.apply(Create.of(KV.of("hello", 37L))).apply(ParDo.of(fn));
+ PAssert.that(output).containsInAnyOrder(3L, 42L);
pipeline.run();
}
@@ -2941,8 +2944,8 @@
public void testEventTimeTimerAlignBounded() throws Exception {
final String timerId = "foo";
- DoFn<KV<String, Integer>, KV<Integer, Instant>> fn =
- new DoFn<KV<String, Integer>, KV<Integer, Instant>>() {
+ DoFn<KV<String, Long>, KV<Long, Instant>> fn =
+ new DoFn<KV<String, Long>, KV<Long, Instant>>() {
@TimerId(timerId)
private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@@ -2951,24 +2954,23 @@
public void processElement(
@TimerId(timerId) Timer timer,
@Timestamp Instant timestamp,
- OutputReceiver<KV<Integer, Instant>> r) {
+ OutputReceiver<KV<Long, Instant>> r) {
timer.align(Duration.standardSeconds(1)).offset(Duration.millis(1)).setRelative();
- r.output(KV.of(3, timestamp));
+ r.output(KV.of(3L, timestamp));
}
@OnTimer(timerId)
- public void onTimer(
- @Timestamp Instant timestamp, OutputReceiver<KV<Integer, Instant>> r) {
- r.output(KV.of(42, timestamp));
+ public void onTimer(@Timestamp Instant timestamp, OutputReceiver<KV<Long, Instant>> r) {
+ r.output(KV.of(42L, timestamp));
}
};
- PCollection<KV<Integer, Instant>> output =
- pipeline.apply(Create.of(KV.of("hello", 37))).apply(ParDo.of(fn));
+ PCollection<KV<Long, Instant>> output =
+ pipeline.apply(Create.of(KV.of("hello", 37L))).apply(ParDo.of(fn));
PAssert.that(output)
.containsInAnyOrder(
- KV.of(3, BoundedWindow.TIMESTAMP_MIN_VALUE),
- KV.of(42, BoundedWindow.TIMESTAMP_MIN_VALUE.plus(1774)));
+ KV.of(3L, BoundedWindow.TIMESTAMP_MIN_VALUE),
+ KV.of(42L, BoundedWindow.TIMESTAMP_MIN_VALUE.plus(1774)));
pipeline.run();
}
@@ -3035,28 +3037,27 @@
public void testEventTimeTimerAbsolute() throws Exception {
final String timerId = "foo";
- DoFn<KV<String, Integer>, Integer> fn =
- new DoFn<KV<String, Integer>, Integer>() {
+ DoFn<KV<String, Long>, Long> fn =
+ new DoFn<KV<String, Long>, Long>() {
@TimerId(timerId)
private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@ProcessElement
public void processElement(
- @TimerId(timerId) Timer timer, BoundedWindow window, OutputReceiver<Integer> r) {
+ @TimerId(timerId) Timer timer, BoundedWindow window, OutputReceiver<Long> r) {
timer.set(window.maxTimestamp());
- r.output(3);
+ r.output(3L);
}
@OnTimer(timerId)
- public void onTimer(OutputReceiver<Integer> r) {
- r.output(42);
+ public void onTimer(OutputReceiver<Long> r) {
+ r.output(42L);
}
};
- PCollection<Integer> output =
- pipeline.apply(Create.of(KV.of("hello", 37))).apply(ParDo.of(fn));
- PAssert.that(output).containsInAnyOrder(3, 42);
+ PCollection<Long> output = pipeline.apply(Create.of(KV.of("hello", 37L))).apply(ParDo.of(fn));
+ PAssert.that(output).containsInAnyOrder(3L, 42L);
pipeline.run();
}
@@ -3122,10 +3123,10 @@
final String stateId = "sizzle";
final int offset = 5000;
- final int timerOutput = 4093;
+ final long timerOutput = 4093;
- DoFn<KV<String, Integer>, KV<String, Integer>> fn =
- new DoFn<KV<String, Integer>, KV<String, Integer>>() {
+ DoFn<KV<String, Long>, KV<String, Long>> fn =
+ new DoFn<KV<String, Long>, KV<String, Long>>() {
@TimerId(timerId)
private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@@ -3148,21 +3149,21 @@
@OnTimer(timerId)
public void onTimer(
- @StateId(stateId) ValueState<String> state, OutputReceiver<KV<String, Integer>> r) {
+ @StateId(stateId) ValueState<String> state, OutputReceiver<KV<String, Long>> r) {
r.output(KV.of(state.read(), timerOutput));
}
};
// Enough keys that we exercise interesting code paths
int numKeys = 50;
- List<KV<String, Integer>> input = new ArrayList<>();
- List<KV<String, Integer>> expectedOutput = new ArrayList<>();
+ List<KV<String, Long>> input = new ArrayList<>();
+ List<KV<String, Long>> expectedOutput = new ArrayList<>();
- for (Integer key = 0; key < numKeys; ++key) {
+ for (Long key = 0L; key < numKeys; ++key) {
// Each key should have just one final output at GC time
expectedOutput.add(KV.of(key.toString(), timerOutput));
- for (int i = 0; i < 15; ++i) {
+ for (long i = 0; i < 15; ++i) {
// Each input should be output with the offset added
input.add(KV.of(key.toString(), i));
expectedOutput.add(KV.of(key.toString(), i + offset));
@@ -3171,8 +3172,7 @@
Collections.shuffle(input);
- PCollection<KV<String, Integer>> output =
- pipeline.apply(Create.of(input)).apply(ParDo.of(fn));
+ PCollection<KV<String, Long>> output = pipeline.apply(Create.of(input)).apply(ParDo.of(fn));
PAssert.that(output).containsInAnyOrder(expectedOutput);
pipeline.run();
}
@@ -3662,10 +3662,10 @@
PCollection<String> results =
pipeline
- .apply(Create.of(KV.of(0, 0)))
+ .apply(Create.of(KV.of(0L, 0L)))
.apply(
ParDo.of(
- new DoFn<KV<Integer, Integer>, String>() {
+ new DoFn<KV<Long, Long>, String>() {
@TimerId(timerId)
private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@@ -3729,6 +3729,83 @@
pipeline.run();
}
+ @Test
+ @Category({
+ ValidatesRunner.class,
+ UsesStatefulParDo.class,
+ UsesTimersInParDo.class,
+ UsesTestStream.class,
+ })
+ public void testOutputTimestamp() {
+ final String timerId = "bar";
+ DoFn<KV<String, Long>, KV<String, Long>> fn1 =
+ new DoFn<KV<String, Long>, KV<String, Long>>() {
+
+ @TimerId(timerId)
+ private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ @ProcessElement
+ public void processElement(
+ @TimerId(timerId) Timer timer, OutputReceiver<KV<String, Long>> o) {
+ timer.withOutputTimestamp(new Instant(5)).set(new Instant(10));
+ // Output a message. This will cause the next DoFn to set a timer as well.
+ o.output(KV.of("foo", 100L));
+ }
+
+ @OnTimer(timerId)
+ public void onTimer(OnTimerContext c, BoundedWindow w) {}
+ };
+
+ DoFn<KV<String, Long>, Long> fn2 =
+ new DoFn<KV<String, Long>, Long>() {
+
+ @TimerId(timerId)
+ private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ @StateId("timerFired")
+ final StateSpec<ValueState<Boolean>> timerFiredState = StateSpecs.value();
+
+ @ProcessElement
+ public void processElement(
+ @TimerId(timerId) Timer timer,
+ @StateId("timerFired") ValueState<Boolean> timerFiredState) {
+ Boolean timerFired = timerFiredState.read();
+ assertTrue(timerFired == null || !timerFired);
+ // Set a timer to 8. This is earlier than the previous DoFn's timer, but after the
+ // previous
+ // DoFn timer's watermark hold. This timer should not fire until the previous timer
+ // fires and removes
+ // the watermark hold.
+ timer.set(new Instant(8));
+ }
+
+ @OnTimer(timerId)
+ public void onTimer(
+ @StateId("timerFired") ValueState<Boolean> timerFiredState,
+ OutputReceiver<Long> o) {
+ timerFiredState.write(true);
+ o.output(100L);
+ }
+ };
+
+ TestStream<KV<String, Long>> stream =
+ TestStream.create(KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()))
+ .advanceWatermarkTo(new Instant(0))
+ // Cause fn2 to set a timer.
+ .addElements(KV.of("key", 1L))
+ // Normally this would case fn2's timer to expire, but it shouldn't here because of
+ // the output timestamp.
+ .advanceWatermarkTo(new Instant(9))
+ // If the timer fired, then this would case fn2 to fail with an assertion error.
+ .addElements(KV.of("key", 1L))
+ .advanceWatermarkToInfinity();
+ PCollection<Long> output =
+ pipeline.apply(stream).apply("first", ParDo.of(fn1)).apply("second", ParDo.of(fn2));
+
+ PAssert.that(output).containsInAnyOrder(100L); // result output
+ pipeline.run();
+ }
+
private static class TwoTimerTest extends PTransform<PBegin, PDone> {
private static PTransform<PBegin, PDone> of(
@@ -4209,4 +4286,133 @@
r.output("It works");
}
}
+
+ /** Tests to validate ParDo timerFamily. */
+ @RunWith(JUnit4.class)
+ public static class TimerFamilyTests extends SharedTestBase implements Serializable {
+
+ @Test
+ @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesTimerMap.class})
+ public void testTimerFamilyEventTime() throws Exception {
+ final String timerFamilyId = "foo";
+
+ DoFn<KV<String, Integer>, String> fn =
+ new DoFn<KV<String, Integer>, String>() {
+
+ @TimerFamily(timerFamilyId)
+ private final TimerSpec spec = TimerSpecs.timerMap(TimeDomain.EVENT_TIME);
+
+ @ProcessElement
+ public void processElement(
+ @TimerFamily(timerFamilyId) TimerMap timers, OutputReceiver<String> r) {
+ timers.set("timer1", new Instant(1));
+ timers.set("timer2", new Instant(2));
+ r.output("process");
+ }
+
+ @OnTimerFamily(timerFamilyId)
+ public void onTimer(
+ @TimerId String timerId,
+ @Timestamp Instant ts,
+ @TimerFamily(timerFamilyId) TimerMap timerMap,
+ OutputReceiver<String> r) {
+ System.out.println("timer Id : " + timerId);
+ System.out.println("timerMap : " + timerMap.toString());
+ r.output(timerId);
+ }
+ };
+
+ PCollection<String> output =
+ pipeline.apply(Create.of(KV.of("hello", 37))).apply(ParDo.of(fn));
+ PAssert.that(output).containsInAnyOrder("process", "timer1", "timer2");
+ pipeline.run();
+ }
+
+ @Test
+ @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesTimerMap.class})
+ public void testTimerWithMultipleTimerFamily() throws Exception {
+ final String timerFamilyId1 = "foo";
+ final String timerFamilyId2 = "bar";
+
+ DoFn<KV<String, Integer>, String> fn =
+ new DoFn<KV<String, Integer>, String>() {
+
+ @TimerFamily(timerFamilyId1)
+ private final TimerSpec spec1 = TimerSpecs.timerMap(TimeDomain.EVENT_TIME);
+
+ @TimerFamily(timerFamilyId2)
+ private final TimerSpec spec2 = TimerSpecs.timerMap(TimeDomain.EVENT_TIME);
+
+ @ProcessElement
+ public void processElement(
+ @TimerFamily(timerFamilyId1) TimerMap timerMap1,
+ @TimerFamily(timerFamilyId2) TimerMap timerMap2,
+ OutputReceiver<String> r) {
+ timerMap1.set("timer", new Instant(1));
+ timerMap2.set("timer", new Instant(2));
+ r.output("process");
+ }
+
+ @OnTimerFamily(timerFamilyId1)
+ public void onTimer1(
+ @TimerId String timerId, @Timestamp Instant ts, OutputReceiver<String> r) {
+ r.output(timerId);
+ }
+
+ @OnTimerFamily(timerFamilyId2)
+ public void onTimer2(
+ @TimerId String timerId, @Timestamp Instant ts, OutputReceiver<String> r) {
+ r.output(timerId);
+ }
+ };
+
+ PCollection<String> output =
+ pipeline.apply(Create.of(KV.of("hello", 37))).apply(ParDo.of(fn));
+ PAssert.that(output).containsInAnyOrder("process", "timer", "timer");
+ pipeline.run();
+ }
+
+ @Test
+ @Category({
+ NeedsRunner.class,
+ UsesTimersInParDo.class,
+ UsesTestStreamWithProcessingTime.class,
+ UsesTimerMap.class
+ })
+ public void testTimerFamilyProcessingTime() throws Exception {
+ final String timerId = "foo";
+
+ DoFn<KV<String, Integer>, Integer> fn =
+ new DoFn<KV<String, Integer>, Integer>() {
+
+ @TimerFamily(timerId)
+ private final TimerSpec spec = TimerSpecs.timerMap(TimeDomain.PROCESSING_TIME);
+
+ @ProcessElement
+ public void processElement(
+ @TimerFamily(timerId) TimerMap timerMap, OutputReceiver<Integer> r) {
+ Timer timer = timerMap.get("timerId1");
+ timer.offset(Duration.standardSeconds(1)).setRelative();
+ r.output(3);
+ }
+
+ @OnTimerFamily(timerId)
+ public void onTimer(TimeDomain timeDomain, OutputReceiver<Integer> r) {
+ if (timeDomain.equals(TimeDomain.PROCESSING_TIME)) {
+ r.output(42);
+ }
+ }
+ };
+
+ TestStream<KV<String, Integer>> stream =
+ TestStream.create(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .addElements(KV.of("hello", 37))
+ .advanceProcessingTime(Duration.standardSeconds(2))
+ .advanceWatermarkToInfinity();
+
+ PCollection<Integer> output = pipeline.apply(stream).apply(ParDo.of(fn));
+ PAssert.that(output).containsInAnyOrder(3, 42);
+ pipeline.run();
+ }
+ }
}
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
index 44e943e..18223e5 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
@@ -113,7 +113,7 @@
}
private void invokeOnTimer(String timerId, DoFn<String, String> fn) {
- DoFnInvokers.invokerFor(fn).invokeOnTimer(timerId, mockArgumentProvider);
+ DoFnInvokers.invokerFor(fn).invokeOnTimer(timerId, timerId, mockArgumentProvider);
}
@Test
@@ -831,7 +831,7 @@
SimpleTimerDoFn fn = new SimpleTimerDoFn();
DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
- invoker.invokeOnTimer(timerId, mockArgumentProvider);
+ invoker.invokeOnTimer(timerId, timerId, mockArgumentProvider);
assertThat(fn.status, equalTo("OK now"));
}
@@ -860,7 +860,7 @@
SimpleTimerDoFn fn = new SimpleTimerDoFn();
DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
- invoker.invokeOnTimer(timerId, mockArgumentProvider);
+ invoker.invokeOnTimer(timerId, timerId, mockArgumentProvider);
assertThat(fn.window, equalTo(testWindow));
}
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java
index 1a19eeb..593c846 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java
@@ -37,9 +37,7 @@
@Test
public void testBadExtraProcessContextType() throws Exception {
thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage(
- "Integer is not a valid context parameter. "
- + "Should be one of [BoundedWindow, RestrictionTracker<?, ?>]");
+ thrown.expectMessage("Integer is not a valid context parameter.");
analyzeProcessElementMethod(
new AnonymousMethod() {
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java
index bac8459..d5efdc5 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java
@@ -19,23 +19,34 @@
import static org.apache.beam.sdk.transforms.reflect.DoFnSignaturesTestUtils.analyzeProcessElementMethod;
import static org.apache.beam.sdk.transforms.reflect.DoFnSignaturesTestUtils.errors;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
+import java.lang.reflect.Method;
import java.util.List;
import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.StructuredCoder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.BoundedPerElement;
+import org.apache.beam.sdk.transforms.DoFn.StateId;
import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignatures.FnAnalysisContext;
import org.apache.beam.sdk.transforms.reflect.DoFnSignaturesTestUtils.AnonymousMethod;
import org.apache.beam.sdk.transforms.reflect.DoFnSignaturesTestUtils.FakeDoFn;
import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Predicates;
+import org.joda.time.Instant;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
@@ -92,19 +103,45 @@
}
@Test
- public void testSplittableProcessElementMustNotHaveOtherParams() throws Exception {
+ public void testSplittableProcessElementMustNotHaveUnsupportedParams() throws Exception {
thrown.expect(IllegalArgumentException.class);
thrown.expectMessage("Illegal parameter");
- thrown.expectMessage("BoundedWindow");
+ thrown.expectMessage("ValueState");
- DoFnSignature.ProcessElementMethod signature =
- analyzeProcessElementMethod(
- new AnonymousMethod() {
- private void method(
- DoFn<Integer, String>.ProcessContext context,
- SomeRestrictionTracker tracker,
- BoundedWindow window) {}
- });
+ DoFn<Integer, String> doFn =
+ new DoFn<Integer, String>() {
+ @StateId("my-state-id")
+ public final StateSpec<ValueState<String>> myStateSpec =
+ StateSpecs.value(StringUtf8Coder.of());
+
+ @ProcessElement
+ public void method(
+ DoFn<Integer, String>.ProcessContext context,
+ SomeRestrictionTracker tracker,
+ @StateId("my-state-id") ValueState<String> myState) {}
+ };
+ Method processElementMethod = null;
+ for (Method method : doFn.getClass().getDeclaredMethods()) {
+ if ("method".equals(method.getName())) {
+ processElementMethod = method;
+ }
+ }
+ checkState(processElementMethod != null);
+
+ FnAnalysisContext context = FnAnalysisContext.create();
+ context.addStateDeclaration(
+ DoFnSignature.StateDeclaration.create(
+ "my-state-id",
+ doFn.getClass().getField("myStateSpec"),
+ new TypeDescriptor<ValueState<String>>() {}));
+
+ DoFnSignatures.analyzeProcessElementMethod(
+ errors(),
+ new TypeDescriptor<DoFn<Integer, String>>() {},
+ processElementMethod,
+ TypeDescriptor.of(Integer.class),
+ TypeDescriptor.of(String.class),
+ context);
}
@Test
@@ -226,7 +263,7 @@
/** Tests a splittable {@link DoFn} that defines all methods in their full form, correctly. */
@Test
- public void testSplittableWithAllFunctions() throws Exception {
+ public void testSplittableWithAllFunctionsAndAllParameters() throws Exception {
class GoodSplittableDoFn extends DoFn<Integer, String> {
@ProcessElement
public ProcessContinuation processElement(
@@ -235,16 +272,32 @@
}
@GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer element) {
+ public SomeRestriction getInitialRestriction(
+ Integer element,
+ PipelineOptions pipelineOptions,
+ BoundedWindow boundedWindow,
+ PaneInfo paneInfo,
+ @Timestamp Instant timestamp) {
return null;
}
@SplitRestriction
public void splitRestriction(
- Integer element, SomeRestriction restriction, OutputReceiver<SomeRestriction> receiver) {}
+ Integer element,
+ SomeRestriction restriction,
+ OutputReceiver<SomeRestriction> receiver,
+ PipelineOptions pipelineOptions,
+ BoundedWindow boundedWindow,
+ PaneInfo paneInfo,
+ @Timestamp Instant timestamp) {}
@NewTracker
- public SomeRestrictionTracker newTracker(SomeRestriction restriction) {
+ public SomeRestrictionTracker newTracker(
+ SomeRestriction restriction,
+ PipelineOptions pipelineOptions,
+ BoundedWindow boundedWindow,
+ PaneInfo paneInfo,
+ @Timestamp Instant timestamp) {
return null;
}
@@ -455,7 +508,9 @@
void method(
Integer element, SomeRestriction restriction, DoFn.OutputReceiver<String> receiver) {}
}.getMethod(),
- TypeDescriptor.of(Integer.class));
+ TypeDescriptor.of(Integer.class),
+ TypeDescriptor.of(String.class),
+ FnAnalysisContext.create());
}
@Test
@@ -476,12 +531,14 @@
SomeRestriction restriction,
DoFn.OutputReceiver<SomeRestriction> receiver) {}
}.getMethod(),
- TypeDescriptor.of(Integer.class));
+ TypeDescriptor.of(Integer.class),
+ TypeDescriptor.of(String.class),
+ FnAnalysisContext.create());
}
@Test
- public void testSplitRestrictionWrongNumArguments() throws Exception {
- thrown.expectMessage("Must have exactly 3 arguments");
+ public void testSplitRestrictionWrongArgumentType() throws Exception {
+ thrown.expectMessage("Object is not a valid context parameter.");
DoFnSignatures.analyzeSplitRestrictionMethod(
errors(),
TypeDescriptor.of(FakeDoFn.class),
@@ -492,7 +549,9 @@
DoFn.OutputReceiver<SomeRestriction> receiver,
Object extra) {}
}.getMethod(),
- TypeDescriptor.of(Integer.class));
+ TypeDescriptor.of(Integer.class),
+ TypeDescriptor.of(String.class),
+ FnAnalysisContext.create());
}
@Test
@@ -563,8 +622,8 @@
}
@Test
- public void testNewTrackerWrongNumArguments() throws Exception {
- thrown.expectMessage("Must have a single argument");
+ public void testNewTrackerWrongArgumentType() throws Exception {
+ thrown.expectMessage("Object is not a valid context parameter.");
DoFnSignatures.analyzeNewTrackerMethod(
errors(),
TypeDescriptor.of(FakeDoFn.class),
@@ -572,7 +631,10 @@
private SomeRestrictionTracker method(SomeRestriction restriction, Object extra) {
return null;
}
- }.getMethod());
+ }.getMethod(),
+ TypeDescriptor.of(Integer.class),
+ TypeDescriptor.of(String.class),
+ FnAnalysisContext.create());
}
@Test
@@ -587,6 +649,9 @@
private SomeRestrictionTracker method(String restriction) {
return null;
}
- }.getMethod());
+ }.getMethod(),
+ TypeDescriptor.of(Integer.class),
+ TypeDescriptor.of(String.class),
+ FnAnalysisContext.create());
}
}
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
index e0427f2..fedccf4 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
@@ -1026,6 +1026,11 @@
assertThat(stateParam.referent(), equalTo(decl));
return null;
}
+
+ @Override
+ public Void dispatch(Parameter.TimerIdParameter p) {
+ return null;
+ }
});
}
diff --git a/sdks/java/extensions/sketching/build.gradle b/sdks/java/extensions/sketching/build.gradle
index d923501..54cd4d2d 100644
--- a/sdks/java/extensions/sketching/build.gradle
+++ b/sdks/java/extensions/sketching/build.gradle
@@ -31,7 +31,6 @@
compile "com.tdunning:t-digest:$tdigest_version"
compile library.java.slf4j_api
testCompile library.java.avro
- testCompile library.java.commons_lang3
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
diff --git a/sdks/java/extensions/sql/build.gradle b/sdks/java/extensions/sql/build.gradle
index d7cc6b1..ee4db16 100644
--- a/sdks/java/extensions/sql/build.gradle
+++ b/sdks/java/extensions/sql/build.gradle
@@ -48,7 +48,6 @@
compile project(path: ":runners:direct-java", configuration: "shadow")
compile library.java.commons_codec
compile library.java.commons_csv
- compile library.java.commons_lang3
compile library.java.jackson_databind
compile library.java.joda_time
compile library.java.vendored_calcite_1_20_0
diff --git a/sdks/java/extensions/sql/perf-tests/build.gradle b/sdks/java/extensions/sql/perf-tests/build.gradle
new file mode 100644
index 0000000..7875a6b
--- /dev/null
+++ b/sdks/java/extensions/sql/perf-tests/build.gradle
@@ -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.
+ */
+
+plugins { id 'org.apache.beam.module' }
+applyJavaNature(automaticModuleName: 'org.apache.beam.sdk.extensions.sql.meta.provider')
+provideIntegrationTestingDependencies()
+enableJavaPerformanceTesting()
+
+description = "Apache Beam :: SDKs :: Java :: Extensions :: SQL :: IO Performance tests"
+ext.summary = "Performance tests for SQL IO sources"
+
+dependencies {
+ testCompile project(path: ":sdks:java:io:google-cloud-platform", configuration: "testRuntime")
+ testCompile project(path: ":sdks:java:extensions:sql", configuration: "testRuntime")
+}
+
+
diff --git a/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryIOPushDownIT.java b/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryIOPushDownIT.java
new file mode 100644
index 0000000..caa5497
--- /dev/null
+++ b/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryIOPushDownIT.java
@@ -0,0 +1,204 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.bigquery;
+
+import static org.apache.beam.sdk.extensions.sql.impl.planner.BeamRuleSets.getRuleSets;
+
+import com.google.cloud.Timestamp;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Function;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv;
+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.extensions.sql.impl.rule.BeamIOPushDownRule;
+import org.apache.beam.sdk.extensions.sql.meta.store.InMemoryMetaStore;
+import org.apache.beam.sdk.io.common.IOITHelper;
+import org.apache.beam.sdk.io.common.IOTestPipelineOptions;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testutils.NamedTestResult;
+import org.apache.beam.sdk.testutils.metrics.IOITMetrics;
+import org.apache.beam.sdk.testutils.metrics.MetricsReader;
+import org.apache.beam.sdk.testutils.metrics.TimeMonitor;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptRule;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.RuleSet;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.RuleSets;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class BigQueryIOPushDownIT {
+ private static final String READ_FROM_TABLE =
+ "apache-beam-testing:beam_performance.hacker_news_full";
+ private static final String NAMESPACE = BigQueryIOPushDownIT.class.getName();
+ private static final String FIELDS_READ_METRIC = "fields_read";
+ private static final String READ_TIME_METRIC = "read_time";
+ private static final String CREATE_TABLE_STATEMENT =
+ "CREATE EXTERNAL TABLE HACKER_NEWS( \n"
+ + " title VARCHAR, \n"
+ + " url VARCHAR, \n"
+ + " text VARCHAR, \n"
+ + " dead BOOLEAN, \n"
+ + " `by` VARCHAR, \n"
+ + " score INTEGER, \n"
+ + " `time` INTEGER, \n"
+ + " `timestamp` TIMESTAMP, \n"
+ + " type VARCHAR, \n"
+ + " id INTEGER, \n"
+ + " parent INTEGER, \n"
+ + " descendants INTEGER, \n"
+ + " ranking INTEGER, \n"
+ + " deleted BOOLEAN \n"
+ + ") \n"
+ + "TYPE 'bigquery' \n"
+ + "LOCATION '"
+ + READ_FROM_TABLE
+ + "' \n"
+ + "TBLPROPERTIES '{ method: \"%s\" }'";
+ private static final String SELECT_STATEMENT =
+ "SELECT `by` as author, type, title, score from HACKER_NEWS where (type='story' or type='job') and score>2";
+
+ private static SQLBigQueryPerfTestOptions options;
+ private static String metricsBigQueryDataset;
+ private static String metricsBigQueryTable;
+ private Pipeline pipeline = Pipeline.create(options);
+ private BeamSqlEnv sqlEnv;
+
+ @BeforeClass
+ public static void setUp() {
+ options = IOITHelper.readIOTestPipelineOptions(SQLBigQueryPerfTestOptions.class);
+ metricsBigQueryDataset = options.getMetricsBigQueryDataset();
+ metricsBigQueryTable = options.getMetricsBigQueryTable();
+ }
+
+ @Before
+ public void before() {
+ sqlEnv = BeamSqlEnv.inMemory(new BigQueryPerfTableProvider(NAMESPACE, FIELDS_READ_METRIC));
+ }
+
+ @Test
+ public void readUsingDirectReadMethodPushDown() {
+ sqlEnv.executeDdl(String.format(CREATE_TABLE_STATEMENT, Method.DIRECT_READ.toString()));
+
+ BeamRelNode beamRelNode = sqlEnv.parseQuery(SELECT_STATEMENT);
+ PCollection<Row> output =
+ BeamSqlRelUtils.toPCollection(pipeline, beamRelNode)
+ .apply(ParDo.of(new TimeMonitor<>(NAMESPACE, READ_TIME_METRIC)));
+
+ PipelineResult result = pipeline.run();
+ result.waitUntilFinish();
+ collectAndPublishMetrics(result, "_directread_pushdown");
+ }
+
+ @Test
+ public void readUsingDirectReadMethod() {
+ List<RelOptRule> ruleList = new ArrayList<>();
+ for (RuleSet x : getRuleSets()) {
+ x.iterator().forEachRemaining(ruleList::add);
+ }
+ // Remove push-down rule
+ ruleList.remove(BeamIOPushDownRule.INSTANCE);
+
+ InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
+ inMemoryMetaStore.registerProvider(
+ new BigQueryPerfTableProvider(NAMESPACE, FIELDS_READ_METRIC));
+ sqlEnv =
+ BeamSqlEnv.builder(inMemoryMetaStore)
+ .setPipelineOptions(PipelineOptionsFactory.create())
+ .setRuleSets(new RuleSet[] {RuleSets.ofList(ruleList)})
+ .build();
+ sqlEnv.executeDdl(String.format(CREATE_TABLE_STATEMENT, Method.DIRECT_READ.toString()));
+
+ BeamRelNode beamRelNode = sqlEnv.parseQuery(SELECT_STATEMENT);
+ PCollection<Row> output =
+ BeamSqlRelUtils.toPCollection(pipeline, beamRelNode)
+ .apply(ParDo.of(new TimeMonitor<>(NAMESPACE, READ_TIME_METRIC)));
+
+ PipelineResult result = pipeline.run();
+ result.waitUntilFinish();
+ collectAndPublishMetrics(result, "_directread");
+ }
+
+ @Test
+ public void readUsingDefaultMethod() {
+ sqlEnv.executeDdl(String.format(CREATE_TABLE_STATEMENT, Method.DEFAULT.toString()));
+
+ BeamRelNode beamRelNode = sqlEnv.parseQuery(SELECT_STATEMENT);
+ PCollection<Row> output =
+ BeamSqlRelUtils.toPCollection(pipeline, beamRelNode)
+ .apply(ParDo.of(new TimeMonitor<>(NAMESPACE, READ_TIME_METRIC)));
+
+ PipelineResult result = pipeline.run();
+ result.waitUntilFinish();
+ collectAndPublishMetrics(result, "_default");
+ }
+
+ private void collectAndPublishMetrics(PipelineResult readResult, String postfix) {
+ String uuid = UUID.randomUUID().toString();
+ String timestamp = Timestamp.now().toString();
+
+ Set<Function<MetricsReader, NamedTestResult>> readSuppliers = getReadSuppliers(uuid, timestamp);
+ IOITMetrics readMetrics =
+ new IOITMetrics(readSuppliers, readResult, NAMESPACE, uuid, timestamp);
+ readMetrics.publish(metricsBigQueryDataset, metricsBigQueryTable + postfix);
+ }
+
+ private Set<Function<MetricsReader, NamedTestResult>> getReadSuppliers(
+ String uuid, String timestamp) {
+ Set<Function<MetricsReader, NamedTestResult>> suppliers = new HashSet<>();
+ suppliers.add(
+ reader -> {
+ long readStart = reader.getStartTimeMetric(READ_TIME_METRIC);
+ long readEnd = reader.getEndTimeMetric(READ_TIME_METRIC);
+ return NamedTestResult.create(
+ uuid, timestamp, READ_TIME_METRIC, (readEnd - readStart) / 1e3);
+ });
+ suppliers.add(
+ reader -> {
+ long fieldsRead = reader.getCounterMetric(FIELDS_READ_METRIC);
+ return NamedTestResult.create(uuid, timestamp, FIELDS_READ_METRIC, fieldsRead);
+ });
+ return suppliers;
+ }
+
+ /** Options for this io performance test. */
+ public interface SQLBigQueryPerfTestOptions extends IOTestPipelineOptions {
+ @Description("BQ dataset for the metrics data")
+ String getMetricsBigQueryDataset();
+
+ void setMetricsBigQueryDataset(String dataset);
+
+ @Description("BQ table for metrics data")
+ String getMetricsBigQueryTable();
+
+ void setMetricsBigQueryTable(String table);
+ }
+}
diff --git a/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryPerfTable.java b/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryPerfTable.java
new file mode 100644
index 0000000..98a1330
--- /dev/null
+++ b/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryPerfTable.java
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.bigquery;
+
+import java.util.List;
+import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTableFilter;
+import org.apache.beam.sdk.extensions.sql.meta.Table;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.ConversionOptions;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.Row;
+
+public class BigQueryPerfTable extends BigQueryTable {
+ private final String namespace;
+ private final String metric;
+
+ BigQueryPerfTable(Table table, ConversionOptions options, String namespace, String metric) {
+ super(table, options);
+ this.namespace = namespace;
+ this.metric = metric;
+ }
+
+ @Override
+ public PCollection<Row> buildIOReader(PBegin begin) {
+ return super.buildIOReader(begin).apply(ParDo.of(new RowMonitor(namespace, metric)));
+ }
+
+ @Override
+ public PCollection<Row> buildIOReader(
+ PBegin begin, BeamSqlTableFilter filters, List<String> fieldNames) {
+ return super.buildIOReader(begin, filters, fieldNames)
+ .apply(ParDo.of(new RowMonitor(namespace, metric)));
+ }
+
+ /** Monitor that records the number of Fields in each Row read from an IO. */
+ private static class RowMonitor extends DoFn<Row, Row> {
+
+ private Counter totalRows;
+
+ RowMonitor(String namespace, String name) {
+ this.totalRows = Metrics.counter(namespace, name);
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext c) {
+ totalRows.inc(c.element().getFieldCount());
+ c.output(c.element());
+ }
+ }
+}
diff --git a/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryPerfTableProvider.java b/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryPerfTableProvider.java
new file mode 100644
index 0000000..9488d7e
--- /dev/null
+++ b/sdks/java/extensions/sql/perf-tests/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryPerfTableProvider.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.bigquery;
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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 org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable;
+import org.apache.beam.sdk.extensions.sql.meta.Table;
+
+/** A test table provider for BigQueryIOPushDownIT. */
+public class BigQueryPerfTableProvider extends BigQueryTableProvider {
+ private final String namespace;
+ private final String metric;
+
+ BigQueryPerfTableProvider(String namespace, String metric) {
+ this.namespace = namespace;
+ this.metric = metric;
+ }
+
+ @Override
+ public BeamSqlTable buildBeamSqlTable(Table table) {
+ return new BigQueryPerfTable(
+ table, getConversionOptions(table.getProperties()), namespace, metric);
+ }
+}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java
index e061632..13186e0 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/SqlTransform.java
@@ -27,6 +27,7 @@
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv.BeamSqlEnvBuilder;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlPipelineOptions;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils;
import org.apache.beam.sdk.extensions.sql.impl.schema.BeamPCollectionTable;
import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable;
@@ -87,6 +88,8 @@
abstract String queryString();
+ abstract QueryParameters queryParameters();
+
abstract List<UdfDefinition> udfDefinitions();
abstract List<UdafDefinition> udafDefinitions();
@@ -122,7 +125,8 @@
sqlEnvBuilder.setPipelineOptions(input.getPipeline().getOptions());
BeamSqlEnv sqlEnv = sqlEnvBuilder.build();
- return BeamSqlRelUtils.toPCollection(input.getPipeline(), sqlEnv.parseQuery(queryString()));
+ return BeamSqlRelUtils.toPCollection(
+ input.getPipeline(), sqlEnv.parseQuery(queryString(), queryParameters()));
}
@SuppressWarnings("unchecked")
@@ -177,6 +181,7 @@
public static SqlTransform query(String queryString) {
return builder()
.setQueryString(queryString)
+ .setQueryParameters(QueryParameters.ofNone())
.setUdafDefinitions(Collections.emptyList())
.setUdfDefinitions(Collections.emptyList())
.setTableProviderMap(Collections.emptyMap())
@@ -194,6 +199,14 @@
return withTableProvider(name, tableProvider).toBuilder().setDefaultTableProvider(name).build();
}
+ public SqlTransform withNamedParameters(Map parameters) {
+ return toBuilder().setQueryParameters(QueryParameters.ofNamed(parameters)).build();
+ }
+
+ public SqlTransform withPositionalParameters(List parameters) {
+ return toBuilder().setQueryParameters(QueryParameters.ofPositional(parameters)).build();
+ }
+
public SqlTransform withAutoUdfUdafLoad(boolean autoUdfUdafLoad) {
return toBuilder().setAutoUdfUdafLoad(autoUdfUdafLoad).build();
}
@@ -245,6 +258,8 @@
abstract static class Builder {
abstract Builder setQueryString(String queryString);
+ abstract Builder setQueryParameters(QueryParameters queryParameters);
+
abstract Builder setUdfDefinitions(List<UdfDefinition> udfDefinitions);
abstract Builder setUdafDefinitions(List<UdafDefinition> udafDefinitions);
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java
index f27cb2e..b8b0e72 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlEnv.java
@@ -31,6 +31,7 @@
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.extensions.sql.BeamSqlUdf;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRuleSets;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
import org.apache.beam.sdk.extensions.sql.impl.udf.BeamBuiltinFunctionProvider;
@@ -100,7 +101,12 @@
}
public BeamRelNode parseQuery(String query) throws ParseException {
- return planner.convertToBeamRel(query);
+ return planner.convertToBeamRel(query, QueryParameters.ofNone());
+ }
+
+ public BeamRelNode parseQuery(String query, QueryParameters queryParameters)
+ throws ParseException {
+ return planner.convertToBeamRel(query, queryParameters);
}
public boolean isDdl(String sqlStatement) throws ParseException {
@@ -122,7 +128,7 @@
public String explain(String sqlString) throws ParseException {
try {
- return RelOptUtil.toString(planner.convertToBeamRel(sqlString));
+ return RelOptUtil.toString(planner.convertToBeamRel(sqlString, QueryParameters.ofNone()));
} catch (Exception e) {
throw new ParseException("Unable to parse statement", e);
}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java
index c367197..bddd3d5 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java
@@ -20,6 +20,7 @@
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters.Kind;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel;
import org.apache.beam.sdk.extensions.sql.impl.planner.RelMdNodeStats;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamLogicalConvention;
@@ -60,6 +61,7 @@
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.RuleSet;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.ValidationException;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.BuiltInMethod;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -133,10 +135,16 @@
return parsed;
}
- /** It parses and validate the input query, then convert into a {@link BeamRelNode} tree. */
+ /**
+ * It parses and validate the input query, then convert into a {@link BeamRelNode} tree. Note that
+ * query parameters are not yet supported.
+ */
@Override
- public BeamRelNode convertToBeamRel(String sqlStatement)
+ public BeamRelNode convertToBeamRel(String sqlStatement, QueryParameters queryParameters)
throws ParseException, SqlConversionException {
+ Preconditions.checkArgument(
+ queryParameters.getKind() == Kind.NONE,
+ "Beam SQL Calcite dialect does not yet support query parameters.");
BeamRelNode beamRelNode;
try {
SqlNode parsed = planner.parse(sqlStatement);
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java
index cec0045..755737a 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/QueryPlanner.java
@@ -17,6 +17,9 @@
*/
package org.apache.beam.sdk.extensions.sql.impl;
+import com.google.auto.value.AutoOneOf;
+import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlNode;
@@ -26,8 +29,38 @@
*/
public interface QueryPlanner {
/** It parses and validate the input query, then convert into a {@link BeamRelNode} tree. */
- BeamRelNode convertToBeamRel(String sqlStatement) throws ParseException, SqlConversionException;
+ BeamRelNode convertToBeamRel(String sqlStatement, QueryParameters queryParameters)
+ throws ParseException, SqlConversionException;
/** Parse input SQL query, and return a {@link SqlNode} as grammar tree. */
SqlNode parse(String sqlStatement) throws ParseException;
+
+ @AutoOneOf(QueryParameters.Kind.class)
+ abstract class QueryParameters {
+ public enum Kind {
+ NONE,
+ NAMED,
+ POSITIONAL
+ }
+
+ public abstract Kind getKind();
+
+ abstract void none();
+
+ public abstract Map named();
+
+ public abstract List positional();
+
+ public static QueryParameters ofNone() {
+ return AutoOneOf_QueryPlanner_QueryParameters.none();
+ }
+
+ public static QueryParameters ofNamed(Map namedParams) {
+ return AutoOneOf_QueryPlanner_QueryParameters.named(namedParams);
+ }
+
+ public static QueryParameters ofPositional(List positionalParams) {
+ return AutoOneOf_QueryPlanner_QueryParameters.positional(positionalParams);
+ }
+ }
}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java
index c3761bb..29d82d3 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java
@@ -31,6 +31,7 @@
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.Schema.TypeName;
import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.ByteString;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.NlsString;
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVParser;
@@ -137,6 +138,10 @@
throw new UnsupportedOperationException(
String.format("Column type %s is not supported yet!", type));
}
+ } else if (type.getTypeName().isPrimitiveType()) {
+ if (TypeName.BYTES.equals(type.getTypeName()) && rawObj instanceof ByteString) {
+ return ((ByteString) rawObj).getBytes();
+ }
}
return rawObj;
}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/udf/BuiltinStringFunctions.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/udf/BuiltinStringFunctions.java
index 1a90bf5..b7f9318 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/udf/BuiltinStringFunctions.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/udf/BuiltinStringFunctions.java
@@ -21,12 +21,12 @@
import com.google.auto.service.AutoService;
import java.util.Arrays;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils;
import org.apache.beam.sdk.schemas.Schema.TypeName;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.function.Strict;
import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Hex;
-import org.apache.commons.lang3.ArrayUtils;
-import org.apache.commons.lang3.StringUtils;
/** BuiltinStringFunctions. */
@AutoService(BeamBuiltinFunctionProvider.class)
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamBigQuerySqlDialect.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamBigQuerySqlDialect.java
index ab847bb..0672162 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamBigQuerySqlDialect.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamBigQuerySqlDialect.java
@@ -21,19 +21,25 @@
import java.util.List;
import java.util.regex.Pattern;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.Casing;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.TimeUnit;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.config.NullCollation;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataType;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataTypeSystem;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlCall;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlDataTypeSpec;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlDialect;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlIdentifier;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlIntervalLiteral;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlIntervalQualifier;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlKind;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlLiteral;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlOperator;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlSetOperator;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlSyntax;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlWriter;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.dialect.BigQuerySqlDialect;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.fun.SqlTrimFunction;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.BasicSqlType;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
@@ -182,6 +188,11 @@
}
@Override
+ public boolean supportsNestedAggregations() {
+ return false;
+ }
+
+ @Override
public void unparseOffsetFetch(SqlWriter writer, SqlNode offset, SqlNode fetch) {
unparseFetchUsingLimit(writer, offset, fetch);
}
@@ -216,11 +227,98 @@
SqlSyntax.BINARY.unparse(writer, INTERSECT_DISTINCT, call, leftPrec, rightPrec);
}
break;
+ case TRIM:
+ unparseTrim(writer, call, leftPrec, rightPrec);
+ break;
default:
super.unparseCall(writer, call, leftPrec, rightPrec);
}
}
+ /** BigQuery interval syntax: INTERVAL int64 time_unit. */
+ @Override
+ public void unparseSqlIntervalLiteral(
+ SqlWriter writer, SqlIntervalLiteral literal, int leftPrec, int rightPrec) {
+ SqlIntervalLiteral.IntervalValue interval =
+ (SqlIntervalLiteral.IntervalValue) literal.getValue();
+ writer.keyword("INTERVAL");
+ if (interval.getSign() == -1) {
+ writer.print("-");
+ }
+ Long intervalValueInLong;
+ try {
+ intervalValueInLong = Long.parseLong(literal.getValue().toString());
+ } catch (NumberFormatException e) {
+ throw new RuntimeException("Only INT64 is supported as the interval value for BigQuery.");
+ }
+ writer.literal(intervalValueInLong.toString());
+ unparseSqlIntervalQualifier(writer, interval.getIntervalQualifier(), RelDataTypeSystem.DEFAULT);
+ }
+
+ @Override
+ public void unparseSqlIntervalQualifier(
+ SqlWriter writer, SqlIntervalQualifier qualifier, RelDataTypeSystem typeSystem) {
+ final String start = validate(qualifier.timeUnitRange.startUnit).name();
+ if (qualifier.timeUnitRange.endUnit == null) {
+ writer.keyword(start);
+ } else {
+ throw new RuntimeException("Range time unit is not supported for BigQuery.");
+ }
+ }
+
+ /**
+ * For usage of TRIM, LTRIM and RTRIM in BQ see <a
+ * href="https://cloud.google.com/bigquery/docs/reference/standard-sql/functions-and-operators#trim">
+ * BQ Trim Function</a>.
+ */
+ private void unparseTrim(SqlWriter writer, SqlCall call, int leftPrec, int rightPrec) {
+ final String operatorName;
+ SqlLiteral trimFlag = call.operand(0);
+ SqlLiteral valueToTrim = call.operand(1);
+ switch (trimFlag.getValueAs(SqlTrimFunction.Flag.class)) {
+ case LEADING:
+ operatorName = "LTRIM";
+ break;
+ case TRAILING:
+ operatorName = "RTRIM";
+ break;
+ default:
+ operatorName = call.getOperator().getName();
+ break;
+ }
+ final SqlWriter.Frame trimFrame = writer.startFunCall(operatorName);
+ call.operand(2).unparse(writer, leftPrec, rightPrec);
+
+ /**
+ * If the trimmed character is non space character then add it to the target sql. eg: TRIM(BOTH
+ * 'A' from 'ABCD' Output Query: TRIM('ABC', 'A')
+ */
+ if (!valueToTrim.toValue().matches("\\s+")) {
+ writer.literal(",");
+ call.operand(1).unparse(writer, leftPrec, rightPrec);
+ }
+ writer.endFunCall(trimFrame);
+ }
+
+ private TimeUnit validate(TimeUnit timeUnit) {
+ switch (timeUnit) {
+ case MICROSECOND:
+ case MILLISECOND:
+ case SECOND:
+ case MINUTE:
+ case HOUR:
+ case DAY:
+ case WEEK:
+ case MONTH:
+ case QUARTER:
+ case YEAR:
+ case ISOYEAR:
+ return timeUnit;
+ default:
+ throw new RuntimeException("Time unit " + timeUnit + " is not supported for BigQuery.");
+ }
+ }
+
/**
* BigQuery data type reference: <a
* href="https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types">Bigquery
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java
index 432b3b5..d66d88c 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java
@@ -20,6 +20,12 @@
import static org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rel2sql.SqlImplementor.POS;
import java.util.function.IntFunction;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.text.translate.CharSequenceTranslator;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.text.translate.EntityArrays;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.text.translate.JavaUnicodeEscaper;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.text.translate.LookupTranslator;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.ByteString;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.TimeUnitRange;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rel2sql.SqlImplementor;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexLiteral;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexNode;
@@ -32,10 +38,28 @@
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.SqlTypeFamily;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.SqlTypeName;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.BitString;
-import org.apache.beam.vendor.calcite.v1_20_0.org.apache.commons.lang.StringEscapeUtils;
public class BeamSqlUnparseContext extends SqlImplementor.SimpleContext {
+ // More about escape sequences here:
+ // https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical
+ // No need to escape: \`, \?, \v, \a, \ooo, \xhh (since this in not a thing in Java)
+ // TODO: Move away from deprecated classes.
+ // TODO: Escaping single quotes, SqlCharStringLiteral (produced by SqlLiteral.createCharString)
+ // introduces extra.
+ private static final CharSequenceTranslator ESCAPE_FOR_ZETA_SQL =
+ // ZetaSQL specific:
+ new LookupTranslator(
+ new String[][] {
+ {"\"", "\\\""},
+ {"\\", "\\\\"},
+ })
+ // \b, \n, \t, \f, \r
+ .with(new LookupTranslator(EntityArrays.JAVA_CTRL_CHARS_ESCAPE()))
+ // TODO(BEAM-9180): Add support for \Uhhhhhhhh
+ // Unicode (only 4 hex digits)
+ .with(JavaUnicodeEscaper.outsideOf(32, 0x7f));
+
public BeamSqlUnparseContext(IntFunction<SqlNode> field) {
super(BeamBigQuerySqlDialect.DEFAULT, field);
}
@@ -46,11 +70,21 @@
final RexLiteral literal = (RexLiteral) rex;
SqlTypeFamily family = literal.getTypeName().getFamily();
if (SqlTypeFamily.BINARY.equals(family)) {
- BitString bitString = BitString.createFromBytes(literal.getValueAs(byte[].class));
+ ByteString byteString = literal.getValueAs(ByteString.class);
+ BitString bitString = BitString.createFromHexString(byteString.toString(16));
return new SqlByteStringLiteral(bitString, POS);
} else if (SqlTypeFamily.CHARACTER.equals(family)) {
- String escaped = StringEscapeUtils.escapeJava(literal.getValueAs(String.class));
+ String escaped = ESCAPE_FOR_ZETA_SQL.translate(literal.getValueAs(String.class));
return SqlLiteral.createCharString(escaped, POS);
+ } else if (SqlTypeName.SYMBOL.equals(literal.getTypeName())) {
+ Enum symbol = literal.getValueAs(Enum.class);
+ if (TimeUnitRange.DOW.equals(symbol)) {
+ return new ReplaceLiteral(literal, POS, "DAYOFWEEK");
+ } else if (TimeUnitRange.DOY.equals(symbol)) {
+ return new ReplaceLiteral(literal, POS, "DAYOFYEAR");
+ } else if (TimeUnitRange.WEEK.equals(symbol)) {
+ return new ReplaceLiteral(literal, POS, "ISOWEEK");
+ }
}
}
@@ -81,4 +115,35 @@
writer.literal(builder.toString());
}
}
+
+ private static class ReplaceLiteral extends SqlLiteral {
+
+ private final String newValue;
+
+ ReplaceLiteral(RexLiteral literal, SqlParserPos pos, String newValue) {
+ super(literal.getValue(), literal.getTypeName(), pos);
+ this.newValue = newValue;
+ }
+
+ @Override
+ public void unparse(SqlWriter writer, int leftPrec, int rightPrec) {
+ writer.literal(newValue);
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (!(obj instanceof ReplaceLiteral)) {
+ return false;
+ }
+ if (!newValue.equals(((ReplaceLiteral) obj).newValue)) {
+ return false;
+ }
+ return super.equals(obj);
+ }
+
+ @Override
+ public int hashCode() {
+ return super.hashCode();
+ }
+ }
}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryTableProvider.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryTableProvider.java
index 9c4266b..b1646aa 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryTableProvider.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BigQueryTableProvider.java
@@ -19,6 +19,7 @@
import static org.apache.beam.vendor.calcite.v1_20_0.com.google.common.base.MoreObjects.firstNonNull;
+import com.alibaba.fastjson.JSONObject;
import com.google.auto.service.AutoService;
import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable;
import org.apache.beam.sdk.extensions.sql.meta.Table;
@@ -52,13 +53,15 @@
@Override
public BeamSqlTable buildBeamSqlTable(Table table) {
- return new BigQueryTable(
- table,
- ConversionOptions.builder()
- .setTruncateTimestamps(
- firstNonNull(table.getProperties().getBoolean("truncateTimestamps"), false)
- ? TruncateTimestamps.TRUNCATE
- : TruncateTimestamps.REJECT)
- .build());
+ return new BigQueryTable(table, getConversionOptions(table.getProperties()));
+ }
+
+ protected static ConversionOptions getConversionOptions(JSONObject properties) {
+ return ConversionOptions.builder()
+ .setTruncateTimestamps(
+ firstNonNull(properties.getBoolean("truncateTimestamps"), false)
+ ? TruncateTimestamps.TRUNCATE
+ : TruncateTimestamps.REJECT)
+ .build();
}
}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreV1Table.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreV1Table.java
new file mode 100644
index 0000000..92b9326
--- /dev/null
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreV1Table.java
@@ -0,0 +1,426 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.datastore;
+
+import static com.google.datastore.v1.client.DatastoreHelper.makeKey;
+import static com.google.datastore.v1.client.DatastoreHelper.makeValue;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+
+import com.alibaba.fastjson.JSONObject;
+import com.google.datastore.v1.Entity;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.Query;
+import com.google.datastore.v1.Value;
+import com.google.datastore.v1.Value.ValueTypeCase;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.InvalidProtocolBufferException;
+import java.io.Serializable;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.function.Supplier;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.extensions.sql.impl.BeamTableStatistics;
+import org.apache.beam.sdk.extensions.sql.meta.SchemaBaseBeamTable;
+import org.apache.beam.sdk.extensions.sql.meta.Table;
+import org.apache.beam.sdk.io.gcp.datastore.DatastoreIO;
+import org.apache.beam.sdk.io.gcp.datastore.DatastoreV1;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.Schema.TypeName;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollection.IsBounded;
+import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.calcite.v1_20_0.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@Internal
+@Experimental
+class DataStoreV1Table extends SchemaBaseBeamTable implements Serializable {
+ public static final String KEY_FIELD_PROPERTY = "keyField";
+ @VisibleForTesting static final String DEFAULT_KEY_FIELD = "__key__";
+ private static final Logger LOGGER = LoggerFactory.getLogger(DataStoreV1Table.class);
+ // Should match: `projectId/kind`.
+ private static final Pattern locationPattern = Pattern.compile("(?<projectId>.+)/(?<kind>.+)");
+ @VisibleForTesting final String keyField;
+ @VisibleForTesting final String projectId;
+ @VisibleForTesting final String kind;
+
+ DataStoreV1Table(Table table) {
+ super(table.getSchema());
+
+ // TODO: allow users to specify a name of the field to store a key value via TableProperties.
+ JSONObject properties = table.getProperties();
+ if (properties.containsKey(KEY_FIELD_PROPERTY)) {
+ String field = properties.getString(KEY_FIELD_PROPERTY);
+ checkArgument(
+ field != null && !field.isEmpty(), "'%s' property cannot be null.", KEY_FIELD_PROPERTY);
+ keyField = field;
+ } else {
+ keyField = DEFAULT_KEY_FIELD;
+ }
+ // TODO: allow users to specify a namespace in a location string.
+ String location = table.getLocation();
+ checkArgument(location != null, "DataStoreV1 location must be set.");
+ Matcher matcher = locationPattern.matcher(location);
+ checkArgument(
+ matcher.matches(),
+ "DataStoreV1 location must be in the following format: 'projectId/kind'");
+
+ this.projectId = matcher.group("projectId");
+ this.kind = matcher.group("kind");
+ }
+
+ @Override
+ public PCollection<Row> buildIOReader(PBegin begin) {
+ Query.Builder q = Query.newBuilder();
+ q.addKindBuilder().setName(kind);
+ Query query = q.build();
+
+ DatastoreV1.Read readInstance =
+ DatastoreIO.v1().read().withProjectId(projectId).withQuery(query);
+
+ return begin
+ .apply("Read Datastore Entities", readInstance)
+ .apply("Convert Datastore Entities to Rows", EntityToRow.create(getSchema(), keyField));
+ }
+
+ @Override
+ public POutput buildIOWriter(PCollection<Row> input) {
+ return input
+ .apply("Convert Rows to Datastore Entities", RowToEntity.create(keyField, kind))
+ .apply("Write Datastore Entities", DatastoreIO.v1().write().withProjectId(projectId));
+ }
+
+ @Override
+ public IsBounded isBounded() {
+ return IsBounded.BOUNDED;
+ }
+
+ @Override
+ public BeamTableStatistics getTableStatistics(PipelineOptions options) {
+ long count =
+ DatastoreIO.v1().read().withProjectId(projectId).getNumEntities(options, kind, null);
+
+ if (count < 0) {
+ return BeamTableStatistics.BOUNDED_UNKNOWN;
+ }
+
+ return BeamTableStatistics.createBoundedTableStatistics((double) count);
+ }
+
+ /**
+ * A {@code PTransform} to perform a conversion of {@code PCollection<Entity>} to {@code
+ * PCollection<Row>}.
+ */
+ public static class EntityToRow extends PTransform<PCollection<Entity>, PCollection<Row>> {
+ private final Schema schema;
+ private final String keyField;
+
+ private EntityToRow(Schema schema, String keyField) {
+ this.schema = schema;
+ this.keyField = keyField;
+
+ if (schema.getFieldNames().contains(keyField)) {
+ if (!schema.getField(keyField).getType().getTypeName().equals(TypeName.BYTES)) {
+ throw new IllegalStateException(
+ "Field `"
+ + keyField
+ + "` should of type `VARBINARY`. Please change the type or specify a field to store the KEY value.");
+ }
+ LOGGER.info("Entity KEY will be stored under `" + keyField + "` field.");
+ }
+ }
+
+ /**
+ * Create a PTransform instance.
+ *
+ * @param schema {@code Schema} of the target row.
+ * @param keyField A name of the row field to store the {@code Key} in.
+ * @return {@code PTransform} instance for Entity to Row conversion.
+ */
+ public static EntityToRow create(Schema schema, String keyField) {
+ return new EntityToRow(schema, keyField);
+ }
+
+ @Override
+ public PCollection<Row> expand(PCollection<Entity> input) {
+ return input.apply(ParDo.of(new EntityToRowConverter())).setRowSchema(schema);
+ }
+
+ @VisibleForTesting
+ class EntityToRowConverter extends DoFn<Entity, Row> {
+
+ @DoFn.ProcessElement
+ public void processElement(ProcessContext context) {
+ Entity entity = context.element();
+ ImmutableMap.Builder<String, Value> mapBuilder = ImmutableMap.builder();
+ mapBuilder.put(keyField, makeValue(entity.getKey()).build());
+ mapBuilder.putAll(entity.getPropertiesMap());
+
+ context.output(extractRowFromProperties(schema, mapBuilder.build()));
+ }
+
+ /**
+ * Convert DataStore {@code Value} to Beam type.
+ *
+ * @param currentFieldType Beam {@code Schema.FieldType} to convert to (used for {@code Row}
+ * and {@code Array}).
+ * @param val DataStore {@code Value}.
+ * @return resulting Beam type.
+ */
+ private Object convertValueToObject(FieldType currentFieldType, Value val) {
+ ValueTypeCase typeCase = val.getValueTypeCase();
+
+ switch (typeCase) {
+ case NULL_VALUE:
+ case VALUETYPE_NOT_SET:
+ return null;
+ case BOOLEAN_VALUE:
+ return val.getBooleanValue();
+ case INTEGER_VALUE:
+ return val.getIntegerValue();
+ case DOUBLE_VALUE:
+ return val.getDoubleValue();
+ case TIMESTAMP_VALUE:
+ com.google.protobuf.Timestamp time = val.getTimestampValue();
+ long millis = time.getSeconds() * 1000 + time.getNanos() / 1000;
+ return Instant.ofEpochMilli(millis).toDateTime();
+ case STRING_VALUE:
+ return val.getStringValue();
+ case KEY_VALUE:
+ return val.getKeyValue().toByteArray();
+ case BLOB_VALUE:
+ return val.getBlobValue().toByteArray();
+ case ENTITY_VALUE:
+ // Recursive mapping for row type.
+ Schema rowSchema = currentFieldType.getRowSchema();
+ assert rowSchema != null;
+ Entity entity = val.getEntityValue();
+ return extractRowFromProperties(rowSchema, entity.getPropertiesMap());
+ case ARRAY_VALUE:
+ // Recursive mapping for collection type.
+ FieldType elementType = currentFieldType.getCollectionElementType();
+ List<Value> valueList = val.getArrayValue().getValuesList();
+ return valueList.stream()
+ .map(v -> convertValueToObject(elementType, v))
+ .collect(Collectors.toList());
+ case GEO_POINT_VALUE:
+ default:
+ throw new IllegalStateException(
+ "No conversion exists from type: "
+ + val.getValueTypeCase().name()
+ + " to Beam type.");
+ }
+ }
+
+ /**
+ * Converts all properties of an {@code Entity} to Beam {@code Row}.
+ *
+ * @param schema Target row {@code Schema}.
+ * @param values A map of property names and values.
+ * @return resulting Beam {@code Row}.
+ */
+ private Row extractRowFromProperties(Schema schema, Map<String, Value> values) {
+ Row.Builder builder = Row.withSchema(schema);
+ // It is not a guarantee that the values will be in the same order as the schema.
+ // Maybe metadata:
+ // https://cloud.google.com/appengine/docs/standard/python/datastore/metadataqueries
+ // TODO: figure out in what order the elements are in (without relying on Beam schema).
+ for (Schema.Field field : schema.getFields()) {
+ Value val = values.get(field.getName());
+ builder.addValue(convertValueToObject(field.getType(), val));
+ }
+ return builder.build();
+ }
+ }
+ }
+
+ /**
+ * A {@code PTransform} to perform a conversion of {@code PCollection<Row>} to {@code
+ * PCollection<Entity>}.
+ */
+ public static class RowToEntity extends PTransform<PCollection<Row>, PCollection<Entity>> {
+ private final Supplier<String> keySupplier;
+ private final String kind;
+ private final String keyField;
+
+ private RowToEntity(Supplier<String> keySupplier, String kind, String keyField) {
+ this.keySupplier = keySupplier;
+ this.kind = kind;
+ this.keyField = keyField;
+ }
+
+ @Override
+ public PCollection<Entity> expand(PCollection<Row> input) {
+ boolean isFieldPresent = input.getSchema().getFieldNames().contains(keyField);
+ if (isFieldPresent) {
+ if (!input.getSchema().getField(keyField).getType().getTypeName().equals(TypeName.BYTES)) {
+ throw new IllegalStateException(
+ "Field `"
+ + keyField
+ + "` should of type `VARBINARY`. Please change the type or specify a field to write the KEY value from via TableProperties.");
+ }
+ LOGGER.info("Field to use as Entity KEY is set to: `" + keyField + "`.");
+ }
+ return input.apply(ParDo.of(new RowToEntityConverter(isFieldPresent)));
+ }
+
+ /**
+ * Create a PTransform instance.
+ *
+ * @param keyField Row field containing a serialized {@code Key}, must be set when using user
+ * specified keys.
+ * @param kind DataStore `Kind` data will be written to (required when generating random {@code
+ * Key}s).
+ * @return {@code PTransform} instance for Row to Entity conversion.
+ */
+ public static RowToEntity create(String keyField, String kind) {
+ return new RowToEntity(
+ (Supplier<String> & Serializable) () -> UUID.randomUUID().toString(), kind, keyField);
+ }
+
+ @VisibleForTesting
+ static RowToEntity createTest(String keyString, String keyField, String kind) {
+ return new RowToEntity((Supplier<String> & Serializable) () -> keyString, kind, keyField);
+ }
+
+ @VisibleForTesting
+ class RowToEntityConverter extends DoFn<Row, Entity> {
+ private final boolean useNonRandomKey;
+
+ RowToEntityConverter(boolean useNonRandomKey) {
+ super();
+ this.useNonRandomKey = useNonRandomKey;
+ }
+
+ @DoFn.ProcessElement
+ public void processElement(ProcessContext context) {
+ Row row = context.element();
+
+ Schema schemaWithoutKeyField =
+ Schema.builder()
+ .addFields(
+ row.getSchema().getFields().stream()
+ .filter(field -> !field.getName().equals(keyField))
+ .collect(Collectors.toList()))
+ .build();
+ Entity.Builder entityBuilder = constructEntityFromRow(schemaWithoutKeyField, row);
+ entityBuilder.setKey(constructKeyFromRow(row));
+
+ context.output(entityBuilder.build());
+ }
+
+ /**
+ * Converts an entire {@code Row} to an appropriate DataStore {@code Entity.Builder}.
+ *
+ * @param row {@code Row} to convert.
+ * @return resulting {@code Entity.Builder}.
+ */
+ private Entity.Builder constructEntityFromRow(Schema schema, Row row) {
+ Entity.Builder entityBuilder = Entity.newBuilder();
+ for (Schema.Field field : schema.getFields()) {
+ Value val = mapObjectToValue(row.getValue(field.getName()));
+ entityBuilder.putProperties(field.getName(), val);
+ }
+ return entityBuilder;
+ }
+
+ /**
+ * Create a random key for a {@code Row} without a keyField or use a user-specified key by
+ * parsing it from byte array when keyField is set.
+ *
+ * @param row {@code Row} to construct a key for.
+ * @return resulting {@code Key}.
+ */
+ private Key constructKeyFromRow(Row row) {
+ if (!useNonRandomKey) {
+ // When key field is not present - use key supplier to generate a random one.
+ return makeKey(kind, keySupplier.get()).build();
+ }
+ byte[] keyBytes = row.getBytes(keyField);
+ try {
+ return Key.parseFrom(keyBytes);
+ } catch (InvalidProtocolBufferException e) {
+ throw new IllegalStateException("Failed to parse DataStore key from bytes.");
+ }
+ }
+
+ /**
+ * Converts a {@code Row} value to an appropriate DataStore {@code Value} object.
+ *
+ * @param value {@code Row} value to convert.
+ * @throws IllegalStateException when no mapping function for object of given type exists.
+ * @return resulting {@code Value}.
+ */
+ private Value mapObjectToValue(Object value) {
+ if (value == null) {
+ return Value.newBuilder().build();
+ }
+
+ if (Boolean.class.equals(value.getClass())) {
+ return makeValue((Boolean) value).build();
+ } else if (Byte.class.equals(value.getClass())) {
+ return makeValue((Byte) value).build();
+ } else if (Long.class.equals(value.getClass())) {
+ return makeValue((Long) value).build();
+ } else if (Short.class.equals(value.getClass())) {
+ return makeValue((Short) value).build();
+ } else if (Integer.class.equals(value.getClass())) {
+ return makeValue((Integer) value).build();
+ } else if (Double.class.equals(value.getClass())) {
+ return makeValue((Double) value).build();
+ } else if (Float.class.equals(value.getClass())) {
+ return makeValue((Float) value).build();
+ } else if (String.class.equals(value.getClass())) {
+ return makeValue((String) value).build();
+ } else if (Instant.class.equals(value.getClass())) {
+ return makeValue(((Instant) value).toDate()).build();
+ } else if (byte[].class.equals(value.getClass())) {
+ return makeValue(ByteString.copyFrom((byte[]) value)).build();
+ } else if (value instanceof Row) {
+ // Recursive conversion to handle nested rows.
+ Row row = (Row) value;
+ return makeValue(constructEntityFromRow(row.getSchema(), row)).build();
+ } else if (value instanceof Collection) {
+ // Recursive to handle nested collections.
+ Collection<Object> collection = (Collection<Object>) value;
+ List<Value> arrayValues =
+ collection.stream().map(this::mapObjectToValue).collect(Collectors.toList());
+ return makeValue(arrayValues).build();
+ }
+ throw new IllegalStateException(
+ "No conversion exists from type: " + value.getClass() + " to DataStove Value.");
+ }
+ }
+ }
+}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreV1TableProvider.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreV1TableProvider.java
new file mode 100644
index 0000000..9ecf668
--- /dev/null
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreV1TableProvider.java
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.datastore;
+
+import com.google.auto.service.AutoService;
+import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable;
+import org.apache.beam.sdk.extensions.sql.meta.Table;
+import org.apache.beam.sdk.extensions.sql.meta.provider.InMemoryMetaTableProvider;
+import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
+
+/**
+ * {@link TableProvider} for {@link DataStoreV1Table}.
+ *
+ * <p>A sample of DataStoreV1Table table is:
+ *
+ * <pre>{@code
+ * CREATE TABLE ORDERS(
+ * name VARCHAR,
+ * favorite_color VARCHAR,
+ * favorite_numbers ARRAY<INTEGER>
+ * )
+ * TYPE 'datastoreV1'
+ * LOCATION 'projectId/kind'
+ * }</pre>
+ */
+@AutoService(TableProvider.class)
+public class DataStoreV1TableProvider extends InMemoryMetaTableProvider {
+
+ @Override
+ public String getTableType() {
+ return "datastoreV1";
+ }
+
+ @Override
+ public BeamSqlTable buildBeamSqlTable(Table table) {
+ return new DataStoreV1Table(table);
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/package-info.java
similarity index 65%
rename from runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
rename to sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/package-info.java
index a114f40..5c743e7 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/utils/FlinkClassloading.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/package-info.java
@@ -15,16 +15,6 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.runners.flink.translation.utils;
-import com.fasterxml.jackson.databind.type.TypeFactory;
-
-/** Utilities for dealing with classloading. */
-public class FlinkClassloading {
-
- public static void deleteStaticCaches() {
- // Clear cache to get rid of any references to the Flink Classloader
- // See https://jira.apache.org/jira/browse/BEAM-6460
- TypeFactory.defaultInstance().clearCache();
- }
-}
+/** Table schema for DataStore. */
+package org.apache.beam.sdk.extensions.sql.meta.provider.datastore;
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbTable.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbTable.java
index 9b06a12..4da3d35e 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbTable.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbTable.java
@@ -17,12 +17,19 @@
*/
package org.apache.beam.sdk.extensions.sql.meta.provider.mongodb;
+import static org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlKind.AND;
+import static org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlKind.COMPARISON;
+import static org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlKind.OR;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import com.mongodb.client.model.Filters;
import java.io.Serializable;
+import java.util.ArrayList;
import java.util.List;
+import java.util.function.IntFunction;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
+import java.util.stream.Collectors;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.extensions.sql.impl.BeamTableStatistics;
import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTableFilter;
@@ -34,7 +41,9 @@
import org.apache.beam.sdk.io.mongodb.MongoDbIO;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
+import org.apache.beam.sdk.schemas.FieldTypeDescriptors;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.utils.SelectHelpers;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.JsonToRow;
@@ -49,12 +58,23 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.calcite.v1_20_0.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexCall;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexInputRef;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexLiteral;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexNode;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlKind;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.bson.Document;
+import org.bson.conversions.Bson;
import org.bson.json.JsonMode;
import org.bson.json.JsonWriterSettings;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
@Experimental
public class MongoDbTable extends SchemaBaseBeamTable implements Serializable {
+ private static final Logger LOGGER = LoggerFactory.getLogger(MongoDbTable.class);
// Should match: mongodb://username:password@localhost:27017/database/collection
@VisibleForTesting
final Pattern locationPattern =
@@ -103,14 +123,23 @@
.resolve(getSchema());
final Schema newSchema = SelectHelpers.getOutputSchema(getSchema(), resolved);
+ FindQuery findQuery = FindQuery.create();
+
if (!(filters instanceof DefaultTableFilter)) {
- throw new AssertionError("Predicate push-down is unsupported, yet received a predicate.");
+ MongoDbFilter mongoFilter = (MongoDbFilter) filters;
+ if (!mongoFilter.getSupported().isEmpty()) {
+ Bson filter = constructPredicate(mongoFilter.getSupported());
+ LOGGER.info("Pushing down the following filter: " + filter.toString());
+ findQuery = findQuery.withFilters(filter);
+ }
}
if (!fieldNames.isEmpty()) {
- readInstance = readInstance.withQueryFn(FindQuery.create().withProjection(fieldNames));
+ findQuery = findQuery.withProjection(fieldNames);
}
+ readInstance = readInstance.withQueryFn(findQuery);
+
return readInstance.expand(begin).apply(DocumentToRow.withSchema(newSchema));
}
@@ -127,6 +156,131 @@
}
@Override
+ public BeamSqlTableFilter constructFilter(List<RexNode> filter) {
+ return MongoDbFilter.create(filter);
+ }
+
+ /**
+ * Given a predicate in a conjunctive normal form (CNF), construct a {@code Bson} filter for
+ * MongoDB find query.
+ *
+ * @param supported A list of {@code RexNode} in CNF.
+ * @return {@code Bson} filter.
+ */
+ private Bson constructPredicate(List<RexNode> supported) {
+ assert !supported.isEmpty();
+ List<Bson> cnf =
+ supported.stream().map(this::translateRexNodeToBson).collect(Collectors.toList());
+ if (cnf.size() == 1) {
+ return cnf.get(0);
+ }
+ return Filters.and(cnf);
+ }
+
+ /**
+ * Recursively translates a single RexNode to MongoDB Bson filter. Supports simple comparison
+ * operations, negation, and nested conjunction/disjunction. Boolean fields are translated as an
+ * `$eq` operation with a boolean `true`.
+ *
+ * @param node {@code RexNode} to translate.
+ * @return {@code Bson} filter.
+ */
+ private Bson translateRexNodeToBson(RexNode node) {
+ final IntFunction<String> fieldIdToName = i -> getSchema().getField(i).getName();
+ // Supported operations are described in MongoDbFilter#isSupported
+ if (node instanceof RexCall) {
+ RexCall compositeNode = (RexCall) node;
+ List<RexLiteral> literals = new ArrayList<>();
+ List<RexInputRef> inputRefs = new ArrayList<>();
+
+ for (RexNode operand : compositeNode.getOperands()) {
+ if (operand instanceof RexLiteral) {
+ literals.add((RexLiteral) operand);
+ } else if (operand instanceof RexInputRef) {
+ inputRefs.add((RexInputRef) operand);
+ }
+ }
+
+ // Operation is a comparison, since one of the operands in a field reference.
+ if (inputRefs.size() == 1) {
+ RexInputRef inputRef = inputRefs.get(0);
+ String inputFieldName = fieldIdToName.apply(inputRef.getIndex());
+ if (literals.size() > 0) {
+ // Convert literal value to the same Java type as the field we are comparing to.
+ Object literal = convertToExpectedType(inputRef, literals.get(0));
+
+ switch (node.getKind()) {
+ case IN:
+ return Filters.in(inputFieldName, convertToExpectedType(inputRef, literals));
+ case EQUALS:
+ return Filters.eq(inputFieldName, literal);
+ case NOT_EQUALS:
+ return Filters.not(Filters.eq(inputFieldName, literal));
+ case LESS_THAN:
+ return Filters.lt(inputFieldName, literal);
+ case GREATER_THAN:
+ return Filters.gt(inputFieldName, literal);
+ case GREATER_THAN_OR_EQUAL:
+ return Filters.gte(inputFieldName, literal);
+ case LESS_THAN_OR_EQUAL:
+ return Filters.lte(inputFieldName, literal);
+ default:
+ // Encountered an unexpected node kind, RuntimeException below.
+ break;
+ }
+ } else if (node.getKind().equals(SqlKind.NOT)) {
+ // Ex: `where not boolean_field`
+ return Filters.not(translateRexNodeToBson(inputRef));
+ } else {
+ throw new RuntimeException(
+ "Cannot create a filter for an unsupported node: " + node.toString());
+ }
+ } else { // Operation is a conjunction/disjunction.
+ switch (node.getKind()) {
+ case AND:
+ // Recursively construct filter for each operand of conjunction.
+ return Filters.and(
+ compositeNode.getOperands().stream()
+ .map(this::translateRexNodeToBson)
+ .collect(Collectors.toList()));
+ case OR:
+ // Recursively construct filter for each operand of disjunction.
+ return Filters.or(
+ compositeNode.getOperands().stream()
+ .map(this::translateRexNodeToBson)
+ .collect(Collectors.toList()));
+ default:
+ // Encountered an unexpected node kind, RuntimeException below.
+ break;
+ }
+ }
+ throw new RuntimeException(
+ "Encountered an unexpected node kind: " + node.getKind().toString());
+ } else if (node instanceof RexInputRef
+ && node.getType().getSqlTypeName().equals(SqlTypeName.BOOLEAN)) {
+ // Boolean field, must be true. Ex: `select * from table where bool_field`
+ return Filters.eq(fieldIdToName.apply(((RexInputRef) node).getIndex()), true);
+ }
+
+ throw new RuntimeException(
+ "Was expecting a RexCall or a boolean RexInputRef, but received: "
+ + node.getClass().getSimpleName());
+ }
+
+ private Object convertToExpectedType(RexInputRef inputRef, RexLiteral literal) {
+ FieldType beamFieldType = getSchema().getField(inputRef.getIndex()).getType();
+
+ return literal.getValueAs(
+ FieldTypeDescriptors.javaTypeForFieldType(beamFieldType).getRawType());
+ }
+
+ private Object convertToExpectedType(RexInputRef inputRef, List<RexLiteral> literals) {
+ return literals.stream()
+ .map(l -> convertToExpectedType(inputRef, l))
+ .collect(Collectors.toList());
+ }
+
+ @Override
public IsBounded isBounded() {
return IsBounded.BOUNDED;
}
@@ -205,4 +359,117 @@
}
}
}
+
+ static class MongoDbFilter implements BeamSqlTableFilter {
+ private List<RexNode> supported;
+ private List<RexNode> unsupported;
+
+ public MongoDbFilter(List<RexNode> supported, List<RexNode> unsupported) {
+ this.supported = supported;
+ this.unsupported = unsupported;
+ }
+
+ @Override
+ public List<RexNode> getNotSupported() {
+ return unsupported;
+ }
+
+ @Override
+ public int numSupported() {
+ return BeamSqlTableFilter.expressionsInFilter(supported);
+ }
+
+ public List<RexNode> getSupported() {
+ return supported;
+ }
+
+ @Override
+ public String toString() {
+ String supStr =
+ "supported{"
+ + supported.stream().map(RexNode::toString).collect(Collectors.joining())
+ + "}";
+ String unsupStr =
+ "unsupported{"
+ + unsupported.stream().map(RexNode::toString).collect(Collectors.joining())
+ + "}";
+
+ return "[" + supStr + ", " + unsupStr + "]";
+ }
+
+ public static MongoDbFilter create(List<RexNode> predicateCNF) {
+ ImmutableList.Builder<RexNode> supported = ImmutableList.builder();
+ ImmutableList.Builder<RexNode> unsupported = ImmutableList.builder();
+
+ for (RexNode node : predicateCNF) {
+ if (!node.getType().getSqlTypeName().equals(SqlTypeName.BOOLEAN)) {
+ throw new RuntimeException(
+ "Predicate node '"
+ + node.getClass().getSimpleName()
+ + "' should be a boolean expression, but was: "
+ + node.getType().getSqlTypeName());
+ }
+
+ if (isSupported(node)) {
+ supported.add(node);
+ } else {
+ unsupported.add(node);
+ }
+ }
+
+ return new MongoDbFilter(supported.build(), unsupported.build());
+ }
+
+ /**
+ * Check whether a {@code RexNode} is supported. To keep things simple:<br>
+ * 1. Support comparison operations in predicate, which compare a single field to literal
+ * values. 2. Support nested Conjunction (AND), Disjunction (OR) as long as child operations are
+ * supported.<br>
+ * 3. Support boolean fields.
+ *
+ * @param node A node to check for predicate push-down support.
+ * @return A boolean whether an expression is supported.
+ */
+ private static boolean isSupported(RexNode node) {
+ if (node instanceof RexCall) {
+ RexCall compositeNode = (RexCall) node;
+
+ if (node.getKind().belongsTo(COMPARISON) || node.getKind().equals(SqlKind.NOT)) {
+ int fields = 0;
+ for (RexNode operand : compositeNode.getOperands()) {
+ if (operand instanceof RexInputRef) {
+ fields++;
+ } else if (operand instanceof RexLiteral) {
+ // RexLiterals are expected, but no action is needed.
+ } else {
+ // Complex predicates are not supported. Ex: `field1+5 == 10`.
+ return false;
+ }
+ }
+ // All comparison operations should have exactly one field reference.
+ // Ex: `field1 == field2` is not supported.
+ // TODO: Can be supported via Filters#where.
+ if (fields == 1) {
+ return true;
+ }
+ } else if (node.getKind().equals(AND) || node.getKind().equals(OR)) {
+ // Nested ANDs and ORs are supported as long as all operands are supported.
+ for (RexNode operand : compositeNode.getOperands()) {
+ if (!isSupported(operand)) {
+ return false;
+ }
+ }
+ return true;
+ }
+ } else if (node instanceof RexInputRef) {
+ // When field is a boolean.
+ return true;
+ } else {
+ throw new RuntimeException(
+ "Encountered an unexpected node type: " + node.getClass().getSimpleName());
+ }
+
+ return false;
+ }
+ }
}
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/udf/BeamSalUhfSpecialTypeAndValueTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/udf/BeamSalUhfSpecialTypeAndValueTest.java
index ad59c8e..1370c62 100644
--- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/udf/BeamSalUhfSpecialTypeAndValueTest.java
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/udf/BeamSalUhfSpecialTypeAndValueTest.java
@@ -19,6 +19,7 @@
import static java.nio.charset.StandardCharsets.UTF_8;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils;
import org.apache.beam.sdk.extensions.sql.BeamSqlDslBase;
import org.apache.beam.sdk.extensions.sql.SqlTransform;
import org.apache.beam.sdk.schemas.Schema;
@@ -26,7 +27,6 @@
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
-import org.apache.commons.lang3.ArrayUtils;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreReadWriteIT.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreReadWriteIT.java
new file mode 100644
index 0000000..3274846
--- /dev/null
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreReadWriteIT.java
@@ -0,0 +1,231 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.datastore;
+
+import static com.google.datastore.v1.client.DatastoreHelper.makeFilter;
+import static com.google.datastore.v1.client.DatastoreHelper.makeKey;
+import static com.google.datastore.v1.client.DatastoreHelper.makeValue;
+import static org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.VARBINARY;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.BOOLEAN;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.DATETIME;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.DOUBLE;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.INT64;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.STRING;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.PropertyFilter.Operator;
+import com.google.datastore.v1.Query;
+import java.util.UUID;
+import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.PipelineResult.State;
+import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv;
+import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils;
+import org.apache.beam.sdk.extensions.sql.meta.provider.datastore.DataStoreV1Table.EntityToRow;
+import org.apache.beam.sdk.extensions.sql.meta.provider.datastore.DataStoreV1Table.RowToEntity;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryOptions;
+import org.apache.beam.sdk.io.gcp.datastore.DatastoreIO;
+import org.apache.beam.sdk.io.gcp.datastore.DatastoreV1;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+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.Row;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.ByteString;
+import org.joda.time.Duration;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class DataStoreReadWriteIT {
+ private static final BigQueryOptions options =
+ TestPipeline.testingPipelineOptions().as(BigQueryOptions.class);
+ private static final Schema SOURCE_SCHEMA =
+ Schema.builder()
+ .addNullableField("__key__", VARBINARY)
+ .addNullableField("content", STRING)
+ .build();
+ private static final Schema SOURCE_SCHEMA_WITHOUT_KEY =
+ Schema.builder().addNullableField("content", STRING).build();
+ private static final String KIND = "writereadtest";
+ private static final String KIND_ALL_TYPES = "writereadalltypestest";
+
+ @Rule public final TestPipeline writePipeline = TestPipeline.create();
+ @Rule public transient TestPipeline readPipeline = TestPipeline.create();
+
+ @Test
+ public void testDataStoreV1SqlWriteRead() {
+ BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new DataStoreV1TableProvider());
+ String projectId = options.getProject();
+
+ String createTableStatement =
+ "CREATE EXTERNAL TABLE TEST( \n"
+ + " `__key__` VARBINARY, \n"
+ + " `content` VARCHAR \n"
+ + ") \n"
+ + "TYPE 'datastoreV1' \n"
+ + "LOCATION '"
+ + projectId
+ + "/"
+ + KIND
+ + "'";
+ sqlEnv.executeDdl(createTableStatement);
+
+ Key ancestor = makeKey(KIND, UUID.randomUUID().toString()).build();
+ Key itemKey = makeKey(ancestor, KIND, UUID.randomUUID().toString()).build();
+ String insertStatement =
+ "INSERT INTO TEST VALUES ( \n" + keyToSqlByteString(itemKey) + ", \n" + "'2000' \n" + ")";
+
+ BeamSqlRelUtils.toPCollection(writePipeline, sqlEnv.parseQuery(insertStatement));
+ writePipeline.run().waitUntilFinish();
+
+ String selectTableStatement = "SELECT * FROM TEST";
+ PCollection<Row> output =
+ BeamSqlRelUtils.toPCollection(readPipeline, sqlEnv.parseQuery(selectTableStatement));
+
+ assertThat(output.getSchema(), equalTo(SOURCE_SCHEMA));
+
+ PipelineResult.State state = readPipeline.run().waitUntilFinish(Duration.standardMinutes(5));
+ assertThat(state, equalTo(State.DONE));
+ }
+
+ @Test
+ public void testDataStoreV1SqlWriteRead_withoutKey() {
+ BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new DataStoreV1TableProvider());
+ String projectId = options.getProject();
+
+ String createTableStatement =
+ "CREATE EXTERNAL TABLE TEST( \n"
+ + " `content` VARCHAR \n"
+ + ") \n"
+ + "TYPE 'datastoreV1' \n"
+ + "LOCATION '"
+ + projectId
+ + "/"
+ + KIND
+ + "'";
+ sqlEnv.executeDdl(createTableStatement);
+
+ String insertStatement = "INSERT INTO TEST VALUES ( '3000' )";
+
+ BeamSqlRelUtils.toPCollection(writePipeline, sqlEnv.parseQuery(insertStatement));
+ writePipeline.run().waitUntilFinish();
+
+ String selectTableStatement = "SELECT * FROM TEST";
+ PCollection<Row> output =
+ BeamSqlRelUtils.toPCollection(readPipeline, sqlEnv.parseQuery(selectTableStatement));
+
+ assertThat(output.getSchema(), equalTo(SOURCE_SCHEMA_WITHOUT_KEY));
+
+ PipelineResult.State state = readPipeline.run().waitUntilFinish(Duration.standardMinutes(5));
+ assertThat(state, equalTo(State.DONE));
+ }
+
+ @Test
+ public void testWriteRead_viaCoreBeamIO() {
+ String projectId = options.getProject();
+ Key ancestor = makeKey(KIND, UUID.randomUUID().toString()).build();
+ Key itemKey =
+ makeKey(ancestor, KIND, UUID.randomUUID().toString())
+ .setPartitionId(PartitionId.newBuilder().setProjectId(projectId).build())
+ .build();
+ Row testWriteRow =
+ Row.withSchema(SOURCE_SCHEMA).addValues(itemKey.toByteArray(), "4000").build();
+
+ writePipeline
+ .apply(Create.of(testWriteRow).withRowSchema(SOURCE_SCHEMA))
+ .apply(RowToEntity.create("__key__", KIND))
+ .apply(DatastoreIO.v1().write().withProjectId(projectId));
+ writePipeline.run().waitUntilFinish();
+
+ Query.Builder query = Query.newBuilder();
+ query.addKindBuilder().setName(KIND);
+ query.setFilter(makeFilter("__key__", Operator.EQUAL, makeValue(itemKey)));
+
+ DatastoreV1.Read read =
+ DatastoreIO.v1().read().withProjectId(projectId).withQuery(query.build());
+ PCollection<Row> rowsRead =
+ readPipeline.apply(read).apply(EntityToRow.create(SOURCE_SCHEMA, "__key__"));
+
+ PAssert.that(rowsRead).containsInAnyOrder(testWriteRow);
+ readPipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testReadAllSupportedTypes() {
+ BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new DataStoreV1TableProvider());
+ String projectId = options.getProject();
+
+ final Schema expectedSchema =
+ Schema.builder()
+ .addNullableField("__key__", VARBINARY)
+ .addNullableField("boolean", BOOLEAN)
+ .addNullableField("datetime", DATETIME)
+ // TODO: flattening of nested fields by Calcite causes some issues.
+ /*.addRowField("embeddedentity",
+ Schema.builder()
+ .addNullableField("property1", STRING)
+ .addNullableField("property2", INT64)
+ .build())*/
+ .addNullableField("floatingnumber", DOUBLE)
+ .addNullableField("integer", INT64)
+ .addNullableField("primitivearray", FieldType.array(STRING))
+ .addNullableField("string", STRING)
+ .addNullableField("text", STRING)
+ .build();
+
+ String createTableStatement =
+ "CREATE EXTERNAL TABLE TEST( \n"
+ + " `__key__` VARBINARY, \n"
+ + " `boolean` BOOLEAN, \n"
+ + " `datetime` TIMESTAMP, \n"
+ // + " `embeddedentity` ROW(`property1` VARCHAR, `property2` BIGINT), \n"
+ + " `floatingnumber` DOUBLE, \n"
+ + " `integer` BIGINT, \n"
+ + " `primitivearray` ARRAY<VARCHAR>, \n"
+ + " `string` VARCHAR, \n"
+ + " `text` VARCHAR"
+ + ") \n"
+ + "TYPE 'datastoreV1' \n"
+ + "LOCATION '"
+ + projectId
+ + "/"
+ + KIND_ALL_TYPES
+ + "'";
+ sqlEnv.executeDdl(createTableStatement);
+
+ String selectTableStatement = "SELECT * FROM TEST";
+ PCollection<Row> output =
+ BeamSqlRelUtils.toPCollection(readPipeline, sqlEnv.parseQuery(selectTableStatement));
+
+ assertThat(output.getSchema(), equalTo(expectedSchema));
+
+ PipelineResult.State state = readPipeline.run().waitUntilFinish(Duration.standardMinutes(5));
+ assertThat(state, equalTo(State.DONE));
+ }
+
+ private static String keyToSqlByteString(Key key) {
+ return "X'" + ByteString.toString(key.toByteArray(), 16) + "'";
+ }
+}
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreTableProviderTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreTableProviderTest.java
new file mode 100644
index 0000000..879add1
--- /dev/null
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreTableProviderTest.java
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.datastore;
+
+import static org.apache.beam.sdk.extensions.sql.meta.provider.datastore.DataStoreV1Table.DEFAULT_KEY_FIELD;
+import static org.apache.beam.sdk.extensions.sql.meta.provider.datastore.DataStoreV1Table.KEY_FIELD_PROPERTY;
+import static org.apache.beam.sdk.schemas.Schema.toSchema;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import com.alibaba.fastjson.JSON;
+import java.util.stream.Stream;
+import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable;
+import org.apache.beam.sdk.extensions.sql.meta.Table;
+import org.apache.beam.sdk.schemas.Schema;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class DataStoreTableProviderTest {
+ private DataStoreV1TableProvider provider = new DataStoreV1TableProvider();
+
+ @Test
+ public void testGetTableType() {
+ assertEquals("datastoreV1", provider.getTableType());
+ }
+
+ @Test
+ public void testBuildBeamSqlTable() {
+ final String location = "projectId/batch_kind";
+ Table table = fakeTable("TEST", location);
+ BeamSqlTable sqlTable = provider.buildBeamSqlTable(table);
+
+ assertNotNull(sqlTable);
+ assertTrue(sqlTable instanceof DataStoreV1Table);
+
+ DataStoreV1Table datastoreTable = (DataStoreV1Table) sqlTable;
+ assertEquals("projectId", datastoreTable.projectId);
+ assertEquals("batch_kind", datastoreTable.kind);
+ assertEquals(DEFAULT_KEY_FIELD, datastoreTable.keyField);
+ }
+
+ @Test
+ public void testTableProperty() {
+ final String location = "projectId/batch_kind";
+ Table table =
+ fakeTableWithProperties("TEST", location, "{ " + KEY_FIELD_PROPERTY + ": \"field_name\" }");
+ BeamSqlTable sqlTable = provider.buildBeamSqlTable(table);
+
+ assertNotNull(sqlTable);
+ assertTrue(sqlTable instanceof DataStoreV1Table);
+
+ DataStoreV1Table datastoreTable = (DataStoreV1Table) sqlTable;
+ assertEquals("projectId", datastoreTable.projectId);
+ assertEquals("batch_kind", datastoreTable.kind);
+ assertEquals("field_name", datastoreTable.keyField);
+ }
+
+ @Test
+ public void testTableProperty_nullValue_throwsException() {
+ final String location = "projectId/batch_kind";
+ Table table = fakeTableWithProperties("TEST", location, "{ " + KEY_FIELD_PROPERTY + ": \"\" }");
+ assertThrows(IllegalArgumentException.class, () -> provider.buildBeamSqlTable(table));
+ }
+
+ private static Table fakeTable(String name, String location) {
+ return Table.builder()
+ .name(name)
+ .comment(name + " table")
+ .location(location)
+ .schema(
+ Stream.of(
+ Schema.Field.nullable("id", Schema.FieldType.INT32),
+ Schema.Field.nullable("name", Schema.FieldType.STRING))
+ .collect(toSchema()))
+ .type("datastoreV1")
+ .build();
+ }
+
+ private static Table fakeTableWithProperties(String name, String location, String properties) {
+ return Table.builder()
+ .name(name)
+ .comment(name + " table")
+ .location(location)
+ .schema(
+ Stream.of(
+ Schema.Field.nullable("id", Schema.FieldType.INT32),
+ Schema.Field.nullable("name", Schema.FieldType.STRING))
+ .collect(toSchema()))
+ .type("datastoreV1")
+ .properties(JSON.parseObject(properties))
+ .build();
+ }
+}
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreTableTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreTableTest.java
new file mode 100644
index 0000000..a682b2a
--- /dev/null
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/datastore/DataStoreTableTest.java
@@ -0,0 +1,190 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.datastore;
+
+import static com.google.datastore.v1.client.DatastoreHelper.makeKey;
+import static com.google.datastore.v1.client.DatastoreHelper.makeValue;
+import static org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.VARBINARY;
+import static org.apache.beam.sdk.extensions.sql.meta.provider.datastore.DataStoreV1Table.DEFAULT_KEY_FIELD;
+import static org.apache.beam.sdk.extensions.sql.utils.DateTimeUtils.parseTimestampWithUTCTimeZone;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.BOOLEAN;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.BYTES;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.DATETIME;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.DOUBLE;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.INT64;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.STRING;
+import static org.apache.beam.sdk.schemas.Schema.FieldType.array;
+
+import com.google.datastore.v1.Entity;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.Value;
+import com.google.protobuf.ByteString;
+import java.nio.charset.Charset;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
+import org.apache.beam.sdk.extensions.sql.meta.provider.datastore.DataStoreV1Table.EntityToRow;
+import org.apache.beam.sdk.extensions.sql.meta.provider.datastore.DataStoreV1Table.RowToEntity;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.Schema.FieldType;
+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.Row;
+import org.joda.time.DateTime;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class DataStoreTableTest {
+ private static final String KIND = "kind";
+ private static final String UUID_VALUE = UUID.randomUUID().toString();
+ private static final Key.Builder KEY = makeKey(KIND, UUID_VALUE);
+ private static final DateTime DATE_TIME = parseTimestampWithUTCTimeZone("2018-05-28 20:17:40");
+
+ private static final Schema NESTED_ROW_SCHEMA =
+ Schema.builder().addNullableField("nestedLong", INT64).build();
+ private static final Schema SCHEMA =
+ Schema.builder()
+ .addNullableField("__key__", VARBINARY)
+ .addNullableField("long", INT64)
+ .addNullableField("bool", BOOLEAN)
+ .addNullableField("datetime", DATETIME)
+ .addNullableField("array", array(STRING))
+ .addNullableField("rowArray", array(FieldType.row(NESTED_ROW_SCHEMA)))
+ .addNullableField("double", DOUBLE)
+ .addNullableField("bytes", BYTES)
+ .addNullableField("string", CalciteUtils.CHAR)
+ .addNullableField("nullable", INT64)
+ .build();
+ private static final Entity NESTED_ENTITY =
+ Entity.newBuilder().putProperties("nestedLong", makeValue(Long.MIN_VALUE).build()).build();
+ private static final Entity ENTITY =
+ Entity.newBuilder()
+ .setKey(KEY)
+ .putProperties("long", makeValue(Long.MAX_VALUE).build())
+ .putProperties("bool", makeValue(true).build())
+ .putProperties("datetime", makeValue(DATE_TIME.toDate()).build())
+ .putProperties("array", makeValue(makeValue("string1"), makeValue("string2")).build())
+ .putProperties(
+ "rowArray",
+ makeValue(Collections.singletonList(makeValue(NESTED_ENTITY).build())).build())
+ .putProperties("double", makeValue(Double.MAX_VALUE).build())
+ .putProperties(
+ "bytes", makeValue(ByteString.copyFrom("hello", Charset.defaultCharset())).build())
+ .putProperties("string", makeValue("string").build())
+ .putProperties("nullable", Value.newBuilder().build())
+ .build();
+ private static final Row ROW =
+ row(
+ SCHEMA,
+ KEY.build().toByteArray(),
+ Long.MAX_VALUE,
+ true,
+ DATE_TIME,
+ Arrays.asList("string1", "string2"),
+ Collections.singletonList(row(NESTED_ROW_SCHEMA, Long.MIN_VALUE)),
+ Double.MAX_VALUE,
+ "hello".getBytes(Charset.defaultCharset()),
+ "string",
+ null);
+
+ @Rule public transient TestPipeline pipeline = TestPipeline.create();
+
+ @Test
+ public void testEntityToRowConverter() {
+ PCollection<Row> result =
+ pipeline.apply(Create.of(ENTITY)).apply(EntityToRow.create(SCHEMA, DEFAULT_KEY_FIELD));
+ PAssert.that(result).containsInAnyOrder(ROW);
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testEntityToRowConverterWithoutKey() {
+ Schema schemaWithoutKey =
+ Schema.builder()
+ .addFields(
+ SCHEMA.getFields().stream()
+ .filter(f -> !f.getName().equals("__key__"))
+ .collect(Collectors.toList()))
+ .build();
+ Row rowWithoutKey =
+ Row.withSchema(schemaWithoutKey)
+ .addValues(
+ schemaWithoutKey.getFieldNames().stream()
+ .map(ROW::getValue)
+ .collect(Collectors.toList()))
+ .build();
+ PCollection<Row> result =
+ pipeline
+ .apply(Create.of(ENTITY))
+ .apply(EntityToRow.create(schemaWithoutKey, DEFAULT_KEY_FIELD));
+ PAssert.that(result).containsInAnyOrder(rowWithoutKey);
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testRowToEntityConverter() {
+ PCollection<Entity> result =
+ pipeline
+ .apply(Create.of(ROW))
+ .setRowSchema(SCHEMA)
+ .apply(RowToEntity.createTest(UUID_VALUE, "__key__", KIND));
+ PAssert.that(result).containsInAnyOrder(ENTITY);
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testRowToEntityConverterWithoutKey() {
+ Schema schemaWithoutKey =
+ Schema.builder()
+ .addFields(
+ SCHEMA.getFields().stream()
+ .filter(f -> !f.getName().equals("__key__"))
+ .collect(Collectors.toList()))
+ .build();
+ Row rowWithoutKey =
+ Row.withSchema(schemaWithoutKey)
+ .addValues(
+ schemaWithoutKey.getFieldNames().stream()
+ .map(ROW::getValue)
+ .collect(Collectors.toList()))
+ .build();
+ PCollection<Entity> result =
+ pipeline
+ .apply(Create.of(rowWithoutKey))
+ .setRowSchema(schemaWithoutKey)
+ .apply(RowToEntity.createTest(UUID_VALUE, "__key__", KIND));
+
+ PAssert.that(result).containsInAnyOrder(ENTITY);
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ private static Row row(Schema schema, Object... values) {
+ return Row.withSchema(schema).addValues(values).build();
+ }
+}
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbFilterTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbFilterTest.java
new file mode 100644
index 0000000..aafaf9b
--- /dev/null
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbFilterTest.java
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.meta.provider.mongodb;
+
+import static org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider.PUSH_DOWN_OPTION;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.instanceOf;
+
+import com.alibaba.fastjson.JSON;
+import java.util.Arrays;
+import java.util.Collection;
+import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv;
+import org.apache.beam.sdk.extensions.sql.impl.rel.BeamCalcRel;
+import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
+import org.apache.beam.sdk.extensions.sql.meta.Table;
+import org.apache.beam.sdk.extensions.sql.meta.provider.mongodb.MongoDbTable.MongoDbFilter;
+import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider;
+import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider.PushDownOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.values.Row;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class MongoDbFilterTest {
+ private static final Schema BASIC_SCHEMA =
+ Schema.builder()
+ .addInt32Field("unused1")
+ .addInt32Field("id")
+ .addStringField("name")
+ .addInt16Field("unused2")
+ .addBooleanField("b")
+ .build();
+ private BeamSqlEnv sqlEnv;
+
+ @Parameters
+ public static Collection<Object[]> data() {
+ return Arrays.asList(
+ new Object[][] {
+ {"select * from TEST where unused1=100", true},
+ {"select * from TEST where unused1 in (100, 200)", true},
+ {"select * from TEST where b", true},
+ {"select * from TEST where not b", true},
+ {
+ "select * from TEST where unused1>100 and unused1<=200 and id<>1 and (name='two' or id=2)",
+ true
+ },
+ {"select * from TEST where name like 'o%e'", false},
+ {"select * from TEST where unused1+10=110", false},
+ {"select * from TEST where unused1=unused2 and id=2", false},
+ {"select * from TEST where unused1+unused2=10", false}
+ });
+ }
+
+ @Parameter public String query;
+
+ @Parameter(1)
+ public boolean isSupported;
+
+ @Rule public TestPipeline pipeline = TestPipeline.create();
+
+ @Before
+ public void buildUp() {
+ TestTableProvider tableProvider = new TestTableProvider();
+ Table table = getTable("TEST", PushDownOptions.NONE);
+ tableProvider.createTable(table);
+ tableProvider.addRows(
+ table.getName(),
+ row(BASIC_SCHEMA, 100, 1, "one", (short) 100, true),
+ row(BASIC_SCHEMA, 200, 2, "two", (short) 200, false));
+
+ sqlEnv =
+ BeamSqlEnv.builder(tableProvider)
+ .setPipelineOptions(PipelineOptionsFactory.create())
+ .build();
+ }
+
+ @Test
+ public void testIsSupported() {
+ BeamRelNode beamRelNode = sqlEnv.parseQuery(query);
+ assertThat(beamRelNode, instanceOf(BeamCalcRel.class));
+ MongoDbFilter filter =
+ MongoDbFilter.create(((BeamCalcRel) beamRelNode).getProgram().split().right);
+
+ assertThat(
+ "Query: '" + query + "' is expected to be " + (isSupported ? "supported." : "unsupported."),
+ filter.getNotSupported().isEmpty() == isSupported);
+ }
+
+ private static Table getTable(String name, PushDownOptions options) {
+ return Table.builder()
+ .name(name)
+ .comment(name + " table")
+ .schema(BASIC_SCHEMA)
+ .properties(
+ JSON.parseObject("{ " + PUSH_DOWN_OPTION + ": " + "\"" + options.toString() + "\" }"))
+ .type("test")
+ .build();
+ }
+
+ private static Row row(Schema schema, Object... objects) {
+ return Row.withSchema(schema).addValues(objects).build();
+ }
+}
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbReadWriteIT.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbReadWriteIT.java
index 0d4296a..47ad96a 100644
--- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbReadWriteIT.java
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/mongodb/MongoDbReadWriteIT.java
@@ -57,6 +57,7 @@
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.bson.Document;
import org.junit.After;
import org.junit.AfterClass;
@@ -145,7 +146,7 @@
public void init() {
sqlEnv = BeamSqlEnv.inMemory(new MongoDbTableProvider());
MongoDatabase db = client.getDatabase(database);
- Document r = db.runCommand(new BasicDBObject().append("profile", 2));
+ db.runCommand(new BasicDBObject().append("profile", 2));
}
@After
@@ -295,6 +296,107 @@
containsInAnyOrder("c_varchar", "c_boolean", "c_integer"));
}
+ @Test
+ public void testPredicatePushDown() {
+ final Document expectedFilter =
+ new Document()
+ .append(
+ "$or",
+ ImmutableList.of(
+ new Document("c_varchar", "varchar"),
+ new Document(
+ "c_varchar", new Document("$not", new Document("$eq", "fakeString"))))
+ .asList())
+ .append("c_boolean", true)
+ .append("c_integer", 2147483647);
+ final Schema expectedSchema =
+ Schema.builder()
+ .addNullableField("c_varchar", STRING)
+ .addNullableField("c_boolean", BOOLEAN)
+ .addNullableField("c_integer", INT32)
+ .build();
+ Row testRow = row(expectedSchema, "varchar", true, 2147483647);
+
+ String createTableStatement =
+ "CREATE EXTERNAL TABLE TEST( \n"
+ + " c_bigint BIGINT, \n "
+ + " c_tinyint TINYINT, \n"
+ + " c_smallint SMALLINT, \n"
+ + " c_integer INTEGER, \n"
+ + " c_float FLOAT, \n"
+ + " c_double DOUBLE, \n"
+ + " c_boolean BOOLEAN, \n"
+ + " c_varchar VARCHAR, \n "
+ + " c_arr ARRAY<VARCHAR> \n"
+ + ") \n"
+ + "TYPE 'mongodb' \n"
+ + "LOCATION '"
+ + mongoSqlUrl
+ + "'";
+ sqlEnv.executeDdl(createTableStatement);
+
+ String insertStatement =
+ "INSERT INTO TEST VALUES ("
+ + "9223372036854775807, "
+ + "127, "
+ + "32767, "
+ + "2147483647, "
+ + "1.0, "
+ + "1.0, "
+ + "TRUE, "
+ + "'varchar', "
+ + "ARRAY['123', '456']"
+ + ")";
+
+ BeamRelNode insertRelNode = sqlEnv.parseQuery(insertStatement);
+ BeamSqlRelUtils.toPCollection(writePipeline, insertRelNode);
+ writePipeline.run().waitUntilFinish();
+
+ BeamRelNode node =
+ sqlEnv.parseQuery(
+ "select c_varchar, c_boolean, c_integer from TEST"
+ + " where (c_varchar='varchar' or c_varchar<>'fakeString') and c_boolean and c_integer=2147483647");
+ // Calc should be dropped, since MongoDb can push-down all predicate operations from a query
+ // above.
+ assertThat(node, instanceOf(BeamPushDownIOSourceRel.class));
+ // Only selected fields are projected.
+ assertThat(
+ node.getRowType().getFieldNames(),
+ containsInAnyOrder("c_varchar", "c_boolean", "c_integer"));
+ PCollection<Row> output = BeamSqlRelUtils.toPCollection(readPipeline, node);
+
+ assertThat(output.getSchema(), equalTo(expectedSchema));
+ PAssert.that(output).containsInAnyOrder(testRow);
+
+ readPipeline.run().waitUntilFinish();
+
+ MongoDatabase db = client.getDatabase(database);
+ MongoCollection coll = db.getCollection("system.profile");
+ // Find the last executed query.
+ Object query =
+ coll.find()
+ .filter(Filters.eq("op", "query"))
+ .sort(new BasicDBObject().append("ts", -1))
+ .iterator()
+ .next();
+
+ // Retrieve a projection parameters.
+ assertThat(query, instanceOf(Document.class));
+ Object command = ((Document) query).get("command");
+ assertThat(command, instanceOf(Document.class));
+ Object filter = ((Document) command).get("filter");
+ assertThat(filter, instanceOf(Document.class));
+ Object projection = ((Document) command).get("projection");
+ assertThat(projection, instanceOf(Document.class));
+
+ // Validate projected fields.
+ assertThat(
+ ((Document) projection).keySet(),
+ containsInAnyOrder("c_varchar", "c_boolean", "c_integer"));
+ // Validate filtered fields.
+ assertThat(((Document) filter), equalTo(expectedFilter));
+ }
+
private Row row(Schema schema, Object... values) {
return Row.withSchema(schema).addValues(values).build();
}
diff --git a/sdks/java/extensions/sql/zetasql/build.gradle b/sdks/java/extensions/sql/zetasql/build.gradle
index 330209b..06b265e 100644
--- a/sdks/java/extensions/sql/zetasql/build.gradle
+++ b/sdks/java/extensions/sql/zetasql/build.gradle
@@ -25,7 +25,7 @@
description = "Apache Beam :: SDKs :: Java :: Extensions :: SQL :: ZetaSQL"
ext.summary = "ZetaSQL to Calcite translator"
-def zetasql_version = "2019.12.1"
+def zetasql_version = "2020.01.1"
dependencies {
compile project(":sdks:java:core")
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 330fb2d..153df25 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
@@ -34,7 +34,6 @@
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.Schema;
-import org.apache.beam.sdk.schemas.Schema.Field;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -67,13 +66,14 @@
private static final SqlDialect DIALECT = BeamBigQuerySqlDialect.DEFAULT;
private final SqlImplementor.Context context;
+ private static String columnName(int i) {
+ return "_" + i;
+ }
+
public BeamZetaSqlCalcRel(
RelOptCluster cluster, RelTraitSet traits, RelNode input, RexProgram program) {
super(cluster, traits, input, program);
- final IntFunction<SqlNode> fn =
- i ->
- new SqlIdentifier(
- getProgram().getInputRowType().getFieldList().get(i).getName(), SqlParserPos.ZERO);
+ final IntFunction<SqlNode> fn = i -> new SqlIdentifier(columnName(i), SqlParserPos.ZERO);
context = new BeamSqlUnparseContext(fn);
}
@@ -146,20 +146,21 @@
@Setup
public void setup() {
AnalyzerOptions options = SqlAnalyzer.initAnalyzerOptions();
- for (Field field : inputSchema.getFields()) {
+ for (int i = 0; i < inputSchema.getFieldCount(); i++) {
options.addExpressionColumn(
- sanitize(field.getName()), ZetaSqlUtils.beamFieldTypeToZetaSqlType(field.getType()));
+ columnName(i),
+ ZetaSqlUtils.beamFieldTypeToZetaSqlType(inputSchema.getField(i).getType()));
}
// TODO[BEAM-8630]: use a single PreparedExpression for all condition and projects
projectExps = new ArrayList<>();
for (String project : projects) {
- PreparedExpression projectExp = new PreparedExpression(sanitize(project));
+ PreparedExpression projectExp = new PreparedExpression(project);
projectExp.prepare(options);
projectExps.add(projectExp);
}
if (condition != null) {
- conditionExp = new PreparedExpression(sanitize(condition));
+ conditionExp = new PreparedExpression(condition);
conditionExp.prepare(options);
}
}
@@ -168,10 +169,11 @@
public void processElement(ProcessContext c) {
Map<String, Value> columns = new HashMap<>();
Row row = c.element();
- for (Field field : inputSchema.getFields()) {
+ for (int i = 0; i < inputSchema.getFieldCount(); i++) {
columns.put(
- sanitize(field.getName()),
- ZetaSqlUtils.javaObjectToZetaSqlValue(row.getValue(field.getName()), field.getType()));
+ columnName(i),
+ ZetaSqlUtils.javaObjectToZetaSqlValue(
+ row.getValue(i), inputSchema.getField(i).getType()));
}
// TODO[BEAM-8630]: support parameters in expression evaluation
@@ -201,12 +203,5 @@
conditionExp.close();
}
}
-
- // Replaces "$" with "_" because "$" is not allowed in a valid ZetaSQL identifier
- // (ZetaSQL identifier syntax: [A-Za-z_][A-Za-z_0-9]*)
- // TODO[BEAM-8630]: check if this is sufficient and correct, or even better fix this in Calcite
- private static String sanitize(String identifier) {
- return identifier.replaceAll("\\$", "_");
- }
}
}
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlAnalyzer.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlAnalyzer.java
index 5c65f4b..63da1fb 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlAnalyzer.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlAnalyzer.java
@@ -31,6 +31,7 @@
import com.google.zetasql.ZetaSQLFunctions.FunctionEnums.Mode;
import com.google.zetasql.ZetaSQLOptions.ErrorMessageMode;
import com.google.zetasql.ZetaSQLOptions.LanguageFeature;
+import com.google.zetasql.ZetaSQLOptions.ParameterMode;
import com.google.zetasql.ZetaSQLOptions.ProductMode;
import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedCreateFunctionStmt;
import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedStatement;
@@ -39,6 +40,8 @@
import java.util.List;
import java.util.Map;
import java.util.Optional;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters.Kind;
import org.apache.beam.sdk.extensions.sql.zetasql.TableResolution.SimpleTableWithPath;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.Context;
@@ -46,7 +49,6 @@
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.SchemaPlus;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
/** Adapter for {@link Analyzer} to simplify the API for parsing the query and resolving the AST. */
@@ -72,9 +74,9 @@
this.builder = builder;
}
- /** Static factory method to create the builder with query parameters. */
- static Builder withQueryParams(Map<String, Value> params) {
- return new Builder().withQueryParams(ImmutableMap.copyOf(params));
+ /** Static factory method to create the builder. */
+ static Builder getBuilder() {
+ return new Builder();
}
/**
@@ -116,11 +118,19 @@
return options;
}
- private static AnalyzerOptions initAnalyzerOptions(Map<String, Value> queryParams) {
+ private static AnalyzerOptions initAnalyzerOptions(QueryParameters queryParams) {
AnalyzerOptions options = initAnalyzerOptions();
- for (Map.Entry<String, Value> entry : queryParams.entrySet()) {
- options.addQueryParameter(entry.getKey(), entry.getValue().getType());
+ if (queryParams.getKind() == Kind.NAMED) {
+ options.setParameterMode(ParameterMode.PARAMETER_NAMED);
+ for (Map.Entry<String, Value> entry : ((Map<String, Value>) queryParams.named()).entrySet()) {
+ options.addQueryParameter(entry.getKey(), entry.getValue().getType());
+ }
+ } else if (queryParams.getKind() == Kind.POSITIONAL) {
+ options.setParameterMode(ParameterMode.PARAMETER_POSITIONAL);
+ for (Value param : (List<Value>) queryParams.positional()) {
+ options.addPositionalQueryParameter(param.getType());
+ }
}
return options;
@@ -233,7 +243,7 @@
/** Builder for SqlAnalyzer. */
static class Builder {
- private Map<String, Value> queryParams;
+ private QueryParameters queryParams;
private QueryTrait queryTrait;
private SchemaPlus topLevelSchema;
private JavaTypeFactory typeFactory;
@@ -241,8 +251,8 @@
private Builder() {}
/** Query parameters. */
- Builder withQueryParams(Map<String, Value> params) {
- this.queryParams = ImmutableMap.copyOf(params);
+ Builder withQueryParams(QueryParameters params) {
+ this.queryParams = params;
return this;
}
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlOperators.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlOperators.java
index 9f0e948..3e9459f 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlOperators.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SqlOperators.java
@@ -31,6 +31,7 @@
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.ScalarFunction;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlIdentifier;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlOperator;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlSyntax;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.FamilyOperandTypeChecker;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.InferTypes;
@@ -74,6 +75,11 @@
// SqlUserDefinedFunction will be able to pass through Calcite codegen and get proper function
// called.
public static SqlUserDefinedFunction createUdfOperator(String name, Method method) {
+ return createUdfOperator(name, method, SqlSyntax.FUNCTION);
+ }
+
+ public static SqlUserDefinedFunction createUdfOperator(
+ String name, Method method, final SqlSyntax syntax) {
Function function = ScalarFunctionImpl.create(method);
final RelDataTypeFactory typeFactory = createTypeFactory();
@@ -96,7 +102,12 @@
InferTypes.explicit(argTypes),
typeChecker,
paramTypes,
- function);
+ function) {
+ @Override
+ public SqlSyntax getSyntax() {
+ return syntax;
+ }
+ };
}
private static RelDataType createSqlType(SqlTypeName typeName, boolean withNullability) {
@@ -171,7 +182,8 @@
public static final SqlOperator ENDS_WITH =
createUdfOperator("ENDS_WITH", BeamBuiltinMethods.ENDS_WITH_METHOD);
- public static final SqlOperator LIKE = createUdfOperator("LIKE", BeamBuiltinMethods.LIKE_METHOD);
+ public static final SqlOperator LIKE =
+ createUdfOperator("LIKE", BeamBuiltinMethods.LIKE_METHOD, SqlSyntax.BINARY);
public static final SqlOperator VALIDATE_TIMESTAMP =
createUdfOperator(
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLPlannerImpl.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLPlannerImpl.java
index ae77cbd..dce0ee2 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLPlannerImpl.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLPlannerImpl.java
@@ -19,13 +19,12 @@
import com.google.zetasql.Analyzer;
import com.google.zetasql.LanguageOptions;
-import com.google.zetasql.Value;
import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedQueryStmt;
import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedStatement;
import java.io.Reader;
import java.util.List;
-import java.util.Map;
import java.util.logging.Logger;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
import org.apache.beam.sdk.extensions.sql.zetasql.translation.ConversionContext;
import org.apache.beam.sdk.extensions.sql.zetasql.translation.ExpressionConverter;
import org.apache.beam.sdk.extensions.sql.zetasql.translation.QueryStatementConverter;
@@ -130,7 +129,7 @@
String.format("%s.rel(SqlNode) is not implemented", this.getClass().getCanonicalName()));
}
- public RelRoot rel(String sql, Map<String, Value> params) {
+ public RelRoot rel(String sql, QueryParameters params) {
this.cluster = RelOptCluster.create(planner, new RexBuilder(typeFactory));
this.expressionConverter = new ExpressionConverter(cluster, params);
@@ -142,7 +141,8 @@
TableResolution.registerTables(this.defaultSchemaPlus, tables);
ResolvedStatement statement =
- SqlAnalyzer.withQueryParams(params)
+ SqlAnalyzer.getBuilder()
+ .withQueryParams(params)
.withQueryTrait(trait)
.withCalciteContext(config.getContext())
.withTopLevelSchema(defaultSchemaPlus)
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLQueryPlanner.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLQueryPlanner.java
index 0f5ec16..03b06a6 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLQueryPlanner.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLQueryPlanner.java
@@ -18,7 +18,6 @@
package org.apache.beam.sdk.extensions.sql.zetasql;
import com.google.zetasql.Value;
-import java.util.Collections;
import java.util.Map;
import org.apache.beam.sdk.extensions.sql.impl.CalciteQueryPlanner.NonCumulativeCostImpl;
import org.apache.beam.sdk.extensions.sql.impl.JdbcConnection;
@@ -30,7 +29,6 @@
import org.apache.beam.sdk.extensions.sql.impl.planner.RelMdNodeStats;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamLogicalConvention;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
-import org.apache.beam.sdk.extensions.sql.impl.rule.BeamCalcRule;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.config.CalciteConnectionConfig;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.jdbc.CalciteSchema;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.ConventionTraitDef;
@@ -42,6 +40,7 @@
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.metadata.ChainedRelMetadataProvider;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.metadata.JaninoRelMetadataProvider;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rules.JoinCommuteRule;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.SchemaPlus;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlOperatorTable;
@@ -65,26 +64,30 @@
}
public ZetaSQLQueryPlanner(JdbcConnection jdbcConnection, RuleSet[] ruleSets) {
- // TODO[BEAM-8630]: uncomment the next lines once we have fully migrated to BeamZetaSqlCalcRel
- // plannerImpl =
- // new ZetaSQLPlannerImpl(defaultConfig(jdbcConnection, replaceBeamCalcRule(ruleSets)));
- plannerImpl = new ZetaSQLPlannerImpl(defaultConfig(jdbcConnection, ruleSets));
+ plannerImpl =
+ new ZetaSQLPlannerImpl(defaultConfig(jdbcConnection, modifyRuleSetsForZetaSql(ruleSets)));
}
public static RuleSet[] getZetaSqlRuleSets() {
- // TODO[BEAM-8630]: uncomment the next line once we have fully migrated to BeamZetaSqlCalcRel
- // return replaceBeamCalcRule(BeamRuleSets.getRuleSets());
- return BeamRuleSets.getRuleSets();
+ return modifyRuleSetsForZetaSql(BeamRuleSets.getRuleSets());
}
- private static RuleSet[] replaceBeamCalcRule(RuleSet[] ruleSets) {
+ private static RuleSet[] modifyRuleSetsForZetaSql(RuleSet[] ruleSets) {
RuleSet[] ret = new RuleSet[ruleSets.length];
for (int i = 0; i < ruleSets.length; i++) {
ImmutableList.Builder<RelOptRule> bd = ImmutableList.builder();
for (RelOptRule rule : ruleSets[i]) {
- if (rule instanceof BeamCalcRule) {
- bd.add(BeamZetaSqlCalcRule.INSTANCE);
- } else {
+ // TODO[BEAM-9075]: Fix join re-ordering for ZetaSQL planner. Currently join re-ordering
+ // requires the JoinCommuteRule, which doesn't work without struct flattening.
+ if (rule instanceof JoinCommuteRule) {
+ continue;
+ }
+ // TODO[BEAM-8630]: uncomment the next block once we have fully migrated to
+ // BeamZetaSqlCalcRel
+ // else if (rule instanceof BeamCalcRule) {
+ // bd.add(BeamZetaSqlCalcRule.INSTANCE);
+ // }
+ else {
bd.add(rule);
}
}
@@ -93,11 +96,15 @@
return ret;
}
+ public BeamRelNode convertToBeamRel(String sqlStatement) {
+ return convertToBeamRel(sqlStatement, QueryParameters.ofNone());
+ }
+
@Override
- public BeamRelNode convertToBeamRel(String sqlStatement)
+ public BeamRelNode convertToBeamRel(String sqlStatement, QueryParameters queryParameters)
throws ParseException, SqlConversionException {
try {
- return parseQuery(sqlStatement);
+ return parseQuery(sqlStatement, queryParameters);
} catch (RelConversionException e) {
throw new SqlConversionException(e.getCause());
}
@@ -114,17 +121,17 @@
public BeamRelNode convertToBeamRel(String sqlStatement, Map<String, Value> queryParams)
throws ParseException, SqlConversionException {
try {
- return parseQuery(sqlStatement, queryParams);
+ return parseQuery(sqlStatement, QueryParameters.ofNamed(queryParams));
} catch (RelConversionException e) {
throw new SqlConversionException(e.getCause());
}
}
public BeamRelNode parseQuery(String sql) throws RelConversionException {
- return parseQuery(sql, Collections.emptyMap());
+ return parseQuery(sql, QueryParameters.ofNone());
}
- public BeamRelNode parseQuery(String sql, Map<String, Value> queryParams)
+ public BeamRelNode parseQuery(String sql, QueryParameters queryParams)
throws RelConversionException {
RelRoot root = plannerImpl.rel(sql, queryParams);
RelTraitSet desiredTraits =
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java
index e55481e..e1079b3 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java
@@ -31,6 +31,7 @@
import static org.apache.beam.sdk.extensions.sql.zetasql.ZetaSQLCastFunctionImpl.ZETASQL_CAST_OP;
import com.google.common.base.Ascii;
+import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
@@ -62,6 +63,8 @@
import java.util.Optional;
import java.util.stream.Collectors;
import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters.Kind;
import org.apache.beam.sdk.extensions.sql.zetasql.SqlOperatorRewriter;
import org.apache.beam.sdk.extensions.sql.zetasql.SqlOperators;
import org.apache.beam.sdk.extensions.sql.zetasql.SqlStdOperatorMappingTable;
@@ -173,9 +176,9 @@
+ INTERVAL_DATE_PART_MSG;
private final RelOptCluster cluster;
- private final Map<String, Value> queryParams;
+ private final QueryParameters queryParams;
- public ExpressionConverter(RelOptCluster cluster, Map<String, Value> params) {
+ public ExpressionConverter(RelOptCluster cluster, QueryParameters params) {
this.cluster = cluster;
this.queryParams = params;
}
@@ -1000,9 +1003,27 @@
}
private RexNode convertResolvedParameter(ResolvedParameter parameter) {
- assert parameter.getType().equals(queryParams.get(parameter.getName()).getType());
- return convertValueToRexNode(
- queryParams.get(parameter.getName()).getType(), queryParams.get(parameter.getName()));
+ if (queryParams.getKind() == Kind.NAMED) {
+ Map<String, Value> queryParameterMap = (Map<String, Value>) queryParams.named();
+ Value value = queryParameterMap.get(parameter.getName());
+ Preconditions.checkState(
+ parameter.getType().equals(value.getType()),
+ String.format(
+ "Expected resolved parameter %s to have type %s, but it has type %s",
+ parameter.getName(), value.getType(), parameter.getType()));
+ return convertValueToRexNode(value.getType(), value);
+ } else if (queryParams.getKind() == Kind.POSITIONAL) {
+ List<Value> queryParameterList = (List<Value>) queryParams.positional();
+ // parameter is 1-indexed, while parameter list is 0-indexed.
+ Value value = queryParameterList.get((int) parameter.getPosition() - 1);
+ Preconditions.checkState(
+ parameter.getType().equals(value.getType()),
+ String.format(
+ "Expected resolved parameter %d to have type %s, but it has type %s",
+ parameter.getPosition(), value.getType(), parameter.getType()));
+ return convertValueToRexNode(value.getType(), value);
+ }
+ throw new IllegalArgumentException("Found unexpected parameter " + parameter);
}
private RexNode convertResolvedStructFieldAccess(ResolvedGetStructField resolvedGetStructField) {
diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java
index e140c02..0fa725c 100644
--- a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java
+++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java
@@ -58,6 +58,7 @@
import java.util.Map;
import org.apache.beam.sdk.extensions.sql.impl.JdbcConnection;
import org.apache.beam.sdk.extensions.sql.impl.JdbcDriver;
+import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils;
@@ -141,6 +142,23 @@
}
@Test
+ public void testByteLiterals() {
+ String sql = "SELECT b'abc'";
+
+ byte[] byteString = new byte[] {'a', 'b', 'c'};
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ final Schema schema = Schema.builder().addNullableField("ColA", FieldType.BYTES).build();
+
+ PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues(byteString).build());
+
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testByteString() {
String sql = "SELECT @p0 IS NULL AS ColA";
@@ -177,7 +195,7 @@
@Test
public void testStringLiterals() {
- String sql = "SELECT 'abc\\n'";
+ String sql = "SELECT '\"America/Los_Angeles\"\\n'";
ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
@@ -185,6 +203,24 @@
final Schema schema = Schema.builder().addNullableField("ColA", FieldType.STRING).build();
+ PAssert.that(stream)
+ .containsInAnyOrder(Row.withSchema(schema).addValues("\"America/Los_Angeles\"\n").build());
+
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testParameterString() {
+ String sql = "SELECT ?";
+ ImmutableList<Value> params = ImmutableList.of(Value.createStringValue("abc\n"));
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode =
+ zetaSQLQueryPlanner.convertToBeamRel(sql, QueryParameters.ofPositional(params));
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ final Schema schema = Schema.builder().addNullableField("ColA", FieldType.STRING).build();
+
PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues("abc\n").build());
pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
@@ -289,6 +325,22 @@
}
@Test
+ public void testEQ6() {
+ String sql = "SELECT ? = ? AS ColA";
+ ImmutableList<Value> params =
+ ImmutableList.of(Value.createInt64Value(4L), Value.createInt64Value(5L));
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode =
+ zetaSQLQueryPlanner.convertToBeamRel(sql, QueryParameters.ofPositional(params));
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ final Schema schema = Schema.builder().addNullableField("field1", FieldType.BOOLEAN).build();
+
+ PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues(false).build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testIsNotNull1() {
String sql = "SELECT @p0 IS NOT NULL AS ColA";
ImmutableMap<String, Value> params =
@@ -369,6 +421,25 @@
}
@Test
+ public void testIfPositional() {
+ String sql = "SELECT IF(?, ?, ?) AS ColA";
+
+ ImmutableList<Value> params =
+ ImmutableList.of(
+ Value.createBoolValue(true), Value.createInt64Value(1), Value.createInt64Value(2));
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode =
+ zetaSQLQueryPlanner.convertToBeamRel(sql, QueryParameters.ofPositional(params));
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ final Schema schema = Schema.builder().addNullableField("field1", FieldType.INT64).build();
+
+ PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues(1L).build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testCoalesceBasic() {
String sql = "SELECT COALESCE(@p0, @p1, @p2) AS ColA";
ImmutableMap<String, Value> params =
@@ -1370,6 +1441,24 @@
}
@Test
+ public void testZetaSQLStructFieldAccessInnerJoin() {
+ String sql =
+ "SELECT A.rowCol.data FROM table_with_struct_two AS A INNER JOIN "
+ + "table_with_struct AS B "
+ + "ON A.rowCol.row_id = B.id";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+ final Schema schema = Schema.builder().addStringField("field1").build();
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(schema).addValue("data1").build(),
+ Row.withSchema(schema).addValue("data2").build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testZetaSQLSelectFromTableWithArrayType() {
String sql = "SELECT array_col FROM table_with_array;";
@@ -2893,7 +2982,7 @@
}
@Test
- public void testConcatParameterQuery() {
+ public void testConcatNamedParameterQuery() {
String sql = "SELECT CONCAT(@p0, @p1) AS ColA";
ImmutableMap<String, Value> params =
ImmutableMap.of("p0", Value.createStringValue(""), "p1", Value.createStringValue("A"));
@@ -2907,6 +2996,24 @@
}
@Test
+ public void testConcatPositionalParameterQuery() {
+ String sql = "SELECT CONCAT(?, ?, ?) AS ColA";
+ ImmutableList<Value> params =
+ ImmutableList.of(
+ Value.createStringValue("a"),
+ Value.createStringValue("b"),
+ Value.createStringValue("c"));
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode =
+ zetaSQLQueryPlanner.convertToBeamRel(sql, QueryParameters.ofPositional(params));
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+ final Schema schema = Schema.builder().addStringField("field1").build();
+ PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues("abc").build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testReplace1() {
String sql = "SELECT REPLACE(@p0, @p1, @p2) AS ColA";
ImmutableMap<String, Value> params =
@@ -3644,7 +3751,6 @@
/** Only sample scenarios are covered here. Excessive testing is done via Compliance tests. */
@Test
- @Ignore("ZetaSQL does not support EnumType to IdentifierLiteral")
public void testExtractTimestamp() {
String sql =
"WITH Timestamps AS (\n"
diff --git a/sdks/java/extensions/zetasketch/src/test/java/org/apache/beam/sdk/extensions/zetasketch/BigQueryHllSketchCompatibilityIT.java b/sdks/java/extensions/zetasketch/src/test/java/org/apache/beam/sdk/extensions/zetasketch/BigQueryHllSketchCompatibilityIT.java
index 462a715..ef66b5f 100644
--- a/sdks/java/extensions/zetasketch/src/test/java/org/apache/beam/sdk/extensions/zetasketch/BigQueryHllSketchCompatibilityIT.java
+++ b/sdks/java/extensions/zetasketch/src/test/java/org/apache/beam/sdk/extensions/zetasketch/BigQueryHllSketchCompatibilityIT.java
@@ -17,6 +17,10 @@
*/
package org.apache.beam.sdk.extensions.zetasketch;
+import static org.apache.beam.sdk.io.gcp.testing.BigqueryMatcher.createQueryUsingStandardSql;
+import static org.apache.beam.sdk.io.gcp.testing.BigqueryMatcher.queryResultHasChecksum;
+import static org.hamcrest.MatcherAssert.assertThat;
+
import com.google.api.services.bigquery.model.Table;
import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableReference;
@@ -36,7 +40,6 @@
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead.Method;
import org.apache.beam.sdk.io.gcp.bigquery.SchemaAndRecord;
import org.apache.beam.sdk.io.gcp.testing.BigqueryClient;
-import org.apache.beam.sdk.io.gcp.testing.BigqueryMatcher;
import org.apache.beam.sdk.options.ApplicationNameOptions;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -234,11 +237,6 @@
TestPipelineOptions options =
TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class);
- // After the pipeline finishes, BigqueryMatcher will send a query to retrieve the estimated
- // count and verifies its correctness using checksum.
- options.setOnSuccessMatcher(
- BigqueryMatcher.createUsingStandardSql(APP_NAME, PROJECT_ID, query, expectedChecksum));
-
Pipeline p = Pipeline.create(options);
p.apply(Create.of(testData).withType(TypeDescriptor.of(String.class)))
.apply(HllCount.Init.forStrings().globally())
@@ -253,5 +251,11 @@
new TableRow().set(SKETCH_FIELD_NAME, sketch.length == 0 ? null : sketch))
.withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE));
p.run().waitUntilFinish();
+
+ // BigqueryMatcher will send a query to retrieve the estimated count and verifies its
+ // correctness using checksum.
+ assertThat(
+ createQueryUsingStandardSql(APP_NAME, PROJECT_ID, query),
+ queryResultHasChecksum(expectedChecksum));
}
}
diff --git a/sdks/java/fn-execution/build.gradle b/sdks/java/fn-execution/build.gradle
index ea46cff..d099333 100644
--- a/sdks/java/fn-execution/build.gradle
+++ b/sdks/java/fn-execution/build.gradle
@@ -27,7 +27,7 @@
compile project(path: ":model:pipeline", configuration: "shadow")
compile project(path: ":model:fn-execution", configuration: "shadow")
compile project(path: ":sdks:java:core", configuration: "shadow")
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
compile library.java.vendored_guava_26_0_jre
compile library.java.slf4j_api
compile library.java.joda_time
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactory.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactory.java
index bc6da0e..c6180a3 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactory.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactory.java
@@ -20,14 +20,14 @@
import java.net.SocketAddress;
import java.util.List;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ClientInterceptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.netty.NettyChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.EpollDomainSocketChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.EpollEventLoopGroup;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.EpollSocketChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.unix.DomainSocketAddress;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ClientInterceptor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.netty.NettyChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.EpollDomainSocketChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.EpollEventLoopGroup;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.EpollSocketChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.unix.DomainSocketAddress;
/** A Factory which creates an underlying {@link ManagedChannel} implementation. */
public abstract class ManagedChannelFactory {
@@ -36,7 +36,7 @@
}
public static ManagedChannelFactory createEpoll() {
- org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.Epoll.ensureAvailability();
+ org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.Epoll.ensureAvailability();
return new Epoll();
}
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/SocketAddressFactory.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/SocketAddressFactory.java
index c77e1bc..b7d9d76 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/SocketAddressFactory.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/channel/SocketAddressFactory.java
@@ -23,7 +23,7 @@
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.unix.DomainSocketAddress;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.unix.DomainSocketAddress;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.net.HostAndPort;
/** Creates a {@link SocketAddress} based upon a supplied string. */
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataBufferingOutboundObserver.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataBufferingOutboundObserver.java
index bbc2916..e741c7e 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataBufferingOutboundObserver.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataBufferingOutboundObserver.java
@@ -23,7 +23,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.options.ExperimentalOptions;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
/**
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java
index 7ed83df..3140616 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java
@@ -28,8 +28,8 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements.Data;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
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.collect.ImmutableList;
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserver.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserver.java
index c0215ae..51d9595 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserver.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserver.java
@@ -20,8 +20,8 @@
import java.io.IOException;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataTimeBasedBufferingOutboundObserver.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataTimeBasedBufferingOutboundObserver.java
index 3595fbd..b339e33 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataTimeBasedBufferingOutboundObserver.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataTimeBasedBufferingOutboundObserver.java
@@ -25,7 +25,7 @@
import java.util.concurrent.TimeUnit;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
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.util.concurrent.ThreadFactoryBuilder;
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortRead.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortRead.java
index 9568b90..f06599e 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortRead.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortRead.java
@@ -23,7 +23,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
/**
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWrite.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWrite.java
index b1c7604..42fd798 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWrite.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWrite.java
@@ -24,7 +24,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
/**
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java
index da7505d..12f042d 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/BufferingStreamObserver.java
@@ -27,8 +27,8 @@
import java.util.concurrent.Phaser;
import java.util.concurrent.TimeUnit;
import javax.annotation.concurrent.ThreadSafe;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
/**
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java
index 3134ea4..140f508 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java
@@ -27,7 +27,7 @@
import java.util.NoSuchElementException;
import java.util.concurrent.BlockingQueue;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.ByteStreams;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.CountingInputStream;
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java
index dce3452..3906318 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DirectStreamObserver.java
@@ -21,8 +21,8 @@
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import javax.annotation.concurrent.ThreadSafe;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java
index a25985d..016cb11 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java
@@ -17,9 +17,9 @@
*/
package org.apache.beam.sdk.fn.stream;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientCallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientResponseObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientCallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientResponseObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* A {@link ClientResponseObserver} which delegates all {@link StreamObserver} calls.
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java
index 83f94f9..6693fee 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactory.java
@@ -18,8 +18,8 @@
package org.apache.beam.sdk.fn.stream;
import java.util.concurrent.ExecutorService;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* Creates factories which determine an underlying {@link StreamObserver} implementation to use in
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/SynchronizedStreamObserver.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/SynchronizedStreamObserver.java
index c960d96..31e9af2 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/SynchronizedStreamObserver.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/SynchronizedStreamObserver.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.sdk.fn.stream;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* A {@link StreamObserver} which provides synchronous access access to an underlying {@link
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/InProcessManagedChannelFactory.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/InProcessManagedChannelFactory.java
index a4c99a1..aad1fd1 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/InProcessManagedChannelFactory.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/InProcessManagedChannelFactory.java
@@ -19,8 +19,8 @@
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.sdk.fn.channel.ManagedChannelFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
/**
* A {@link ManagedChannelFactory} that uses in-process channels.
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/TestStreams.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/TestStreams.java
index b76997e..cd8b977 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/TestStreams.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/test/TestStreams.java
@@ -19,8 +19,8 @@
import java.util.function.Consumer;
import java.util.function.Supplier;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/** Utility methods which enable testing of {@link StreamObserver}s. */
public class TestStreams {
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindow.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindow.java
index 94d3400..3e9b21c 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindow.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindow.java
@@ -25,7 +25,7 @@
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.VarInt;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.ByteStreams;
import org.joda.time.Instant;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactoryTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactoryTest.java
index 3e60697..fc0c813 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactoryTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/ManagedChannelFactoryTest.java
@@ -21,7 +21,7 @@
import static org.junit.Assume.assumeTrue;
import org.apache.beam.model.pipeline.v1.Endpoints;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
@@ -45,7 +45,7 @@
@Test
public void testEpollHostPortChannel() {
- assumeTrue(org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.Epoll.isAvailable());
+ assumeTrue(org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.Epoll.isAvailable());
Endpoints.ApiServiceDescriptor apiServiceDescriptor =
Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:123").build();
ManagedChannel channel =
@@ -56,7 +56,7 @@
@Test
public void testEpollDomainSocketChannel() throws Exception {
- assumeTrue(org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.epoll.Epoll.isAvailable());
+ assumeTrue(org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.epoll.Epoll.isAvailable());
Endpoints.ApiServiceDescriptor apiServiceDescriptor =
Endpoints.ApiServiceDescriptor.newBuilder()
.setUrl("unix://" + tmpFolder.newFile().getAbsolutePath())
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/SocketAddressFactoryTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/SocketAddressFactoryTest.java
index 0107a7b..91c1e17 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/SocketAddressFactoryTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/channel/SocketAddressFactoryTest.java
@@ -23,7 +23,7 @@
import java.io.File;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
-import org.apache.beam.vendor.grpc.v1p21p0.io.netty.channel.unix.DomainSocketAddress;
+import org.apache.beam.vendor.grpc.v1p26p0.io.netty.channel.unix.DomainSocketAddress;
import org.hamcrest.Matchers;
import org.junit.Rule;
import org.junit.Test;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java
index bf1b1d3..51301e4 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java
@@ -31,7 +31,7 @@
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.fn.test.TestStreams;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
import org.junit.Test;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserverTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserverTest.java
index ed2f700..a75b9fa 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserverTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataSizeBasedBufferingOutboundObserverTest.java
@@ -40,7 +40,7 @@
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortReadTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortReadTest.java
index c1b2175..97aebaf 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortReadTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortReadTest.java
@@ -24,7 +24,7 @@
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.model.pipeline.v1.Endpoints.OAuth2ClientCredentialsGrant;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWriteTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWriteTest.java
index c4be16b..1775728 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWriteTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/data/RemoteGrpcPortWriteTest.java
@@ -24,7 +24,7 @@
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.model.pipeline.v1.Endpoints.OAuth2ClientCredentialsGrant;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java
index 3e66d50..8dd5819 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java
@@ -41,7 +41,7 @@
import org.apache.beam.sdk.fn.stream.DataStreams.DataStreamDecoder;
import org.apache.beam.sdk.fn.stream.DataStreams.ElementDelimitedOutputStream;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.ByteStreams;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.CountingOutputStream;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java
index 97fc2da..5e0e4b5 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java
@@ -21,9 +21,9 @@
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientCallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientResponseObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientCallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientResponseObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java
index 60cd8b0..de56d01 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java
@@ -22,8 +22,8 @@
import static org.junit.Assert.assertThat;
import java.util.concurrent.Executors;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindowTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindowTest.java
index 18d6896..ecf5bc1 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindowTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/windowing/EncodedBoundedWindowTest.java
@@ -19,7 +19,7 @@
import org.apache.beam.sdk.fn.windowing.EncodedBoundedWindow.Coder;
import org.apache.beam.sdk.testing.CoderProperties;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/sdks/java/harness/build.gradle b/sdks/java/harness/build.gradle
index 51c65bb..bfc9112 100644
--- a/sdks/java/harness/build.gradle
+++ b/sdks/java/harness/build.gradle
@@ -62,12 +62,13 @@
shadowTest library.java.powermock_mockito
compile library.java.joda_time
compile library.java.slf4j_api
- compile library.java.vendored_grpc_1_21_0
+ compile library.java.vendored_grpc_1_26_0
provided library.java.error_prone_annotations
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
testCompile library.java.junit
testCompile library.java.mockito_core
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
+ testCompile project(":runners:core-construction-java")
shadowTestRuntimeClasspath library.java.slf4j_jdk14
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
index e93b2ba..65f11d9 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BeamFnDataReadRunner.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.fn.harness;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables.getOnlyElement;
import com.google.auto.service.AutoService;
@@ -24,12 +25,16 @@
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Supplier;
+import org.apache.beam.fn.harness.HandlesSplits.SplitResult;
import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest.DesiredSplit;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -47,6 +52,7 @@
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.Ints;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -131,6 +137,11 @@
private final BeamFnDataClient beamFnDataClient;
private final Coder<WindowedValue<OutputT>> coder;
+ private final Object splittingLock = new Object();
+ // 0-based index of the current element being processed
+ private long index = -1;
+ // 0-based index of the first element to not process, aka the first element of the residual
+ private long stopIndex = Long.MAX_VALUE;
private InboundDataClient readFuture;
BeamFnDataReadRunner(
@@ -170,7 +181,109 @@
apiServiceDescriptor,
LogicalEndpoint.of(processBundleInstructionIdSupplier.get(), pTransformId),
coder,
- consumer);
+ this::forwardElementToConsumer);
+ }
+
+ public void forwardElementToConsumer(WindowedValue<OutputT> element) throws Exception {
+ synchronized (splittingLock) {
+ if (index == stopIndex - 1) {
+ return;
+ }
+ index += 1;
+ }
+ consumer.accept(element);
+ }
+
+ public void split(
+ ProcessBundleSplitRequest request, ProcessBundleSplitResponse.Builder response) {
+ DesiredSplit desiredSplit = request.getDesiredSplitsMap().get(pTransformId);
+ if (desiredSplit == null) {
+ return;
+ }
+
+ long totalBufferSize = desiredSplit.getEstimatedInputElements();
+
+ HandlesSplits splittingConsumer = null;
+ if (consumer instanceof HandlesSplits) {
+ splittingConsumer = ((HandlesSplits) consumer);
+ }
+
+ synchronized (splittingLock) {
+ // Since we hold the splittingLock, we guarantee that we will not pass the next element
+ // to the downstream consumer. We still have a race where the downstream consumer may
+ // have yet to see the element or has completed processing the element by the time
+ // we ask it to split (even after we have asked for its progress).
+
+ // If the split request we received was delayed and is less then the known number of elements
+ // then use "index + 1" as the total size. Similarly, if we have already split and the
+ // split request is bounded incorrectly, use the stop index as the upper bound.
+ if (totalBufferSize < index + 1) {
+ totalBufferSize = index + 1;
+ } else if (totalBufferSize > stopIndex) {
+ totalBufferSize = stopIndex;
+ }
+
+ // In the case where we have yet to process an element, set the current element progress to 1.
+ double currentElementProgress = 1;
+
+ // If we have started processing at least one element, attempt to get the downstream
+ // progress defaulting to 0.5 if no progress was able to get fetched.
+ if (index >= 0) {
+ if (splittingConsumer != null) {
+ currentElementProgress = splittingConsumer.getProgress();
+ } else {
+ currentElementProgress = 0.5;
+ }
+ }
+
+ checkArgument(
+ desiredSplit.getAllowedSplitPointsList().isEmpty(),
+ "TODO: BEAM-3836, support split point restrictions.");
+
+ // Now figure out where to split.
+ //
+ // The units here (except for keepOfElementRemainder) are all in terms of number or
+ // (possibly fractional) elements.
+
+ // Compute the amount of "remaining" work that we know of.
+ double remainder = totalBufferSize - index - currentElementProgress;
+ // Compute the number of elements (including fractional elements) that we should "keep".
+ double keep = remainder * desiredSplit.getFractionOfRemainder();
+
+ // If the downstream operator says the progress is less than 1 then the element could be
+ // splittable.
+ if (currentElementProgress < 1) {
+ // See if the amount we need to keep falls within the current element's remainder and if
+ // so, attempt to split it.
+ double keepOfElementRemainder = keep / (1 - currentElementProgress);
+ if (keepOfElementRemainder < 1) {
+ SplitResult splitResult =
+ splittingConsumer != null ? splittingConsumer.trySplit(keepOfElementRemainder) : null;
+ if (splitResult != null) {
+ stopIndex = index + 1;
+ response
+ .addPrimaryRoots(splitResult.getPrimaryRoot())
+ .addResidualRoots(splitResult.getResidualRoot())
+ .addChannelSplitsBuilder()
+ .setLastPrimaryElement(index - 1)
+ .setFirstResidualElement(stopIndex);
+ return;
+ }
+ }
+ }
+
+ // Otherwise, split at the closest element boundary.
+ int newStopIndex =
+ Ints.checkedCast(index + Math.max(1, Math.round(currentElementProgress + keep)));
+ if (newStopIndex < stopIndex) {
+ stopIndex = newStopIndex;
+ response
+ .addChannelSplitsBuilder()
+ .setLastPrimaryElement(stopIndex - 1)
+ .setFirstResidualElement(stopIndex);
+ return;
+ }
+ }
}
public void blockTillReadFinishes() throws Exception {
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java
index a632aa2..a06c002 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/BoundedSourceRunner.java
@@ -43,7 +43,7 @@
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
index 5240e0c..6bd16be 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/CombineRunners.java
@@ -31,8 +31,7 @@
import org.apache.beam.model.pipeline.v1.RunnerApi.CombinePayload;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms;
-import org.apache.beam.runners.core.construction.BeamUrns;
+import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.RehydratedComponents;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
@@ -59,13 +58,13 @@
@Override
public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
return ImmutableMap.of(
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_PRECOMBINE),
+ PTransformTranslation.COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN,
new PrecombineFactory(),
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_MERGE_ACCUMULATORS),
+ PTransformTranslation.COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN,
MapFnRunners.forValueMapFnFactory(CombineRunners::createMergeAccumulatorsMapFunction),
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_PER_KEY_EXTRACT_OUTPUTS),
+ PTransformTranslation.COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN,
MapFnRunners.forValueMapFnFactory(CombineRunners::createExtractOutputsMapFunction),
- BeamUrns.getUrn(StandardPTransforms.CombineComponents.COMBINE_GROUPED_VALUES),
+ PTransformTranslation.COMBINE_GROUPED_VALUES_TRANSFORM_URN,
MapFnRunners.forValueMapFnFactory(CombineRunners::createCombineGroupedValuesMapFunction));
}
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/DoFnPTransformRunnerFactory.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/DoFnPTransformRunnerFactory.java
deleted file mode 100644
index 433572b..0000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/DoFnPTransformRunnerFactory.java
+++ /dev/null
@@ -1,300 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT 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.fn.harness;
-
-import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
-
-import java.io.IOException;
-import java.util.Map;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-import org.apache.beam.fn.harness.control.BundleSplitListener;
-import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
-import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
-import org.apache.beam.fn.harness.state.BeamFnStateClient;
-import org.apache.beam.fn.harness.state.SideInputSpec;
-import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
-import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
-import org.apache.beam.runners.core.construction.PCollectionViewTranslation;
-import org.apache.beam.runners.core.construction.ParDoTranslation;
-import org.apache.beam.runners.core.construction.RehydratedComponents;
-import org.apache.beam.runners.core.construction.Timer;
-import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.sdk.function.ThrowingRunnable;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.schemas.SchemaCoder;
-import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.Materializations;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
-import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.WindowingStrategy;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableListMultimap;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ListMultimap;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
-
-/** A {@link PTransformRunnerFactory} for transforms invoking a {@link DoFn}. */
-abstract class DoFnPTransformRunnerFactory<
- TransformInputT,
- FnInputT,
- OutputT,
- RunnerT extends DoFnPTransformRunnerFactory.DoFnPTransformRunner<TransformInputT>>
- implements PTransformRunnerFactory<RunnerT> {
- interface DoFnPTransformRunner<T> {
- void startBundle() throws Exception;
-
- void processElement(WindowedValue<T> input) throws Exception;
-
- void processTimer(
- String timerId, TimeDomain timeDomain, WindowedValue<KV<Object, Timer>> input);
-
- void finishBundle() throws Exception;
-
- void tearDown() throws Exception;
- }
-
- @Override
- public final RunnerT createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- BeamFnStateClient beamFnStateClient,
- String pTransformId,
- PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- PCollectionConsumerRegistry pCollectionConsumerRegistry,
- PTransformFunctionRegistry startFunctionRegistry,
- PTransformFunctionRegistry finishFunctionRegistry,
- Consumer<ThrowingRunnable> tearDownFunctions,
- BundleSplitListener splitListener) {
- Context<FnInputT, OutputT> context =
- new Context<>(
- pipelineOptions,
- beamFnStateClient,
- pTransformId,
- pTransform,
- processBundleInstructionId,
- pCollections,
- coders,
- windowingStrategies,
- pCollectionConsumerRegistry,
- splitListener);
-
- RunnerT runner = createRunner(context);
-
- // Register the appropriate handlers.
- startFunctionRegistry.register(pTransformId, runner::startBundle);
- Iterable<String> mainInput =
- Sets.difference(
- pTransform.getInputsMap().keySet(),
- Sets.union(
- context.parDoPayload.getSideInputsMap().keySet(),
- context.parDoPayload.getTimerSpecsMap().keySet()));
- for (String localInputName : mainInput) {
- pCollectionConsumerRegistry.register(
- pTransform.getInputsOrThrow(localInputName),
- pTransformId,
- (FnDataReceiver) (FnDataReceiver<WindowedValue<TransformInputT>>) runner::processElement);
- }
-
- // Register as a consumer for each timer PCollection.
- for (String localName : context.parDoPayload.getTimerSpecsMap().keySet()) {
- TimeDomain timeDomain =
- DoFnSignatures.getTimerSpecOrThrow(
- context.doFnSignature.timerDeclarations().get(localName), context.doFn)
- .getTimeDomain();
- pCollectionConsumerRegistry.register(
- pTransform.getInputsOrThrow(localName),
- pTransformId,
- (FnDataReceiver)
- timer ->
- runner.processTimer(
- localName, timeDomain, (WindowedValue<KV<Object, Timer>>) timer));
- }
-
- finishFunctionRegistry.register(pTransformId, runner::finishBundle);
- tearDownFunctions.accept(runner::tearDown);
- return runner;
- }
-
- abstract RunnerT createRunner(Context<FnInputT, OutputT> context);
-
- static class Context<InputT, OutputT> {
- final PipelineOptions pipelineOptions;
- final BeamFnStateClient beamFnStateClient;
- final String ptransformId;
- final PTransform pTransform;
- final Supplier<String> processBundleInstructionId;
- final RehydratedComponents rehydratedComponents;
- final DoFn<InputT, OutputT> doFn;
- final DoFnSignature doFnSignature;
- final TupleTag<OutputT> mainOutputTag;
- final Coder<?> inputCoder;
- final SchemaCoder<InputT> schemaCoder;
- final Coder<?> keyCoder;
- final SchemaCoder<OutputT> mainOutputSchemaCoder;
- final Coder<? extends BoundedWindow> windowCoder;
- final WindowingStrategy<InputT, ?> windowingStrategy;
- final Map<TupleTag<?>, SideInputSpec> tagToSideInputSpecMap;
- Map<TupleTag<?>, Coder<?>> outputCoders;
- final ParDoPayload parDoPayload;
- final ListMultimap<String, FnDataReceiver<WindowedValue<?>>> localNameToConsumer;
- final BundleSplitListener splitListener;
-
- Context(
- PipelineOptions pipelineOptions,
- BeamFnStateClient beamFnStateClient,
- String ptransformId,
- PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
- PCollectionConsumerRegistry pCollectionConsumerRegistry,
- BundleSplitListener splitListener) {
- this.pipelineOptions = pipelineOptions;
- this.beamFnStateClient = beamFnStateClient;
- this.ptransformId = ptransformId;
- this.pTransform = pTransform;
- this.processBundleInstructionId = processBundleInstructionId;
- ImmutableMap.Builder<TupleTag<?>, SideInputSpec> tagToSideInputSpecMapBuilder =
- ImmutableMap.builder();
- try {
- rehydratedComponents =
- RehydratedComponents.forComponents(
- RunnerApi.Components.newBuilder()
- .putAllCoders(coders)
- .putAllPcollections(pCollections)
- .putAllWindowingStrategies(windowingStrategies)
- .build())
- .withPipeline(Pipeline.create());
- parDoPayload = ParDoPayload.parseFrom(pTransform.getSpec().getPayload());
- doFn = (DoFn) ParDoTranslation.getDoFn(parDoPayload);
- doFnSignature = DoFnSignatures.signatureForDoFn(doFn);
- mainOutputTag = (TupleTag) ParDoTranslation.getMainOutputTag(parDoPayload);
- String mainInputTag =
- Iterables.getOnlyElement(
- Sets.difference(
- pTransform.getInputsMap().keySet(),
- Sets.union(
- parDoPayload.getSideInputsMap().keySet(),
- parDoPayload.getTimerSpecsMap().keySet())));
- PCollection mainInput = pCollections.get(pTransform.getInputsOrThrow(mainInputTag));
- inputCoder = rehydratedComponents.getCoder(mainInput.getCoderId());
- if (inputCoder instanceof KvCoder
- // TODO: Stop passing windowed value coders within PCollections.
- || (inputCoder instanceof WindowedValue.WindowedValueCoder
- && (((WindowedValueCoder) inputCoder).getValueCoder() instanceof KvCoder))) {
- this.keyCoder =
- inputCoder instanceof WindowedValueCoder
- ? ((KvCoder) ((WindowedValueCoder) inputCoder).getValueCoder()).getKeyCoder()
- : ((KvCoder) inputCoder).getKeyCoder();
- } else {
- this.keyCoder = null;
- }
- if (inputCoder instanceof SchemaCoder
- // TODO: Stop passing windowed value coders within PCollections.
- || (inputCoder instanceof WindowedValue.WindowedValueCoder
- && (((WindowedValueCoder) inputCoder).getValueCoder() instanceof SchemaCoder))) {
- this.schemaCoder =
- inputCoder instanceof WindowedValueCoder
- ? (SchemaCoder<InputT>) ((WindowedValueCoder) inputCoder).getValueCoder()
- : ((SchemaCoder<InputT>) inputCoder);
- } else {
- this.schemaCoder = null;
- }
-
- windowingStrategy =
- (WindowingStrategy)
- rehydratedComponents.getWindowingStrategy(mainInput.getWindowingStrategyId());
- windowCoder = windowingStrategy.getWindowFn().windowCoder();
-
- outputCoders = Maps.newHashMap();
- for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) {
- TupleTag<?> outputTag = new TupleTag<>(entry.getKey());
- RunnerApi.PCollection outputPCollection = pCollections.get(entry.getValue());
- Coder<?> outputCoder = rehydratedComponents.getCoder(outputPCollection.getCoderId());
- if (outputCoder instanceof WindowedValueCoder) {
- outputCoder = ((WindowedValueCoder) outputCoder).getValueCoder();
- }
- outputCoders.put(outputTag, outputCoder);
- }
- Coder<OutputT> outputCoder = (Coder<OutputT>) outputCoders.get(mainOutputTag);
- mainOutputSchemaCoder =
- (outputCoder instanceof SchemaCoder) ? (SchemaCoder<OutputT>) outputCoder : null;
-
- // Build the map from tag id to side input specification
- for (Map.Entry<String, RunnerApi.SideInput> entry :
- parDoPayload.getSideInputsMap().entrySet()) {
- String sideInputTag = entry.getKey();
- RunnerApi.SideInput sideInput = entry.getValue();
- checkArgument(
- Materializations.MULTIMAP_MATERIALIZATION_URN.equals(
- sideInput.getAccessPattern().getUrn()),
- "This SDK is only capable of dealing with %s materializations "
- + "but was asked to handle %s for PCollectionView with tag %s.",
- Materializations.MULTIMAP_MATERIALIZATION_URN,
- sideInput.getAccessPattern().getUrn(),
- sideInputTag);
-
- PCollection sideInputPCollection =
- pCollections.get(pTransform.getInputsOrThrow(sideInputTag));
- WindowingStrategy sideInputWindowingStrategy =
- rehydratedComponents.getWindowingStrategy(
- sideInputPCollection.getWindowingStrategyId());
- tagToSideInputSpecMapBuilder.put(
- new TupleTag<>(entry.getKey()),
- SideInputSpec.create(
- rehydratedComponents.getCoder(sideInputPCollection.getCoderId()),
- sideInputWindowingStrategy.getWindowFn().windowCoder(),
- PCollectionViewTranslation.viewFnFromProto(entry.getValue().getViewFn()),
- PCollectionViewTranslation.windowMappingFnFromProto(
- entry.getValue().getWindowMappingFn())));
- }
- } catch (IOException exn) {
- throw new IllegalArgumentException("Malformed ParDoPayload", exn);
- }
-
- ImmutableListMultimap.Builder<String, FnDataReceiver<WindowedValue<?>>>
- localNameToConsumerBuilder = ImmutableListMultimap.builder();
- for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) {
- localNameToConsumerBuilder.putAll(
- entry.getKey(), pCollectionConsumerRegistry.getMultiplexingConsumer(entry.getValue()));
- }
- localNameToConsumer = localNameToConsumerBuilder.build();
- tagToSideInputSpecMap = tagToSideInputSpecMapBuilder.build();
- this.splitListener = splitListener;
- }
- }
-}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
index a44629f..fa83614 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
@@ -22,45 +22,82 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
import com.google.auto.service.AutoService;
+import com.google.auto.value.AutoValue;
import java.io.IOException;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
-import org.apache.beam.fn.harness.DoFnPTransformRunnerFactory.Context;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+import org.apache.beam.fn.harness.control.BundleSplitListener;
+import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
+import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
+import org.apache.beam.fn.harness.state.BeamFnStateClient;
import org.apache.beam.fn.harness.state.FnApiStateAccessor;
+import org.apache.beam.fn.harness.state.SideInputSpec;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
+import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
+import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
import org.apache.beam.runners.core.DoFnRunner;
import org.apache.beam.runners.core.LateDataUtils;
+import org.apache.beam.runners.core.construction.PCollectionViewTranslation;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.ParDoTranslation;
+import org.apache.beam.runners.core.construction.RehydratedComponents;
import org.apache.beam.runners.core.construction.Timer;
+import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.function.ThrowingRunnable;
import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.TimerMap;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
import org.apache.beam.sdk.transforms.DoFnOutputReceivers;
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
+import org.apache.beam.sdk.transforms.Materializations;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.StateDeclaration;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerDeclaration;
import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.Sizes.HasSize;
+import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.util.Durations;
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.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableListMultimap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ListMultimap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
import org.joda.time.DateTimeUtils;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -70,23 +107,275 @@
* abstraction caused by StateInternals/TimerInternals since they model state and timer concepts
* differently.
*/
-public class FnApiDoFnRunner<InputT, OutputT>
- implements DoFnPTransformRunnerFactory.DoFnPTransformRunner<InputT> {
+public class FnApiDoFnRunner<InputT, RestrictionT, PositionT, OutputT> {
/** A registrar which provides a factory to handle Java {@link DoFn}s. */
@AutoService(PTransformRunnerFactory.Registrar.class)
public static class Registrar implements PTransformRunnerFactory.Registrar {
@Override
public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
- return ImmutableMap.of(PTransformTranslation.PAR_DO_TRANSFORM_URN, new Factory());
+ Factory factory = new Factory();
+ return ImmutableMap.<String, PTransformRunnerFactory>builder()
+ .put(PTransformTranslation.PAR_DO_TRANSFORM_URN, factory)
+ .put(PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN, factory)
+ .put(PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN, factory)
+ .put(PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN, factory)
+ .put(PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN, factory)
+ .put(
+ PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN, factory)
+ .build();
}
}
- static class Factory<InputT, OutputT>
- extends DoFnPTransformRunnerFactory<
- InputT, InputT, OutputT, FnApiDoFnRunner<InputT, OutputT>> {
+ static class Context<InputT, OutputT> {
+ final PipelineOptions pipelineOptions;
+ final BeamFnStateClient beamFnStateClient;
+ final String ptransformId;
+ final PTransform pTransform;
+ final Supplier<String> processBundleInstructionId;
+ final RehydratedComponents rehydratedComponents;
+ final DoFn<InputT, OutputT> doFn;
+ final DoFnSignature doFnSignature;
+ final TupleTag<OutputT> mainOutputTag;
+ final Coder<?> inputCoder;
+ final SchemaCoder<InputT> schemaCoder;
+ final Coder<?> keyCoder;
+ final SchemaCoder<OutputT> mainOutputSchemaCoder;
+ final Coder<? extends BoundedWindow> windowCoder;
+ final WindowingStrategy<InputT, ?> windowingStrategy;
+ final Map<TupleTag<?>, SideInputSpec> tagToSideInputSpecMap;
+ Map<TupleTag<?>, Coder<?>> outputCoders;
+ final ParDoPayload parDoPayload;
+ final ListMultimap<String, FnDataReceiver<WindowedValue<?>>> localNameToConsumer;
+ final BundleSplitListener splitListener;
+
+ Context(
+ PipelineOptions pipelineOptions,
+ BeamFnStateClient beamFnStateClient,
+ String ptransformId,
+ PTransform pTransform,
+ Supplier<String> processBundleInstructionId,
+ Map<String, PCollection> pCollections,
+ Map<String, RunnerApi.Coder> coders,
+ Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
+ BundleSplitListener splitListener) {
+ this.pipelineOptions = pipelineOptions;
+ this.beamFnStateClient = beamFnStateClient;
+ this.ptransformId = ptransformId;
+ this.pTransform = pTransform;
+ this.processBundleInstructionId = processBundleInstructionId;
+ ImmutableMap.Builder<TupleTag<?>, SideInputSpec> tagToSideInputSpecMapBuilder =
+ ImmutableMap.builder();
+ try {
+ rehydratedComponents =
+ RehydratedComponents.forComponents(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(coders)
+ .putAllPcollections(pCollections)
+ .putAllWindowingStrategies(windowingStrategies)
+ .build())
+ .withPipeline(Pipeline.create());
+ parDoPayload = ParDoPayload.parseFrom(pTransform.getSpec().getPayload());
+ doFn = (DoFn) ParDoTranslation.getDoFn(parDoPayload);
+ doFnSignature = DoFnSignatures.signatureForDoFn(doFn);
+ switch (pTransform.getSpec().getUrn()) {
+ case PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN:
+ case PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
+ case PTransformTranslation.PAR_DO_TRANSFORM_URN:
+ mainOutputTag = (TupleTag) ParDoTranslation.getMainOutputTag(parDoPayload);
+ break;
+ case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
+ case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
+ case PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN:
+ mainOutputTag =
+ new TupleTag(Iterables.getOnlyElement(pTransform.getOutputsMap().keySet()));
+ break;
+ default:
+ throw new IllegalStateException(
+ String.format("Unknown urn: %s", pTransform.getSpec().getUrn()));
+ }
+ String mainInputTag =
+ Iterables.getOnlyElement(
+ Sets.difference(
+ pTransform.getInputsMap().keySet(),
+ Sets.union(
+ parDoPayload.getSideInputsMap().keySet(),
+ parDoPayload.getTimerSpecsMap().keySet())));
+ PCollection mainInput = pCollections.get(pTransform.getInputsOrThrow(mainInputTag));
+ inputCoder = rehydratedComponents.getCoder(mainInput.getCoderId());
+ if (inputCoder instanceof KvCoder
+ // TODO: Stop passing windowed value coders within PCollections.
+ || (inputCoder instanceof WindowedValue.WindowedValueCoder
+ && (((WindowedValueCoder) inputCoder).getValueCoder() instanceof KvCoder))) {
+ this.keyCoder =
+ inputCoder instanceof WindowedValueCoder
+ ? ((KvCoder) ((WindowedValueCoder) inputCoder).getValueCoder()).getKeyCoder()
+ : ((KvCoder) inputCoder).getKeyCoder();
+ } else {
+ this.keyCoder = null;
+ }
+ if (inputCoder instanceof SchemaCoder
+ // TODO: Stop passing windowed value coders within PCollections.
+ || (inputCoder instanceof WindowedValue.WindowedValueCoder
+ && (((WindowedValueCoder) inputCoder).getValueCoder() instanceof SchemaCoder))) {
+ this.schemaCoder =
+ inputCoder instanceof WindowedValueCoder
+ ? (SchemaCoder<InputT>) ((WindowedValueCoder) inputCoder).getValueCoder()
+ : ((SchemaCoder<InputT>) inputCoder);
+ } else {
+ this.schemaCoder = null;
+ }
+
+ windowingStrategy =
+ (WindowingStrategy)
+ rehydratedComponents.getWindowingStrategy(mainInput.getWindowingStrategyId());
+ windowCoder = windowingStrategy.getWindowFn().windowCoder();
+
+ outputCoders = Maps.newHashMap();
+ for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) {
+ TupleTag<?> outputTag = new TupleTag<>(entry.getKey());
+ RunnerApi.PCollection outputPCollection = pCollections.get(entry.getValue());
+ Coder<?> outputCoder = rehydratedComponents.getCoder(outputPCollection.getCoderId());
+ if (outputCoder instanceof WindowedValueCoder) {
+ outputCoder = ((WindowedValueCoder) outputCoder).getValueCoder();
+ }
+ outputCoders.put(outputTag, outputCoder);
+ }
+ Coder<OutputT> outputCoder = (Coder<OutputT>) outputCoders.get(mainOutputTag);
+ mainOutputSchemaCoder =
+ (outputCoder instanceof SchemaCoder) ? (SchemaCoder<OutputT>) outputCoder : null;
+
+ // Build the map from tag id to side input specification
+ for (Map.Entry<String, RunnerApi.SideInput> entry :
+ parDoPayload.getSideInputsMap().entrySet()) {
+ String sideInputTag = entry.getKey();
+ RunnerApi.SideInput sideInput = entry.getValue();
+ checkArgument(
+ Materializations.MULTIMAP_MATERIALIZATION_URN.equals(
+ sideInput.getAccessPattern().getUrn()),
+ "This SDK is only capable of dealing with %s materializations "
+ + "but was asked to handle %s for PCollectionView with tag %s.",
+ Materializations.MULTIMAP_MATERIALIZATION_URN,
+ sideInput.getAccessPattern().getUrn(),
+ sideInputTag);
+
+ PCollection sideInputPCollection =
+ pCollections.get(pTransform.getInputsOrThrow(sideInputTag));
+ WindowingStrategy sideInputWindowingStrategy =
+ rehydratedComponents.getWindowingStrategy(
+ sideInputPCollection.getWindowingStrategyId());
+ tagToSideInputSpecMapBuilder.put(
+ new TupleTag<>(entry.getKey()),
+ SideInputSpec.create(
+ rehydratedComponents.getCoder(sideInputPCollection.getCoderId()),
+ sideInputWindowingStrategy.getWindowFn().windowCoder(),
+ PCollectionViewTranslation.viewFnFromProto(entry.getValue().getViewFn()),
+ PCollectionViewTranslation.windowMappingFnFromProto(
+ entry.getValue().getWindowMappingFn())));
+ }
+ } catch (IOException exn) {
+ throw new IllegalArgumentException("Malformed ParDoPayload", exn);
+ }
+
+ ImmutableListMultimap.Builder<String, FnDataReceiver<WindowedValue<?>>>
+ localNameToConsumerBuilder = ImmutableListMultimap.builder();
+ for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) {
+ localNameToConsumerBuilder.putAll(
+ entry.getKey(), pCollectionConsumerRegistry.getMultiplexingConsumer(entry.getValue()));
+ }
+ localNameToConsumer = localNameToConsumerBuilder.build();
+ tagToSideInputSpecMap = tagToSideInputSpecMapBuilder.build();
+ this.splitListener = splitListener;
+ }
+ }
+
+ static class Factory<InputT, RestrictionT, PositionT, OutputT>
+ implements PTransformRunnerFactory<
+ FnApiDoFnRunner<InputT, RestrictionT, PositionT, OutputT>> {
+
@Override
- public FnApiDoFnRunner<InputT, OutputT> createRunner(Context<InputT, OutputT> context) {
- return new FnApiDoFnRunner<>(context);
+ public final FnApiDoFnRunner<InputT, RestrictionT, PositionT, OutputT>
+ createRunnerForPTransform(
+ PipelineOptions pipelineOptions,
+ BeamFnDataClient beamFnDataClient,
+ BeamFnStateClient beamFnStateClient,
+ String pTransformId,
+ PTransform pTransform,
+ Supplier<String> processBundleInstructionId,
+ Map<String, PCollection> pCollections,
+ Map<String, RunnerApi.Coder> coders,
+ Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
+ PCollectionConsumerRegistry pCollectionConsumerRegistry,
+ PTransformFunctionRegistry startFunctionRegistry,
+ PTransformFunctionRegistry finishFunctionRegistry,
+ Consumer<ThrowingRunnable> tearDownFunctions,
+ BundleSplitListener splitListener) {
+ Context<InputT, OutputT> context =
+ new Context<>(
+ pipelineOptions,
+ beamFnStateClient,
+ pTransformId,
+ pTransform,
+ processBundleInstructionId,
+ pCollections,
+ coders,
+ windowingStrategies,
+ pCollectionConsumerRegistry,
+ splitListener);
+
+ FnApiDoFnRunner<InputT, RestrictionT, PositionT, OutputT> runner =
+ new FnApiDoFnRunner<>(context);
+
+ // Register the appropriate handlers.
+ startFunctionRegistry.register(pTransformId, runner::startBundle);
+ String mainInput;
+ try {
+ mainInput = ParDoTranslation.getMainInputName(pTransform);
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ FnDataReceiver<WindowedValue> mainInputConsumer;
+ switch (pTransform.getSpec().getUrn()) {
+ case PTransformTranslation.PAR_DO_TRANSFORM_URN:
+ mainInputConsumer = runner::processElementForParDo;
+ break;
+ case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
+ mainInputConsumer = runner::processElementForPairWithRestriction;
+ break;
+ case PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN:
+ case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
+ mainInputConsumer = runner::processElementForSplitRestriction;
+ break;
+ case PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN:
+ mainInputConsumer = runner::processElementForElementAndRestriction;
+ break;
+ case PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
+ mainInputConsumer = runner::processElementForSizedElementAndRestriction;
+ break;
+ default:
+ throw new IllegalStateException("Unknown urn: " + pTransform.getSpec().getUrn());
+ }
+ pCollectionConsumerRegistry.register(
+ pTransform.getInputsOrThrow(mainInput), pTransformId, (FnDataReceiver) mainInputConsumer);
+
+ // Register as a consumer for each timer PCollection.
+ for (String localName : context.parDoPayload.getTimerSpecsMap().keySet()) {
+ TimeDomain timeDomain =
+ DoFnSignatures.getTimerSpecOrThrow(
+ context.doFnSignature.timerDeclarations().get(localName), context.doFn)
+ .getTimeDomain();
+ pCollectionConsumerRegistry.register(
+ pTransform.getInputsOrThrow(localName),
+ pTransformId,
+ (FnDataReceiver)
+ timer ->
+ runner.processTimer(
+ localName, timeDomain, (WindowedValue<KV<Object, Timer>>) timer));
+ }
+
+ finishFunctionRegistry.register(pTransformId, runner::finishBundle);
+ tearDownFunctions.accept(runner::tearDown);
+ return runner;
}
}
@@ -94,32 +383,58 @@
private final Context<InputT, OutputT> context;
private final Collection<FnDataReceiver<WindowedValue<OutputT>>> mainOutputConsumers;
+ private final String mainInputId;
private FnApiStateAccessor stateAccessor;
private final DoFnInvoker<InputT, OutputT> doFnInvoker;
private final DoFn<InputT, OutputT>.StartBundleContext startBundleContext;
private final ProcessBundleContext processContext;
private final OnTimerContext onTimerContext;
private final DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext;
+ /**
+ * Only set for {@link PTransformTranslation#SPLITTABLE_SPLIT_RESTRICTION_URN} and {@link
+ * PTransformTranslation#SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN} transforms. Can only be
+ * invoked from within {@code processElement...} methods.
+ */
+ private final OutputReceiver<RestrictionT> outputSplitRestrictionReceiver;
+ /**
+ * Only set for {@link PTransformTranslation#SPLITTABLE_PROCESS_ELEMENTS_URN} and {@link
+ * PTransformTranslation#SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN} transforms. Can
+ * only be invoked from within {@code processElement...} methods.
+ */
+ private final Function<SplitResult<RestrictionT>, WindowedSplitResult>
+ convertSplitResultToWindowedSplitResult;
- /** Only valid during {@link #processElement}, null otherwise. */
+ private final DoFnSchemaInformation doFnSchemaInformation;
+ private final Map<String, PCollectionView<?>> sideInputMapping;
+
+ // The member variables below are only valid for the lifetime of certain methods.
+ /** Only valid during {@code processElement...} methods, null otherwise. */
private WindowedValue<InputT> currentElement;
- /** Only valid during {@link #processElement} and {@link #processTimer}, null otherwise. */
+ /**
+ * Only valid during {@code processElement...} and {@link #processTimer} methods, null otherwise.
+ */
private BoundedWindow currentWindow;
+ /**
+ * Only valid during {@link #processElementForElementAndRestriction} and {@link
+ * #processElementForSizedElementAndRestriction}, null otherwise.
+ */
+ private RestrictionTracker<RestrictionT, PositionT> currentTracker;
+
/** Only valid during {@link #processTimer}, null otherwise. */
private WindowedValue<KV<Object, Timer>> currentTimer;
/** Only valid during {@link #processTimer}, null otherwise. */
private TimeDomain currentTimeDomain;
- private DoFnSchemaInformation doFnSchemaInformation;
-
- private Map<String, PCollectionView<?>> sideInputMapping;
-
FnApiDoFnRunner(Context<InputT, OutputT> context) {
this.context = context;
-
+ try {
+ this.mainInputId = ParDoTranslation.getMainInputName(context.pTransform);
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
this.mainOutputConsumers =
(Collection<FnDataReceiver<WindowedValue<OutputT>>>)
(Collection) context.localNameToConsumer.get(context.mainOutputTag.getId());
@@ -162,9 +477,133 @@
outputTo(consumers, WindowedValue.of(output, timestamp, window, PaneInfo.NO_FIRING));
}
};
+ switch (context.pTransform.getSpec().getUrn()) {
+ case PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN:
+ this.outputSplitRestrictionReceiver =
+ new OutputReceiver<RestrictionT>() {
+
+ @Override
+ public void output(RestrictionT output) {
+ outputTo(
+ mainOutputConsumers,
+ (WindowedValue<OutputT>)
+ currentElement.withValue(KV.of(currentElement.getValue(), output)));
+ }
+
+ @Override
+ public void outputWithTimestamp(RestrictionT output, Instant timestamp) {
+ outputTo(
+ mainOutputConsumers,
+ (WindowedValue<OutputT>)
+ WindowedValue.of(
+ KV.of(currentElement.getValue(), output),
+ timestamp,
+ currentWindow,
+ currentElement.getPane()));
+ }
+ };
+ break;
+ case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
+ this.outputSplitRestrictionReceiver =
+ new OutputReceiver<RestrictionT>() {
+
+ @Override
+ public void output(RestrictionT output) {
+ RestrictionTracker<RestrictionT, PositionT> outputTracker =
+ doFnInvoker.invokeNewTracker(output);
+ outputTo(
+ mainOutputConsumers,
+ (WindowedValue<OutputT>)
+ currentElement.withValue(
+ KV.of(
+ KV.of(currentElement.getValue(), output),
+ outputTracker instanceof HasSize
+ ? ((HasSize) outputTracker).getSize()
+ : 1.0)));
+ }
+
+ @Override
+ public void outputWithTimestamp(RestrictionT output, Instant timestamp) {
+ RestrictionTracker<RestrictionT, PositionT> outputTracker =
+ doFnInvoker.invokeNewTracker(output);
+ outputTo(
+ mainOutputConsumers,
+ (WindowedValue<OutputT>)
+ WindowedValue.of(
+ KV.of(
+ KV.of(currentElement.getValue(), output),
+ outputTracker instanceof HasSize
+ ? ((HasSize) outputTracker).getSize()
+ : 1.0),
+ timestamp,
+ currentWindow,
+ currentElement.getPane()));
+ }
+ };
+ break;
+ default:
+ this.outputSplitRestrictionReceiver =
+ new OutputReceiver<RestrictionT>() {
+ @Override
+ public void output(RestrictionT output) {
+ throw new IllegalStateException(
+ String.format(
+ "Unimplemented split output handler for %s.",
+ context.pTransform.getSpec().getUrn()));
+ }
+
+ @Override
+ public void outputWithTimestamp(RestrictionT output, Instant timestamp) {
+ throw new IllegalStateException(
+ String.format(
+ "Unimplemented split output handler for %s.",
+ context.pTransform.getSpec().getUrn()));
+ }
+ };
+ }
+ switch (context.pTransform.getSpec().getUrn()) {
+ case PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN:
+ this.convertSplitResultToWindowedSplitResult =
+ (splitResult) ->
+ WindowedSplitResult.forRoots(
+ currentElement.withValue(
+ KV.of(currentElement.getValue(), splitResult.getPrimary())),
+ currentElement.withValue(
+ KV.of(currentElement.getValue(), splitResult.getResidual())));
+ break;
+ case PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
+ this.convertSplitResultToWindowedSplitResult =
+ (splitResult) -> {
+ RestrictionTracker<RestrictionT, PositionT> primaryTracker =
+ doFnInvoker.invokeNewTracker(splitResult.getPrimary());
+ RestrictionTracker<RestrictionT, PositionT> residualTracker =
+ doFnInvoker.invokeNewTracker(splitResult.getResidual());
+ return WindowedSplitResult.forRoots(
+ currentElement.withValue(
+ KV.of(
+ KV.of(currentElement.getValue(), splitResult.getPrimary()),
+ primaryTracker instanceof HasSize
+ ? ((HasSize) primaryTracker).getSize()
+ : 1.0)),
+ currentElement.withValue(
+ KV.of(
+ KV.of(currentElement.getValue(), splitResult.getResidual()),
+ residualTracker instanceof HasSize
+ ? ((HasSize) residualTracker).getSize()
+ : 1.0)));
+ };
+ break;
+ default:
+ this.convertSplitResultToWindowedSplitResult =
+ (splitResult) -> {
+ throw new IllegalStateException(
+ String.format(
+ "Unimplemented split conversion handler for %s.",
+ context.pTransform.getSpec().getUrn()));
+ };
+ }
}
- @Override
public void startBundle() {
this.stateAccessor =
new FnApiStateAccessor(
@@ -181,8 +620,7 @@
doFnInvoker.invokeStartBundle(startBundleContext);
}
- @Override
- public void processElement(WindowedValue<InputT> elem) {
+ public void processElementForParDo(WindowedValue<InputT> elem) {
currentElement = elem;
try {
Iterator<BoundedWindow> windowIterator =
@@ -197,7 +635,129 @@
}
}
- @Override
+ public void processElementForPairWithRestriction(WindowedValue<InputT> elem) {
+ currentElement = elem;
+ try {
+ Iterator<BoundedWindow> windowIterator =
+ (Iterator<BoundedWindow>) elem.getWindows().iterator();
+ while (windowIterator.hasNext()) {
+ currentWindow = windowIterator.next();
+ outputTo(
+ mainOutputConsumers,
+ (WindowedValue)
+ elem.withValue(
+ KV.of(
+ elem.getValue(),
+ doFnInvoker.invokeGetInitialRestriction(elem.getValue()))));
+ }
+ } finally {
+ currentElement = null;
+ currentWindow = null;
+ }
+ }
+
+ public void processElementForSplitRestriction(WindowedValue<KV<InputT, RestrictionT>> elem) {
+ currentElement = elem.withValue(elem.getValue().getKey());
+ try {
+ Iterator<BoundedWindow> windowIterator =
+ (Iterator<BoundedWindow>) elem.getWindows().iterator();
+ while (windowIterator.hasNext()) {
+ currentWindow = windowIterator.next();
+ doFnInvoker.invokeSplitRestriction(
+ elem.getValue().getKey(),
+ elem.getValue().getValue(),
+ this.outputSplitRestrictionReceiver);
+ }
+ } finally {
+ currentElement = null;
+ currentWindow = null;
+ }
+ }
+
+ /** Internal class to hold the primary and residual roots when converted to an input element. */
+ @AutoValue
+ abstract static class WindowedSplitResult {
+ public static WindowedSplitResult forRoots(
+ WindowedValue primaryRoot, WindowedValue residualRoot) {
+ return new AutoValue_FnApiDoFnRunner_WindowedSplitResult(primaryRoot, residualRoot);
+ }
+
+ public abstract WindowedValue getPrimaryRoot();
+
+ public abstract WindowedValue getResidualRoot();
+ }
+
+ public void processElementForSizedElementAndRestriction(
+ WindowedValue<KV<KV<InputT, RestrictionT>, Double>> elem) {
+ processElementForElementAndRestriction(elem.withValue(elem.getValue().getKey()));
+ }
+
+ public void processElementForElementAndRestriction(WindowedValue<KV<InputT, RestrictionT>> elem) {
+ currentElement = elem.withValue(elem.getValue().getKey());
+ try {
+ Iterator<BoundedWindow> windowIterator =
+ (Iterator<BoundedWindow>) elem.getWindows().iterator();
+ while (windowIterator.hasNext()) {
+ currentTracker = doFnInvoker.invokeNewTracker(elem.getValue().getValue());
+ currentWindow = windowIterator.next();
+ DoFn.ProcessContinuation continuation = doFnInvoker.invokeProcessElement(processContext);
+ // Ensure that all the work is done if the user tells us that they don't want to
+ // resume processing.
+ if (!continuation.shouldResume()) {
+ currentTracker.checkDone();
+ continue;
+ }
+
+ SplitResult<RestrictionT> result = currentTracker.trySplit(0);
+ // After the user has chosen to resume processing later, the Runner may have stolen
+ // the remainder of work through a split call so the above trySplit may fail. If so,
+ // the current restriction must be done.
+ if (result == null) {
+ currentTracker.checkDone();
+ continue;
+ }
+
+ // Otherwise we have a successful self checkpoint.
+ WindowedSplitResult windowedSplitResult =
+ convertSplitResultToWindowedSplitResult.apply(result);
+ ByteString.Output primaryBytes = ByteString.newOutput();
+ ByteString.Output residualBytes = ByteString.newOutput();
+ try {
+ Coder fullInputCoder =
+ WindowedValue.getFullCoder(context.inputCoder, context.windowCoder);
+ fullInputCoder.encode(windowedSplitResult.getPrimaryRoot(), primaryBytes);
+ fullInputCoder.encode(windowedSplitResult.getResidualRoot(), residualBytes);
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ BundleApplication primaryApplication =
+ BundleApplication.newBuilder()
+ .setTransformId(context.ptransformId)
+ .setInputId(mainInputId)
+ .setElement(primaryBytes.toByteString())
+ .build();
+ BundleApplication residualApplication =
+ BundleApplication.newBuilder()
+ .setTransformId(context.ptransformId)
+ .setInputId(mainInputId)
+ .setElement(residualBytes.toByteString())
+ .build();
+ context.splitListener.split(
+ ImmutableList.of(primaryApplication),
+ ImmutableList.of(
+ DelayedBundleApplication.newBuilder()
+ .setApplication(residualApplication)
+ .setRequestedTimeDelay(
+ Durations.fromMillis(continuation.resumeDelay().getMillis()))
+ .build()));
+ }
+ } finally {
+ currentElement = null;
+ currentWindow = null;
+ currentTracker = null;
+ }
+ }
+
public void processTimer(
String timerId, TimeDomain timeDomain, WindowedValue<KV<Object, Timer>> timer) {
currentTimer = timer;
@@ -207,7 +767,7 @@
(Iterator<BoundedWindow>) timer.getWindows().iterator();
while (windowIterator.hasNext()) {
currentWindow = windowIterator.next();
- doFnInvoker.invokeOnTimer(timerId, onTimerContext);
+ doFnInvoker.invokeOnTimer(timerId, timerId, onTimerContext);
}
} finally {
currentTimer = null;
@@ -216,7 +776,6 @@
}
}
- @Override
public void finishBundle() {
doFnInvoker.invokeFinishBundle(finishBundleContext);
@@ -225,7 +784,6 @@
this.stateAccessor = null;
}
- @Override
public void tearDown() {
doFnInvoker.invokeTeardown();
}
@@ -248,6 +806,7 @@
private final Instant currentTimestamp;
private final Duration allowedLateness;
private final WindowedValue<?> currentElementOrTimer;
+ @Nullable private Instant currentOutputTimestamp;
private Duration period = Duration.ZERO;
private Duration offset = Duration.ZERO;
@@ -338,6 +897,20 @@
return this;
}
+ @Override
+ public org.apache.beam.sdk.state.Timer withOutputTimestamp(Instant outputTime) {
+ Instant windowExpiry = LateDataUtils.garbageCollectionTime(currentWindow, allowedLateness);
+ checkArgument(
+ !outputTime.isAfter(windowExpiry),
+ "Attempted to set timer with output timestamp %s but that is after"
+ + " the expiration of window %s",
+ outputTime,
+ windowExpiry);
+
+ this.currentOutputTimestamp = outputTime;
+ return this;
+ }
+
/**
* For event time timers the target time should be prior to window GC time. So it returns
* min(time to set, GC Time of window).
@@ -357,7 +930,32 @@
Collection<FnDataReceiver<WindowedValue<KV<Object, Timer>>>> consumers =
(Collection) context.localNameToConsumer.get(timerId);
- outputTo(consumers, currentElementOrTimer.withValue(KV.of(key, Timer.of(scheduledTime))));
+ if (currentOutputTimestamp == null) {
+ if (TimeDomain.EVENT_TIME.equals(timeDomain)) {
+ currentOutputTimestamp = scheduledTime;
+ } else {
+ currentOutputTimestamp = currentElementOrTimer.getTimestamp();
+ }
+ }
+ outputTo(
+ consumers,
+ WindowedValue.of(
+ KV.of(key, Timer.of(scheduledTime)),
+ currentOutputTimestamp,
+ currentElementOrTimer.getWindows(),
+ currentElementOrTimer.getPane()));
+ }
+ }
+
+ private static class FnApiTimerMap implements TimerMap {
+ FnApiTimerMap() {}
+
+ @Override
+ public void set(String timerId, Instant absoluteTime) {}
+
+ @Override
+ public org.apache.beam.sdk.state.Timer get(String timerId) {
+ return null;
}
}
@@ -421,6 +1019,12 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException(
+ "Cannot access timerId as parameter outside of @OnTimer method.");
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access time domain outside of @ProcessTimer method.");
@@ -449,7 +1053,7 @@
@Override
public RestrictionTracker<?, ?> restrictionTracker() {
- throw new UnsupportedOperationException("RestrictionTracker parameters are not supported.");
+ return currentTracker;
}
@Override
@@ -476,6 +1080,12 @@
}
@Override
+ public TimerMap timerFamily(String tagId) {
+ // TODO: implement timerFamily
+ return null;
+ }
+
+ @Override
public PipelineOptions getPipelineOptions() {
return context.pipelineOptions;
}
@@ -609,6 +1219,11 @@
}
@Override
+ public String timerId(DoFn<InputT, OutputT> doFn) {
+ throw new UnsupportedOperationException("TimerId parameters are not supported.");
+ }
+
+ @Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
return timeDomain();
}
@@ -662,6 +1277,12 @@
}
@Override
+ public TimerMap timerFamily(String tagId) {
+ // TODO: implement timerFamily
+ throw new UnsupportedOperationException("TimerFamily parameters are not supported.");
+ }
+
+ @Override
public PipelineOptions getPipelineOptions() {
return context.pipelineOptions;
}
@@ -723,6 +1344,11 @@
}
@Override
+ public Instant fireTimestamp() {
+ return currentTimer.getValue().getValue().getTimestamp();
+ }
+
+ @Override
public Instant timestamp() {
return currentTimer.getTimestamp();
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
index 6ec1673..579d447 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnHarness.java
@@ -44,7 +44,7 @@
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.options.ExperimentalOptions;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
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.collect.ImmutableList;
import org.slf4j.Logger;
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/HandlesSplits.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/HandlesSplits.java
new file mode 100644
index 0000000..a2ac123
--- /dev/null
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/HandlesSplits.java
@@ -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.
+ */
+package org.apache.beam.fn.harness;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+
+public interface HandlesSplits {
+ SplitResult trySplit(double fractionOfRemainder);
+
+ double getProgress();
+
+ @AutoValue
+ abstract class SplitResult {
+ public static SplitResult of(
+ BeamFnApi.BundleApplication primaryRoot, BeamFnApi.DelayedBundleApplication residualRoot) {
+ return new AutoValue_HandlesSplits_SplitResult(primaryRoot, residualRoot);
+ }
+
+ public abstract BeamFnApi.BundleApplication getPrimaryRoot();
+
+ public abstract BeamFnApi.DelayedBundleApplication getResidualRoot();
+ }
+}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableProcessElementsRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableProcessElementsRunner.java
deleted file mode 100644
index b32f28a..0000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/SplittableProcessElementsRunner.java
+++ /dev/null
@@ -1,282 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT 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.fn.harness;
-
-import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
-
-import com.google.auto.service.AutoService;
-import java.io.IOException;
-import java.util.Collection;
-import java.util.Map;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import org.apache.beam.fn.harness.DoFnPTransformRunnerFactory.Context;
-import org.apache.beam.fn.harness.state.FnApiStateAccessor;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
-import org.apache.beam.runners.core.OutputAndTimeBoundedSplittableProcessElementInvoker;
-import org.apache.beam.runners.core.OutputWindowedValue;
-import org.apache.beam.runners.core.SplittableProcessElementInvoker;
-import org.apache.beam.runners.core.construction.PTransformTranslation;
-import org.apache.beam.runners.core.construction.Timer;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.PaneInfo;
-import org.apache.beam.sdk.util.UserCodeException;
-import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
-import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.util.Durations;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-import org.joda.time.Duration;
-import org.joda.time.Instant;
-
-/** Runs the {@link PTransformTranslation#SPLITTABLE_PROCESS_ELEMENTS_URN} transform. */
-public class SplittableProcessElementsRunner<InputT, RestrictionT, OutputT>
- implements DoFnPTransformRunnerFactory.DoFnPTransformRunner<KV<InputT, RestrictionT>> {
- /** A registrar which provides a factory to handle Java {@link DoFn}s. */
- @AutoService(PTransformRunnerFactory.Registrar.class)
- public static class Registrar implements PTransformRunnerFactory.Registrar {
- @Override
- public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
- return ImmutableMap.of(PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN, new Factory());
- }
- }
-
- static class Factory<InputT, RestrictionT, OutputT>
- extends DoFnPTransformRunnerFactory<
- KV<InputT, RestrictionT>,
- InputT,
- OutputT,
- SplittableProcessElementsRunner<InputT, RestrictionT, OutputT>> {
-
- @Override
- SplittableProcessElementsRunner<InputT, RestrictionT, OutputT> createRunner(
- Context<InputT, OutputT> context) {
- Coder<WindowedValue<KV<InputT, RestrictionT>>> windowedCoder =
- FullWindowedValueCoder.of(
- (Coder<KV<InputT, RestrictionT>>) context.inputCoder, context.windowCoder);
-
- return new SplittableProcessElementsRunner<>(
- context,
- windowedCoder,
- (Collection<FnDataReceiver<WindowedValue<OutputT>>>)
- (Collection) context.localNameToConsumer.get(context.mainOutputTag.getId()),
- Iterables.getOnlyElement(context.pTransform.getInputsMap().keySet()));
- }
- }
-
- //////////////////////////////////////////////////////////////////////////////////////////////////
-
- private final Context<InputT, OutputT> context;
- private final String mainInputId;
- private final Coder<WindowedValue<KV<InputT, RestrictionT>>> inputCoder;
- private final Collection<FnDataReceiver<WindowedValue<OutputT>>> mainOutputConsumers;
- private final DoFnInvoker<InputT, OutputT> doFnInvoker;
- private final ScheduledExecutorService executor;
-
- private FnApiStateAccessor stateAccessor;
-
- private final DoFn<InputT, OutputT>.StartBundleContext startBundleContext;
- private final DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext;
-
- SplittableProcessElementsRunner(
- Context<InputT, OutputT> context,
- Coder<WindowedValue<KV<InputT, RestrictionT>>> inputCoder,
- Collection<FnDataReceiver<WindowedValue<OutputT>>> mainOutputConsumers,
- String mainInputId) {
- this.context = context;
- this.mainInputId = mainInputId;
- this.inputCoder = inputCoder;
- this.mainOutputConsumers = mainOutputConsumers;
- this.doFnInvoker = DoFnInvokers.invokerFor(context.doFn);
- this.doFnInvoker.invokeSetup();
- this.executor = Executors.newSingleThreadScheduledExecutor();
-
- this.startBundleContext =
- context.doFn.new StartBundleContext() {
- @Override
- public PipelineOptions getPipelineOptions() {
- return context.pipelineOptions;
- }
- };
- this.finishBundleContext =
- context.doFn.new FinishBundleContext() {
- @Override
- public PipelineOptions getPipelineOptions() {
- return context.pipelineOptions;
- }
-
- @Override
- public void output(OutputT output, Instant timestamp, BoundedWindow window) {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public <T> void output(
- TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
- throw new UnsupportedOperationException();
- }
- };
- }
-
- @Override
- public void startBundle() {
- doFnInvoker.invokeStartBundle(startBundleContext);
- }
-
- @Override
- public void processElement(WindowedValue<KV<InputT, RestrictionT>> elem) {
- processElementTyped(elem);
- }
-
- private <PositionT> void processElementTyped(WindowedValue<KV<InputT, RestrictionT>> elem) {
- checkArgument(
- elem.getWindows().size() == 1,
- "SPLITTABLE_PROCESS_ELEMENTS expects its input to be in 1 window, but got %s windows",
- elem.getWindows().size());
- WindowedValue<InputT> element = elem.withValue(elem.getValue().getKey());
- BoundedWindow window = elem.getWindows().iterator().next();
- this.stateAccessor =
- new FnApiStateAccessor(
- context.pipelineOptions,
- context.ptransformId,
- context.processBundleInstructionId,
- context.tagToSideInputSpecMap,
- context.beamFnStateClient,
- context.keyCoder,
- (Coder<BoundedWindow>) context.windowCoder,
- () -> elem,
- () -> window);
- RestrictionTracker<RestrictionT, PositionT> tracker =
- doFnInvoker.invokeNewTracker(elem.getValue().getValue());
- OutputAndTimeBoundedSplittableProcessElementInvoker<InputT, OutputT, RestrictionT, PositionT>
- processElementInvoker =
- new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
- context.doFn,
- context.pipelineOptions,
- new OutputWindowedValue<OutputT>() {
- @Override
- public void outputWindowedValue(
- OutputT output,
- Instant timestamp,
- Collection<? extends BoundedWindow> windows,
- PaneInfo pane) {
- outputTo(
- mainOutputConsumers, WindowedValue.of(output, timestamp, windows, pane));
- }
-
- @Override
- public <AdditionalOutputT> void outputWindowedValue(
- TupleTag<AdditionalOutputT> tag,
- AdditionalOutputT output,
- Instant timestamp,
- Collection<? extends BoundedWindow> windows,
- PaneInfo pane) {
- Collection<FnDataReceiver<WindowedValue<AdditionalOutputT>>> consumers =
- (Collection) context.localNameToConsumer.get(tag.getId());
- if (consumers == null) {
- throw new IllegalArgumentException(
- String.format("Unknown output tag %s", tag));
- }
- outputTo(consumers, WindowedValue.of(output, timestamp, windows, pane));
- }
- },
- stateAccessor,
- executor,
- 10000,
- Duration.standardSeconds(10));
- SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, PositionT>.Result result =
- processElementInvoker.invokeProcessElement(doFnInvoker, element, tracker);
- this.stateAccessor = null;
-
- if (result.getContinuation().shouldResume()) {
- WindowedValue<KV<InputT, RestrictionT>> primary =
- element.withValue(KV.of(element.getValue(), tracker.currentRestriction()));
- WindowedValue<KV<InputT, RestrictionT>> residual =
- element.withValue(KV.of(element.getValue(), result.getResidualRestriction()));
- ByteString.Output primaryBytes = ByteString.newOutput();
- ByteString.Output residualBytes = ByteString.newOutput();
- try {
- inputCoder.encode(primary, primaryBytes);
- inputCoder.encode(residual, residualBytes);
- } catch (IOException e) {
- throw new RuntimeException(e);
- }
- BundleApplication primaryApplication =
- BundleApplication.newBuilder()
- .setTransformId(context.ptransformId)
- .setInputId(mainInputId)
- .setElement(primaryBytes.toByteString())
- .build();
- BundleApplication residualApplication =
- BundleApplication.newBuilder()
- .setTransformId(context.ptransformId)
- .setInputId(mainInputId)
- .setElement(residualBytes.toByteString())
- .build();
- context.splitListener.split(
- ImmutableList.of(primaryApplication),
- ImmutableList.of(
- DelayedBundleApplication.newBuilder()
- .setApplication(residualApplication)
- .setRequestedTimeDelay(
- Durations.fromMillis(result.getContinuation().resumeDelay().getMillis()))
- .build()));
- }
- }
-
- @Override
- public void processTimer(
- String timerId, TimeDomain timeDomain, WindowedValue<KV<Object, Timer>> input) {
- throw new UnsupportedOperationException("Timers are unsupported in a SplittableDoFn.");
- }
-
- @Override
- public void finishBundle() {
- doFnInvoker.invokeFinishBundle(finishBundleContext);
- }
-
- @Override
- public void tearDown() {
- doFnInvoker.invokeTeardown();
- }
-
- /** Outputs the given element to the specified set of consumers wrapping any exceptions. */
- private <T> void outputTo(
- Collection<FnDataReceiver<WindowedValue<T>>> consumers, WindowedValue<T> output) {
- try {
- for (FnDataReceiver<WindowedValue<T>> consumer : consumers) {
- consumer.accept(output);
- }
- } catch (Throwable t) {
- throw UserCodeException.wrap(t);
- }
- }
-}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMappingFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMappingFnRunner.java
index bdebcb4..ba9d3a2 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMappingFnRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMappingFnRunner.java
@@ -22,9 +22,8 @@
import java.util.Map;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms;
-import org.apache.beam.runners.core.construction.BeamUrns;
import org.apache.beam.runners.core.construction.PCollectionViewTranslation;
+import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.sdk.function.ThrowingFunction;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
@@ -45,7 +44,7 @@
* to associate each input with its output. The nonce is represented as an opaque set of bytes.
*/
public class WindowMappingFnRunner {
- static final String URN = BeamUrns.getUrn(StandardPTransforms.Primitives.MAP_WINDOWS);
+ static final String URN = PTransformTranslation.MAP_WINDOWS_TRANSFORM_URN;
/**
* A registrar which provides a factory to handle mapping main input windows onto side input
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMergingFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMergingFnRunner.java
index ec79163..1cdc374 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMergingFnRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/WindowMergingFnRunner.java
@@ -26,8 +26,7 @@
import java.util.Map;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.model.pipeline.v1.RunnerApi.StandardPTransforms;
-import org.apache.beam.runners.core.construction.BeamUrns;
+import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.WindowingStrategyTranslation;
import org.apache.beam.sdk.function.ThrowingFunction;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
@@ -55,7 +54,7 @@
* its output. The nonce is represented as an opaque set of bytes.
*/
public abstract class WindowMergingFnRunner<T, W extends BoundedWindow> {
- static final String URN = BeamUrns.getUrn(StandardPTransforms.Primitives.MERGE_WINDOWS);
+ static final String URN = PTransformTranslation.MERGE_WINDOWS_TRANSFORM_URN;
/**
* A registrar which provides a factory to handle merging windows based upon the {@link WindowFn}.
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/AddHarnessIdInterceptor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/AddHarnessIdInterceptor.java
index 7fec44e..f949716 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/AddHarnessIdInterceptor.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/AddHarnessIdInterceptor.java
@@ -19,10 +19,10 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ClientInterceptor;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Metadata;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Metadata.Key;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.MetadataUtils;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ClientInterceptor;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Metadata;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Metadata.Key;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.MetadataUtils;
/** A {@link ClientInterceptor} that attaches a provided SDK Harness ID to outgoing messages. */
public class AddHarnessIdInterceptor {
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 a6a0211..bfaddda 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
@@ -32,8 +32,8 @@
import org.apache.beam.sdk.fn.channel.ManagedChannelFactory;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.function.ThrowingFunction;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
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 a258e01..00882b5 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
@@ -64,8 +64,8 @@
import org.apache.beam.sdk.function.ThrowingRunnable;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.common.ReflectHelpers;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Message;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Message;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
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.collect.ArrayListMultimap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.HashMultimap;
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java
index 6a02c7d..bfa0980 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java
@@ -25,7 +25,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RegisterResponse;
import org.apache.beam.model.pipeline.v1.RunnerApi;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Message;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Message;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
index 61c4580..e21fee9 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
@@ -33,7 +33,7 @@
import org.apache.beam.sdk.fn.data.LogicalEndpoint;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiver.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiver.java
deleted file mode 100644
index f234844..0000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiver.java
+++ /dev/null
@@ -1,69 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT 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.fn.harness.data;
-
-import java.io.Closeable;
-import java.util.HashMap;
-import org.apache.beam.runners.core.metrics.LabeledMetrics;
-import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
-import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
-import org.apache.beam.runners.core.metrics.MonitoringInfoConstants.Labels;
-import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.sdk.metrics.Counter;
-import org.apache.beam.sdk.metrics.MetricsContainer;
-import org.apache.beam.sdk.metrics.MetricsEnvironment;
-import org.apache.beam.sdk.util.WindowedValue;
-
-/**
- * A wrapping {@code FnDataReceiver<WindowedValue<T>>} which counts the number of elements consumed
- * by the original {@code FnDataReceiver<WindowedValue<T>>}.
- *
- * @param <T> - The receiving type of the PTransform.
- */
-public class ElementCountFnDataReceiver<T> implements FnDataReceiver<WindowedValue<T>> {
-
- private FnDataReceiver<WindowedValue<T>> original;
- private Counter counter;
- private MetricsContainer unboundMetricContainer;
-
- public ElementCountFnDataReceiver(
- FnDataReceiver<WindowedValue<T>> original,
- String pCollection,
- MetricsContainerStepMap metricContainerRegistry) {
- this.original = original;
- HashMap<String, String> labels = new HashMap<String, String>();
- labels.put(Labels.PCOLLECTION, pCollection);
- MonitoringInfoMetricName metricName =
- MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.ELEMENT_COUNT, labels);
- this.counter = LabeledMetrics.counter(metricName);
- // Collect the metric in a metric container which is not bound to the step name.
- // This is required to count elements from impulse steps, which will produce elements outside
- // of a pTransform context.
- this.unboundMetricContainer = metricContainerRegistry.getUnboundContainer();
- }
-
- @Override
- public void accept(WindowedValue<T> input) throws Exception {
- try (Closeable close = MetricsEnvironment.scopedMetricsContainer(this.unboundMetricContainer)) {
- // Increment the counter for each window the element occurs in.
- this.counter.inc(input.getWindows().size());
- this.original.accept(input);
- }
- }
-}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/MultiplexingFnDataReceiver.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/MultiplexingFnDataReceiver.java
deleted file mode 100644
index 65f75b0..0000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/MultiplexingFnDataReceiver.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT 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.fn.harness.data;
-
-import java.util.Collection;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-
-/**
- * A {@link FnDataReceiver} which forwards all received inputs to a collection of {@link
- * FnDataReceiver receivers}.
- */
-public class MultiplexingFnDataReceiver<T> implements FnDataReceiver<T> {
- public static <T> FnDataReceiver<T> forConsumers(Collection<FnDataReceiver<T>> consumers) {
- if (consumers.size() == 1) {
- return Iterables.getOnlyElement(consumers);
- }
- return new MultiplexingFnDataReceiver<>(consumers);
- }
-
- private final Collection<FnDataReceiver<T>> consumers;
-
- private MultiplexingFnDataReceiver(Collection<FnDataReceiver<T>> consumers) {
- this.consumers = consumers;
- }
-
- @Override
- public void accept(T input) throws Exception {
- for (FnDataReceiver<T> consumer : consumers) {
- consumer.accept(input);
- }
- }
-}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java
index 80d270f..3276c46 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistry.java
@@ -17,23 +17,32 @@
*/
package org.apache.beam.fn.harness.data;
+import com.google.auto.value.AutoValue;
import java.io.Closeable;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import org.apache.beam.fn.harness.HandlesSplits;
import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
+import org.apache.beam.runners.core.metrics.LabeledMetrics;
import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants.Labels;
+import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName;
import org.apache.beam.runners.core.metrics.SimpleExecutionState;
import org.apache.beam.runners.core.metrics.SimpleStateRegistry;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.MetricsContainer;
import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.util.WindowedValue;
+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.collect.ArrayListMultimap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ListMultimap;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
/**
* The {@code PCollectionConsumerRegistry} is used to maintain a collection of consuming
@@ -43,8 +52,24 @@
*/
public class PCollectionConsumerRegistry {
- private ListMultimap<String, FnDataReceiver<WindowedValue<?>>> pCollectionIdsToConsumers;
- private Map<String, ElementCountFnDataReceiver> pCollectionIdsToWrappedConsumer;
+ /** Stores metadata about each consumer so that the appropriate metrics tracking can occur. */
+ @AutoValue
+ abstract static class ConsumerAndMetadata {
+ public static ConsumerAndMetadata forConsumer(
+ FnDataReceiver consumer, String pTransformId, SimpleExecutionState state) {
+ return new AutoValue_PCollectionConsumerRegistry_ConsumerAndMetadata(
+ consumer, pTransformId, state);
+ }
+
+ public abstract FnDataReceiver getConsumer();
+
+ public abstract String getPTransformId();
+
+ public abstract SimpleExecutionState getExecutionState();
+ }
+
+ private ListMultimap<String, ConsumerAndMetadata> pCollectionIdsToConsumers;
+ private Map<String, FnDataReceiver> pCollectionIdsToWrappedConsumer;
private MetricsContainerStepMap metricsContainerRegistry;
private ExecutionStateTracker stateTracker;
private SimpleStateRegistry executionStates = new SimpleStateRegistry();
@@ -54,7 +79,7 @@
this.metricsContainerRegistry = metricsContainerRegistry;
this.stateTracker = stateTracker;
this.pCollectionIdsToConsumers = ArrayListMultimap.create();
- this.pCollectionIdsToWrappedConsumer = new HashMap<String, ElementCountFnDataReceiver>();
+ this.pCollectionIdsToWrappedConsumer = new HashMap<>();
}
/**
@@ -77,15 +102,13 @@
// Just save these consumers for now, but package them up later with an
// ElementCountFnDataReceiver and possibly a MultiplexingFnDataReceiver
// if there are multiple consumers.
- ElementCountFnDataReceiver wrappedConsumer =
- pCollectionIdsToWrappedConsumer.getOrDefault(pCollectionId, null);
- if (wrappedConsumer != null) {
+ if (pCollectionIdsToWrappedConsumer.containsKey(pCollectionId)) {
throw new RuntimeException(
"New consumers for a pCollectionId cannot be register()-d after "
+ "calling getMultiplexingConsumer.");
}
- HashMap<String, String> labelsMetadata = new HashMap<String, String>();
+ HashMap<String, String> labelsMetadata = new HashMap<>();
labelsMetadata.put(MonitoringInfoConstants.Labels.PTRANSFORM, pTransformId);
SimpleExecutionState state =
new SimpleExecutionState(
@@ -93,20 +116,9 @@
MonitoringInfoConstants.Urns.PROCESS_BUNDLE_MSECS,
labelsMetadata);
executionStates.register(state);
- // Wrap the consumer with extra logic to set the metric container with the appropriate
- // PTransform context. This ensures that user metrics obtain the pTransform ID when they are
- // created. Also use the ExecutionStateTracker and enter an appropriate state to track the
- // Process Bundle Execution time metric.
- FnDataReceiver<WindowedValue<T>> wrapAndEnableMetricContainer =
- (WindowedValue<T> input) -> {
- MetricsContainerImpl container = metricsContainerRegistry.getContainer(pTransformId);
- try (Closeable closeable = MetricsEnvironment.scopedMetricsContainer(container)) {
- try (Closeable trackerCloseable = this.stateTracker.enterState(state)) {
- consumer.accept(input);
- }
- }
- };
- pCollectionIdsToConsumers.put(pCollectionId, (FnDataReceiver) wrapAndEnableMetricContainer);
+
+ pCollectionIdsToConsumers.put(
+ pCollectionId, ConsumerAndMetadata.forConsumer(consumer, pTransformId, state));
}
/** Reset the execution states of the registered functions. */
@@ -121,24 +133,28 @@
/**
* New consumers should not be register()-ed after calling this method. This will cause a
- * RuntimeException, as this would fail to properly wrap the late-added consumer to the
- * ElementCountFnDataReceiver.
+ * RuntimeException, as this would fail to properly wrap the late-added consumer to track metrics.
*
- * @return A single ElementCountFnDataReceiver which directly wraps all the registered consumers.
+ * @return A {@link FnDataReceiver} which directly wraps all the registered consumers.
*/
public FnDataReceiver<WindowedValue<?>> getMultiplexingConsumer(String pCollectionId) {
- ElementCountFnDataReceiver wrappedConsumer =
- pCollectionIdsToWrappedConsumer.getOrDefault(pCollectionId, null);
- if (wrappedConsumer == null) {
- List<FnDataReceiver<WindowedValue<?>>> consumers =
- pCollectionIdsToConsumers.get(pCollectionId);
- FnDataReceiver<WindowedValue<?>> consumer =
- MultiplexingFnDataReceiver.forConsumers(consumers);
- wrappedConsumer =
- new ElementCountFnDataReceiver(consumer, pCollectionId, metricsContainerRegistry);
- pCollectionIdsToWrappedConsumer.put(pCollectionId, wrappedConsumer);
- }
- return wrappedConsumer;
+ return pCollectionIdsToWrappedConsumer.computeIfAbsent(
+ pCollectionId,
+ pcId -> {
+ List<ConsumerAndMetadata> consumerAndMetadatas = pCollectionIdsToConsumers.get(pcId);
+ if (consumerAndMetadatas == null) {
+ throw new IllegalArgumentException(
+ String.format("Unknown PCollectionId %s", pCollectionId));
+ } else if (consumerAndMetadatas.size() == 1) {
+ if (consumerAndMetadatas.get(0).getConsumer() instanceof HandlesSplits) {
+ return new SplittingMetricTrackingFnDataReceiver(pcId, consumerAndMetadatas.get(0));
+ }
+ return new MetricTrackingFnDataReceiver(pcId, consumerAndMetadatas.get(0));
+ } else {
+ /* TODO(SDF), Consider supporting splitting each consumer individually. This would never come up in the existing SDF expansion, but might be useful to support fused SDF nodes. This would require dedicated delivery of the split results to each of the consumers separately. */
+ return new MultiplexingMetricTrackingFnDataReceiver(pcId, consumerAndMetadatas);
+ }
+ });
}
/** @return Execution Time MonitoringInfos based on the tracked start or finish function. */
@@ -146,11 +162,141 @@
return executionStates.getExecutionTimeMonitoringInfos();
}
+ /** @return the underlying consumers for a pCollectionId, some tests may wish to check this. */
+ @VisibleForTesting
+ public List<FnDataReceiver> getUnderlyingConsumers(String pCollectionId) {
+ return Lists.transform(
+ pCollectionIdsToConsumers.get(pCollectionId), input -> input.getConsumer());
+ }
+
/**
- * @return the number of underlying consumers for a pCollectionId, some tests may wish to check
- * this.
+ * A wrapping {@code FnDataReceiver<WindowedValue<T>>} which counts the number of elements
+ * consumed by the original {@code FnDataReceiver<WindowedValue<T>> consumer} and sets up metrics
+ * for tracking PTransform processing time.
+ *
+ * @param <T> - The receiving type of the PTransform.
*/
- public List<FnDataReceiver<WindowedValue<?>>> getUnderlyingConsumers(String pCollectionId) {
- return pCollectionIdsToConsumers.get(pCollectionId);
+ private class MetricTrackingFnDataReceiver<T> implements FnDataReceiver<WindowedValue<T>> {
+ private final FnDataReceiver<WindowedValue<T>> delegate;
+ private final String pTransformId;
+ private final SimpleExecutionState state;
+ private final Counter counter;
+ private final MetricsContainer unboundMetricContainer;
+
+ public MetricTrackingFnDataReceiver(
+ String pCollectionId, ConsumerAndMetadata consumerAndMetadata) {
+ this.delegate = consumerAndMetadata.getConsumer();
+ this.state = consumerAndMetadata.getExecutionState();
+ this.pTransformId = consumerAndMetadata.getPTransformId();
+ HashMap<String, String> labels = new HashMap<String, String>();
+ labels.put(Labels.PCOLLECTION, pCollectionId);
+ MonitoringInfoMetricName metricName =
+ MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.ELEMENT_COUNT, labels);
+ this.counter = LabeledMetrics.counter(metricName);
+ // Collect the metric in a metric container which is not bound to the step name.
+ // This is required to count elements from impulse steps, which will produce elements outside
+ // of a pTransform context.
+ this.unboundMetricContainer = metricsContainerRegistry.getUnboundContainer();
+ }
+
+ @Override
+ public void accept(WindowedValue<T> input) throws Exception {
+ try (Closeable close =
+ MetricsEnvironment.scopedMetricsContainer(this.unboundMetricContainer)) {
+ // Increment the counter for each window the element occurs in.
+ this.counter.inc(input.getWindows().size());
+
+ // Wrap the consumer with extra logic to set the metric container with the appropriate
+ // PTransform context. This ensures that user metrics obtain the pTransform ID when they are
+ // created. Also use the ExecutionStateTracker and enter an appropriate state to track the
+ // Process Bundle Execution time metric.
+ MetricsContainerImpl container = metricsContainerRegistry.getContainer(pTransformId);
+ try (Closeable closeable = MetricsEnvironment.scopedMetricsContainer(container)) {
+ try (Closeable trackerCloseable = stateTracker.enterState(state)) {
+ this.delegate.accept(input);
+ }
+ }
+ }
+ }
+ }
+
+ /**
+ * A wrapping {@code FnDataReceiver<WindowedValue<T>>} which counts the number of elements
+ * consumed by the original {@code FnDataReceiver<WindowedValue<T>> consumers} and sets up metrics
+ * for tracking PTransform processing time.
+ *
+ * @param <T> - The receiving type of the PTransform.
+ */
+ private class MultiplexingMetricTrackingFnDataReceiver<T>
+ implements FnDataReceiver<WindowedValue<T>> {
+ private final List<ConsumerAndMetadata> consumerAndMetadatas;
+ private final Counter counter;
+ private final MetricsContainer unboundMetricContainer;
+
+ public MultiplexingMetricTrackingFnDataReceiver(
+ String pCollectionId, List<ConsumerAndMetadata> consumerAndMetadatas) {
+ this.consumerAndMetadatas = consumerAndMetadatas;
+ HashMap<String, String> labels = new HashMap<String, String>();
+ labels.put(Labels.PCOLLECTION, pCollectionId);
+ MonitoringInfoMetricName metricName =
+ MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.ELEMENT_COUNT, labels);
+ this.counter = LabeledMetrics.counter(metricName);
+ // Collect the metric in a metric container which is not bound to the step name.
+ // This is required to count elements from impulse steps, which will produce elements outside
+ // of a pTransform context.
+ this.unboundMetricContainer = metricsContainerRegistry.getUnboundContainer();
+ }
+
+ @Override
+ public void accept(WindowedValue<T> input) throws Exception {
+ try (Closeable close =
+ MetricsEnvironment.scopedMetricsContainer(this.unboundMetricContainer)) {
+ // Increment the counter for each window the element occurs in.
+ this.counter.inc(input.getWindows().size());
+
+ // Wrap the consumer with extra logic to set the metric container with the appropriate
+ // PTransform context. This ensures that user metrics obtain the pTransform ID when they are
+ // created. Also use the ExecutionStateTracker and enter an appropriate state to track the
+ // Process Bundle Execution time metric.
+ for (ConsumerAndMetadata consumerAndMetadata : consumerAndMetadatas) {
+ MetricsContainerImpl container =
+ metricsContainerRegistry.getContainer(consumerAndMetadata.getPTransformId());
+ try (Closeable closeable = MetricsEnvironment.scopedMetricsContainer(container)) {
+ try (Closeable trackerCloseable =
+ stateTracker.enterState(consumerAndMetadata.getExecutionState())) {
+ consumerAndMetadata.getConsumer().accept(input);
+ }
+ }
+ }
+ }
+ }
+ }
+
+ /**
+ * A wrapping {@code FnDataReceiver<WindowedValue<T>>} which counts the number of elements
+ * consumed by the original {@code FnDataReceiver<WindowedValue<T>> consumer} and forwards split
+ * and progress requests to the original consumer.
+ *
+ * @param <T> - The receiving type of the PTransform.
+ */
+ private class SplittingMetricTrackingFnDataReceiver<T> extends MetricTrackingFnDataReceiver<T>
+ implements HandlesSplits {
+ private final HandlesSplits delegate;
+
+ public SplittingMetricTrackingFnDataReceiver(
+ String pCollection, ConsumerAndMetadata consumerAndMetadata) {
+ super(pCollection, consumerAndMetadata);
+ this.delegate = (HandlesSplits) consumerAndMetadata.getConsumer();
+ }
+
+ @Override
+ public SplitResult trySplit(double fractionOfRemainder) {
+ return delegate.trySplit(fractionOfRemainder);
+ }
+
+ @Override
+ public double getProgress() {
+ return delegate.getProgress();
+ }
}
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClient.java
index 01cf0c7..8fae554 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClient.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClient.java
@@ -21,7 +21,6 @@
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import org.apache.beam.fn.harness.control.ProcessBundleHandler;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.sdk.coders.Coder;
@@ -96,7 +95,7 @@
*
* <p>This method is NOT thread safe. This should only be invoked by a single thread, and is
* intended for use with a newly constructed QueueingBeamFnDataClient in {@link
- * ProcessBundleHandler#processBundle(InstructionRequest)}.
+ * ProcessBundleHandler#processBundle}.
*/
public void drainAndBlock() throws Exception {
while (true) {
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClient.java
index 1941a10..e17f5eb 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClient.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClient.java
@@ -46,12 +46,12 @@
import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.SdkHarnessOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Timestamp;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientCallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.ClientResponseObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Timestamp;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientCallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.ClientResponseObserver;
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.collect.ImmutableMap;
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 b3e6f64..38e2027 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
@@ -27,7 +27,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.stream.DataStreams;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
/**
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 f85b3c8..bb4a661 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
@@ -31,8 +31,8 @@
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.sdk.fn.IdGenerator;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
index 26b0dfa..1d2dc1f 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
@@ -54,7 +54,7 @@
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java
index 996b87e..d795d44 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java
@@ -22,7 +22,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.stream.DataStreams;
import org.apache.beam.sdk.transforms.Materializations.MultimapView;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
/**
* An implementation of a multimap side input that utilizes the Beam Fn State API to fetch values.
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 1ebadb5..a7cb38d 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
@@ -24,7 +24,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateGetRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Throwables;
/**
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactories.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactories.java
index 7f21991..8707c3b 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactories.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactories.java
@@ -22,7 +22,7 @@
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.options.ExperimentalOptions;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
/**
* Uses {@link PipelineOptions} to configure which underlying {@link StreamObserver} implementation
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
index e6e5297..c27400f 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/AssignWindowsRunnerTest.java
@@ -184,7 +184,7 @@
null /* pipelineOptions */,
null /* beamFnDataClient */,
null /* beamFnStateClient */,
- null /* pTransformId */,
+ "ptransform",
PTransform.newBuilder()
.putInputs("in", "input")
.putOutputs("out", "output")
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
index e8a814c..39e1f9e 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BeamFnDataReadRunnerTest.java
@@ -21,9 +21,11 @@
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.empty;
+import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;
import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyDouble;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
@@ -39,12 +41,18 @@
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
+import org.apache.beam.fn.harness.HandlesSplits.SplitResult;
import org.apache.beam.fn.harness.PTransformRunnerFactory.Registrar;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.data.MultiplexingFnDataReceiver;
import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitRequest.DesiredSplit;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse.ChannelSplit;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents;
@@ -65,7 +73,6 @@
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
@@ -203,10 +210,8 @@
when(mockBeamFnDataClient.receive(any(), any(), any(), any()))
.thenReturn(bundle1Future)
.thenReturn(bundle2Future);
- List<WindowedValue<String>> valuesA = new ArrayList<>();
- List<WindowedValue<String>> valuesB = new ArrayList<>();
- FnDataReceiver<WindowedValue<String>> consumers =
- MultiplexingFnDataReceiver.forConsumers(ImmutableList.of(valuesA::add, valuesB::add));
+ List<WindowedValue<String>> values = new ArrayList<>();
+ FnDataReceiver<WindowedValue<String>> consumers = values::add;
AtomicReference<String> bundleId = new AtomicReference<>("0");
BeamFnDataReadRunner<String> readRunner =
new BeamFnDataReadRunner<>(
@@ -245,13 +250,11 @@
readRunner.blockTillReadFinishes();
future.get();
- assertThat(valuesA, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF")));
- assertThat(valuesB, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF")));
+ assertThat(values, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF")));
// Process for bundle id 1
bundleId.set("1");
- valuesA.clear();
- valuesB.clear();
+ values.clear();
readRunner.registerInputLocation();
verify(mockBeamFnDataClient)
@@ -278,8 +281,7 @@
readRunner.blockTillReadFinishes();
future.get();
- assertThat(valuesA, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
- assertThat(valuesB, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
+ assertThat(values, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
verifyNoMoreInteractions(mockBeamFnDataClient);
}
@@ -296,4 +298,185 @@
}
fail("Expected registrar not found.");
}
+
+ @Test
+ public void testSplittingWhenNoElementsProcessed() throws Exception {
+ List<WindowedValue<String>> outputValues = new ArrayList<>();
+ BeamFnDataReadRunner<String> readRunner = createReadRunner(outputValues::add);
+
+ ProcessBundleSplitRequest request =
+ ProcessBundleSplitRequest.newBuilder()
+ .putDesiredSplits(
+ "pTransformId",
+ DesiredSplit.newBuilder()
+ .setEstimatedInputElements(10)
+ .setFractionOfRemainder(0.5)
+ .build())
+ .build();
+ ProcessBundleSplitResponse.Builder responseBuilder = ProcessBundleSplitResponse.newBuilder();
+ readRunner.split(request, responseBuilder);
+
+ ProcessBundleSplitResponse expected =
+ ProcessBundleSplitResponse.newBuilder()
+ .addChannelSplits(
+ ChannelSplit.newBuilder()
+ .setLastPrimaryElement(4)
+ .setFirstResidualElement(5)
+ .build())
+ .build();
+ assertEquals(expected, responseBuilder.build());
+
+ // Ensure that we process the correct number of elements after splitting.
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("A"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("B"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("C"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("D"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("E"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("F"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("G"));
+ assertThat(
+ outputValues,
+ contains(
+ valueInGlobalWindow("A"),
+ valueInGlobalWindow("B"),
+ valueInGlobalWindow("C"),
+ valueInGlobalWindow("D"),
+ valueInGlobalWindow("E")));
+ }
+
+ @Test
+ public void testSplittingWhenSomeElementsProcessed() throws Exception {
+ List<WindowedValue<String>> outputValues = new ArrayList<>();
+ BeamFnDataReadRunner<String> readRunner = createReadRunner(outputValues::add);
+
+ ProcessBundleSplitRequest request =
+ ProcessBundleSplitRequest.newBuilder()
+ .putDesiredSplits(
+ "pTransformId",
+ DesiredSplit.newBuilder()
+ .setEstimatedInputElements(10)
+ .setFractionOfRemainder(0.5)
+ .build())
+ .build();
+ ProcessBundleSplitResponse.Builder responseBuilder = ProcessBundleSplitResponse.newBuilder();
+
+ // Process 2 elements then split
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("A"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("B"));
+ readRunner.split(request, responseBuilder);
+
+ ProcessBundleSplitResponse expected =
+ ProcessBundleSplitResponse.newBuilder()
+ .addChannelSplits(
+ ChannelSplit.newBuilder()
+ .setLastPrimaryElement(5)
+ .setFirstResidualElement(6)
+ .build())
+ .build();
+ assertEquals(expected, responseBuilder.build());
+
+ // Ensure that we process the correct number of elements after splitting.
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("C"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("D"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("E"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("F"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("G"));
+ assertThat(
+ outputValues,
+ contains(
+ valueInGlobalWindow("A"),
+ valueInGlobalWindow("B"),
+ valueInGlobalWindow("C"),
+ valueInGlobalWindow("D"),
+ valueInGlobalWindow("E"),
+ valueInGlobalWindow("F")));
+ }
+
+ @Test
+ public void testSplittingDownstreamReceiver() throws Exception {
+ SplitResult splitResult =
+ SplitResult.of(
+ BundleApplication.newBuilder().setInputId("primary").build(),
+ DelayedBundleApplication.newBuilder()
+ .setApplication(BundleApplication.newBuilder().setInputId("residual").build())
+ .build());
+ SplittingReceiver splittingReceiver = mock(SplittingReceiver.class);
+ when(splittingReceiver.getProgress()).thenReturn(0.3);
+ when(splittingReceiver.trySplit(anyDouble())).thenReturn(splitResult);
+ BeamFnDataReadRunner<String> readRunner = createReadRunner(splittingReceiver);
+
+ ProcessBundleSplitRequest request =
+ ProcessBundleSplitRequest.newBuilder()
+ .putDesiredSplits(
+ "pTransformId",
+ DesiredSplit.newBuilder()
+ .setEstimatedInputElements(10)
+ .setFractionOfRemainder(0.05)
+ .build())
+ .build();
+ ProcessBundleSplitResponse.Builder responseBuilder = ProcessBundleSplitResponse.newBuilder();
+
+ // We will be "processing" the 'C' element, aka 2nd index
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("A"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("B"));
+ readRunner.forwardElementToConsumer(valueInGlobalWindow("C"));
+ readRunner.split(request, responseBuilder);
+
+ ProcessBundleSplitResponse expected =
+ ProcessBundleSplitResponse.newBuilder()
+ .addPrimaryRoots(splitResult.getPrimaryRoot())
+ .addResidualRoots(splitResult.getResidualRoot())
+ .addChannelSplits(
+ ChannelSplit.newBuilder()
+ .setLastPrimaryElement(1)
+ .setFirstResidualElement(3)
+ .build())
+ .build();
+ assertEquals(expected, responseBuilder.build());
+ }
+
+ private abstract static class SplittingReceiver
+ implements FnDataReceiver<WindowedValue<String>>, HandlesSplits {}
+
+ private BeamFnDataReadRunner<String> createReadRunner(
+ FnDataReceiver<WindowedValue<String>> consumer) throws Exception {
+ String bundleId = "57";
+
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ String localOutputId = "outputPC";
+ String pTransformId = "pTransformId";
+ consumers.register(localOutputId, pTransformId, consumer);
+ PTransformFunctionRegistry startFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "start");
+ PTransformFunctionRegistry finishFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "finish");
+ List<ThrowingRunnable> teardownFunctions = new ArrayList<>();
+
+ RunnerApi.PTransform pTransform =
+ RemoteGrpcPortRead.readFromPort(PORT_SPEC, localOutputId).toPTransform();
+
+ return new BeamFnDataReadRunner.Factory<String>()
+ .createRunnerForPTransform(
+ PipelineOptionsFactory.create(),
+ mockBeamFnDataClient,
+ null /* beamFnStateClient */,
+ pTransformId,
+ pTransform,
+ Suppliers.ofInstance(bundleId)::get,
+ ImmutableMap.of(
+ localOutputId,
+ RunnerApi.PCollection.newBuilder().setCoderId(ELEMENT_CODER_SPEC_ID).build()),
+ COMPONENTS.getCodersMap(),
+ COMPONENTS.getWindowingStrategiesMap(),
+ consumers,
+ startFunctionRegistry,
+ finishFunctionRegistry,
+ teardownFunctions::add,
+ null /* splitListener */);
+ }
}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java
index bc31e9c..0f461d2 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/BoundedSourceRunnerTest.java
@@ -44,7 +44,7 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
index 4533283..81fa48d 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
@@ -25,6 +25,7 @@
import static org.hamcrest.Matchers.hasSize;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
@@ -34,16 +35,22 @@
import java.util.ArrayList;
import java.util.List;
import java.util.ServiceLoader;
+import org.apache.beam.fn.harness.control.BundleSplitListener;
import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry;
import org.apache.beam.fn.harness.data.PTransformFunctionRegistry;
import org.apache.beam.fn.harness.state.FakeBeamFnStateClient;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey;
import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Environment;
import org.apache.beam.runners.core.construction.PTransformTranslation;
+import org.apache.beam.runners.core.construction.ParDoTranslation;
import org.apache.beam.runners.core.construction.PipelineTranslation;
import org.apache.beam.runners.core.construction.SdkComponents;
+import org.apache.beam.runners.core.construction.graph.ProtoOverrides;
+import org.apache.beam.runners.core.construction.graph.SplittableParDoExpander;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
import org.apache.beam.runners.core.metrics.MetricUpdates.MetricUpdate;
import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
@@ -54,6 +61,7 @@
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.function.ThrowingRunnable;
+import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.MetricKey;
import org.apache.beam.sdk.metrics.MetricName;
@@ -76,6 +84,8 @@
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -90,7 +100,7 @@
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Suppliers;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
@@ -277,7 +287,6 @@
.put(bagUserStateKey("combine", "Y"), encode("Y1Y2"))
.build(),
fakeClient.getData());
- mainOutputValues.clear();
}
/** Produces a bag user {@link StateKey} for the test PTransform id in the global window. */
@@ -446,7 +455,6 @@
// Assert that state data did not change
assertEquals(stateData, fakeClient.getData());
- mainOutputValues.clear();
}
private static class TestSideInputIsAccessibleForDownstreamCallersDoFn
@@ -736,7 +744,7 @@
@TimerId("processing") Timer processingTimeTimer) {
context.output("main" + context.element().getKey() + Iterables.toString(bagState.read()));
bagState.add(context.element().getValue());
- eventTimeTimer.set(context.timestamp().plus(1L));
+ eventTimeTimer.withOutputTimestamp(context.timestamp()).set(context.timestamp().plus(1L));
processingTimeTimer.offset(Duration.millis(2L));
processingTimeTimer.setRelative();
}
@@ -749,7 +757,9 @@
@TimerId("processing") Timer processingTimeTimer) {
context.output("event" + Iterables.toString(bagState.read()));
bagState.add("event");
- eventTimeTimer.set(context.timestamp().plus(11L));
+ eventTimeTimer
+ .withOutputTimestamp(context.timestamp())
+ .set(context.fireTimestamp().plus(11L));
processingTimeTimer.offset(Duration.millis(12L));
processingTimeTimer.setRelative();
}
@@ -762,7 +772,7 @@
@TimerId("processing") Timer processingTimeTimer) {
context.output("processing" + Iterables.toString(bagState.read()));
bagState.add("processing");
- eventTimeTimer.set(context.timestamp().plus(21L));
+ eventTimeTimer.withOutputTimestamp(context.timestamp()).set(context.timestamp().plus(21L));
processingTimeTimer.offset(Duration.millis(22L));
processingTimeTimer.setRelative();
}
@@ -921,9 +931,9 @@
timerInGlobalWindow("Y", new Instant(1100L), new Instant(1101L)),
timerInGlobalWindow("X", new Instant(1200L), new Instant(1201L)),
timerInGlobalWindow("Y", new Instant(1300L), new Instant(1301L)),
- timerInGlobalWindow("A", new Instant(1400L), new Instant(1411L)),
- timerInGlobalWindow("B", new Instant(1500L), new Instant(1511L)),
- timerInGlobalWindow("A", new Instant(1600L), new Instant(1611L)),
+ timerInGlobalWindow("A", new Instant(1400L), new Instant(2411L)),
+ timerInGlobalWindow("B", new Instant(1500L), new Instant(2511L)),
+ timerInGlobalWindow("A", new Instant(1600L), new Instant(2611L)),
timerInGlobalWindow("X", new Instant(1700L), new Instant(1721L)),
timerInGlobalWindow("C", new Instant(1800L), new Instant(1821L)),
timerInGlobalWindow("B", new Instant(1900L), new Instant(1921L))));
@@ -957,7 +967,6 @@
.put(bagUserStateKey("bag", "C"), encode("C0", "processing"))
.build(),
fakeClient.getData());
- mainOutputValues.clear();
}
private <T> WindowedValue<T> valueInWindow(T value, BoundedWindow window) {
@@ -1020,4 +1029,350 @@
}
fail("Expected registrar not found.");
}
+
+ static class TestSplittableDoFn extends DoFn<String, String> {
+ private final PCollectionView<String> singletonSideInput;
+
+ private TestSplittableDoFn(PCollectionView<String> singletonSideInput) {
+ this.singletonSideInput = singletonSideInput;
+ }
+
+ @ProcessElement
+ public ProcessContinuation processElement(
+ ProcessContext context, RestrictionTracker<OffsetRange, Long> tracker) {
+ int upperBound = Integer.parseInt(context.sideInput(singletonSideInput));
+ for (int i = 0; i < upperBound; ++i) {
+ if (tracker.tryClaim((long) i)) {
+ context.output(context.element() + ":" + i);
+ }
+ }
+ if (tracker.currentRestriction().getTo() > upperBound) {
+ return ProcessContinuation.resume().withResumeDelay(Duration.millis(42L));
+ } else {
+ return ProcessContinuation.stop();
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange restriction(String element) {
+ return new OffsetRange(0, Integer.parseInt(element));
+ }
+
+ @NewTracker
+ public RestrictionTracker<OffsetRange, Long> newTracker(OffsetRange restriction) {
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @SplitRestriction
+ public void splitRange(
+ String element, OffsetRange range, OutputReceiver<OffsetRange> receiver) {
+ receiver.output(new OffsetRange(range.getFrom(), (range.getFrom() + range.getTo()) / 2));
+ receiver.output(new OffsetRange((range.getFrom() + range.getTo()) / 2, range.getTo()));
+ }
+ }
+
+ @Test
+ public void testProcessElementForSizedElementAndRestriction() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ PCollectionView<String> singletonSideInputView = valuePCollection.apply(View.asSingleton());
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID,
+ ParDo.of(new TestSplittableDoFn(singletonSideInputView))
+ .withSideInputs(singletonSideInputView));
+
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createSizedReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(
+ PTransformTranslation
+ .SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ String outputPCollectionId = pTransform.getOutputsOrThrow("output");
+
+ ImmutableMap<StateKey, ByteString> stateData =
+ ImmutableMap.of(
+ multimapSideInputKey(singletonSideInputView.getTagInternal().getId(), ByteString.EMPTY),
+ encode("3"));
+
+ FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient(stateData);
+
+ List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ consumers.register(
+ outputPCollectionId,
+ TEST_TRANSFORM_ID,
+ (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>) mainOutputValues::add);
+ PTransformFunctionRegistry startFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "start");
+ PTransformFunctionRegistry finishFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "finish");
+ List<ThrowingRunnable> teardownFunctions = new ArrayList<>();
+ List<BundleApplication> primarySplits = new ArrayList<>();
+ List<DelayedBundleApplication> residualSplits = new ArrayList<>();
+
+ new FnApiDoFnRunner.Factory<>()
+ .createRunnerForPTransform(
+ PipelineOptionsFactory.create(),
+ null /* beamFnDataClient */,
+ fakeClient,
+ TEST_TRANSFORM_ID,
+ pTransform,
+ Suppliers.ofInstance("57L")::get,
+ pProto.getComponents().getPcollectionsMap(),
+ pProto.getComponents().getCodersMap(),
+ pProto.getComponents().getWindowingStrategiesMap(),
+ consumers,
+ startFunctionRegistry,
+ finishFunctionRegistry,
+ teardownFunctions::add,
+ new BundleSplitListener() {
+ @Override
+ public void split(
+ List<BundleApplication> primaryRoots,
+ List<DelayedBundleApplication> residualRoots) {
+ primarySplits.addAll(primaryRoots);
+ residualSplits.addAll(residualRoots);
+ }
+ });
+
+ Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
+ mainOutputValues.clear();
+
+ assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId, outputPCollectionId));
+
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ consumers.getMultiplexingConsumer(inputPCollectionId);
+ mainInput.accept(valueInGlobalWindow(KV.of(KV.of("5", new OffsetRange(0, 5)), 5.0)));
+ BundleApplication primaryRoot = Iterables.getOnlyElement(primarySplits);
+ DelayedBundleApplication residualRoot = Iterables.getOnlyElement(residualSplits);
+ assertEquals(ParDoTranslation.getMainInputName(pTransform), primaryRoot.getInputId());
+ assertEquals(TEST_TRANSFORM_ID, primaryRoot.getTransformId());
+ assertEquals(
+ ParDoTranslation.getMainInputName(pTransform), residualRoot.getApplication().getInputId());
+ assertEquals(TEST_TRANSFORM_ID, residualRoot.getApplication().getTransformId());
+ primarySplits.clear();
+ residualSplits.clear();
+
+ mainInput.accept(valueInGlobalWindow(KV.of(KV.of("2", new OffsetRange(0, 2)), 2.0)));
+ assertThat(
+ mainOutputValues,
+ contains(
+ valueInGlobalWindow("5:0"),
+ valueInGlobalWindow("5:1"),
+ valueInGlobalWindow("5:2"),
+ valueInGlobalWindow("2:0"),
+ valueInGlobalWindow("2:1")));
+ assertTrue(primarySplits.isEmpty());
+ assertTrue(residualSplits.isEmpty());
+ mainOutputValues.clear();
+
+ Iterables.getOnlyElement(finishFunctionRegistry.getFunctions()).run();
+ assertThat(mainOutputValues, empty());
+
+ Iterables.getOnlyElement(teardownFunctions).run();
+ assertThat(mainOutputValues, empty());
+
+ // Assert that state data did not change
+ assertEquals(stateData, fakeClient.getData());
+ }
+
+ @Test
+ public void testProcessElementForPairWithRestriction() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ PCollectionView<String> singletonSideInputView = valuePCollection.apply(View.asSingleton());
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID,
+ ParDo.of(new TestSplittableDoFn(singletonSideInputView))
+ .withSideInputs(singletonSideInputView));
+
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createSizedReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
+
+ FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient(ImmutableMap.of());
+
+ List<WindowedValue<KV<String, OffsetRange>>> mainOutputValues = new ArrayList<>();
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ consumers.register(outputPCollectionId, TEST_TRANSFORM_ID, ((List) mainOutputValues)::add);
+ PTransformFunctionRegistry startFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "start");
+ PTransformFunctionRegistry finishFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "finish");
+ List<ThrowingRunnable> teardownFunctions = new ArrayList<>();
+
+ new FnApiDoFnRunner.Factory<>()
+ .createRunnerForPTransform(
+ PipelineOptionsFactory.create(),
+ null /* beamFnDataClient */,
+ fakeClient,
+ TEST_TRANSFORM_ID,
+ pTransform,
+ Suppliers.ofInstance("57L")::get,
+ pProto.getComponents().getPcollectionsMap(),
+ pProto.getComponents().getCodersMap(),
+ pProto.getComponents().getWindowingStrategiesMap(),
+ consumers,
+ startFunctionRegistry,
+ finishFunctionRegistry,
+ teardownFunctions::add,
+ null /* bundleSplitListener */);
+
+ Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
+ mainOutputValues.clear();
+
+ assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId, outputPCollectionId));
+
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ consumers.getMultiplexingConsumer(inputPCollectionId);
+ mainInput.accept(valueInGlobalWindow("5"));
+ mainInput.accept(valueInGlobalWindow("2"));
+ assertThat(
+ mainOutputValues,
+ contains(
+ valueInGlobalWindow(KV.of("5", new OffsetRange(0, 5))),
+ valueInGlobalWindow(KV.of("2", new OffsetRange(0, 2)))));
+ mainOutputValues.clear();
+
+ Iterables.getOnlyElement(finishFunctionRegistry.getFunctions()).run();
+ assertThat(mainOutputValues, empty());
+
+ Iterables.getOnlyElement(teardownFunctions).run();
+ assertThat(mainOutputValues, empty());
+ }
+
+ @Test
+ public void testProcessElementForSplitAndSizeRestriction() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<String> valuePCollection = p.apply(Create.of("unused"));
+ PCollectionView<String> singletonSideInputView = valuePCollection.apply(View.asSingleton());
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID,
+ ParDo.of(new TestSplittableDoFn(singletonSideInputView))
+ .withSideInputs(singletonSideInputView));
+
+ RunnerApi.Pipeline pProto =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ PipelineTranslation.toProto(p, SdkComponents.create(p.getOptions()), true),
+ SplittableParDoExpander.createSizedReplacement());
+ String expandedTransformId =
+ Iterables.find(
+ pProto.getComponents().getTransformsMap().entrySet(),
+ entry ->
+ entry
+ .getValue()
+ .getSpec()
+ .getUrn()
+ .equals(
+ PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN)
+ && entry.getValue().getUniqueName().contains(TEST_TRANSFORM_ID))
+ .getKey();
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(expandedTransformId);
+ String inputPCollectionId =
+ pTransform.getInputsOrThrow(ParDoTranslation.getMainInputName(pTransform));
+ String outputPCollectionId = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
+
+ FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient(ImmutableMap.of());
+
+ List<WindowedValue<KV<KV<String, OffsetRange>, Double>>> mainOutputValues = new ArrayList<>();
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ consumers.register(outputPCollectionId, TEST_TRANSFORM_ID, ((List) mainOutputValues)::add);
+ PTransformFunctionRegistry startFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "start");
+ PTransformFunctionRegistry finishFunctionRegistry =
+ new PTransformFunctionRegistry(
+ mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "finish");
+ List<ThrowingRunnable> teardownFunctions = new ArrayList<>();
+
+ new FnApiDoFnRunner.Factory<>()
+ .createRunnerForPTransform(
+ PipelineOptionsFactory.create(),
+ null /* beamFnDataClient */,
+ fakeClient,
+ TEST_TRANSFORM_ID,
+ pTransform,
+ Suppliers.ofInstance("57L")::get,
+ pProto.getComponents().getPcollectionsMap(),
+ pProto.getComponents().getCodersMap(),
+ pProto.getComponents().getWindowingStrategiesMap(),
+ consumers,
+ startFunctionRegistry,
+ finishFunctionRegistry,
+ teardownFunctions::add,
+ null /* bundleSplitListener */);
+
+ Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
+ mainOutputValues.clear();
+
+ assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId, outputPCollectionId));
+
+ FnDataReceiver<WindowedValue<?>> mainInput =
+ consumers.getMultiplexingConsumer(inputPCollectionId);
+ mainInput.accept(valueInGlobalWindow(KV.of("5", new OffsetRange(0, 5))));
+ mainInput.accept(valueInGlobalWindow(KV.of("2", new OffsetRange(0, 2))));
+ assertThat(
+ mainOutputValues,
+ contains(
+ valueInGlobalWindow(KV.of(KV.of("5", new OffsetRange(0, 2)), 2.0)),
+ valueInGlobalWindow(KV.of(KV.of("5", new OffsetRange(2, 5)), 3.0)),
+ valueInGlobalWindow(KV.of(KV.of("2", new OffsetRange(0, 1)), 1.0)),
+ valueInGlobalWindow(KV.of(KV.of("2", new OffsetRange(1, 2)), 1.0))));
+ mainOutputValues.clear();
+
+ Iterables.getOnlyElement(finishFunctionRegistry.getFunctions()).run();
+ assertThat(mainOutputValues, empty());
+
+ Iterables.getOnlyElement(teardownFunctions).run();
+ assertThat(mainOutputValues, empty());
+ }
}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnHarnessTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnHarnessTest.java
index eaf0b07..90d6848 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnHarnessTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnHarnessTest.java
@@ -42,10 +42,10 @@
import org.apache.beam.sdk.harness.JvmInitializer;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.TextFormat;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.TextFormat;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
import org.junit.Test;
import org.junit.runner.RunWith;
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 36a4779..5adb001 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
@@ -43,10 +43,10 @@
import org.apache.beam.sdk.fn.test.InProcessManagedChannelFactory;
import org.apache.beam.sdk.fn.test.TestStreams;
import org.apache.beam.sdk.function.ThrowingFunction;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
import org.junit.Test;
import org.junit.runner.RunWith;
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 a7a233d..a3e959b 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
@@ -77,8 +77,8 @@
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Message;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Message;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Multimap;
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java
index deb6218..7e5dda9 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java
@@ -46,13 +46,13 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataInboundObserverTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataInboundObserverTest.java
index aa45df6..f51f006 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataInboundObserverTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataInboundObserverTest.java
@@ -36,7 +36,7 @@
import org.apache.beam.sdk.fn.data.InboundDataClient;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiverTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiverTest.java
deleted file mode 100644
index ace16c6..0000000
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/ElementCountFnDataReceiverTest.java
+++ /dev/null
@@ -1,98 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT 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.fn.harness.data;
-
-import static junit.framework.TestCase.assertEquals;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.withSettings;
-import static org.powermock.api.mockito.PowerMockito.mockStatic;
-
-import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
-import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
-import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
-import org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.sdk.metrics.MetricsEnvironment;
-import org.apache.beam.sdk.util.WindowedValue;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.powermock.api.mockito.PowerMockito;
-import org.powermock.core.classloader.annotations.PrepareForTest;
-import org.powermock.modules.junit4.PowerMockRunner;
-
-/** Tests for {@link ElementCountFnDataReceiver}. */
-@RunWith(PowerMockRunner.class)
-@PrepareForTest(MetricsEnvironment.class)
-public class ElementCountFnDataReceiverTest {
- /**
- * Test that the elements are counted, and a MonitoringInfo can be extracted from a
- * metricsContainer, if it is in scope.
- *
- * @throws Exception
- */
- @Test
- public void testCountsElements() throws Exception {
- final String pCollectionA = "pCollectionA";
-
- MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
-
- FnDataReceiver<WindowedValue<String>> consumer = mock(FnDataReceiver.class);
- ElementCountFnDataReceiver<String> wrapperConsumer =
- new ElementCountFnDataReceiver(consumer, pCollectionA, metricsContainerRegistry);
- WindowedValue<String> element = WindowedValue.valueInGlobalWindow("elem");
- int numElements = 20;
- for (int i = 0; i < numElements; i++) {
- wrapperConsumer.accept(element);
- }
- verify(consumer, times(numElements)).accept(element);
-
- SimpleMonitoringInfoBuilder builder = new SimpleMonitoringInfoBuilder();
- builder.setUrn(MonitoringInfoConstants.Urns.ELEMENT_COUNT);
- builder.setLabel(MonitoringInfoConstants.Labels.PCOLLECTION, pCollectionA);
- builder.setInt64Value(numElements);
- MonitoringInfo expected = builder.build();
-
- // Clear the timestamp before comparison.
- MonitoringInfo first = metricsContainerRegistry.getMonitoringInfos().iterator().next();
- MonitoringInfo result = SimpleMonitoringInfoBuilder.copyAndClearTimestamp(first);
- assertEquals(expected, result);
- }
-
- @Test
- public void testScopedMetricContainerInvokedUponAccept() throws Exception {
- mockStatic(MetricsEnvironment.class, withSettings().verboseLogging());
- final String pCollectionA = "pCollectionA";
-
- MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
-
- FnDataReceiver<WindowedValue<String>> consumer =
- mock(FnDataReceiver.class, withSettings().verboseLogging());
- ElementCountFnDataReceiver<String> wrapperConsumer =
- new ElementCountFnDataReceiver(consumer, pCollectionA, metricsContainerRegistry);
- WindowedValue<String> element = WindowedValue.valueInGlobalWindow("elem");
- wrapperConsumer.accept(element);
-
- verify(consumer, times(1)).accept(element);
-
- // Verify that static scopedMetricsContainer is called with unbound container.
- PowerMockito.verifyStatic(MetricsEnvironment.class, times(1));
- MetricsEnvironment.scopedMetricsContainer(metricsContainerRegistry.getUnboundContainer());
- }
-}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/MultiplexingFnDataReceiverTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/MultiplexingFnDataReceiverTest.java
deleted file mode 100644
index f3e21c2..0000000
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/MultiplexingFnDataReceiverTest.java
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT 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.fn.harness.data;
-
-import static org.hamcrest.Matchers.contains;
-import static org.hamcrest.Matchers.containsInAnyOrder;
-import static org.junit.Assert.assertThat;
-
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-import org.apache.beam.sdk.fn.data.FnDataReceiver;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** Tests for {@link MultiplexingFnDataReceiver}. */
-@RunWith(JUnit4.class)
-public class MultiplexingFnDataReceiverTest {
- @Rule public ExpectedException thrown = ExpectedException.none();
-
- @Test
- public void singleConsumer() throws Exception {
- List<String> consumer = new ArrayList<>();
- FnDataReceiver<String> multiplexer =
- MultiplexingFnDataReceiver.forConsumers(
- ImmutableList.<FnDataReceiver<String>>of(consumer::add));
-
- multiplexer.accept("foo");
- multiplexer.accept("bar");
-
- assertThat(consumer, contains("foo", "bar"));
- }
-
- @Test
- public void singleConsumerException() throws Exception {
- String message = "my_exception";
- FnDataReceiver<Integer> multiplexer =
- MultiplexingFnDataReceiver.forConsumers(
- ImmutableList.<FnDataReceiver<Integer>>of(
- (Integer i) -> {
- if (i > 1) {
- throw new Exception(message);
- }
- }));
-
- multiplexer.accept(0);
- multiplexer.accept(1);
- thrown.expectMessage(message);
- thrown.expect(Exception.class);
- multiplexer.accept(2);
- }
-
- @Test
- public void multipleConsumers() throws Exception {
- List<String> consumer = new ArrayList<>();
- Set<String> otherConsumer = new HashSet<>();
- FnDataReceiver<String> multiplexer =
- MultiplexingFnDataReceiver.forConsumers(
- ImmutableList.<FnDataReceiver<String>>of(consumer::add, otherConsumer::add));
-
- multiplexer.accept("foo");
- multiplexer.accept("bar");
- multiplexer.accept("foo");
-
- assertThat(consumer, contains("foo", "bar", "foo"));
- assertThat(otherConsumer, containsInAnyOrder("foo", "bar"));
- }
-
- @Test
- public void multipleConsumersException() throws Exception {
- String message = "my_exception";
- List<Integer> consumer = new ArrayList<>();
- FnDataReceiver<Integer> multiplexer =
- MultiplexingFnDataReceiver.forConsumers(
- ImmutableList.<FnDataReceiver<Integer>>of(
- consumer::add,
- (Integer i) -> {
- if (i > 1) {
- throw new Exception(message);
- }
- }));
-
- multiplexer.accept(0);
- multiplexer.accept(1);
- assertThat(consumer, containsInAnyOrder(0, 1));
-
- thrown.expectMessage(message);
- thrown.expect(Exception.class);
- multiplexer.accept(2);
- }
-}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java
index e8b377b..dac1fbe 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/PCollectionConsumerRegistryTest.java
@@ -17,18 +17,30 @@
*/
package org.apache.beam.fn.harness.data;
+import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow;
import static org.hamcrest.Matchers.contains;
+import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.withSettings;
import static org.powermock.api.mockito.PowerMockito.mockStatic;
+import org.apache.beam.fn.harness.HandlesSplits;
+import org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo;
import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants.Labels;
+import org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
@@ -44,6 +56,72 @@
@Rule public ExpectedException expectedException = ExpectedException.none();
+ @Test
+ public void singleConsumer() throws Exception {
+ final String pCollectionA = "pCollectionA";
+ final String pTransformIdA = "pTransformIdA";
+
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ FnDataReceiver<WindowedValue<String>> consumerA1 = mock(FnDataReceiver.class);
+
+ consumers.register(pCollectionA, pTransformIdA, consumerA1);
+
+ FnDataReceiver<WindowedValue<String>> wrapperConsumer =
+ (FnDataReceiver<WindowedValue<String>>)
+ (FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
+
+ WindowedValue<String> element = valueInGlobalWindow("elem");
+ int numElements = 20;
+ for (int i = 0; i < numElements; i++) {
+ wrapperConsumer.accept(element);
+ }
+
+ // Check that the underlying consumers are each invoked per element.
+ verify(consumerA1, times(numElements)).accept(element);
+ assertThat(consumers.keySet(), contains(pCollectionA));
+
+ SimpleMonitoringInfoBuilder builder = new SimpleMonitoringInfoBuilder();
+ builder.setUrn(MonitoringInfoConstants.Urns.ELEMENT_COUNT);
+ builder.setLabel(MonitoringInfoConstants.Labels.PCOLLECTION, pCollectionA);
+ builder.setInt64Value(numElements);
+ MonitoringInfo expected = builder.build();
+
+ // Clear the timestamp before comparison.
+ MonitoringInfo pCollectionCount =
+ Iterables.find(
+ metricsContainerRegistry.getMonitoringInfos(),
+ monitoringInfo -> monitoringInfo.containsLabels(Labels.PCOLLECTION));
+ MonitoringInfo result = SimpleMonitoringInfoBuilder.copyAndClearTimestamp(pCollectionCount);
+ assertEquals(expected, result);
+ }
+
+ @Test
+ public void singleConsumerException() throws Exception {
+ final String pCollectionA = "pCollectionA";
+ final String pTransformId = "pTransformId";
+ final String message = "testException";
+
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ FnDataReceiver<WindowedValue<String>> consumer = mock(FnDataReceiver.class);
+
+ consumers.register(pCollectionA, pTransformId, consumer);
+
+ FnDataReceiver<WindowedValue<String>> wrapperConsumer =
+ (FnDataReceiver<WindowedValue<String>>)
+ (FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
+ doThrow(new Exception(message)).when(consumer).accept(any());
+
+ expectedException.expectMessage(message);
+ expectedException.expect(Exception.class);
+ wrapperConsumer.accept(valueInGlobalWindow("elem"));
+ }
+
/**
* Test that the counter increments only once when multiple consumers of same pCollection read the
* same element.
@@ -51,7 +129,54 @@
@Test
public void multipleConsumersSamePCollection() throws Exception {
final String pCollectionA = "pCollectionA";
+ final String pTransformIdA = "pTransformIdA";
+ final String pTransformIdB = "pTransformIdB";
+
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ FnDataReceiver<WindowedValue<String>> consumerA1 = mock(FnDataReceiver.class);
+ FnDataReceiver<WindowedValue<String>> consumerA2 = mock(FnDataReceiver.class);
+
+ consumers.register(pCollectionA, pTransformIdA, consumerA1);
+ consumers.register(pCollectionA, pTransformIdB, consumerA2);
+
+ FnDataReceiver<WindowedValue<String>> wrapperConsumer =
+ (FnDataReceiver<WindowedValue<String>>)
+ (FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
+
+ WindowedValue<String> element = valueInGlobalWindow("elem");
+ int numElements = 20;
+ for (int i = 0; i < numElements; i++) {
+ wrapperConsumer.accept(element);
+ }
+
+ // Check that the underlying consumers are each invoked per element.
+ verify(consumerA1, times(numElements)).accept(element);
+ verify(consumerA2, times(numElements)).accept(element);
+ assertThat(consumers.keySet(), contains(pCollectionA));
+
+ SimpleMonitoringInfoBuilder builder = new SimpleMonitoringInfoBuilder();
+ builder.setUrn(MonitoringInfoConstants.Urns.ELEMENT_COUNT);
+ builder.setLabel(MonitoringInfoConstants.Labels.PCOLLECTION, pCollectionA);
+ builder.setInt64Value(numElements);
+ MonitoringInfo expected = builder.build();
+
+ // Clear the timestamp before comparison.
+ MonitoringInfo pCollectionCount =
+ Iterables.find(
+ metricsContainerRegistry.getMonitoringInfos(),
+ monitoringInfo -> monitoringInfo.containsLabels(Labels.PCOLLECTION));
+ MonitoringInfo result = SimpleMonitoringInfoBuilder.copyAndClearTimestamp(pCollectionCount);
+ assertEquals(expected, result);
+ }
+
+ @Test
+ public void multipleConsumersSamePCollectionException() throws Exception {
+ final String pCollectionA = "pCollectionA";
final String pTransformId = "pTransformId";
+ final String message = "testException";
MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
PCollectionConsumerRegistry consumers =
@@ -66,17 +191,11 @@
FnDataReceiver<WindowedValue<String>> wrapperConsumer =
(FnDataReceiver<WindowedValue<String>>)
(FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
+ doThrow(new Exception(message)).when(consumerA2).accept(any());
- WindowedValue<String> element = WindowedValue.valueInGlobalWindow("elem");
- int numElements = 20;
- for (int i = 0; i < numElements; i++) {
- wrapperConsumer.accept(element);
- }
-
- // Check that the underlying consumers are each invoked per element.
- verify(consumerA1, times(numElements)).accept(element);
- verify(consumerA2, times(numElements)).accept(element);
- assertThat(consumers.keySet(), contains(pCollectionA));
+ expectedException.expectMessage(message);
+ expectedException.expect(Exception.class);
+ wrapperConsumer.accept(valueInGlobalWindow("elem"));
}
@Test
@@ -118,7 +237,7 @@
(FnDataReceiver<WindowedValue<String>>)
(FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
- WindowedValue<String> element = WindowedValue.valueInGlobalWindow("elem");
+ WindowedValue<String> element = valueInGlobalWindow("elem");
wrapperConsumer.accept(element);
// Verify that static scopedMetricsContainer is called with pTransformA's container.
@@ -129,4 +248,61 @@
PowerMockito.verifyStatic(MetricsEnvironment.class, times(1));
MetricsEnvironment.scopedMetricsContainer(metricsContainerRegistry.getContainer("pTransformB"));
}
+
+ @Test
+ public void testScopedMetricContainerInvokedUponAccept() throws Exception {
+ mockStatic(MetricsEnvironment.class, withSettings().verboseLogging());
+ final String pCollectionA = "pCollectionA";
+ final String pTransformIdA = "pTransformIdA";
+
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ FnDataReceiver<WindowedValue<String>> consumer =
+ mock(FnDataReceiver.class, withSettings().verboseLogging());
+
+ consumers.register(pCollectionA, pTransformIdA, consumer);
+
+ FnDataReceiver<WindowedValue<String>> wrapperConsumer =
+ (FnDataReceiver<WindowedValue<String>>)
+ (FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
+
+ WindowedValue<String> element = WindowedValue.valueInGlobalWindow("elem");
+ wrapperConsumer.accept(element);
+
+ verify(consumer, times(1)).accept(element);
+
+ // Verify that static scopedMetricsContainer is called with unbound container.
+ PowerMockito.verifyStatic(MetricsEnvironment.class, times(1));
+ MetricsEnvironment.scopedMetricsContainer(metricsContainerRegistry.getUnboundContainer());
+ }
+
+ @Test
+ public void testHandlesSplitsPassedToOriginalConsumer() throws Exception {
+ final String pCollectionA = "pCollectionA";
+ final String pTransformIdA = "pTransformIdA";
+
+ MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
+ PCollectionConsumerRegistry consumers =
+ new PCollectionConsumerRegistry(
+ metricsContainerRegistry, mock(ExecutionStateTracker.class));
+ SplittingReceiver consumerA1 = mock(SplittingReceiver.class);
+
+ consumers.register(pCollectionA, pTransformIdA, consumerA1);
+
+ FnDataReceiver<WindowedValue<String>> wrapperConsumer =
+ (FnDataReceiver<WindowedValue<String>>)
+ (FnDataReceiver) consumers.getMultiplexingConsumer(pCollectionA);
+
+ assertTrue(wrapperConsumer instanceof HandlesSplits);
+
+ ((HandlesSplits) wrapperConsumer).getProgress();
+ verify(consumerA1).getProgress();
+
+ ((HandlesSplits) wrapperConsumer).trySplit(0.3);
+ verify(consumerA1).trySplit(0.3);
+ }
+
+ private abstract static class SplittingReceiver<T> implements FnDataReceiver<T>, HandlesSplits {}
}
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClientTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClientTest.java
index 8bcacfa..094d9f7 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClientTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClientTest.java
@@ -47,13 +47,13 @@
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClientTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClientTest.java
index e3a4266..dc49275 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClientTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/logging/BeamFnLoggingClientTest.java
@@ -37,14 +37,14 @@
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.sdk.fn.test.TestStreams;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Timestamp;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.Timestamp;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java
index 5b01c0f..dbf9885 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java
@@ -25,7 +25,7 @@
import java.io.IOException;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.junit.Rule;
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 e1feac1..e8f1780 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
@@ -36,14 +36,14 @@
import org.apache.beam.sdk.fn.IdGenerators;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.fn.test.TestStreams;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Status;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.StatusRuntimeException;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessServerBuilder;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ManagedChannel;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Status;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.StatusRuntimeException;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessChannelBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.inprocess.InProcessServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.Uninterruptibles;
import org.junit.After;
import org.junit.Before;
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 7762e66..e634652 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
@@ -32,7 +32,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest.RequestCase;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
/** A fake implementation of a {@link BeamFnStateClient} to aid with testing. */
public class FakeBeamFnStateClient implements BeamFnStateClient {
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java
index 9705267..635c111 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java
@@ -22,7 +22,7 @@
import java.io.IOException;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.junit.Test;
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 630627d..d184ca2 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
@@ -24,7 +24,7 @@
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateGetResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactoriesTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactoriesTest.java
index d8f5872..76294ee 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactoriesTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/stream/HarnessStreamObserverFactoriesTest.java
@@ -25,8 +25,8 @@
import org.apache.beam.sdk.fn.stream.DirectStreamObserver;
import org.apache.beam.sdk.fn.stream.ForwardingClientResponseObserver;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.CallStreamObserver;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.CallStreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
diff --git a/sdks/java/io/bigquery-io-perf-tests/src/test/java/org/apache/beam/sdk/bigqueryioperftests/BigQueryIOIT.java b/sdks/java/io/bigquery-io-perf-tests/src/test/java/org/apache/beam/sdk/bigqueryioperftests/BigQueryIOIT.java
index 843ba78..a12dd41 100644
--- a/sdks/java/io/bigquery-io-perf-tests/src/test/java/org/apache/beam/sdk/bigqueryioperftests/BigQueryIOIT.java
+++ b/sdks/java/io/bigquery-io-perf-tests/src/test/java/org/apache/beam/sdk/bigqueryioperftests/BigQueryIOIT.java
@@ -163,7 +163,7 @@
BigQueryIO.Write.Method method = BigQueryIO.Write.Method.valueOf(options.getWriteMethod());
pipeline
.apply("Read from source", Read.from(new SyntheticBoundedSource(sourceOptions)))
- .apply("Gather time", ParDo.of(new TimeMonitor<>(NAMESPACE, WRITE_TIME_METRIC_NAME)))
+ .apply("Gather time", ParDo.of(new TimeMonitor<>(NAMESPACE, metricName)))
.apply("Map records", ParDo.of(new MapKVToV()))
.apply(
"Write to BQ",
diff --git a/sdks/java/io/cassandra/build.gradle b/sdks/java/io/cassandra/build.gradle
index 36dbede..6f02904 100644
--- a/sdks/java/io/cassandra/build.gradle
+++ b/sdks/java/io/cassandra/build.gradle
@@ -43,6 +43,7 @@
// for embedded cassandra
testCompile group: 'info.archinnov', name: 'achilles-junit', version: "$achilles_version"
+ testCompile library.java.commons_io
testCompile library.java.jackson_jaxb_annotations
testRuntimeOnly library.java.slf4j_jdk14
testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-2/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-2/build.gradle
index 0c4411d..3b64569 100644
--- a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-2/build.gradle
+++ b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-2/build.gradle
@@ -42,7 +42,6 @@
testCompile "org.apache.logging.log4j:log4j-core:$log4j_version"
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
- testCompile library.java.commons_io_1x
testCompile library.java.junit
testCompile "org.elasticsearch.client:elasticsearch-rest-client:5.6.3"
testCompile "org.elasticsearch:elasticsearch:$elastic_search_version"
diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-5/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-5/build.gradle
index 2e13700..c1560ac 100644
--- a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-5/build.gradle
+++ b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-5/build.gradle
@@ -61,7 +61,6 @@
testCompile "net.java.dev.jna:jna:$jna_version"
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
- testCompile library.java.commons_io_1x
testCompile library.java.junit
testCompile "org.elasticsearch.client:elasticsearch-rest-client:$elastic_search_version"
testRuntimeOnly library.java.slf4j_jdk14
diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-6/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-6/build.gradle
index b7bf6d0..6d7ae8a 100644
--- a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-6/build.gradle
+++ b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-6/build.gradle
@@ -61,7 +61,6 @@
testCompile "net.java.dev.jna:jna:$jna_version"
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
- testCompile library.java.commons_io_1x
testCompile library.java.junit
testCompile "org.elasticsearch.client:elasticsearch-rest-client:$elastic_search_version"
testRuntimeOnly library.java.slf4j_jdk14
diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-7/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-7/build.gradle
index 924b3fc..f731cd0 100644
--- a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-7/build.gradle
+++ b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-7/build.gradle
@@ -61,7 +61,6 @@
testCompile "net.java.dev.jna:jna:$jna_version"
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
- testCompile library.java.commons_io_1x
testCompile library.java.junit
testCompile "org.elasticsearch.client:elasticsearch-rest-client:$elastic_search_version"
testRuntimeOnly library.java.slf4j_jdk14
diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-common/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-common/build.gradle
index a77a29a..830abac 100644
--- a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-common/build.gradle
+++ b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-common/build.gradle
@@ -49,7 +49,6 @@
testCompile "org.apache.logging.log4j:log4j-core:$log4j_version"
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
- testCompile library.java.commons_io_1x
testCompile library.java.junit
testCompile "org.elasticsearch.client:elasticsearch-rest-client:$elastic_search_version"
testRuntimeOnly library.java.slf4j_jdk14
diff --git a/sdks/java/io/file-based-io-tests/src/test/java/org/apache/beam/sdk/io/tfrecord/TFRecordIOIT.java b/sdks/java/io/file-based-io-tests/src/test/java/org/apache/beam/sdk/io/tfrecord/TFRecordIOIT.java
index 59ac591..5473926 100644
--- a/sdks/java/io/file-based-io-tests/src/test/java/org/apache/beam/sdk/io/tfrecord/TFRecordIOIT.java
+++ b/sdks/java/io/file-based-io-tests/src/test/java/org/apache/beam/sdk/io/tfrecord/TFRecordIOIT.java
@@ -80,6 +80,12 @@
public class TFRecordIOIT {
private static final String TFRECORD_NAMESPACE = TFRecordIOIT.class.getName();
+ // Metric names
+ private static final String WRITE_TIME = "write_time";
+ private static final String READ_TIME = "read_time";
+ private static final String DATASET_SIZE = "dataset_size";
+ private static final String RUN_TIME = "run_time";
+
private static String filenamePrefix;
private static String bigQueryDataset;
private static String bigQueryTable;
@@ -111,7 +117,7 @@
// TODO: There are two pipelines due to: https://issues.apache.org/jira/browse/BEAM-3267
@Test
public void writeThenReadAll() {
- TFRecordIO.Write writeTransform =
+ final TFRecordIO.Write writeTransform =
TFRecordIO.write()
.to(filenamePrefix)
.withCompression(compressionType)
@@ -125,10 +131,11 @@
.apply("Transform strings to bytes", MapElements.via(new StringToByteArray()))
.apply(
"Record time before writing",
- ParDo.of(new TimeMonitor<>(TFRECORD_NAMESPACE, "writeTime")))
+ ParDo.of(new TimeMonitor<>(TFRECORD_NAMESPACE, WRITE_TIME)))
.apply("Write content to files", writeTransform);
- writePipeline.run().waitUntilFinish();
+ final PipelineResult writeResult = writePipeline.run();
+ writeResult.waitUntilFinish();
String filenamePattern = createFilenamePattern();
PCollection<String> consolidatedHashcode =
@@ -136,7 +143,7 @@
.apply(TFRecordIO.read().from(filenamePattern).withCompression(AUTO))
.apply(
"Record time after reading",
- ParDo.of(new TimeMonitor<>(TFRECORD_NAMESPACE, "readTime")))
+ ParDo.of(new TimeMonitor<>(TFRECORD_NAMESPACE, READ_TIME)))
.apply("Transform bytes to strings", MapElements.via(new ByteArrayToString()))
.apply("Calculate hashcode", Combine.globally(new HashingFn()))
.apply(Reshuffle.viaRandomKey());
@@ -149,54 +156,62 @@
"Delete test files",
ParDo.of(new DeleteFileFn())
.withSideInputs(consolidatedHashcode.apply(View.asSingleton())));
- PipelineResult result = readPipeline.run();
- result.waitUntilFinish();
- collectAndPublishMetrics(result);
+ final PipelineResult readResult = readPipeline.run();
+ readResult.waitUntilFinish();
+ collectAndPublishMetrics(writeResult, readResult);
}
- private void collectAndPublishMetrics(PipelineResult result) {
- String uuid = UUID.randomUUID().toString();
- String timestamp = Timestamp.now().toString();
+ private void collectAndPublishMetrics(
+ final PipelineResult writeResults, final PipelineResult readResults) {
+ final String uuid = UUID.randomUUID().toString();
+ final String timestamp = Timestamp.now().toString();
+ final Set<NamedTestResult> results = new HashSet<>();
- Set<Function<MetricsReader, NamedTestResult>> metricSuppliers =
- fillMetricSuppliers(uuid, timestamp);
- new IOITMetrics(metricSuppliers, result, TFRECORD_NAMESPACE, uuid, timestamp)
- .publish(bigQueryDataset, bigQueryTable);
+ results.add(
+ NamedTestResult.create(uuid, timestamp, RUN_TIME, getRunTime(writeResults, readResults)));
+ results.addAll(
+ MetricsReader.ofResults(writeResults, TFRECORD_NAMESPACE)
+ .readAll(getWriteMetricSuppliers(uuid, timestamp)));
+ results.addAll(
+ MetricsReader.ofResults(readResults, TFRECORD_NAMESPACE)
+ .readAll(getReadMetricSuppliers(uuid, timestamp)));
+
+ IOITMetrics.publish(uuid, timestamp, bigQueryDataset, bigQueryTable, results);
}
- private Set<Function<MetricsReader, NamedTestResult>> fillMetricSuppliers(
- String uuid, String timestamp) {
- Set<Function<MetricsReader, NamedTestResult>> suppliers = new HashSet<>();
- suppliers.add(
- reader -> {
- long writeStart = reader.getStartTimeMetric("writeTime");
- long writeEnd = reader.getEndTimeMetric("writeTime");
- double writeTime = (writeEnd - writeStart) / 1e3;
- return NamedTestResult.create(uuid, timestamp, "write_time", writeTime);
- });
-
- suppliers.add(
- reader -> {
- long readStart = reader.getStartTimeMetric("readTime");
- long readEnd = reader.getEndTimeMetric("readTime");
- double readTime = (readEnd - readStart) / 1e3;
- return NamedTestResult.create(uuid, timestamp, "read_time", readTime);
- });
-
- suppliers.add(
- reader -> {
- long writeStart = reader.getStartTimeMetric("writeTime");
- long readEnd = reader.getEndTimeMetric("readTime");
- double runTime = (readEnd - writeStart) / 1e3;
- return NamedTestResult.create(uuid, timestamp, "run_time", runTime);
- });
- if (datasetSize != null) {
- suppliers.add(
- (ignored) -> NamedTestResult.create(uuid, timestamp, "dataset_size", datasetSize));
- }
+ private static Set<Function<MetricsReader, NamedTestResult>> getWriteMetricSuppliers(
+ final String uuid, final String timestamp) {
+ final Set<Function<MetricsReader, NamedTestResult>> suppliers = new HashSet<>();
+ suppliers.add(getTimeMetric(uuid, timestamp, WRITE_TIME));
+ suppliers.add(ignored -> NamedTestResult.create(uuid, timestamp, DATASET_SIZE, datasetSize));
return suppliers;
}
+ private static Set<Function<MetricsReader, NamedTestResult>> getReadMetricSuppliers(
+ final String uuid, final String timestamp) {
+ final Set<Function<MetricsReader, NamedTestResult>> suppliers = new HashSet<>();
+ suppliers.add(getTimeMetric(uuid, timestamp, READ_TIME));
+ return suppliers;
+ }
+
+ private static Function<MetricsReader, NamedTestResult> getTimeMetric(
+ final String uuid, final String timestamp, final String metricName) {
+ return reader -> {
+ final long startTime = reader.getStartTimeMetric(metricName);
+ final long endTime = reader.getEndTimeMetric(metricName);
+ return NamedTestResult.create(uuid, timestamp, metricName, (endTime - startTime) / 1e3);
+ };
+ }
+
+ private static double getRunTime(
+ final PipelineResult writeResults, final PipelineResult readResult) {
+ final long startTime =
+ MetricsReader.ofResults(writeResults, TFRECORD_NAMESPACE).getStartTimeMetric(WRITE_TIME);
+ final long endTime =
+ MetricsReader.ofResults(readResult, TFRECORD_NAMESPACE).getEndTimeMetric(READ_TIME);
+ return (endTime - startTime) / 1e3;
+ }
+
static class StringToByteArray extends SimpleFunction<String, byte[]> {
@Override
public byte[] apply(String input) {
diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle
index 0c1befd..559488f 100644
--- a/sdks/java/io/google-cloud-platform/build.gradle
+++ b/sdks/java/io/google-cloud-platform/build.gradle
@@ -33,6 +33,7 @@
compile project(":sdks:java:extensions:protobuf")
compile library.java.avro
compile library.java.bigdataoss_util
+ compile library.java.gax
compile library.java.gax_grpc
compile library.java.google_api_client
compile library.java.google_api_services_bigquery
@@ -50,6 +51,7 @@
compile library.java.grpc_all
compile library.java.grpc_auth
compile library.java.grpc_core
+ compile library.java.grpc_context
compile library.java.grpc_netty
compile library.java.grpc_stub
compile library.java.grpc_google_cloud_pubsub_v1
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 f097e47..d477b08 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
@@ -136,6 +136,7 @@
private final Coder<ElementT> elementCoder;
private final RowWriterFactory<ElementT, DestinationT> rowWriterFactory;
private String kmsKey;
+ private boolean clusteringEnabled;
// The maximum number of times to retry failed load or copy jobs.
private int maxRetryJobs = DEFAULT_MAX_RETRY_JOBS;
@@ -151,7 +152,8 @@
boolean ignoreUnknownValues,
Coder<ElementT> elementCoder,
RowWriterFactory<ElementT, DestinationT> rowWriterFactory,
- @Nullable String kmsKey) {
+ @Nullable String kmsKey,
+ boolean clusteringEnabled) {
bigQueryServices = new BigQueryServicesImpl();
this.writeDisposition = writeDisposition;
this.createDisposition = createDisposition;
@@ -170,6 +172,7 @@
this.elementCoder = elementCoder;
this.kmsKey = kmsKey;
this.rowWriterFactory = rowWriterFactory;
+ this.clusteringEnabled = clusteringEnabled;
schemaUpdateOptions = Collections.emptySet();
}
@@ -319,6 +322,9 @@
.withOutputTags(multiPartitionsTag, TupleTagList.of(singlePartitionTag)));
PCollection<KV<TableDestination, String>> tempTables =
writeTempTables(partitions.get(multiPartitionsTag), loadJobIdPrefixView);
+
+ Coder<TableDestination> tableDestinationCoder =
+ clusteringEnabled ? TableDestinationCoderV3.of() : TableDestinationCoderV2.of();
tempTables
// Now that the load job has happened, we want the rename to happen immediately.
.apply(
@@ -326,8 +332,7 @@
.triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1))))
.apply(WithKeys.of((Void) null))
.setCoder(
- KvCoder.of(
- VoidCoder.of(), KvCoder.of(TableDestinationCoderV2.of(), StringUtf8Coder.of())))
+ KvCoder.of(VoidCoder.of(), KvCoder.of(tableDestinationCoder, StringUtf8Coder.of())))
.apply(GroupByKey.create())
.apply(Values.create())
.apply(
@@ -391,9 +396,11 @@
PCollection<KV<TableDestination, String>> tempTables =
writeTempTables(partitions.get(multiPartitionsTag), loadJobIdPrefixView);
+ Coder<TableDestination> tableDestinationCoder =
+ clusteringEnabled ? TableDestinationCoderV3.of() : TableDestinationCoderV2.of();
tempTables
.apply("ReifyRenameInput", new ReifyAsIterable<>())
- .setCoder(IterableCoder.of(KvCoder.of(TableDestinationCoderV2.of(), StringUtf8Coder.of())))
+ .setCoder(IterableCoder.of(KvCoder.of(tableDestinationCoder, StringUtf8Coder.of())))
.apply(
"WriteRenameUntriggered",
ParDo.of(
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 3bd9d8c..1fa8408 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
@@ -1643,6 +1643,7 @@
.setExtendedErrorInfo(false)
.setSkipInvalidRows(false)
.setIgnoreUnknownValues(false)
+ .setIgnoreInsertIds(false)
.setMaxFilesPerPartition(BatchLoads.DEFAULT_MAX_FILES_PER_PARTITION)
.setMaxBytesPerPartition(BatchLoads.DEFAULT_MAX_BYTES_PER_PARTITION)
.setOptimizeWrites(false)
@@ -1774,6 +1775,8 @@
abstract Boolean getIgnoreUnknownValues();
+ abstract Boolean getIgnoreInsertIds();
+
@Nullable
abstract String getKmsKey();
@@ -1846,6 +1849,8 @@
abstract Builder<T> setIgnoreUnknownValues(Boolean ignoreUnknownValues);
+ abstract Builder<T> setIgnoreInsertIds(Boolean ignoreInsertIds);
+
abstract Builder<T> setKmsKey(String kmsKey);
abstract Builder<T> setOptimizeWrites(Boolean optimizeWrites);
@@ -2241,6 +2246,15 @@
return toBuilder().setIgnoreUnknownValues(true).build();
}
+ /**
+ * Setting this option to true disables insertId based data deduplication offered by BigQuery.
+ * For more information, please see
+ * https://cloud.google.com/bigquery/streaming-data-into-bigquery#disabling_best_effort_de-duplication.
+ */
+ public Write<T> ignoreInsertIds() {
+ return toBuilder().setIgnoreInsertIds(true).build();
+ }
+
public Write<T> withKmsKey(String kmsKey) {
return toBuilder().setKmsKey(kmsKey).build();
}
@@ -2302,11 +2316,22 @@
return toBuilder().setMaxFilesPerPartition(maxFilesPerPartition).build();
}
- @VisibleForTesting
- Write<T> withMaxBytesPerPartition(long maxBytesPerPartition) {
+ /**
+ * Control how much data will be assigned to a single BigQuery load job. If the amount of data
+ * flowing into one {@code BatchLoads} partition exceeds this value, that partition will be
+ * handled via multiple load jobs.
+ *
+ * <p>The default value (11 TiB) respects BigQuery's maximum size per load job limit and is
+ * appropriate for most use cases. Reducing the value of this parameter can improve stability
+ * when loading to tables with complex schemas containing thousands of fields.
+ *
+ * @see <a href="https://cloud.google.com/bigquery/quotas#load_jobs">BigQuery Load Job
+ * Limits</a>
+ */
+ public Write<T> withMaxBytesPerPartition(long maxBytesPerPartition) {
checkArgument(
maxBytesPerPartition > 0,
- "maxFilesPerPartition must be > 0, but was: %s",
+ "maxBytesPerPartition must be > 0, but was: %s",
maxBytesPerPartition);
return toBuilder().setMaxBytesPerPartition(maxBytesPerPartition).build();
}
@@ -2600,6 +2625,7 @@
.withExtendedErrorInfo(getExtendedErrorInfo())
.withSkipInvalidRows(getSkipInvalidRows())
.withIgnoreUnknownValues(getIgnoreUnknownValues())
+ .withIgnoreInsertIds(getIgnoreInsertIds())
.withKmsKey(getKmsKey());
return input.apply(streamingInserts);
} else {
@@ -2619,7 +2645,8 @@
getIgnoreUnknownValues(),
elementCoder,
rowWriterFactory,
- getKmsKey());
+ getKmsKey(),
+ getClustering() != null);
batchLoads.setTestServices(getBigQueryServices());
if (getSchemaUpdateOptions() != null) {
batchLoads.setSchemaUpdateOptions(getSchemaUpdateOptions());
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java
index ecd4a85..ce02423 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java
@@ -158,7 +158,8 @@
List<ValueInSingleWindow<T>> failedInserts,
ErrorContainer<T> errorContainer,
boolean skipInvalidRows,
- boolean ignoreUnknownValues)
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds)
throws IOException, InterruptedException;
/** Patch BigQuery {@link Table} description. */
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
index 218de47..bf2bd3c 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java
@@ -713,7 +713,8 @@
List<ValueInSingleWindow<T>> failedInserts,
ErrorContainer<T> errorContainer,
boolean skipInvalidRows,
- boolean ignoreUnkownValues)
+ boolean ignoreUnkownValues,
+ boolean ignoreInsertIds)
throws IOException, InterruptedException {
checkNotNull(ref, "ref");
if (executor == null) {
@@ -733,7 +734,10 @@
// These lists contain the rows to publish. Initially the contain the entire list.
// If there are failures, they will contain only the failed rows to be retried.
List<ValueInSingleWindow<TableRow>> rowsToPublish = rowList;
- List<String> idsToPublish = insertIdList;
+ List<String> idsToPublish = null;
+ if (!ignoreInsertIds) {
+ idsToPublish = insertIdList;
+ }
while (true) {
List<ValueInSingleWindow<TableRow>> retryRows = new ArrayList<>();
List<String> retryIds = (idsToPublish != null) ? new ArrayList<>() : null;
@@ -871,7 +875,8 @@
List<ValueInSingleWindow<T>> failedInserts,
ErrorContainer<T> errorContainer,
boolean skipInvalidRows,
- boolean ignoreUnknownValues)
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds)
throws IOException, InterruptedException {
return insertAll(
ref,
@@ -883,7 +888,8 @@
failedInserts,
errorContainer,
skipInvalidRows,
- ignoreUnknownValues);
+ ignoreUnknownValues,
+ ignoreInsertIds);
}
@Override
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java
index 6dad989..d00adbb 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingInserts.java
@@ -38,6 +38,7 @@
private boolean extendedErrorInfo;
private final boolean skipInvalidRows;
private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
private final String kmsKey;
private final Coder<ElementT> elementCoder;
private final SerializableFunction<ElementT, TableRow> toTableRow;
@@ -56,6 +57,7 @@
false,
false,
false,
+ false,
elementCoder,
toTableRow,
null);
@@ -70,6 +72,7 @@
boolean extendedErrorInfo,
boolean skipInvalidRows,
boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
Coder<ElementT> elementCoder,
SerializableFunction<ElementT, TableRow> toTableRow,
String kmsKey) {
@@ -80,6 +83,7 @@
this.extendedErrorInfo = extendedErrorInfo;
this.skipInvalidRows = skipInvalidRows;
this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
this.elementCoder = elementCoder;
this.toTableRow = toTableRow;
this.kmsKey = kmsKey;
@@ -96,6 +100,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow,
kmsKey);
@@ -111,6 +116,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow,
kmsKey);
@@ -125,6 +131,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow,
kmsKey);
@@ -139,6 +146,22 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
+ elementCoder,
+ toTableRow,
+ kmsKey);
+ }
+
+ StreamingInserts<DestinationT, ElementT> withIgnoreInsertIds(boolean ignoreInsertIds) {
+ return new StreamingInserts<>(
+ createDisposition,
+ dynamicDestinations,
+ bigQueryServices,
+ retryPolicy,
+ extendedErrorInfo,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow,
kmsKey);
@@ -153,6 +176,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow,
kmsKey);
@@ -167,6 +191,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow,
kmsKey);
@@ -188,6 +213,7 @@
.withExtendedErrorInfo(extendedErrorInfo)
.withSkipInvalidRows(skipInvalidRows)
.withIgnoreUnknownValues(ignoreUnknownValues)
+ .withIgnoreInsertIds(ignoreInsertIds)
.withElementCoder(elementCoder)
.withToTableRow(toTableRow));
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
index f56cf01..4e12018 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
@@ -49,6 +49,7 @@
private final ErrorContainer<ErrorT> errorContainer;
private final boolean skipInvalidRows;
private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
private final SerializableFunction<ElementT, TableRow> toTableRow;
/** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
@@ -67,6 +68,7 @@
ErrorContainer<ErrorT> errorContainer,
boolean skipInvalidRows,
boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
SerializableFunction<ElementT, TableRow> toTableRow) {
this.bqServices = bqServices;
this.retryPolicy = retryPolicy;
@@ -74,6 +76,7 @@
this.errorContainer = errorContainer;
this.skipInvalidRows = skipInvalidRows;
this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
this.toTableRow = toTableRow;
}
@@ -145,7 +148,8 @@
failedInserts,
errorContainer,
skipInvalidRows,
- ignoreUnknownValues);
+ ignoreUnknownValues,
+ ignoreInsertIds);
byteCounter.inc(totalBytes);
} catch (IOException e) {
throw new RuntimeException(e);
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
index ea2c020..81f097a 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
@@ -53,6 +53,7 @@
private static final String FAILED_INSERTS_TAG_ID = "failedInserts";
private final boolean skipInvalidRows;
private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
private final Coder<ElementT> elementCoder;
private final SerializableFunction<ElementT, TableRow> toTableRow;
@@ -63,6 +64,7 @@
false, // extendedErrorInfo
false, // skipInvalidRows
false, // ignoreUnknownValues
+ false, // ignoreInsertIds
null, // elementCoder
null); // toTableRow
}
@@ -73,6 +75,7 @@
boolean extendedErrorInfo,
boolean skipInvalidRows,
boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
Coder<ElementT> elementCoder,
SerializableFunction<ElementT, TableRow> toTableRow) {
this.bigQueryServices = bigQueryServices;
@@ -80,6 +83,7 @@
this.extendedErrorInfo = extendedErrorInfo;
this.skipInvalidRows = skipInvalidRows;
this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
this.elementCoder = elementCoder;
this.toTableRow = toTableRow;
}
@@ -91,6 +95,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow);
}
@@ -102,6 +107,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow);
}
@@ -113,6 +119,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow);
}
@@ -124,6 +131,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow);
}
@@ -135,6 +143,19 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
+ elementCoder,
+ toTableRow);
+ }
+
+ StreamingWriteTables<ElementT> withIgnoreInsertIds(boolean ignoreInsertIds) {
+ return new StreamingWriteTables<>(
+ bigQueryServices,
+ retryPolicy,
+ extendedErrorInfo,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow);
}
@@ -146,6 +167,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow);
}
@@ -158,6 +180,7 @@
extendedErrorInfo,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
elementCoder,
toTableRow);
}
@@ -240,6 +263,7 @@
errorContainer,
skipInvalidRows,
ignoreUnknownValues,
+ ignoreInsertIds,
toTableRow))
.withOutputTags(mainOutputTag, TupleTagList.of(failedInsertsTag)));
PCollection<T> failedInserts = tuple.get(failedInsertsTag);
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 acc40e1..f16c616 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
@@ -368,18 +368,9 @@
return entity.getProperties().get("timestamp").getTimestampValue().getSeconds() * 1000000;
}
- /**
- * Get the estimated size of the data returned by the given query.
- *
- * <p>Cloud Datastore provides no way to get a good estimate of how large the result of a query
- * entity kind being queried, using the __Stat_Kind__ system table, assuming exactly 1 kind is
- * specified in the query.
- *
- * <p>See https://cloud.google.com/datastore/docs/concepts/stats.
- */
- static long getEstimatedSizeBytes(Datastore datastore, Query query, @Nullable String namespace)
- throws DatastoreException {
- String ourKind = query.getKind(0).getName();
+ /** Retrieve latest table statistics for a given kind, namespace, and datastore. */
+ private static Entity getLatestTableStats(
+ String ourKind, @Nullable String namespace, Datastore datastore) throws DatastoreException {
long latestTimestamp = queryLatestStatisticsTimestamp(datastore, namespace);
LOG.info("Latest stats timestamp for kind {} is {}", ourKind, latestTimestamp);
@@ -406,7 +397,22 @@
throw new NoSuchElementException(
"Datastore statistics for kind " + ourKind + " unavailable");
}
- Entity entity = batch.getEntityResults(0).getEntity();
+ return batch.getEntityResults(0).getEntity();
+ }
+
+ /**
+ * Get the estimated size of the data returned by the given query.
+ *
+ * <p>Cloud Datastore provides no way to get a good estimate of how large the result of a query
+ * entity kind being queried, using the __Stat_Kind__ system table, assuming exactly 1 kind is
+ * specified in the query.
+ *
+ * <p>See https://cloud.google.com/datastore/docs/concepts/stats.
+ */
+ static long getEstimatedSizeBytes(Datastore datastore, Query query, @Nullable String namespace)
+ throws DatastoreException {
+ String ourKind = query.getKind(0).getName();
+ Entity entity = getLatestTableStats(ourKind, namespace, datastore);
return entity.getProperties().get("entity_bytes").getIntegerValue();
}
@@ -600,6 +606,23 @@
return toBuilder().setLocalhost(localhost).build();
}
+ /** Returns Number of entities available for reading. */
+ public long getNumEntities(
+ PipelineOptions options, String ourKind, @Nullable String namespace) {
+ try {
+ V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost());
+ V1DatastoreFactory datastoreFactory = new V1DatastoreFactory();
+ Datastore datastore =
+ datastoreFactory.getDatastore(
+ options, v1Options.getProjectId(), v1Options.getLocalhost());
+
+ Entity entity = getLatestTableStats(ourKind, namespace, datastore);
+ return entity.getProperties().get("count").getIntegerValue();
+ } catch (Exception e) {
+ return -1;
+ }
+ }
+
@Override
public PCollection<Entity> expand(PBegin input) {
checkArgument(getProjectId() != null, "projectId provider cannot be null");
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java
index 07d6da6..6f0f54d 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java
@@ -21,10 +21,12 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
import com.google.api.client.util.DateTime;
+import com.google.auto.value.AutoValue;
+import com.google.protobuf.ByteString;
+import com.google.pubsub.v1.PubsubMessage;
import java.io.Closeable;
import java.io.IOException;
import java.io.Serializable;
-import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
@@ -298,59 +300,37 @@
* <p>NOTE: This class is {@link Serializable} only to support the {@link PubsubTestClient}. Java
* serialization is never used for non-test clients.
*/
- public static class OutgoingMessage implements Serializable {
- /** Underlying (encoded) element. */
- public final byte[] elementBytes;
+ @AutoValue
+ public abstract static class OutgoingMessage implements Serializable {
- public final Map<String, String> attributes;
+ /** Underlying Message. May not have publish timestamp set. */
+ public abstract PubsubMessage message();
/** Timestamp for element (ms since epoch). */
- public final long timestampMsSinceEpoch;
+ public abstract long timestampMsSinceEpoch();
/**
* If using an id attribute, the record id to associate with this record's metadata so the
* receiver can reject duplicates. Otherwise {@literal null}.
*/
- @Nullable public final String recordId;
+ @Nullable
+ public abstract String recordId();
- public OutgoingMessage(
- byte[] elementBytes,
- Map<String, String> attributes,
+ public static OutgoingMessage of(
+ PubsubMessage message, long timestampMsSinceEpoch, @Nullable String recordId) {
+ return new AutoValue_PubsubClient_OutgoingMessage(message, timestampMsSinceEpoch, recordId);
+ }
+
+ public static OutgoingMessage of(
+ org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage message,
long timestampMsSinceEpoch,
@Nullable String recordId) {
- this.elementBytes = elementBytes;
- this.attributes = attributes;
- this.timestampMsSinceEpoch = timestampMsSinceEpoch;
- this.recordId = recordId;
- }
-
- @Override
- public String toString() {
- return String.format(
- "OutgoingMessage(%db, %dms)", elementBytes.length, timestampMsSinceEpoch);
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
+ PubsubMessage.Builder builder =
+ PubsubMessage.newBuilder().setData(ByteString.copyFrom(message.getPayload()));
+ if (message.getAttributeMap() != null) {
+ builder.putAllAttributes(message.getAttributeMap());
}
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- OutgoingMessage that = (OutgoingMessage) o;
-
- return timestampMsSinceEpoch == that.timestampMsSinceEpoch
- && Arrays.equals(elementBytes, that.elementBytes)
- && Objects.equal(attributes, that.attributes)
- && Objects.equal(recordId, that.recordId);
- }
-
- @Override
- public int hashCode() {
- return Objects.hashCode(
- Arrays.hashCode(elementBytes), attributes, timestampMsSinceEpoch, recordId);
+ return of(builder.build(), timestampMsSinceEpoch, recordId);
}
}
@@ -360,86 +340,35 @@
* <p>NOTE: This class is {@link Serializable} only to support the {@link PubsubTestClient}. Java
* serialization is never used for non-test clients.
*/
- static class IncomingMessage implements Serializable {
- /** Underlying (encoded) element. */
- public final byte[] elementBytes;
+ @AutoValue
+ abstract static class IncomingMessage implements Serializable {
- public Map<String, String> attributes;
+ /** Underlying Message. */
+ public abstract PubsubMessage message();
/**
* Timestamp for element (ms since epoch). Either Pubsub's processing time, or the custom
* timestamp associated with the message.
*/
- public final long timestampMsSinceEpoch;
+ public abstract long timestampMsSinceEpoch();
/** Timestamp (in system time) at which we requested the message (ms since epoch). */
- public final long requestTimeMsSinceEpoch;
+ public abstract long requestTimeMsSinceEpoch();
/** Id to pass back to Pubsub to acknowledge receipt of this message. */
- public final String ackId;
+ public abstract String ackId();
/** Id to pass to the runner to distinguish this message from all others. */
- public final String recordId;
+ public abstract String recordId();
- public IncomingMessage(
- byte[] elementBytes,
- Map<String, String> attributes,
+ public static IncomingMessage of(
+ PubsubMessage message,
long timestampMsSinceEpoch,
long requestTimeMsSinceEpoch,
String ackId,
String recordId) {
- this.elementBytes = elementBytes;
- this.attributes = attributes;
- this.timestampMsSinceEpoch = timestampMsSinceEpoch;
- this.requestTimeMsSinceEpoch = requestTimeMsSinceEpoch;
- this.ackId = ackId;
- this.recordId = recordId;
- }
-
- public IncomingMessage withRequestTime(long requestTimeMsSinceEpoch) {
- return new IncomingMessage(
- elementBytes,
- attributes,
- timestampMsSinceEpoch,
- requestTimeMsSinceEpoch,
- ackId,
- recordId);
- }
-
- @Override
- public String toString() {
- return String.format(
- "IncomingMessage(%db, %dms)", elementBytes.length, timestampMsSinceEpoch);
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- IncomingMessage that = (IncomingMessage) o;
-
- return timestampMsSinceEpoch == that.timestampMsSinceEpoch
- && requestTimeMsSinceEpoch == that.requestTimeMsSinceEpoch
- && ackId.equals(that.ackId)
- && recordId.equals(that.recordId)
- && Arrays.equals(elementBytes, that.elementBytes)
- && Objects.equal(attributes, that.attributes);
- }
-
- @Override
- public int hashCode() {
- return Objects.hashCode(
- Arrays.hashCode(elementBytes),
- attributes,
- timestampMsSinceEpoch,
- requestTimeMsSinceEpoch,
- ackId,
- recordId);
+ return new AutoValue_PubsubClient_IncomingMessage(
+ message, timestampMsSinceEpoch, requestTimeMsSinceEpoch, ackId, recordId);
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java
index ae3fa02..a3b6b8d 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java
@@ -20,7 +20,6 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
import com.google.auth.Credentials;
-import com.google.protobuf.ByteString;
import com.google.protobuf.Timestamp;
import com.google.pubsub.v1.AcknowledgeRequest;
import com.google.pubsub.v1.DeleteSubscriptionRequest;
@@ -213,21 +212,15 @@
public int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages) throws IOException {
PublishRequest.Builder request = PublishRequest.newBuilder().setTopic(topic.getPath());
for (OutgoingMessage outgoingMessage : outgoingMessages) {
- PubsubMessage.Builder message =
- PubsubMessage.newBuilder().setData(ByteString.copyFrom(outgoingMessage.elementBytes));
-
- if (outgoingMessage.attributes != null) {
- message.putAllAttributes(outgoingMessage.attributes);
- }
+ PubsubMessage.Builder message = outgoingMessage.message().toBuilder();
if (timestampAttribute != null) {
- message
- .getMutableAttributes()
- .put(timestampAttribute, String.valueOf(outgoingMessage.timestampMsSinceEpoch));
+ message.putAttributes(
+ timestampAttribute, String.valueOf(outgoingMessage.timestampMsSinceEpoch()));
}
- if (idAttribute != null && !Strings.isNullOrEmpty(outgoingMessage.recordId)) {
- message.getMutableAttributes().put(idAttribute, outgoingMessage.recordId);
+ if (idAttribute != null && !Strings.isNullOrEmpty(outgoingMessage.recordId())) {
+ message.putAttributes(idAttribute, outgoingMessage.recordId());
}
request.addMessages(message);
@@ -259,9 +252,6 @@
PubsubMessage pubsubMessage = message.getMessage();
@Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
- // Payload.
- byte[] elementBytes = pubsubMessage.getData().toByteArray();
-
// Timestamp.
String pubsubTimestampString = null;
Timestamp timestampProto = pubsubMessage.getPublishTime();
@@ -287,13 +277,8 @@
}
incomingMessages.add(
- new IncomingMessage(
- elementBytes,
- attributes,
- timestampMsSinceEpoch,
- requestTimeMsSinceEpoch,
- ackId,
- recordId));
+ IncomingMessage.of(
+ pubsubMessage, timestampMsSinceEpoch, requestTimeMsSinceEpoch, ackId, recordId));
}
return incomingMessages;
}
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 da5266f..5f6d044 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
@@ -1303,7 +1303,14 @@
}
// NOTE: The record id is always null.
- output.add(new OutgoingMessage(payload, attributes, c.timestamp().getMillis(), null));
+ output.add(
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFrom(payload))
+ .putAllAttributes(attributes)
+ .build(),
+ c.timestamp().getMillis(),
+ null));
currentOutputBytes += payload.length;
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java
index 136b1d2..1ae5a55 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java
@@ -39,8 +39,10 @@
import com.google.auth.Credentials;
import com.google.auth.http.HttpCredentialsAdapter;
import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;
+import com.google.protobuf.ByteString;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
@@ -123,8 +125,12 @@
public int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages) throws IOException {
List<PubsubMessage> pubsubMessages = new ArrayList<>(outgoingMessages.size());
for (OutgoingMessage outgoingMessage : outgoingMessages) {
- PubsubMessage pubsubMessage = new PubsubMessage().encodeData(outgoingMessage.elementBytes);
+ PubsubMessage pubsubMessage =
+ new PubsubMessage().encodeData(outgoingMessage.message().getData().toByteArray());
pubsubMessage.setAttributes(getMessageAttributes(outgoingMessage));
+ if (!outgoingMessage.message().getOrderingKey().isEmpty()) {
+ pubsubMessage.put("orderingKey", outgoingMessage.message().getOrderingKey());
+ }
pubsubMessages.add(pubsubMessage);
}
PublishRequest request = new PublishRequest().setMessages(pubsubMessages);
@@ -135,16 +141,16 @@
private Map<String, String> getMessageAttributes(OutgoingMessage outgoingMessage) {
Map<String, String> attributes = null;
- if (outgoingMessage.attributes == null) {
+ if (outgoingMessage.message().getAttributesMap() == null) {
attributes = new TreeMap<>();
} else {
- attributes = new TreeMap<>(outgoingMessage.attributes);
+ attributes = new TreeMap<>(outgoingMessage.message().getAttributesMap());
}
if (timestampAttribute != null) {
- attributes.put(timestampAttribute, String.valueOf(outgoingMessage.timestampMsSinceEpoch));
+ attributes.put(timestampAttribute, String.valueOf(outgoingMessage.timestampMsSinceEpoch()));
}
- if (idAttribute != null && !Strings.isNullOrEmpty(outgoingMessage.recordId)) {
- attributes.put(idAttribute, outgoingMessage.recordId);
+ if (idAttribute != null && !Strings.isNullOrEmpty(outgoingMessage.recordId())) {
+ attributes.put(idAttribute, outgoingMessage.recordId());
}
return attributes;
}
@@ -166,7 +172,12 @@
List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessages().size());
for (ReceivedMessage message : response.getReceivedMessages()) {
PubsubMessage pubsubMessage = message.getMessage();
- @Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
+ Map<String, String> attributes;
+ if (pubsubMessage.getAttributes() != null) {
+ attributes = pubsubMessage.getAttributes();
+ } else {
+ attributes = new HashMap<>();
+ }
// Payload.
byte[] elementBytes = pubsubMessage.getData() == null ? null : pubsubMessage.decodeData();
@@ -184,7 +195,7 @@
// Record id, if any.
@Nullable String recordId = null;
- if (idAttribute != null && attributes != null) {
+ if (idAttribute != null) {
recordId = attributes.get(idAttribute);
}
if (Strings.isNullOrEmpty(recordId)) {
@@ -192,10 +203,15 @@
recordId = pubsubMessage.getMessageId();
}
+ com.google.pubsub.v1.PubsubMessage.Builder protoMessage =
+ com.google.pubsub.v1.PubsubMessage.newBuilder();
+ protoMessage.setData(ByteString.copyFrom(elementBytes));
+ protoMessage.putAllAttributes(attributes);
+ protoMessage.setOrderingKey(
+ (String) pubsubMessage.getUnknownKeys().getOrDefault("orderingKey", ""));
incomingMessages.add(
- new IncomingMessage(
- elementBytes,
- attributes,
+ IncomingMessage.of(
+ protoMessage.build(),
timestampMsSinceEpoch,
requestTimeMsSinceEpoch,
ackId,
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java
index 6b20b56..c3b915d 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java
@@ -309,12 +309,17 @@
IncomingMessage incomingMessage = pendItr.next();
pendItr.remove();
IncomingMessage incomingMessageWithRequestTime =
- incomingMessage.withRequestTime(requestTimeMsSinceEpoch);
+ IncomingMessage.of(
+ incomingMessage.message(),
+ incomingMessage.timestampMsSinceEpoch(),
+ requestTimeMsSinceEpoch,
+ incomingMessage.ackId(),
+ incomingMessage.recordId());
incomingMessages.add(incomingMessageWithRequestTime);
STATE.pendingAckIncomingMessages.put(
- incomingMessageWithRequestTime.ackId, incomingMessageWithRequestTime);
+ incomingMessageWithRequestTime.ackId(), incomingMessageWithRequestTime);
STATE.ackDeadline.put(
- incomingMessageWithRequestTime.ackId,
+ incomingMessageWithRequestTime.ackId(),
requestTimeMsSinceEpoch + STATE.ackTimeoutSec * 1000);
if (incomingMessages.size() >= batchSize) {
break;
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java
index 1258d0b..8be8c56 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java
@@ -30,7 +30,6 @@
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.BigEndianLongCoder;
-import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.KvCoder;
@@ -38,6 +37,7 @@
import org.apache.beam.sdk.coders.NullableCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.extensions.protobuf.ProtoCoder;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.OutgoingMessage;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.PubsubClientFactory;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
@@ -101,19 +101,18 @@
@Override
public void encode(OutgoingMessage value, OutputStream outStream)
throws CoderException, IOException {
- ByteArrayCoder.of().encode(value.elementBytes, outStream);
- ATTRIBUTES_CODER.encode(value.attributes, outStream);
- BigEndianLongCoder.of().encode(value.timestampMsSinceEpoch, outStream);
- RECORD_ID_CODER.encode(value.recordId, outStream);
+ ProtoCoder.of(com.google.pubsub.v1.PubsubMessage.class).encode(value.message(), outStream);
+ BigEndianLongCoder.of().encode(value.timestampMsSinceEpoch(), outStream);
+ RECORD_ID_CODER.encode(value.recordId(), outStream);
}
@Override
public OutgoingMessage decode(InputStream inStream) throws CoderException, IOException {
- byte[] elementBytes = ByteArrayCoder.of().decode(inStream);
- Map<String, String> attributes = ATTRIBUTES_CODER.decode(inStream);
+ com.google.pubsub.v1.PubsubMessage message =
+ ProtoCoder.of(com.google.pubsub.v1.PubsubMessage.class).decode(inStream);
long timestampMsSinceEpoch = BigEndianLongCoder.of().decode(inStream);
@Nullable String recordId = RECORD_ID_CODER.decode(inStream);
- return new OutgoingMessage(elementBytes, attributes, timestampMsSinceEpoch, recordId);
+ return OutgoingMessage.of(message, timestampMsSinceEpoch, recordId);
}
}
@@ -154,7 +153,6 @@
elementCounter.inc();
PubsubMessage message = c.element();
byte[] elementBytes = message.getPayload();
- Map<String, String> attributes = message.getAttributeMap();
long timestampMsSinceEpoch = c.timestamp().getMillis();
@Nullable String recordId = null;
@@ -175,7 +173,7 @@
c.output(
KV.of(
ThreadLocalRandom.current().nextInt(numShards),
- new OutgoingMessage(elementBytes, attributes, timestampMsSinceEpoch, recordId)));
+ OutgoingMessage.of(message, timestampMsSinceEpoch, recordId)));
}
@Override
@@ -246,7 +244,8 @@
List<OutgoingMessage> pubsubMessages = new ArrayList<>(publishBatchSize);
int bytes = 0;
for (OutgoingMessage message : c.element().getValue()) {
- if (!pubsubMessages.isEmpty() && bytes + message.elementBytes.length > publishBatchBytes) {
+ if (!pubsubMessages.isEmpty()
+ && bytes + message.message().getData().size() > publishBatchBytes) {
// Break large (in bytes) batches into smaller.
// (We've already broken by batch size using the trigger below, though that may
// run slightly over the actual PUBLISH_BATCH_SIZE. We'll consider that ok since
@@ -257,7 +256,7 @@
bytes = 0;
}
pubsubMessages.add(message);
- bytes += message.elementBytes.length;
+ bytes += message.message().getData().size();
}
if (!pubsubMessages.isEmpty()) {
// BLOCKS until published.
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java
index d8abfe1..230161c 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java
@@ -727,18 +727,18 @@
// Capture the received messages.
for (PubsubClient.IncomingMessage incomingMessage : receivedMessages) {
notYetRead.add(incomingMessage);
- notYetReadBytes += incomingMessage.elementBytes.length;
+ notYetReadBytes += incomingMessage.message().getData().size();
inFlight.put(
- incomingMessage.ackId,
+ incomingMessage.ackId(),
new InFlightState(requestTimeMsSinceEpoch, deadlineMsSinceEpoch));
numReceived++;
numReceivedRecently.add(requestTimeMsSinceEpoch, 1L);
minReceivedTimestampMsSinceEpoch.add(
- requestTimeMsSinceEpoch, incomingMessage.timestampMsSinceEpoch);
+ requestTimeMsSinceEpoch, incomingMessage.timestampMsSinceEpoch());
maxReceivedTimestampMsSinceEpoch.add(
- requestTimeMsSinceEpoch, incomingMessage.timestampMsSinceEpoch);
+ requestTimeMsSinceEpoch, incomingMessage.timestampMsSinceEpoch());
minUnreadTimestampMsSinceEpoch.add(
- requestTimeMsSinceEpoch, incomingMessage.timestampMsSinceEpoch);
+ requestTimeMsSinceEpoch, incomingMessage.timestampMsSinceEpoch());
}
}
@@ -837,7 +837,7 @@
if (current != null) {
// Current is consumed. It can no longer contribute to holding back the watermark.
- minUnreadTimestampMsSinceEpoch.remove(current.requestTimeMsSinceEpoch);
+ minUnreadTimestampMsSinceEpoch.remove(current.requestTimeMsSinceEpoch());
current = null;
}
@@ -864,18 +864,18 @@
// Try again later.
return false;
}
- notYetReadBytes -= current.elementBytes.length;
+ notYetReadBytes -= current.message().getData().size();
checkState(notYetReadBytes >= 0);
long nowMsSinceEpoch = now();
- numReadBytes.add(nowMsSinceEpoch, current.elementBytes.length);
- minReadTimestampMsSinceEpoch.add(nowMsSinceEpoch, current.timestampMsSinceEpoch);
- if (current.timestampMsSinceEpoch < lastWatermarkMsSinceEpoch) {
+ numReadBytes.add(nowMsSinceEpoch, current.message().getData().size());
+ minReadTimestampMsSinceEpoch.add(nowMsSinceEpoch, current.timestampMsSinceEpoch());
+ if (current.timestampMsSinceEpoch() < lastWatermarkMsSinceEpoch) {
numLateMessages.add(nowMsSinceEpoch, 1L);
}
// Current message can be considered 'read' and will be persisted by the next
// checkpoint. So it is now safe to ACK back to Pubsub.
- safeToAckIds.add(current.ackId);
+ safeToAckIds.add(current.ackId());
return true;
}
@@ -884,7 +884,10 @@
if (current == null) {
throw new NoSuchElementException();
}
- return new PubsubMessage(current.elementBytes, current.attributes, current.recordId);
+ return new PubsubMessage(
+ current.message().getData().toByteArray(),
+ current.message().getAttributesMap(),
+ current.recordId());
}
@Override
@@ -892,7 +895,7 @@
if (current == null) {
throw new NoSuchElementException();
}
- return new Instant(current.timestampMsSinceEpoch);
+ return new Instant(current.timestampMsSinceEpoch());
}
@Override
@@ -900,7 +903,7 @@
if (current == null) {
throw new NoSuchElementException();
}
- return current.recordId.getBytes(StandardCharsets.UTF_8);
+ return current.recordId().getBytes(StandardCharsets.UTF_8);
}
/**
@@ -984,7 +987,7 @@
List<String> snapshotSafeToAckIds = Lists.newArrayList(safeToAckIds);
List<String> snapshotNotYetReadIds = new ArrayList<>(notYetRead.size());
for (PubsubClient.IncomingMessage incomingMessage : notYetRead) {
- snapshotNotYetReadIds.add(incomingMessage.ackId);
+ snapshotNotYetReadIds.add(incomingMessage.ackId());
}
if (outer.subscriptionPath == null) {
// need to include the subscription in case we resume, as it's not stored in the source.
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsub.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsub.java
index 1e75d43..e1e8711 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsub.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsub.java
@@ -22,12 +22,14 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsInAnyOrder;
+import com.google.protobuf.ByteString;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeoutException;
import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.ProjectPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
@@ -205,11 +207,16 @@
if (!messages.isEmpty()) {
pubsub.acknowledge(
subscriptionPath,
- messages.stream().map(msg -> msg.ackId).collect(ImmutableList.toImmutableList()));
+ messages.stream().map(IncomingMessage::ackId).collect(ImmutableList.toImmutableList()));
}
return messages.stream()
- .map(msg -> new PubsubMessage(msg.elementBytes, msg.attributes, msg.recordId))
+ .map(
+ msg ->
+ new PubsubMessage(
+ msg.message().getData().toByteArray(),
+ msg.message().getAttributesMap(),
+ msg.recordId()))
.collect(ImmutableList.toImmutableList());
}
@@ -292,7 +299,12 @@
}
private PubsubClient.OutgoingMessage toOutgoingMessage(PubsubMessage message) {
- return new PubsubClient.OutgoingMessage(
- message.getPayload(), message.getAttributeMap(), DateTime.now().getMillis(), null);
+ return PubsubClient.OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFrom(message.getPayload()))
+ .putAllAttributes(message.getAttributeMap())
+ .build(),
+ DateTime.now().getMillis(),
+ null);
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsubSignal.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsubSignal.java
index f4f8b18..de4a715 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsubSignal.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/TestPubsubSignal.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.sdk.io.gcp.pubsub;
-import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.stream.Collectors.toList;
import static org.apache.beam.sdk.io.gcp.pubsub.TestPubsub.createTopicName;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState;
@@ -30,6 +29,7 @@
import java.util.concurrent.ThreadLocalRandom;
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
import org.apache.beam.sdk.state.BagState;
@@ -251,7 +251,7 @@
try {
signal = pubsub.pull(DateTime.now().getMillis(), signalSubscriptionPath, 1, false);
pubsub.acknowledge(
- signalSubscriptionPath, signal.stream().map(m -> m.ackId).collect(toList()));
+ signalSubscriptionPath, signal.stream().map(IncomingMessage::ackId).collect(toList()));
break;
} catch (StatusRuntimeException e) {
if (!Status.DEADLINE_EXCEEDED.equals(e.getStatus())) {
@@ -271,7 +271,7 @@
signalSubscriptionPath, duration.getStandardSeconds()));
}
- return new String(signal.get(0).elementBytes, UTF_8);
+ return signal.get(0).message().getData().toStringUtf8();
}
private void sleep(long t) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcher.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcher.java
index 5312e23..85d83fd 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcher.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcher.java
@@ -22,6 +22,7 @@
import com.google.api.services.bigquery.model.QueryResponse;
import com.google.api.services.bigquery.model.TableCell;
import com.google.api.services.bigquery.model.TableRow;
+import com.google.auto.value.AutoValue;
import java.io.IOException;
import java.io.InterruptedIOException;
import java.nio.charset.StandardCharsets;
@@ -30,8 +31,8 @@
import java.util.Objects;
import javax.annotation.Nonnull;
import javax.annotation.concurrent.NotThreadSafe;
-import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.io.gcp.testing.BigqueryMatcher.TableAndQuery;
import org.apache.beam.sdk.testing.SerializableMatcher;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
@@ -54,60 +55,70 @@
*/
@NotThreadSafe
@Experimental
-public class BigqueryMatcher extends TypeSafeMatcher<PipelineResult>
- implements SerializableMatcher<PipelineResult> {
+public class BigqueryMatcher extends TypeSafeMatcher<TableAndQuery>
+ implements SerializableMatcher<TableAndQuery> {
private static final Logger LOG = LoggerFactory.getLogger(BigqueryMatcher.class);
// The total number of rows in query response to be formatted for debugging purpose
private static final int TOTAL_FORMATTED_ROWS = 20;
- private final String projectId;
- private final String query;
- private final boolean usingStandardSql;
private final String expectedChecksum;
private String actualChecksum;
private transient QueryResponse response;
private BigqueryClient bigqueryClient;
- public BigqueryMatcher(
- String applicationName, String projectId, String query, String expectedChecksum) {
- this(applicationName, projectId, query, false, expectedChecksum);
- }
-
- private BigqueryMatcher(
- String applicationName,
- String projectId,
- String query,
- boolean usingStandardSql,
- String expectedChecksum) {
- validateArgument("applicationName", applicationName);
- validateArgument("projectId", projectId);
- validateArgument("query", query);
+ private BigqueryMatcher(String expectedChecksum) {
validateArgument("expectedChecksum", expectedChecksum);
- this.projectId = projectId;
- this.query = query;
- this.usingStandardSql = usingStandardSql;
this.expectedChecksum = expectedChecksum;
- this.bigqueryClient = BigqueryClient.getClient(applicationName);
}
- public static BigqueryMatcher createUsingStandardSql(
- String applicationName, String projectId, String query, String expectedChecksum) {
- return new BigqueryMatcher(applicationName, projectId, query, true, expectedChecksum);
+ public static BigqueryMatcher queryResultHasChecksum(String checksum) {
+ return new BigqueryMatcher(checksum);
+ }
+
+ public static TableAndQuery createQuery(String applicationName, String projectId, String query) {
+ return TableAndQuery.create(applicationName, projectId, query, false);
+ }
+
+ public static TableAndQuery createQueryUsingStandardSql(
+ String applicationName, String projectId, String query) {
+ return TableAndQuery.create(applicationName, projectId, query, true);
+ }
+
+ @AutoValue
+ public abstract static class TableAndQuery {
+ public static TableAndQuery create(
+ String applicationName, String projectId, String query, Boolean usingStandardSql) {
+ return new AutoValue_BigqueryMatcher_TableAndQuery(
+ applicationName, projectId, query, usingStandardSql);
+ }
+
+ public abstract String getApplicationName();
+
+ public abstract String getProjectId();
+
+ public abstract String getQuery();
+
+ public abstract Boolean getUsingStandardSql();
}
@Override
- protected boolean matchesSafely(PipelineResult pipelineResult) {
+ protected boolean matchesSafely(TableAndQuery tableAndQuery) {
+ bigqueryClient = BigqueryClient.getClient(tableAndQuery.getApplicationName());
+
LOG.info("Verifying Bigquery data");
// execute query
- LOG.debug("Executing query: {}", query);
+ LOG.debug("Executing query: {}", tableAndQuery.getQuery());
try {
- if (usingStandardSql) {
- response = bigqueryClient.queryWithRetriesUsingStandardSql(query, this.projectId);
+ if (tableAndQuery.getUsingStandardSql()) {
+ response =
+ bigqueryClient.queryWithRetriesUsingStandardSql(
+ tableAndQuery.getQuery(), tableAndQuery.getProjectId());
} else {
- response = bigqueryClient.queryWithRetries(query, this.projectId);
+ response =
+ bigqueryClient.queryWithRetries(tableAndQuery.getQuery(), tableAndQuery.getProjectId());
}
} catch (IOException | InterruptedException e) {
if (e instanceof InterruptedIOException) {
@@ -151,7 +162,7 @@
}
@Override
- public void describeMismatchSafely(PipelineResult pResult, Description description) {
+ public void describeMismatchSafely(TableAndQuery tableAndQuery, Description description) {
String info;
if (!response.getJobComplete()) {
// query job not complete
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java
index 7916513..2ab1fdc 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java
@@ -32,7 +32,6 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.concurrent.ThreadLocalRandom;
import java.util.regex.Pattern;
import javax.annotation.Nullable;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers;
@@ -90,6 +89,13 @@
}
}
+ public List<String> getAllIds(String projectId, String datasetId, String tableId)
+ throws InterruptedException, IOException {
+ synchronized (tables) {
+ return getTableContainer(projectId, datasetId, tableId).getIds();
+ }
+ }
+
private TableContainer getTableContainer(String projectId, String datasetId, String tableId)
throws InterruptedException, IOException {
synchronized (tables) {
@@ -215,7 +221,15 @@
PaneInfo.ON_TIME_AND_ONLY_FIRING));
}
return insertAll(
- ref, windowedRows, insertIdList, InsertRetryPolicy.alwaysRetry(), null, null, false, false);
+ ref,
+ windowedRows,
+ insertIdList,
+ InsertRetryPolicy.alwaysRetry(),
+ null,
+ null,
+ false,
+ false,
+ false);
}
@Override
@@ -227,17 +241,17 @@
List<ValueInSingleWindow<T>> failedInserts,
ErrorContainer<T> errorContainer,
boolean skipInvalidRows,
- boolean ignoreUnknownValues)
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds)
throws IOException, InterruptedException {
Map<TableRow, List<TableDataInsertAllResponse.InsertErrors>> insertErrors = getInsertErrors();
synchronized (tables) {
+ if (ignoreInsertIds) {
+ insertIdList = null;
+ }
+
if (insertIdList != null) {
assertEquals(rowList.size(), insertIdList.size());
- } else {
- insertIdList = Lists.newArrayListWithExpectedSize(rowList.size());
- for (int i = 0; i < rowList.size(); ++i) {
- insertIdList.add(Integer.toString(ThreadLocalRandom.current().nextInt()));
- }
}
long dataSize = 0;
@@ -258,7 +272,11 @@
}
}
if (shouldInsert) {
- dataSize += tableContainer.addRow(row, insertIdList.get(i));
+ if (insertIdList == null) {
+ dataSize += tableContainer.addRow(row, null);
+ } else {
+ dataSize += tableContainer.addRow(row, insertIdList.get(i));
+ }
} else {
errorContainer.add(
failedInserts, allErrors.get(allErrors.size() - 1), ref, rowList.get(i));
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/TableContainer.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/TableContainer.java
index 402da86..46c9f37 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/TableContainer.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/TableContainer.java
@@ -40,7 +40,9 @@
long addRow(TableRow row, String id) {
rows.add(row);
- ids.add(id);
+ if (id != null) {
+ ids.add(id);
+ }
long tableSize = table.getNumBytes() == null ? 0L : table.getNumBytes();
try {
long rowSize = TableRowJsonCoder.of().getEncodedElementByteSize(row);
@@ -58,4 +60,8 @@
List<TableRow> getRows() {
return rows;
}
+
+ List<String> getIds() {
+ return ids;
+ }
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
index 50f6548..bc146ce 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
@@ -67,6 +67,7 @@
classesInPackage("com.google.cloud.bigtable.config"),
classesInPackage("com.google.cloud.bigtable.data"),
classesInPackage("com.google.spanner.v1"),
+ classesInPackage("com.google.pubsub.v1"),
Matchers.equalTo(com.google.api.gax.rpc.ApiException.class),
Matchers.<Class<?>>equalTo(com.google.api.gax.longrunning.OperationFuture.class),
Matchers.<Class<?>>equalTo(com.google.api.gax.longrunning.OperationSnapshot.class),
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 99e3fba..624a62d 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
@@ -684,6 +684,35 @@
p.run();
}
+ @Test
+ public void testWriteWithoutInsertId() throws Exception {
+ TableRow row1 = new TableRow().set("name", "a").set("number", 1);
+ TableRow row2 = new TableRow().set("name", "b").set("number", 2);
+ TableRow row3 = new TableRow().set("name", "c").set("number", 3);
+ p.apply(Create.of(row1, row2, row3).withCoder(TableRowJsonCoder.of()))
+ .apply(
+ BigQueryIO.writeTableRows()
+ .to("project-id:dataset-id.table-id")
+ .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
+ .withMethod(BigQueryIO.Write.Method.STREAMING_INSERTS)
+ .withSchema(
+ new TableSchema()
+ .setFields(
+ ImmutableList.of(
+ new TableFieldSchema().setName("name").setType("STRING"),
+ new TableFieldSchema().setName("number").setType("INTEGER"))))
+ .withTestServices(fakeBqServices)
+ .ignoreInsertIds()
+ .withoutValidation());
+ p.run();
+ assertThat(
+ fakeDatasetService.getAllRows("project-id", "dataset-id", "table-id"),
+ containsInAnyOrder(row1, row2, row3));
+ // Verify no insert id is added.
+ assertThat(
+ fakeDatasetService.getAllIds("project-id", "dataset-id", "table-id"), containsInAnyOrder());
+ }
+
@AutoValue
abstract static class InputRecord implements Serializable {
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
index fb88bb4..bac6c23 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
@@ -517,6 +517,7 @@
null,
null,
false,
+ false,
false);
verify(response, times(2)).getStatusCode();
verify(response, times(2)).getContent();
@@ -551,6 +552,7 @@
null,
null,
false,
+ false,
false);
verify(response, times(2)).getStatusCode();
verify(response, times(2)).getContent();
@@ -601,6 +603,7 @@
null,
null,
false,
+ false,
false);
verify(response, times(2)).getStatusCode();
verify(response, times(2)).getContent();
@@ -646,6 +649,7 @@
null,
null,
false,
+ false,
false);
fail();
} catch (IOException e) {
@@ -693,6 +697,7 @@
null,
null,
false,
+ false,
false);
verify(response, times(2)).getStatusCode();
verify(response, times(2)).getContent();
@@ -761,17 +766,18 @@
failedInserts,
ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
false,
+ false,
false);
assertEquals(1, failedInserts.size());
expectedLogs.verifyInfo("Retrying 1 failed inserts to BigQuery");
}
/**
- * Tests that {@link DatasetServiceImpl#insertAll} respects the skipInvalidRows and
- * ignoreUnknownValues parameters.
+ * Tests that {@link DatasetServiceImpl#insertAll} respects the skipInvalidRows,
+ * ignoreUnknownValues and ignoreInsertIds parameters.
*/
@Test
- public void testSkipInvalidRowsIgnoreUnknownValuesStreaming()
+ public void testSkipInvalidRowsIgnoreUnknownIgnoreInsertIdsValuesStreaming()
throws InterruptedException, IOException {
TableReference ref =
new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
@@ -790,7 +796,7 @@
DatasetServiceImpl dataService =
new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
- // First, test with both flags disabled
+ // First, test with all flags disabled
dataService.insertAll(
ref,
rows,
@@ -801,6 +807,7 @@
Lists.newArrayList(),
ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
false,
+ false,
false);
TableDataInsertAllRequest parsedRequest =
@@ -809,7 +816,7 @@
assertFalse(parsedRequest.getSkipInvalidRows());
assertFalse(parsedRequest.getIgnoreUnknownValues());
- // Then with both enabled
+ // Then with all enabled
dataService.insertAll(
ref,
rows,
@@ -820,12 +827,15 @@
Lists.newArrayList(),
ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
true,
+ true,
true);
parsedRequest = fromString(request.getContentAsString(), TableDataInsertAllRequest.class);
assertTrue(parsedRequest.getSkipInvalidRows());
assertTrue(parsedRequest.getIgnoreUnknownValues());
+ assertNull(parsedRequest.getRows().get(0).getInsertId());
+ assertNull(parsedRequest.getRows().get(1).getInsertId());
}
/** A helper to convert a string response back to a {@link GenericJson} subclass. */
@@ -1002,6 +1012,7 @@
failedInserts,
ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
false,
+ false,
false);
assertThat(failedInserts, is(rows));
@@ -1056,6 +1067,7 @@
failedInserts,
ErrorContainer.BIG_QUERY_INSERT_ERROR_ERROR_CONTAINER,
false,
+ false,
false);
assertThat(failedInserts, is(expected));
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilTest.java
index 4dddf0d..98b8bc8 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilTest.java
@@ -205,7 +205,7 @@
try {
totalBytes =
datasetService.insertAll(
- ref, rows, ids, InsertRetryPolicy.alwaysRetry(), null, null, false, false);
+ ref, rows, ids, InsertRetryPolicy.alwaysRetry(), null, null, false, false, false);
} finally {
verifyInsertAll(5);
// Each of the 25 rows has 1 byte for length and 30 bytes: '{"f":[{"v":"foo"},{"v":1234}]}'
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java
index 7c53170..4dd719b 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java
@@ -142,11 +142,11 @@
List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
assertEquals(1, acutalMessages.size());
IncomingMessage actualMessage = acutalMessages.get(0);
- assertEquals(ACK_ID, actualMessage.ackId);
- assertEquals(DATA, new String(actualMessage.elementBytes, StandardCharsets.UTF_8));
- assertEquals(RECORD_ID, actualMessage.recordId);
- assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch);
- assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch);
+ assertEquals(ACK_ID, actualMessage.ackId());
+ assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
+ assertEquals(RECORD_ID, actualMessage.recordId());
+ assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch());
+ assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch());
assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
} finally {
server.shutdownNow();
@@ -187,8 +187,13 @@
InProcessServerBuilder.forName(channelName).addService(publisherImplBase).build().start();
try {
OutgoingMessage actualMessage =
- new OutgoingMessage(
- DATA.getBytes(StandardCharsets.UTF_8), ATTRIBUTES, MESSAGE_TIME, RECORD_ID);
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .putAllAttributes(ATTRIBUTES)
+ .build(),
+ MESSAGE_TIME,
+ RECORD_ID);
int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
assertEquals(1, n);
assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOExternalTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOExternalTest.java
index 50f7528..abecf89 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOExternalTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOExternalTest.java
@@ -36,8 +36,8 @@
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.hamcrest.Matchers;
import org.junit.Test;
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
index 65b89a7..0dc910f 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
@@ -29,6 +29,7 @@
import static org.junit.Assert.assertThat;
import com.google.api.client.util.Clock;
+import com.google.protobuf.ByteString;
import java.io.IOException;
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
@@ -391,9 +392,10 @@
})
.map(
ba ->
- new IncomingMessage(
- ba,
- null,
+ IncomingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFrom(ba))
+ .build(),
1234L,
0,
UUID.randomUUID().toString(),
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java
index f7fc0f3..aad9729 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java
@@ -34,6 +34,7 @@
import com.google.api.services.pubsub.model.ReceivedMessage;
import com.google.api.services.pubsub.model.Subscription;
import com.google.api.services.pubsub.model.Topic;
+import com.google.protobuf.ByteString;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
@@ -73,6 +74,7 @@
private static final String DATA = "testData";
private static final String RECORD_ID = "testRecordId";
private static final String ACK_ID = "testAckId";
+ private static final String ORDERING_KEY = "testOrderingKey";
@Before
public void setup() {
@@ -98,7 +100,8 @@
.setPublishTime(String.valueOf(PUB_TIME))
.setAttributes(
ImmutableMap.of(
- TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME), ID_ATTRIBUTE, RECORD_ID));
+ TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME), ID_ATTRIBUTE, RECORD_ID))
+ .set("orderingKey", ORDERING_KEY);
ReceivedMessage expectedReceivedMessage =
new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
PullResponse expectedResponse =
@@ -113,11 +116,42 @@
List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
assertEquals(1, acutalMessages.size());
IncomingMessage actualMessage = acutalMessages.get(0);
- assertEquals(ACK_ID, actualMessage.ackId);
- assertEquals(DATA, new String(actualMessage.elementBytes, StandardCharsets.UTF_8));
- assertEquals(RECORD_ID, actualMessage.recordId);
- assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch);
- assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch);
+ assertEquals(ACK_ID, actualMessage.ackId());
+ assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
+ assertEquals(RECORD_ID, actualMessage.recordId());
+ assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch());
+ assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch());
+ assertEquals(ORDERING_KEY, actualMessage.message().getOrderingKey());
+ }
+
+ @Test
+ public void pullOneMessageEmptyAttributes() throws IOException {
+ client = new PubsubJsonClient(null, null, mockPubsub);
+ String expectedSubscription = SUBSCRIPTION.getPath();
+ PullRequest expectedRequest = new PullRequest().setReturnImmediately(true).setMaxMessages(10);
+ PubsubMessage expectedPubsubMessage =
+ new PubsubMessage()
+ .setMessageId(MESSAGE_ID)
+ .encodeData(DATA.getBytes(StandardCharsets.UTF_8))
+ .setPublishTime(String.valueOf(PUB_TIME));
+ ReceivedMessage expectedReceivedMessage =
+ new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
+ PullResponse expectedResponse =
+ new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
+ when((Object)
+ (mockPubsub
+ .projects()
+ .subscriptions()
+ .pull(expectedSubscription, expectedRequest)
+ .execute()))
+ .thenReturn(expectedResponse);
+ List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
+ assertEquals(1, acutalMessages.size());
+ IncomingMessage actualMessage = acutalMessages.get(0);
+ assertEquals(ACK_ID, actualMessage.ackId());
+ assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
+ assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch());
+ assertEquals(PUB_TIME, actualMessage.timestampMsSinceEpoch());
}
@Test
@@ -146,7 +180,7 @@
List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
assertEquals(1, acutalMessages.size());
IncomingMessage actualMessage = acutalMessages.get(0);
- assertArrayEquals(new byte[0], actualMessage.elementBytes);
+ assertArrayEquals(new byte[0], actualMessage.message().getData().toByteArray());
}
@Test
@@ -160,7 +194,8 @@
.put(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME))
.put(ID_ATTRIBUTE, RECORD_ID)
.put("k", "v")
- .build());
+ .build())
+ .set("orderingKey", ORDERING_KEY);
PublishRequest expectedRequest =
new PublishRequest().setMessages(ImmutableList.of(expectedPubsubMessage));
PublishResponse expectedResponse =
@@ -171,7 +206,14 @@
Map<String, String> attrs = new HashMap<>();
attrs.put("k", "v");
OutgoingMessage actualMessage =
- new OutgoingMessage(DATA.getBytes(StandardCharsets.UTF_8), attrs, MESSAGE_TIME, RECORD_ID);
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .putAllAttributes(attrs)
+ .setOrderingKey(ORDERING_KEY)
+ .build(),
+ MESSAGE_TIME,
+ RECORD_ID);
int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
assertEquals(1, n);
}
@@ -195,8 +237,12 @@
(mockPubsub.projects().topics().publish(expectedTopic, expectedRequest).execute()))
.thenReturn(expectedResponse);
OutgoingMessage actualMessage =
- new OutgoingMessage(
- DATA.getBytes(StandardCharsets.UTF_8), ImmutableMap.of(), MESSAGE_TIME, RECORD_ID);
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .build(),
+ MESSAGE_TIME,
+ RECORD_ID);
int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
assertEquals(1, n);
}
@@ -222,7 +268,13 @@
Map<String, String> attrs = new HashMap<>();
attrs.put("k", "v");
OutgoingMessage actualMessage =
- new OutgoingMessage(DATA.getBytes(StandardCharsets.UTF_8), attrs, MESSAGE_TIME, RECORD_ID);
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .putAllAttributes(attrs)
+ .build(),
+ MESSAGE_TIME,
+ RECORD_ID);
int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
assertEquals(1, n);
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java
index 2b698f0..6b920e8 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java
@@ -20,8 +20,9 @@
import static org.junit.Assert.assertEquals;
import com.google.api.client.util.Clock;
+import com.google.protobuf.ByteString;
+import com.google.pubsub.v1.PubsubMessage;
import java.io.IOException;
-import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage;
@@ -54,9 +55,8 @@
final AtomicLong now = new AtomicLong();
Clock clock = now::get;
IncomingMessage expectedIncomingMessage =
- new IncomingMessage(
- DATA.getBytes(StandardCharsets.UTF_8),
- null,
+ IncomingMessage.of(
+ PubsubMessage.newBuilder().setData(ByteString.copyFromUtf8(DATA)).build(),
MESSAGE_TIME,
REQ_TIME,
ACK_ID,
@@ -75,7 +75,14 @@
client.advance();
incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
assertEquals(1, incomingMessages.size());
- assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
+ assertEquals(
+ IncomingMessage.of(
+ expectedIncomingMessage.message(),
+ expectedIncomingMessage.timestampMsSinceEpoch(),
+ now.get(),
+ expectedIncomingMessage.ackId(),
+ expectedIncomingMessage.recordId()),
+ incomingMessages.get(0));
now.addAndGet(10 * 1000);
client.advance();
// Extend ack
@@ -85,7 +92,14 @@
client.advance();
incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
assertEquals(1, incomingMessages.size());
- assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
+ assertEquals(
+ IncomingMessage.of(
+ expectedIncomingMessage.message(),
+ expectedIncomingMessage.timestampMsSinceEpoch(),
+ now.get(),
+ expectedIncomingMessage.ackId(),
+ expectedIncomingMessage.recordId()),
+ incomingMessages.get(0));
// Extend ack
client.modifyAckDeadline(SUBSCRIPTION, ImmutableList.of(ACK_ID), 20);
// Ack
@@ -99,7 +113,10 @@
@Test
public void publishOneMessage() throws IOException {
OutgoingMessage expectedOutgoingMessage =
- new OutgoingMessage(DATA.getBytes(StandardCharsets.UTF_8), null, MESSAGE_TIME, MESSAGE_ID);
+ OutgoingMessage.of(
+ PubsubMessage.newBuilder().setData(ByteString.copyFromUtf8(DATA)).build(),
+ MESSAGE_TIME,
+ MESSAGE_ID);
try (PubsubTestClientFactory factory =
PubsubTestClient.createFactoryForPublish(
TOPIC, Sets.newHashSet(expectedOutgoingMessage), ImmutableList.of())) {
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java
index f588e05..f8cd86e 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.gcp.pubsub;
+import com.google.protobuf.ByteString;
import java.io.IOException;
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
@@ -83,8 +84,12 @@
@Test
public void saneCoder() throws Exception {
OutgoingMessage message =
- new OutgoingMessage(
- DATA.getBytes(StandardCharsets.UTF_8), ImmutableMap.of(), TIMESTAMP, getRecordId(DATA));
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .build(),
+ TIMESTAMP,
+ getRecordId(DATA));
CoderProperties.coderDecodeEncodeEqual(PubsubUnboundedSink.CODER, message);
CoderProperties.coderSerializable(PubsubUnboundedSink.CODER);
}
@@ -93,8 +98,13 @@
public void sendOneMessage() throws IOException {
List<OutgoingMessage> outgoing =
ImmutableList.of(
- new OutgoingMessage(
- DATA.getBytes(StandardCharsets.UTF_8), ATTRIBUTES, TIMESTAMP, getRecordId(DATA)));
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .putAllAttributes(ATTRIBUTES)
+ .build(),
+ TIMESTAMP,
+ getRecordId(DATA)));
int batchSize = 1;
int batchBytes = 1;
try (PubsubTestClientFactory factory =
@@ -121,9 +131,10 @@
public void sendOneMessageWithoutAttributes() throws IOException {
List<OutgoingMessage> outgoing =
ImmutableList.of(
- new OutgoingMessage(
- DATA.getBytes(StandardCharsets.UTF_8),
- null /* attributes */,
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .build(),
TIMESTAMP,
getRecordId(DATA)));
try (PubsubTestClientFactory factory =
@@ -157,9 +168,10 @@
for (int i = 0; i < batchSize * 10; i++) {
String str = String.valueOf(i);
outgoing.add(
- new OutgoingMessage(
- str.getBytes(StandardCharsets.UTF_8),
- ImmutableMap.of(),
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(str))
+ .build(),
TIMESTAMP,
getRecordId(str)));
data.add(str);
@@ -198,9 +210,10 @@
}
String str = sb.toString();
outgoing.add(
- new OutgoingMessage(
- str.getBytes(StandardCharsets.UTF_8),
- ImmutableMap.of(),
+ OutgoingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(str))
+ .build(),
TIMESTAMP,
getRecordId(str)));
data.add(str);
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java
index b2dacf0..43ecbdc 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java
@@ -31,6 +31,7 @@
import static org.junit.Assert.assertTrue;
import com.google.api.client.util.Clock;
+import com.google.protobuf.ByteString;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
@@ -100,8 +101,14 @@
private void setupOneMessage() {
setupOneMessage(
ImmutableList.of(
- new IncomingMessage(
- DATA.getBytes(StandardCharsets.UTF_8), null, TIMESTAMP, 0, ACK_ID, RECORD_ID)));
+ IncomingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(DATA))
+ .build(),
+ TIMESTAMP,
+ 0,
+ ACK_ID,
+ RECORD_ID)));
}
@After
@@ -219,8 +226,14 @@
String data = String.format("data_%d", i);
String ackid = String.format("ackid_%d", i);
incoming.add(
- new IncomingMessage(
- data.getBytes(StandardCharsets.UTF_8), null, TIMESTAMP, 0, ackid, RECORD_ID));
+ IncomingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(data))
+ .build(),
+ TIMESTAMP,
+ 0,
+ ackid,
+ RECORD_ID));
}
setupOneMessage(incoming);
PubsubReader reader = primSource.createReader(p.getOptions(), null);
@@ -279,9 +292,10 @@
String recid = String.format("recordid_%d", messageNum);
String ackId = String.format("ackid_%d", messageNum);
incoming.add(
- new IncomingMessage(
- data.getBytes(StandardCharsets.UTF_8),
- null,
+ IncomingMessage.of(
+ com.google.pubsub.v1.PubsubMessage.newBuilder()
+ .setData(ByteString.copyFromUtf8(data))
+ .build(),
messageNumToTimestamp(messageNum),
0,
ackId,
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/storage/GcsKmsKeyIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/storage/GcsKmsKeyIT.java
index 9f092e7..66fe8bf 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/storage/GcsKmsKeyIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/storage/GcsKmsKeyIT.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.gcp.storage;
+import static org.apache.beam.sdk.testing.FileChecksumMatcher.fileContentsHaveChecksum;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.Assert.assertNotNull;
@@ -36,10 +37,10 @@
import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
-import org.apache.beam.sdk.testing.FileChecksumMatcher;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
import org.apache.beam.sdk.testing.UsesKms;
+import org.apache.beam.sdk.util.NumberedShardedFile;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.junit.Test;
import org.junit.experimental.categories.Category;
@@ -87,7 +88,7 @@
assertThat(state, equalTo(State.DONE));
String filePattern = filenamePrefix + "*-of-*";
- assertThat(result, new FileChecksumMatcher(EXPECTED_CHECKSUM, filePattern));
+ assertThat(new NumberedShardedFile(filePattern), fileContentsHaveChecksum(EXPECTED_CHECKSUM));
// Verify objects have KMS key set.
try {
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcherTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcherTest.java
index 6bf6f92..508dfec 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcherTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/testing/BigqueryMatcherTest.java
@@ -26,7 +26,6 @@
import com.google.api.services.bigquery.model.TableCell;
import com.google.api.services.bigquery.model.TableRow;
import java.math.BigInteger;
-import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.junit.Before;
import org.junit.Rule;
@@ -49,7 +48,6 @@
@Rule public ExpectedException thrown = ExpectedException.none();
@Mock private BigqueryClient mockBigqueryClient;
- @Mock private PipelineResult mockResult;
@Before
public void setUp() {
@@ -61,19 +59,16 @@
@Test
public void testBigqueryMatcherThatSucceeds() throws Exception {
BigqueryMatcher matcher =
- spy(
- new BigqueryMatcher(
- appName, projectId, query, "9bb47f5c90d2a99cad526453dff5ed5ec74650dc"));
+ spy(BigqueryMatcher.queryResultHasChecksum("9bb47f5c90d2a99cad526453dff5ed5ec74650dc"));
when(mockBigqueryClient.queryWithRetries(anyString(), anyString()))
.thenReturn(createResponseContainingTestData());
- assertThat(mockResult, matcher);
+ assertThat(BigqueryMatcher.createQuery(appName, projectId, query), matcher);
}
@Test
public void testBigqueryMatcherFailsForChecksumMismatch() throws Exception {
- BigqueryMatcher matcher =
- spy(new BigqueryMatcher(appName, projectId, query, "incorrect-checksum"));
+ BigqueryMatcher matcher = spy(BigqueryMatcher.queryResultHasChecksum("incorrect-checksum"));
when(mockBigqueryClient.queryWithRetries(anyString(), anyString()))
.thenReturn(createResponseContainingTestData());
@@ -82,12 +77,12 @@
thrown.expectMessage("Total number of rows are: 1");
thrown.expectMessage("abc");
- assertThat(mockResult, matcher);
+ assertThat(BigqueryMatcher.createQuery(appName, projectId, query), matcher);
}
@Test
public void testBigqueryMatcherFailsWhenQueryJobNotComplete() throws Exception {
- BigqueryMatcher matcher = spy(new BigqueryMatcher(appName, projectId, query, "some-checksum"));
+ BigqueryMatcher matcher = spy(BigqueryMatcher.queryResultHasChecksum("some-checksum"));
when(mockBigqueryClient.queryWithRetries(anyString(), anyString()))
.thenReturn(new QueryResponse().setJobComplete(false));
@@ -95,7 +90,7 @@
thrown.expectMessage("The query job hasn't completed.");
thrown.expectMessage("jobComplete=false");
- assertThat(mockResult, matcher);
+ assertThat(BigqueryMatcher.createQuery(appName, projectId, query), matcher);
}
private QueryResponse createResponseContainingTestData() {
diff --git a/sdks/java/io/hadoop-common/build.gradle b/sdks/java/io/hadoop-common/build.gradle
index 08f60c6..c4d6cad 100644
--- a/sdks/java/io/hadoop-common/build.gradle
+++ b/sdks/java/io/hadoop-common/build.gradle
@@ -27,7 +27,6 @@
provided library.java.hadoop_client
provided library.java.hadoop_common
provided library.java.hadoop_mapreduce_client_core
- testCompile library.java.commons_lang3
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
testCompile library.java.junit
diff --git a/sdks/java/io/hadoop-common/src/test/java/org/apache/beam/sdk/io/hadoop/SerializableConfigurationTest.java b/sdks/java/io/hadoop-common/src/test/java/org/apache/beam/sdk/io/hadoop/SerializableConfigurationTest.java
index f85cf6b..07b5b7f 100644
--- a/sdks/java/io/hadoop-common/src/test/java/org/apache/beam/sdk/io/hadoop/SerializableConfigurationTest.java
+++ b/sdks/java/io/hadoop-common/src/test/java/org/apache/beam/sdk/io/hadoop/SerializableConfigurationTest.java
@@ -20,7 +20,7 @@
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
-import org.apache.commons.lang3.SerializationUtils;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.SerializationUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.Job;
import org.junit.Rule;
diff --git a/sdks/java/io/hadoop-file-system/src/main/java/org/apache/beam/sdk/io/hdfs/HadoopResourceId.java b/sdks/java/io/hadoop-file-system/src/main/java/org/apache/beam/sdk/io/hdfs/HadoopResourceId.java
index 4e5f6e2..b68839b 100644
--- a/sdks/java/io/hadoop-file-system/src/main/java/org/apache/beam/sdk/io/hdfs/HadoopResourceId.java
+++ b/sdks/java/io/hadoop-file-system/src/main/java/org/apache/beam/sdk/io/hdfs/HadoopResourceId.java
@@ -65,6 +65,10 @@
@Override
public String getFilename() {
+ if (isDirectory()) {
+ Path parentPath = new Path(uri).getParent();
+ return parentPath == null ? null : parentPath.getName();
+ }
return new Path(uri).getName();
}
diff --git a/sdks/java/io/hadoop-file-system/src/test/java/org/apache/beam/sdk/io/hdfs/HadoopResourceIdTest.java b/sdks/java/io/hadoop-file-system/src/test/java/org/apache/beam/sdk/io/hdfs/HadoopResourceIdTest.java
index 4d7fb8d..1726a3e 100644
--- a/sdks/java/io/hadoop-file-system/src/test/java/org/apache/beam/sdk/io/hdfs/HadoopResourceIdTest.java
+++ b/sdks/java/io/hadoop-file-system/src/test/java/org/apache/beam/sdk/io/hdfs/HadoopResourceIdTest.java
@@ -17,6 +17,9 @@
*/
package org.apache.beam.sdk.io.hdfs;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
import java.net.URI;
import java.util.Collections;
import org.apache.beam.sdk.io.FileSystems;
@@ -65,4 +68,16 @@
"hdfs://" + hdfsClusterBaseUri.getPath(), true /* isDirectory */);
ResourceIdTester.runResourceIdBattery(baseDirectory);
}
+
+ @Test
+ public void testGetFilename() {
+ assertNull(toResourceIdentifier("").getFilename());
+ assertEquals("abc", toResourceIdentifier("/dirA/abc").getFilename());
+ assertEquals("abc", toResourceIdentifier("/dirA/abc/").getFilename());
+ assertEquals("xyz.txt", toResourceIdentifier("/dirA/abc/xyz.txt").getFilename());
+ }
+
+ private ResourceId toResourceIdentifier(String path) {
+ return new HadoopResourceId(hdfsClusterBaseUri.resolve(path));
+ }
}
diff --git a/sdks/java/io/hadoop-format/build.gradle b/sdks/java/io/hadoop-format/build.gradle
index 1e0cf94..8adf06e 100644
--- a/sdks/java/io/hadoop-format/build.gradle
+++ b/sdks/java/io/hadoop-format/build.gradle
@@ -75,9 +75,10 @@
}
// elasticsearch-hadoop 5.0.0 uses commons-httpclient's URIException
testCompile "commons-httpclient:commons-httpclient:3.1"
+ testCompile library.java.commons_io
testCompile library.java.cassandra_driver_core
testCompile library.java.cassandra_driver_mapping
- testCompile "org.apache.cassandra:cassandra-all:3.11.3"
+ testCompile "org.apache.cassandra:cassandra-all:3.11.5"
testCompile library.java.postgres
testCompile "org.apache.logging.log4j:log4j-core:$log4j_version"
testCompile library.java.junit
@@ -85,7 +86,6 @@
testCompile library.java.hamcrest_library
testRuntimeOnly library.java.slf4j_jdk14
testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
- compile library.java.commons_io_2x
delegate.add("sparkRunner", project(":sdks:java:io:hadoop-format"))
delegate.add("sparkRunner", project(path: ":sdks:java:io:hadoop-format", configuration: "testRuntime"))
diff --git a/sdks/java/io/hbase/build.gradle b/sdks/java/io/hbase/build.gradle
index e2dd902..d33e0ca 100644
--- a/sdks/java/io/hbase/build.gradle
+++ b/sdks/java/io/hbase/build.gradle
@@ -44,7 +44,6 @@
compile "org.apache.hbase:hbase-shaded-client:$hbase_version"
testCompile project(path: ":sdks:java:io:common", configuration: "testRuntime")
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
- testCompile library.java.commons_lang3
testCompile library.java.junit
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
diff --git a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java
index 5f29185..eb8cec7 100644
--- a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java
+++ b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java
@@ -31,6 +31,7 @@
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.StringUtils;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.hbase.HBaseIO.HBaseSource;
@@ -43,7 +44,6 @@
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.HBaseTestingUtility;
diff --git a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/SerializableScanTest.java b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/SerializableScanTest.java
index 0085ea1..a5258cb 100644
--- a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/SerializableScanTest.java
+++ b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/SerializableScanTest.java
@@ -21,7 +21,7 @@
import static org.junit.Assert.assertNotNull;
import java.nio.charset.StandardCharsets;
-import org.apache.commons.lang3.SerializationUtils;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.SerializationUtils;
import org.apache.hadoop.hbase.client.Scan;
import org.junit.Rule;
import org.junit.Test;
diff --git a/sdks/java/io/hcatalog/build.gradle b/sdks/java/io/hcatalog/build.gradle
index e32277d..e0b43a2 100644
--- a/sdks/java/io/hcatalog/build.gradle
+++ b/sdks/java/io/hcatalog/build.gradle
@@ -59,7 +59,7 @@
exclude group: "com.google.protobuf", module: "protobuf-java"
}
testCompile project(":sdks:java:io:common").sourceSets.test.output
- testCompile library.java.commons_io_2x
+ testCompile library.java.commons_io
testCompile library.java.junit
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOExternalTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOExternalTest.java
index a7b7f8a..500673c 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOExternalTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOExternalTest.java
@@ -39,8 +39,8 @@
import org.apache.beam.sdk.transforms.Impulse;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver;
+import org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.stub.StreamObserver;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java
index 37cf836..2c0ad51 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.sdk.io.kinesis;
-import static org.apache.commons.lang.builder.HashCodeBuilder.reflectionHashCode;
+import static org.apache.commons.lang3.builder.HashCodeBuilder.reflectionHashCode;
import com.amazonaws.services.kinesis.clientlibrary.types.ExtendedSequenceNumber;
import com.amazonaws.services.kinesis.clientlibrary.types.UserRecord;
diff --git a/sdks/java/io/kudu/build.gradle b/sdks/java/io/kudu/build.gradle
index b32a02d..0090e13 100644
--- a/sdks/java/io/kudu/build.gradle
+++ b/sdks/java/io/kudu/build.gradle
@@ -32,7 +32,7 @@
}
}
-def kudu_version = "1.9.0"
+def kudu_version = "1.11.1"
dependencies {
compile library.java.vendored_guava_26_0_jre
diff --git a/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java b/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
index 726acf6..bef77e4 100644
--- a/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
+++ b/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
@@ -28,6 +28,7 @@
import java.nio.channels.WritableByteChannel;
import javax.annotation.Nullable;
import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericRecord;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AvroCoder;
@@ -73,6 +74,11 @@
*
* <p>As {@link Read} is based on {@link FileIO}, it supports any filesystem (hdfs, ...).
*
+ * <p>When using schemas created via reflection, it may be useful to generate {@link GenericRecord}
+ * instances rather than instances of the class associated with the schema. {@link Read} and {@link
+ * ReadFiles} provide {@link ParquetIO.Read#withAvroDataModel(GenericData)} allowing implementations
+ * to set the data model associated with the {@link AvroParquetReader}
+ *
* <p>For more advanced use cases, like reading each file in a {@link PCollection} of {@link
* FileIO.ReadableFile}, use the {@link ReadFiles} transform.
*
@@ -144,6 +150,9 @@
@Nullable
abstract Schema getSchema();
+ @Nullable
+ abstract GenericData getAvroDataModel();
+
abstract Builder toBuilder();
@AutoValue.Builder
@@ -152,6 +161,8 @@
abstract Builder setSchema(Schema schema);
+ abstract Builder setAvroDataModel(GenericData model);
+
abstract Read build();
}
@@ -165,6 +176,13 @@
return from(ValueProvider.StaticValueProvider.of(filepattern));
}
+ /**
+ * Define the Avro data model; see {@link AvroParquetReader.Builder#withDataModel(GenericData)}.
+ */
+ public Read withAvroDataModel(GenericData model) {
+ return toBuilder().setAvroDataModel(model).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PBegin input) {
checkNotNull(getFilepattern(), "Filepattern cannot be null.");
@@ -173,7 +191,7 @@
.apply("Create filepattern", Create.ofProvider(getFilepattern(), StringUtf8Coder.of()))
.apply(FileIO.matchAll())
.apply(FileIO.readMatches())
- .apply(readFiles(getSchema()));
+ .apply(readFiles(getSchema()).withAvroDataModel(getAvroDataModel()));
}
@Override
@@ -192,21 +210,43 @@
@Nullable
abstract Schema getSchema();
+ @Nullable
+ abstract GenericData getAvroDataModel();
+
+ abstract Builder toBuilder();
+
@AutoValue.Builder
abstract static class Builder {
abstract Builder setSchema(Schema schema);
+ abstract Builder setAvroDataModel(GenericData model);
+
abstract ReadFiles build();
}
+ /**
+ * Define the Avro data model; see {@link AvroParquetReader.Builder#withDataModel(GenericData)}.
+ */
+ public ReadFiles withAvroDataModel(GenericData model) {
+ return toBuilder().setAvroDataModel(model).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PCollection<FileIO.ReadableFile> input) {
checkNotNull(getSchema(), "Schema can not be null");
- return input.apply(ParDo.of(new ReadFn())).setCoder(AvroCoder.of(getSchema()));
+ return input
+ .apply(ParDo.of(new ReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
}
static class ReadFn extends DoFn<FileIO.ReadableFile, GenericRecord> {
+ private Class<? extends GenericData> modelClass;
+
+ ReadFn(GenericData model) {
+ this.modelClass = model != null ? model.getClass() : null;
+ }
+
@ProcessElement
public void processElement(ProcessContext processContext) throws Exception {
FileIO.ReadableFile file = processContext.element();
@@ -218,9 +258,14 @@
SeekableByteChannel seekableByteChannel = file.openSeekable();
- try (ParquetReader<GenericRecord> reader =
- AvroParquetReader.<GenericRecord>builder(new BeamParquetInputFile(seekableByteChannel))
- .build()) {
+ AvroParquetReader.Builder builder =
+ AvroParquetReader.<GenericRecord>builder(new BeamParquetInputFile(seekableByteChannel));
+ if (modelClass != null) {
+ // all GenericData implementations have a static get method
+ builder = builder.withDataModel((GenericData) modelClass.getMethod("get").invoke(null));
+ }
+
+ try (ParquetReader<GenericRecord> reader = builder.build()) {
GenericRecord read;
while ((read = reader.read()) != null) {
processContext.output(read);
diff --git a/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java b/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java
index ed80f90..6840260 100644
--- a/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java
+++ b/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java
@@ -23,8 +23,10 @@
import java.util.ArrayList;
import java.util.List;
import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.generic.GenericRecordBuilder;
+import org.apache.avro.reflect.ReflectData;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.io.FileIO;
import org.apache.beam.sdk.testing.PAssert;
@@ -126,4 +128,57 @@
Assert.assertThat(displayData, hasDisplayItem("filePattern", "foo.parquet"));
}
+
+ public static class TestRecord {
+ String name;
+
+ public TestRecord(String name) {
+ this.name = name;
+ }
+ }
+
+ @Test(expected = org.apache.beam.sdk.Pipeline.PipelineExecutionException.class)
+ public void testWriteAndReadUsingReflectDataSchemaWithoutDataModelThrowsException() {
+ Schema testRecordSchema = ReflectData.get().getSchema(TestRecord.class);
+
+ List<GenericRecord> records = generateGenericRecords(1000);
+ mainPipeline
+ .apply(Create.of(records).withCoder(AvroCoder.of(testRecordSchema)))
+ .apply(
+ FileIO.<GenericRecord>write()
+ .via(ParquetIO.sink(testRecordSchema))
+ .to(temporaryFolder.getRoot().getAbsolutePath()));
+ mainPipeline.run().waitUntilFinish();
+
+ PCollection<GenericRecord> readBack =
+ readPipeline.apply(
+ ParquetIO.read(testRecordSchema)
+ .from(temporaryFolder.getRoot().getAbsolutePath() + "/*"));
+
+ PAssert.that(readBack).containsInAnyOrder(records);
+ readPipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testWriteAndReadUsingReflectDataSchemaWithDataModel() {
+ Schema testRecordSchema = ReflectData.get().getSchema(TestRecord.class);
+
+ List<GenericRecord> records = generateGenericRecords(1000);
+ mainPipeline
+ .apply(Create.of(records).withCoder(AvroCoder.of(testRecordSchema)))
+ .apply(
+ FileIO.<GenericRecord>write()
+ .via(ParquetIO.sink(testRecordSchema))
+ .to(temporaryFolder.getRoot().getAbsolutePath()));
+ mainPipeline.run().waitUntilFinish();
+
+ PCollection<GenericRecord> readBack =
+ readPipeline.apply(
+ ParquetIO.read(testRecordSchema)
+ .withAvroDataModel(GenericData.get())
+ .from(temporaryFolder.getRoot().getAbsolutePath() + "/*"));
+
+ PAssert.that(readBack).containsInAnyOrder(records);
+ readPipeline.run().waitUntilFinish();
+ }
}
diff --git a/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java b/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
index 9612487..57617c0 100644
--- a/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
+++ b/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
@@ -114,6 +114,7 @@
.setConnectionConfiguration(RedisConnectionConfiguration.create())
.setKeyPattern("*")
.setBatchSize(1000)
+ .setOutputParallelization(true)
.build();
}
@@ -125,6 +126,7 @@
return new AutoValue_RedisIO_ReadAll.Builder()
.setConnectionConfiguration(RedisConnectionConfiguration.create())
.setBatchSize(1000)
+ .setOutputParallelization(true)
.build();
}
@@ -150,7 +152,9 @@
abstract int batchSize();
- abstract Builder builder();
+ abstract boolean outputParallelization();
+
+ abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
@@ -162,41 +166,53 @@
abstract Builder setBatchSize(int batchSize);
+ abstract Builder setOutputParallelization(boolean outputParallelization);
+
abstract Read build();
}
public Read withEndpoint(String host, int port) {
checkArgument(host != null, "host can not be null");
checkArgument(0 < port && port < 65536, "port must be a positive integer less than 65536");
- return builder()
+ return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withHost(host).withPort(port))
.build();
}
public Read withAuth(String auth) {
checkArgument(auth != null, "auth can not be null");
- return builder().setConnectionConfiguration(connectionConfiguration().withAuth(auth)).build();
+ return toBuilder()
+ .setConnectionConfiguration(connectionConfiguration().withAuth(auth))
+ .build();
}
public Read withTimeout(int timeout) {
checkArgument(timeout >= 0, "timeout can not be negative");
- return builder()
+ return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withTimeout(timeout))
.build();
}
public Read withKeyPattern(String keyPattern) {
checkArgument(keyPattern != null, "keyPattern can not be null");
- return builder().setKeyPattern(keyPattern).build();
+ return toBuilder().setKeyPattern(keyPattern).build();
}
public Read withConnectionConfiguration(RedisConnectionConfiguration connection) {
checkArgument(connection != null, "connection can not be null");
- return builder().setConnectionConfiguration(connection).build();
+ return toBuilder().setConnectionConfiguration(connection).build();
}
public Read withBatchSize(int batchSize) {
- return builder().setBatchSize(batchSize).build();
+ return toBuilder().setBatchSize(batchSize).build();
+ }
+
+ /**
+ * Whether to reshuffle the resulting PCollection so results are distributed to all workers. The
+ * default is to parallelize and should only be changed if this is known to be unnecessary.
+ */
+ public Read withOutputParallelization(boolean outputParallelization) {
+ return toBuilder().setOutputParallelization(outputParallelization).build();
}
@Override
@@ -214,7 +230,8 @@
.apply(
RedisIO.readAll()
.withConnectionConfiguration(connectionConfiguration())
- .withBatchSize(batchSize()));
+ .withBatchSize(batchSize())
+ .withOutputParallelization(outputParallelization()));
}
}
@@ -228,14 +245,18 @@
abstract int batchSize();
- abstract ReadAll.Builder builder();
+ abstract boolean outputParallelization();
+
+ abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
@Nullable
- abstract ReadAll.Builder setConnectionConfiguration(RedisConnectionConfiguration connection);
+ abstract Builder setConnectionConfiguration(RedisConnectionConfiguration connection);
- abstract ReadAll.Builder setBatchSize(int batchSize);
+ abstract Builder setBatchSize(int batchSize);
+
+ abstract Builder setOutputParallelization(boolean outputParallelization);
abstract ReadAll build();
}
@@ -243,44 +264,57 @@
public ReadAll withEndpoint(String host, int port) {
checkArgument(host != null, "host can not be null");
checkArgument(port > 0, "port can not be negative or 0");
- return builder()
+ return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withHost(host).withPort(port))
.build();
}
public ReadAll withAuth(String auth) {
checkArgument(auth != null, "auth can not be null");
- return builder().setConnectionConfiguration(connectionConfiguration().withAuth(auth)).build();
+ return toBuilder()
+ .setConnectionConfiguration(connectionConfiguration().withAuth(auth))
+ .build();
}
public ReadAll withTimeout(int timeout) {
checkArgument(timeout >= 0, "timeout can not be negative");
- return builder()
+ return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withTimeout(timeout))
.build();
}
public ReadAll withConnectionConfiguration(RedisConnectionConfiguration connection) {
checkArgument(connection != null, "connection can not be null");
- return builder().setConnectionConfiguration(connection).build();
+ return toBuilder().setConnectionConfiguration(connection).build();
}
public ReadAll withBatchSize(int batchSize) {
- return builder().setBatchSize(batchSize).build();
+ return toBuilder().setBatchSize(batchSize).build();
+ }
+
+ /**
+ * Whether to reshuffle the resulting PCollection so results are distributed to all workers. The
+ * default is to parallelize and should only be changed if this is known to be unnecessary.
+ */
+ public ReadAll withOutputParallelization(boolean outputParallelization) {
+ return toBuilder().setOutputParallelization(outputParallelization).build();
}
@Override
public PCollection<KV<String, String>> expand(PCollection<String> input) {
checkArgument(connectionConfiguration() != null, "withConnectionConfiguration() is required");
-
- return input
- .apply(ParDo.of(new ReadFn(connectionConfiguration(), batchSize())))
- .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
- .apply(new Reparallelize());
+ PCollection<KV<String, String>> output =
+ input
+ .apply(ParDo.of(new ReadFn(connectionConfiguration(), batchSize())))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
+ if (outputParallelization()) {
+ output = output.apply(new Reparallelize());
+ }
+ return output;
}
}
- abstract static class BaseReadFn<T> extends DoFn<String, T> {
+ private abstract static class BaseReadFn<T> extends DoFn<String, T> {
protected final RedisConnectionConfiguration connectionConfiguration;
transient Jedis jedis;
@@ -307,9 +341,9 @@
}
@ProcessElement
- public void processElement(ProcessContext processContext) throws Exception {
+ public void processElement(ProcessContext c) {
ScanParams scanParams = new ScanParams();
- scanParams.match(processContext.element());
+ scanParams.match(c.element());
String cursor = ScanParams.SCAN_POINTER_START;
boolean finished = false;
@@ -317,7 +351,7 @@
ScanResult<String> scanResult = jedis.scan(cursor, scanParams);
List<String> keys = scanResult.getResult();
for (String k : keys) {
- processContext.output(k);
+ c.output(k);
}
cursor = scanResult.getCursor();
if (cursor.equals(ScanParams.SCAN_POINTER_START)) {
@@ -326,42 +360,52 @@
}
}
}
+
/** A {@link DoFn} requesting Redis server to get key/value pairs. */
private static class ReadFn extends BaseReadFn<KV<String, String>> {
@Nullable transient Multimap<BoundedWindow, String> bundles = null;
@Nullable AtomicInteger batchCount = null;
private final int batchSize;
- @StartBundle
- public void startBundle(StartBundleContext context) {
- bundles = ArrayListMultimap.create();
- batchCount = new AtomicInteger();
- }
-
ReadFn(RedisConnectionConfiguration connectionConfiguration, int batchSize) {
super(connectionConfiguration);
this.batchSize = batchSize;
}
- private int getBatchSize() {
- return batchSize;
+ @StartBundle
+ public void startBundle() {
+ bundles = ArrayListMultimap.create();
+ batchCount = new AtomicInteger();
}
@ProcessElement
- public void processElement(ProcessContext processContext, BoundedWindow window)
- throws Exception {
- String key = processContext.element();
+ public void processElement(ProcessContext c, BoundedWindow window) {
+ String key = c.element();
bundles.put(window, key);
if (batchCount.incrementAndGet() > getBatchSize()) {
Multimap<BoundedWindow, KV<String, String>> kvs = fetchAndFlush();
for (BoundedWindow w : kvs.keySet()) {
for (KV<String, String> kv : kvs.get(w)) {
- processContext.output(kv);
+ c.output(kv);
}
}
}
}
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) {
+ Multimap<BoundedWindow, KV<String, String>> kvs = fetchAndFlush();
+ for (BoundedWindow w : kvs.keySet()) {
+ for (KV<String, String> kv : kvs.get(w)) {
+ context.output(kv, w.maxTimestamp(), w);
+ }
+ }
+ }
+
+ private int getBatchSize() {
+ return batchSize;
+ }
+
private Multimap<BoundedWindow, KV<String, String>> fetchAndFlush() {
Multimap<BoundedWindow, KV<String, String>> kvs = ArrayListMultimap.create();
for (BoundedWindow w : bundles.keySet()) {
@@ -378,16 +422,6 @@
batchCount.set(0);
return kvs;
}
-
- @FinishBundle
- public void finishBundle(FinishBundleContext context) throws Exception {
- Multimap<BoundedWindow, KV<String, String>> kvs = fetchAndFlush();
- for (BoundedWindow w : kvs.keySet()) {
- for (KV<String, String> kv : kvs.get(w)) {
- context.output(kv, w.maxTimestamp(), w);
- }
- }
- }
}
private static class Reparallelize
@@ -395,8 +429,7 @@
@Override
public PCollection<KV<String, String>> expand(PCollection<KV<String, String>> input) {
- // reparallelize mimics the same behavior as in JdbcIO
- // breaking fusion
+ // reparallelize mimics the same behavior as in JdbcIO, used to break fusion
PCollectionView<Iterable<KV<String, String>>> empty =
input
.apply("Consume", Filter.by(SerializableFunctions.constant(false)))
@@ -407,8 +440,8 @@
ParDo.of(
new DoFn<KV<String, String>, KV<String, String>>() {
@ProcessElement
- public void processElement(ProcessContext context) {
- context.output(context.element());
+ public void processElement(ProcessContext c) {
+ c.output(c.element());
}
})
.withSideInputs(empty));
@@ -468,7 +501,7 @@
@Nullable
abstract Long expireTime();
- abstract Builder builder();
+ abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
@@ -486,37 +519,39 @@
public Write withEndpoint(String host, int port) {
checkArgument(host != null, "host can not be null");
checkArgument(port > 0, "port can not be negative or 0");
- return builder()
+ return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withHost(host).withPort(port))
.build();
}
public Write withAuth(String auth) {
checkArgument(auth != null, "auth can not be null");
- return builder().setConnectionConfiguration(connectionConfiguration().withAuth(auth)).build();
+ return toBuilder()
+ .setConnectionConfiguration(connectionConfiguration().withAuth(auth))
+ .build();
}
public Write withTimeout(int timeout) {
checkArgument(timeout >= 0, "timeout can not be negative");
- return builder()
+ return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withTimeout(timeout))
.build();
}
public Write withConnectionConfiguration(RedisConnectionConfiguration connection) {
checkArgument(connection != null, "connection can not be null");
- return builder().setConnectionConfiguration(connection).build();
+ return toBuilder().setConnectionConfiguration(connection).build();
}
public Write withMethod(Method method) {
checkArgument(method != null, "method can not be null");
- return builder().setMethod(method).build();
+ return toBuilder().setMethod(method).build();
}
public Write withExpireTime(Long expireTimeMillis) {
checkArgument(expireTimeMillis != null, "expireTimeMillis can not be null");
checkArgument(expireTimeMillis > 0, "expireTimeMillis can not be negative or 0");
- return builder().setExpireTime(expireTimeMillis).build();
+ return toBuilder().setExpireTime(expireTimeMillis).build();
}
@Override
@@ -555,8 +590,8 @@
}
@ProcessElement
- public void processElement(ProcessContext processContext) {
- KV<String, String> record = processContext.element();
+ public void processElement(ProcessContext c) {
+ KV<String, String> record = c.element();
writeRecord(record);
diff --git a/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java b/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
index bcb3fca..badf039 100644
--- a/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
+++ b/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
@@ -175,7 +175,7 @@
}
@Test
- public void testWriteUsingINCRBY() throws Exception {
+ public void testWriteUsingINCRBY() {
String key = "key_incr";
List<String> values = Arrays.asList("0", "1", "2", "-3", "2", "4", "0", "5");
List<KV<String, String>> data = buildConstantKeyList(key, values);
@@ -190,7 +190,7 @@
}
@Test
- public void testWriteUsingDECRBY() throws Exception {
+ public void testWriteUsingDECRBY() {
String key = "key_decr";
List<String> values = Arrays.asList("-10", "1", "2", "-3", "2", "4", "0", "5");
diff --git a/sdks/java/testing/expansion-service/src/test/java/org/apache/beam/sdk/expansion/TestExpansionService.java b/sdks/java/testing/expansion-service/src/test/java/org/apache/beam/sdk/expansion/TestExpansionService.java
index 68c0e71..d1a832f 100644
--- a/sdks/java/testing/expansion-service/src/test/java/org/apache/beam/sdk/expansion/TestExpansionService.java
+++ b/sdks/java/testing/expansion-service/src/test/java/org/apache/beam/sdk/expansion/TestExpansionService.java
@@ -31,8 +31,8 @@
import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.Server;
-import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ServerBuilder;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.Server;
+import org.apache.beam.vendor.grpc.v1p26p0.io.grpc.ServerBuilder;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
/**
@@ -57,11 +57,7 @@
Schema schema = new Schema.Parser().parse(rawSchema);
return ImmutableMap.of(
TEST_COUNT_URN, spec -> Count.perElement(),
- TEST_FILTER_URN,
- spec ->
- Filter.lessThanEq(
- // TODO(BEAM-6587): Use strings directly rather than longs.
- (long) spec.getPayload().toStringUtf8().charAt(0)),
+ TEST_FILTER_URN, spec -> Filter.lessThanEq(spec.getPayload().toStringUtf8()),
TEST_PARQUET_READ_URN,
spec ->
new PTransform<PBegin, PCollection<GenericRecord>>() {
diff --git a/sdks/java/testing/nexmark/build.gradle b/sdks/java/testing/nexmark/build.gradle
index c294c8e..290a62c 100644
--- a/sdks/java/testing/nexmark/build.gradle
+++ b/sdks/java/testing/nexmark/build.gradle
@@ -67,7 +67,6 @@
compile library.java.avro
compile library.java.joda_time
compile library.java.slf4j_api
- compile library.java.commons_lang3
compile library.java.kafka_clients
provided library.java.junit
provided library.java.hamcrest_core
diff --git a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkConfiguration.java b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkConfiguration.java
index 2b91e07..2ff5744 100644
--- a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkConfiguration.java
+++ b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkConfiguration.java
@@ -205,14 +205,7 @@
debug = options.getDebug();
}
if (options.getQuery() != null) {
- try {
- query = NexmarkQueryName.valueOf(options.getQuery());
- } catch (IllegalArgumentException exc) {
- query = NexmarkQueryName.fromNumber(Integer.parseInt(options.getQuery()));
- }
- if (query == null) {
- throw new IllegalArgumentException("Unknown query: " + query);
- }
+ query = NexmarkQueryName.fromId(options.getQuery());
}
if (options.getSourceType() != null) {
sourceType = options.getSourceType();
diff --git a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java
index 12109ec..794d154 100644
--- a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java
+++ b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkLauncher.java
@@ -27,8 +27,10 @@
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
import javax.annotation.Nullable;
import org.apache.beam.sdk.Pipeline;
@@ -93,10 +95,12 @@
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Splitter;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.LongDeserializer;
@@ -159,10 +163,14 @@
@Nullable private String pubsubSubscription;
@Nullable private PubsubHelper pubsubHelper;
+ private final Map<NexmarkQueryName, NexmarkQuery> queries;
+ private final Map<NexmarkQueryName, NexmarkQueryModel> models;
public NexmarkLauncher(OptionT options, NexmarkConfiguration configuration) {
this.options = options;
this.configuration = configuration;
+ queries = createQueries();
+ models = createQueryModels();
}
/** Is this query running in streaming mode? */
@@ -1193,12 +1201,10 @@
}
private NexmarkQueryModel getNexmarkQueryModel() {
- Map<NexmarkQueryName, NexmarkQueryModel> models = createQueryModels();
return models.get(configuration.query);
}
private NexmarkQuery<?> getNexmarkQuery() {
- Map<NexmarkQueryName, NexmarkQuery> queries = createQueries();
return queries.get(configuration.query);
}
@@ -1228,7 +1234,22 @@
}
private Map<NexmarkQueryName, NexmarkQuery> createQueries() {
- return isSql() ? createSqlQueries() : createJavaQueries();
+ Map<NexmarkQueryName, NexmarkQuery> defaultQueries =
+ isSql() ? createSqlQueries() : createJavaQueries();
+ Set<NexmarkQueryName> skippableQueries = getSkippableQueries();
+ return ImmutableMap.copyOf(
+ Maps.filterKeys(defaultQueries, query -> !skippableQueries.contains(query)));
+ }
+
+ private Set<NexmarkQueryName> getSkippableQueries() {
+ Set<NexmarkQueryName> skipQueries = new LinkedHashSet<>();
+ if (options.getSkipQueries() != null && !options.getSkipQueries().trim().equals("")) {
+ Iterable<String> queries = Splitter.on(',').split(options.getSkipQueries());
+ for (String query : queries) {
+ skipQueries.add(NexmarkQueryName.fromId(query.trim()));
+ }
+ }
+ return skipQueries;
}
private Map<NexmarkQueryName, NexmarkQuery> createSqlQueries() {
diff --git a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkOptions.java b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkOptions.java
index 681316a..1245c6f 100644
--- a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkOptions.java
+++ b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkOptions.java
@@ -77,6 +77,12 @@
void setQuery(String query);
+ @Description("Skip the execution of the given queries (comma separated)")
+ @Nullable
+ String getSkipQueries();
+
+ void setSkipQueries(String queries);
+
@Description("Prefix for output files if using text output for results or running Query 10.")
@Nullable
String getOutputPath();
diff --git a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkQueryName.java b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkQueryName.java
index 7159c86..27a488d 100644
--- a/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkQueryName.java
+++ b/sdks/java/testing/nexmark/src/main/java/org/apache/beam/sdk/nexmark/NexmarkQueryName.java
@@ -81,4 +81,21 @@
}
return null;
}
+
+ /**
+ * @return The given {@link NexmarkQueryName} for the id. The id can be the query number (for
+ * backwards compatibility) or its name.
+ */
+ public static NexmarkQueryName fromId(String id) {
+ NexmarkQueryName query;
+ try {
+ query = NexmarkQueryName.valueOf(id);
+ } catch (IllegalArgumentException exc) {
+ query = NexmarkQueryName.fromNumber(Integer.parseInt(id));
+ }
+ if (query == null) {
+ throw new IllegalArgumentException("Unknown query: " + id);
+ }
+ return query;
+ }
}
diff --git a/sdks/java/testing/test-utils/src/main/java/org/apache/beam/sdk/testutils/metrics/MetricsReader.java b/sdks/java/testing/test-utils/src/main/java/org/apache/beam/sdk/testutils/metrics/MetricsReader.java
index 986d568..7fee6a2 100644
--- a/sdks/java/testing/test-utils/src/main/java/org/apache/beam/sdk/testutils/metrics/MetricsReader.java
+++ b/sdks/java/testing/test-utils/src/main/java/org/apache/beam/sdk/testutils/metrics/MetricsReader.java
@@ -17,6 +17,9 @@
*/
package org.apache.beam.sdk.testutils.metrics;
+import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+
+import com.sun.istack.internal.NotNull;
import java.util.Collection;
import java.util.NoSuchElementException;
import java.util.Optional;
@@ -164,4 +167,19 @@
private boolean isCredible(long value) {
return (Math.abs(value - now) <= Duration.standardDays(10000).getMillis());
}
+
+ /**
+ * Factory method. Returns instance of {@link MetricsReader}.
+ *
+ * @param results pipeline from which metrics are read from
+ * @param namespace namespace is required due to fact methods uses it internally
+ * @return instance of {@link MetricsReader}
+ */
+ @NotNull
+ public static MetricsReader ofResults(
+ @NotNull final PipelineResult results, @NotNull final String namespace) {
+ checkNotNull(results);
+ checkNotNull(namespace);
+ return new MetricsReader(results, namespace);
+ }
}
diff --git a/sdks/python/.pylintrc b/sdks/python/.pylintrc
index 26b39df..fb226b7 100644
--- a/sdks/python/.pylintrc
+++ b/sdks/python/.pylintrc
@@ -179,7 +179,7 @@
ignore-long-lines=(?x)
(^\s*(import|from)\s
|^\s*(\#\ )?<?(https?|ftp):\/\/[^\s\/$.?#].[^\s]*>?$
- |# type:
+ |^.*\#\ type\:
)
[VARIABLES]
diff --git a/sdks/python/apache_beam/coders/avro_record.py b/sdks/python/apache_beam/coders/avro_record.py
index a5b8b60..c551c90 100644
--- a/sdks/python/apache_beam/coders/avro_record.py
+++ b/sdks/python/apache_beam/coders/avro_record.py
@@ -17,6 +17,8 @@
"""AvroRecord for AvroGenericCoder."""
+# pytype: skip-file
+
from __future__ import absolute_import
__all__ = ['AvroRecord']
diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py
index 379156a..8f78b95 100644
--- a/sdks/python/apache_beam/coders/coder_impl.py
+++ b/sdks/python/apache_beam/coders/coder_impl.py
@@ -31,6 +31,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -907,7 +909,8 @@
buffer = create_OutputStream()
if (self._write_state is not None
and out.size() - start_size > self._write_state_threshold):
- tail = (value_iter[index + 1:] if isinstance(value, (list, tuple))
+ tail = (value_iter[index + 1:]
+ if isinstance(value_iter, (list, tuple))
else value_iter)
state_token = self._write_state(tail, self._elem_coder)
out.write_var_int64(-1)
diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py
index c6f9f1d..bff6dd0 100644
--- a/sdks/python/apache_beam/coders/coders.py
+++ b/sdks/python/apache_beam/coders/coders.py
@@ -19,6 +19,8 @@
Only those coders listed in __all__ are part of the public API of this module.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
diff --git a/sdks/python/apache_beam/coders/coders_test.py b/sdks/python/apache_beam/coders/coders_test.py
index 9b39962..a151a4a 100644
--- a/sdks/python/apache_beam/coders/coders_test.py
+++ b/sdks/python/apache_beam/coders/coders_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py
index cdab818..f57407b 100644
--- a/sdks/python/apache_beam/coders/coders_test_common.py
+++ b/sdks/python/apache_beam/coders/coders_test_common.py
@@ -16,6 +16,8 @@
#
"""Tests common to all coder implementations."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/coders/fast_coders_test.py b/sdks/python/apache_beam/coders/fast_coders_test.py
index 6247a60..c2f795c 100644
--- a/sdks/python/apache_beam/coders/fast_coders_test.py
+++ b/sdks/python/apache_beam/coders/fast_coders_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for compiled implementation of coder impls."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/coders/observable.py b/sdks/python/apache_beam/coders/observable.py
index 3d0a7fc..b744f5e 100644
--- a/sdks/python/apache_beam/coders/observable.py
+++ b/sdks/python/apache_beam/coders/observable.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/coders/observable_test.py b/sdks/python/apache_beam/coders/observable_test.py
index fc3c410..7bf5ab8 100644
--- a/sdks/python/apache_beam/coders/observable_test.py
+++ b/sdks/python/apache_beam/coders/observable_test.py
@@ -16,6 +16,8 @@
#
"""Tests for the Observable mixin class."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/coders/row_coder.py b/sdks/python/apache_beam/coders/row_coder.py
index a259f36..bc5fd69 100644
--- a/sdks/python/apache_beam/coders/row_coder.py
+++ b/sdks/python/apache_beam/coders/row_coder.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import itertools
@@ -68,7 +70,7 @@
def as_cloud_object(self, coders_context=None):
raise NotImplementedError("as_cloud_object not supported for RowCoder")
- __hash__ = None
+ __hash__ = None # type: ignore[assignment]
def __eq__(self, other):
return type(self) == type(other) and self.schema == other.schema
diff --git a/sdks/python/apache_beam/coders/row_coder_test.py b/sdks/python/apache_beam/coders/row_coder_test.py
index dbdc5fc..76088f6 100644
--- a/sdks/python/apache_beam/coders/row_coder_test.py
+++ b/sdks/python/apache_beam/coders/row_coder_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/coders/slow_coders_test.py b/sdks/python/apache_beam/coders/slow_coders_test.py
index 2ddc46e..9d3da5d 100644
--- a/sdks/python/apache_beam/coders/slow_coders_test.py
+++ b/sdks/python/apache_beam/coders/slow_coders_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for uncompiled implementation of coder impls."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/coders/slow_stream.py b/sdks/python/apache_beam/coders/slow_stream.py
index 08a6c83..d8501d4 100644
--- a/sdks/python/apache_beam/coders/slow_stream.py
+++ b/sdks/python/apache_beam/coders/slow_stream.py
@@ -19,6 +19,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import struct
diff --git a/sdks/python/apache_beam/coders/standard_coders_test.py b/sdks/python/apache_beam/coders/standard_coders_test.py
index 810481b..bd8ea16 100644
--- a/sdks/python/apache_beam/coders/standard_coders_test.py
+++ b/sdks/python/apache_beam/coders/standard_coders_test.py
@@ -17,6 +17,8 @@
"""Unit tests for coders that must be consistent across all Beam SDKs.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/coders/stream_test.py b/sdks/python/apache_beam/coders/stream_test.py
index e627ebb..d82a816 100644
--- a/sdks/python/apache_beam/coders/stream_test.py
+++ b/sdks/python/apache_beam/coders/stream_test.py
@@ -16,6 +16,8 @@
#
"""Tests for the stream implementations."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/coders/typecoders.py b/sdks/python/apache_beam/coders/typecoders.py
index b957f67..cd694af 100644
--- a/sdks/python/apache_beam/coders/typecoders.py
+++ b/sdks/python/apache_beam/coders/typecoders.py
@@ -63,6 +63,9 @@
See apache_beam.typehints.decorators module for more details.
"""
+
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/coders/typecoders_test.py b/sdks/python/apache_beam/coders/typecoders_test.py
index 52e32fb..3bc8aef 100644
--- a/sdks/python/apache_beam/coders/typecoders_test.py
+++ b/sdks/python/apache_beam/coders/typecoders_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the typecoders module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/error.py b/sdks/python/apache_beam/error.py
index 47fec08..3165842 100644
--- a/sdks/python/apache_beam/error.py
+++ b/sdks/python/apache_beam/error.py
@@ -17,6 +17,8 @@
"""Python Dataflow error classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
diff --git a/sdks/python/apache_beam/examples/avro_bitcoin.py b/sdks/python/apache_beam/examples/avro_bitcoin.py
index f6ab89e..917c8c4 100644
--- a/sdks/python/apache_beam/examples/avro_bitcoin.py
+++ b/sdks/python/apache_beam/examples/avro_bitcoin.py
@@ -24,6 +24,8 @@
--compress --fastavro --output fastavro-compressed
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/complete/autocomplete.py b/sdks/python/apache_beam/examples/complete/autocomplete.py
index 03b8500..6be41ad 100644
--- a/sdks/python/apache_beam/examples/complete/autocomplete.py
+++ b/sdks/python/apache_beam/examples/complete/autocomplete.py
@@ -17,6 +17,8 @@
"""A workflow emitting the top k most common words for each prefix."""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/complete/autocomplete_test.py b/sdks/python/apache_beam/examples/complete/autocomplete_test.py
index 29d8203..863ed51 100644
--- a/sdks/python/apache_beam/examples/complete/autocomplete_test.py
+++ b/sdks/python/apache_beam/examples/complete/autocomplete_test.py
@@ -17,6 +17,8 @@
"""Test for the autocomplete example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/examples/complete/distribopt.py b/sdks/python/apache_beam/examples/complete/distribopt.py
index 42c7b83..20c0a8c 100644
--- a/sdks/python/apache_beam/examples/complete/distribopt.py
+++ b/sdks/python/apache_beam/examples/complete/distribopt.py
@@ -49,6 +49,8 @@
- Selecting the mapping with the lowest cost.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/examples/complete/distribopt_test.py b/sdks/python/apache_beam/examples/complete/distribopt_test.py
index ffdbd99..bce76a1 100644
--- a/sdks/python/apache_beam/examples/complete/distribopt_test.py
+++ b/sdks/python/apache_beam/examples/complete/distribopt_test.py
@@ -17,6 +17,8 @@
"""Test for the distrib_optimization example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi.py b/sdks/python/apache_beam/examples/complete/estimate_pi.py
index aa41d02..9870075 100644
--- a/sdks/python/apache_beam/examples/complete/estimate_pi.py
+++ b/sdks/python/apache_beam/examples/complete/estimate_pi.py
@@ -24,6 +24,8 @@
we multiply our counts ratio by four to estimate π.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi_test.py b/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
index 78376b6..2ea4da2 100644
--- a/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
+++ b/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
@@ -17,6 +17,8 @@
"""Test for the estimate_pi example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 8f446e6..f0ab8fc 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats.py
@@ -70,6 +70,8 @@
--temp_location gs://$BUCKET/user_score/temp
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -184,11 +186,6 @@
return ', '.join(
'%s:%s' % (col, self.schema[col]) for col in self.schema)
- def get_schema(self):
- """Build the output table schema."""
- return ', '.join(
- '%s:%s' % (col, self.schema[col]) for col in self.schema)
-
def expand(self, pcoll):
return (
pcoll
diff --git a/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py b/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
index 70dafb0..96e57b1 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats_it_test.py
@@ -30,6 +30,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/game/game_stats_test.py b/sdks/python/apache_beam/examples/complete/game/game_stats_test.py
index 209f0cf..e3058d8 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats_test.py
@@ -17,6 +17,8 @@
"""Test for the game_stats example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 e0a5c47..caa7a90 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
@@ -64,6 +64,8 @@
--temp_location gs://$BUCKET/user_score/temp
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py
index 8d86f18..4499d66 100644
--- a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_it_test.py
@@ -30,6 +30,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py
index 8c2497a..79db0bc 100644
--- a/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/hourly_team_score_test.py
@@ -17,6 +17,8 @@
"""Test for the user_score example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 2288d16..c297532 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board.py
@@ -78,6 +78,8 @@
--temp_location gs://$BUCKET/user_score/temp
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py b/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
index af2f2e6..d718f70 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board_it_test.py
@@ -30,6 +30,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/game/leader_board_test.py b/sdks/python/apache_beam/examples/complete/game/leader_board_test.py
index 3aad052..08ac76b 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board_test.py
@@ -17,6 +17,8 @@
"""Test for the leader_board example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 74b47ba..962104b 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score.py
@@ -53,6 +53,8 @@
--temp_location gs://$BUCKET/user_score/temp
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py b/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
index 5e5ba97..00b6bfe 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score_it_test.py
@@ -30,6 +30,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score_test.py b/sdks/python/apache_beam/examples/complete/game/user_score_test.py
index f41006f..3dbe6df 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score_test.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score_test.py
@@ -17,6 +17,8 @@
"""Test for the user_score example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py
index f861e48..9b55708 100644
--- a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py
+++ b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py
@@ -20,6 +20,8 @@
We use the quadratic polinomial f(z) = z*z + c, with c = -.62772 +.42193i
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py
index a7bc0f9..1b0b41a 100644
--- a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py
+++ b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py
@@ -17,6 +17,8 @@
"""Test for the juliaset example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/juliaset/juliaset_main.py b/sdks/python/apache_beam/examples/complete/juliaset/juliaset_main.py
index 9d72772..32361cf 100644
--- a/sdks/python/apache_beam/examples/complete/juliaset/juliaset_main.py
+++ b/sdks/python/apache_beam/examples/complete/juliaset/juliaset_main.py
@@ -47,6 +47,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/complete/juliaset/setup.py b/sdks/python/apache_beam/examples/complete/juliaset/setup.py
index ee02a16..ab4306a 100644
--- a/sdks/python/apache_beam/examples/complete/juliaset/setup.py
+++ b/sdks/python/apache_beam/examples/complete/juliaset/setup.py
@@ -25,6 +25,8 @@
when running the workflow for remote execution.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/complete/tfidf.py b/sdks/python/apache_beam/examples/complete/tfidf.py
index 77ee4c1..6069416 100644
--- a/sdks/python/apache_beam/examples/complete/tfidf.py
+++ b/sdks/python/apache_beam/examples/complete/tfidf.py
@@ -21,6 +21,8 @@
http://en.wikipedia.org/wiki/Tf-idf
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/examples/complete/tfidf_test.py b/sdks/python/apache_beam/examples/complete/tfidf_test.py
index 4b19269..5a9b456 100644
--- a/sdks/python/apache_beam/examples/complete/tfidf_test.py
+++ b/sdks/python/apache_beam/examples/complete/tfidf_test.py
@@ -17,6 +17,8 @@
"""Test for the TF-IDF example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 6b04a00..4e65525 100644
--- a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
+++ b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
@@ -39,6 +39,8 @@
be overridden with --input.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py
index fcef981..8befc86 100644
--- a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py
+++ b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions_test.py
@@ -17,6 +17,8 @@
"""Test for the top wikipedia sessions example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_schema.py b/sdks/python/apache_beam/examples/cookbook/bigquery_schema.py
index c7b1ccd..7d3ab49 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_schema.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_schema.py
@@ -22,6 +22,8 @@
nested and repeated fields.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py
index fb7ee42..8abb8f4 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py
@@ -27,6 +27,8 @@
a word that should be ignored when forming groups.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
index 031eeb3..97bcbd6 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
@@ -17,6 +17,8 @@
"""Test for the BigQuery side input example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes.py b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes.py
index c7c837b..5ba995f 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes.py
@@ -32,6 +32,8 @@
represents table rows as plain Python dictionaries.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_it_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_it_test.py
index f7eb93b..4beb214 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_it_test.py
@@ -17,6 +17,8 @@
"""End-to-end test for Bigquery tornadoes example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
index cad34d8..581ef56 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
@@ -17,6 +17,8 @@
"""Test for the BigQuery tornadoes example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 3c1204c..5c06295 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
@@ -16,6 +16,8 @@
#
"""Unittest for GCP Bigtable testing."""
+# pytype: skip-file
+
from __future__ import absolute_import
import datetime
diff --git a/sdks/python/apache_beam/examples/cookbook/coders.py b/sdks/python/apache_beam/examples/cookbook/coders.py
index b0f2a2b..68845cf 100644
--- a/sdks/python/apache_beam/examples/cookbook/coders.py
+++ b/sdks/python/apache_beam/examples/cookbook/coders.py
@@ -28,6 +28,8 @@
[TEAM_NAME, POINTS]
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/coders_test.py b/sdks/python/apache_beam/examples/cookbook/coders_test.py
index 7a3b7f8..b706123 100644
--- a/sdks/python/apache_beam/examples/cookbook/coders_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/coders_test.py
@@ -17,6 +17,8 @@
"""Test for the coders example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/combiners_test.py b/sdks/python/apache_beam/examples/cookbook/combiners_test.py
index 15714c0..587978f 100644
--- a/sdks/python/apache_beam/examples/cookbook/combiners_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/combiners_test.py
@@ -23,6 +23,8 @@
checked directly on the last PCollection produced.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py b/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py
index db86003..4d0a64c 100644
--- a/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py
+++ b/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py
@@ -20,6 +20,8 @@
These example show the different ways you can write custom PTransforms.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py b/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py
index 7620dae..c56a31a 100644
--- a/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/custom_ptransform_test.py
@@ -17,6 +17,8 @@
"""Tests for the various custom Count implementation examples."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py b/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py
index 3ee39b7..8b110b5 100644
--- a/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py
+++ b/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py
@@ -61,6 +61,8 @@
https://github.com/googleapis/googleapis/tree/master/google/datastore/v1
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/datastore_wordcount_it_test.py b/sdks/python/apache_beam/examples/cookbook/datastore_wordcount_it_test.py
index 5a4fea9..08305bc 100644
--- a/sdks/python/apache_beam/examples/cookbook/datastore_wordcount_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/datastore_wordcount_it_test.py
@@ -17,6 +17,8 @@
"""End-to-end test for Datastore Wordcount example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/filters.py b/sdks/python/apache_beam/examples/cookbook/filters.py
index d1b0201..06ec8a2 100644
--- a/sdks/python/apache_beam/examples/cookbook/filters.py
+++ b/sdks/python/apache_beam/examples/cookbook/filters.py
@@ -24,6 +24,8 @@
as well as global aggregates computed during pipeline execution.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/filters_test.py b/sdks/python/apache_beam/examples/cookbook/filters_test.py
index 5187a2f..eebbc10 100644
--- a/sdks/python/apache_beam/examples/cookbook/filters_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/filters_test.py
@@ -17,6 +17,8 @@
"""Test for the filters example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -39,33 +41,31 @@
{'year': 2011, 'month': 3, 'day': 3, 'mean_temp': 5, 'removed': 'a'},
]
- def _get_result_for_month(self, month):
- p = TestPipeline()
- rows = (p | 'create' >> beam.Create(self.input_data))
-
+ def _get_result_for_month(self, pipeline, month):
+ rows = (pipeline | 'create' >> beam.Create(self.input_data))
results = filters.filter_cold_days(rows, month)
return results
def test_basic(self):
"""Test that the correct result is returned for a simple dataset."""
- results = self._get_result_for_month(1)
- assert_that(
- results,
- equal_to([{'year': 2010, 'month': 1, 'day': 1, 'mean_temp': 3},
- {'year': 2012, 'month': 1, 'day': 2, 'mean_temp': 3}]))
- results.pipeline.run()
+ with TestPipeline() as p:
+ results = self._get_result_for_month(p, 1)
+ assert_that(
+ results,
+ equal_to([{'year': 2010, 'month': 1, 'day': 1, 'mean_temp': 3},
+ {'year': 2012, 'month': 1, 'day': 2, 'mean_temp': 3}]))
def test_basic_empty(self):
"""Test that the correct empty result is returned for a simple dataset."""
- results = self._get_result_for_month(3)
- assert_that(results, equal_to([]))
- results.pipeline.run()
+ with TestPipeline() as p:
+ results = self._get_result_for_month(p, 3)
+ assert_that(results, equal_to([]))
def test_basic_empty_missing(self):
"""Test that the correct empty result is returned for a missing month."""
- results = self._get_result_for_month(4)
- assert_that(results, equal_to([]))
- results.pipeline.run()
+ with TestPipeline() as p:
+ results = self._get_result_for_month(p, 4)
+ assert_that(results, equal_to([]))
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/examples/cookbook/group_with_coder.py b/sdks/python/apache_beam/examples/cookbook/group_with_coder.py
index 2202b8c..171eada 100644
--- a/sdks/python/apache_beam/examples/cookbook/group_with_coder.py
+++ b/sdks/python/apache_beam/examples/cookbook/group_with_coder.py
@@ -25,6 +25,8 @@
and score.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/group_with_coder_test.py b/sdks/python/apache_beam/examples/cookbook/group_with_coder_test.py
index 59b70f9..f73726d 100644
--- a/sdks/python/apache_beam/examples/cookbook/group_with_coder_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/group_with_coder_test.py
@@ -17,6 +17,8 @@
"""Test for the custom coders example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/mergecontacts.py b/sdks/python/apache_beam/examples/cookbook/mergecontacts.py
index 9bd6a96..7d427c4 100644
--- a/sdks/python/apache_beam/examples/cookbook/mergecontacts.py
+++ b/sdks/python/apache_beam/examples/cookbook/mergecontacts.py
@@ -29,6 +29,8 @@
Non-linear pipelines (i.e., pipelines with branches)
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/mergecontacts_test.py b/sdks/python/apache_beam/examples/cookbook/mergecontacts_test.py
index 23f22bc..8d9e85d 100644
--- a/sdks/python/apache_beam/examples/cookbook/mergecontacts_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/mergecontacts_test.py
@@ -17,6 +17,8 @@
"""Test for the mergecontacts example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py b/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py
index 7896027..e8cdc7f 100644
--- a/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py
+++ b/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py
@@ -49,6 +49,8 @@
--output gs://YOUR_OUTPUT_PREFIX
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo_test.py b/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo_test.py
index afe350d..01a5768 100644
--- a/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo_test.py
@@ -17,6 +17,8 @@
"""Test for the multiple_output_pardo example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/fastavro_it_test.py b/sdks/python/apache_beam/examples/fastavro_it_test.py
index 48dfe11..35afa52 100644
--- a/sdks/python/apache_beam/examples/fastavro_it_test.py
+++ b/sdks/python/apache_beam/examples/fastavro_it_test.py
@@ -42,6 +42,8 @@
"
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -161,43 +163,42 @@
result.wait_until_finish()
assert result.state == PipelineState.DONE
- fastavro_read_pipeline = TestPipeline(is_integration_test=True)
+ with TestPipeline(is_integration_test=True) as fastavro_read_pipeline:
- fastavro_records = \
- fastavro_read_pipeline \
- | 'create-fastavro' >> Create(['%s*' % fastavro_output]) \
- | 'read-fastavro' >> ReadAllFromAvro(use_fastavro=True) \
- | Map(lambda rec: (rec['number'], rec))
+ fastavro_records = \
+ fastavro_read_pipeline \
+ | 'create-fastavro' >> Create(['%s*' % fastavro_output]) \
+ | 'read-fastavro' >> ReadAllFromAvro(use_fastavro=True) \
+ | Map(lambda rec: (rec['number'], rec))
- avro_records = \
- fastavro_read_pipeline \
- | 'create-avro' >> Create(['%s*' % avro_output]) \
- | 'read-avro' >> ReadAllFromAvro(use_fastavro=False) \
- | Map(lambda rec: (rec['number'], rec))
+ avro_records = \
+ fastavro_read_pipeline \
+ | 'create-avro' >> Create(['%s*' % avro_output]) \
+ | 'read-avro' >> ReadAllFromAvro(use_fastavro=False) \
+ | Map(lambda rec: (rec['number'], rec))
- def check(elem):
- v = elem[1]
+ def check(elem):
+ v = elem[1]
- def assertEqual(l, r):
- if l != r:
- raise BeamAssertException('Assertion failed: %s == %s' % (l, r))
+ def assertEqual(l, r):
+ if l != r:
+ raise BeamAssertException('Assertion failed: %s == %s' % (l, r))
- assertEqual(v.keys(), ['avro', 'fastavro'])
- avro_values = v['avro']
- fastavro_values = v['fastavro']
- assertEqual(avro_values, fastavro_values)
- assertEqual(len(avro_values), 1)
+ assertEqual(v.keys(), ['avro', 'fastavro'])
+ avro_values = v['avro']
+ fastavro_values = v['fastavro']
+ assertEqual(avro_values, fastavro_values)
+ assertEqual(len(avro_values), 1)
- # pylint: disable=expression-not-assigned
- {
- 'avro': avro_records,
- 'fastavro': fastavro_records
- } \
- | CoGroupByKey() \
- | Map(check)
+ # pylint: disable=expression-not-assigned
+ {
+ 'avro': avro_records,
+ 'fastavro': fastavro_records
+ } \
+ | CoGroupByKey() \
+ | Map(check)
- self.addCleanup(delete_files, [self.output])
- fastavro_read_pipeline.run().wait_until_finish()
+ self.addCleanup(delete_files, [self.output])
assert result.state == PipelineState.DONE
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 0cfaf5d..47c3836 100644
--- a/sdks/python/apache_beam/examples/flink/flink_streaming_impulse.py
+++ b/sdks/python/apache_beam/examples/flink/flink_streaming_impulse.py
@@ -20,6 +20,8 @@
This can only be used with the Flink portable runner.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
@@ -73,24 +75,22 @@
pipeline_options = PipelineOptions(pipeline_args)
- p = beam.Pipeline(options=pipeline_options)
+ with beam.Pipeline(options=pipeline_options) as p:
- messages = (p | FlinkStreamingImpulseSource()
- .set_message_count(known_args.count)
- .set_interval_ms(known_args.interval_ms))
+ messages = (p | FlinkStreamingImpulseSource()
+ .set_message_count(known_args.count)
+ .set_interval_ms(known_args.interval_ms))
- _ = (messages | 'decode' >> beam.Map(lambda x: ('', 1))
- | 'window' >> beam.WindowInto(window.GlobalWindows(),
- trigger=Repeatedly(
- AfterProcessingTime(5 * 1000)),
- accumulation_mode=
- AccumulationMode.DISCARDING)
- | 'group' >> beam.GroupByKey()
- | 'count' >> beam.Map(count)
- | 'log' >> beam.Map(lambda x: logging.info("%d" % x[1])))
+ _ = (messages | 'decode' >> beam.Map(lambda x: ('', 1))
+ | 'window' >> beam.WindowInto(window.GlobalWindows(),
+ trigger=Repeatedly(
+ AfterProcessingTime(5 * 1000)),
+ accumulation_mode=
+ AccumulationMode.DISCARDING)
+ | 'group' >> beam.GroupByKey()
+ | 'count' >> beam.Map(count)
+ | 'log' >> beam.Map(lambda x: logging.info("%d" % x[1])))
- result = p.run()
- result.wait_until_finish()
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/examples/snippets/snippets.py b/sdks/python/apache_beam/examples/snippets/snippets.py
index 02e24d9..aaa99e9 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets.py
@@ -29,6 +29,8 @@
prefix the PATH_TO_HTML where they are included followed by a descriptive
string. The tags can contain only letters, digits and _.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -109,31 +111,29 @@
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
- p = beam.Pipeline(options=PipelineOptions())
- # [END pipelines_constructing_creating]
+ with beam.Pipeline(options=PipelineOptions()) as p:
+ pass # build your pipeline here
+ # [END pipelines_constructing_creating]
- p = TestPipeline() # Use TestPipeline for testing.
+ with TestPipeline() as p: # Use TestPipeline for testing.
+ # pylint: disable=line-too-long
- # [START pipelines_constructing_reading]
- lines = p | 'ReadMyFile' >> beam.io.ReadFromText('gs://some/inputData.txt')
- # [END pipelines_constructing_reading]
+ # [START pipelines_constructing_reading]
+ lines = p | 'ReadMyFile' >> beam.io.ReadFromText('gs://some/inputData.txt')
+ # [END pipelines_constructing_reading]
- # [START pipelines_constructing_applying]
- words = lines | beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x))
- reversed_words = words | ReverseWords()
- # [END pipelines_constructing_applying]
+ # [START pipelines_constructing_applying]
+ words = lines | beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x))
+ reversed_words = words | ReverseWords()
+ # [END pipelines_constructing_applying]
- # [START pipelines_constructing_writing]
- filtered_words = reversed_words | 'FilterWords' >> beam.Filter(filter_words)
- filtered_words | 'WriteMyFile' >> beam.io.WriteToText(
- 'gs://some/outputData.txt')
- # [END pipelines_constructing_writing]
+ # [START pipelines_constructing_writing]
+ filtered_words = reversed_words | 'FilterWords' >> beam.Filter(filter_words)
+ filtered_words | 'WriteMyFile' >> beam.io.WriteToText(
+ 'gs://some/outputData.txt')
+ # [END pipelines_constructing_writing]
- p.visit(SnippetUtils.RenameFiles(renames))
-
- # [START pipelines_constructing_running]
- p.run()
- # [END pipelines_constructing_running]
+ p.visit(SnippetUtils.RenameFiles(renames))
def model_pipelines(argv):
@@ -247,12 +247,10 @@
my_input = my_options.input
my_output = my_options.output
- p = TestPipeline() # Use TestPipeline for testing.
+ with TestPipeline() as p: # Use TestPipeline for testing.
- lines = p | beam.io.ReadFromText(my_input)
- lines | beam.io.WriteToText(my_output)
-
- p.run()
+ lines = p | beam.io.ReadFromText(my_input)
+ lines | beam.io.WriteToText(my_output)
def pipeline_options_local(argv):
@@ -284,13 +282,12 @@
# [START pipeline_options_local]
# Create and set your Pipeline Options.
options = PipelineOptions()
- p = Pipeline(options=options)
- # [END pipeline_options_local]
+ with Pipeline(options=options) as p:
+ # [END pipeline_options_local]
- p = TestPipeline() # Use TestPipeline for testing.
- lines = p | beam.io.ReadFromText(my_input)
- lines | beam.io.WriteToText(my_output)
- p.run()
+ with TestPipeline() as p: # Use TestPipeline for testing.
+ lines = p | beam.io.ReadFromText(my_input)
+ lines | beam.io.WriteToText(my_output)
def pipeline_options_command_line(argv):
@@ -539,30 +536,28 @@
required=True,
help='Output file to write results to.')
pipeline_options = PipelineOptions(['--output', 'some/output_path'])
- p = beam.Pipeline(options=pipeline_options)
+ with beam.Pipeline(options=pipeline_options) as p:
- wordcount_options = pipeline_options.view_as(WordcountTemplatedOptions)
- lines = p | 'Read' >> ReadFromText(wordcount_options.input)
- # [END example_wordcount_templated]
+ wordcount_options = pipeline_options.view_as(WordcountTemplatedOptions)
+ lines = p | 'Read' >> ReadFromText(wordcount_options.input)
+ # [END example_wordcount_templated]
- def format_result(word_count):
- (word, count) = word_count
- return '%s: %s' % (word, count)
+ def format_result(word_count):
+ (word, count) = word_count
+ return '%s: %s' % (word, count)
- (
- lines
- | 'ExtractWords' >> beam.FlatMap(
- lambda x: re.findall(r'[A-Za-z\']+', x))
- | 'PairWithOnes' >> beam.Map(lambda x: (x, 1))
- | 'Group' >> beam.GroupByKey()
- | 'Sum' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1])))
- | 'Format' >> beam.Map(format_result)
- | 'Write' >> WriteToText(wordcount_options.output)
- )
+ (
+ lines
+ | 'ExtractWords' >> beam.FlatMap(
+ lambda x: re.findall(r'[A-Za-z\']+', x))
+ | 'PairWithOnes' >> beam.Map(lambda x: (x, 1))
+ | 'Group' >> beam.GroupByKey()
+ | 'Sum' >> beam.Map(lambda word_ones: (word_ones[0], sum(word_ones[1])))
+ | 'Format' >> beam.Map(format_result)
+ | 'Write' >> WriteToText(wordcount_options.output)
+ )
- p.visit(SnippetUtils.RenameFiles(renames))
- result = p.run()
- result.wait_until_finish()
+ p.visit(SnippetUtils.RenameFiles(renames))
def examples_wordcount_debugging(renames):
@@ -711,25 +706,23 @@
yield self.templated_int.get() + an_int
pipeline_options = PipelineOptions()
- p = beam.Pipeline(options=pipeline_options)
+ with beam.Pipeline(options=pipeline_options) as p:
- user_options = pipeline_options.view_as(TemplatedUserOptions)
- my_sum_fn = MySumFn(user_options.templated_int)
- sum = (p
- | 'ReadCollection' >> beam.io.ReadFromText(
- 'gs://some/integer_collection')
- | 'StringToInt' >> beam.Map(lambda w: int(w))
- | 'AddGivenInt' >> beam.ParDo(my_sum_fn)
- | 'WriteResultingCollection' >> WriteToText('some/output_path'))
- # [END examples_ptransforms_templated]
+ user_options = pipeline_options.view_as(TemplatedUserOptions)
+ my_sum_fn = MySumFn(user_options.templated_int)
+ sum = (p
+ | 'ReadCollection' >> beam.io.ReadFromText(
+ 'gs://some/integer_collection')
+ | 'StringToInt' >> beam.Map(lambda w: int(w))
+ | 'AddGivenInt' >> beam.ParDo(my_sum_fn)
+ | 'WriteResultingCollection' >> WriteToText('some/output_path'))
+ # [END examples_ptransforms_templated]
- # Templates are not supported by DirectRunner (only by DataflowRunner)
- # so a value must be provided at graph-construction time
- my_sum_fn.templated_int = StaticValueProvider(int, 10)
+ # Templates are not supported by DirectRunner (only by DataflowRunner)
+ # so a value must be provided at graph-construction time
+ my_sum_fn.templated_int = StaticValueProvider(int, 10)
- p.visit(SnippetUtils.RenameFiles(renames))
- result = p.run()
- result.wait_until_finish()
+ p.visit(SnippetUtils.RenameFiles(renames))
# Defining a new source.
@@ -833,16 +826,15 @@
['line ' + str(number) for number in range(0, count)]))
# [START model_custom_source_use_ptransform]
- p = beam.Pipeline(options=PipelineOptions())
- numbers = p | 'ProduceNumbers' >> ReadFromCountingSource(count)
- # [END model_custom_source_use_ptransform]
+ with beam.Pipeline(options=PipelineOptions()) as p:
+ numbers = p | 'ProduceNumbers' >> ReadFromCountingSource(count)
+ # [END model_custom_source_use_ptransform]
- lines = numbers | beam.core.Map(lambda number: 'line %d' % number)
- assert_that(
- lines, equal_to(
- ['line ' + str(number) for number in range(0, count)]))
+ lines = numbers | beam.core.Map(lambda number: 'line %d' % number)
+ assert_that(
+ lines, equal_to(
+ ['line ' + str(number) for number in range(0, count)]))
- p.run().wait_until_finish()
# Defining the new sink.
@@ -1400,20 +1392,19 @@
pipeline_options = PipelineOptions()
# Create pipeline.
- p = beam.Pipeline(options=pipeline_options)
+ with beam.Pipeline(options=pipeline_options) as p:
- my_options = pipeline_options.view_as(MyOptions)
- # Add a branch for logging the ValueProvider value.
- _ = (p
- | beam.Create([None])
- | 'LogValueProvs' >> beam.ParDo(
- LogValueProvidersFn(my_options.string_value)))
+ my_options = pipeline_options.view_as(MyOptions)
+ # Add a branch for logging the ValueProvider value.
+ _ = (p
+ | beam.Create([None])
+ | 'LogValueProvs' >> beam.ParDo(
+ LogValueProvidersFn(my_options.string_value)))
- # The main pipeline.
- result_pc = (p
- | "main_pc" >> beam.Create([1, 2, 3])
- | beam.combiners.Sum.Globally())
+ # The main pipeline.
+ result_pc = (p
+ | "main_pc" >> beam.Create([1, 2, 3])
+ | beam.combiners.Sum.Globally())
- p.run().wait_until_finish()
# [END AccessingValueProviderInfoAfterRunSnip1]
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py b/sdks/python/apache_beam/examples/snippets/snippets_test.py
index 38bcb88..1bcaaca 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_test.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_test.py
@@ -17,6 +17,8 @@
#
"""Tests for all code snippets used in public docs."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey.py
index c507e03..c1ddf8d 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py
index ff86628..fc1b0a7 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally.py
index f9c097a..a198f19 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py
index 5990d2e..e849377 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey.py
index 2fba8e4..fc864f6 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py
index e5fc2ac..ec23167 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues.py
index 92fc6c6..7c52113 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py
index 97d0b3d..3e233c9 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count.py
index 22b1140..ab1600a 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py
index 123d5fb..fff2406 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct.py
index 930fdbe..8062e8a 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py
index ea1a7b2..70d1341 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey.py
index 83e4f87..72c5437 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py
index 4d8283a..3ea278c 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest.py
index 1d3003a..fdc228c 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py
index 5a4c8af..6f9bffd 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max.py
index 7de4050..37f3481 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py
index c229adec..af43781 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean.py
index 36fa5b5..bda260d 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py
index 38b27a0..796bbbf 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min.py
new file mode 100644
index 0000000..d004cf2
--- /dev/null
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min.py
@@ -0,0 +1,60 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+
+def min_globally(test=None):
+ # [START min_globally]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ min_element = (
+ pipeline
+ | 'Create numbers' >> beam.Create([3, 4, 1, 2])
+ | 'Get min value' >> beam.CombineGlobally(
+ lambda elements: min(elements or [-1]))
+ | beam.Map(print)
+ )
+ # [END min_globally]
+ if test:
+ test(min_element)
+
+
+def min_per_key(test=None):
+ # [START min_per_key]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ elements_with_min_value_per_key = (
+ pipeline
+ | 'Create produce' >> beam.Create([
+ ('🥕', 3),
+ ('🥕', 2),
+ ('🍆', 1),
+ ('🍅', 4),
+ ('🍅', 5),
+ ('🍅', 3),
+ ])
+ | 'Get min value per key' >> beam.CombinePerKey(min)
+ | beam.Map(print)
+ )
+ # [END min_per_key]
+ if test:
+ test(elements_with_min_value_per_key)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min_test.py
new file mode 100644
index 0000000..321fd12
--- /dev/null
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min_test.py
@@ -0,0 +1,60 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+import unittest
+
+import mock
+
+from apache_beam.examples.snippets.util import assert_matches_stdout
+from apache_beam.testing.test_pipeline import TestPipeline
+
+from . import min as beam_min
+
+
+def check_min_element(actual):
+ expected = '''[START min_element]
+1
+[END min_element]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+def check_elements_with_min_value_per_key(actual):
+ expected = '''[START elements_with_min_value_per_key]
+('🥕', 2)
+('🍆', 1)
+('🍅', 3)
+[END elements_with_min_value_per_key]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+@mock.patch('apache_beam.Pipeline', TestPipeline)
+@mock.patch(
+ 'apache_beam.examples.snippets.transforms.aggregation.min.print', str)
+class MinTest(unittest.TestCase):
+ def test_min_globally(self):
+ beam_min.min_globally(check_min_element)
+
+ def test_min_per_key(self):
+ beam_min.min_per_key(check_elements_with_min_value_per_key)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample.py
index d5abc37..a81ab67 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py
index 22cd656..bb337ee 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum.py
new file mode 100644
index 0000000..26094a5
--- /dev/null
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum.py
@@ -0,0 +1,59 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+
+def sum_globally(test=None):
+ # [START sum_globally]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ total = (
+ pipeline
+ | 'Create numbers' >> beam.Create([3, 4, 1, 2])
+ | 'Sum values' >> beam.CombineGlobally(sum)
+ | beam.Map(print)
+ )
+ # [END sum_globally]
+ if test:
+ test(total)
+
+
+def sum_per_key(test=None):
+ # [START sum_per_key]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ totals_per_key = (
+ pipeline
+ | 'Create produce' >> beam.Create([
+ ('🥕', 3),
+ ('🥕', 2),
+ ('🍆', 1),
+ ('🍅', 4),
+ ('🍅', 5),
+ ('🍅', 3),
+ ])
+ | 'Sum values per key' >> beam.CombinePerKey(sum)
+ | beam.Map(print)
+ )
+ # [END sum_per_key]
+ if test:
+ test(totals_per_key)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py
new file mode 100644
index 0000000..dd59770
--- /dev/null
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py
@@ -0,0 +1,60 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+import unittest
+
+import mock
+
+from apache_beam.examples.snippets.util import assert_matches_stdout
+from apache_beam.testing.test_pipeline import TestPipeline
+
+from . import sum as beam_sum
+
+
+def check_total(actual):
+ expected = '''[START total]
+10
+[END total]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+def check_totals_per_key(actual):
+ expected = '''[START totals_per_key]
+('🥕', 5)
+('🍆', 1)
+('🍅', 12)
+[END totals_per_key]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+@mock.patch('apache_beam.Pipeline', TestPipeline)
+@mock.patch(
+ 'apache_beam.examples.snippets.transforms.aggregation.sum.print', str)
+class SumTest(unittest.TestCase):
+ def test_sum_globally(self):
+ beam_sum.sum_globally(check_total)
+
+ def test_sum_per_key(self):
+ beam_sum.sum_per_key(check_totals_per_key)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py
new file mode 100644
index 0000000..0e496f7
--- /dev/null
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py
@@ -0,0 +1,153 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+
+def top_largest(test=None):
+ # [START top_largest]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ largest_elements = (
+ pipeline
+ | 'Create numbers' >> beam.Create([3, 4, 1, 2])
+ | 'Largest N values' >> beam.combiners.Top.Largest(2)
+ | beam.Map(print)
+ )
+ # [END top_largest]
+ if test:
+ test(largest_elements)
+
+
+def top_largest_per_key(test=None):
+ # [START top_largest_per_key]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ largest_elements_per_key = (
+ pipeline
+ | 'Create produce' >> beam.Create([
+ ('🥕', 3),
+ ('🥕', 2),
+ ('🍆', 1),
+ ('🍅', 4),
+ ('🍅', 5),
+ ('🍅', 3),
+ ])
+ | 'Largest N values per key' >> beam.combiners.Top.LargestPerKey(2)
+ | beam.Map(print)
+ )
+ # [END top_largest_per_key]
+ if test:
+ test(largest_elements_per_key)
+
+
+def top_smallest(test=None):
+ # [START top_smallest]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ smallest_elements = (
+ pipeline
+ | 'Create numbers' >> beam.Create([3, 4, 1, 2])
+ | 'Smallest N values' >> beam.combiners.Top.Smallest(2)
+ | beam.Map(print)
+ )
+ # [END top_smallest]
+ if test:
+ test(smallest_elements)
+
+
+def top_smallest_per_key(test=None):
+ # [START top_smallest_per_key]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ smallest_elements_per_key = (
+ pipeline
+ | 'Create produce' >> beam.Create([
+ ('🥕', 3),
+ ('🥕', 2),
+ ('🍆', 1),
+ ('🍅', 4),
+ ('🍅', 5),
+ ('🍅', 3),
+ ])
+ | 'Smallest N values per key' >> beam.combiners.Top.SmallestPerKey(2)
+ | beam.Map(print)
+ )
+ # [END top_smallest_per_key]
+ if test:
+ test(smallest_elements_per_key)
+
+
+def top_of(test=None):
+ # [START top_of]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ shortest_elements = (
+ pipeline
+ | 'Create produce names' >> beam.Create([
+ '🍓 Strawberry',
+ '🥕 Carrot',
+ '🍏 Green apple',
+ '🍆 Eggplant',
+ '🌽 Corn',
+ ])
+ | 'Shortest names' >> beam.combiners.Top.Of(
+ 2, # number of elements
+ key=len, # optional, defaults to the element itself
+ reverse=True, # optional, defaults to False (largest/descending)
+ )
+ | beam.Map(print)
+ )
+ # [END top_of]
+ if test:
+ test(shortest_elements)
+
+
+def top_per_key(test=None):
+ # [START top_per_key]
+ import apache_beam as beam
+
+ with beam.Pipeline() as pipeline:
+ shortest_elements_per_key = (
+ pipeline
+ | 'Create produce names' >> beam.Create([
+ ('spring', '🥕 Carrot'),
+ ('spring', '🍓 Strawberry'),
+ ('summer', '🥕 Carrot'),
+ ('summer', '🌽 Corn'),
+ ('summer', '🍏 Green apple'),
+ ('fall', '🥕 Carrot'),
+ ('fall', '🍏 Green apple'),
+ ('winter', '🍆 Eggplant'),
+ ])
+ | 'Shortest names per key' >> beam.combiners.Top.PerKey(
+ 2, # number of elements
+ key=len, # optional, defaults to the value itself
+ reverse=True, # optional, defaults to False (largest/descending)
+ )
+ | beam.Map(print)
+ )
+ # [END top_per_key]
+ if test:
+ test(shortest_elements_per_key)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py
new file mode 100644
index 0000000..0489f6d
--- /dev/null
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py
@@ -0,0 +1,111 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+import sys
+import unittest
+
+import mock
+
+from apache_beam.examples.snippets.util import assert_matches_stdout
+from apache_beam.testing.test_pipeline import TestPipeline
+
+from . import top
+
+
+def check_largest_elements(actual):
+ expected = '''[START largest_elements]
+[4, 3]
+[END largest_elements]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+def check_largest_elements_per_key(actual):
+ expected = '''[START largest_elements_per_key]
+('🥕', [3, 2])
+('🍆', [1])
+('🍅', [5, 4])
+[END largest_elements_per_key]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+def check_smallest_elements(actual):
+ expected = '''[START smallest_elements]
+[1, 2]
+[END smallest_elements]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+def check_smallest_elements_per_key(actual):
+ expected = '''[START smallest_elements_per_key]
+('🥕', [2, 3])
+('🍆', [1])
+('🍅', [3, 4])
+[END smallest_elements_per_key]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+def check_shortest_elements(actual):
+ expected = '''[START shortest_elements]
+['🌽 Corn', '🥕 Carrot']
+[END shortest_elements]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+def check_shortest_elements_per_key(actual):
+ expected = '''[START shortest_elements_per_key]
+('spring', ['🥕 Carrot', '🍓 Strawberry'])
+('summer', ['🌽 Corn', '🥕 Carrot'])
+('fall', ['🥕 Carrot', '🍏 Green apple'])
+('winter', ['🍆 Eggplant'])
+[END shortest_elements_per_key]'''.splitlines()[1:-1]
+ assert_matches_stdout(actual, expected)
+
+
+@mock.patch('apache_beam.Pipeline', TestPipeline)
+@mock.patch(
+ 'apache_beam.examples.snippets.transforms.aggregation.top.print', str)
+class TopTest(unittest.TestCase):
+ def test_top_largest(self):
+ top.top_largest(check_largest_elements)
+
+ def test_top_largest_per_key(self):
+ top.top_largest_per_key(check_largest_elements_per_key)
+
+ def test_top_smallest(self):
+ top.top_smallest(check_smallest_elements)
+
+ def test_top_smallest_per_key(self):
+ top.top_smallest_per_key(check_smallest_elements_per_key)
+
+ def test_top_of(self):
+ top.top_of(check_shortest_elements)
+
+ # TODO: Remove this after Python 2 deprecation.
+ # https://issues.apache.org/jira/browse/BEAM-8124
+ @unittest.skipIf(sys.version_info[0] == 2,
+ 'nosetests in Python 2 uses ascii instead of utf-8 in '
+ 'the Top.PerKey transform and causes this to fail')
+ def test_top_per_key(self):
+ top.top_per_key(check_shortest_elements_per_key)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter.py
index 44b11b8..ddb58bd 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py
index 724b1b9..10ad79d 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap.py
index 50ffe7a..ef2856c 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py
index 5c326e9..9e67e74 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys.py
index 01c9d6b..2a975be 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py
index e4a843b..9455836 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap.py
index 2107fd5..e6646bf 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py
index 83f211d..c3eed27 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map.py
index 9defd47..19d30d7 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py
index eb77675..2f34a62 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo.py
index 4ecd74d..8e39556 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py
index cbf4903..ae8515d 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition.py
index 5633607..9820bbf 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py
index 4f98ab1..db81891 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex.py
index b39b534..6f9e230 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py
index 9df9f62..d783b1f 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring.py
index 1d0b7dd..3275433 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py
index 04939a7..e142278 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values.py
index 8504ff4..a1ee8b2 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py
index 7a3b8f3..6c2caaa 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps.py
index 79a9c44..bf4a4d4 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py
index ad8c31b..191d114 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/examples/snippets/util.py b/sdks/python/apache_beam/examples/snippets/util.py
index 60c2c7e..a14cd36 100644
--- a/sdks/python/apache_beam/examples/snippets/util.py
+++ b/sdks/python/apache_beam/examples/snippets/util.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import ast
diff --git a/sdks/python/apache_beam/examples/snippets/util_test.py b/sdks/python/apache_beam/examples/snippets/util_test.py
index fcf3955..756f857 100644
--- a/sdks/python/apache_beam/examples/snippets/util_test.py
+++ b/sdks/python/apache_beam/examples/snippets/util_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount.py b/sdks/python/apache_beam/examples/streaming_wordcount.py
index f0db06a..cdfb6a1 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount.py
@@ -18,6 +18,8 @@
"""A streaming word-counting workflow.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
@@ -56,50 +58,48 @@
pipeline_options = PipelineOptions(pipeline_args)
pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
pipeline_options.view_as(StandardOptions).streaming = True
- p = beam.Pipeline(options=pipeline_options)
+ with beam.Pipeline(options=pipeline_options) as p:
- # Read from PubSub into a PCollection.
- if known_args.input_subscription:
- messages = (p
- | beam.io.ReadFromPubSub(
- subscription=known_args.input_subscription)
- .with_output_types(bytes))
- else:
- messages = (p
- | beam.io.ReadFromPubSub(topic=known_args.input_topic)
- .with_output_types(bytes))
+ # Read from PubSub into a PCollection.
+ if known_args.input_subscription:
+ messages = (p
+ | beam.io.ReadFromPubSub(
+ subscription=known_args.input_subscription)
+ .with_output_types(bytes))
+ else:
+ messages = (p
+ | beam.io.ReadFromPubSub(topic=known_args.input_topic)
+ .with_output_types(bytes))
- lines = messages | 'decode' >> beam.Map(lambda x: x.decode('utf-8'))
+ lines = messages | 'decode' >> beam.Map(lambda x: x.decode('utf-8'))
- # Count the occurrences of each word.
- def count_ones(word_ones):
- (word, ones) = word_ones
- return (word, sum(ones))
+ # Count the occurrences of each word.
+ def count_ones(word_ones):
+ (word, ones) = word_ones
+ return (word, sum(ones))
- counts = (lines
- | 'split' >> (beam.ParDo(WordExtractingDoFn())
- .with_output_types(unicode))
- | 'pair_with_one' >> beam.Map(lambda x: (x, 1))
- | beam.WindowInto(window.FixedWindows(15, 0))
- | 'group' >> beam.GroupByKey()
- | 'count' >> beam.Map(count_ones))
+ counts = (lines
+ | 'split' >> (beam.ParDo(WordExtractingDoFn())
+ .with_output_types(unicode))
+ | 'pair_with_one' >> beam.Map(lambda x: (x, 1))
+ | beam.WindowInto(window.FixedWindows(15, 0))
+ | 'group' >> beam.GroupByKey()
+ | 'count' >> beam.Map(count_ones))
- # Format the counts into a PCollection of strings.
- def format_result(word_count):
- (word, count) = word_count
- return '%s: %d' % (word, count)
+ # Format the counts into a PCollection of strings.
+ def format_result(word_count):
+ (word, count) = word_count
+ return '%s: %d' % (word, count)
- output = (counts
- | 'format' >> beam.Map(format_result)
- | 'encode' >> beam.Map(lambda x: x.encode('utf-8'))
- .with_output_types(bytes))
+ output = (counts
+ | 'format' >> beam.Map(format_result)
+ | 'encode' >> beam.Map(lambda x: x.encode('utf-8'))
+ .with_output_types(bytes))
- # Write to PubSub.
- # pylint: disable=expression-not-assigned
- output | beam.io.WriteToPubSub(known_args.output_topic)
+ # Write to PubSub.
+ # pylint: disable=expression-not-assigned
+ output | beam.io.WriteToPubSub(known_args.output_topic)
- result = p.run()
- result.wait_until_finish()
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py b/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
index edaedb5..79eecea 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
@@ -32,6 +32,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
@@ -101,82 +103,80 @@
pipeline_options = PipelineOptions(pipeline_args)
pipeline_options.view_as(SetupOptions).save_main_session = True
pipeline_options.view_as(StandardOptions).streaming = True
- p = beam.Pipeline(options=pipeline_options)
+ with beam.Pipeline(options=pipeline_options) as p:
- # Read from PubSub into a PCollection.
- if known_args.input_subscription:
- lines = p | beam.io.ReadFromPubSub(
- subscription=known_args.input_subscription)
- else:
- lines = p | beam.io.ReadFromPubSub(topic=known_args.input_topic)
+ # Read from PubSub into a PCollection.
+ if known_args.input_subscription:
+ lines = p | beam.io.ReadFromPubSub(
+ subscription=known_args.input_subscription)
+ else:
+ lines = p | beam.io.ReadFromPubSub(topic=known_args.input_topic)
- # Count the occurrences of each word.
- def count_ones(word_ones):
- (word, ones) = word_ones
- return (word, sum(ones))
+ # Count the occurrences of each word.
+ def count_ones(word_ones):
+ (word, ones) = word_ones
+ return (word, sum(ones))
- counts = (lines
- | 'AddTimestampFn' >> beam.ParDo(AddTimestampFn())
- | 'After AddTimestampFn' >> ParDo(PrintFn('After AddTimestampFn'))
- | 'Split' >> (beam.ParDo(WordExtractingDoFn())
- .with_output_types(unicode))
- | 'PairWithOne' >> beam.Map(lambda x: (x, 1))
- | beam.WindowInto(window.FixedWindows(5, 0))
- | 'GroupByKey' >> beam.GroupByKey()
- | 'CountOnes' >> beam.Map(count_ones))
+ counts = (lines
+ | 'AddTimestampFn' >> beam.ParDo(AddTimestampFn())
+ | 'After AddTimestampFn' >> ParDo(PrintFn('After AddTimestampFn'))
+ | 'Split' >> (beam.ParDo(WordExtractingDoFn())
+ .with_output_types(unicode))
+ | 'PairWithOne' >> beam.Map(lambda x: (x, 1))
+ | beam.WindowInto(window.FixedWindows(5, 0))
+ | 'GroupByKey' >> beam.GroupByKey()
+ | 'CountOnes' >> beam.Map(count_ones))
- # Format the counts into a PCollection of strings.
- def format_result(word_count):
- (word, count) = word_count
- return '%s: %d' % (word, count)
+ # Format the counts into a PCollection of strings.
+ def format_result(word_count):
+ (word, count) = word_count
+ return '%s: %d' % (word, count)
- output = counts | 'format' >> beam.Map(format_result)
+ output = counts | 'format' >> beam.Map(format_result)
- # Write to PubSub.
- # pylint: disable=expression-not-assigned
- output | beam.io.WriteStringsToPubSub(known_args.output_topic)
+ # Write to PubSub.
+ # pylint: disable=expression-not-assigned
+ output | beam.io.WriteStringsToPubSub(known_args.output_topic)
- def check_gbk_format():
- # A matcher that checks that the output of GBK is of the form word: count.
- def matcher(elements):
- # pylint: disable=unused-variable
- actual_elements_in_window, window = elements
- for elm in actual_elements_in_window:
- assert re.match(r'\S+:\s+\d+', elm) is not None
- return matcher
+ def check_gbk_format():
+ # A matcher that checks that the output of GBK is of the form word: count.
+ def matcher(elements):
+ # pylint: disable=unused-variable
+ actual_elements_in_window, window = elements
+ for elm in actual_elements_in_window:
+ assert re.match(r'\S+:\s+\d+', elm) is not None
+ return matcher
- # Check that the format of the output is correct.
- assert_that(
- output,
- check_gbk_format(),
- use_global_window=False,
- label='Assert word:count format.')
+ # Check that the format of the output is correct.
+ assert_that(
+ output,
+ check_gbk_format(),
+ use_global_window=False,
+ label='Assert word:count format.')
- # Check also that elements are ouput in the right window.
- # This expects exactly 1 occurrence of any subset of the elements
- # 150, 151, 152, 153, 154 in the window [150, 155)
- # or exactly 1 occurrence of any subset of the elements
- # 210, 211, 212, 213, 214 in the window [210, 215).
- expected_window_to_elements = {
- window.IntervalWindow(150, 155): [
- ('150: 1'), ('151: 1'), ('152: 1'), ('153: 1'), ('154: 1'),
- ],
- window.IntervalWindow(210, 215): [
- ('210: 1'), ('211: 1'), ('212: 1'), ('213: 1'), ('214: 1'),
- ],
- }
+ # Check also that elements are ouput in the right window.
+ # This expects exactly 1 occurrence of any subset of the elements
+ # 150, 151, 152, 153, 154 in the window [150, 155)
+ # or exactly 1 occurrence of any subset of the elements
+ # 210, 211, 212, 213, 214 in the window [210, 215).
+ expected_window_to_elements = {
+ window.IntervalWindow(150, 155): [
+ ('150: 1'), ('151: 1'), ('152: 1'), ('153: 1'), ('154: 1'),
+ ],
+ window.IntervalWindow(210, 215): [
+ ('210: 1'), ('211: 1'), ('212: 1'), ('213: 1'), ('214: 1'),
+ ],
+ }
- # To make it pass, publish numbers in [150-155) or [210-215) with no repeats.
- # To make it fail, publish a repeated number in the range above range.
- # For example: '210 213 151 213'
- assert_that(
- output,
- equal_to_per_window(expected_window_to_elements),
- use_global_window=False,
- label='Assert correct streaming windowing.')
+ # To pass, publish numbers in [150-155) or [210-215) with no repeats.
+ # To fail, publish a repeated number in the range above range.
+ # For example: '210 213 151 213'
+ assert_that(
+ output,
+ equal_to_per_window(expected_window_to_elements),
+ use_global_window=False,
+ label='Assert correct streaming windowing.')
- result = p.run()
- result.wait_until_finish()
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
index d87d0f4..96a6331 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
@@ -17,6 +17,8 @@
"""End-to-end test for the streaming wordcount example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/windowed_wordcount.py b/sdks/python/apache_beam/examples/windowed_wordcount.py
index 5eb05c0..c59c83b 100644
--- a/sdks/python/apache_beam/examples/windowed_wordcount.py
+++ b/sdks/python/apache_beam/examples/windowed_wordcount.py
@@ -21,6 +21,8 @@
and is not yet available for use.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/wordcount.py b/sdks/python/apache_beam/examples/wordcount.py
index a8f17e3..6fd53b1 100644
--- a/sdks/python/apache_beam/examples/wordcount.py
+++ b/sdks/python/apache_beam/examples/wordcount.py
@@ -17,6 +17,8 @@
"""A word-counting workflow."""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/wordcount_debugging.py b/sdks/python/apache_beam/examples/wordcount_debugging.py
index 389bdd6..bd35bfbb 100644
--- a/sdks/python/apache_beam/examples/wordcount_debugging.py
+++ b/sdks/python/apache_beam/examples/wordcount_debugging.py
@@ -39,6 +39,8 @@
--output gs://YOUR_OUTPUT_PREFIX
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/wordcount_debugging_test.py b/sdks/python/apache_beam/examples/wordcount_debugging_test.py
index 124b680..51f150a 100644
--- a/sdks/python/apache_beam/examples/wordcount_debugging_test.py
+++ b/sdks/python/apache_beam/examples/wordcount_debugging_test.py
@@ -17,6 +17,8 @@
"""Test for the debugging wordcount example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/wordcount_it_test.py b/sdks/python/apache_beam/examples/wordcount_it_test.py
index 8a9b2c5..bf12ba3 100644
--- a/sdks/python/apache_beam/examples/wordcount_it_test.py
+++ b/sdks/python/apache_beam/examples/wordcount_it_test.py
@@ -17,6 +17,8 @@
"""End-to-end test for the wordcount example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/examples/wordcount_minimal.py b/sdks/python/apache_beam/examples/wordcount_minimal.py
index 2bfb6ec..92982bd 100644
--- a/sdks/python/apache_beam/examples/wordcount_minimal.py
+++ b/sdks/python/apache_beam/examples/wordcount_minimal.py
@@ -44,6 +44,8 @@
pipeline. You can see the results in your output bucket in the GCS browser.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/examples/wordcount_minimal_test.py b/sdks/python/apache_beam/examples/wordcount_minimal_test.py
index 9a772d5..a3912a1 100644
--- a/sdks/python/apache_beam/examples/wordcount_minimal_test.py
+++ b/sdks/python/apache_beam/examples/wordcount_minimal_test.py
@@ -17,6 +17,8 @@
"""Test for the minimal wordcount example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/examples/wordcount_test.py b/sdks/python/apache_beam/examples/wordcount_test.py
index 84b14f2..0f140d4 100644
--- a/sdks/python/apache_beam/examples/wordcount_test.py
+++ b/sdks/python/apache_beam/examples/wordcount_test.py
@@ -18,6 +18,8 @@
"""Test for the wordcount example."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/examples/wordcount_xlang.py b/sdks/python/apache_beam/examples/wordcount_xlang.py
index fe1994d..b8353bb 100644
--- a/sdks/python/apache_beam/examples/wordcount_xlang.py
+++ b/sdks/python/apache_beam/examples/wordcount_xlang.py
@@ -17,6 +17,8 @@
"""A cross-language word-counting workflow."""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
@@ -25,6 +27,7 @@
import subprocess
import grpc
+from past.builtins import unicode
import apache_beam as beam
from apache_beam.io import ReadFromText
@@ -53,22 +56,16 @@
The processed element.
"""
text_line = element.strip()
- # Using bytes type to match input and output coders between Python
- # and Java SDKs. Any element type can be used for crossing the language
- # boundary if a matching coder implementation exists in both SDKs.
- # TODO(BEAM-6587): Use strings once they're understood by the
- # Java SDK.
- words = [bytes(x) for x in re.findall(r'[\w\']+', text_line)]
- return words
+ return re.findall(r'[\w\']+', text_line)
-def run(p, input_file, output_file):
+def build_pipeline(p, input_file, output_file):
# Read the text file[pattern] into a PCollection.
lines = p | 'read' >> ReadFromText(input_file)
counts = (lines
| 'split' >> (beam.ParDo(WordExtractingDoFn())
- .with_output_types(bytes))
+ .with_output_types(unicode))
| 'count' >> beam.ExternalTransform(
'beam:transforms:xlang:count', None, EXPANSION_SERVICE_ADDR))
@@ -83,9 +80,6 @@
# pylint: disable=expression-not-assigned
output | 'write' >> WriteToText(output_file)
- result = p.run()
- result.wait_until_finish()
-
def main():
logging.getLogger().setLevel(logging.INFO)
@@ -115,10 +109,6 @@
# workflow rely on global context (e.g., a module imported at module level).
pipeline_options.view_as(SetupOptions).save_main_session = True
- p = beam.Pipeline(options=pipeline_options)
- # Preemptively start due to BEAM-6666.
- p.runner.create_job_service(pipeline_options)
-
try:
server = subprocess.Popen([
'java', '-jar', known_args.expansion_service_jar,
@@ -127,7 +117,11 @@
with grpc.insecure_channel(EXPANSION_SERVICE_ADDR) as channel:
grpc.channel_ready_future(channel).result()
- run(p, known_args.input, known_args.output)
+ with beam.Pipeline(options=pipeline_options) as p:
+ # Preemptively start due to BEAM-6666.
+ p.runner.create_job_service(pipeline_options)
+
+ build_pipeline(p, known_args.input, known_args.output)
finally:
server.kill()
diff --git a/sdks/python/apache_beam/internal/gcp/auth.py b/sdks/python/apache_beam/internal/gcp/auth.py
index 8a94acf..3921f73 100644
--- a/sdks/python/apache_beam/internal/gcp/auth.py
+++ b/sdks/python/apache_beam/internal/gcp/auth.py
@@ -17,6 +17,8 @@
"""Dataflow credentials and authentication."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -126,7 +128,8 @@
'https://www.googleapis.com/auth/cloud-platform',
'https://www.googleapis.com/auth/devstorage.full_control',
'https://www.googleapis.com/auth/userinfo.email',
- 'https://www.googleapis.com/auth/datastore'
+ 'https://www.googleapis.com/auth/datastore',
+ 'https://www.googleapis.com/auth/spanner'
]
try:
credentials = GoogleCredentials.get_application_default()
diff --git a/sdks/python/apache_beam/internal/gcp/json_value.py b/sdks/python/apache_beam/internal/gcp/json_value.py
index c02b639..469da80 100644
--- a/sdks/python/apache_beam/internal/gcp/json_value.py
+++ b/sdks/python/apache_beam/internal/gcp/json_value.py
@@ -17,6 +17,8 @@
"""JSON conversion utility functions."""
+# pytype: skip-file
+
from __future__ import absolute_import
from past.builtins import long
diff --git a/sdks/python/apache_beam/internal/gcp/json_value_test.py b/sdks/python/apache_beam/internal/gcp/json_value_test.py
index 5605d41..d473c30 100644
--- a/sdks/python/apache_beam/internal/gcp/json_value_test.py
+++ b/sdks/python/apache_beam/internal/gcp/json_value_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the json_value module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/internal/http_client.py b/sdks/python/apache_beam/internal/http_client.py
index 1263687..c86c78f 100644
--- a/sdks/python/apache_beam/internal/http_client.py
+++ b/sdks/python/apache_beam/internal/http_client.py
@@ -19,6 +19,8 @@
For internal use only. No backwards compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/internal/http_client_test.py b/sdks/python/apache_beam/internal/http_client_test.py
index c3c0f83..98fc3f2 100644
--- a/sdks/python/apache_beam/internal/http_client_test.py
+++ b/sdks/python/apache_beam/internal/http_client_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the http_client module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import os
diff --git a/sdks/python/apache_beam/internal/module_test.py b/sdks/python/apache_beam/internal/module_test.py
index 8016c6f..45f5792 100644
--- a/sdks/python/apache_beam/internal/module_test.py
+++ b/sdks/python/apache_beam/internal/module_test.py
@@ -17,6 +17,8 @@
"""Module used to define functions and classes used by the coder unit tests."""
+# pytype: skip-file
+
from __future__ import absolute_import
import re
diff --git a/sdks/python/apache_beam/internal/pickler.py b/sdks/python/apache_beam/internal/pickler.py
index 0296840..072add2 100644
--- a/sdks/python/apache_beam/internal/pickler.py
+++ b/sdks/python/apache_beam/internal/pickler.py
@@ -28,6 +28,8 @@
the coders.*PickleCoder classes should be used instead.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
diff --git a/sdks/python/apache_beam/internal/pickler_test.py b/sdks/python/apache_beam/internal/pickler_test.py
index e18c726..e70c28c 100644
--- a/sdks/python/apache_beam/internal/pickler_test.py
+++ b/sdks/python/apache_beam/internal/pickler_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the pickler module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
diff --git a/sdks/python/apache_beam/internal/util.py b/sdks/python/apache_beam/internal/util.py
index 6b13a37..a161440 100644
--- a/sdks/python/apache_beam/internal/util.py
+++ b/sdks/python/apache_beam/internal/util.py
@@ -20,6 +20,8 @@
For internal use only. No backwards compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/internal/util_test.py b/sdks/python/apache_beam/internal/util_test.py
index c3ae191..27f1362 100644
--- a/sdks/python/apache_beam/internal/util_test.py
+++ b/sdks/python/apache_beam/internal/util_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the util module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/io/avroio.py b/sdks/python/apache_beam/io/avroio.py
index b1cf6bd..6f68494 100644
--- a/sdks/python/apache_beam/io/avroio.py
+++ b/sdks/python/apache_beam/io/avroio.py
@@ -41,6 +41,8 @@
that can be used to write a given ``PCollection`` of Python objects to an
Avro file.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import io
diff --git a/sdks/python/apache_beam/io/avroio_test.py b/sdks/python/apache_beam/io/avroio_test.py
index 1b27782..147be62 100644
--- a/sdks/python/apache_beam/io/avroio_test.py
+++ b/sdks/python/apache_beam/io/avroio_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py b/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py
index 4aa8416..9750c9b 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from apache_beam.io.aws.clients.s3 import messages
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/client_test.py b/sdks/python/apache_beam/io/aws/clients/s3/client_test.py
index e78c484..434d173 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/client_test.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/client_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
import logging
import os
import unittest
@@ -68,8 +70,8 @@
# Test nonexistent object
object = self.test_path + 'nonexistent_file_doesnt_exist'
request = messages.GetRequest(self.test_bucket, object)
- self.assertRaises(messages.S3ClientError,
- self.client.get_range,
+ self.assertRaises(messages.S3ClientError,
+ self.client.get_range,
request, 0, 10)
try:
@@ -128,9 +130,9 @@
src_key = self.test_path + 'not_a_real_file_does_not_exist'
dest_key = self.test_path + 'destination_file_location'
- request = messages.CopyRequest(self.test_bucket,
- src_key,
- self.test_bucket,
+ request = messages.CopyRequest(self.test_bucket,
+ src_key,
+ self.test_bucket,
dest_key)
with self.assertRaises(messages.S3ClientError) as e:
@@ -150,14 +152,14 @@
upload_id = response.upload_id
part_number = 0.5
- request = messages.UploadPartRequest(self.test_bucket,
- object,
- upload_id,
- part_number,
+ request = messages.UploadPartRequest(self.test_bucket,
+ object,
+ upload_id,
+ part_number,
contents)
self.assertRaises(messages.S3ClientError,
- self.client.upload_part,
+ self.client.upload_part,
request)
try:
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py b/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py
index 0c7e6c3..597c0e3 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import datetime
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/messages.py b/sdks/python/apache_beam/io/aws/clients/s3/messages.py
index 530f454..40e548d 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/messages.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/messages.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
diff --git a/sdks/python/apache_beam/io/aws/s3filesystem.py b/sdks/python/apache_beam/io/aws/s3filesystem.py
index 4fc7fb1..d1d79c8 100644
--- a/sdks/python/apache_beam/io/aws/s3filesystem.py
+++ b/sdks/python/apache_beam/io/aws/s3filesystem.py
@@ -16,6 +16,8 @@
#
"""S3 file system implementation for accessing files on AWS S3."""
+# pytype: skip-file
+
from __future__ import absolute_import
from future.utils import iteritems
@@ -269,7 +271,7 @@
paths: list of paths that give the file objects to be deleted
"""
results = s3io.S3IO().delete_paths(paths)
- exceptions = {path: error for (path, error) in results
+ exceptions = {path: error for (path, error) in results.items()
if error is not None}
if exceptions:
raise BeamIOError("Delete operation failed", exceptions)
diff --git a/sdks/python/apache_beam/io/aws/s3filesystem_test.py b/sdks/python/apache_beam/io/aws/s3filesystem_test.py
index 4bddbf2..5439411 100644
--- a/sdks/python/apache_beam/io/aws/s3filesystem_test.py
+++ b/sdks/python/apache_beam/io/aws/s3filesystem_test.py
@@ -18,6 +18,8 @@
"""Unit tests for the S3 File System"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -129,7 +131,7 @@
with self.assertRaises(BeamIOError) as error:
self.fs.match(['s3://bucket/'])
- self.assertTrue('Match operation failed' in str(error.exception))
+ self.assertIn('Match operation failed', str(error.exception))
s3io_mock.list_prefix.assert_called_once_with('s3://bucket/')
@mock.patch('apache_beam.io.aws.s3filesystem.s3io')
@@ -225,11 +227,12 @@
problematic_directory = 's3://nonexistent-bucket/tree/'
exception = messages.S3ClientError('Not found', 404)
- s3io_mock.delete_paths.return_value = [
- (problematic_directory, exception),
- ('s3://bucket/object1', None),
- ('s3://bucket/object2', None)
- ]
+ s3io_mock.delete_paths.return_value = {
+ problematic_directory: exception,
+ 's3://bucket/object1': None,
+ 's3://bucket/object2': None,
+ }
+
s3io_mock.size.return_value = 0
files = [
problematic_directory,
@@ -241,7 +244,7 @@
# Issue batch delete.
with self.assertRaises(BeamIOError) as error:
self.fs.delete(files)
- self.assertTrue('Delete operation failed' in str(error.exception))
+ self.assertIn('Delete operation failed', str(error.exception))
self.assertEqual(error.exception.exception_details, expected_results)
s3io_mock.delete_paths.assert_called()
diff --git a/sdks/python/apache_beam/io/aws/s3io.py b/sdks/python/apache_beam/io/aws/s3io.py
index f5f8e0a..30be71c 100644
--- a/sdks/python/apache_beam/io/aws/s3io.py
+++ b/sdks/python/apache_beam/io/aws/s3io.py
@@ -17,6 +17,8 @@
"""AWS S3 client
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import errno
@@ -118,7 +120,18 @@
logging.info("Starting the size estimation of the input")
while True:
- response = self.client.list(request)
+ #The list operation will raise an exception
+ #when trying to list a nonexistent S3 path.
+ #This should not be an issue here.
+ #Ignore this exception or it will break the procedure.
+ try:
+ response = self.client.list(request)
+ except messages.S3ClientError as e:
+ if e.code == 404:
+ break
+ else:
+ raise e
+
for item in response.items:
file_name = 's3://%s/%s' % (bucket, item.key)
file_sizes[file_name] = item.size
diff --git a/sdks/python/apache_beam/io/aws/s3io_test.py b/sdks/python/apache_beam/io/aws/s3io_test.py
index 1f55db4..fd5411c 100644
--- a/sdks/python/apache_beam/io/aws/s3io_test.py
+++ b/sdks/python/apache_beam/io/aws/s3io_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
"""Tests for S3 client."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -243,9 +245,12 @@
to_path + 'fake_directory_1/',
to_path + 'fake_directory_2'
]
-
result = self.aws.copy_paths(list(zip(sources, destinations)))
- self.assertEqual(len(result), len(sources))
+
+ # The copy_paths function of class S3IO does not return one single
+ # result when copying a directory. Instead, it returns the results
+ # of copying every file in the source directory.
+ self.assertEqual(len(result), len(sources) - 1)
for _, _, err in result[:n_real_files]:
self.assertTrue(err is None)
@@ -253,7 +258,9 @@
for _, _, err in result[n_real_files:]:
self.assertIsInstance(err, messages.S3ClientError)
- self.assertEqual(result[-3][2].code, 404)
+ # For the same reason of copy_paths function of S3IO above
+ # skip this assert.
+ #self.assertEqual(result[-3][2].code, 404)
self.assertEqual(result[-2][2].code, 404)
self.assertEqual(result[-1][2].code, 400)
diff --git a/sdks/python/apache_beam/io/concat_source.py b/sdks/python/apache_beam/io/concat_source.py
index ddf3a77..4446e60 100644
--- a/sdks/python/apache_beam/io/concat_source.py
+++ b/sdks/python/apache_beam/io/concat_source.py
@@ -19,6 +19,8 @@
Concat Source, which reads the union of several other sources.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/concat_source_test.py b/sdks/python/apache_beam/io/concat_source_test.py
index cb504f8..41e4e63 100644
--- a/sdks/python/apache_beam/io/concat_source_test.py
+++ b/sdks/python/apache_beam/io/concat_source_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the sources framework."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -223,11 +225,10 @@
RangeSource(10, 100),
RangeSource(100, 1000),
])
- pipeline = TestPipeline()
- pcoll = pipeline | beam.io.Read(source)
- assert_that(pcoll, equal_to(list(range(1000))))
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | beam.io.Read(source)
+ assert_that(pcoll, equal_to(list(range(1000))))
- pipeline.run()
def test_conact_source_exhaustive(self):
source = ConcatSource([RangeSource(0, 10),
diff --git a/sdks/python/apache_beam/io/external/gcp/pubsub.py b/sdks/python/apache_beam/io/external/gcp/pubsub.py
index f0988ed..d417b42 100644
--- a/sdks/python/apache_beam/io/external/gcp/pubsub.py
+++ b/sdks/python/apache_beam/io/external/gcp/pubsub.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import typing
diff --git a/sdks/python/apache_beam/io/external/generate_sequence.py b/sdks/python/apache_beam/io/external/generate_sequence.py
index a17ec7b..47f9297 100644
--- a/sdks/python/apache_beam/io/external/generate_sequence.py
+++ b/sdks/python/apache_beam/io/external/generate_sequence.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from apache_beam.transforms.external import ExternalTransform
diff --git a/sdks/python/apache_beam/io/external/generate_sequence_test.py b/sdks/python/apache_beam/io/external/generate_sequence_test.py
index 95d6fbc..060ce28 100644
--- a/sdks/python/apache_beam/io/external/generate_sequence_test.py
+++ b/sdks/python/apache_beam/io/external/generate_sequence_test.py
@@ -17,6 +17,8 @@
"""Unit tests for cross-language generate sequence."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
@@ -39,12 +41,11 @@
"EXPANSION_PORT environment var is not provided.")
class XlangGenerateSequenceTest(unittest.TestCase):
def test_generate_sequence(self):
- test_pipeline = TestPipeline()
port = os.environ.get('EXPANSION_PORT')
address = 'localhost:%s' % port
try:
- with test_pipeline as p:
+ with TestPipeline() as p:
res = (
p
| GenerateSequence(start=1, stop=10,
diff --git a/sdks/python/apache_beam/io/external/kafka.py b/sdks/python/apache_beam/io/external/kafka.py
index 04d91a7..44bce78 100644
--- a/sdks/python/apache_beam/io/external/kafka.py
+++ b/sdks/python/apache_beam/io/external/kafka.py
@@ -35,6 +35,8 @@
- https://beam.apache.org/roadmap/portability/
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import typing
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 434bb3b..aee35a0 100644
--- a/sdks/python/apache_beam/io/external/xlang_parquetio_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_parquetio_test.py
@@ -17,6 +17,8 @@
"""Unit tests for cross-language parquet io read/write."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
@@ -51,11 +53,10 @@
port = os.environ.get('EXPANSION_PORT')
address = 'localhost:%s' % port
try:
- test_pipeline = TestPipeline()
- test_pipeline.get_pipeline_options().view_as(
- DebugOptions).experiments.append('jar_packages='+expansion_jar)
- test_pipeline.not_use_test_runner_api = True
- with test_pipeline as p:
+ with TestPipeline() as p:
+ p.get_pipeline_options().view_as(
+ DebugOptions).experiments.append('jar_packages='+expansion_jar)
+ p.not_use_test_runner_api = True
_ = p \
| beam.Create([
AvroRecord({"name": "abc"}), AvroRecord({"name": "def"}),
diff --git a/sdks/python/apache_beam/io/filebasedsink.py b/sdks/python/apache_beam/io/filebasedsink.py
index b8f0e37..e9dfa05 100644
--- a/sdks/python/apache_beam/io/filebasedsink.py
+++ b/sdks/python/apache_beam/io/filebasedsink.py
@@ -17,6 +17,8 @@
"""File-based sink."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/filebasedsink_test.py b/sdks/python/apache_beam/io/filebasedsink_test.py
index 4c5ef6b..6bfebd5 100644
--- a/sdks/python/apache_beam/io/filebasedsink_test.py
+++ b/sdks/python/apache_beam/io/filebasedsink_test.py
@@ -18,6 +18,8 @@
"""Unit tests for file sinks."""
+# pytype: skip-file
+
from __future__ import absolute_import
import glob
diff --git a/sdks/python/apache_beam/io/filebasedsource.py b/sdks/python/apache_beam/io/filebasedsource.py
index 0f87d25..422b15b 100644
--- a/sdks/python/apache_beam/io/filebasedsource.py
+++ b/sdks/python/apache_beam/io/filebasedsource.py
@@ -26,6 +26,8 @@
:class:`~apache_beam.io._AvroSource`.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from typing import Callable
diff --git a/sdks/python/apache_beam/io/filebasedsource_test.py b/sdks/python/apache_beam/io/filebasedsource_test.py
index e777e9f..3c9adbd 100644
--- a/sdks/python/apache_beam/io/filebasedsource_test.py
+++ b/sdks/python/apache_beam/io/filebasedsource_test.py
@@ -14,6 +14,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -434,11 +436,10 @@
self.assertCountEqual(expected_data, read_data)
def _run_source_test(self, pattern, expected_data, splittable=True):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- pattern, splittable=splittable))
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ pattern, splittable=splittable))
+ assert_that(pcoll, equal_to(expected_data))
def test_source_file(self):
file_name, expected_data = write_data(100)
@@ -474,13 +475,12 @@
with bz2.BZ2File(filename, 'wb') as f:
f.write(b'\n'.join(lines))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- filename,
- splittable=False,
- compression_type=CompressionTypes.BZIP2))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ filename,
+ splittable=False,
+ compression_type=CompressionTypes.BZIP2))
+ assert_that(pcoll, equal_to(lines))
def test_read_file_gzip(self):
_, lines = write_data(10)
@@ -489,13 +489,12 @@
with gzip.GzipFile(filename, 'wb') as f:
f.write(b'\n'.join(lines))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- filename,
- splittable=False,
- compression_type=CompressionTypes.GZIP))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ filename,
+ splittable=False,
+ compression_type=CompressionTypes.GZIP))
+ assert_that(pcoll, equal_to(lines))
def test_read_pattern_bzip2(self):
_, lines = write_data(200)
@@ -507,13 +506,12 @@
compressed_chunks.append(
compressobj.compress(b'\n'.join(c)) + compressobj.flush())
file_pattern = write_prepared_pattern(compressed_chunks)
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- file_pattern,
- splittable=False,
- compression_type=CompressionTypes.BZIP2))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ file_pattern,
+ splittable=False,
+ compression_type=CompressionTypes.BZIP2))
+ assert_that(pcoll, equal_to(lines))
def test_read_pattern_gzip(self):
_, lines = write_data(200)
@@ -526,13 +524,12 @@
f.write(b'\n'.join(c))
compressed_chunks.append(out.getvalue())
file_pattern = write_prepared_pattern(compressed_chunks)
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- file_pattern,
- splittable=False,
- compression_type=CompressionTypes.GZIP))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ file_pattern,
+ splittable=False,
+ compression_type=CompressionTypes.GZIP))
+ assert_that(pcoll, equal_to(lines))
def test_read_auto_single_file_bzip2(self):
_, lines = write_data(10)
@@ -541,12 +538,11 @@
with bz2.BZ2File(filename, 'wb') as f:
f.write(b'\n'.join(lines))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- filename,
- compression_type=CompressionTypes.AUTO))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ filename,
+ compression_type=CompressionTypes.AUTO))
+ assert_that(pcoll, equal_to(lines))
def test_read_auto_single_file_gzip(self):
_, lines = write_data(10)
@@ -555,12 +551,11 @@
with gzip.GzipFile(filename, 'wb') as f:
f.write(b'\n'.join(lines))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- filename,
- compression_type=CompressionTypes.AUTO))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ filename,
+ compression_type=CompressionTypes.AUTO))
+ assert_that(pcoll, equal_to(lines))
def test_read_auto_pattern(self):
_, lines = write_data(200)
@@ -574,12 +569,11 @@
compressed_chunks.append(out.getvalue())
file_pattern = write_prepared_pattern(
compressed_chunks, suffixes=['.gz']*len(chunks))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- file_pattern,
- compression_type=CompressionTypes.AUTO))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ file_pattern,
+ compression_type=CompressionTypes.AUTO))
+ assert_that(pcoll, equal_to(lines))
def test_read_auto_pattern_compressed_and_uncompressed(self):
_, lines = write_data(200)
@@ -596,12 +590,11 @@
chunks_to_write.append(b'\n'.join(c))
file_pattern = write_prepared_pattern(chunks_to_write,
suffixes=(['.gz', '']*3))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
- file_pattern,
- compression_type=CompressionTypes.AUTO))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> beam.io.Read(LineSource(
+ file_pattern,
+ compression_type=CompressionTypes.AUTO))
+ assert_that(pcoll, equal_to(lines))
def test_splits_get_coder_from_fbs(self):
class DummyCoder(object):
diff --git a/sdks/python/apache_beam/io/fileio.py b/sdks/python/apache_beam/io/fileio.py
index 76b125f..ce8f0bc 100644
--- a/sdks/python/apache_beam/io/fileio.py
+++ b/sdks/python/apache_beam/io/fileio.py
@@ -88,6 +88,8 @@
No backward compatibility guarantees. Everything in this module is experimental.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py
index c83c8d7..f8803d1 100644
--- a/sdks/python/apache_beam/io/fileio_test.py
+++ b/sdks/python/apache_beam/io/fileio_test.py
@@ -17,6 +17,8 @@
"""Tests for transforms defined in apache_beam.io.fileio."""
+# pytype: skip-file
+
from __future__ import absolute_import
import csv
@@ -93,7 +95,7 @@
assert_that(files_pc, equal_to(files))
- def test_match_files_one_directory_failure(self):
+ def test_match_files_one_directory_failure1(self):
directories = [
'%s%s' % (self._new_tempdir(), os.sep),
'%s%s' % (self._new_tempdir(), os.sep)]
@@ -112,7 +114,7 @@
assert_that(files_pc, equal_to(files))
- def test_match_files_one_directory_failure(self):
+ def test_match_files_one_directory_failure2(self):
directories = [
'%s%s' % (self._new_tempdir(), os.sep),
'%s%s' % (self._new_tempdir(), os.sep)]
diff --git a/sdks/python/apache_beam/io/filesystem.py b/sdks/python/apache_beam/io/filesystem.py
index aa7552d..6643443 100644
--- a/sdks/python/apache_beam/io/filesystem.py
+++ b/sdks/python/apache_beam/io/filesystem.py
@@ -21,6 +21,8 @@
LocalFileSystem, which gets unix-style paths in the form /foo/bar.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/filesystem_test.py b/sdks/python/apache_beam/io/filesystem_test.py
index fbc094f..6880163 100644
--- a/sdks/python/apache_beam/io/filesystem_test.py
+++ b/sdks/python/apache_beam/io/filesystem_test.py
@@ -17,6 +17,8 @@
#
"""Unit tests for filesystem module."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/filesystemio.py b/sdks/python/apache_beam/io/filesystemio.py
index ae71aff..6197046 100644
--- a/sdks/python/apache_beam/io/filesystemio.py
+++ b/sdks/python/apache_beam/io/filesystemio.py
@@ -16,6 +16,8 @@
#
"""Utilities for ``FileSystem`` implementations."""
+# pytype: skip-file
+
from __future__ import absolute_import
import abc
diff --git a/sdks/python/apache_beam/io/filesystemio_test.py b/sdks/python/apache_beam/io/filesystemio_test.py
index 7797eb8..90a397e 100644
--- a/sdks/python/apache_beam/io/filesystemio_test.py
+++ b/sdks/python/apache_beam/io/filesystemio_test.py
@@ -16,6 +16,8 @@
#
"""Tests for filesystemio."""
+# pytype: skip-file
+
from __future__ import absolute_import
import io
diff --git a/sdks/python/apache_beam/io/filesystems.py b/sdks/python/apache_beam/io/filesystems.py
index 39e6100..e3a072d 100644
--- a/sdks/python/apache_beam/io/filesystems.py
+++ b/sdks/python/apache_beam/io/filesystems.py
@@ -17,6 +17,8 @@
"""FileSystems interface class for accessing the correct filesystem"""
+# pytype: skip-file
+
from __future__ import absolute_import
import re
@@ -96,9 +98,10 @@
systems = [fs for fs in FileSystem.get_all_subclasses()
if fs.scheme() == path_scheme]
if len(systems) == 0:
- raise ValueError('Unable to get filesystem from specified path, please use the correct path '
- 'or ensure the required dependency is installed, e.g., pip install apache_beam[gcp]. '
- 'Path specified: %s' % path)
+ raise ValueError(
+ 'Unable to get filesystem from specified path, please use the '
+ 'correct path or ensure the required dependency is installed, '
+ 'e.g., pip install apache_beam[gcp]. Path specified: %s' % path)
elif len(systems) == 1:
# Pipeline options could come either from the Pipeline itself (using
# direct runner), or via RuntimeValueProvider (other runners).
diff --git a/sdks/python/apache_beam/io/filesystems_test.py b/sdks/python/apache_beam/io/filesystems_test.py
index 17cec46..298644f 100644
--- a/sdks/python/apache_beam/io/filesystems_test.py
+++ b/sdks/python/apache_beam/io/filesystems_test.py
@@ -18,12 +18,15 @@
"""Unit tests for LocalFileSystem."""
+# pytype: skip-file
+
from __future__ import absolute_import
import filecmp
import logging
import os
import shutil
+import sys
import tempfile
import unittest
@@ -47,6 +50,12 @@
class FileSystemsTest(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ # Method has been renamed in Python 3
+ if sys.version_info[0] < 3:
+ cls.assertCountEqual = cls.assertItemsEqual
+
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
@@ -130,7 +139,7 @@
FileSystems.match([None])
self.assertEqual(list(error.exception.exception_details), [None])
- def test_match_directory(self):
+ def test_match_directory_with_files(self):
path1 = os.path.join(self.tmpdir, 'f1')
path2 = os.path.join(self.tmpdir, 'f2')
open(path1, 'a').close()
@@ -140,7 +149,7 @@
path = os.path.join(self.tmpdir, '*')
result = FileSystems.match([path])[0]
files = [f.path for f in result.metadata_list]
- self.assertEqual(files, [path1, path2])
+ self.assertCountEqual(files, [path1, path2])
def test_match_directory(self):
result = FileSystems.match([self.tmpdir])[0]
diff --git a/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source.py b/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source.py
index a151409..05b2f2b 100644
--- a/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source.py
+++ b/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source.py
@@ -20,6 +20,8 @@
This can only be used with the flink runner.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
diff --git a/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py b/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py
index d357946..91c100e 100644
--- a/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py
@@ -18,6 +18,8 @@
Integration test for Google Cloud BigQuery.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
diff --git a/sdks/python/apache_beam/io/gcp/big_query_query_to_table_pipeline.py b/sdks/python/apache_beam/io/gcp/big_query_query_to_table_pipeline.py
index fec9f6b..6034aaa 100644
--- a/sdks/python/apache_beam/io/gcp/big_query_query_to_table_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/big_query_query_to_table_pipeline.py
@@ -21,6 +21,8 @@
big query table at the end of the pipeline.
"""
+# pytype: skip-file
+
# pylint: disable=wrong-import-order, wrong-import-position
from __future__ import absolute_import
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py
index e685c8b..ba79e70 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -106,13 +106,13 @@
]))
table_names = (p | beam.Create([
- ('error', 'my_project.dataset1.error_table_for_today'),
- ('user_log', 'my_project.dataset1.query_table_for_today'),
+ ('error', 'my_project:dataset1.error_table_for_today'),
+ ('user_log', 'my_project:dataset1.query_table_for_today'),
])
table_names_dict = beam.pvalue.AsDict(table_names)
- elements | beam.io.gcp.WriteToBigQuery(
+ elements | beam.io.gcp.bigquery.WriteToBigQuery(
table=lambda row, table_dict: table_dict[row['type']],
table_side_inputs=(table_names_dict,))
@@ -146,7 +146,7 @@
{'type': 'user_log', 'timestamp': '12:34:59', 'query': 'flu symptom'},
]))
- elements | beam.io.gcp.WriteToBigQuery(
+ elements | beam.io.gcp.bigquery.WriteToBigQuery(
table=compute_table_name,
schema=lambda table: (errors_schema
if 'errors' in table
@@ -183,8 +183,8 @@
{'country': 'canada', 'timestamp': '12:34:59', 'query': 'influenza'},
]))
- elements | beam.io.gcp.WriteToBigQuery(
- table='project_name1.dataset_2.query_events_table',
+ elements | beam.io.gcp.bigquery.WriteToBigQuery(
+ table='project_name1:dataset_2.query_events_table',
additional_bq_parameters=additional_bq_parameters)
Much like the schema case, the parameter with `additional_bq_parameters` can
@@ -229,6 +229,8 @@
returned as base64-encoded bytes.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
index 5864a83..fdce60b 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
@@ -26,6 +26,8 @@
NOTHING IN THIS FILE HAS BACKWARDS COMPATIBILITY GUARANTEES.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import datetime
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
index 92e37c6..eb673a6 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
@@ -17,6 +17,8 @@
"""Unit tests for BigQuery file loads utilities."""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_io_read_it_test.py b/sdks/python/apache_beam/io/gcp/bigquery_io_read_it_test.py
index 72e1bb7..386d5d7 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_io_read_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_io_read_it_test.py
@@ -21,6 +21,8 @@
Can be configured to simulate slow reading for a given number of rows.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -43,17 +45,22 @@
"1G": 11110839,
"1T": 11110839000,}
- def run_bigquery_io_read_pipeline(self, input_size):
+ def run_bigquery_io_read_pipeline(self, input_size, beam_bq_source=False):
test_pipeline = TestPipeline(is_integration_test=True)
pipeline_verifiers = [PipelineStateMatcher(),]
extra_opts = {'input_table': self.DEFAULT_DATASET + "." +
self.DEFAULT_TABLE_PREFIX + input_size,
'num_records': self.NUM_RECORDS[input_size],
+ 'beam_bq_source': str(beam_bq_source),
'on_success_matcher': all_of(*pipeline_verifiers)}
bigquery_io_read_pipeline.run(test_pipeline.get_full_options_as_args(
**extra_opts))
@attr('IT')
+ def test_bigquery_read_custom_1M_python(self):
+ self.run_bigquery_io_read_pipeline('1M', True)
+
+ @attr('IT')
def test_bigquery_read_1M_python(self):
self.run_bigquery_io_read_pipeline('1M')
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py b/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py
index 142182a..7e1dc94 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py
@@ -20,6 +20,8 @@
Can be configured to simulate slow reading for a given number of rows.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
@@ -28,6 +30,8 @@
import time
import apache_beam as beam
+from apache_beam.io.gcp.bigquery import _ReadFromBigQuery
+from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
@@ -58,20 +62,27 @@
parser.add_argument('--num_slow', default=0,
help=('Percentage of rows that will be slow. '
'Must be in the range [0, 100)'))
+ parser.add_argument('--beam_bq_source', default=False, type=bool,
+ help=('Whether to use the new _ReadFromBigQuery'
+ ' transform, or the BigQuerySource.'))
known_args, pipeline_args = parser.parse_known_args(argv)
- p = TestPipeline(options=PipelineOptions(pipeline_args))
+ options = PipelineOptions(pipeline_args)
+ with TestPipeline(options=options) as p:
+ if known_args.beam_bq_source:
+ reader = _ReadFromBigQuery(
+ table='%s:%s' % (options.view_as(GoogleCloudOptions).project,
+ known_args.input_table))
+ else:
+ reader = beam.io.Read(beam.io.BigQuerySource(known_args.input_table))
- # pylint: disable=expression-not-assigned
- count = (p | 'read' >> beam.io.Read(beam.io.BigQuerySource(
- known_args.input_table))
- | 'row to string' >> beam.ParDo(RowToStringWithSlowDown(),
- num_slow=known_args.num_slow)
- | 'count' >> beam.combiners.Count.Globally())
+ # pylint: disable=expression-not-assigned
+ count = (p | 'read' >> reader
+ | 'row to string' >> beam.ParDo(RowToStringWithSlowDown(),
+ num_slow=known_args.num_slow)
+ | 'count' >> beam.combiners.Count.Globally())
- assert_that(count, equal_to([known_args.num_records]))
-
- p.run()
+ assert_that(count, equal_to([known_args.num_records]))
if __name__ == '__main__':
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 31ca577..1b37d49 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
@@ -18,6 +18,8 @@
#
"""Unit tests for BigQuery sources and sinks."""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
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 5b48287..d93cd88 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
@@ -49,6 +49,8 @@
--tests apache_beam.io.gcp.bigquery_read_perf_test
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
@@ -115,17 +117,17 @@
# of the part
return {'data': base64.b64encode(record[1])}
- p = TestPipeline()
- # pylint: disable=expression-not-assigned
- (p
- | 'Produce rows' >> Read(SyntheticSource(self.parseTestPipelineOptions()))
- | 'Format' >> Map(format_record)
- | 'Write to BigQuery' >> WriteToBigQuery(
- dataset=self.input_dataset, table=self.input_table,
- schema=SCHEMA,
- create_disposition=BigQueryDisposition.CREATE_IF_NEEDED,
- write_disposition=BigQueryDisposition.WRITE_EMPTY))
- p.run().wait_until_finish()
+ with TestPipeline() as p:
+ # pylint: disable=expression-not-assigned
+ (p
+ | 'Produce rows' >> Read(SyntheticSource(
+ self.parseTestPipelineOptions()))
+ | 'Format' >> Map(format_record)
+ | 'Write to BigQuery' >> WriteToBigQuery(
+ dataset=self.input_dataset, table=self.input_table,
+ schema=SCHEMA,
+ create_disposition=BigQueryDisposition.CREATE_IF_NEEDED,
+ write_disposition=BigQueryDisposition.WRITE_EMPTY))
def test(self):
output = (self.pipeline
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py
index 2125260..7564edb 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for BigQuery sources and sinks."""
+# pytype: skip-file
+
from __future__ import absolute_import
import decimal
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index ba69bb0..b74c967 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -25,6 +25,8 @@
NOTHING IN THIS FILE HAS BACKWARDS COMPATIBILITY GUARANTEES.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import datetime
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
index 7cfea2f..e8eaaa0 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import datetime
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py b/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
index ae56e35..759e319 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
@@ -18,6 +18,8 @@
#
"""Unit tests for BigQuery sources and sinks."""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
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 c8dd304..e6ced26 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
@@ -48,6 +48,8 @@
This setup will result in a table of 1MB size.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py
index ccb10c5..928e6d8 100644
--- a/sdks/python/apache_beam/io/gcp/bigtableio.py
+++ b/sdks/python/apache_beam/io/gcp/bigtableio.py
@@ -35,6 +35,8 @@
instance_id,
table_id))
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import apache_beam as beam
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler.py b/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler.py
index f6c65a5..9b38545 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler.py
@@ -19,6 +19,8 @@
#
# For internal use only; no backwards-compatibility guarantees.
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler_test.py
index e3ccb92..9d7fdfd 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/adaptive_throttler_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio.py b/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio.py
index 9af7674..fa3c869 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio.py
@@ -22,6 +22,8 @@
``apache_beam.io.gcp.datastore.v1new.datastoreio`` will replace it in the
next Beam major release.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio_test.py
index ca9a129..6c77204 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/datastoreio_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py b/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py
index 054df9d..535af10 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import uuid
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py b/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py
index 4ea2898..dbb39e2 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import errno
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/helper_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1/helper_test.py
index e71255a..31e94b3 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/helper_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/helper_test.py
@@ -16,6 +16,8 @@
#
"""Tests for datastore helper."""
+# pytype: skip-file
+
from __future__ import absolute_import
import errno
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter.py b/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter.py
index 32fd8e7..f660324 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter.py
@@ -16,6 +16,8 @@
#
"""Implements a Cloud Datastore query splitter."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py
index 6243c4d..1791d8d 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py
@@ -17,6 +17,8 @@
"""Cloud Datastore query splitter test."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/util.py b/sdks/python/apache_beam/io/gcp/datastore/v1/util.py
index 5595a5d..5d6b597 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/util.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/util.py
@@ -19,6 +19,8 @@
#
# For internal use only; no backwards-compatibility guarantees.
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/util_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1/util_test.py
index 4669106..738c5b4 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1/util_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1/util_test.py
@@ -16,6 +16,8 @@
#
"""Tests for util.py."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py
index efe80c8..ea17f3f 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py
@@ -27,6 +27,8 @@
5. Query the written Entities, verify no results.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py
index 8b577aa..69aea0e 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py
@@ -25,6 +25,8 @@
results in the pipeline.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 a70ea95..633ed52 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
@@ -27,6 +27,8 @@
This module is experimental, no backwards compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
index 4a5f939..9eb0db0 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
@@ -17,6 +17,8 @@
"""Unit tests for datastoreio."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/helper.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/helper.py
index 202e9ce..2dcf7fa 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/helper.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/helper.py
@@ -21,6 +21,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import os
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 6db4ca4..4f4dc0d 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
@@ -20,6 +20,8 @@
For internal use only. No backwards compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter_test.py
index 437282b..6f13f7f 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter_test.py
@@ -17,6 +17,8 @@
"""Cloud Datastore query splitter test."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
index a6e77da..497b76a 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
@@ -21,6 +21,8 @@
This module is experimental, no backwards compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import copy
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py
index 2881de2..0625418 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/types_test.py
@@ -17,6 +17,8 @@
"""Unit tests for types module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import datetime
diff --git a/sdks/python/apache_beam/io/gcp/datastore_write_it_pipeline.py b/sdks/python/apache_beam/io/gcp/datastore_write_it_pipeline.py
index 2d0be8f..9950fa1 100644
--- a/sdks/python/apache_beam/io/gcp/datastore_write_it_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/datastore_write_it_pipeline.py
@@ -27,6 +27,8 @@
5. Query the written Entities, verify no results.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/io/gcp/datastore_write_it_test.py b/sdks/python/apache_beam/io/gcp/datastore_write_it_test.py
index ff611fd..0a32ea1 100644
--- a/sdks/python/apache_beam/io/gcp/datastore_write_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore_write_it_test.py
@@ -25,6 +25,8 @@
results in the pipeline.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/experimental/__init__.py b/sdks/python/apache_beam/io/gcp/experimental/__init__.py
new file mode 100644
index 0000000..f4f43cb
--- /dev/null
+++ b/sdks/python/apache_beam/io/gcp/experimental/__init__.py
@@ -0,0 +1,17 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from __future__ import absolute_import
diff --git a/sdks/python/apache_beam/io/gcp/experimental/spannerio.py b/sdks/python/apache_beam/io/gcp/experimental/spannerio.py
new file mode 100644
index 0000000..21a2f8f
--- /dev/null
+++ b/sdks/python/apache_beam/io/gcp/experimental/spannerio.py
@@ -0,0 +1,583 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Google Cloud Spanner IO
+
+Experimental; no backwards-compatibility guarantees.
+
+This is an experimental module for reading and writing data from Google Cloud
+Spanner. Visit: https://cloud.google.com/spanner for more details.
+
+To read from Cloud Spanner apply ReadFromSpanner transformation. It will
+return a PCollection, where each element represents an individual row returned
+from the read operation. Both Query and Read APIs are supported.
+
+ReadFromSpanner relies on the ReadOperation objects which is exposed by the
+SpannerIO API. ReadOperation holds the immutable data which is responsible to
+execute batch and naive reads on Cloud Spanner. This is done for more
+convenient programming.
+
+ReadFromSpanner reads from Cloud Spanner by providing either an 'sql' param
+in the constructor or 'table' name with 'columns' as list. For example:::
+
+ records = (pipeline
+ | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME,
+ sql='Select * from users'))
+
+ records = (pipeline
+ | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME,
+ table='users', columns=['id', 'name', 'email']))
+
+You can also perform multiple reads by providing a list of ReadOperations
+to the ReadFromSpanner transform constructor. ReadOperation exposes two static
+methods. Use 'query' to perform sql based reads, 'table' to perform read from
+table name. For example:::
+
+ read_operations = [
+ ReadOperation.table(table='customers', columns=['name',
+ 'email']),
+ ReadOperation.table(table='vendors', columns=['name',
+ 'email']),
+ ]
+ all_users = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME,
+ read_operations=read_operations)
+
+ ...OR...
+
+ read_operations = [
+ ReadOperation.query(sql='Select name, email from
+ customers'),
+ ReadOperation.query(
+ sql='Select * from users where id <= @user_id',
+ params={'user_id': 100},
+ params_type={'user_id': param_types.INT64}
+ ),
+ ]
+ # `params_types` are instance of `google.cloud.spanner.param_types`
+ all_users = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME,
+ read_operations=read_operations)
+
+For more information, please review the docs on class ReadOperation.
+
+User can also able to provide the ReadOperation in form of PCollection via
+pipeline. For example:::
+
+ users = (pipeline
+ | beam.Create([ReadOperation...])
+ | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME))
+
+User may also create cloud spanner transaction from the transform called
+`create_transaction` which is available in the SpannerIO API.
+
+The transform is guaranteed to be executed on a consistent snapshot of data,
+utilizing the power of read only transactions. Staleness of data can be
+controlled by providing the `read_timestamp` or `exact_staleness` param values
+in the constructor.
+
+This transform requires root of the pipeline (PBegin) and returns PTransform
+which is passed later to the `ReadFromSpanner` constructor. `ReadFromSpanner`
+pass this transaction PTransform as a singleton side input to the
+`_NaiveSpannerReadDoFn` containing 'session_id' and 'transaction_id'.
+For example:::
+
+ transaction = (pipeline | create_transaction(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ DB_NAME))
+
+ users = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME,
+ sql='Select * from users', transaction=transaction)
+
+ tweets = pipeline | ReadFromSpanner(PROJECT_ID, INSTANCE_ID, DB_NAME,
+ sql='Select * from tweets', transaction=transaction)
+
+For further details of this transform, please review the docs on the
+:meth:`create_transaction` method available in the SpannerIO API.
+
+ReadFromSpanner takes this transform in the constructor and pass this to the
+read pipeline as the singleton side input.
+"""
+from __future__ import absolute_import
+
+import typing
+from collections import namedtuple
+
+from apache_beam import Create
+from apache_beam import DoFn
+from apache_beam import ParDo
+from apache_beam import Reshuffle
+from apache_beam.pvalue import AsSingleton
+from apache_beam.pvalue import PBegin
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import ptransform_fn
+from apache_beam.transforms.display import DisplayDataItem
+from apache_beam.typehints import with_input_types
+from apache_beam.typehints import with_output_types
+from apache_beam.utils.annotations import experimental
+
+try:
+ from google.cloud.spanner import Client
+ from google.cloud.spanner import KeySet
+ from google.cloud.spanner_v1.database import BatchSnapshot
+except ImportError:
+ Client = None
+ KeySet = None
+ BatchSnapshot = None
+
+__all__ = ['create_transaction', 'ReadFromSpanner', 'ReadOperation']
+
+
+class _SPANNER_TRANSACTION(namedtuple("SPANNER_TRANSACTION", ["transaction"])):
+ """
+ Holds the spanner transaction details.
+ """
+
+ __slots__ = ()
+
+
+class ReadOperation(namedtuple("ReadOperation", ["is_sql", "is_table",
+ "read_operation", "kwargs"])):
+ """
+ Encapsulates a spanner read operation.
+ """
+
+ __slots__ = ()
+
+ @classmethod
+ def query(cls, sql, params=None, param_types=None):
+ """
+ A convenient method to construct ReadOperation from sql query.
+
+ Args:
+ sql: SQL query statement
+ params: (optional) values for parameter replacement. Keys must match the
+ names used in sql
+ param_types: (optional) maps explicit types for one or more param values;
+ required if parameters are passed.
+ """
+
+ if params:
+ assert param_types is not None
+
+ return cls(
+ is_sql=True,
+ is_table=False,
+ read_operation="process_query_batch",
+ kwargs={'sql': sql, 'params': params, 'param_types': param_types}
+ )
+
+ @classmethod
+ def table(cls, table, columns, index="", keyset=None):
+ """
+ A convenient method to construct ReadOperation from table.
+
+ Args:
+ table: name of the table from which to fetch data.
+ columns: names of columns to be retrieved.
+ index: (optional) name of index to use, rather than the table's primary
+ key.
+ keyset: (optional) `KeySet` keys / ranges identifying rows to be
+ retrieved.
+ """
+ keyset = keyset or KeySet(all_=True)
+ if not isinstance(keyset, KeySet):
+ raise ValueError("keyset must be an instance of class "
+ "google.cloud.spanner.KeySet")
+ return cls(
+ is_sql=False,
+ is_table=True,
+ read_operation="process_read_batch",
+ kwargs={'table': table, 'columns': columns, 'index': index,
+ 'keyset': keyset}
+ )
+
+
+class _BeamSpannerConfiguration(namedtuple(
+ "_BeamSpannerConfiguration", ["project", "instance", "database",
+ "credentials", "pool",
+ "snapshot_read_timestamp",
+ "snapshot_exact_staleness"])):
+ """
+ A namedtuple holds the immutable data of the connection string to the cloud
+ spanner.
+ """
+
+ @property
+ def snapshot_options(self):
+ snapshot_options = {}
+ if self.snapshot_exact_staleness:
+ snapshot_options['exact_staleness'] = self.snapshot_exact_staleness
+ if self.snapshot_read_timestamp:
+ snapshot_options['read_timestamp'] = self.snapshot_read_timestamp
+ return snapshot_options
+
+@with_input_types(ReadOperation, typing.Dict[typing.Any, typing.Any])
+@with_output_types(typing.List[typing.Any])
+class _NaiveSpannerReadDoFn(DoFn):
+
+ def __init__(self, spanner_configuration):
+ """
+ A naive version of Spanner read which uses the transaction API of the
+ cloud spanner.
+ https://googleapis.dev/python/spanner/latest/transaction-api.html
+ In Naive reads, this transform performs single reads, where as the
+ Batch reads use the spanner partitioning query to create batches.
+
+ Args:
+ spanner_configuration: (_BeamSpannerConfiguration) Connection details to
+ connect with cloud spanner.
+ """
+ self._spanner_configuration = spanner_configuration
+ self._snapshot = None
+ self._session = None
+
+ def _get_session(self):
+ if self._session is None:
+ session = self._session = self._database.session()
+ session.create()
+ return self._session
+
+ def _close_session(self):
+ if self._session is not None:
+ self._session.delete()
+
+ def setup(self):
+ # setting up client to connect with cloud spanner
+ spanner_client = Client(self._spanner_configuration.project)
+ instance = spanner_client.instance(self._spanner_configuration.instance)
+ self._database = instance.database(self._spanner_configuration.database,
+ pool=self._spanner_configuration.pool)
+
+ def process(self, element, spanner_transaction):
+ # `spanner_transaction` should be the instance of the _SPANNER_TRANSACTION
+ # object.
+ if not isinstance(spanner_transaction, _SPANNER_TRANSACTION):
+ raise ValueError("Invalid transaction object: %s. It should be instance "
+ "of SPANNER_TRANSACTION object created by "
+ "spannerio.create_transaction transform."
+ % type(spanner_transaction))
+
+ transaction_info = spanner_transaction.transaction
+
+ # We used batch snapshot to reuse the same transaction passed through the
+ # side input
+ self._snapshot = BatchSnapshot.from_dict(self._database, transaction_info)
+
+ # getting the transaction from the snapshot's session to run read operation.
+ # with self._snapshot.session().transaction() as transaction:
+ with self._get_session().transaction() as transaction:
+ if element.is_sql is True:
+ transaction_read = transaction.execute_sql
+ elif element.is_table is True:
+ transaction_read = transaction.read
+ else:
+ raise ValueError("ReadOperation is improperly configure: %s" % str(
+ element))
+
+ for row in transaction_read(**element.kwargs):
+ yield row
+
+
+@with_input_types(ReadOperation)
+@with_output_types(typing.Dict[typing.Any, typing.Any])
+class _CreateReadPartitions(DoFn):
+ """
+ A DoFn to create partitions. Uses the Partitioning API (PartitionRead /
+ PartitionQuery) request to start a partitioned query operation. Returns a
+ list of batch information needed to perform the actual queries.
+
+ If the element is the instance of :class:`ReadOperation` is to perform sql
+ query, `PartitionQuery` API is used the create partitions and returns mappings
+ of information used perform actual partitioned reads via
+ :meth:`process_query_batch`.
+
+ If the element is the instance of :class:`ReadOperation` is to perform read
+ from table, `PartitionRead` API is used the create partitions and returns
+ mappings of information used perform actual partitioned reads via
+ :meth:`process_read_batch`.
+ """
+
+ def __init__(self, spanner_configuration):
+ self._spanner_configuration = spanner_configuration
+
+ def setup(self):
+ spanner_client = Client(project=self._spanner_configuration.project,
+ credentials=self._spanner_configuration.credentials)
+ instance = spanner_client.instance(self._spanner_configuration.instance)
+ self._database = instance.database(self._spanner_configuration.database,
+ pool=self._spanner_configuration.pool)
+ self._snapshot = self._database.batch_snapshot(**self._spanner_configuration
+ .snapshot_options)
+ self._snapshot_dict = self._snapshot.to_dict()
+
+ def process(self, element):
+ if element.is_sql is True:
+ partitioning_action = self._snapshot.generate_query_batches
+ elif element.is_table is True:
+ partitioning_action = self._snapshot.generate_read_batches
+ else:
+ raise ValueError("ReadOperation is improperly configure: %s" % str(
+ element))
+
+ for p in partitioning_action(**element.kwargs):
+ yield {"is_sql": element.is_sql, "is_table": element.is_table,
+ "read_operation": element.read_operation, "partitions": p,
+ "transaction_info": self._snapshot_dict}
+
+
+@with_input_types(int)
+@with_output_types(typing.Dict[typing.Any, typing.Any])
+class _CreateTransactionFn(DoFn):
+ """
+ A DoFn to create the transaction of cloud spanner.
+ It connects to the database and and returns the transaction_id and session_id
+ by using the batch_snapshot.to_dict() method available in the google cloud
+ spanner sdk.
+
+ https://googleapis.dev/python/spanner/latest/database-api.html?highlight=
+ batch_snapshot#google.cloud.spanner_v1.database.BatchSnapshot.to_dict
+ """
+
+ def __init__(self, project_id, instance_id, database_id, credentials,
+ pool, read_timestamp,
+ exact_staleness):
+ self._project_id = project_id
+ self._instance_id = instance_id
+ self._database_id = database_id
+ self._credentials = credentials
+ self._pool = pool
+
+ self._snapshot_options = {}
+ if read_timestamp:
+ self._snapshot_options['read_timestamp'] = read_timestamp
+ if exact_staleness:
+ self._snapshot_options['exact_staleness'] = exact_staleness
+ self._snapshot = None
+
+ def setup(self):
+ self._spanner_client = Client(project=self._project_id,
+ credentials=self._credentials)
+ self._instance = self._spanner_client.instance(self._instance_id)
+ self._database = self._instance.database(self._database_id, pool=self._pool)
+
+ def process(self, element, *args, **kwargs):
+ self._snapshot = self._database.batch_snapshot(**self._snapshot_options)
+ return [_SPANNER_TRANSACTION(self._snapshot.to_dict())]
+
+
+@ptransform_fn
+def create_transaction(pbegin, project_id, instance_id, database_id,
+ credentials=None, pool=None, read_timestamp=None,
+ exact_staleness=None):
+ """
+ A PTransform method to create a batch transaction.
+
+ Args:
+ pbegin: Root of the pipeline
+ project_id: Cloud spanner project id. Be sure to use the Project ID,
+ not the Project Number.
+ instance_id: Cloud spanner instance id.
+ database_id: Cloud spanner database id.
+ credentials: (optional) The authorization credentials to attach to requests.
+ These credentials identify this application to the service.
+ If none are specified, the client will attempt to ascertain
+ the credentials from the environment.
+ pool: (optional) session pool to be used by database. If not passed,
+ Spanner Cloud SDK uses the BurstyPool by default.
+ `google.cloud.spanner.BurstyPool`. Ref:
+ https://googleapis.dev/python/spanner/latest/database-api.html?#google.
+ cloud.spanner_v1.database.Database
+ read_timestamp: (optional) An instance of the `datetime.datetime` object to
+ execute all reads at the given timestamp.
+ exact_staleness: (optional) An instance of the `datetime.timedelta`
+ object. These timestamp bounds execute reads at a user-specified
+ timestamp.
+ """
+
+ assert isinstance(pbegin, PBegin)
+
+ return (pbegin | Create([1]) | ParDo(_CreateTransactionFn(
+ project_id, instance_id, database_id, credentials,
+ pool, read_timestamp,
+ exact_staleness)))
+
+@with_input_types(typing.Dict[typing.Any, typing.Any])
+@with_output_types(typing.List[typing.Any])
+class _ReadFromPartitionFn(DoFn):
+ """
+ A DoFn to perform reads from the partition.
+ """
+
+ def __init__(self, spanner_configuration):
+ self._spanner_configuration = spanner_configuration
+
+ def setup(self):
+ spanner_client = Client(self._spanner_configuration.project)
+ instance = spanner_client.instance(self._spanner_configuration.instance)
+ self._database = instance.database(self._spanner_configuration.database,
+ pool=self._spanner_configuration.pool)
+ self._snapshot = self._database.batch_snapshot(**self._spanner_configuration
+ .snapshot_options)
+
+ def process(self, element):
+ self._snapshot = BatchSnapshot.from_dict(
+ self._database,
+ element['transaction_info']
+ )
+
+ if element['is_sql'] is True:
+ read_action = self._snapshot.process_query_batch
+ elif element['is_table'] is True:
+ read_action = self._snapshot.process_read_batch
+ else:
+ raise ValueError("ReadOperation is improperly configure: %s" % str(
+ element))
+
+ for row in read_action(element['partitions']):
+ yield row
+
+ def teardown(self):
+ if self._snapshot:
+ self._snapshot.close()
+
+
+@experimental(extra_message="No backwards-compatibility guarantees.")
+class ReadFromSpanner(PTransform):
+ """
+ A PTransform to perform reads from cloud spanner.
+ ReadFromSpanner uses BatchAPI to perform all read operations.
+ """
+
+ def __init__(self, project_id, instance_id, database_id, pool=None,
+ read_timestamp=None, exact_staleness=None, credentials=None,
+ sql=None, params=None, param_types=None, # with_query
+ table=None, columns=None, index="", keyset=None, # with_table
+ read_operations=None, # for read all
+ transaction=None
+ ):
+ """
+ A PTransform that uses Spanner Batch API to perform reads.
+
+ Args:
+ project_id: Cloud spanner project id. Be sure to use the Project ID,
+ not the Project Number.
+ instance_id: Cloud spanner instance id.
+ database_id: Cloud spanner database id.
+ pool: (optional) session pool to be used by database. If not passed,
+ Spanner Cloud SDK uses the BurstyPool by default.
+ `google.cloud.spanner.BurstyPool`. Ref:
+ https://googleapis.dev/python/spanner/latest/database-api.html?#google.
+ cloud.spanner_v1.database.Database
+ read_timestamp: (optional) An instance of the `datetime.datetime` object
+ to execute all reads at the given timestamp. By default, set to `None`.
+ exact_staleness: (optional) An instance of the `datetime.timedelta`
+ object. These timestamp bounds execute reads at a user-specified
+ timestamp. By default, set to `None`.
+ credentials: (optional) The authorization credentials to attach to
+ requests. These credentials identify this application to the service.
+ If none are specified, the client will attempt to ascertain
+ the credentials from the environment. By default, set to `None`.
+ sql: (optional) SQL query statement.
+ params: (optional) Values for parameter replacement. Keys must match the
+ names used in sql. By default, set to `None`.
+ param_types: (optional) maps explicit types for one or more param values;
+ required if params are passed. By default, set to `None`.
+ table: (optional) Name of the table from which to fetch data. By
+ default, set to `None`.
+ columns: (optional) List of names of columns to be retrieved; required if
+ the table is passed. By default, set to `None`.
+ index: (optional) name of index to use, rather than the table's primary
+ key. By default, set to `None`.
+ keyset: (optional) keys / ranges identifying rows to be retrieved. By
+ default, set to `None`.
+ read_operations: (optional) List of the objects of :class:`ReadOperation`
+ to perform read all. By default, set to `None`.
+ transaction: (optional) PTransform of the :meth:`create_transaction` to
+ perform naive read on cloud spanner. By default, set to `None`.
+ """
+ self._configuration = _BeamSpannerConfiguration(
+ project=project_id, instance=instance_id, database=database_id,
+ credentials=credentials, pool=pool,
+ snapshot_read_timestamp=read_timestamp,
+ snapshot_exact_staleness=exact_staleness
+ )
+
+ self._read_operations = read_operations
+ self._transaction = transaction
+
+ if self._read_operations is None:
+ if table is not None:
+ if columns is None:
+ raise ValueError("Columns are required with the table name.")
+ self._read_operations = [ReadOperation.table(
+ table=table, columns=columns, index=index, keyset=keyset)]
+ elif sql is not None:
+ self._read_operations = [ReadOperation.query(
+ sql=sql, params=params, param_types=param_types)]
+
+ def expand(self, pbegin):
+ if self._read_operations is not None and isinstance(pbegin,
+ PBegin):
+ pcoll = pbegin.pipeline | Create(self._read_operations)
+ elif not isinstance(pbegin, PBegin):
+ if self._read_operations is not None:
+ raise ValueError("Read operation in the constructor only works with "
+ "the root of the pipeline.")
+ pcoll = pbegin
+ else:
+ raise ValueError("Spanner required read operation, sql or table "
+ "with columns.")
+
+ if self._transaction is None:
+ # reading as batch read using the spanner partitioning query to create
+ # batches.
+ p = (pcoll
+ | 'Generate Partitions' >> ParDo(_CreateReadPartitions(
+ spanner_configuration=self._configuration))
+ | 'Reshuffle' >> Reshuffle()
+ | 'Read From Partitions' >> ParDo(_ReadFromPartitionFn(
+ spanner_configuration=self._configuration)))
+ else:
+ # reading as naive read, in which we don't make batches and execute the
+ # queries as a single read.
+ p = (pcoll
+ | 'Reshuffle' >> Reshuffle().with_input_types(ReadOperation)
+ | 'Perform Read' >> ParDo(
+ _NaiveSpannerReadDoFn(spanner_configuration=self._configuration),
+ AsSingleton(self._transaction)))
+ return p
+
+ def display_data(self):
+ res = dict()
+ sql = []
+ table = []
+ if self._read_operations is not None:
+ for ro in self._read_operations:
+ if ro.is_sql is True:
+ sql.append(ro.kwargs)
+ elif ro.is_table is True:
+ table.append(ro.kwargs)
+
+ if sql:
+ res['sql'] = DisplayDataItem(str(sql), label='Sql')
+ if table:
+ res['table'] = DisplayDataItem(str(table), label='Table')
+
+ if self._transaction:
+ res['transaction'] = DisplayDataItem(str(self._transaction),
+ label='transaction')
+
+ return res
diff --git a/sdks/python/apache_beam/io/gcp/experimental/spannerio_test.py b/sdks/python/apache_beam/io/gcp/experimental/spannerio_test.py
new file mode 100644
index 0000000..be838f4
--- /dev/null
+++ b/sdks/python/apache_beam/io/gcp/experimental/spannerio_test.py
@@ -0,0 +1,338 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import absolute_import
+
+import datetime
+import logging
+import random
+import string
+import unittest
+
+import mock
+
+import apache_beam as beam
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+
+# Protect against environments where spanner library is not available.
+# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
+try:
+ from google.cloud import spanner
+ from apache_beam.io.gcp.experimental.spannerio import (create_transaction,
+ ReadOperation,
+ ReadFromSpanner) # pylint: disable=unused-import
+ # disable=unused-import
+except ImportError:
+ spanner = None
+# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
+
+
+MAX_DB_NAME_LENGTH = 30
+TEST_PROJECT_ID = 'apache-beam-testing'
+TEST_INSTANCE_ID = 'beam-test'
+TEST_DATABASE_PREFIX = 'spanner-testdb-'
+FAKE_TRANSACTION_INFO = {"session_id": "qwerty", "transaction_id": "qwerty"}
+FAKE_ROWS = [[1, 'Alice'], [2, 'Bob'], [3, 'Carl'], [4, 'Dan'], [5, 'Evan'],
+ [6, 'Floyd']]
+
+
+def _generate_database_name():
+ mask = string.ascii_lowercase + string.digits
+ length = MAX_DB_NAME_LENGTH - 1 - len(TEST_DATABASE_PREFIX)
+ return TEST_DATABASE_PREFIX + ''.join(random.choice(mask) for i in range(
+ length))
+
+
+def _generate_test_data():
+ mask = string.ascii_lowercase + string.digits
+ length = 100
+ return [('users', ['Key', 'Value'], [(x, ''.join(
+ random.choice(mask) for _ in range(length))) for x in range(1, 5)])]
+
+
+@unittest.skipIf(spanner is None, 'GCP dependencies are not installed.')
+@mock.patch('apache_beam.io.gcp.experimental.spannerio.Client')
+@mock.patch('apache_beam.io.gcp.experimental.spannerio.BatchSnapshot')
+class SpannerReadTest(unittest.TestCase):
+
+ def test_read_with_query_batch(self, mock_batch_snapshot_class,
+ mock_client_class):
+ mock_snapshot = mock.MagicMock()
+
+ mock_snapshot.generate_query_batches.return_value = [
+ {'query': {'sql': 'SELECT * FROM users'},
+ 'partition': 'test_partition'} for _ in range(3)]
+ mock_snapshot.process_query_batch.side_effect = [
+ FAKE_ROWS[0:2], FAKE_ROWS[2:4], FAKE_ROWS[4:]]
+
+ ro = [ReadOperation.query("Select * from users")]
+ pipeline = TestPipeline()
+
+ read = (pipeline
+ | 'read' >> ReadFromSpanner(TEST_PROJECT_ID, TEST_INSTANCE_ID,
+ _generate_database_name(),
+ sql="SELECT * FROM users"))
+
+ readall = (pipeline
+ | 'read all' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name(),
+ read_operations=ro))
+
+ readpipeline = (pipeline
+ | 'create reads' >> beam.Create(ro)
+ | 'reads' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name()))
+
+ pipeline.run()
+ assert_that(read, equal_to(FAKE_ROWS), label='checkRead')
+ assert_that(readall, equal_to(FAKE_ROWS), label='checkReadAll')
+ assert_that(readpipeline, equal_to(FAKE_ROWS), label='checkReadPipeline')
+
+ def test_read_with_table_batch(self, mock_batch_snapshot_class,
+ mock_client_class):
+ mock_snapshot = mock.MagicMock()
+ mock_snapshot.generate_read_batches.return_value = [{
+ 'read': {'table': 'users', 'keyset': {'all': True},
+ 'columns': ['Key', 'Value'], 'index': ''},
+ 'partition': 'test_partition'} for _ in range(3)]
+ mock_snapshot.process_read_batch.side_effect = [
+ FAKE_ROWS[0:2], FAKE_ROWS[2:4], FAKE_ROWS[4:]]
+
+ ro = [ReadOperation.table("users", ["Key", "Value"])]
+ pipeline = TestPipeline()
+
+ read = (pipeline
+ | 'read' >> ReadFromSpanner(TEST_PROJECT_ID, TEST_INSTANCE_ID,
+ _generate_database_name(),
+ table="users",
+ columns=["Key", "Value"]))
+
+ readall = (pipeline
+ | 'read all' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name(),
+ read_operations=ro))
+
+ readpipeline = (pipeline
+ | 'create reads' >> beam.Create(ro)
+ | 'reads' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name()))
+
+ pipeline.run()
+ assert_that(read, equal_to(FAKE_ROWS), label='checkRead')
+ assert_that(readall, equal_to(FAKE_ROWS), label='checkReadAll')
+ assert_that(readpipeline, equal_to(FAKE_ROWS), label='checkReadPipeline')
+
+ with self.assertRaises(ValueError):
+ # Test the exception raised when user passes the read operations in the
+ # constructor and also in the pipeline.
+ _ = (pipeline | 'reads error' >> ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ table="users"
+ ))
+ pipeline.run()
+
+ def test_read_with_index(self, mock_batch_snapshot_class,
+ mock_client_class):
+ mock_snapshot = mock.MagicMock()
+ mock_snapshot.generate_read_batches.return_value = [{
+ 'read': {'table': 'users', 'keyset': {'all': True},
+ 'columns': ['Key', 'Value'], 'index': ''},
+ 'partition': 'test_partition'} for _ in range(3)]
+ mock_snapshot.process_read_batch.side_effect = [
+ FAKE_ROWS[0:2], FAKE_ROWS[2:4], FAKE_ROWS[4:]]
+ ro = [ReadOperation.table("users", ["Key", "Value"], index="Key")]
+ pipeline = TestPipeline()
+ read = (pipeline
+ | 'read' >> ReadFromSpanner(TEST_PROJECT_ID, TEST_INSTANCE_ID,
+ _generate_database_name(),
+ table="users",
+ columns=["Key", "Value"]))
+ readall = (pipeline
+ | 'read all' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name(),
+ read_operations=ro))
+ readpipeline = (pipeline
+ | 'create reads' >> beam.Create(ro)
+ | 'reads' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name()))
+ pipeline.run()
+ assert_that(read, equal_to(FAKE_ROWS), label='checkRead')
+ assert_that(readall, equal_to(FAKE_ROWS), label='checkReadAll')
+ assert_that(readpipeline, equal_to(FAKE_ROWS), label='checkReadPipeline')
+ with self.assertRaises(ValueError):
+ # Test the exception raised when user passes the read operations in the
+ # constructor and also in the pipeline.
+ _ = (pipeline | 'reads error' >> ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ table="users"
+ ))
+ pipeline.run()
+
+ def test_read_with_transaction(self, mock_batch_snapshot_class,
+ mock_client_class):
+ mock_client = mock.MagicMock()
+ mock_instance = mock.MagicMock()
+ mock_database = mock.MagicMock()
+ mock_snapshot = mock.MagicMock()
+
+ mock_client_class.return_value = mock_client
+ mock_client.instance.return_value = mock_instance
+ mock_instance.database.return_value = mock_database
+ mock_database.batch_snapshot.return_value = mock_snapshot
+ mock_batch_snapshot_class.return_value = mock_snapshot
+ mock_batch_snapshot_class.from_dict.return_value = mock_snapshot
+ mock_snapshot.to_dict.return_value = FAKE_TRANSACTION_INFO
+
+ mock_session = mock.MagicMock()
+ mock_transaction_ctx = mock.MagicMock()
+ mock_transaction = mock.MagicMock()
+
+ mock_snapshot._get_session.return_value = mock_session
+ mock_session.transaction.return_value = mock_transaction
+ mock_transaction.__enter__.return_value = mock_transaction_ctx
+ mock_transaction_ctx.execute_sql.return_value = FAKE_ROWS
+
+ ro = [ReadOperation.query("Select * from users")]
+ p = TestPipeline()
+
+ transaction = (p | create_transaction(
+ project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ exact_staleness=datetime.timedelta(seconds=10)))
+
+ read_query = (p | 'with query' >> ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ transaction=transaction,
+ sql="Select * from users"
+ ))
+
+ read_table = (p | 'with table' >> ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ transaction=transaction,
+ table="users",
+ columns=["Key", "Value"]
+ ))
+
+ read_indexed_table = (p | 'with index' >> ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ transaction=transaction,
+ table="users",
+ index="Key",
+ columns=["Key", "Value"]
+ ))
+
+ read = (p | 'read all' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name(),
+ transaction=transaction,
+ read_operations=ro))
+
+ read_pipeline = (p
+ | 'create read operations' >> beam.Create(ro)
+ | 'reads' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name(),
+ transaction=transaction))
+
+ p.run()
+
+ assert_that(read_query, equal_to(FAKE_ROWS), label='checkQuery')
+ assert_that(read_table, equal_to(FAKE_ROWS), label='checkTable')
+ assert_that(read_indexed_table, equal_to(FAKE_ROWS),
+ label='checkTableIndex')
+ assert_that(read, equal_to(FAKE_ROWS), label='checkReadAll')
+ assert_that(read_pipeline, equal_to(FAKE_ROWS), label='checkReadPipeline')
+
+ with self.assertRaises(ValueError):
+ # Test the exception raised when user passes the read operations in the
+ # constructor and also in the pipeline.
+ _ = (p
+ | 'create read operations2' >> beam.Create(ro)
+ | 'reads with error' >> ReadFromSpanner(TEST_PROJECT_ID,
+ TEST_INSTANCE_ID,
+ _generate_database_name(),
+ transaction=transaction,
+ read_operations=ro))
+ p.run()
+
+ def test_invalid_transaction(self, mock_batch_snapshot_class,
+ mock_client_class):
+ with self.assertRaises(ValueError):
+ p = TestPipeline()
+ transaction = (p | beam.Create([{"invalid": "transaction"}]))
+ _ = (p | 'with query' >> ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ transaction=transaction,
+ sql="Select * from users"
+ ))
+ p.run()
+
+
+ def test_display_data(self, *args):
+ dd_sql = ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ sql="Select * from users"
+ ).display_data()
+
+ dd_table = ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ table="users",
+ columns=['id', 'name']
+ ).display_data()
+
+ dd_transaction = ReadFromSpanner(
+ project_id=TEST_PROJECT_ID,
+ instance_id=TEST_INSTANCE_ID,
+ database_id=_generate_database_name(),
+ table="users",
+ columns=['id', 'name'],
+ transaction={"transaction_id": "test123", "session_id": "test456"}
+ ).display_data()
+
+ self.assertTrue("sql" in dd_sql)
+ self.assertTrue("table" in dd_table)
+ self.assertTrue("table" in dd_transaction)
+ self.assertTrue("transaction" in dd_transaction)
+
+
+if __name__ == '__main__':
+ logging.getLogger().setLevel(logging.INFO)
+ unittest.main()
diff --git a/sdks/python/apache_beam/io/gcp/gcsfilesystem.py b/sdks/python/apache_beam/io/gcp/gcsfilesystem.py
index 30c052a..d6dc838 100644
--- a/sdks/python/apache_beam/io/gcp/gcsfilesystem.py
+++ b/sdks/python/apache_beam/io/gcp/gcsfilesystem.py
@@ -16,6 +16,8 @@
#
"""GCS file system implementation for accessing files on GCS."""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import zip
diff --git a/sdks/python/apache_beam/io/gcp/gcsfilesystem_test.py b/sdks/python/apache_beam/io/gcp/gcsfilesystem_test.py
index cc1a4ed..31d324e 100644
--- a/sdks/python/apache_beam/io/gcp/gcsfilesystem_test.py
+++ b/sdks/python/apache_beam/io/gcp/gcsfilesystem_test.py
@@ -18,6 +18,8 @@
"""Unit tests for GCS File System."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py
index c1e0314..180ca41 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -20,6 +20,8 @@
https://github.com/GoogleCloudPlatform/appengine-gcs-client.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import errno
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 9b1d1b5..9ba8f61 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py
@@ -36,6 +36,8 @@
-DkmsKeyName=KMS_KEY_NAME
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/gcsio_overrides.py b/sdks/python/apache_beam/io/gcp/gcsio_overrides.py
index 1be587d..72fde36 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio_overrides.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio_overrides.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/gcsio_test.py b/sdks/python/apache_beam/io/gcp/gcsio_test.py
index 5dcfbd8..fa71dd4 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio_test.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
"""Tests for Google Cloud Storage client."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
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 25d3e52..5cae1c1 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
@@ -17,6 +17,7 @@
"""Generated client library for bigquery version v2."""
# NOTE: This file is autogenerated and should not be edited by hand.
+
from __future__ import absolute_import
from apitools.base.py import base_api
diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
index 61c64e5..6929149 100644
--- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
+++ b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py
@@ -20,6 +20,7 @@
A data platform for customers to create, manage, share and query data.
"""
# NOTE: This file is autogenerated and should not be edited by hand.
+
from __future__ import absolute_import
from apitools.base.protorpclite import message_types as _message_types
diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py b/sdks/python/apache_beam/io/gcp/pubsub.py
index b2db2eb..0496707 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub.py
@@ -22,6 +22,8 @@
This API is currently under development and is subject to change.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import re
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
index 2c43786..6910d24 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
@@ -17,6 +17,8 @@
"""
Integration test for Google Cloud Pub/Sub.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py b/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py
index 8a8c8b4..37e2e71 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_it_pipeline.py
@@ -18,6 +18,8 @@
Test pipeline for use by pubsub_integration_test.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py b/sdks/python/apache_beam/io/gcp/pubsub_test.py
index 7377cf2..a4b9e62 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py
@@ -18,6 +18,8 @@
"""Unit tests for PubSub sources and sinks."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -352,12 +354,11 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- pcoll = (p
- | ReadFromPubSub('projects/fakeprj/topics/a_topic',
- None, None, with_attributes=True))
- assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
- p.run()
+ with TestPipeline(options=options) as p:
+ pcoll = (p
+ | ReadFromPubSub('projects/fakeprj/topics/a_topic',
+ None, None, with_attributes=True))
+ assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
mock_pubsub.return_value.acknowledge.assert_has_calls([
mock.call(mock.ANY, [ack_id])])
@@ -376,12 +377,11 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- pcoll = (p
- | ReadStringsFromPubSub('projects/fakeprj/topics/a_topic',
- None, None))
- assert_that(pcoll, equal_to(expected_elements))
- p.run()
+ with TestPipeline(options=options) as p:
+ pcoll = (p
+ | ReadStringsFromPubSub('projects/fakeprj/topics/a_topic',
+ None, None))
+ assert_that(pcoll, equal_to(expected_elements))
mock_pubsub.return_value.acknowledge.assert_has_calls([
mock.call(mock.ANY, [ack_id])])
@@ -398,11 +398,10 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- pcoll = (p
- | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, None))
- assert_that(pcoll, equal_to(expected_elements))
- p.run()
+ with TestPipeline(options=options) as p:
+ pcoll = (p
+ | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, None))
+ assert_that(pcoll, equal_to(expected_elements))
mock_pubsub.return_value.acknowledge.assert_has_calls([
mock.call(mock.ANY, [ack_id])])
@@ -429,13 +428,12 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- pcoll = (p
- | ReadFromPubSub(
- 'projects/fakeprj/topics/a_topic', None, None,
- with_attributes=True, timestamp_attribute='time'))
- assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
- p.run()
+ with TestPipeline(options=options) as p:
+ pcoll = (p
+ | ReadFromPubSub(
+ 'projects/fakeprj/topics/a_topic', None, None,
+ with_attributes=True, timestamp_attribute='time'))
+ assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
mock_pubsub.return_value.acknowledge.assert_has_calls([
mock.call(mock.ANY, [ack_id])])
@@ -462,13 +460,12 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- pcoll = (p
- | ReadFromPubSub(
- 'projects/fakeprj/topics/a_topic', None, None,
- with_attributes=True, timestamp_attribute='time'))
- assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
- p.run()
+ with TestPipeline(options=options) as p:
+ pcoll = (p
+ | ReadFromPubSub(
+ 'projects/fakeprj/topics/a_topic', None, None,
+ with_attributes=True, timestamp_attribute='time'))
+ assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
mock_pubsub.return_value.acknowledge.assert_has_calls([
mock.call(mock.ANY, [ack_id])])
@@ -496,13 +493,12 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- pcoll = (p
- | ReadFromPubSub(
- 'projects/fakeprj/topics/a_topic', None, None,
- with_attributes=True, timestamp_attribute='nonexistent'))
- assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
- p.run()
+ with TestPipeline(options=options) as p:
+ pcoll = (p
+ | ReadFromPubSub(
+ 'projects/fakeprj/topics/a_topic', None, None,
+ with_attributes=True, timestamp_attribute='nonexistent'))
+ assert_that(pcoll, equal_to(expected_elements), reify_windows=True)
mock_pubsub.return_value.acknowledge.assert_has_calls([
mock.call(mock.ANY, [ack_id])])
@@ -539,11 +535,11 @@
# id_label is unsupported in DirectRunner.
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- _ = (p | ReadFromPubSub('projects/fakeprj/topics/a_topic', None, 'a_label'))
with self.assertRaisesRegex(NotImplementedError,
r'id_label is not supported'):
- p.run()
+ with TestPipeline(options=options) as p:
+ _ = (p | ReadFromPubSub(
+ 'projects/fakeprj/topics/a_topic', None, 'a_label'))
@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
@@ -556,12 +552,11 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- _ = (p
- | Create(payloads)
- | WriteToPubSub('projects/fakeprj/topics/a_topic',
- with_attributes=False))
- p.run()
+ with TestPipeline(options=options) as p:
+ _ = (p
+ | Create(payloads)
+ | WriteToPubSub('projects/fakeprj/topics/a_topic',
+ with_attributes=False))
mock_pubsub.return_value.publish.assert_has_calls([
mock.call(mock.ANY, data)])
@@ -571,11 +566,10 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- _ = (p
- | Create(payloads)
- | WriteStringsToPubSub('projects/fakeprj/topics/a_topic'))
- p.run()
+ with TestPipeline(options=options) as p:
+ _ = (p
+ | Create(payloads)
+ | WriteStringsToPubSub('projects/fakeprj/topics/a_topic'))
mock_pubsub.return_value.publish.assert_has_calls([
mock.call(mock.ANY, data)])
@@ -586,12 +580,11 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- _ = (p
- | Create(payloads)
- | WriteToPubSub('projects/fakeprj/topics/a_topic',
- with_attributes=True))
- p.run()
+ with TestPipeline(options=options) as p:
+ _ = (p
+ | Create(payloads)
+ | WriteToPubSub('projects/fakeprj/topics/a_topic',
+ with_attributes=True))
mock_pubsub.return_value.publish.assert_has_calls([
mock.call(mock.ANY, data, **attributes)])
@@ -602,14 +595,13 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- _ = (p
- | Create(payloads)
- | WriteToPubSub('projects/fakeprj/topics/a_topic',
- with_attributes=True))
with self.assertRaisesRegex(AttributeError,
r'str.*has no attribute.*data'):
- p.run()
+ with TestPipeline(options=options) as p:
+ _ = (p
+ | Create(payloads)
+ | WriteToPubSub('projects/fakeprj/topics/a_topic',
+ with_attributes=True))
def test_write_messages_unsupported_features(self, mock_pubsub):
data = b'data'
@@ -618,24 +610,23 @@
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- _ = (p
- | Create(payloads)
- | WriteToPubSub('projects/fakeprj/topics/a_topic',
- id_label='a_label'))
with self.assertRaisesRegex(NotImplementedError,
r'id_label is not supported'):
- p.run()
+ with TestPipeline(options=options) as p:
+ _ = (p
+ | Create(payloads)
+ | WriteToPubSub('projects/fakeprj/topics/a_topic',
+ id_label='a_label'))
+
options = PipelineOptions([])
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- _ = (p
- | Create(payloads)
- | WriteToPubSub('projects/fakeprj/topics/a_topic',
- timestamp_attribute='timestamp'))
with self.assertRaisesRegex(NotImplementedError,
r'timestamp_attribute is not supported'):
- p.run()
+ with TestPipeline(options=options) as p:
+ _ = (p
+ | Create(payloads)
+ | WriteToPubSub('projects/fakeprj/topics/a_topic',
+ timestamp_attribute='timestamp'))
if __name__ == '__main__':
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 c99d60c..57e9d83 100644
--- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
+++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
@@ -17,6 +17,8 @@
"""Bigquery data verifier for end-to-end test."""
+# pytype: skip-file
+
from __future__ import absolute_import
import concurrent
diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
index f6f4394..cfd562f 100644
--- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
+++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
@@ -17,6 +17,8 @@
"""Unit test for Bigquery verifier"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py
index 53ebae4..0b9622c 100644
--- a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py
+++ b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher.py
@@ -17,6 +17,8 @@
"""PubSub verifier used for end-to-end test."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py
index 643fdc4..f8adc8b 100644
--- a/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py
+++ b/sdks/python/apache_beam/io/gcp/tests/pubsub_matcher_test.py
@@ -17,6 +17,8 @@
"""Unit test for PubSub verifier."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/tests/utils.py b/sdks/python/apache_beam/io/gcp/tests/utils.py
index dbf8ac9..0664a00 100644
--- a/sdks/python/apache_beam/io/gcp/tests/utils.py
+++ b/sdks/python/apache_beam/io/gcp/tests/utils.py
@@ -18,6 +18,8 @@
"""Utility methods for testing on GCP."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/gcp/tests/utils_test.py b/sdks/python/apache_beam/io/gcp/tests/utils_test.py
index db547d9..de69f96 100644
--- a/sdks/python/apache_beam/io/gcp/tests/utils_test.py
+++ b/sdks/python/apache_beam/io/gcp/tests/utils_test.py
@@ -17,6 +17,8 @@
"""Unittest for GCP testing utils."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/hadoopfilesystem.py b/sdks/python/apache_beam/io/hadoopfilesystem.py
index ba80cb06..dde807b 100644
--- a/sdks/python/apache_beam/io/hadoopfilesystem.py
+++ b/sdks/python/apache_beam/io/hadoopfilesystem.py
@@ -18,6 +18,8 @@
""":class:`~apache_beam.io.filesystem.FileSystem` implementation for accessing
Hadoop Distributed File System files."""
+# pytype: skip-file
+
from __future__ import absolute_import
import io
diff --git a/sdks/python/apache_beam/io/hadoopfilesystem_test.py b/sdks/python/apache_beam/io/hadoopfilesystem_test.py
index 98fafcf..5c1414e 100644
--- a/sdks/python/apache_beam/io/hadoopfilesystem_test.py
+++ b/sdks/python/apache_beam/io/hadoopfilesystem_test.py
@@ -17,6 +17,8 @@
"""Unit tests for :class:`HadoopFileSystem`."""
+# pytype: skip-file
+
from __future__ import absolute_import
import io
diff --git a/sdks/python/apache_beam/io/hdfs_integration_test/Dockerfile b/sdks/python/apache_beam/io/hdfs_integration_test/Dockerfile
index 788b8d2..1c78181 100644
--- a/sdks/python/apache_beam/io/hdfs_integration_test/Dockerfile
+++ b/sdks/python/apache_beam/io/hdfs_integration_test/Dockerfile
@@ -24,22 +24,13 @@
WORKDIR /app
ENV HDFSCLI_CONFIG /app/sdks/python/apache_beam/io/hdfs_integration_test/hdfscli.cfg
-RUN pip install --no-cache-dir holdup gsutil
-RUN gsutil cp gs://dataflow-samples/shakespeare/kinglear.txt .
-# Install Beam and dependencies.
-ADD sdks/python /app/sdks/python
-ADD model /app/model
-RUN cd sdks/python && \
- python setup.py sdist && \
- pip install --no-cache-dir $(ls dist/apache-beam-*.tar.gz | tail -n1)[gcp]
+# Add Beam SDK sources.
+COPY sdks/python /app/sdks/python
+COPY model /app/model
+
+# This step should look like setupVirtualenv minus virtualenv creation.
+RUN pip install --no-cache-dir tox==3.11.1 -r sdks/python/build-requirements.txt
# Run wordcount, and write results to HDFS.
-CMD holdup -t 45 http://namenode:50070 http://datanode:50075 && \
- echo "Waiting for safe mode to end." && \
- sleep 45 && \
- hdfscli -v -v -v upload -f kinglear.txt / && \
- python -m apache_beam.examples.wordcount \
- --input hdfs://kinglear* \
- --output hdfs://py-wordcount-integration \
- --hdfs_host namenode --hdfs_port 50070 --hdfs_user root
+CMD cd sdks/python && tox -e hdfs_integration_test
diff --git a/sdks/python/apache_beam/io/iobase.py b/sdks/python/apache_beam/io/iobase.py
index 4fbc1f0..4779066a 100644
--- a/sdks/python/apache_beam/io/iobase.py
+++ b/sdks/python/apache_beam/io/iobase.py
@@ -29,6 +29,8 @@
the sink.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -897,10 +899,13 @@
def _infer_output_coder(self, input_type=None, input_coder=None):
# type: (...) -> Optional[coders.Coder]
+ from apache_beam.runners.dataflow.native_io import iobase as dataflow_io
if isinstance(self.source, BoundedSource):
return self.source.default_output_coder()
- else:
+ elif isinstance(self.source, dataflow_io.NativeSource):
return self.source.coder
+ else:
+ return None
def display_data(self):
return {'source': DisplayDataItem(self.source.__class__,
@@ -1244,6 +1249,7 @@
"""
def __init__(self, restriction_tracker):
+ # type: (RestrictionTracker) -> None
if not isinstance(restriction_tracker, RestrictionTracker):
raise ValueError(
'Initialize ThreadsafeRestrictionTracker requires'
@@ -1377,6 +1383,7 @@
@property
def completed_work(self):
+ # type: () -> float
if self._completed:
return self._completed
elif self._remaining and self._fraction:
@@ -1384,6 +1391,7 @@
@property
def remaining_work(self):
+ # type: () -> float
if self._remaining:
return self._remaining
elif self._completed:
@@ -1391,10 +1399,12 @@
@property
def total_work(self):
+ # type: () -> float
return self.completed_work + self.remaining_work
@property
def fraction_completed(self):
+ # type: () -> float
if self._fraction is not None:
return self._fraction
else:
@@ -1402,6 +1412,7 @@
@property
def fraction_remaining(self):
+ # type: () -> float
if self._fraction is not None:
return 1 - self._fraction
else:
diff --git a/sdks/python/apache_beam/io/iobase_test.py b/sdks/python/apache_beam/io/iobase_test.py
index a574d20..886ade2 100644
--- a/sdks/python/apache_beam/io/iobase_test.py
+++ b/sdks/python/apache_beam/io/iobase_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the SDFRestrictionProvider module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import time
diff --git a/sdks/python/apache_beam/io/localfilesystem.py b/sdks/python/apache_beam/io/localfilesystem.py
index 132e93c..99e894e 100644
--- a/sdks/python/apache_beam/io/localfilesystem.py
+++ b/sdks/python/apache_beam/io/localfilesystem.py
@@ -16,6 +16,8 @@
#
"""Local File system implementation for accessing files on disk."""
+# pytype: skip-file
+
from __future__ import absolute_import
import os
diff --git a/sdks/python/apache_beam/io/localfilesystem_test.py b/sdks/python/apache_beam/io/localfilesystem_test.py
index 20fa593..c22c657 100644
--- a/sdks/python/apache_beam/io/localfilesystem_test.py
+++ b/sdks/python/apache_beam/io/localfilesystem_test.py
@@ -18,6 +18,8 @@
"""Unit tests for LocalFileSystem."""
+# pytype: skip-file
+
from __future__ import absolute_import
import filecmp
diff --git a/sdks/python/apache_beam/io/mongodbio.py b/sdks/python/apache_beam/io/mongodbio.py
index fba5b43..8ec8df4 100644
--- a/sdks/python/apache_beam/io/mongodbio.py
+++ b/sdks/python/apache_beam/io/mongodbio.py
@@ -51,6 +51,8 @@
No backward compatibility guarantees. Everything in this module is experimental.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/mongodbio_it_test.py b/sdks/python/apache_beam/io/mongodbio_it_test.py
index 6868135..c8fe2e0 100644
--- a/sdks/python/apache_beam/io/mongodbio_it_test.py
+++ b/sdks/python/apache_beam/io/mongodbio_it_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/io/mongodbio_test.py b/sdks/python/apache_beam/io/mongodbio_test.py
index 72ae0ad..0c9c56a 100644
--- a/sdks/python/apache_beam/io/mongodbio_test.py
+++ b/sdks/python/apache_beam/io/mongodbio_test.py
@@ -14,6 +14,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/parquetio.py b/sdks/python/apache_beam/io/parquetio.py
index a4e894cd..062b46f 100644
--- a/sdks/python/apache_beam/io/parquetio.py
+++ b/sdks/python/apache_beam/io/parquetio.py
@@ -27,6 +27,8 @@
that can be used to write a given ``PCollection`` of Python objects to a
Parquet file.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from functools import partial
diff --git a/sdks/python/apache_beam/io/parquetio_it_test.py b/sdks/python/apache_beam/io/parquetio_it_test.py
index 0b58f47..ecc1957 100644
--- a/sdks/python/apache_beam/io/parquetio_it_test.py
+++ b/sdks/python/apache_beam/io/parquetio_it_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -66,12 +68,10 @@
file_prefix = "parquet_it_test"
init_size = 10
data_size = 20000
- p = TestPipeline(is_integration_test=True)
- pcol = self._generate_data(
- p, file_prefix, init_size, data_size)
- self._verify_data(pcol, init_size, data_size)
- result = p.run()
- result.wait_until_finish()
+ with TestPipeline(is_integration_test=True) as p:
+ pcol = self._generate_data(
+ p, file_prefix, init_size, data_size)
+ self._verify_data(pcol, init_size, data_size)
@staticmethod
def _sum_verifier(init_size, data_size, x):
diff --git a/sdks/python/apache_beam/io/parquetio_test.py b/sdks/python/apache_beam/io/parquetio_test.py
index 719bf55..195d52d 100644
--- a/sdks/python/apache_beam/io/parquetio_test.py
+++ b/sdks/python/apache_beam/io/parquetio_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import json
diff --git a/sdks/python/apache_beam/io/range_trackers.py b/sdks/python/apache_beam/io/range_trackers.py
index d4845fb..4deb5f4 100644
--- a/sdks/python/apache_beam/io/range_trackers.py
+++ b/sdks/python/apache_beam/io/range_trackers.py
@@ -17,6 +17,8 @@
"""iobase.RangeTracker implementations provided with Apache Beam.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/range_trackers_test.py b/sdks/python/apache_beam/io/range_trackers_test.py
index e80401f..782807d 100644
--- a/sdks/python/apache_beam/io/range_trackers_test.py
+++ b/sdks/python/apache_beam/io/range_trackers_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the range_trackers module."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/restriction_trackers.py b/sdks/python/apache_beam/io/restriction_trackers.py
index d601a44..36b2214 100644
--- a/sdks/python/apache_beam/io/restriction_trackers.py
+++ b/sdks/python/apache_beam/io/restriction_trackers.py
@@ -16,6 +16,8 @@
#
"""`iobase.RestrictionTracker` implementations provided with Apache Beam."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/io/restriction_trackers_test.py b/sdks/python/apache_beam/io/restriction_trackers_test.py
index 4a57d98..8c2c7f7 100644
--- a/sdks/python/apache_beam/io/restriction_trackers_test.py
+++ b/sdks/python/apache_beam/io/restriction_trackers_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the range_trackers module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -139,13 +141,6 @@
self.assertFalse(tracker.try_claim(220))
tracker.check_done()
- def test_check_done_after_try_claim_past_end_of_range(self):
- tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
- self.assertTrue(tracker.try_claim(150))
- self.assertTrue(tracker.try_claim(175))
- self.assertFalse(tracker.try_claim(200))
- tracker.check_done()
-
def test_check_done_after_try_claim_right_before_end_of_range(self):
tracker = OffsetRestrictionTracker(OffsetRange(100, 200))
self.assertTrue(tracker.try_claim(150))
diff --git a/sdks/python/apache_beam/io/source_test_utils.py b/sdks/python/apache_beam/io/source_test_utils.py
index 7291786..3afc650 100644
--- a/sdks/python/apache_beam/io/source_test_utils.py
+++ b/sdks/python/apache_beam/io/source_test_utils.py
@@ -43,6 +43,8 @@
* apache_beam.io.source_test_utils_test.py
* apache_beam.io.avroio_test.py
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
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 a8c3d82..3ee00be 100644
--- a/sdks/python/apache_beam/io/source_test_utils_test.py
+++ b/sdks/python/apache_beam/io/source_test_utils_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/sources_test.py b/sdks/python/apache_beam/io/sources_test.py
index 8908681..fb34a98 100644
--- a/sdks/python/apache_beam/io/sources_test.py
+++ b/sdks/python/apache_beam/io/sources_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the sources framework."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -123,11 +125,10 @@
def test_run_direct(self):
file_name = self._create_temp_file(b'aaaa\nbbbb\ncccc\ndddd')
- pipeline = TestPipeline()
- pcoll = pipeline | beam.io.Read(LineSource(file_name))
- assert_that(pcoll, equal_to([b'aaaa', b'bbbb', b'cccc', b'dddd']))
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | beam.io.Read(LineSource(file_name))
+ assert_that(pcoll, equal_to([b'aaaa', b'bbbb', b'cccc', b'dddd']))
- pipeline.run()
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/io/textio.py b/sdks/python/apache_beam/io/textio.py
index 612164d..0514992 100644
--- a/sdks/python/apache_beam/io/textio.py
+++ b/sdks/python/apache_beam/io/textio.py
@@ -18,6 +18,8 @@
"""A source and a sink for reading from and writing to text files."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/io/textio_test.py b/sdks/python/apache_beam/io/textio_test.py
index ad336c5..32765f0 100644
--- a/sdks/python/apache_beam/io/textio_test.py
+++ b/sdks/python/apache_beam/io/textio_test.py
@@ -16,6 +16,8 @@
#
"""Tests for textio module."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -422,28 +424,25 @@
def test_read_from_text_single_file(self):
file_name, expected_data = write_data(5)
assert len(expected_data) == 5
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(file_name)
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(file_name)
+ assert_that(pcoll, equal_to(expected_data))
def test_read_from_text_with_file_name_single_file(self):
file_name, data = write_data(5)
expected_data = [(file_name, el) for el in data]
assert len(expected_data) == 5
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(file_name)
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(file_name)
+ assert_that(pcoll, equal_to(expected_data))
def test_read_all_single_file(self):
file_name, expected_data = write_data(5)
assert len(expected_data) == 5
- pipeline = TestPipeline()
- pcoll = pipeline | 'Create' >> Create(
- [file_name]) |'ReadAll' >> ReadAllFromText()
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Create' >> Create(
+ [file_name]) |'ReadAll' >> ReadAllFromText()
+ assert_that(pcoll, equal_to(expected_data))
def test_read_all_many_single_files(self):
file_name1, expected_data1 = write_data(5)
@@ -456,11 +455,10 @@
expected_data.extend(expected_data1)
expected_data.extend(expected_data2)
expected_data.extend(expected_data3)
- pipeline = TestPipeline()
- pcoll = pipeline | 'Create' >> Create(
- [file_name1, file_name2, file_name3]) |'ReadAll' >> ReadAllFromText()
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Create' >> Create(
+ [file_name1, file_name2, file_name3]) |'ReadAll' >> ReadAllFromText()
+ assert_that(pcoll, equal_to(expected_data))
def test_read_all_unavailable_files_ignored(self):
file_name1, expected_data1 = write_data(5)
@@ -474,13 +472,12 @@
expected_data.extend(expected_data1)
expected_data.extend(expected_data2)
expected_data.extend(expected_data3)
- pipeline = TestPipeline()
- pcoll = (pipeline
- | 'Create' >> Create(
- [file_name1, file_name2, file_name3, file_name4])
- |'ReadAll' >> ReadAllFromText())
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = (pipeline
+ | 'Create' >> Create(
+ [file_name1, file_name2, file_name3, file_name4])
+ |'ReadAll' >> ReadAllFromText())
+ assert_that(pcoll, equal_to(expected_data))
def test_read_from_text_single_file_with_coder(self):
class DummyCoder(coders.Coder):
@@ -492,37 +489,33 @@
file_name, expected_data = write_data(5)
assert len(expected_data) == 5
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(file_name, coder=DummyCoder())
- assert_that(pcoll, equal_to([record * 2 for record in expected_data]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(file_name, coder=DummyCoder())
+ assert_that(pcoll, equal_to([record * 2 for record in expected_data]))
def test_read_from_text_file_pattern(self):
pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4])
assert len(expected_data) == 40
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(pattern)
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(pattern)
+ assert_that(pcoll, equal_to(expected_data))
def test_read_from_text_with_file_name_file_pattern(self):
pattern, expected_data = write_pattern(
lines_per_file=[5, 5], return_filenames=True)
assert len(expected_data) == 10
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(pattern)
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(pattern)
+ assert_that(pcoll, equal_to(expected_data))
def test_read_all_file_pattern(self):
pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4])
assert len(expected_data) == 40
- pipeline = TestPipeline()
- pcoll = (pipeline
- | 'Create' >> Create([pattern])
- |'ReadAll' >> ReadAllFromText())
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = (pipeline
+ | 'Create' >> Create([pattern])
+ |'ReadAll' >> ReadAllFromText())
+ assert_that(pcoll, equal_to(expected_data))
def test_read_all_many_file_patterns(self):
pattern1, expected_data1 = write_pattern([5, 3, 12, 8, 8, 4])
@@ -535,11 +528,10 @@
expected_data.extend(expected_data1)
expected_data.extend(expected_data2)
expected_data.extend(expected_data3)
- pipeline = TestPipeline()
- pcoll = pipeline | 'Create' >> Create(
- [pattern1, pattern2, pattern3]) |'ReadAll' >> ReadAllFromText()
- assert_that(pcoll, equal_to(expected_data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Create' >> Create(
+ [pattern1, pattern2, pattern3]) |'ReadAll' >> ReadAllFromText()
+ assert_that(pcoll, equal_to(expected_data))
def test_read_auto_bzip2(self):
_, lines = write_data(15)
@@ -548,10 +540,9 @@
with bz2.BZ2File(file_name, 'wb') as f:
f.write('\n'.join(lines).encode('utf-8'))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(file_name)
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(file_name)
+ assert_that(pcoll, equal_to(lines))
def test_read_auto_deflate(self):
_, lines = write_data(15)
@@ -560,10 +551,9 @@
with open(file_name, 'wb') as f:
f.write(zlib.compress('\n'.join(lines).encode('utf-8')))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(file_name)
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(file_name)
+ assert_that(pcoll, equal_to(lines))
def test_read_auto_gzip(self):
_, lines = write_data(15)
@@ -573,10 +563,9 @@
with gzip.GzipFile(file_name, 'wb') as f:
f.write('\n'.join(lines).encode('utf-8'))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(file_name)
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(file_name)
+ assert_that(pcoll, equal_to(lines))
def test_read_bzip2(self):
_, lines = write_data(15)
@@ -585,12 +574,11 @@
with bz2.BZ2File(file_name, 'wb') as f:
f.write('\n'.join(lines).encode('utf-8'))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- compression_type=CompressionTypes.BZIP2)
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ compression_type=CompressionTypes.BZIP2)
+ assert_that(pcoll, equal_to(lines))
def test_read_corrupted_bzip2_fails(self):
_, lines = write_data(15)
@@ -602,13 +590,12 @@
with open(file_name, 'wb') as f:
f.write(b'corrupt')
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- compression_type=CompressionTypes.BZIP2)
- assert_that(pcoll, equal_to(lines))
with self.assertRaises(Exception):
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ compression_type=CompressionTypes.BZIP2)
+ assert_that(pcoll, equal_to(lines))
def test_read_bzip2_concat(self):
with TempDir() as tempdir:
@@ -643,14 +630,13 @@
final_bzip2_file, 'ab') as dst:
dst.writelines(src.readlines())
- pipeline = TestPipeline()
- lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(
- final_bzip2_file,
- compression_type=beam.io.filesystem.CompressionTypes.BZIP2)
+ with TestPipeline() as pipeline:
+ lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(
+ final_bzip2_file,
+ compression_type=beam.io.filesystem.CompressionTypes.BZIP2)
- expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']
- assert_that(lines, equal_to(expected))
- pipeline.run()
+ expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']
+ assert_that(lines, equal_to(expected))
def test_read_deflate(self):
_, lines = write_data(15)
@@ -659,13 +645,12 @@
with open(file_name, 'wb') as f:
f.write(zlib.compress('\n'.join(lines).encode('utf-8')))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- 0, CompressionTypes.DEFLATE,
- True, coders.StrUtf8Coder())
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ 0, CompressionTypes.DEFLATE,
+ True, coders.StrUtf8Coder())
+ assert_that(pcoll, equal_to(lines))
def test_read_corrupted_deflate_fails(self):
_, lines = write_data(15)
@@ -677,15 +662,13 @@
with open(file_name, 'wb') as f:
f.write(b'corrupt')
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- 0, CompressionTypes.DEFLATE,
- True, coders.StrUtf8Coder())
- assert_that(pcoll, equal_to(lines))
-
with self.assertRaises(Exception):
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ 0, CompressionTypes.DEFLATE,
+ True, coders.StrUtf8Coder())
+ assert_that(pcoll, equal_to(lines))
def test_read_deflate_concat(self):
with TempDir() as tempdir:
@@ -720,13 +703,13 @@
open(final_deflate_file, 'ab') as dst:
dst.writelines(src.readlines())
- pipeline = TestPipeline()
- lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(
- final_deflate_file,
- compression_type=beam.io.filesystem.CompressionTypes.DEFLATE)
+ with TestPipeline() as pipeline:
+ lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(
+ final_deflate_file,
+ compression_type=beam.io.filesystem.CompressionTypes.DEFLATE)
- expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']
- assert_that(lines, equal_to(expected))
+ expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']
+ assert_that(lines, equal_to(expected))
def test_read_gzip(self):
_, lines = write_data(15)
@@ -735,13 +718,12 @@
with gzip.GzipFile(file_name, 'wb') as f:
f.write('\n'.join(lines).encode('utf-8'))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- 0, CompressionTypes.GZIP,
- True, coders.StrUtf8Coder())
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ 0, CompressionTypes.GZIP,
+ True, coders.StrUtf8Coder())
+ assert_that(pcoll, equal_to(lines))
def test_read_corrupted_gzip_fails(self):
_, lines = write_data(15)
@@ -753,15 +735,13 @@
with open(file_name, 'wb') as f:
f.write(b'corrupt')
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- 0, CompressionTypes.GZIP,
- True, coders.StrUtf8Coder())
- assert_that(pcoll, equal_to(lines))
-
with self.assertRaises(Exception):
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ 0, CompressionTypes.GZIP,
+ True, coders.StrUtf8Coder())
+ assert_that(pcoll, equal_to(lines))
def test_read_gzip_concat(self):
with TempDir() as tempdir:
@@ -796,13 +776,13 @@
open(final_gzip_file, 'ab') as dst:
dst.writelines(src.readlines())
- pipeline = TestPipeline()
- lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(
- final_gzip_file,
- compression_type=beam.io.filesystem.CompressionTypes.GZIP)
+ with TestPipeline() as pipeline:
+ lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(
+ final_gzip_file,
+ compression_type=beam.io.filesystem.CompressionTypes.GZIP)
- expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']
- assert_that(lines, equal_to(expected))
+ expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']
+ assert_that(lines, equal_to(expected))
def test_read_all_gzip(self):
_, lines = write_data(100)
@@ -810,13 +790,12 @@
file_name = tempdir.create_temp_file()
with gzip.GzipFile(file_name, 'wb') as f:
f.write('\n'.join(lines).encode('utf-8'))
- pipeline = TestPipeline()
- pcoll = (pipeline
- | Create([file_name])
- | 'ReadAll' >> ReadAllFromText(
- compression_type=CompressionTypes.GZIP))
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = (pipeline
+ | Create([file_name])
+ | 'ReadAll' >> ReadAllFromText(
+ compression_type=CompressionTypes.GZIP))
+ assert_that(pcoll, equal_to(lines))
def test_read_gzip_large(self):
_, lines = write_data(10000)
@@ -826,13 +805,12 @@
with gzip.GzipFile(file_name, 'wb') as f:
f.write('\n'.join(lines).encode('utf-8'))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- 0, CompressionTypes.GZIP,
- True, coders.StrUtf8Coder())
- assert_that(pcoll, equal_to(lines))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ 0, CompressionTypes.GZIP,
+ True, coders.StrUtf8Coder())
+ assert_that(pcoll, equal_to(lines))
def test_read_gzip_large_after_splitting(self):
_, lines = write_data(10000)
@@ -859,13 +837,12 @@
def test_read_gzip_empty_file(self):
with TempDir() as tempdir:
file_name = tempdir.create_temp_file()
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name,
- 0, CompressionTypes.GZIP,
- True, coders.StrUtf8Coder())
- assert_that(pcoll, equal_to([]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name,
+ 0, CompressionTypes.GZIP,
+ True, coders.StrUtf8Coder())
+ assert_that(pcoll, equal_to([]))
def _remove_lines(self, lines, sublist_lengths, num_to_remove):
"""Utility function to remove num_to_remove lines from each sublist.
@@ -948,12 +925,11 @@
with gzip.GzipFile(file_name, 'wb') as f:
f.write('\n'.join(lines).encode('utf-8'))
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromText(
- file_name, 0, CompressionTypes.GZIP,
- True, coders.StrUtf8Coder(), skip_header_lines=2)
- assert_that(pcoll, equal_to(lines[2:]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromText(
+ file_name, 0, CompressionTypes.GZIP,
+ True, coders.StrUtf8Coder(), skip_header_lines=2)
+ assert_that(pcoll, equal_to(lines[2:]))
def test_read_after_splitting_skip_header(self):
file_name, expected_data = write_data(100)
@@ -1103,10 +1079,9 @@
self.assertEqual(f.read().splitlines(), header.splitlines())
def test_write_dataflow(self):
- pipeline = TestPipeline()
- pcoll = pipeline | beam.core.Create(self.lines)
- pcoll | 'Write' >> WriteToText(self.path) # pylint: disable=expression-not-assigned
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | beam.core.Create(self.lines)
+ pcoll | 'Write' >> WriteToText(self.path) # pylint: disable=expression-not-assigned
read_result = []
for file_name in glob.glob(self.path + '*'):
@@ -1116,10 +1091,9 @@
self.assertEqual(sorted(read_result), sorted(self.lines))
def test_write_dataflow_auto_compression(self):
- pipeline = TestPipeline()
- pcoll = pipeline | beam.core.Create(self.lines)
- pcoll | 'Write' >> WriteToText(self.path, file_name_suffix='.gz') # pylint: disable=expression-not-assigned
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | beam.core.Create(self.lines)
+ pcoll | 'Write' >> WriteToText(self.path, file_name_suffix='.gz') # pylint: disable=expression-not-assigned
read_result = []
for file_name in glob.glob(self.path + '*'):
@@ -1129,13 +1103,12 @@
self.assertEqual(sorted(read_result), sorted(self.lines))
def test_write_dataflow_auto_compression_unsharded(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Create' >> beam.core.Create(self.lines)
- pcoll | 'Write' >> WriteToText( # pylint: disable=expression-not-assigned
- self.path + '.gz',
- shard_name_template='')
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Create' >> beam.core.Create(self.lines)
+ pcoll | 'Write' >> WriteToText( # pylint: disable=expression-not-assigned
+ self.path + '.gz',
+ shard_name_template='')
- pipeline.run()
read_result = []
for file_name in glob.glob(self.path + '*'):
@@ -1145,14 +1118,13 @@
self.assertEqual(sorted(read_result), sorted(self.lines))
def test_write_dataflow_header(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Create' >> beam.core.Create(self.lines)
- header_text = 'foo'
- pcoll | 'Write' >> WriteToText( # pylint: disable=expression-not-assigned
- self.path + '.gz',
- shard_name_template='',
- header=header_text)
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Create' >> beam.core.Create(self.lines)
+ header_text = 'foo'
+ pcoll | 'Write' >> WriteToText( # pylint: disable=expression-not-assigned
+ self.path + '.gz',
+ shard_name_template='',
+ header=header_text)
read_result = []
for file_name in glob.glob(self.path + '*'):
diff --git a/sdks/python/apache_beam/io/tfrecordio.py b/sdks/python/apache_beam/io/tfrecordio.py
index 4a62b91..734c598 100644
--- a/sdks/python/apache_beam/io/tfrecordio.py
+++ b/sdks/python/apache_beam/io/tfrecordio.py
@@ -16,6 +16,8 @@
#
"""TFRecord sources and sinks."""
+# pytype: skip-file
+
from __future__ import absolute_import
import codecs
diff --git a/sdks/python/apache_beam/io/tfrecordio_test.py b/sdks/python/apache_beam/io/tfrecordio_test.py
index dfb154a..90f484a 100644
--- a/sdks/python/apache_beam/io/tfrecordio_test.py
+++ b/sdks/python/apache_beam/io/tfrecordio_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import binascii
@@ -288,7 +290,7 @@
validate=True))
assert_that(result, equal_to([b'foo', b'bar']))
- def test_process_gzip(self):
+ def test_process_gzip_with_coder(self):
with TempDir() as temp_dir:
path = temp_dir.create_temp_file('result')
_write_file_gzip(path, FOO_BAR_RECORD_BASE64)
@@ -301,6 +303,17 @@
validate=True))
assert_that(result, equal_to([b'foo', b'bar']))
+ def test_process_gzip_without_coder(self):
+ with TempDir() as temp_dir:
+ path = temp_dir.create_temp_file('result')
+ _write_file_gzip(path, FOO_BAR_RECORD_BASE64)
+ with TestPipeline() as p:
+ result = (p
+ | ReadFromTFRecord(
+ path,
+ compression_type=CompressionTypes.GZIP))
+ assert_that(result, equal_to([b'foo', b'bar']))
+
def test_process_auto(self):
with TempDir() as temp_dir:
path = temp_dir.create_temp_file('result.gz')
@@ -314,16 +327,6 @@
validate=True))
assert_that(result, equal_to([b'foo', b'bar']))
- def test_process_gzip(self):
- with TempDir() as temp_dir:
- path = temp_dir.create_temp_file('result')
- _write_file_gzip(path, FOO_BAR_RECORD_BASE64)
- with TestPipeline() as p:
- result = (p
- | ReadFromTFRecord(
- path, compression_type=CompressionTypes.GZIP))
- assert_that(result, equal_to([b'foo', b'bar']))
-
def test_process_gzip_auto(self):
with TempDir() as temp_dir:
path = temp_dir.create_temp_file('result.gz')
diff --git a/sdks/python/apache_beam/io/utils.py b/sdks/python/apache_beam/io/utils.py
index d5912ae..8b12fea 100644
--- a/sdks/python/apache_beam/io/utils.py
+++ b/sdks/python/apache_beam/io/utils.py
@@ -19,6 +19,8 @@
on transforms.ptransform_test.test_read_metrics.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import range
diff --git a/sdks/python/apache_beam/io/utils_test.py b/sdks/python/apache_beam/io/utils_test.py
index 94003cc..99b4b49 100644
--- a/sdks/python/apache_beam/io/utils_test.py
+++ b/sdks/python/apache_beam/io/utils_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/io/vcfio.py b/sdks/python/apache_beam/io/vcfio.py
index 1260ec1..7cff80a 100644
--- a/sdks/python/apache_beam/io/vcfio.py
+++ b/sdks/python/apache_beam/io/vcfio.py
@@ -20,6 +20,8 @@
The 4.2 spec is available at https://samtools.github.io/hts-specs/VCFv4.2.pdf.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -179,9 +181,6 @@
return self < other or self == other
- def __ne__(self, other):
- return not self == other
-
def __gt__(self, other):
if not isinstance(other, Variant):
return NotImplemented
diff --git a/sdks/python/apache_beam/io/vcfio_test.py b/sdks/python/apache_beam/io/vcfio_test.py
index 9a4b793..8091ab6 100644
--- a/sdks/python/apache_beam/io/vcfio_test.py
+++ b/sdks/python/apache_beam/io/vcfio_test.py
@@ -17,6 +17,8 @@
"""Tests for vcfio module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -494,26 +496,23 @@
with TempDir() as tempdir:
file_name = self._create_temp_vcf_file(_SAMPLE_HEADER_LINES +
_SAMPLE_TEXT_LINES, tempdir)
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromVcf(file_name)
- assert_that(pcoll, _count_equals_to(len(_SAMPLE_TEXT_LINES)))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromVcf(file_name)
+ assert_that(pcoll, _count_equals_to(len(_SAMPLE_TEXT_LINES)))
@unittest.skipIf(VCF_FILE_DIR_MISSING, 'VCF test file directory is missing')
def test_pipeline_read_single_file_large(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromVcf(
- get_full_file_path('valid-4.0.vcf'))
- assert_that(pcoll, _count_equals_to(5))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromVcf(
+ get_full_file_path('valid-4.0.vcf'))
+ assert_that(pcoll, _count_equals_to(5))
@unittest.skipIf(VCF_FILE_DIR_MISSING, 'VCF test file directory is missing')
def test_pipeline_read_file_pattern_large(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Read' >> ReadFromVcf(
- os.path.join(get_full_dir(), 'valid-*.vcf'))
- assert_that(pcoll, _count_equals_to(9900))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Read' >> ReadFromVcf(
+ os.path.join(get_full_dir(), 'valid-*.vcf'))
+ assert_that(pcoll, _count_equals_to(9900))
def test_read_reentrant_without_splitting(self):
with TempDir() as tempdir:
diff --git a/sdks/python/apache_beam/metrics/cells.py b/sdks/python/apache_beam/metrics/cells.py
index 96dd973..8f91288 100644
--- a/sdks/python/apache_beam/metrics/cells.py
+++ b/sdks/python/apache_beam/metrics/cells.py
@@ -21,12 +21,15 @@
context.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
import threading
import time
from builtins import object
+from typing import Optional
from apache_beam.portability.api import beam_fn_api_pb2
from apache_beam.portability.api import metrics_pb2
@@ -84,6 +87,7 @@
self.value = CounterAggregator.identity_element()
def combine(self, other):
+ # type: (CounterCell) -> CounterCell
result = CounterCell()
result.inc(self.value + other.value)
return result
@@ -104,6 +108,7 @@
self.value += value
def get_cumulative(self):
+ # type: () -> int
with self._lock:
return self.value
@@ -142,6 +147,7 @@
self.data = DistributionAggregator.identity_element()
def combine(self, other):
+ # type: (DistributionCell) -> DistributionCell
result = DistributionCell()
result.data = self.data.combine(other.data)
return result
@@ -167,6 +173,7 @@
self.data.max = ivalue
def get_cumulative(self):
+ # type: () -> DistributionData
with self._lock:
return self.data.get_cumulative()
@@ -202,6 +209,7 @@
self.data = GaugeAggregator.identity_element()
def combine(self, other):
+ # type: (GaugeCell) -> GaugeCell
result = GaugeCell()
result.data = self.data.combine(other.data)
return result
@@ -218,6 +226,7 @@
self.data.timestamp = time.time()
def get_cumulative(self):
+ # type: () -> GaugeData
with self._lock:
return self.data.get_cumulative()
@@ -237,6 +246,7 @@
class DistributionResult(object):
"""The result of a Distribution metric."""
def __init__(self, data):
+ # type: (DistributionData) -> None
self.data = data
def __eq__(self, other):
@@ -288,6 +298,7 @@
class GaugeResult(object):
def __init__(self, data):
+ # type: (GaugeData) -> None
self.data = data
def __eq__(self, other):
@@ -347,9 +358,11 @@
self.timestamp)
def get_cumulative(self):
+ # type: () -> GaugeData
return GaugeData(self.value, timestamp=self.timestamp)
def combine(self, other):
+ # type: (Optional[GaugeData]) -> GaugeData
if other is None:
return self
@@ -360,6 +373,7 @@
@staticmethod
def singleton(value, timestamp=None):
+ # type: (...) -> GaugeData
return GaugeData(value, timestamp=timestamp)
def to_runner_api(self):
@@ -425,9 +439,11 @@
self.max)
def get_cumulative(self):
+ # type: () -> DistributionData
return DistributionData(self.sum, self.count, self.min, self.max)
def combine(self, other):
+ # type: (Optional[DistributionData]) -> DistributionData
if other is None:
return self
@@ -472,7 +488,7 @@
"""
raise NotImplementedError
- def combine(self, updates):
+ def combine(self, x, y):
raise NotImplementedError
def result(self, x):
@@ -488,12 +504,15 @@
"""
@staticmethod
def identity_element():
+ # type: () -> int
return 0
def combine(self, x, y):
+ # type: (...) -> int
return int(x) + int(y)
def result(self, x):
+ # type: (...) -> int
return int(x)
@@ -506,12 +525,15 @@
"""
@staticmethod
def identity_element():
+ # type: () -> DistributionData
return DistributionData(0, 0, 2**63 - 1, -2**63)
def combine(self, x, y):
+ # type: (DistributionData, DistributionData) -> DistributionData
return x.combine(y)
def result(self, x):
+ # type: (DistributionData) -> DistributionResult
return DistributionResult(x.get_cumulative())
@@ -524,11 +546,14 @@
"""
@staticmethod
def identity_element():
+ # type: () -> GaugeData
return GaugeData(None, timestamp=0)
def combine(self, x, y):
+ # type: (GaugeData, GaugeData) -> GaugeData
result = x.combine(y)
return result
def result(self, x):
+ # type: (GaugeData) -> GaugeResult
return GaugeResult(x.get_cumulative())
diff --git a/sdks/python/apache_beam/metrics/cells_test.py b/sdks/python/apache_beam/metrics/cells_test.py
index d50cc9c..ca59695 100644
--- a/sdks/python/apache_beam/metrics/cells_test.py
+++ b/sdks/python/apache_beam/metrics/cells_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import threading
diff --git a/sdks/python/apache_beam/metrics/execution.py b/sdks/python/apache_beam/metrics/execution.py
index 6918914..ac08407 100644
--- a/sdks/python/apache_beam/metrics/execution.py
+++ b/sdks/python/apache_beam/metrics/execution.py
@@ -30,6 +30,8 @@
unit-of-commit (bundle).
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/metrics/execution_test.py b/sdks/python/apache_beam/metrics/execution_test.py
index fc363a4..ebaa106 100644
--- a/sdks/python/apache_beam/metrics/execution_test.py
+++ b/sdks/python/apache_beam/metrics/execution_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/metrics/metric.py b/sdks/python/apache_beam/metrics/metric.py
index 8bbe191..edefae5 100644
--- a/sdks/python/apache_beam/metrics/metric.py
+++ b/sdks/python/apache_beam/metrics/metric.py
@@ -24,6 +24,8 @@
- Metrics - This class lets pipeline and transform writers create and access
metric objects such as counters, distributions, etc.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import inspect
diff --git a/sdks/python/apache_beam/metrics/metric_test.py b/sdks/python/apache_beam/metrics/metric_test.py
index cb18dc7..ce5bc2e 100644
--- a/sdks/python/apache_beam/metrics/metric_test.py
+++ b/sdks/python/apache_beam/metrics/metric_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/metrics/metricbase.py b/sdks/python/apache_beam/metrics/metricbase.py
index 420bfef..af78771 100644
--- a/sdks/python/apache_beam/metrics/metricbase.py
+++ b/sdks/python/apache_beam/metrics/metricbase.py
@@ -32,6 +32,8 @@
- MetricName - Namespace and name used to refer to a Metric.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/metrics/monitoring_infos.py b/sdks/python/apache_beam/metrics/monitoring_infos.py
index ce281e4..fac301c 100644
--- a/sdks/python/apache_beam/metrics/monitoring_infos.py
+++ b/sdks/python/apache_beam/metrics/monitoring_infos.py
@@ -18,6 +18,8 @@
# cython: language_level=3
# cython: profile=True
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/metrics/monitoring_infos_test.py b/sdks/python/apache_beam/metrics/monitoring_infos_test.py
index 466969a..696d629 100644
--- a/sdks/python/apache_beam/metrics/monitoring_infos_test.py
+++ b/sdks/python/apache_beam/metrics/monitoring_infos_test.py
@@ -14,6 +14,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py
index 15a5d96..52ff363 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -17,6 +17,8 @@
"""Pipeline options obtained from command line parsing."""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
@@ -286,8 +288,10 @@
_LOGGER.warning("Discarding unparseable args: %s", unknown_args)
result = vars(known_args)
+ overrides = self._all_options.copy()
# Apply the overrides if any
for k in list(result):
+ overrides.pop(k, None)
if k in self._all_options:
result[k] = self._all_options[k]
if (drop_default and
@@ -295,6 +299,9 @@
not isinstance(parser.get_default(k), ValueProvider)):
del result[k]
+ if overrides:
+ _LOGGER.warning("Discarding invalid overrides: %s", overrides)
+
return result
def display_data(self):
@@ -438,6 +445,12 @@
type=int,
default=1,
help='number of parallel running workers.')
+ parser.add_argument(
+ '--direct_running_mode',
+ default='in_memory',
+ choices=['in_memory', 'multi_threading', 'multi_processing'],
+ help='Workers running environment.'
+ )
class GoogleCloudOptions(PipelineOptions):
diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py
index fdedfc4..e2d7418 100644
--- a/sdks/python/apache_beam/options/pipeline_options_test.py
+++ b/sdks/python/apache_beam/options/pipeline_options_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the pipeline options module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -54,6 +56,20 @@
'mock_option': None,
'mock_multi_option': None},
'display_data': [DisplayDataItemMatcher('direct_num_workers', 5)]},
+ {'flags': ['--direct_running_mode', 'multi_threading'],
+ 'expected': {'direct_running_mode': 'multi_threading',
+ 'mock_flag': False,
+ 'mock_option': None,
+ 'mock_multi_option': None},
+ 'display_data': [DisplayDataItemMatcher('direct_running_mode',
+ 'multi_threading')]},
+ {'flags': ['--direct_running_mode', 'multi_processing'],
+ 'expected': {'direct_running_mode': 'multi_processing',
+ 'mock_flag': False,
+ 'mock_option': None,
+ 'mock_multi_option': None},
+ 'display_data': [DisplayDataItemMatcher('direct_running_mode',
+ 'multi_processing')]},
{
'flags': [
'--profile_cpu', '--profile_location', 'gs://bucket/', 'ignored'],
@@ -248,6 +264,18 @@
self.assertEqual(options.get_all_options()['num_workers'], 5)
self.assertTrue(options.get_all_options()['mock_flag'])
+ def test_override_init_options(self):
+ base_flags = ['--num_workers', '5']
+ options = PipelineOptions(base_flags, mock_flag=True)
+ self.assertEqual(options.get_all_options()['num_workers'], 5)
+ self.assertEqual(options.get_all_options()['mock_flag'], True)
+
+ def test_invalid_override_init_options(self):
+ base_flags = ['--num_workers', '5']
+ options = PipelineOptions(base_flags, mock_invalid_flag=True)
+ self.assertEqual(options.get_all_options()['num_workers'], 5)
+ self.assertEqual(options.get_all_options()['mock_flag'], False)
+
def test_experiments(self):
options = PipelineOptions(['--experiment', 'abc', '--experiment', 'def'])
self.assertEqual(
diff --git a/sdks/python/apache_beam/options/pipeline_options_validator.py b/sdks/python/apache_beam/options/pipeline_options_validator.py
index 33c35b3..900a617 100644
--- a/sdks/python/apache_beam/options/pipeline_options_validator.py
+++ b/sdks/python/apache_beam/options/pipeline_options_validator.py
@@ -19,6 +19,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import re
diff --git a/sdks/python/apache_beam/options/pipeline_options_validator_test.py b/sdks/python/apache_beam/options/pipeline_options_validator_test.py
index f380f9e..e7b474d 100644
--- a/sdks/python/apache_beam/options/pipeline_options_validator_test.py
+++ b/sdks/python/apache_beam/options/pipeline_options_validator_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the pipeline options validator module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/options/value_provider.py b/sdks/python/apache_beam/options/value_provider.py
index 380dd7b..4ff5ad7 100644
--- a/sdks/python/apache_beam/options/value_provider.py
+++ b/sdks/python/apache_beam/options/value_provider.py
@@ -19,6 +19,8 @@
and dynamically provided values.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/options/value_provider_test.py b/sdks/python/apache_beam/options/value_provider_test.py
index 7ac8670..6c81bb6 100644
--- a/sdks/python/apache_beam/options/value_provider_test.py
+++ b/sdks/python/apache_beam/options/value_provider_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the ValueProvider class."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py
index f42f44a..2dcd14c 100644
--- a/sdks/python/apache_beam/pipeline.py
+++ b/sdks/python/apache_beam/pipeline.py
@@ -44,12 +44,13 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import abc
import logging
import os
-import re
import shutil
import tempfile
from builtins import object
@@ -78,7 +79,10 @@
from apache_beam.portability import common_urns
from apache_beam.runners import PipelineRunner
from apache_beam.runners import create_runner
+from apache_beam.transforms import ParDo
from apache_beam.transforms import ptransform
+from apache_beam.transforms.core import RunnerAPIPTransformHolder
+from apache_beam.transforms.sideinputs import get_sideinput_index
#from apache_beam.transforms import external
from apache_beam.typehints import TypeCheckError
from apache_beam.typehints import typehints
@@ -89,6 +93,7 @@
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.runners.pipeline_context import PipelineContext
from apache_beam.runners.runner import PipelineResult
+ from apache_beam.transforms import environments
__all__ = ['Pipeline', 'PTransformOverride']
@@ -677,7 +682,7 @@
return_context=False,
context=None, # type: Optional[PipelineContext]
use_fake_coders=False,
- default_environment=None # type: Optional[beam_runner_api_pb2.Environment]
+ default_environment=None # type: Optional[environments.Environment]
):
# type: (...) -> beam_runner_api_pb2.Pipeline
"""For internal use only; no backwards-compatibility guarantees."""
@@ -819,7 +824,7 @@
def __init__(self,
parent,
- transform, # type: ptransform.PTransform
+ transform, # type: Optional[ptransform.PTransform]
full_label, # type: str
inputs, # type: Optional[Sequence[Union[pvalue.PBegin, pvalue.PCollection]]]
environment_id=None # type: Optional[str]
@@ -907,17 +912,18 @@
# type: (...) -> None
"""Visits all nodes reachable from the current node."""
- for pval in self.inputs:
- if pval not in visited and not isinstance(pval, pvalue.PBegin):
- if pval.producer is not None:
- pval.producer.visit(visitor, pipeline, visited)
+ for in_pval in self.inputs:
+ if in_pval not in visited and not isinstance(in_pval, pvalue.PBegin):
+ if in_pval.producer is not None:
+ in_pval.producer.visit(visitor, pipeline, visited)
# The value should be visited now since we visit outputs too.
- assert pval in visited, pval
+ assert in_pval in visited, in_pval
# Visit side inputs.
- for pval in self.side_inputs:
- if isinstance(pval, pvalue.AsSideInput) and pval.pvalue not in visited:
- pval = pval.pvalue # Unpack marker-object-wrapped pvalue.
+ for side_input in self.side_inputs:
+ if isinstance(side_input, pvalue.AsSideInput) \
+ and side_input.pvalue not in visited:
+ pval = side_input.pvalue # Unpack marker-object-wrapped pvalue.
if pval.producer is not None:
pval.producer.visit(visitor, pipeline, visited)
# The value should be visited now since we visit outputs too.
@@ -942,11 +948,11 @@
# output of such a transform is the containing DoOutputsTuple, not the
# PCollection inside it. Without the code below a tagged PCollection will
# not be marked as visited while visiting its producer.
- for pval in self.outputs.values():
- if isinstance(pval, pvalue.DoOutputsTuple):
- pvals = (v for v in pval)
+ for out_pval in self.outputs.values():
+ if isinstance(out_pval, pvalue.DoOutputsTuple):
+ pvals = (v for v in out_pval)
else:
- pvals = (pval,)
+ pvals = (out_pval,)
for v in pvals:
if v not in visited:
visited.add(v)
@@ -1017,15 +1023,17 @@
main_inputs = [context.pcollections.get_by_id(id)
for tag, id in proto.inputs.items()
if not is_side_input(tag)]
+
# Ordering is important here.
- indexed_side_inputs = [(int(re.match('side([0-9]+)(-.*)?$', tag).group(1)),
+ indexed_side_inputs = [(get_sideinput_index(tag),
context.pcollections.get_by_id(id))
for tag, id in proto.inputs.items()
if is_side_input(tag)]
side_inputs = [si for _, si in sorted(indexed_side_inputs)]
+ transform = ptransform.PTransform.from_runner_api(proto.spec, context)
result = AppliedPTransform(
parent=None,
- transform=ptransform.PTransform.from_runner_api(proto.spec, context),
+ transform=transform,
full_label=proto.unique_name,
inputs=main_inputs,
environment_id=proto.environment_id)
@@ -1043,6 +1051,9 @@
for tag, id in proto.outputs.items()}
# This annotation is expected by some runners.
if proto.spec.urn == common_urns.primitives.PAR_DO.urn:
+ # TODO(BEAM-9168): Figure out what to do for RunnerAPIPTransformHolder.
+ assert isinstance(result.transform, (ParDo, RunnerAPIPTransformHolder)),\
+ type(result.transform)
result.transform.output_tags = set(proto.outputs.keys()).difference(
{'None'})
if not result.parts:
diff --git a/sdks/python/apache_beam/pipeline_test.py b/sdks/python/apache_beam/pipeline_test.py
index 0e11a9a..7415baf 100644
--- a/sdks/python/apache_beam/pipeline_test.py
+++ b/sdks/python/apache_beam/pipeline_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the Pipeline class."""
+# pytype: skip-file
+
from __future__ import absolute_import
import copy
@@ -152,86 +154,83 @@
self.leave_composite.append(transform_node)
def test_create(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'label1' >> Create([1, 2, 3])
- assert_that(pcoll, equal_to([1, 2, 3]))
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'label1' >> Create([1, 2, 3])
+ assert_that(pcoll, equal_to([1, 2, 3]))
- # Test if initial value is an iterator object.
- pcoll2 = pipeline | 'label2' >> Create(iter((4, 5, 6)))
- pcoll3 = pcoll2 | 'do' >> FlatMap(lambda x: [x + 10])
- assert_that(pcoll3, equal_to([14, 15, 16]), label='pcoll3')
- pipeline.run()
+ # Test if initial value is an iterator object.
+ pcoll2 = pipeline | 'label2' >> Create(iter((4, 5, 6)))
+ pcoll3 = pcoll2 | 'do' >> FlatMap(lambda x: [x + 10])
+ assert_that(pcoll3, equal_to([14, 15, 16]), label='pcoll3')
def test_flatmap_builtin(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'label1' >> Create([1, 2, 3])
- assert_that(pcoll, equal_to([1, 2, 3]))
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'label1' >> Create([1, 2, 3])
+ assert_that(pcoll, equal_to([1, 2, 3]))
- pcoll2 = pcoll | 'do' >> FlatMap(lambda x: [x + 10])
- assert_that(pcoll2, equal_to([11, 12, 13]), label='pcoll2')
+ pcoll2 = pcoll | 'do' >> FlatMap(lambda x: [x + 10])
+ assert_that(pcoll2, equal_to([11, 12, 13]), label='pcoll2')
- pcoll3 = pcoll2 | 'm1' >> Map(lambda x: [x, 12])
- assert_that(pcoll3,
- equal_to([[11, 12], [12, 12], [13, 12]]), label='pcoll3')
+ pcoll3 = pcoll2 | 'm1' >> Map(lambda x: [x, 12])
+ assert_that(pcoll3,
+ equal_to([[11, 12], [12, 12], [13, 12]]), label='pcoll3')
- pcoll4 = pcoll3 | 'do2' >> FlatMap(set)
- assert_that(pcoll4, equal_to([11, 12, 12, 12, 13]), label='pcoll4')
- pipeline.run()
+ pcoll4 = pcoll3 | 'do2' >> FlatMap(set)
+ assert_that(pcoll4, equal_to([11, 12, 12, 12, 13]), label='pcoll4')
def test_maptuple_builtin(self):
- pipeline = TestPipeline()
- pcoll = pipeline | Create([('e1', 'e2')])
- side1 = beam.pvalue.AsSingleton(pipeline | 'side1' >> Create(['s1']))
- side2 = beam.pvalue.AsSingleton(pipeline | 'side2' >> Create(['s2']))
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | Create([('e1', 'e2')])
+ side1 = beam.pvalue.AsSingleton(pipeline | 'side1' >> Create(['s1']))
+ side2 = beam.pvalue.AsSingleton(pipeline | 'side2' >> Create(['s2']))
- # A test function with a tuple input, an auxiliary parameter,
- # and some side inputs.
- fn = lambda e1, e2, t=DoFn.TimestampParam, s1=None, s2=None: (
- e1, e2, t, s1, s2)
- assert_that(pcoll | 'NoSides' >> beam.core.MapTuple(fn),
- equal_to([('e1', 'e2', MIN_TIMESTAMP, None, None)]),
- label='NoSidesCheck')
- assert_that(pcoll | 'StaticSides' >> beam.core.MapTuple(fn, 's1', 's2'),
- equal_to([('e1', 'e2', MIN_TIMESTAMP, 's1', 's2')]),
- label='StaticSidesCheck')
- assert_that(pcoll | 'DynamicSides' >> beam.core.MapTuple(fn, side1, side2),
- equal_to([('e1', 'e2', MIN_TIMESTAMP, 's1', 's2')]),
- label='DynamicSidesCheck')
- assert_that(pcoll | 'MixedSides' >> beam.core.MapTuple(fn, s2=side2),
- equal_to([('e1', 'e2', MIN_TIMESTAMP, None, 's2')]),
- label='MixedSidesCheck')
- pipeline.run()
+ # A test function with a tuple input, an auxiliary parameter,
+ # and some side inputs.
+ fn = lambda e1, e2, t=DoFn.TimestampParam, s1=None, s2=None: (
+ e1, e2, t, s1, s2)
+ assert_that(pcoll | 'NoSides' >> beam.core.MapTuple(fn),
+ equal_to([('e1', 'e2', MIN_TIMESTAMP, None, None)]),
+ label='NoSidesCheck')
+ assert_that(pcoll | 'StaticSides' >> beam.core.MapTuple(fn, 's1', 's2'),
+ equal_to([('e1', 'e2', MIN_TIMESTAMP, 's1', 's2')]),
+ label='StaticSidesCheck')
+ assert_that(pcoll | 'DynamicSides' >> beam.core.MapTuple(
+ fn, side1, side2),
+ equal_to([('e1', 'e2', MIN_TIMESTAMP, 's1', 's2')]),
+ label='DynamicSidesCheck')
+ assert_that(pcoll | 'MixedSides' >> beam.core.MapTuple(fn, s2=side2),
+ equal_to([('e1', 'e2', MIN_TIMESTAMP, None, 's2')]),
+ label='MixedSidesCheck')
def test_flatmaptuple_builtin(self):
- pipeline = TestPipeline()
- pcoll = pipeline | Create([('e1', 'e2')])
- side1 = beam.pvalue.AsSingleton(pipeline | 'side1' >> Create(['s1']))
- side2 = beam.pvalue.AsSingleton(pipeline | 'side2' >> Create(['s2']))
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | Create([('e1', 'e2')])
+ side1 = beam.pvalue.AsSingleton(pipeline | 'side1' >> Create(['s1']))
+ side2 = beam.pvalue.AsSingleton(pipeline | 'side2' >> Create(['s2']))
- # A test function with a tuple input, an auxiliary parameter,
- # and some side inputs.
- fn = lambda e1, e2, t=DoFn.TimestampParam, s1=None, s2=None: (
- e1, e2, t, s1, s2)
- assert_that(pcoll | 'NoSides' >> beam.core.FlatMapTuple(fn),
- equal_to(['e1', 'e2', MIN_TIMESTAMP, None, None]),
- label='NoSidesCheck')
- assert_that(pcoll | 'StaticSides' >> beam.core.FlatMapTuple(fn, 's1', 's2'),
- equal_to(['e1', 'e2', MIN_TIMESTAMP, 's1', 's2']),
- label='StaticSidesCheck')
- assert_that(pcoll
- | 'DynamicSides' >> beam.core.FlatMapTuple(fn, side1, side2),
- equal_to(['e1', 'e2', MIN_TIMESTAMP, 's1', 's2']),
- label='DynamicSidesCheck')
- assert_that(pcoll | 'MixedSides' >> beam.core.FlatMapTuple(fn, s2=side2),
- equal_to(['e1', 'e2', MIN_TIMESTAMP, None, 's2']),
- label='MixedSidesCheck')
- pipeline.run()
+ # A test function with a tuple input, an auxiliary parameter,
+ # and some side inputs.
+ fn = lambda e1, e2, t=DoFn.TimestampParam, s1=None, s2=None: (
+ e1, e2, t, s1, s2)
+ assert_that(pcoll | 'NoSides' >> beam.core.FlatMapTuple(fn),
+ equal_to(['e1', 'e2', MIN_TIMESTAMP, None, None]),
+ label='NoSidesCheck')
+ assert_that(pcoll | 'StaticSides' >> beam.core.FlatMapTuple(
+ fn, 's1', 's2'),
+ equal_to(['e1', 'e2', MIN_TIMESTAMP, 's1', 's2']),
+ label='StaticSidesCheck')
+ assert_that(pcoll
+ | 'DynamicSides' >> beam.core.FlatMapTuple(fn, side1, side2),
+ equal_to(['e1', 'e2', MIN_TIMESTAMP, 's1', 's2']),
+ label='DynamicSidesCheck')
+ assert_that(pcoll | 'MixedSides' >> beam.core.FlatMapTuple(fn, s2=side2),
+ equal_to(['e1', 'e2', MIN_TIMESTAMP, None, 's2']),
+ label='MixedSidesCheck')
def test_create_singleton_pcollection(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'label' >> Create([[1, 2, 3]])
- assert_that(pcoll, equal_to([[1, 2, 3]]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'label' >> Create([[1, 2, 3]])
+ assert_that(pcoll, equal_to([[1, 2, 3]]))
# TODO(BEAM-1555): Test is failing on the service, with FakeSource.
# @attr('ValidatesRunner')
@@ -247,10 +246,9 @@
self.assertEqual(outputs_counter.committed, 6)
def test_fake_read(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'read' >> Read(FakeSource([1, 2, 3]))
- assert_that(pcoll, equal_to([1, 2, 3]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'read' >> Read(FakeSource([1, 2, 3]))
+ assert_that(pcoll, equal_to([1, 2, 3]))
def test_visit_entire_graph(self):
pipeline = Pipeline()
@@ -272,11 +270,10 @@
self.assertEqual(visitor.leave_composite[0].transform, transform)
def test_apply_custom_transform(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'pcoll' >> Create([1, 2, 3])
- result = pcoll | PipelineTest.CustomTransform()
- assert_that(result, equal_to([2, 3, 4]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'pcoll' >> Create([1, 2, 3])
+ result = pcoll | PipelineTest.CustomTransform()
+ assert_that(result, equal_to([2, 3, 4]))
def test_reuse_custom_transform_instance(self):
pipeline = Pipeline()
@@ -293,15 +290,14 @@
'pvalue | "label" >> transform')
def test_reuse_cloned_custom_transform_instance(self):
- pipeline = TestPipeline()
- pcoll1 = pipeline | 'pc1' >> Create([1, 2, 3])
- pcoll2 = pipeline | 'pc2' >> Create([4, 5, 6])
- transform = PipelineTest.CustomTransform()
- result1 = pcoll1 | transform
- result2 = pcoll2 | 'new_label' >> transform
- assert_that(result1, equal_to([2, 3, 4]), label='r1')
- assert_that(result2, equal_to([5, 6, 7]), label='r2')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll1 = pipeline | 'pc1' >> Create([1, 2, 3])
+ pcoll2 = pipeline | 'pc2' >> Create([4, 5, 6])
+ transform = PipelineTest.CustomTransform()
+ result1 = pcoll1 | transform
+ result2 = pcoll2 | 'new_label' >> transform
+ assert_that(result1, equal_to([2, 3, 4]), label='r1')
+ assert_that(result2, equal_to([5, 6, 7]), label='r2')
def test_transform_no_super_init(self):
class AddSuffix(PTransform):
@@ -345,24 +341,23 @@
# TODO(robertwb): reduce memory usage of FnApiRunner so that this test
# passes.
- pipeline = TestPipeline(runner='BundleBasedDirectRunner')
+ with TestPipeline(runner='BundleBasedDirectRunner') as pipeline:
- # Consumed memory should not be proportional to the number of maps.
- memory_threshold = (
- get_memory_usage_in_bytes() + (5 * len_elements * num_elements))
+ # Consumed memory should not be proportional to the number of maps.
+ memory_threshold = (
+ get_memory_usage_in_bytes() + (5 * len_elements * num_elements))
- # Plus small additional slack for memory fluctuations during the test.
- memory_threshold += 10 * (2 ** 20)
+ # Plus small additional slack for memory fluctuations during the test.
+ memory_threshold += 10 * (2 ** 20)
- biglist = pipeline | 'oom:create' >> Create(
- ['x' * len_elements] * num_elements)
- for i in range(num_maps):
- biglist = biglist | ('oom:addone-%d' % i) >> Map(lambda x: x + 'y')
- result = biglist | 'oom:check' >> Map(check_memory, memory_threshold)
- assert_that(result, equal_to(
- ['x' * len_elements + 'y' * num_maps] * num_elements))
+ biglist = pipeline | 'oom:create' >> Create(
+ ['x' * len_elements] * num_elements)
+ for i in range(num_maps):
+ biglist = biglist | ('oom:addone-%d' % i) >> Map(lambda x: x + 'y')
+ result = biglist | 'oom:check' >> Map(check_memory, memory_threshold)
+ assert_that(result, equal_to(
+ ['x' * len_elements + 'y' * num_maps] * num_elements))
- pipeline.run()
def test_aggregator_empty_input(self):
actual = [] | CombineGlobally(max).without_defaults()
@@ -471,27 +466,27 @@
else:
yield TaggedOutput('letters', x)
- p = TestPipeline()
- multi = (p
- | beam.Create([1, 2, 3, 'a', 'b', 'c'])
- | 'MyMultiOutput' >> beam.ParDo(mux_input).with_outputs())
- letters = multi.letters | 'MyLetters' >> beam.Map(lambda x: x)
- numbers = multi.numbers | 'MyNumbers' >> beam.Map(lambda x: x)
+ with TestPipeline() as p:
+ multi = (p
+ | beam.Create([1, 2, 3, 'a', 'b', 'c'])
+ | 'MyMultiOutput' >> beam.ParDo(mux_input).with_outputs())
+ letters = multi.letters | 'MyLetters' >> beam.Map(lambda x: x)
+ numbers = multi.numbers | 'MyNumbers' >> beam.Map(lambda x: x)
- # Assert that the PCollection replacement worked correctly and that elements
- # are flowing through. The replacement transform first multiples by 2 then
- # the leaf nodes inside the composite multiply by an additional 3 and 5. Use
- # prime numbers to ensure that each transform is getting executed once.
- assert_that(letters,
- equal_to(['a'*2*3, 'b'*2*3, 'c'*2*3]),
- label='assert letters')
- assert_that(numbers,
- equal_to([1*2*5, 2*2*5, 3*2*5]),
- label='assert numbers')
+ # Assert that the PCollection replacement worked correctly and that
+ # elements are flowing through. The replacement transform first
+ # multiples by 2 then the leaf nodes inside the composite multiply by
+ # an additional 3 and 5. Use prime numbers to ensure that each
+ # transform is getting executed once.
+ assert_that(letters,
+ equal_to(['a'*2*3, 'b'*2*3, 'c'*2*3]),
+ label='assert letters')
+ assert_that(numbers,
+ equal_to([1*2*5, 2*2*5, 3*2*5]),
+ label='assert numbers')
- # Do the replacement and run the element assertions.
- p.replace_all([MultiOutputOverride()])
- p.run()
+ # Do the replacement and run the element assertions.
+ p.replace_all([MultiOutputOverride()])
# The following checks the graph to make sure the replacement occurred.
visitor = PipelineTest.Visitor(visited=[])
@@ -533,20 +528,18 @@
def process(self, element, counter=DoFn.StateParam(BYTES_STATE)):
return self.return_recursive(1)
- p = TestPipeline()
- pcoll = (p
- | beam.Create([(1, 1), (2, 2), (3, 3)])
- | beam.GroupByKey()
- | beam.ParDo(StatefulDoFn()))
- p.run()
+ with TestPipeline() as p:
+ pcoll = (p
+ | beam.Create([(1, 1), (2, 2), (3, 3)])
+ | beam.GroupByKey()
+ | beam.ParDo(StatefulDoFn()))
self.assertEqual(pcoll.element_type, typehints.Any)
- p = TestPipeline()
- pcoll = (p
- | beam.Create([(1, 1), (2, 2), (3, 3)])
- | beam.GroupByKey()
- | beam.ParDo(StatefulDoFn()).with_output_types(str))
- p.run()
+ with TestPipeline() as p:
+ pcoll = (p
+ | beam.Create([(1, 1), (2, 2), (3, 3)])
+ | beam.GroupByKey()
+ | beam.ParDo(StatefulDoFn()).with_output_types(str))
self.assertEqual(pcoll.element_type, str)
def test_track_pcoll_unbounded(self):
@@ -607,40 +600,37 @@
def process(self, element):
yield element + 10
- pipeline = TestPipeline()
- pcoll = pipeline | 'Create' >> Create([1, 2]) | 'Do' >> ParDo(TestDoFn())
- assert_that(pcoll, equal_to([11, 12]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Create' >> Create([1, 2]) | 'Do' >> ParDo(TestDoFn())
+ assert_that(pcoll, equal_to([11, 12]))
def test_side_input_no_tag(self):
class TestDoFn(DoFn):
def process(self, element, prefix, suffix):
return ['%s-%s-%s' % (prefix, element, suffix)]
- pipeline = TestPipeline()
- words_list = ['aa', 'bb', 'cc']
- words = pipeline | 'SomeWords' >> Create(words_list)
- prefix = 'zyx'
- suffix = pipeline | 'SomeString' >> Create(['xyz']) # side in
- result = words | 'DecorateWordsDoFnNoTag' >> ParDo(
- TestDoFn(), prefix, suffix=AsSingleton(suffix))
- assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ words_list = ['aa', 'bb', 'cc']
+ words = pipeline | 'SomeWords' >> Create(words_list)
+ prefix = 'zyx'
+ suffix = pipeline | 'SomeString' >> Create(['xyz']) # side in
+ result = words | 'DecorateWordsDoFnNoTag' >> ParDo(
+ TestDoFn(), prefix, suffix=AsSingleton(suffix))
+ assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list]))
def test_side_input_tagged(self):
class TestDoFn(DoFn):
def process(self, element, prefix, suffix=DoFn.SideInputParam):
return ['%s-%s-%s' % (prefix, element, suffix)]
- pipeline = TestPipeline()
- words_list = ['aa', 'bb', 'cc']
- words = pipeline | 'SomeWords' >> Create(words_list)
- prefix = 'zyx'
- suffix = pipeline | 'SomeString' >> Create(['xyz']) # side in
- result = words | 'DecorateWordsDoFnNoTag' >> ParDo(
- TestDoFn(), prefix, suffix=AsSingleton(suffix))
- assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ words_list = ['aa', 'bb', 'cc']
+ words = pipeline | 'SomeWords' >> Create(words_list)
+ prefix = 'zyx'
+ suffix = pipeline | 'SomeString' >> Create(['xyz']) # side in
+ result = words | 'DecorateWordsDoFnNoTag' >> ParDo(
+ TestDoFn(), prefix, suffix=AsSingleton(suffix))
+ assert_that(result, equal_to(['zyx-%s-xyz' % x for x in words_list]))
@attr('ValidatesRunner')
def test_element_param(self):
@@ -666,32 +656,30 @@
def process(self, element, window=DoFn.WindowParam):
yield (element, (float(window.start), float(window.end)))
- pipeline = TestPipeline()
- pcoll = (pipeline
- | Create([1, 7])
- | Map(lambda x: TimestampedValue(x, x))
- | WindowInto(windowfn=SlidingWindows(10, 5))
- | ParDo(TestDoFn()))
- assert_that(pcoll, equal_to([(1, (-5, 5)), (1, (0, 10)),
- (7, (0, 10)), (7, (5, 15))]))
- pcoll2 = pcoll | 'Again' >> ParDo(TestDoFn())
- assert_that(
- pcoll2,
- equal_to([
- ((1, (-5, 5)), (-5, 5)), ((1, (0, 10)), (0, 10)),
- ((7, (0, 10)), (0, 10)), ((7, (5, 15)), (5, 15))]),
- label='doubled windows')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = (pipeline
+ | Create([1, 7])
+ | Map(lambda x: TimestampedValue(x, x))
+ | WindowInto(windowfn=SlidingWindows(10, 5))
+ | ParDo(TestDoFn()))
+ assert_that(pcoll, equal_to([(1, (-5, 5)), (1, (0, 10)),
+ (7, (0, 10)), (7, (5, 15))]))
+ pcoll2 = pcoll | 'Again' >> ParDo(TestDoFn())
+ assert_that(
+ pcoll2,
+ equal_to([
+ ((1, (-5, 5)), (-5, 5)), ((1, (0, 10)), (0, 10)),
+ ((7, (0, 10)), (0, 10)), ((7, (5, 15)), (5, 15))]),
+ label='doubled windows')
def test_timestamp_param(self):
class TestDoFn(DoFn):
def process(self, element, timestamp=DoFn.TimestampParam):
yield timestamp
- pipeline = TestPipeline()
- pcoll = pipeline | 'Create' >> Create([1, 2]) | 'Do' >> ParDo(TestDoFn())
- assert_that(pcoll, equal_to([MIN_TIMESTAMP, MIN_TIMESTAMP]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Create' >> Create([1, 2]) | 'Do' >> ParDo(TestDoFn())
+ assert_that(pcoll, equal_to([MIN_TIMESTAMP, MIN_TIMESTAMP]))
def test_timestamp_param_map(self):
with TestPipeline() as p:
@@ -731,12 +719,11 @@
# Ensure that we don't use default values in a context where they must be
# comparable (see BEAM-8301).
- pipeline = TestPipeline()
- pcoll = (pipeline
- | beam.Create([None])
- | Map(lambda e, x=IncomparableType(): (e, type(x).__name__)))
- assert_that(pcoll, equal_to([(None, 'IncomparableType')]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = (pipeline
+ | beam.Create([None])
+ | Map(lambda e, x=IncomparableType(): (e, type(x).__name__)))
+ assert_that(pcoll, equal_to([(None, 'IncomparableType')]))
class Bacon(PipelineOptions):
diff --git a/sdks/python/apache_beam/portability/common_urns.py b/sdks/python/apache_beam/portability/common_urns.py
index c1164e3..6f2943e 100644
--- a/sdks/python/apache_beam/portability/common_urns.py
+++ b/sdks/python/apache_beam/portability/common_urns.py
@@ -17,72 +17,40 @@
""" Accessors for URNs of common Beam entities. """
+# pytype: skip-file
+
from __future__ import absolute_import
-from builtins import object
+from apache_beam.portability.api.beam_runner_api_pb2_urns import BeamConstants
+from apache_beam.portability.api.beam_runner_api_pb2_urns import StandardCoders
+from apache_beam.portability.api.beam_runner_api_pb2_urns import StandardEnvironments
+from apache_beam.portability.api.beam_runner_api_pb2_urns import StandardPTransforms
+from apache_beam.portability.api.beam_runner_api_pb2_urns import StandardSideInputTypes
+from apache_beam.portability.api.metrics_pb2_urns import MonitoringInfo
+from apache_beam.portability.api.metrics_pb2_urns import MonitoringInfoSpecs
+from apache_beam.portability.api.metrics_pb2_urns import MonitoringInfoTypeUrns
+from apache_beam.portability.api.standard_window_fns_pb2_urns import FixedWindowsPayload
+from apache_beam.portability.api.standard_window_fns_pb2_urns import GlobalWindowsPayload
+from apache_beam.portability.api.standard_window_fns_pb2_urns import SessionsPayload
+from apache_beam.portability.api.standard_window_fns_pb2_urns import SlidingWindowsPayload
-from apache_beam.portability.api import beam_runner_api_pb2
-from apache_beam.portability.api import metrics_pb2
-from apache_beam.portability.api import standard_window_fns_pb2
+primitives = StandardPTransforms.Primitives
+deprecated_primitives = StandardPTransforms.DeprecatedPrimitives
+composites = StandardPTransforms.Composites
+combine_components = StandardPTransforms.CombineComponents
+sdf_components = StandardPTransforms.SplittableParDoComponents
+side_inputs = StandardSideInputTypes.Enum
+coders = StandardCoders.Enum
+constants = BeamConstants.Constants
-class PropertiesFromEnumValue(object):
- def __init__(self, value_descriptor):
- self.urn = (value_descriptor.GetOptions().Extensions[
- beam_runner_api_pb2.beam_urn])
- self.constant = (value_descriptor.GetOptions().Extensions[
- beam_runner_api_pb2.beam_constant])
- self.spec = (value_descriptor.GetOptions().Extensions[
- metrics_pb2.monitoring_info_spec])
- self.label_props = (value_descriptor.GetOptions().Extensions[
- metrics_pb2.label_props])
+environments = StandardEnvironments.Environments
+global_windows = GlobalWindowsPayload.Enum.PROPERTIES
+fixed_windows = FixedWindowsPayload.Enum.PROPERTIES
+sliding_windows = SlidingWindowsPayload.Enum.PROPERTIES
+session_windows = SessionsPayload.Enum.PROPERTIES
-class PropertiesFromEnumType(object):
- def __init__(self, enum_type):
- for v in enum_type.DESCRIPTOR.values:
- setattr(self, v.name, PropertiesFromEnumValue(v))
-
-
-primitives = PropertiesFromEnumType(
- beam_runner_api_pb2.StandardPTransforms.Primitives)
-deprecated_primitives = PropertiesFromEnumType(
- beam_runner_api_pb2.StandardPTransforms.DeprecatedPrimitives)
-composites = PropertiesFromEnumType(
- beam_runner_api_pb2.StandardPTransforms.Composites)
-combine_components = PropertiesFromEnumType(
- beam_runner_api_pb2.StandardPTransforms.CombineComponents)
-sdf_components = PropertiesFromEnumType(
- beam_runner_api_pb2.StandardPTransforms.SplittableParDoComponents)
-
-side_inputs = PropertiesFromEnumType(
- beam_runner_api_pb2.StandardSideInputTypes.Enum)
-
-coders = PropertiesFromEnumType(beam_runner_api_pb2.StandardCoders.Enum)
-
-constants = PropertiesFromEnumType(
- beam_runner_api_pb2.BeamConstants.Constants)
-
-environments = PropertiesFromEnumType(
- beam_runner_api_pb2.StandardEnvironments.Environments)
-
-
-def PropertiesFromPayloadType(payload_type):
- return PropertiesFromEnumType(payload_type.Enum).PROPERTIES
-
-
-global_windows = PropertiesFromPayloadType(
- standard_window_fns_pb2.GlobalWindowsPayload)
-fixed_windows = PropertiesFromPayloadType(
- standard_window_fns_pb2.FixedWindowsPayload)
-sliding_windows = PropertiesFromPayloadType(
- standard_window_fns_pb2.SlidingWindowsPayload)
-session_windows = PropertiesFromPayloadType(
- standard_window_fns_pb2.SessionsPayload)
-
-monitoring_info_specs = PropertiesFromEnumType(
- metrics_pb2.MonitoringInfoSpecs.Enum)
-monitoring_info_types = PropertiesFromEnumType(
- metrics_pb2.MonitoringInfoTypeUrns.Enum)
-monitoring_info_labels = PropertiesFromEnumType(
- metrics_pb2.MonitoringInfo.MonitoringInfoLabels)
+monitoring_info_specs = MonitoringInfoSpecs.Enum
+monitoring_info_types = MonitoringInfoTypeUrns.Enum
+monitoring_info_labels = MonitoringInfo.MonitoringInfoLabels
diff --git a/sdks/python/apache_beam/portability/utils.py b/sdks/python/apache_beam/portability/utils.py
new file mode 100644
index 0000000..7c60598
--- /dev/null
+++ b/sdks/python/apache_beam/portability/utils.py
@@ -0,0 +1,34 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""For internal use only; no backwards-compatibility guarantees."""
+from __future__ import absolute_import
+
+from typing import TYPE_CHECKING
+from typing import NamedTuple
+
+if TYPE_CHECKING:
+ from apache_beam.portability.api import metrics_pb2
+
+
+PropertiesFromEnumValue = NamedTuple(
+ 'PropertiesFromEnumValue', [
+ ('urn', str),
+ ('constant', str),
+ ('spec', 'metrics_pb2.MonitoringInfoSpec'),
+ ('label_props', 'metrics_pb2.MonitoringInfoLabelProps'),
+ ])
diff --git a/sdks/python/apache_beam/pvalue.py b/sdks/python/apache_beam/pvalue.py
index 5fa8ece..f072e41 100644
--- a/sdks/python/apache_beam/pvalue.py
+++ b/sdks/python/apache_beam/pvalue.py
@@ -24,6 +24,8 @@
produced when the pipeline gets executed.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
@@ -158,6 +160,7 @@
def windowing(self):
# type: () -> Windowing
if not hasattr(self, '_windowing'):
+ assert self.producer is not None and self.producer.transform is not None
self._windowing = self.producer.transform.get_windowing(
self.producer.inputs)
return self._windowing
@@ -199,10 +202,14 @@
@staticmethod
def from_runner_api(proto, context):
# type: (beam_runner_api_pb2.PCollection, PipelineContext) -> PCollection
- # Producer and tag will be filled in later, the key point is that the
- # same object is returned for the same pcollection id.
+ # Producer and tag will be filled in later, the key point is that the same
+ # object is returned for the same pcollection id.
+ # We pass None for the PCollection's Pipeline to avoid a cycle during
+ # deserialization. It will be populated soon after this call, in
+ # Pipeline.from_runner_api(). This brief period is the only time that
+ # PCollection.pipeline is allowed to be None.
return PCollection(
- None,
+ None, # type: ignore[arg-type]
element_type=context.element_type_from_coder_id(proto.coder_id),
windowing=context.windowing_strategies.get_by_id(
proto.windowing_strategy_id),
@@ -247,7 +254,7 @@
# gets applied.
self.producer = None # type: Optional[AppliedPTransform]
# Dictionary of PCollections already associated with tags.
- self._pcolls = {} # type: Dict[Optional[str], PValue]
+ self._pcolls = {} # type: Dict[Optional[str], PCollection]
def __str__(self):
return '<%s>' % self._str_internal()
@@ -260,15 +267,15 @@
self.__class__.__name__, self._main_tag, self._tags, self._transform)
def __iter__(self):
- # type: () -> Iterator[PValue]
- """Iterates over tags returning for each call a (tag, pvalue) pair."""
+ # type: () -> Iterator[PCollection]
+ """Iterates over tags returning for each call a (tag, pcollection) pair."""
if self._main_tag is not None:
yield self[self._main_tag]
for tag in self._tags:
yield self[tag]
def __getattr__(self, tag):
- # type: (str) -> PValue
+ # type: (str) -> PCollection
# Special methods which may be accessed before the object is
# fully constructed (e.g. in unpickling).
if tag[:2] == tag[-2:] == '__':
@@ -276,7 +283,7 @@
return self[tag]
def __getitem__(self, tag):
- # type: (Union[int, str, None]) -> PValue
+ # type: (Union[int, str, None]) -> PCollection
# Accept int tags so that we can look at Partition tags with the
# same ints that we used in the partition function.
# TODO(gildea): Consider requiring string-based tags everywhere.
@@ -297,7 +304,7 @@
assert self.producer is not None
if tag is not None:
self._transform.output_tags.add(tag)
- pcoll = PCollection(self._pipeline, tag=tag, element_type=typehints.Any) # type: PValue
+ pcoll = PCollection(self._pipeline, tag=tag, element_type=typehints.Any)
# Transfer the producer from the DoOutputsTuple to the resulting
# PCollection.
pcoll.producer = self.producer.parts[0]
@@ -308,7 +315,10 @@
self.producer.add_output(pcoll, tag)
else:
# Main output is output of inner ParDo.
- pcoll = self.producer.parts[0].outputs[None]
+ pval = self.producer.parts[0].outputs[None]
+ assert isinstance(pval, PCollection), (
+ "DoOutputsTuple should follow a ParDo.")
+ pcoll = pval
self._pcolls[tag] = pcoll
return pcoll
diff --git a/sdks/python/apache_beam/pvalue_test.py b/sdks/python/apache_beam/pvalue_test.py
index b3f02e3..de689a0 100644
--- a/sdks/python/apache_beam/pvalue_test.py
+++ b/sdks/python/apache_beam/pvalue_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the PValue and PCollection classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/common.pxd b/sdks/python/apache_beam/runners/common.pxd
index 37e05bf..00840fe 100644
--- a/sdks/python/apache_beam/runners/common.pxd
+++ b/sdks/python/apache_beam/runners/common.pxd
@@ -69,7 +69,6 @@
cpdef invoke_process(self, WindowedValue windowed_value,
restriction_tracker=*,
- OutputProcessor output_processor=*,
additional_args=*, additional_kwargs=*)
cpdef invoke_start_bundle(self)
cpdef invoke_finish_bundle(self)
diff --git a/sdks/python/apache_beam/runners/common.py b/sdks/python/apache_beam/runners/common.py
index 1f116ae..a6145ac 100644
--- a/sdks/python/apache_beam/runners/common.py
+++ b/sdks/python/apache_beam/runners/common.py
@@ -21,6 +21,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import traceback
@@ -338,7 +340,7 @@
represented by a given DoFnSignature."""
def __init__(self,
- output_processor, # type: Optional[_OutputProcessor]
+ output_processor, # type: OutputProcessor
signature # type: DoFnSignature
):
# type: (...) -> None
@@ -357,7 +359,7 @@
@staticmethod
def create_invoker(
signature, # type: DoFnSignature
- output_processor=None, # type: Optional[_OutputProcessor]
+ output_processor, # type: OutputProcessor
context=None, # type: Optional[DoFnContext]
side_inputs=None, # type: Optional[List[sideinputs.SideInputMap]]
input_args=None, input_kwargs=None,
@@ -398,6 +400,8 @@
if use_simple_invoker:
return SimpleInvoker(output_processor, signature)
else:
+ if context is None:
+ raise TypeError("Must provide context when not using SimpleInvoker")
return PerWindowInvoker(
output_processor,
signature, context, side_inputs, input_args, input_kwargs,
@@ -406,7 +410,6 @@
def invoke_process(self,
windowed_value, # type: WindowedValue
restriction_tracker=None, # type: Optional[iobase.RestrictionTracker]
- output_processor=None, # type: Optional[OutputProcessor]
additional_args=None,
additional_kwargs=None
):
@@ -417,7 +420,7 @@
windowed_value: a WindowedValue object that gives the element for which
process() method should be invoked along with the window
the element belongs to.
- output_procesor: if provided given OutputProcessor will be used.
+ output_processor: if provided given OutputProcessor will be used.
additional_args: additional arguments to be passed to the current
`DoFn.process()` invocation, usually as side inputs.
additional_kwargs: additional keyword arguments to be passed to the
@@ -435,14 +438,16 @@
# type: () -> None
"""Invokes the DoFn.start_bundle() method.
"""
- self.output_processor.start_bundle_outputs(
+ # self.output_processor is Optional, but in practice it won't be None here
+ self.output_processor.start_bundle_outputs( # type: ignore[union-attr]
self.signature.start_bundle_method.method_value())
def invoke_finish_bundle(self):
# type: () -> None
"""Invokes the DoFn.finish_bundle() method.
"""
- self.output_processor.finish_bundle_outputs(
+ # self.output_processor is Optional, but in practice it won't be None here
+ self.output_processor.finish_bundle_outputs( # type: ignore[union-attr]
self.signature.finish_bundle_method.method_value())
def invoke_teardown(self):
@@ -452,7 +457,8 @@
self.signature.teardown_lifecycle_method.method_value()
def invoke_user_timer(self, timer_spec, key, window, timestamp):
- self.output_processor.process_outputs(
+ # self.output_processor is Optional, but in practice it won't be None here
+ self.output_processor.process_outputs( # type: ignore[union-attr]
WindowedValue(None, timestamp, (window,)),
self.signature.timer_methods[timer_spec].invoke_timer_callback(
self.user_state_context, key, window, timestamp))
@@ -474,7 +480,7 @@
"""An invoker that processes elements ignoring windowing information."""
def __init__(self,
- output_processor, # type: Optional[_OutputProcessor]
+ output_processor, # type: OutputProcessor
signature # type: DoFnSignature
):
# type: (...) -> None
@@ -484,14 +490,11 @@
def invoke_process(self,
windowed_value, # type: WindowedValue
restriction_tracker=None, # type: Optional[iobase.RestrictionTracker]
- output_processor=None, # type: Optional[OutputProcessor]
additional_args=None,
additional_kwargs=None
):
# type: (...) -> None
- if not output_processor:
- output_processor = self.output_processor
- output_processor.process_outputs(
+ self.output_processor.process_outputs(
windowed_value, self.process_method(windowed_value.value))
@@ -499,7 +502,7 @@
"""An invoker that processes elements considering windowing information."""
def __init__(self,
- output_processor, # type: Optional[_OutputProcessor]
+ output_processor, # type: OutputProcessor
signature, # type: DoFnSignature
context, # type: DoFnContext
side_inputs, # type: Iterable[sideinputs.SideInputMap]
@@ -523,8 +526,8 @@
self.watermark_estimator_param = (
self.signature.process_method.watermark_estimator_arg_name
if self.watermark_estimator else None)
- self.threadsafe_restriction_tracker = None
- self.current_windowed_value = None
+ self.threadsafe_restriction_tracker = None # type: Optional[iobase.ThreadsafeRestrictionTracker]
+ self.current_windowed_value = None # type: Optional[WindowedValue]
self.bundle_finalizer_param = bundle_finalizer_param
self.is_key_param_required = False
@@ -606,7 +609,6 @@
def invoke_process(self,
windowed_value, # type: WindowedValue
restriction_tracker=None,
- output_processor=None, # type: Optional[OutputProcessor]
additional_args=None,
additional_kwargs=None
):
@@ -616,8 +618,6 @@
if not additional_kwargs:
additional_kwargs = {}
- if not output_processor:
- output_processor = self.output_processor
self.context.set_element(windowed_value)
# Call for the process function for each window if has windowed side inputs
# or if the process accesses the window parameter. We can just call it once
@@ -653,8 +653,7 @@
try:
self.current_windowed_value = windowed_value
return self._invoke_process_per_window(
- windowed_value, additional_args, additional_kwargs,
- output_processor)
+ windowed_value, additional_args, additional_kwargs)
finally:
self.threadsafe_restriction_tracker = None
self.current_windowed_value = windowed_value
@@ -663,16 +662,16 @@
for w in windowed_value.windows:
self._invoke_process_per_window(
WindowedValue(windowed_value.value, windowed_value.timestamp, (w,)),
- additional_args, additional_kwargs, output_processor)
+ additional_args, additional_kwargs)
else:
self._invoke_process_per_window(
- windowed_value, additional_args, additional_kwargs, output_processor)
+ windowed_value, additional_args, additional_kwargs)
+ return None
def _invoke_process_per_window(self,
windowed_value, # type: WindowedValue
additional_args,
additional_kwargs,
- output_processor # type: OutputProcessor
):
# type: (...) -> Optional[Tuple[WindowedValue, Timestamp]]
if self.has_windowed_inputs:
@@ -723,9 +722,11 @@
elif core.DoFn.PaneInfoParam == p:
args_for_process[i] = windowed_value.pane_info
elif isinstance(p, core.DoFn.StateParam):
+ assert self.user_state_context is not None
args_for_process[i] = (
self.user_state_context.get_state(p.state_spec, key, window))
elif isinstance(p, core.DoFn.TimerParam):
+ assert self.user_state_context is not None
args_for_process[i] = (
self.user_state_context.get_timer(p.timer_spec, key, window))
elif core.DoFn.BundleFinalizerParam == p:
@@ -739,14 +740,15 @@
kwargs_for_process[key] = additional_kwargs[key]
if kwargs_for_process:
- output_processor.process_outputs(
+ self.output_processor.process_outputs(
windowed_value,
self.process_method(*args_for_process, **kwargs_for_process))
else:
- output_processor.process_outputs(
+ self.output_processor.process_outputs(
windowed_value, self.process_method(*args_for_process))
if self.is_splittable:
+ assert self.threadsafe_restriction_tracker is not None
# TODO: Consider calling check_done right after SDF.Process() finishing.
# In order to do this, we need to know that current invoking dofn is
# ProcessSizedElementAndRestriction.
@@ -763,11 +765,11 @@
return ((
windowed_value.with_value(((element, deferred_restriction), size)),
output_watermark), deferred_watermark)
+ return None
def try_split(self, fraction):
- restriction_tracker = self.threadsafe_restriction_tracker
- current_windowed_value = self.current_windowed_value
- if restriction_tracker and current_windowed_value:
+ # type: (...) -> Optional[Tuple[SplitResultType, SplitResultType]]
+ if self.threadsafe_restriction_tracker and self.current_windowed_value:
# Temporary workaround for [BEAM-7473]: get current_watermark before
# split, in case watermark gets advanced before getting split results.
# In worst case, current_watermark is always stale, which is ok.
@@ -775,7 +777,7 @@
current_watermark = self.watermark_estimator.current_watermark()
else:
current_watermark = None
- split = restriction_tracker.try_split(fraction)
+ split = self.threadsafe_restriction_tracker.try_split(fraction)
if split:
primary, residual = split
element = self.current_windowed_value.value
@@ -793,6 +795,8 @@
restriction_tracker = self.threadsafe_restriction_tracker
if restriction_tracker:
return restriction_tracker.current_progress()
+ else:
+ return None
class DoFnRunner(Receiver):
@@ -807,7 +811,7 @@
kwargs,
side_inputs, # type: Iterable[sideinputs.SideInputMap]
windowing,
- tagged_receivers=None, # type: Mapping[Optional[str], Receiver]
+ tagged_receivers, # type: Mapping[Optional[str], Receiver]
step_name=None, # type: Optional[str]
logging_context=None,
state=None,
@@ -879,6 +883,7 @@
return self.do_fn_invoker.invoke_process(windowed_value)
except BaseException as exn:
self._reraise_augmented(exn)
+ return None
def process_with_sized_restriction(self, windowed_value):
# type: (WindowedValue) -> Optional[Tuple[WindowedValue, Timestamp]]
@@ -893,6 +898,7 @@
def current_element_progress(self):
# type: () -> Optional[iobase.RestrictionProgress]
+ assert isinstance(self.do_fn_invoker, PerWindowInvoker)
return self.do_fn_invoker.current_element_progress()
def process_user_timer(self, timer_spec, key, window, timestamp):
diff --git a/sdks/python/apache_beam/runners/common_test.py b/sdks/python/apache_beam/runners/common_test.py
index 9377708..62e7147 100644
--- a/sdks/python/apache_beam/runners/common_test.py
+++ b/sdks/python/apache_beam/runners/common_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline.py b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline.py
index f0d6df9..eb8963e 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline.py
@@ -17,6 +17,8 @@
"""A word-counting workflow."""
+# pytype: skip-file
+
from __future__ import absolute_import
import time
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline_test.py b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline_test.py
index d1afbcf..acbb5ae 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_metrics_pipeline_test.py
@@ -17,6 +17,8 @@
"""A word-counting workflow."""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline.py b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline.py
index b0e75e6..aff56cf 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline.py
@@ -17,6 +17,8 @@
"""A word-counting workflow."""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline_test.py b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline_test.py
index 2a351e2..3cfe244 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_exercise_streaming_metrics_pipeline_test.py
@@ -17,6 +17,8 @@
"""A word-counting workflow."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py b/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py
index d518bb0..16a14ef 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_metrics.py
@@ -21,6 +21,8 @@
service.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_metrics_test.py b/sdks/python/apache_beam/runners/dataflow/dataflow_metrics_test.py
index 9899176..db1f5a7 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_metrics_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_metrics_test.py
@@ -19,6 +19,8 @@
the DataflowMetrics class.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import types
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
index 6786760..d7c4cb46 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -20,6 +20,8 @@
The runner will create a JSON description of the job graph and then submit it
to the Dataflow Service for remote execution by a worker.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -50,6 +52,7 @@
from apache_beam.options.pipeline_options import WorkerOptions
from apache_beam.portability import common_urns
from apache_beam.pvalue import AsSideInput
+from apache_beam.runners.common import DoFnSignature
from apache_beam.runners.dataflow.internal import names
from apache_beam.runners.dataflow.internal.clients import dataflow as dataflow_api
from apache_beam.runners.dataflow.internal.names import PropertyNames
@@ -299,7 +302,7 @@
return SetPDoneVisitor(pipeline)
@staticmethod
- def side_input_visitor():
+ def side_input_visitor(use_unified_worker=False):
# Imported here to avoid circular dependencies.
# pylint: disable=wrong-import-order, wrong-import-position
from apache_beam.pipeline import PipelineVisitor
@@ -317,24 +320,33 @@
for ix, side_input in enumerate(transform_node.side_inputs):
access_pattern = side_input._side_input_data().access_pattern
if access_pattern == common_urns.side_inputs.ITERABLE.urn:
- # Add a map to ('', value) as Dataflow currently only handles
- # keyed side inputs.
- pipeline = side_input.pvalue.pipeline
- new_side_input = _DataflowIterableSideInput(side_input)
- new_side_input.pvalue = beam.pvalue.PCollection(
- pipeline,
- element_type=typehints.KV[
- bytes, side_input.pvalue.element_type],
- is_bounded=side_input.pvalue.is_bounded)
- parent = transform_node.parent or pipeline._root_transform()
- map_to_void_key = beam.pipeline.AppliedPTransform(
- pipeline,
- beam.Map(lambda x: (b'', x)),
- transform_node.full_label + '/MapToVoidKey%s' % ix,
- (side_input.pvalue,))
- new_side_input.pvalue.producer = map_to_void_key
- map_to_void_key.add_output(new_side_input.pvalue)
- parent.add_part(map_to_void_key)
+ if use_unified_worker:
+ # TODO(BEAM-9173): Stop patching up the access pattern to
+ # appease Dataflow when using the UW and hardcode the output
+ # type to be Any since the Dataflow JSON and pipeline proto
+ # can differ in coders which leads to encoding/decoding issues
+ # within the runner.
+ side_input.pvalue.element_type = typehints.Any
+ new_side_input = _DataflowIterableSideInput(side_input)
+ else:
+ # Add a map to ('', value) as Dataflow currently only handles
+ # keyed side inputs when using the JRH.
+ pipeline = side_input.pvalue.pipeline
+ new_side_input = _DataflowIterableAsMultimapSideInput(
+ side_input)
+ new_side_input.pvalue = beam.pvalue.PCollection(
+ pipeline,
+ element_type=typehints.KV[bytes,
+ side_input.pvalue.element_type],
+ is_bounded=side_input.pvalue.is_bounded)
+ parent = transform_node.parent or pipeline._root_transform()
+ map_to_void_key = beam.pipeline.AppliedPTransform(
+ pipeline, beam.Map(lambda x: (b'', x)),
+ transform_node.full_label + '/MapToVoidKey%s' % ix,
+ (side_input.pvalue,))
+ new_side_input.pvalue.producer = map_to_void_key
+ map_to_void_key.add_output(new_side_input.pvalue)
+ parent.add_part(map_to_void_key)
elif access_pattern == common_urns.side_inputs.MULTIMAP.urn:
# Ensure the input coder is a KV coder and patch up the
# access pattern to appease Dataflow.
@@ -394,7 +406,8 @@
# Convert all side inputs into a form acceptable to Dataflow.
if apiclient._use_fnapi(options):
- pipeline.visit(self.side_input_visitor())
+ pipeline.visit(
+ self.side_input_visitor(apiclient._use_unified_worker(options)))
# Performing configured PTransform overrides. Note that this is currently
# done before Runner API serialization, since the new proto needs to contain
@@ -498,7 +511,9 @@
test_options = options.view_as(TestOptions)
# If it is a dry run, return without submitting the job.
if test_options.dry_run:
- return None
+ result = PipelineResult(PipelineState.DONE)
+ result.wait_until_finish = lambda duration=None: None
+ return result
# Get a Dataflow API client and set its options
self.dataflow_client = apiclient.DataflowApplicationClient(options)
@@ -948,6 +963,10 @@
step.add_property(PropertyNames.RESTRICTION_ENCODING,
self._get_cloud_encoding(restriction_coder))
+ if options.view_as(StandardOptions).streaming and DoFnSignature(
+ transform.dofn).is_stateful_dofn():
+ step.add_property(PropertyNames.USES_KEYED_STATE, "true")
+
@staticmethod
def _pardo_fn_data(transform_node, get_label):
transform = transform_node.transform
@@ -1125,7 +1144,6 @@
coders.registry.get_coder(transform_node.outputs[None].element_type),
coders.coders.GlobalWindowCoder())
- from apache_beam.runners.dataflow.internal import apiclient
step.encoding = self._get_cloud_encoding(coder)
step.add_property(
PropertyNames.OUTPUT_INFO,
@@ -1211,7 +1229,6 @@
# correct coder.
coder = coders.WindowedValueCoder(transform.sink.coder,
coders.coders.GlobalWindowCoder())
- from apache_beam.runners.dataflow.internal import apiclient
step.encoding = self._get_cloud_encoding(coder)
step.add_property(PropertyNames.ENCODING, step.encoding)
step.add_property(
@@ -1313,12 +1330,12 @@
return self._data
-class _DataflowIterableSideInput(_DataflowSideInput):
+class _DataflowIterableAsMultimapSideInput(_DataflowSideInput):
"""Wraps an iterable side input as dataflow-compatible side input."""
- def __init__(self, iterable_side_input):
+ def __init__(self, side_input):
# pylint: disable=protected-access
- side_input_data = iterable_side_input._side_input_data()
+ side_input_data = side_input._side_input_data()
assert (
side_input_data.access_pattern == common_urns.side_inputs.ITERABLE.urn)
iterable_view_fn = side_input_data.view_fn
@@ -1328,6 +1345,20 @@
lambda multimap: iterable_view_fn(multimap[b'']))
+class _DataflowIterableSideInput(_DataflowSideInput):
+ """Wraps an iterable side input as dataflow-compatible side input."""
+
+ def __init__(self, side_input):
+ # pylint: disable=protected-access
+ self.pvalue = side_input.pvalue
+ side_input_data = side_input._side_input_data()
+ assert (
+ side_input_data.access_pattern == common_urns.side_inputs.ITERABLE.urn)
+ self._data = beam.pvalue.SideInputData(common_urns.side_inputs.ITERABLE.urn,
+ side_input_data.window_mapping_fn,
+ side_input_data.view_fn)
+
+
class _DataflowMultimapSideInput(_DataflowSideInput):
"""Wraps a multimap side input as dataflow-compatible side input."""
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 39bbab6..d00066c 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the DataflowRunner class."""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
@@ -208,12 +210,12 @@
self.default_properties.append('--experiments=beam_fn_api')
self.default_properties.append('--worker_harness_container_image=FOO')
remote_runner = DataflowRunner()
- p = Pipeline(remote_runner,
- options=PipelineOptions(self.default_properties))
- (p | ptransform.Create([1, 2, 3]) # pylint: disable=expression-not-assigned
- | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
- | ptransform.GroupByKey())
- p.run()
+ with Pipeline(
+ remote_runner,
+ options=PipelineOptions(self.default_properties)) as p:
+ (p | ptransform.Create([1, 2, 3]) # pylint: disable=expression-not-assigned
+ | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
+ | ptransform.GroupByKey())
self.assertEqual(
list(remote_runner.proto_pipeline.components.environments.values()),
[beam_runner_api_pb2.Environment(
@@ -223,20 +225,19 @@
def test_remote_runner_translation(self):
remote_runner = DataflowRunner()
- p = Pipeline(remote_runner,
- options=PipelineOptions(self.default_properties))
+ with Pipeline(
+ remote_runner,
+ options=PipelineOptions(self.default_properties)) as p:
- (p | ptransform.Create([1, 2, 3]) # pylint: disable=expression-not-assigned
- | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
- | ptransform.GroupByKey())
- p.run()
+ (p | ptransform.Create([1, 2, 3]) # pylint: disable=expression-not-assigned
+ | 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
+ | ptransform.GroupByKey())
def test_streaming_create_translation(self):
remote_runner = DataflowRunner()
self.default_properties.append("--streaming")
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
+ 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))
self.assertEqual(len(job_dict[u'steps']), 3)
@@ -250,11 +251,12 @@
def test_biqquery_read_streaming_fail(self):
remote_runner = DataflowRunner()
self.default_properties.append("--streaming")
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- _ = p | beam.io.Read(beam.io.BigQuerySource('some.table'))
with self.assertRaisesRegex(ValueError,
r'source is not currently available'):
- p.run()
+ with Pipeline(
+ remote_runner,
+ PipelineOptions(self.default_properties)) as p:
+ _ = p | beam.io.Read(beam.io.BigQuerySource('some.table'))
# TODO(BEAM-8095): Segfaults in Python 3.7 with xdist.
@pytest.mark.no_xdist
@@ -420,9 +422,8 @@
remote_runner = DataflowRunner()
self.default_properties.append('--min_cpu_platform=Intel Haswell')
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
+ with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p:
+ p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
self.assertIn('min_cpu_platform=Intel Haswell',
remote_runner.job.options.view_as(DebugOptions).experiments)
@@ -432,9 +433,8 @@
self.default_properties.append('--enable_streaming_engine')
self.default_properties.append('--experiment=some_other_experiment')
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
+ with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p:
+ p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
experiments_for_job = (
remote_runner.job.options.view_as(DebugOptions).experiments)
@@ -447,9 +447,8 @@
self.default_properties.append('--experiment=some_other_experiment')
self.default_properties.append('--dataflow_worker_jar=test.jar')
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
+ with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p:
+ p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
experiments_for_job = (
remote_runner.job.options.view_as(DebugOptions).experiments)
@@ -461,9 +460,8 @@
self.default_properties.append('--experiment=beam_fn_api')
self.default_properties.append('--dataflow_worker_jar=test.jar')
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
+ with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p:
+ p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
experiments_for_job = (
remote_runner.job.options.view_as(DebugOptions).experiments)
@@ -473,9 +471,8 @@
def test_use_fastavro_experiment_is_added_on_py3_and_onwards(self):
remote_runner = DataflowRunner()
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
+ with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p:
+ p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
self.assertEqual(
sys.version_info[0] > 2,
@@ -486,9 +483,8 @@
remote_runner = DataflowRunner()
self.default_properties.append('--experiment=use_avro')
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
+ with Pipeline(remote_runner, PipelineOptions(self.default_properties)) as p:
+ p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
debug_options = remote_runner.job.options.view_as(DebugOptions)
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
index da37813..200dbcc 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
@@ -19,6 +19,8 @@
Dataflow client utility functions."""
+# pytype: skip-file
+
from __future__ import absolute_import
import codecs
@@ -479,8 +481,9 @@
@retry.no_retries # Using no_retries marks this as an integration point.
def _gcs_file_copy(self, from_path, to_path):
to_folder, to_name = os.path.split(to_path)
+ total_size = os.path.getsize(from_path)
with open(from_path, 'rb') as f:
- self.stage_file(to_folder, to_name, f)
+ self.stage_file(to_folder, to_name, f, total_size=total_size)
def _stage_resources(self, options):
google_cloud_options = options.view_as(GoogleCloudOptions)
@@ -497,7 +500,7 @@
return resources
def stage_file(self, gcs_or_local_path, file_name, stream,
- mime_type='application/octet-stream'):
+ mime_type='application/octet-stream', total_size=None):
"""Stages a file at a GCS or local path with stream-supplied contents."""
if not gcs_or_local_path.startswith('gs://'):
local_path = FileSystems.join(gcs_or_local_path, file_name)
@@ -512,7 +515,7 @@
bucket=bucket, name=name)
start_time = time.time()
_LOGGER.info('Starting GCS upload to %s...', gcs_location)
- upload = storage.Upload(stream, mime_type)
+ upload = storage.Upload(stream, mime_type, total_size)
try:
response = self._storage_client.objects.Insert(request, upload=upload)
except exceptions.HttpError as e:
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 df5eb06..cbaed0b 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the apiclient module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers.py b/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers.py
index 8389e62..6749418 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from future.utils import iteritems
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers_test.py b/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers_test.py
index 3e6b6d7..3d18013 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/clients/dataflow/message_matchers_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/names.py b/sdks/python/apache_beam/runners/dataflow/internal/names.py
index 621e8fa..e9b34d45 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/names.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/names.py
@@ -20,6 +20,8 @@
# All constants are for internal use only; no backwards-compatibility
# guarantees.
+# pytype: skip-file
+
from __future__ import absolute_import
# Standard file names used for staging files.
@@ -121,6 +123,7 @@
USE_INDEXED_FORMAT = 'use_indexed_format'
USER_FN = 'user_fn'
USER_NAME = 'user_name'
+ USES_KEYED_STATE = 'uses_keyed_state'
VALIDATE_SINK = 'validate_sink'
VALIDATE_SOURCE = 'validate_source'
VALUE = 'value'
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 619ed54..8114190 100644
--- a/sdks/python/apache_beam/runners/dataflow/native_io/iobase.py
+++ b/sdks/python/apache_beam/runners/dataflow/native_io/iobase.py
@@ -20,16 +20,23 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
from builtins import object
+from typing import TYPE_CHECKING
+from typing import Optional
from apache_beam import pvalue
from apache_beam.io import iobase
from apache_beam.transforms import ptransform
from apache_beam.transforms.display import HasDisplayData
+if TYPE_CHECKING:
+ from apache_beam import coders
+
_LOGGER = logging.getLogger(__name__)
@@ -56,6 +63,7 @@
This class is deprecated and should not be used to define new sources.
"""
+ coder = None # type: Optional[coders.Coder]
def reader(self):
"""Returns a NativeSourceReader instance associated with this source."""
diff --git a/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py b/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py
index a0a6541..ec69ee1 100644
--- a/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py
@@ -17,6 +17,8 @@
"""Tests corresponding to Dataflow's iobase module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
@@ -183,10 +185,9 @@
def Write(self, value):
self.written_values.append(value)
- p = TestPipeline()
- sink = FakeSink()
- p | Create(['a', 'b', 'c']) | _NativeWrite(sink) # pylint: disable=expression-not-assigned
- p.run()
+ with TestPipeline() as p:
+ sink = FakeSink()
+ p | Create(['a', 'b', 'c']) | _NativeWrite(sink) # pylint: disable=expression-not-assigned
self.assertEqual(['a', 'b', 'c'], sorted(sink.written_values))
diff --git a/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py b/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
index e3e76a5..c6753ea 100644
--- a/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
+++ b/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
@@ -17,6 +17,8 @@
"""Ptransform overrides for DataflowRunner."""
+# pytype: skip-file
+
from __future__ import absolute_import
from apache_beam.pipeline import PTransformOverride
diff --git a/sdks/python/apache_beam/runners/dataflow/template_runner_test.py b/sdks/python/apache_beam/runners/dataflow/template_runner_test.py
index c152f44..9edb2db 100644
--- a/sdks/python/apache_beam/runners/dataflow/template_runner_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/template_runner_test.py
@@ -17,6 +17,8 @@
"""Unit tests for templated pipelines."""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
@@ -52,19 +54,19 @@
dummy_dir = tempfile.mkdtemp()
remote_runner = DataflowRunner()
- pipeline = Pipeline(remote_runner,
- options=PipelineOptions([
- '--dataflow_endpoint=ignored',
- '--sdk_location=' + dummy_file_name,
- '--job_name=test-job',
- '--project=test-project',
- '--staging_location=' + dummy_dir,
- '--temp_location=/dev/null',
- '--template_location=' + dummy_file_name,
- '--no_auth']))
+ with Pipeline(remote_runner,
+ options=PipelineOptions([
+ '--dataflow_endpoint=ignored',
+ '--sdk_location=' + dummy_file_name,
+ '--job_name=test-job',
+ '--project=test-project',
+ '--staging_location=' + dummy_dir,
+ '--temp_location=/dev/null',
+ '--template_location=' + dummy_file_name,
+ '--no_auth'])) as pipeline:
- pipeline | beam.Create([1, 2, 3]) | beam.Map(lambda x: x) # pylint: disable=expression-not-assigned
- pipeline.run().wait_until_finish()
+ pipeline | beam.Create([1, 2, 3]) | beam.Map(lambda x: x) # pylint: disable=expression-not-assigned
+
with open(dummy_file_name) as template_file:
saved_job_dict = json.load(template_file)
self.assertEqual(
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 2e92bef..ffffe4e 100644
--- a/sdks/python/apache_beam/runners/dataflow/test_dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/test_dataflow_runner.py
@@ -17,6 +17,8 @@
"""Wrapper of Beam runners that's built for running and verifying e2e tests."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/direct/bundle_factory.py b/sdks/python/apache_beam/runners/direct/bundle_factory.py
index d1d959e..6958fb9 100644
--- a/sdks/python/apache_beam/runners/direct/bundle_factory.py
+++ b/sdks/python/apache_beam/runners/direct/bundle_factory.py
@@ -17,6 +17,8 @@
"""A factory that creates UncommittedBundles."""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
@@ -24,6 +26,7 @@
from typing import Iterator
from typing import List
from typing import Union
+from typing import cast
from apache_beam import pvalue
from apache_beam.runners import common
@@ -143,9 +146,11 @@
or as a list of copied WindowedValues.
"""
if not self._stacked:
+ # we can safely assume self._elements contains only WindowedValues
+ elements = cast('List[WindowedValue]', self._elements)
if self._committed and not make_copy:
- return self._elements
- return list(self._elements)
+ return elements
+ return list(elements)
def iterable_stacked_or_elements(elements):
for e in elements:
diff --git a/sdks/python/apache_beam/runners/direct/clock.py b/sdks/python/apache_beam/runners/direct/clock.py
index 6dbf8b2..54b5701 100644
--- a/sdks/python/apache_beam/runners/direct/clock.py
+++ b/sdks/python/apache_beam/runners/direct/clock.py
@@ -19,6 +19,8 @@
For internal use only. No backwards compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import time
diff --git a/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor.py b/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor.py
index df3c1f8..eaa2f59 100644
--- a/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor.py
+++ b/sdks/python/apache_beam/runners/direct/consumer_tracking_pipeline_visitor.py
@@ -17,6 +17,8 @@
"""ConsumerTrackingPipelineVisitor, a PipelineVisitor object."""
+# pytype: skip-file
+
from __future__ import absolute_import
from typing import TYPE_CHECKING
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 6d21e55..3961fd4 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
@@ -16,6 +16,8 @@
#
"""Tests for consumer_tracking_pipeline_visitor."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/direct/direct_metrics.py b/sdks/python/apache_beam/runners/direct/direct_metrics.py
index 91ee03d..accf642 100644
--- a/sdks/python/apache_beam/runners/direct/direct_metrics.py
+++ b/sdks/python/apache_beam/runners/direct/direct_metrics.py
@@ -20,6 +20,8 @@
responding to queries of current metrics, but also of keeping the common
state consistent.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import threading
diff --git a/sdks/python/apache_beam/runners/direct/direct_metrics_test.py b/sdks/python/apache_beam/runners/direct/direct_metrics_test.py
index 3ce42c1..f40a508 100644
--- a/sdks/python/apache_beam/runners/direct/direct_metrics_test.py
+++ b/sdks/python/apache_beam/runners/direct/direct_metrics_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py
index ef81a27..b6902ff 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -21,6 +21,8 @@
graph of transformations belonging to a pipeline on the local machine.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import itertools
@@ -116,6 +118,10 @@
return outputs
+ # We must mark this method as not a test or else its name is a matcher for
+ # nosetest tests.
+ apply_TestStream.__test__ = False
+
def run_pipeline(self, pipeline, options):
from apache_beam.pipeline import PipelineVisitor
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner_test.py b/sdks/python/apache_beam/runners/direct/direct_runner_test.py
index 95df81d..50164b9 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner_test.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import threading
diff --git a/sdks/python/apache_beam/runners/direct/direct_userstate.py b/sdks/python/apache_beam/runners/direct/direct_userstate.py
index 42afaa3..dab4b71 100644
--- a/sdks/python/apache_beam/runners/direct/direct_userstate.py
+++ b/sdks/python/apache_beam/runners/direct/direct_userstate.py
@@ -16,6 +16,8 @@
#
"""Support for user state in the BundleBasedDirectRunner."""
+# pytype: skip-file
+
from __future__ import absolute_import
import itertools
diff --git a/sdks/python/apache_beam/runners/direct/evaluation_context.py b/sdks/python/apache_beam/runners/direct/evaluation_context.py
index e3b67c1..c004ce5 100644
--- a/sdks/python/apache_beam/runners/direct/evaluation_context.py
+++ b/sdks/python/apache_beam/runners/direct/evaluation_context.py
@@ -17,6 +17,8 @@
"""EvaluationContext tracks global state, triggers and watermarks."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/runners/direct/executor.py b/sdks/python/apache_beam/runners/direct/executor.py
index f961952..e1a4692 100644
--- a/sdks/python/apache_beam/runners/direct/executor.py
+++ b/sdks/python/apache_beam/runners/direct/executor.py
@@ -17,6 +17,8 @@
"""An executor that schedules and executes applied ptransforms."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/runners/direct/helper_transforms.py b/sdks/python/apache_beam/runners/direct/helper_transforms.py
index 2cdff58..da16d94 100644
--- a/sdks/python/apache_beam/runners/direct/helper_transforms.py
+++ b/sdks/python/apache_beam/runners/direct/helper_transforms.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/runners/direct/sdf_direct_runner.py b/sdks/python/apache_beam/runners/direct/sdf_direct_runner.py
index e54224e..7659a65 100644
--- a/sdks/python/apache_beam/runners/direct/sdf_direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/sdf_direct_runner.py
@@ -18,6 +18,8 @@
"""This module contains Splittable DoFn logic that is specific to DirectRunner.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import uuid
@@ -259,9 +261,11 @@
self._restriction_tag = _ValueStateTag('restriction')
self.watermark_hold_tag = _ValueStateTag('watermark_hold')
self._process_element_invoker = None
+ self._output_processor = _OutputProcessor()
self.sdf_invoker = DoFnInvoker.create_invoker(
DoFnSignature(self.sdf), context=DoFnContext('unused_context'),
+ output_processor=self._output_processor,
input_args=args_for_invoker, input_kwargs=kwargs_for_invoker)
self._step_context = None
@@ -327,7 +331,8 @@
SDFProcessElementInvoker)
output_values = self._process_element_invoker.invoke_process_element(
- self.sdf_invoker, windowed_element, tracker, *args, **kwargs)
+ self.sdf_invoker, self._output_processor, windowed_element, tracker,
+ *args, **kwargs)
sdf_result = None
for output in output_values:
@@ -422,7 +427,7 @@
raise ValueError
def invoke_process_element(
- self, sdf_invoker, element, tracker, *args, **kwargs):
+ self, sdf_invoker, output_processor, element, tracker, *args, **kwargs):
"""Invokes `process()` method of a Splittable `DoFn` for a given element.
Args:
@@ -451,10 +456,10 @@
checkpoint_state.residual_restriction = tracker.checkpoint()
checkpoint_state.checkpointed = object()
- output_processor = _OutputProcessor()
+ output_processor.reset()
Timer(self._max_duration, initiate_checkpoint).start()
sdf_invoker.invoke_process(
- element, restriction_tracker=tracker, output_processor=output_processor,
+ element, restriction_tracker=tracker,
additional_args=args, additional_kwargs=kwargs)
assert output_processor.output_iter is not None
@@ -503,3 +508,6 @@
def process_outputs(self, windowed_input_element, output_iter):
# type: (WindowedValue, Iterable[Any]) -> None
self.output_iter = output_iter
+
+ def reset(self):
+ self.output_iter = None
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 d9d68cc..40c761e 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
@@ -17,6 +17,8 @@
"""Unit tests for SDF implementation for DirectRunner."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
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 04dbe50..81f5870 100644
--- a/sdks/python/apache_beam/runners/direct/test_direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/test_direct_runner.py
@@ -17,6 +17,8 @@
"""Wrapper of Beam runners that's built for running and verifying e2e tests."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/direct/test_stream_impl.py b/sdks/python/apache_beam/runners/direct/test_stream_impl.py
index aa1405d..c1930fa 100644
--- a/sdks/python/apache_beam/runners/direct/test_stream_impl.py
+++ b/sdks/python/apache_beam/runners/direct/test_stream_impl.py
@@ -23,6 +23,8 @@
tagged PCollection.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from apache_beam import coders
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
index ff73aab..71fd133 100644
--- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
@@ -17,6 +17,8 @@
"""An evaluator of a specific application of a transform."""
+# pytype: skip-file
+
from __future__ import absolute_import
import atexit
@@ -584,6 +586,7 @@
bundles = [bundle]
else:
bundles = []
+ assert self._applied_ptransform.transform is not None
if self._applied_ptransform.inputs:
input_pvalue = self._applied_ptransform.inputs[0] # type: Union[pvalue.PBegin, pvalue.PCollection]
else:
diff --git a/sdks/python/apache_beam/runners/direct/util.py b/sdks/python/apache_beam/runners/direct/util.py
index 57650ac..478a981 100644
--- a/sdks/python/apache_beam/runners/direct/util.py
+++ b/sdks/python/apache_beam/runners/direct/util.py
@@ -20,6 +20,8 @@
For internal use only. No backwards compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/runners/direct/watermark_manager.py b/sdks/python/apache_beam/runners/direct/watermark_manager.py
index 6288a87..8856ee5 100644
--- a/sdks/python/apache_beam/runners/direct/watermark_manager.py
+++ b/sdks/python/apache_beam/runners/direct/watermark_manager.py
@@ -17,6 +17,8 @@
"""Manages watermarks of PCollections and AppliedPTransforms."""
+# pytype: skip-file
+
from __future__ import absolute_import
import threading
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 f157f9b..d46317a 100644
--- a/sdks/python/apache_beam/runners/interactive/background_caching_job.py
+++ b/sdks/python/apache_beam/runners/interactive/background_caching_job.py
@@ -35,6 +35,8 @@
invalidated.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import apache_beam as beam
diff --git a/sdks/python/apache_beam/runners/interactive/background_caching_job_test.py b/sdks/python/apache_beam/runners/interactive/background_caching_job_test.py
index 19e4889..409d338 100644
--- a/sdks/python/apache_beam/runners/interactive/background_caching_job_test.py
+++ b/sdks/python/apache_beam/runners/interactive/background_caching_job_test.py
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.interactive.background_caching_job."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/interactive/cache_manager.py b/sdks/python/apache_beam/runners/interactive/cache_manager.py
index 02926aa..1531e6d 100644
--- a/sdks/python/apache_beam/runners/interactive/cache_manager.py
+++ b/sdks/python/apache_beam/runners/interactive/cache_manager.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/interactive/cache_manager_test.py b/sdks/python/apache_beam/runners/interactive/cache_manager_test.py
index f097cf5..6ab51b4 100644
--- a/sdks/python/apache_beam/runners/interactive/cache_manager_test.py
+++ b/sdks/python/apache_beam/runners/interactive/cache_manager_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
index 2348c78..d401cd0 100644
--- a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
+++ b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from apache_beam.portability.api.beam_runner_api_pb2 import TestStreamPayload
diff --git a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py
index 5b65992..1fb2d7c 100644
--- a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py
+++ b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/interactive/display/display_manager.py b/sdks/python/apache_beam/runners/interactive/display/display_manager.py
index 79bb010..5cd27f1 100644
--- a/sdks/python/apache_beam/runners/interactive/display/display_manager.py
+++ b/sdks/python/apache_beam/runners/interactive/display/display_manager.py
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
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 c482561..aec3b36 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
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization.py b/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization.py
index bdd34ca..9f98f7c 100644
--- a/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization.py
+++ b/sdks/python/apache_beam/runners/interactive/display/pcoll_visualization.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
Only works with Python 3.5+.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import base64
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 b2d330e..28411e2 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
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.interactive.display.pcoll_visualization."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
diff --git a/sdks/python/apache_beam/runners/interactive/display/pipeline_graph.py b/sdks/python/apache_beam/runners/interactive/display/pipeline_graph.py
index 6063a72..b2d9b3a 100644
--- a/sdks/python/apache_beam/runners/interactive/display/pipeline_graph.py
+++ b/sdks/python/apache_beam/runners/interactive/display/pipeline_graph.py
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_renderer.py b/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_renderer.py
index 1ada4e8..5bc637e 100644
--- a/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_renderer.py
+++ b/sdks/python/apache_beam/runners/interactive/display/pipeline_graph_renderer.py
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
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 7e4aa13..780e70c 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
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.interactive.display.pipeline_graph."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_beam.py b/sdks/python/apache_beam/runners/interactive/interactive_beam.py
index a7a7584..d35c1d8 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_beam.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_beam.py
@@ -28,6 +28,9 @@
Note: If you want backward-compatibility, only invoke interfaces provided by
this module in your notebook or application code.
"""
+
+# pytype: skip-file
+
from __future__ import absolute_import
from apache_beam.runners.interactive import interactive_environment as ie
@@ -62,10 +65,9 @@
class Foo(object)
def run_pipeline(self):
- p = beam.Pipeline()
- init_pcoll = p | 'Init Create' >> beam.Create(range(10))
- watch(locals())
- p.run()
+ with beam.Pipeline() as p:
+ init_pcoll = p | 'Init Create' >> beam.Create(range(10))
+ watch(locals())
return init_pcoll
init_pcoll = Foo().run_pipeline()
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 7660b1a..720f443 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.interactive.interactive_beam."""
+# pytype: skip-file
+
from __future__ import absolute_import
import importlib
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_environment.py b/sdks/python/apache_beam/runners/interactive/interactive_environment.py
index 31bb736..d5f850d 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_environment.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_environment.py
@@ -22,6 +22,8 @@
External Interactive Beam users please use interactive_beam module in
application code or notebook.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import atexit
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 ca9ac68..85f546d 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_environment_test.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_environment_test.py
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.interactive.interactive_environment."""
+# pytype: skip-file
+
from __future__ import absolute_import
import importlib
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_runner.py b/sdks/python/apache_beam/runners/interactive/interactive_runner.py
index fd2b236..5e05cb5 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_runner.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_runner.py
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
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 9b5abf3..9ef8b13 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_runner_test.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_runner_test.py
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_analyzer.py b/sdks/python/apache_beam/runners/interactive/pipeline_analyzer.py
index ab56921..bc72465 100644
--- a/sdks/python/apache_beam/runners/interactive/pipeline_analyzer.py
+++ b/sdks/python/apache_beam/runners/interactive/pipeline_analyzer.py
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_analyzer_test.py b/sdks/python/apache_beam/runners/interactive/pipeline_analyzer_test.py
index b0433ff..9aaeaba 100644
--- a/sdks/python/apache_beam/runners/interactive/pipeline_analyzer_test.py
+++ b/sdks/python/apache_beam/runners/interactive/pipeline_analyzer_test.py
@@ -20,6 +20,8 @@
This module is experimental. No backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_instrument.py b/sdks/python/apache_beam/runners/interactive/pipeline_instrument.py
index 084e2ac..01224c8 100644
--- a/sdks/python/apache_beam/runners/interactive/pipeline_instrument.py
+++ b/sdks/python/apache_beam/runners/interactive/pipeline_instrument.py
@@ -21,6 +21,8 @@
This module accesses current interactive environment and analyzes given pipeline
to transform original pipeline into a one-shot pipeline with interactivity.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import apache_beam as beam
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 f06c271..eab2172 100644
--- a/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py
+++ b/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.interactive.pipeline_instrument."""
+# pytype: skip-file
+
from __future__ import absolute_import
import tempfile
@@ -288,7 +290,7 @@
ib.watch({'irrelevant_user_pipeline': irrelevant_user_pipeline})
# Build instrument from the runner pipeline.
pipeline_instrument = instr.pin(runner_pipeline)
- self.assertTrue(pipeline_instrument.user_pipeline is user_pipeline)
+ self.assertIs(pipeline_instrument.user_pipeline, user_pipeline)
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/runners/job/manager.py b/sdks/python/apache_beam/runners/job/manager.py
index 991d2fe..ae8eb94 100644
--- a/sdks/python/apache_beam/runners/job/manager.py
+++ b/sdks/python/apache_beam/runners/job/manager.py
@@ -18,6 +18,8 @@
"""A object to control to the Job API Co-Process
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/job/utils.py b/sdks/python/apache_beam/runners/job/utils.py
index 4c7c965..1a90adb 100644
--- a/sdks/python/apache_beam/runners/job/utils.py
+++ b/sdks/python/apache_beam/runners/job/utils.py
@@ -18,6 +18,8 @@
"""Utility functions for efficiently processing with the job API
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
diff --git a/sdks/python/apache_beam/runners/pipeline_context.py b/sdks/python/apache_beam/runners/pipeline_context.py
index d8821bc..6b6c3a0 100644
--- a/sdks/python/apache_beam/runners/pipeline_context.py
+++ b/sdks/python/apache_beam/runners/pipeline_context.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
@@ -134,14 +136,6 @@
Used for accessing and constructing the referenced objects of a Pipeline.
"""
- _COMPONENT_TYPES = {
- 'transforms': pipeline.AppliedPTransform,
- 'pcollections': pvalue.PCollection,
- 'coders': coders.Coder,
- 'windowing_strategies': core.Windowing,
- 'environments': environments.Environment,
- }
-
def __init__(self,
proto=None, # type: Optional[Union[beam_runner_api_pb2.Components, beam_fn_api_pb2.ProcessBundleDescriptor]]
default_environment=None, # type: Optional[environments.Environment]
@@ -156,13 +150,26 @@
coders=dict(proto.coders.items()),
windowing_strategies=dict(proto.windowing_strategies.items()),
environments=dict(proto.environments.items()))
- for name, cls in self._COMPONENT_TYPES.items():
- setattr(
- self, name, _PipelineContextMap(
- self, cls, namespace, getattr(proto, name, None)))
+
+ self.transforms = _PipelineContextMap(
+ self, pipeline.AppliedPTransform, namespace,
+ proto.transforms if proto is not None else None)
+ self.pcollections = _PipelineContextMap(
+ self, pvalue.PCollection, namespace,
+ proto.pcollections if proto is not None else None)
+ self.coders = _PipelineContextMap(
+ self, coders.Coder, namespace,
+ proto.coders if proto is not None else None)
+ self.windowing_strategies = _PipelineContextMap(
+ self, core.Windowing, namespace,
+ proto.windowing_strategies if proto is not None else None)
+ self.environments = _PipelineContextMap(
+ self, environments.Environment, namespace,
+ proto.environments if proto is not None else None)
+
if default_environment:
self._default_environment_id = self.environments.get_id(
- default_environment, label='default_environment')
+ default_environment, label='default_environment') # type: Optional[str]
else:
self._default_environment_id = None
self.use_fake_coders = use_fake_coders
@@ -177,7 +184,7 @@
def coder_id_from_element_type(self, element_type):
# type: (Any) -> str
if self.use_fake_coders:
- return pickler.dumps(element_type)
+ return pickler.dumps(element_type).decode('ascii')
else:
return self.coders.get_id(coders.registry.get_coder(element_type))
@@ -197,8 +204,13 @@
def to_runner_api(self):
# type: () -> beam_runner_api_pb2.Components
context_proto = beam_runner_api_pb2.Components()
- for name in self._COMPONENT_TYPES:
- getattr(self, name).populate_map(getattr(context_proto, name))
+
+ self.transforms.populate_map(context_proto.transforms)
+ self.pcollections.populate_map(context_proto.pcollections)
+ self.coders.populate_map(context_proto.coders)
+ self.windowing_strategies.populate_map(context_proto.windowing_strategies)
+ self.environments.populate_map(context_proto.environments)
+
return context_proto
def default_environment_id(self):
diff --git a/sdks/python/apache_beam/runners/pipeline_context_test.py b/sdks/python/apache_beam/runners/pipeline_context_test.py
index 6f1ec74..697ea6d 100644
--- a/sdks/python/apache_beam/runners/pipeline_context_test.py
+++ b/sdks/python/apache_beam/runners/pipeline_context_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the windowing classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
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 df7f9d22..8a97b73 100644
--- a/sdks/python/apache_beam/runners/portability/abstract_job_service.py
+++ b/sdks/python/apache_beam/runners/portability/abstract_job_service.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import itertools
@@ -29,6 +31,7 @@
from typing import Dict
from typing import Iterator
from typing import Optional
+from typing import Tuple
from typing import Union
import grpc
@@ -48,6 +51,7 @@
_LOGGER = logging.getLogger(__name__)
+StateEvent = Tuple[int, Union[timestamp_pb2.Timestamp, Timestamp]]
def make_state_event(state, timestamp):
if isinstance(timestamp, Timestamp):
@@ -127,7 +131,7 @@
request, # type: beam_job_api_pb2.GetJobStateRequest
context=None
):
- # type: (...) -> beam_job_api_pb2.GetJobStateResponse
+ # type: (...) -> beam_job_api_pb2.JobStateEvent
return beam_job_api_pb2.JobStateEvent(
state=self._jobs[request.job_id].get_state())
@@ -151,7 +155,7 @@
state=self._jobs[request.job_id].get_state())
def GetStateStream(self, request, context=None, timeout=None):
- # type: (...) -> Iterator[beam_job_api_pb2.GetJobStateResponse]
+ # type: (...) -> Iterator[beam_job_api_pb2.JobStateEvent]
"""Yields state transitions since the stream started.
"""
if request.job_id not in self._jobs:
@@ -216,11 +220,11 @@
raise NotImplementedError(self)
def get_state_stream(self):
- # type: () -> Iterator[Optional[beam_job_api_pb2.JobState.Enum]]
+ # type: () -> Iterator[StateEvent]
raise NotImplementedError(self)
def get_message_stream(self):
- # type: () -> Iterator[Union[int, Optional[beam_job_api_pb2.JobMessage]]]
+ # type: () -> Iterator[Union[StateEvent, Optional[beam_job_api_pb2.JobMessage]]]
raise NotImplementedError(self)
diff --git a/sdks/python/apache_beam/runners/portability/artifact_service.py b/sdks/python/apache_beam/runners/portability/artifact_service.py
index e7da20a..708aceb 100644
--- a/sdks/python/apache_beam/runners/portability/artifact_service.py
+++ b/sdks/python/apache_beam/runners/portability/artifact_service.py
@@ -18,6 +18,8 @@
The staging service here can be backed by any beam filesystem.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -148,7 +150,7 @@
with self._open(artifact.uri, 'r') as fin:
# This value is not emitted, but lets us yield a single empty
# chunk on an empty file.
- chunk = True
+ chunk = b'1'
while chunk:
chunk = fin.read(self._chunk_size)
yield beam_artifact_api_pb2.ArtifactChunk(data=chunk)
diff --git a/sdks/python/apache_beam/runners/portability/artifact_service_test.py b/sdks/python/apache_beam/runners/portability/artifact_service_test.py
index 6efb60d..9c59a6b 100644
--- a/sdks/python/apache_beam/runners/portability/artifact_service_test.py
+++ b/sdks/python/apache_beam/runners/portability/artifact_service_test.py
@@ -15,6 +15,8 @@
#
"""Test cases for :module:`artifact_service_client`."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/expansion_service.py b/sdks/python/apache_beam/runners/portability/expansion_service.py
index 55a526d..d47e59e 100644
--- a/sdks/python/apache_beam/runners/portability/expansion_service.py
+++ b/sdks/python/apache_beam/runners/portability/expansion_service.py
@@ -17,6 +17,8 @@
"""A PipelineExpansion service.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/expansion_service_test.py b/sdks/python/apache_beam/runners/portability/expansion_service_test.py
index 8426311..2919d2f 100644
--- a/sdks/python/apache_beam/runners/portability/expansion_service_test.py
+++ b/sdks/python/apache_beam/runners/portability/expansion_service_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/runners/portability/flink_runner.py b/sdks/python/apache_beam/runners/portability/flink_runner.py
index 4941cec..8096a3a 100644
--- a/sdks/python/apache_beam/runners/portability/flink_runner.py
+++ b/sdks/python/apache_beam/runners/portability/flink_runner.py
@@ -17,6 +17,8 @@
"""A runner for executing portable pipelines on Flink."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
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 5102290..7d750df 100644
--- a/sdks/python/apache_beam/runners/portability/flink_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/flink_runner_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
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 8a035f8..5f0ae78 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
@@ -17,6 +17,8 @@
"""A job server submitting portable pipelines as uber jars to Flink."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py
index 2c24cbd..ca5bfef 100644
--- a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py
+++ b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner.py b/sdks/python/apache_beam/runners/portability/fn_api_runner.py
index 7c45847..c5e1dac 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner.py
@@ -17,6 +17,8 @@
"""A PipelineRunner using the SDK harness.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
@@ -48,6 +50,7 @@
from typing import Type
from typing import TypeVar
from typing import Union
+from typing import overload
import grpc
@@ -112,7 +115,8 @@
DataSideInput = Dict[Tuple[str, str],
Tuple[bytes, beam_runner_api_pb2.FunctionSpec]]
DataOutput = Dict[str, bytes]
-BundleProcessResult = Tuple[beam_fn_api_pb2.InstructionResponse, List[beam_fn_api_pb2.ProcessBundleSplitResponse]]
+BundleProcessResult = Tuple[beam_fn_api_pb2.InstructionResponse,
+ List[beam_fn_api_pb2.ProcessBundleSplitResponse]]
# This module is experimental. No backwards-compatibility guarantees.
@@ -149,8 +153,17 @@
for data in self._input:
self._futures_by_id.pop(data.instruction_id).set(data)
+ @overload
def push(self, req):
- # type: (...) -> Optional[ControlFuture]
+ # type: (BeamFnControlServicer.DoneMarker) -> None
+ pass
+
+ @overload
+ def push(self, req):
+ # type: (beam_fn_api_pb2.InstructionRequest) -> ControlFuture
+ pass
+
+ def push(self, req):
if req == BeamFnControlServicer._DONE_MARKER:
self._push_queue.put(req)
return None
@@ -192,7 +205,10 @@
STARTED_STATE = 'started'
DONE_STATE = 'done'
- _DONE_MARKER = object()
+ class DoneMarker(object):
+ pass
+
+ _DONE_MARKER = DoneMarker()
def __init__(self):
self._lock = threading.Lock()
@@ -364,8 +380,9 @@
# type: (bytes) -> None
input_stream = create_InputStream(elements_data)
while input_stream.size() > 0:
- windowed_value = self._windowed_value_coder.get_impl(
- ).decode_from_stream(input_stream, True)
+ windowed_val_coder_impl = self._windowed_value_coder.get_impl() # type: WindowedValueCoderImpl
+ windowed_value = windowed_val_coder_impl.decode_from_stream(
+ input_stream, True)
key, value = self._kv_extractor(windowed_value.value)
for window in windowed_value.windows:
self._values_by_window[key, window].append(value)
@@ -387,7 +404,7 @@
def __init__(
self,
- default_environment=None, # type: Optional[beam_runner_api_pb2.Environment]
+ default_environment=None, # type: Optional[environments.Environment]
bundle_repeat=0,
use_state_iterables=False,
provision_info=None, # type: Optional[ExtendedProvisionInfo]
@@ -449,6 +466,19 @@
pipeline_options.DirectOptions).direct_runner_bundle_repeat
self._num_workers = options.view_as(
pipeline_options.DirectOptions).direct_num_workers or self._num_workers
+
+ # set direct workers running mode if it is defined with pipeline options.
+ running_mode = \
+ options.view_as(pipeline_options.DirectOptions).direct_running_mode
+ if running_mode == 'multi_threading':
+ self._default_environment = environments.EmbeddedPythonGrpcEnvironment()
+ elif running_mode == 'multi_processing':
+ command_string = '%s -m apache_beam.runners.worker.sdk_worker_main' \
+ % sys.executable
+ self._default_environment = environments.SubprocessSDKEnvironment(
+ command_string=command_string
+ )
+
self._profiler_factory = profiler.Profile.factory_from_options(
options.view_as(pipeline_options.ProfilingOptions))
@@ -803,9 +833,10 @@
pipeline_components.windowing_strategies.items()),
environments=dict(pipeline_components.environments.items()))
- if worker_handler.state_api_service_descriptor():
+ state_api_service_descriptor = worker_handler.state_api_service_descriptor()
+ if state_api_service_descriptor:
process_bundle_descriptor.state_api_service_descriptor.url = (
- worker_handler.state_api_service_descriptor().url)
+ state_api_service_descriptor.url)
# Store the required side inputs into state so it is accessible for the
# worker when it runs this bundle.
@@ -914,7 +945,8 @@
# merged results. Without residual_roots, pipeline stops earlier and we
# may miss some data.
bundle_manager._num_workers = 1
- bundle_manager._skip_registration = True
+ # TODO(BEAM-8486): this should be changed to _registered
+ bundle_manager._skip_registration = True # type: ignore[attr-defined]
last_result, splits = bundle_manager.process_bundle(
deferred_inputs, data_output)
last_sent = deferred_inputs
@@ -991,7 +1023,8 @@
# These classes are used to interact with the worker.
- class StateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer):
+ class StateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer,
+ sdk_worker.StateHandler):
class CopyOnWriteState(object):
def __init__(self, underlying):
@@ -1073,11 +1106,11 @@
else:
token_base, index = continuation_token.split(':')
ix = int(index)
- full_state = self._continuations[token_base]
- if ix == len(full_state):
+ full_state_cont = self._continuations[token_base]
+ if ix == len(full_state_cont):
return b'', None
else:
- return full_state[ix], '%s:%d' % (token_base, ix + 1)
+ return full_state_cont[ix], '%s:%d' % (token_base, ix + 1)
else:
assert not continuation_token
return b''.join(full_state), None
@@ -1147,16 +1180,16 @@
"""A singleton cache for a StateServicer."""
def __init__(self, state_handler):
- # type: (sdk_worker.StateHandler) -> None
+ # type: (sdk_worker.CachingStateHandler) -> None
self._state_handler = state_handler
def create_state_handler(self, api_service_descriptor):
- # type: (endpoints_pb2.ApiServiceDescriptor) -> sdk_worker.StateHandler
+ # type: (endpoints_pb2.ApiServiceDescriptor) -> sdk_worker.CachingStateHandler
"""Returns the singleton state handler."""
return self._state_handler
def close(self):
- # type: (...) -> None
+ # type: () -> None
"""Does nothing."""
pass
@@ -1213,6 +1246,9 @@
_worker_id_counter = -1
_lock = threading.Lock()
+ control_conn = None # type: ControlConnection
+ data_conn = None # type: data_plane._GrpcDataChannel
+
def __init__(self,
control_handler,
data_plane_handler,
@@ -1294,7 +1330,7 @@
def __init__(self,
unused_payload, # type: None
- state,
+ state, # type: sdk_worker.StateHandler
provision_info, # type: Optional[ExtendedProvisionInfo]
unused_grpc_server=None
):
@@ -1428,10 +1464,10 @@
# If we have provision info, serve these off the control port as well.
if self.provision_info:
if self.provision_info.provision_info:
- provision_info = self.provision_info.provision_info
- if not provision_info.worker_id:
- provision_info = copy.copy(provision_info)
- provision_info.worker_id = str(uuid.uuid4())
+ provision_info_proto = self.provision_info.provision_info
+ if not provision_info_proto.worker_id:
+ provision_info_proto = copy.copy(provision_info_proto)
+ provision_info_proto.worker_id = str(uuid.uuid4())
beam_provision_api_pb2_grpc.add_ProvisionServiceServicer_to_server(
BasicProvisionService(self.provision_info.provision_info),
self.control_server)
@@ -1856,6 +1892,7 @@
read_transform_id, # type: str
byte_streams
):
+ assert self._worker_handler is not None
data_out = self._worker_handler.data_conn.output_stream(
process_bundle_id, read_transform_id)
for byte_stream in byte_streams:
@@ -1867,6 +1904,7 @@
if self._registered:
registration_future = None
else:
+ assert self._worker_handler is not None
process_bundle_registration = beam_fn_api_pb2.InstructionRequest(
register=beam_fn_api_pb2.RegisterRequest(
process_bundle_descriptor=[self._bundle_descriptor]))
@@ -1914,6 +1952,8 @@
self._send_input_to_worker(
process_bundle_id, read_transform_id, [byte_stream])
+ assert self._worker_handler is not None
+
# Execute the requested splits.
while not done:
if split_fraction is None:
@@ -2051,13 +2091,18 @@
merged_result = None # type: Optional[beam_fn_api_pb2.InstructionResponse]
split_result_list = [] # type: List[beam_fn_api_pb2.ProcessBundleSplitResponse]
- with UnboundedThreadPoolExecutor() as executor:
- for result, split_result in executor.map(lambda part: BundleManager(
+
+ def execute(part_map):
+ # type: (...) -> BundleProcessResult
+ bundle_manager = BundleManager(
self._worker_handler_list, self._get_buffer,
self._get_input_coder_impl, self._bundle_descriptor,
self._progress_frequency, self._registered,
- cache_token_generator=self._cache_token_generator).process_bundle(
- part, expected_outputs), part_inputs):
+ cache_token_generator=self._cache_token_generator)
+ return bundle_manager.process_bundle(part_map, expected_outputs)
+
+ with UnboundedThreadPoolExecutor() as executor:
+ for result, split_result in executor.map(execute, part_inputs):
split_result_list += split_result
if merged_result is None:
@@ -2070,6 +2115,7 @@
result.process_bundle.monitoring_infos,
merged_result.process_bundle.monitoring_infos))),
error=result.error or merged_result.error)
+ assert merged_result is not None
return merged_result, split_result_list
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
index e678f65..e82f31c 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
@@ -564,6 +566,12 @@
p | 'd' >> beam.Create(additional)) | beam.Flatten()
assert_that(res, equal_to(['a', 'b', 'c'] + additional))
+ def test_flatten_same_pcollections(self, with_transcoding=True):
+ with self.create_pipeline() as p:
+ pc = p | beam.Create(['a', 'b'])
+ assert_that((pc, pc, pc) | beam.Flatten(), equal_to(['a', 'b'] * 3))
+
+
def test_combine_per_key(self):
with self.create_pipeline() as p:
res = (p
@@ -1183,10 +1191,10 @@
class FnApiRunnerTestWithGrpcAndMultiWorkers(FnApiRunnerTest):
def create_pipeline(self):
- pipeline_options = PipelineOptions(direct_num_workers=2)
+ pipeline_options = PipelineOptions(direct_num_workers=2,
+ direct_running_mode='multi_threading')
p = beam.Pipeline(
- runner=fn_api_runner.FnApiRunner(
- default_environment=environments.EmbeddedPythonGrpcEnvironment()),
+ runner=fn_api_runner.FnApiRunner(),
options=pipeline_options)
#TODO(BEAM-8444): Fix these tests..
p.options.view_as(DebugOptions).experiments.remove('beam_fn_api')
@@ -1543,10 +1551,10 @@
class FnApiRunnerSplitTestWithMultiWorkers(FnApiRunnerSplitTest):
def create_pipeline(self):
- pipeline_options = PipelineOptions(direct_num_workers=2)
+ pipeline_options = PipelineOptions(direct_num_workers=2,
+ direct_running_mode='multi_threading')
p = beam.Pipeline(
- runner=fn_api_runner.FnApiRunner(
- default_environment=environments.EmbeddedPythonGrpcEnvironment()),
+ runner=fn_api_runner.FnApiRunner(),
options=pipeline_options)
#TODO(BEAM-8444): Fix these tests..
p.options.view_as(DebugOptions).experiments.remove('beam_fn_api')
@@ -1586,7 +1594,7 @@
self.assertRegex(
''.join(logs.output),
- '.*There has been a processing lull of over.*',
+ '.*Operation ongoing for over.*',
'Unable to find a lull logged for this job.')
class StateBackedTestElementType(object):
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner_transforms.py b/sdks/python/apache_beam/runners/portability/fn_api_runner_transforms.py
index 6ff8689..d206122 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner_transforms.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner_transforms.py
@@ -17,6 +17,8 @@
"""Pipeline transformations for the FnApiRunner.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
@@ -275,7 +277,7 @@
stage_components.transforms[side.transform_id].inputs[side.local_name]
for side in side_inputs
}, main_input=main_input_id)
- payload = beam_runner_api_pb2.ExecutableStagePayload(
+ exec_payload = beam_runner_api_pb2.ExecutableStagePayload(
environment=components.environments[self.environment],
input=main_input_id,
outputs=external_outputs,
@@ -289,7 +291,7 @@
unique_name=unique_name(None, self.name),
spec=beam_runner_api_pb2.FunctionSpec(
urn='beam:runner:executable_stage:v1',
- payload=payload.SerializeToString()),
+ payload=exec_payload.SerializeToString()),
inputs=named_inputs,
outputs={'output_%d' % ix: pcoll
for ix, pcoll in enumerate(external_outputs)},)
@@ -315,7 +317,8 @@
class TransformContext(object):
_KNOWN_CODER_URNS = set(
- value.urn for value in common_urns.coders.__dict__.values())
+ value.urn for key, value in common_urns.coders.__dict__.items()
+ if not key.startswith('_'))
def __init__(self,
components, # type: beam_runner_api_pb2.Components
@@ -578,15 +581,18 @@
This representation is also amenable to simple recomputation on fusion.
"""
consumers = collections.defaultdict(list) # type: DefaultDict[str, List[Stage]]
+ def get_all_side_inputs():
+ # type: () -> Set[str]
+ all_side_inputs = set() # type: Set[str]
+ for stage in stages:
+ for transform in stage.transforms:
+ for input in transform.inputs.values():
+ consumers[input].append(stage)
+ for si in stage.side_inputs():
+ all_side_inputs.add(si)
+ return all_side_inputs
- all_side_inputs = set()
- for stage in stages:
- for transform in stage.transforms:
- for input in transform.inputs.values():
- consumers[input].append(stage)
- for si in stage.side_inputs():
- all_side_inputs.add(si)
- all_side_inputs = frozenset(all_side_inputs)
+ all_side_inputs = frozenset(get_all_side_inputs())
downstream_side_inputs_by_stage = {} # type: Dict[Stage, FrozenSet[str]]
@@ -1438,4 +1444,5 @@
def split_buffer_id(buffer_id):
# type: (bytes) -> Tuple[str, str]
"""A buffer id is "kind:pcollection_id". Split into (kind, pcoll_id). """
- return buffer_id.decode('utf-8').split(':', 1)
+ kind, pcoll_id = buffer_id.decode('utf-8').split(':', 1)
+ return kind, pcoll_id
diff --git a/sdks/python/apache_beam/runners/portability/job_server.py b/sdks/python/apache_beam/runners/portability/job_server.py
index 8e0b396..66d1558 100644
--- a/sdks/python/apache_beam/runners/portability/job_server.py
+++ b/sdks/python/apache_beam/runners/portability/job_server.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import atexit
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 0e22d06..7fecdc7 100644
--- a/sdks/python/apache_beam/runners/portability/local_job_service.py
+++ b/sdks/python/apache_beam/runners/portability/local_job_service.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/portability/local_job_service_main.py b/sdks/python/apache_beam/runners/portability/local_job_service_main.py
index aa33263..ac5a178 100644
--- a/sdks/python/apache_beam/runners/portability/local_job_service_main.py
+++ b/sdks/python/apache_beam/runners/portability/local_job_service_main.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/runners/portability/local_job_service_test.py b/sdks/python/apache_beam/runners/portability/local_job_service_test.py
index 475d2fe..4ff382f 100644
--- a/sdks/python/apache_beam/runners/portability/local_job_service_test.py
+++ b/sdks/python/apache_beam/runners/portability/local_job_service_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/portable_metrics.py b/sdks/python/apache_beam/runners/portability/portable_metrics.py
index 0c88b73..956d4d2 100644
--- a/sdks/python/apache_beam/runners/portability/portable_metrics.py
+++ b/sdks/python/apache_beam/runners/portability/portable_metrics.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/portability/portable_runner.py b/sdks/python/apache_beam/runners/portability/portable_runner.py
index 0221068..b5dd1d4 100644
--- a/sdks/python/apache_beam/runners/portability/portable_runner.py
+++ b/sdks/python/apache_beam/runners/portability/portable_runner.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import functools
@@ -86,7 +88,7 @@
@staticmethod
def _create_environment(options):
- # type: (PipelineOptions) -> beam_runner_api_pb2.Environment
+ # type: (PipelineOptions) -> environments.Environment
portable_options = options.view_as(PortableOptions)
# Do not set a Runner. Otherwise this can cause problems in Java's
# PipelineOptions, i.e. ClassNotFoundException, if the corresponding Runner
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 81e8cce..49b4231 100644
--- a/sdks/python/apache_beam/runners/portability/portable_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/portable_runner_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/portable_stager.py b/sdks/python/apache_beam/runners/portability/portable_stager.py
index 328025e..b9447fe 100644
--- a/sdks/python/apache_beam/runners/portability/portable_stager.py
+++ b/sdks/python/apache_beam/runners/portability/portable_stager.py
@@ -16,6 +16,8 @@
"""A :class:`FileHandler` to work with :class:`ArtifactStagingServiceStub`.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/portable_stager_test.py b/sdks/python/apache_beam/runners/portability/portable_stager_test.py
index fd86819..77ffda5 100644
--- a/sdks/python/apache_beam/runners/portability/portable_stager_test.py
+++ b/sdks/python/apache_beam/runners/portability/portable_stager_test.py
@@ -15,6 +15,8 @@
#
"""Test cases for :module:`artifact_service_client`."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/spark_runner.py b/sdks/python/apache_beam/runners/portability/spark_runner.py
index ce33acd..608d2b6 100644
--- a/sdks/python/apache_beam/runners/portability/spark_runner.py
+++ b/sdks/python/apache_beam/runners/portability/spark_runner.py
@@ -17,6 +17,8 @@
"""A runner for executing portable pipelines on Spark."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
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 1ac5e6c..fa7e795 100644
--- a/sdks/python/apache_beam/runners/portability/spark_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/spark_runner_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
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 88f7cb7..146ab44 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
@@ -17,6 +17,8 @@
"""A job server submitting portable pipelines as uber jars to Spark."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
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 7106403..e999480 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
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/portability/stager.py b/sdks/python/apache_beam/runners/portability/stager.py
index 786e4f7..0f989dc 100644
--- a/sdks/python/apache_beam/runners/portability/stager.py
+++ b/sdks/python/apache_beam/runners/portability/stager.py
@@ -44,6 +44,8 @@
TODO(silviuc): Should we allow several setup packages?
TODO(silviuc): We should allow customizing the exact command for setup build.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import glob
diff --git a/sdks/python/apache_beam/runners/portability/stager_test.py b/sdks/python/apache_beam/runners/portability/stager_test.py
index 2a39f08..180ba8a 100644
--- a/sdks/python/apache_beam/runners/portability/stager_test.py
+++ b/sdks/python/apache_beam/runners/portability/stager_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the stager module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/runner.py b/sdks/python/apache_beam/runners/runner.py
index 6912902..bb74e62 100644
--- a/sdks/python/apache_beam/runners/runner.py
+++ b/sdks/python/apache_beam/runners/runner.py
@@ -17,6 +17,8 @@
"""PipelineRunner, an abstract base runner object."""
+# pytype: skip-file
+
from __future__ import absolute_import
import importlib
diff --git a/sdks/python/apache_beam/runners/runner_test.py b/sdks/python/apache_beam/runners/runner_test.py
index 914fa12..b26134f 100644
--- a/sdks/python/apache_beam/runners/runner_test.py
+++ b/sdks/python/apache_beam/runners/runner_test.py
@@ -22,6 +22,8 @@
caching and clearing values that are not tested elsewhere.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py
index 95c67cf..0612e5e 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -17,6 +17,8 @@
"""SDK harness for executing Python Fns via the Fn API."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -26,7 +28,6 @@
import json
import logging
import random
-import re
import threading
from builtins import next
from builtins import object
@@ -158,7 +159,7 @@
"""A source-like operation that gathers input from the runner."""
def __init__(self,
- operation_name, # type: str
+ operation_name, # type: Union[str, common.NameContext]
step_name,
consumers, # type: Mapping[Any, Iterable[operations.Operation]]
counter_factory,
@@ -244,6 +245,7 @@
if stop_index < self.stop:
self.stop = stop_index
return self.stop - 1, None, None, self.stop
+ return None
def progress_metrics(self):
# type: () -> beam_fn_api_pb2.Metrics.PTransform
@@ -263,7 +265,7 @@
class _StateBackedIterable(object):
def __init__(self,
- state_handler,
+ state_handler, # type: sdk_worker.CachingStateHandler
state_key, # type: beam_fn_api_pb2.StateKey
coder_or_impl, # type: Union[coders.Coder, coder_impl.CoderImpl]
is_cached=False
@@ -292,7 +294,7 @@
class StateBackedSideInputMap(object):
def __init__(self,
- state_handler,
+ state_handler, # type: sdk_worker.CachingStateHandler
transform_id, # type: str
tag, # type: Optional[str]
side_input_data, # type: pvalue.SideInputData
@@ -432,7 +434,7 @@
class SynchronousBagRuntimeState(userstate.BagRuntimeState):
def __init__(self,
- state_handler,
+ state_handler, # type: sdk_worker.CachingStateHandler
state_key, # type: beam_fn_api_pb2.StateKey
value_coder # type: coders.Coder
):
@@ -479,7 +481,7 @@
class SynchronousSetRuntimeState(userstate.SetRuntimeState):
def __init__(self,
- state_handler,
+ state_handler, # type: sdk_worker.CachingStateHandler
state_key, # type: beam_fn_api_pb2.StateKey
value_coder # type: coders.Coder
):
@@ -576,7 +578,7 @@
"""Interface for state and timers from SDK to Fn API servicer of state.."""
def __init__(self,
- state_handler,
+ state_handler, # type: sdk_worker.CachingStateHandler
transform_id, # type: str
key_coder, # type: coders.Coder
window_coder, # type: coders.Coder
@@ -614,6 +616,7 @@
window # type: windowed_value.BoundedWindow
):
# type: (...) -> OutputTimer
+ assert self._timer_receivers is not None
return OutputTimer(
key, window, self._timer_receivers[timer_spec.name])
@@ -693,7 +696,7 @@
def __init__(self,
process_bundle_descriptor, # type: beam_fn_api_pb2.ProcessBundleDescriptor
- state_handler, # type: Union[FnApiRunner.StateServicer, GrpcStateHandler]
+ state_handler, # type: sdk_worker.CachingStateHandler
data_channel_factory # type: data_plane.DataChannelFactory
):
# type: (...) -> None
@@ -866,6 +869,7 @@
deferred_remainder # type: Tuple[windowed_value.WindowedValue, Timestamp]
):
# type: (...) -> beam_fn_api_pb2.DelayedBundleApplication
+ assert op.input_info is not None
# TODO(SDF): For non-root nodes, need main_input_coder + residual_coder.
((element_and_restriction, output_watermark),
deferred_watermark) = deferred_remainder
@@ -1009,7 +1013,7 @@
data_channel_factory, # type: data_plane.DataChannelFactory
counter_factory,
state_sampler, # type: statesampler.StateSampler
- state_handler
+ state_handler # type: sdk_worker.CachingStateHandler
):
self.descriptor = descriptor
self.data_channel_factory = data_channel_factory
@@ -1128,19 +1132,26 @@
@BeamTransformFactory.register_urn(
DATA_INPUT_URN, beam_fn_api_pb2.RemoteGrpcPort)
-def create(factory, transform_id, transform_proto, grpc_port, consumers):
+def create_source_runner(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ grpc_port, # type: beam_fn_api_pb2.RemoteGrpcPort
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> DataInputOperation
# Timers are the one special case where we don't want to call the
# (unlabeled) operation.process() method, which we detect here.
# TODO(robertwb): Consider generalizing if there are any more cases.
output_pcoll = only_element(transform_proto.outputs.values())
output_consumers = only_element(consumers.values())
- if (len(output_consumers) == 1
- and isinstance(only_element(output_consumers), operations.DoOperation)):
+ if len(output_consumers) == 1:
do_op = only_element(output_consumers)
- for tag, pcoll_id in do_op.timer_inputs.items():
- if pcoll_id == output_pcoll:
- output_consumers[:] = [TimerConsumer(tag, do_op)]
- break
+ if isinstance(do_op, operations.DoOperation):
+ for tag, pcoll_id in do_op.timer_inputs.items():
+ if pcoll_id == output_pcoll:
+ output_consumers[:] = [TimerConsumer(tag, do_op)]
+ break
if grpc_port.coder_id:
output_coder = factory.get_coder(grpc_port.coder_id)
@@ -1163,7 +1174,14 @@
@BeamTransformFactory.register_urn(
DATA_OUTPUT_URN, beam_fn_api_pb2.RemoteGrpcPort)
-def create(factory, transform_id, transform_proto, grpc_port, consumers):
+def create_sink_runner(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ grpc_port, # type: beam_fn_api_pb2.RemoteGrpcPort
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> DataOutputOperation
if grpc_port.coder_id:
output_coder = factory.get_coder(grpc_port.coder_id)
else:
@@ -1184,7 +1202,14 @@
@BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_READ_URN, None)
-def create(factory, transform_id, transform_proto, parameter, consumers):
+def create_source_java(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ parameter,
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> operations.ReadOperation
# The Dataflow runner harness strips the base64 encoding.
source = pickler.loads(base64.b64encode(parameter))
spec = operation_specs.WorkerRead(
@@ -1226,7 +1251,14 @@
@BeamTransformFactory.register_urn(
python_urns.IMPULSE_READ_TRANSFORM, beam_runner_api_pb2.ReadPayload)
-def create(factory, transform_id, transform_proto, parameter, consumers):
+def create_read_from_impulse_python(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ parameter, # type: beam_runner_api_pb2.ReadPayload
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> operations.ImpulseReadOperation
return operations.ImpulseReadOperation(
common.NameContext(transform_proto.unique_name, transform_id),
factory.counter_factory,
@@ -1238,7 +1270,13 @@
@BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_PARDO_URN, None)
-def create(factory, transform_id, transform_proto, serialized_fn, consumers):
+def create_dofn_javasdk(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ serialized_fn,
+ consumers # type: Dict[str, List[operations.Operation]]
+):
return _create_pardo_operation(
factory, transform_id, transform_proto, consumers, serialized_fn)
@@ -1246,7 +1284,7 @@
@BeamTransformFactory.register_urn(
common_urns.sdf_components.PAIR_WITH_RESTRICTION.urn,
beam_runner_api_pb2.ParDoPayload)
-def create(*args):
+def create_pair_with_restriction(*args):
class PairWithRestriction(beam.DoFn):
def __init__(self, fn, restriction_provider):
@@ -1267,7 +1305,7 @@
@BeamTransformFactory.register_urn(
common_urns.sdf_components.SPLIT_AND_SIZE_RESTRICTIONS.urn,
beam_runner_api_pb2.ParDoPayload)
-def create(*args):
+def create_split_and_size_restrictions(*args):
class SplitAndSizeRestrictions(beam.DoFn):
def __init__(self, fn, restriction_provider):
@@ -1285,7 +1323,13 @@
@BeamTransformFactory.register_urn(
common_urns.sdf_components.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn,
beam_runner_api_pb2.ParDoPayload)
-def create(factory, transform_id, transform_proto, parameter, consumers):
+def create_process_sized_elements_and_restrictions(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ parameter, # type: beam_runner_api_pb2.ParDoPayload
+ consumers # type: Dict[str, List[operations.Operation]]
+):
assert parameter.do_fn.urn == python_urns.PICKLED_DOFN_INFO
serialized_fn = parameter.do_fn.payload
return _create_pardo_operation(
@@ -1310,7 +1354,14 @@
@BeamTransformFactory.register_urn(
common_urns.primitives.PAR_DO.urn, beam_runner_api_pb2.ParDoPayload)
-def create(factory, transform_id, transform_proto, parameter, consumers):
+def create_par_do(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ parameter, # type: beam_runner_api_pb2.ParDoPayload
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> operations.DoOperation
assert parameter.do_fn.urn == python_urns.PICKLED_DOFN_INFO
serialized_fn = parameter.do_fn.payload
return _create_pardo_operation(
@@ -1334,9 +1385,7 @@
(tag, beam.pvalue.SideInputData.from_runner_api(si, factory.context))
for tag, si in pardo_proto.side_inputs.items()]
tagged_side_inputs.sort(
- key=lambda tag_si: int(re.match('side([0-9]+)(-.*)?$',
- tag_si[0],
- re.DOTALL).group(1)))
+ key=lambda tag_si: sideinputs.get_sideinput_index(tag_si[0]))
side_input_maps = [
StateBackedSideInputMap(
factory.state_handler,
@@ -1445,7 +1494,13 @@
@BeamTransformFactory.register_urn(
common_urns.primitives.ASSIGN_WINDOWS.urn,
beam_runner_api_pb2.WindowingStrategy)
-def create(factory, transform_id, transform_proto, parameter, consumers):
+def create_assign_windows(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ parameter, # type: beam_runner_api_pb2.WindowingStrategy
+ consumers # type: Dict[str, List[operations.Operation]]
+):
class WindowIntoDoFn(beam.DoFn):
def __init__(self, windowing):
self.windowing = windowing
@@ -1464,7 +1519,14 @@
@BeamTransformFactory.register_urn(IDENTITY_DOFN_URN, None)
-def create(factory, transform_id, transform_proto, unused_parameter, consumers):
+def create_identity_dofn(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ parameter,
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> operations.FlattenOperation
return factory.augment_oldstyle_op(
operations.FlattenOperation(
common.NameContext(transform_proto.unique_name, transform_id),
@@ -1479,7 +1541,14 @@
@BeamTransformFactory.register_urn(
common_urns.combine_components.COMBINE_PER_KEY_PRECOMBINE.urn,
beam_runner_api_pb2.CombinePayload)
-def create(factory, transform_id, transform_proto, payload, consumers):
+def create_combine_per_key_precombine(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ payload, # type: beam_runner_api_pb2.CombinePayload
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> operations.PGBKCVOperation
serialized_combine_fn = pickler.dumps(
(beam.CombineFn.from_runner_api(payload.combine_fn, factory.context),
[], {}))
@@ -1500,7 +1569,13 @@
@BeamTransformFactory.register_urn(
common_urns.combine_components.COMBINE_PER_KEY_MERGE_ACCUMULATORS.urn,
beam_runner_api_pb2.CombinePayload)
-def create(factory, transform_id, transform_proto, payload, consumers):
+def create_combbine_per_key_merge_accumulators(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ payload, # type: beam_runner_api_pb2.CombinePayload
+ consumers # type: Dict[str, List[operations.Operation]]
+):
return _create_combine_phase_operation(
factory, transform_id, transform_proto, payload, consumers, 'merge')
@@ -1508,7 +1583,13 @@
@BeamTransformFactory.register_urn(
common_urns.combine_components.COMBINE_PER_KEY_EXTRACT_OUTPUTS.urn,
beam_runner_api_pb2.CombinePayload)
-def create(factory, transform_id, transform_proto, payload, consumers):
+def create_combine_per_key_extract_outputs(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ payload, # type: beam_runner_api_pb2.CombinePayload
+ consumers # type: Dict[str, List[operations.Operation]]
+):
return _create_combine_phase_operation(
factory, transform_id, transform_proto, payload, consumers, 'extract')
@@ -1516,7 +1597,13 @@
@BeamTransformFactory.register_urn(
common_urns.combine_components.COMBINE_GROUPED_VALUES.urn,
beam_runner_api_pb2.CombinePayload)
-def create(factory, transform_id, transform_proto, payload, consumers):
+def create_combine_grouped_values(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ payload, # type: beam_runner_api_pb2.CombinePayload
+ consumers # type: Dict[str, List[operations.Operation]]
+):
return _create_combine_phase_operation(
factory, transform_id, transform_proto, payload, consumers, 'all')
@@ -1542,7 +1629,14 @@
@BeamTransformFactory.register_urn(common_urns.primitives.FLATTEN.urn, None)
-def create(factory, transform_id, transform_proto, unused_parameter, consumers):
+def create_flatten(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ payload,
+ consumers # type: Dict[str, List[operations.Operation]]
+):
+ # type: (...) -> operations.FlattenOperation
return factory.augment_oldstyle_op(
operations.FlattenOperation(
common.NameContext(transform_proto.unique_name, transform_id),
@@ -1558,7 +1652,13 @@
@BeamTransformFactory.register_urn(
common_urns.primitives.MAP_WINDOWS.urn,
beam_runner_api_pb2.FunctionSpec)
-def create(factory, transform_id, transform_proto, mapping_fn_spec, consumers):
+def create_map_windows(
+ factory, # type: BeamTransformFactory
+ transform_id, # type: str
+ transform_proto, # type: beam_runner_api_pb2.PTransform
+ mapping_fn_spec, # type: beam_runner_api_pb2.SdkFunctionSpec
+ consumers # type: Dict[str, List[operations.Operation]]
+):
assert mapping_fn_spec.urn == python_urns.PICKLED_WINDOW_MAPPING_FN
window_mapping_fn = pickler.loads(mapping_fn_spec.payload)
diff --git a/sdks/python/apache_beam/runners/worker/channel_factory.py b/sdks/python/apache_beam/runners/worker/channel_factory.py
index d0823fa..29d1504 100644
--- a/sdks/python/apache_beam/runners/worker/channel_factory.py
+++ b/sdks/python/apache_beam/runners/worker/channel_factory.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
"""Factory to create grpc channel."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py
index 6fb00f3..d2fddf6 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane.py
@@ -17,6 +17,8 @@
"""Implementation of ``DataChannel``s to communicate across the data plane."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/worker/data_plane_test.py b/sdks/python/apache_beam/runners/worker/data_plane_test.py
index 677758b..64b4ae7 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane_test.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane_test.py
@@ -17,6 +17,8 @@
"""Tests for apache_beam.runners.worker.data_plane."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/worker/log_handler.py b/sdks/python/apache_beam/runners/worker/log_handler.py
index 4ecd342..f1df332 100644
--- a/sdks/python/apache_beam/runners/worker/log_handler.py
+++ b/sdks/python/apache_beam/runners/worker/log_handler.py
@@ -16,6 +16,8 @@
#
"""Beam fn API log handler."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/worker/log_handler_test.py b/sdks/python/apache_beam/runners/worker/log_handler_test.py
index c79ccf9..1fa5fc7 100644
--- a/sdks/python/apache_beam/runners/worker/log_handler_test.py
+++ b/sdks/python/apache_beam/runners/worker/log_handler_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/worker/logger.py b/sdks/python/apache_beam/runners/worker/logger.py
index 72f24ff..be67e47 100644
--- a/sdks/python/apache_beam/runners/worker/logger.py
+++ b/sdks/python/apache_beam/runners/worker/logger.py
@@ -19,6 +19,8 @@
"""Python worker logging."""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
diff --git a/sdks/python/apache_beam/runners/worker/logger_test.py b/sdks/python/apache_beam/runners/worker/logger_test.py
index c131775..e04e1b5 100644
--- a/sdks/python/apache_beam/runners/worker/logger_test.py
+++ b/sdks/python/apache_beam/runners/worker/logger_test.py
@@ -17,6 +17,8 @@
"""Tests for worker logging utilities."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import unicode_literals
diff --git a/sdks/python/apache_beam/runners/worker/opcounters.py b/sdks/python/apache_beam/runners/worker/opcounters.py
index 80528d3..9a1a30a 100644
--- a/sdks/python/apache_beam/runners/worker/opcounters.py
+++ b/sdks/python/apache_beam/runners/worker/opcounters.py
@@ -20,6 +20,8 @@
"""Counters collect the progress of the Worker for reporting to the service."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/runners/worker/opcounters_test.py b/sdks/python/apache_beam/runners/worker/opcounters_test.py
index 13e78b2..62e8b5c 100644
--- a/sdks/python/apache_beam/runners/worker/opcounters_test.py
+++ b/sdks/python/apache_beam/runners/worker/opcounters_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/runners/worker/operation_specs.py b/sdks/python/apache_beam/runners/worker/operation_specs.py
index 464558d..1f53d5a 100644
--- a/sdks/python/apache_beam/runners/worker/operation_specs.py
+++ b/sdks/python/apache_beam/runners/worker/operation_specs.py
@@ -21,6 +21,8 @@
source, write to a sink, parallel do, etc.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py
index b4601f1..aa95057 100644
--- a/sdks/python/apache_beam/runners/worker/operations.py
+++ b/sdks/python/apache_beam/runners/worker/operations.py
@@ -20,6 +20,8 @@
"""Worker operations executor."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
@@ -655,6 +657,7 @@
with self.scoped_process_state:
delayed_application = self.dofn_receiver.receive(o)
if delayed_application:
+ assert self.execution_context is not None
self.execution_context.delayed_applications.append(
(self, delayed_application))
@@ -746,6 +749,7 @@
def process(self, o):
# type: (WindowedValue) -> None
+ assert self.tagged_receivers is not None
with self.scoped_process_state:
try:
with self.lock:
@@ -756,6 +760,7 @@
# the lock.
delayed_application = self.dofn_runner.process_with_sized_restriction(o)
if delayed_application:
+ assert self.execution_context is not None
self.execution_context.delayed_applications.append(
(self, delayed_application))
finally:
@@ -782,6 +787,7 @@
metrics = super(SdfProcessSizedElements, self).progress_metrics()
current_element_progress = self.current_element_progress()
if current_element_progress:
+ assert self.input_info is not None
metrics.active_elements.measured.input_element_counts[
self.input_info[1]] = 1
metrics.active_elements.fraction_remaining = (
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py
index 3eefb17..193ac1c 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py
@@ -16,6 +16,8 @@
#
"""SDK harness for executing Python Fns via the Fn API."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -31,9 +33,11 @@
from builtins import object
from concurrent import futures
from typing import TYPE_CHECKING
+from typing import Any
from typing import Callable
from typing import DefaultDict
from typing import Dict
+from typing import Iterable
from typing import Iterator
from typing import List
from typing import Optional
@@ -423,7 +427,7 @@
step_name = sampler_info.state_name.step_name
state_name = sampler_info.state_name.name
state_lull_log = (
- 'There has been a processing lull of over %.2f seconds in state %s'
+ 'Operation ongoing for over %.2f seconds in state %s'
% (sampler_info.time_since_transition / 1e9, state_name))
step_name_log = (' in step %s ' % step_name) if step_name else ''
@@ -436,7 +440,8 @@
stack_trace = '-NOT AVAILABLE-'
_LOGGER.warning(
- '%s%s. Traceback:\n%s', state_lull_log, step_name_log, stack_trace)
+ '%s%s without returning. Current Traceback:\n%s',
+ state_lull_log, step_name_log, stack_trace)
def process_bundle_progress(self,
request, # type: beam_fn_api_pb2.ProcessBundleProgressRequest
@@ -488,16 +493,43 @@
yield
-class StateHandlerFactory(with_metaclass(abc.ABCMeta, object)):
+class StateHandler(with_metaclass(abc.ABCMeta, object)): # type: ignore[misc]
+ """An abstract object representing a ``StateHandler``."""
+
+ @abc.abstractmethod
+ def get_raw(self,
+ state_key, # type: beam_fn_api_pb2.StateKey
+ continuation_token=None # type: Optional[bytes]
+ ):
+ # type: (...) -> Tuple[bytes, Optional[bytes]]
+ raise NotImplementedError(type(self))
+
+ @abc.abstractmethod
+ def append_raw(self,
+ state_key, # type: beam_fn_api_pb2.StateKey
+ data # type: bytes
+ ):
+ # type: (...) -> _Future
+ raise NotImplementedError(type(self))
+
+ @abc.abstractmethod
+ def clear(self, state_key):
+ # type: (beam_fn_api_pb2.StateKey) -> _Future
+ raise NotImplementedError(type(self))
+
+
+class StateHandlerFactory(with_metaclass(abc.ABCMeta, object)): # type: ignore[misc]
"""An abstract factory for creating ``DataChannel``."""
@abc.abstractmethod
def create_state_handler(self, api_service_descriptor):
+ # type: (endpoints_pb2.ApiServiceDescriptor) -> CachingStateHandler
"""Returns a ``StateHandler`` from the given ApiServiceDescriptor."""
raise NotImplementedError(type(self))
@abc.abstractmethod
def close(self):
+ # type: () -> None
"""Close all channels that this factory owns."""
raise NotImplementedError(type(self))
@@ -509,14 +541,14 @@
"""
def __init__(self, state_cache, credentials=None):
- self._state_handler_cache = {} # type: Dict[str, GrpcStateHandler]
+ self._state_handler_cache = {} # type: Dict[str, CachingStateHandler]
self._lock = threading.Lock()
self._throwing_state_handler = ThrowingStateHandler()
self._credentials = credentials
self._state_cache = state_cache
def create_state_handler(self, api_service_descriptor):
- # type: (endpoints_pb2.ApiServiceDescriptor) -> GrpcStateHandler
+ # type: (endpoints_pb2.ApiServiceDescriptor) -> CachingStateHandler
if not api_service_descriptor:
return self._throwing_state_handler
url = api_service_descriptor.url
@@ -547,6 +579,7 @@
return self._state_handler_cache[url]
def close(self):
+ # type: () -> None
_LOGGER.info('Closing all cached gRPC state handlers.')
for _, state_handler in self._state_handler_cache.items():
state_handler.done()
@@ -554,15 +587,15 @@
self._state_cache.evict_all()
-class ThrowingStateHandler(object):
+class ThrowingStateHandler(StateHandler):
"""A state handler that errors on any requests."""
- def blocking_get(self, state_key, coder):
+ def get_raw(self, state_key, coder):
raise RuntimeError(
'Unable to handle state requests for ProcessBundleDescriptor without '
'state ApiServiceDescriptor for state key %s.' % state_key)
- def append(self, state_key, coder, elements):
+ def append_raw(self, state_key, coder, elements):
raise RuntimeError(
'Unable to handle state requests for ProcessBundleDescriptor without '
'state ApiServiceDescriptor for state key %s.' % state_key)
@@ -573,7 +606,7 @@
'state ApiServiceDescriptor for state key %s.' % state_key)
-class GrpcStateHandler(object):
+class GrpcStateHandler(StateHandler):
_DONE = object()
@@ -671,6 +704,7 @@
return future
def _blocking_request(self, request):
+ # type: (beam_fn_api_pb2.StateRequest) -> beam_fn_api_pb2.StateResponse
req_future = self._request(request)
while not req_future.wait(timeout=1):
if self._exc_info:
@@ -699,7 +733,10 @@
class CachingStateHandler(object):
""" A State handler which retrieves and caches state. """
- def __init__(self, global_state_cache, underlying_state):
+ def __init__(self,
+ global_state_cache, # type: StateCache
+ underlying_state # type: StateHandler
+ ):
self._underlying = underlying_state
self._state_cache = global_state_cache
self._context = threading.local()
@@ -725,7 +762,12 @@
finally:
self._context.cache_token = None
- def blocking_get(self, state_key, coder, is_cached=False):
+ def blocking_get(self,
+ state_key, # type: beam_fn_api_pb2.StateKey
+ coder, # type: coder_impl.CoderImpl
+ is_cached=False
+ ):
+ # type: (...) -> Iterator[Any]
if not self._should_be_cached(is_cached):
# Cache disabled / no cache token. Can't do a lookup/store in the cache.
# Fall back to lazily materializing the state, one element at a time.
@@ -766,6 +808,7 @@
return self._underlying.append_raw(state_key, out.get())
def clear(self, state_key, is_cached=False):
+ # type: (beam_fn_api_pb2.StateKey, bool) -> _Future
if self._should_be_cached(is_cached):
cache_key = self._convert_to_cache_key(state_key)
self._state_cache.clear(cache_key, self._context.cache_token)
@@ -775,7 +818,11 @@
# type: () -> None
self._underlying.done()
- def _materialize_iter(self, state_key, coder):
+ def _materialize_iter(self,
+ state_key, # type: beam_fn_api_pb2.StateKey
+ coder # type: coder_impl.CoderImpl
+ ):
+ # type: (...) -> Iterator[Any]
"""Materializes the state lazily, one element at a time.
:return A generator which returns the next element if advanced.
"""
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
index 99921b8..3569e13 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
@@ -16,6 +16,8 @@
#
"""SDK Fn Harness entry point."""
+# pytype: skip-file
+
from __future__ import absolute_import
import http.server
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py
index cae65a2..19285f9 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.worker.sdk_worker_main."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py
index e14b3f4..0aa102d 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py
@@ -16,6 +16,8 @@
#
"""Tests for apache_beam.runners.worker.sdk_worker."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/worker/sideinputs.py b/sdks/python/apache_beam/runners/worker/sideinputs.py
index 7c1d649..7d84323 100644
--- a/sdks/python/apache_beam/runners/worker/sideinputs.py
+++ b/sdks/python/apache_beam/runners/worker/sideinputs.py
@@ -17,6 +17,8 @@
"""Utilities for handling side inputs."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/runners/worker/sideinputs_test.py b/sdks/python/apache_beam/runners/worker/sideinputs_test.py
index 4a8f7c8..e7491f7 100644
--- a/sdks/python/apache_beam/runners/worker/sideinputs_test.py
+++ b/sdks/python/apache_beam/runners/worker/sideinputs_test.py
@@ -17,6 +17,8 @@
"""Tests for side input utilities."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/worker/statecache.py b/sdks/python/apache_beam/runners/worker/statecache.py
index 0c96f3f..1e304c4 100644
--- a/sdks/python/apache_beam/runners/worker/statecache.py
+++ b/sdks/python/apache_beam/runners/worker/statecache.py
@@ -16,22 +16,31 @@
#
"""A module for caching state reads/writes in Beam applications."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
import logging
import threading
+from typing import Callable
+from typing import DefaultDict
+from typing import Hashable
+from typing import Set
+from typing import TypeVar
from apache_beam.metrics import monitoring_infos
_LOGGER = logging.getLogger(__name__)
+CallableT = TypeVar('CallableT', bound='Callable')
+
class Metrics(object):
"""Metrics container for state cache metrics."""
# A set of all registered metrics
- ALL_METRICS = set()
+ ALL_METRICS = set() # type: Set[Hashable]
PREFIX = "beam:metric:statecache:"
def __init__(self):
@@ -42,12 +51,14 @@
"""
if hasattr(self._context, 'metrics'):
return # Already initialized
- self._context.metrics = collections.defaultdict(int)
+ self._context.metrics = collections.defaultdict(int) # type: DefaultDict[Hashable, int]
def count(self, name):
+ # type: (str) -> None
self._context.metrics[name] += 1
def hit_miss(self, total_name, hit_miss_name):
+ # type: (str, str) -> None
self._context.metrics[total_name] += 1
self._context.metrics[hit_miss_name] += 1
@@ -78,6 +89,7 @@
@staticmethod
def counter_hit_miss(total_name, hit_name, miss_name):
+ # type: (str, str, str) -> Callable[[CallableT], CallableT]
"""Decorator for counting function calls and whether
the return value equals None (=miss) or not (=hit)."""
Metrics.ALL_METRICS.update([total_name, hit_name, miss_name])
@@ -98,6 +110,7 @@
@staticmethod
def counter(metric_name):
+ # type: (str) -> Callable[[CallableT], CallableT]
"""Decorator for counting function calls."""
Metrics.ALL_METRICS.add(metric_name)
diff --git a/sdks/python/apache_beam/runners/worker/statecache_test.py b/sdks/python/apache_beam/runners/worker/statecache_test.py
index 00ae852..730d149 100644
--- a/sdks/python/apache_beam/runners/worker/statecache_test.py
+++ b/sdks/python/apache_beam/runners/worker/statecache_test.py
@@ -16,6 +16,8 @@
#
"""Tests for state caching."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/runners/worker/statesampler.py b/sdks/python/apache_beam/runners/worker/statesampler.py
index 950755c..36a6568 100644
--- a/sdks/python/apache_beam/runners/worker/statesampler.py
+++ b/sdks/python/apache_beam/runners/worker/statesampler.py
@@ -17,13 +17,15 @@
# This module is experimental. No backwards-compatibility guarantees.
+# pytype: skip-file
+
from __future__ import absolute_import
import contextlib
import threading
-from collections import namedtuple
from typing import TYPE_CHECKING
from typing import Dict
+from typing import NamedTuple
from typing import Optional
from typing import Union
@@ -80,12 +82,12 @@
return get_current_tracker()
-StateSamplerInfo = namedtuple(
+StateSamplerInfo = NamedTuple(
'StateSamplerInfo',
- ['state_name',
- 'transition_count',
- 'time_since_transition',
- 'tracked_thread'])
+ [('state_name', CounterName),
+ ('transition_count', int),
+ ('time_since_transition', int),
+ ('tracked_thread', Optional[threading.Thread])])
# Default period for sampling current state of pipeline execution.
diff --git a/sdks/python/apache_beam/runners/worker/statesampler_slow.py b/sdks/python/apache_beam/runners/worker/statesampler_slow.py
index cb41fe0..0034046 100644
--- a/sdks/python/apache_beam/runners/worker/statesampler_slow.py
+++ b/sdks/python/apache_beam/runners/worker/statesampler_slow.py
@@ -17,6 +17,8 @@
# This module is experimental. No backwards-compatibility guarantees.
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
@@ -78,8 +80,7 @@
def reset(self):
# type: () -> None
- for state in self._states_by_name.values():
- state.nsecs = 0
+ pass
class ScopedState(object):
diff --git a/sdks/python/apache_beam/runners/worker/statesampler_test.py b/sdks/python/apache_beam/runners/worker/statesampler_test.py
index ed51ae1..3b5caed 100644
--- a/sdks/python/apache_beam/runners/worker/statesampler_test.py
+++ b/sdks/python/apache_beam/runners/worker/statesampler_test.py
@@ -16,6 +16,8 @@
#
"""Tests for state sampler."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py b/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py
index be24eb4..1ba11b6 100644
--- a/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py
+++ b/sdks/python/apache_beam/runners/worker/worker_id_interceptor.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
"""Client Interceptor to inject worker_id"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py b/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py
index 411e309..e23a50d 100644
--- a/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py
+++ b/sdks/python/apache_beam/runners/worker/worker_id_interceptor_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
"""Test for WorkerIdInterceptor"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
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 117f38b..baa1f3a 100644
--- a/sdks/python/apache_beam/runners/worker/worker_pool_main.py
+++ b/sdks/python/apache_beam/runners/worker/worker_pool_main.py
@@ -26,6 +26,8 @@
This entry point is used by the Python SDK container in worker pool mode.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/preprocess.py b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/preprocess.py
index e36ffb7..c98047a 100644
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/preprocess.py
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/preprocess.py
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Preprocessor applying tf.transform to the chicago_taxi data."""
+# pytype: skip-file
+
from __future__ import absolute_import, division, print_function
import argparse
@@ -24,6 +26,7 @@
from tensorflow_transform.tf_metadata import dataset_metadata, dataset_schema
import apache_beam as beam
+from apache_beam.io.gcp.bigquery import _ReadFromBigQuery as ReadFromBigQuery
from apache_beam.metrics.metric import MetricsFilter
from apache_beam.testing.load_tests.load_test_metrics_utils import (
MeasureTime, MetricsReader)
@@ -125,6 +128,7 @@
metrics_monitor = None
if publish_to_bq:
metrics_monitor = MetricsReader(
+ publish_to_bq=publish_to_bq,
project_name=project,
bq_table=metrics_table,
bq_dataset=metrics_dataset,
@@ -140,9 +144,8 @@
query = taxi.make_sql(input_handle, max_rows, for_eval=False)
raw_data = (
pipeline
- | 'ReadBigQuery' >> beam.io.Read(
- beam.io.BigQuerySource(query=query,
- use_standard_sql=True))
+ | 'ReadBigQuery' >> ReadFromBigQuery(query=query, project=project,
+ use_standard_sql=True)
| 'Measure time: start' >> beam.ParDo(MeasureTime(namespace)))
decode_transform = beam.Map(
taxi.clean_raw_data_dict, raw_feature_spec=raw_feature_spec)
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/process_tfma.py b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/process_tfma.py
index 11fc335..53d8532 100644
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/process_tfma.py
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/process_tfma.py
@@ -13,6 +13,8 @@
# limitations under the License.
"""Runs a batch job for performing Tensorflow Model Analysis."""
+# pytype: skip-file
+
from __future__ import absolute_import, division, print_function
import argparse
@@ -22,6 +24,7 @@
from tensorflow_model_analysis.evaluators import evaluator
import apache_beam as beam
+from apache_beam.io.gcp.bigquery import _ReadFromBigQuery as ReadFromBigQuery
from apache_beam.metrics.metric import MetricsFilter
from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime
from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader
@@ -80,6 +83,7 @@
metrics_monitor = None
if publish_to_bq:
metrics_monitor = MetricsReader(
+ publish_to_bq=publish_to_bq,
project_name=project,
bq_table=metrics_table,
bq_dataset=metrics_dataset,
@@ -92,8 +96,8 @@
raw_feature_spec = taxi.get_raw_feature_spec(schema)
raw_data = (
pipeline
- | 'ReadBigQuery' >> beam.io.Read(
- beam.io.BigQuerySource(query=query, use_standard_sql=True))
+ | 'ReadBigQuery' >> ReadFromBigQuery(query=query, project=project,
+ use_standard_sql=True)
| 'Measure time: Start' >> beam.ParDo(MeasureTime(metrics_namespace))
| 'CleanData' >> beam.Map(lambda x: (
taxi.clean_raw_data_dict(x, raw_feature_spec))))
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/run_chicago.sh b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/run_chicago.sh
index 1fa0a87..c93fa0e 100755
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/run_chicago.sh
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/run_chicago.sh
@@ -15,12 +15,16 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-
-# This script builds a Docker container with the user specified requirements on top of
-# an existing Python worker Docker container (either one you build from source as
-# described in CONTAINERS.md or from a released Docker container).
-set -euo pipefail
-echo Starting distributed TFDV stats computation and schema generation...
+#
+# The Chicago Taxi example demonstrates the end-to-end workflow and steps
+# of how to analyze, validate and transform data, train a model, analyze
+# and serve it.
+#
+# Example usage:
+# ./run_chicago.sh gs://my-gcs-bucket DataflowRunner \
+# --sdk_location=\"apache-beam.tar.gz\" --region=\"us-central1\"
+#
+set -eo pipefail
if [[ -z "$1" ]]; then
echo "GCS bucket name required"
@@ -32,14 +36,21 @@
exit 1
fi
-if [[ -z "$3" ]]; then
- echo "SDK location needed"
- exit 1
-fi
-
GCS_BUCKET=$1
RUNNER=$2
-SDK_LOCATION=$3
+PIPELINE_OPTIONS=$3
+
+if [[ "$RUNNER" == "PortableRunner" ]]; then
+ METRICS_TABLE_SUFFIX='_flink'
+fi
+
+# Loop through pipeline options and append
+shift
+while [[ $# -gt 2 ]]
+do
+ PIPELINE_OPTIONS=${PIPELINE_OPTIONS}" "$3
+ shift
+done
JOB_ID="chicago-taxi-tfdv-$(date +%Y%m%d-%H%M%S)"
JOB_OUTPUT_PATH=${GCS_BUCKET}/${JOB_ID}/chicago_taxi_output
@@ -62,6 +73,7 @@
# Analyze and validate
# Compute stats and generate a schema based on the stats.
+echo Starting distributed TFDV stats computation and schema generation...
python tfdv_analyze_and_validate.py \
--input bigquery-public-data.chicago_taxi_trips.taxi_trips \
@@ -78,10 +90,10 @@
--max_rows=${MAX_ROWS} \
--publish_to_big_query=true \
--metrics_dataset='beam_performance' \
- --metrics_table='tfdv_analyze' \
+ --metrics_table='tfdv_analyze'${METRICS_TABLE_SUFFIX} \
--metric_reporting_project ${GCP_PROJECT} \
- --sdk_location=${SDK_LOCATION} \
- --setup_file ./setup.py
+ --setup_file ./setup.py \
+ ${PIPELINE_OPTIONS}
EVAL_JOB_ID=${JOB_ID}-eval
@@ -103,10 +115,10 @@
--max_rows=${MAX_ROWS} \
--publish_to_big_query=true \
--metrics_dataset='beam_performance' \
- --metrics_table='chicago_taxi_tfdv_validate' \
- --sdk_location=${SDK_LOCATION} \
+ --metrics_table='chicago_taxi_tfdv_validate'${METRICS_TABLE_SUFFIX} \
--metric_reporting_project ${GCP_PROJECT} \
- --setup_file ./setup.py
+ --setup_file ./setup.py \
+ ${PIPELINE_OPTIONS}
# End analyze and validate
echo Preprocessing train data...
@@ -125,10 +137,10 @@
--max_rows ${MAX_ROWS} \
--publish_to_big_query=true \
--metrics_dataset='beam_performance' \
- --metrics_table='chicago_taxi_preprocess' \
- --sdk_location=${SDK_LOCATION} \
+ --metrics_table='chicago_taxi_preprocess'${METRICS_TABLE_SUFFIX} \
--metric_reporting_project ${GCP_PROJECT} \
- --setup_file ./setup.py
+ --setup_file ./setup.py \
+ ${PIPELINE_OPTIONS}
#Train ML engine
TRAINER_JOB_ID="chicago_taxi_trainer_$(date +%Y%m%d_%H%M%S)"
@@ -186,7 +198,7 @@
--max_eval_rows=${MAX_ROWS} \
--publish_to_big_query=true \
--metrics_dataset='beam_performance' \
- --metrics_table='chicago_taxi_process_tfma' \
- --sdk_location=${SDK_LOCATION} \
+ --metrics_table='chicago_taxi_process_tfma'${METRICS_TABLE_SUFFIX} \
--metric_reporting_project ${GCP_PROJECT} \
- --setup_file ./setup.py
+ --setup_file ./setup.py \
+ ${PIPELINE_OPTIONS}
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/setup.py b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/setup.py
index 7832420..22f6150 100644
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/setup.py
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/setup.py
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Setup dependencies for local and cloud deployment."""
+# pytype: skip-file
+
from __future__ import absolute_import
import setuptools
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/tfdv_analyze_and_validate.py b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/tfdv_analyze_and_validate.py
index 2e76c0b..b73fafd 100644
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/tfdv_analyze_and_validate.py
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/tfdv_analyze_and_validate.py
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Compute stats, infer schema, and validate stats for chicago taxi example."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -25,6 +27,7 @@
from tensorflow_metadata.proto.v0 import statistics_pb2
import apache_beam as beam
+from apache_beam.io.gcp.bigquery import _ReadFromBigQuery as ReadFromBigQuery
from apache_beam.metrics.metric import MetricsFilter
from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime
from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader
@@ -95,6 +98,7 @@
metrics_monitor = None
if publish_to_bq:
metrics_monitor = MetricsReader(
+ publish_to_bq=publish_to_bq,
project_name=project,
bq_table=metrics_table,
bq_dataset=metrics_dataset,
@@ -105,8 +109,8 @@
table_name=input_handle, max_rows=max_rows, for_eval=for_eval)
raw_data = (
pipeline
- | 'ReadBigQuery' >> beam.io.Read(
- beam.io.BigQuerySource(query=query, use_standard_sql=True))
+ | 'ReadBigQuery' >> ReadFromBigQuery(query=query, project=project,
+ use_standard_sql=True)
| 'Measure time: Start' >> beam.ParDo(MeasureTime(namespace))
| 'ConvertToTFDVInput' >> beam.Map(
lambda x: {key: np.asarray([x[key]])
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/model.py b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/model.py
index f04d0a8..a726dfd 100644
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/model.py
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/model.py
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Defines the model used to predict who will tip in the Chicago Taxi demo."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/task.py b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/task.py
index 7b80056..aef7d3c 100644
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/task.py
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/task.py
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Trainer for the chicago_taxi demo."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/taxi.py b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/taxi.py
index e1ec44f..540f0ca 100644
--- a/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/taxi.py
+++ b/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/trainer/taxi.py
@@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utility and schema methods for the chicago_taxi sample."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_launcher.py b/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_launcher.py
index 2779cb4..9977cc6 100644
--- a/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_launcher.py
+++ b/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_launcher.py
@@ -56,6 +56,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_util.py b/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_util.py
index 916faa4..c55f937 100644
--- a/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_util.py
+++ b/sdks/python/apache_beam/testing/benchmarks/nexmark/nexmark_util.py
@@ -32,6 +32,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query0.py b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query0.py
index 3df848d..5c56178 100644
--- a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query0.py
+++ b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query0.py
@@ -25,6 +25,8 @@
to verify the infrastructure.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import apache_beam as beam
diff --git a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query1.py b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query1.py
index 24b8579..88df169 100644
--- a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query1.py
+++ b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query1.py
@@ -23,6 +23,8 @@
This query converts bid prices from dollars to euros.
It illustrates a simple map.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import apache_beam as beam
diff --git a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query2.py b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query2.py
index 33ee3f3..fe45503 100644
--- a/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query2.py
+++ b/sdks/python/apache_beam/testing/benchmarks/nexmark/queries/query2.py
@@ -24,6 +24,8 @@
It illustrates a simple filter.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import apache_beam as beam
diff --git a/sdks/python/apache_beam/testing/datatype_inference.py b/sdks/python/apache_beam/testing/datatype_inference.py
index b1a689c..fbf3380 100644
--- a/sdks/python/apache_beam/testing/datatype_inference.py
+++ b/sdks/python/apache_beam/testing/datatype_inference.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import array
diff --git a/sdks/python/apache_beam/testing/datatype_inference_test.py b/sdks/python/apache_beam/testing/datatype_inference_test.py
index 131eafb..0cbd6f9 100644
--- a/sdks/python/apache_beam/testing/datatype_inference_test.py
+++ b/sdks/python/apache_beam/testing/datatype_inference_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/testing/extra_assertions.py b/sdks/python/apache_beam/testing/extra_assertions.py
index 53a9eeb..b67f814 100644
--- a/sdks/python/apache_beam/testing/extra_assertions.py
+++ b/sdks/python/apache_beam/testing/extra_assertions.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
diff --git a/sdks/python/apache_beam/testing/extra_assertions_test.py b/sdks/python/apache_beam/testing/extra_assertions_test.py
index 8948f4e..8b078af 100644
--- a/sdks/python/apache_beam/testing/extra_assertions_test.py
+++ b/sdks/python/apache_beam/testing/extra_assertions_test.py
@@ -15,6 +15,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 57425a2..51b5d98 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
@@ -142,6 +142,8 @@
-Prunner=TestDataflowRunner :sdks:python:apache_beam:testing:load-tests:run
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
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 9c2f2d0..068e95e 100644
--- a/sdks/python/apache_beam/testing/load_tests/combine_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/combine_test.py
@@ -117,6 +117,8 @@
TestDataflowRunner :sdks:python:apache_beam:testing:load-tests:run
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 d19e7f7..bc15bbb 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
@@ -118,6 +118,8 @@
-Prunner=TestDataflowRunner :sdks:python:apache_beam:testing:load-tests:run
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/testing/load_tests/load_test.py b/sdks/python/apache_beam/testing/load_tests/load_test.py
index 71aa3a8..0a7eb37 100644
--- a/sdks/python/apache_beam/testing/load_tests/load_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/load_test.py
@@ -14,6 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+# pytype: skip-file
+
from __future__ import absolute_import
import json
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 ec5cd77..d843e24 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
@@ -27,6 +27,8 @@
* total_bytes_count
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 7c05422..ce9c264 100644
--- a/sdks/python/apache_beam/testing/load_tests/pardo_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/pardo_test.py
@@ -117,6 +117,8 @@
-Prunner=TestDataflowRunner :sdks:python:apache_beam:testing:load-tests:run
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
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 4143f83..2e0dc189 100644
--- a/sdks/python/apache_beam/testing/load_tests/sideinput_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/sideinput_test.py
@@ -112,6 +112,8 @@
-Prunner=TestDataflowRunner :sdks:python:apache_beam:testing:load-tests:run
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_pipeline.py b/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_pipeline.py
index 3fc277e..88ccaed 100644
--- a/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_pipeline.py
+++ b/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_pipeline.py
@@ -25,6 +25,8 @@
Values have to be reparsed again to bytes
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
diff --git a/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_test.py b/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_test.py
index 6042e68..56ad43f 100644
--- a/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_test.py
+++ b/sdks/python/apache_beam/testing/load_tests/streaming/group_by_key_streaming_test.py
@@ -34,6 +34,8 @@
* --metrics_table=gbk_stream
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/testing/metric_result_matchers.py b/sdks/python/apache_beam/testing/metric_result_matchers.py
index 126ba3d..15dc6ff 100644
--- a/sdks/python/apache_beam/testing/metric_result_matchers.py
+++ b/sdks/python/apache_beam/testing/metric_result_matchers.py
@@ -41,6 +41,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from hamcrest import equal_to
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 dc674ac..6f6a131 100644
--- a/sdks/python/apache_beam/testing/metric_result_matchers_test.py
+++ b/sdks/python/apache_beam/testing/metric_result_matchers_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the metric_result_matchers."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/testing/pipeline_verifiers.py b/sdks/python/apache_beam/testing/pipeline_verifiers.py
index cf99541..3e286a4 100644
--- a/sdks/python/apache_beam/testing/pipeline_verifiers.py
+++ b/sdks/python/apache_beam/testing/pipeline_verifiers.py
@@ -22,6 +22,8 @@
`hamcrest.core.base_matcher.BaseMatcher` and override _matches.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/testing/pipeline_verifiers_test.py b/sdks/python/apache_beam/testing/pipeline_verifiers_test.py
index 2b2db66..f7b2f07 100644
--- a/sdks/python/apache_beam/testing/pipeline_verifiers_test.py
+++ b/sdks/python/apache_beam/testing/pipeline_verifiers_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the test pipeline verifiers"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/testing/synthetic_pipeline.py b/sdks/python/apache_beam/testing/synthetic_pipeline.py
index fbef112..653281d 100644
--- a/sdks/python/apache_beam/testing/synthetic_pipeline.py
+++ b/sdks/python/apache_beam/testing/synthetic_pipeline.py
@@ -32,6 +32,8 @@
data for the pipeline.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/testing/synthetic_pipeline_test.py b/sdks/python/apache_beam/testing/synthetic_pipeline_test.py
index 18d4bdd..b74e650 100644
--- a/sdks/python/apache_beam/testing/synthetic_pipeline_test.py
+++ b/sdks/python/apache_beam/testing/synthetic_pipeline_test.py
@@ -17,6 +17,8 @@
"""Tests for apache_beam.testing.synthetic_pipeline."""
+# pytype: skip-file
+
from __future__ import absolute_import
import glob
diff --git a/sdks/python/apache_beam/testing/test_pipeline.py b/sdks/python/apache_beam/testing/test_pipeline.py
index 7a2d575..0ca81ec 100644
--- a/sdks/python/apache_beam/testing/test_pipeline.py
+++ b/sdks/python/apache_beam/testing/test_pipeline.py
@@ -17,12 +17,13 @@
"""Test Pipeline, a wrapper of Pipeline for test purpose"""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
import shlex
-
-from nose.plugins.skip import SkipTest
+from unittest import SkipTest
from apache_beam.internal import pickler
from apache_beam.options.pipeline_options import PipelineOptions
@@ -56,10 +57,9 @@
For example, use assert_that for test validation::
- pipeline = TestPipeline()
- pcoll = ...
- assert_that(pcoll, equal_to(...))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = ...
+ assert_that(pcoll, equal_to(...))
"""
def __init__(self,
diff --git a/sdks/python/apache_beam/testing/test_pipeline_test.py b/sdks/python/apache_beam/testing/test_pipeline_test.py
index 8efd8c6..59779cd 100644
--- a/sdks/python/apache_beam/testing/test_pipeline_test.py
+++ b/sdks/python/apache_beam/testing/test_pipeline_test.py
@@ -17,6 +17,8 @@
"""Unit test for the TestPipeline class"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -104,16 +106,16 @@
self.assertEqual(test_pipeline.get_option(name), value)
def test_skip_IT(self):
- test_pipeline = TestPipeline(is_integration_test=True)
- test_pipeline.run()
- # Note that this will never be reached since it should be skipped above.
+ with TestPipeline(is_integration_test=True) as _:
+ # Note that this will never be reached since it should be skipped above.
+ pass
self.fail()
@mock.patch('apache_beam.testing.test_pipeline.Pipeline.run', autospec=True)
def test_not_use_test_runner_api(self, mock_run):
- test_pipeline = TestPipeline(argv=['--not-use-test-runner-api'],
- blocking=False)
- test_pipeline.run()
+ with TestPipeline(argv=['--not-use-test-runner-api'],
+ blocking=False) as test_pipeline:
+ pass
mock_run.assert_called_once_with(test_pipeline, test_runner_api=False)
diff --git a/sdks/python/apache_beam/testing/test_stream.py b/sdks/python/apache_beam/testing/test_stream.py
index ff54fe2..7f026df 100644
--- a/sdks/python/apache_beam/testing/test_stream.py
+++ b/sdks/python/apache_beam/testing/test_stream.py
@@ -19,6 +19,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from abc import ABCMeta
diff --git a/sdks/python/apache_beam/testing/test_stream_it_test.py b/sdks/python/apache_beam/testing/test_stream_it_test.py
index 8e724f6..1ef6ecd 100644
--- a/sdks/python/apache_beam/testing/test_stream_it_test.py
+++ b/sdks/python/apache_beam/testing/test_stream_it_test.py
@@ -17,6 +17,8 @@
"""Integration tests for the test_stream module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
@@ -47,8 +49,8 @@
@wraps(fn)
def wrapped(self):
if self.runner_name not in runners:
- self.skipTest('The "{}", does not support the TestStream transform. Supported runners: {}'.format(
- self.runner_name, runners))
+ self.skipTest('The "{}", does not support the TestStream transform. '
+ 'Supported runners: {}'.format(self.runner_name, runners))
else:
return fn(self)
return wrapped
diff --git a/sdks/python/apache_beam/testing/test_stream_service.py b/sdks/python/apache_beam/testing/test_stream_service.py
index 8819073..28ec9ca 100644
--- a/sdks/python/apache_beam/testing/test_stream_service.py
+++ b/sdks/python/apache_beam/testing/test_stream_service.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from concurrent.futures import ThreadPoolExecutor
diff --git a/sdks/python/apache_beam/testing/test_stream_service_test.py b/sdks/python/apache_beam/testing/test_stream_service_test.py
index 2b25100..fd0f897 100644
--- a/sdks/python/apache_beam/testing/test_stream_service_test.py
+++ b/sdks/python/apache_beam/testing/test_stream_service_test.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/testing/test_stream_test.py b/sdks/python/apache_beam/testing/test_stream_test.py
index bfadb5e..ba599bd 100644
--- a/sdks/python/apache_beam/testing/test_stream_test.py
+++ b/sdks/python/apache_beam/testing/test_stream_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the test_stream module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
@@ -112,20 +114,19 @@
options = PipelineOptions()
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- my_record_fn = RecordFn()
- records = p | test_stream | beam.ParDo(my_record_fn)
+ with TestPipeline(options=options) as p:
+ my_record_fn = RecordFn()
+ records = p | test_stream | beam.ParDo(my_record_fn)
- assert_that(records, equal_to([
- ('a', timestamp.Timestamp(10)),
- ('b', timestamp.Timestamp(10)),
- ('c', timestamp.Timestamp(10)),
- ('d', timestamp.Timestamp(20)),
- ('e', timestamp.Timestamp(20)),
- ('late', timestamp.Timestamp(12)),
- ('last', timestamp.Timestamp(310)),]))
+ assert_that(records, equal_to([
+ ('a', timestamp.Timestamp(10)),
+ ('b', timestamp.Timestamp(10)),
+ ('c', timestamp.Timestamp(10)),
+ ('d', timestamp.Timestamp(20)),
+ ('e', timestamp.Timestamp(20)),
+ ('late', timestamp.Timestamp(12)),
+ ('last', timestamp.Timestamp(310)),]))
- p.run()
def test_multiple_outputs(self):
"""Tests that the TestStream supports emitting to multiple PCollections."""
@@ -416,33 +417,31 @@
def test_basic_execution_sideinputs(self):
options = PipelineOptions()
options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
+ with TestPipeline(options=options) as p:
- main_stream = (p
- | 'main TestStream' >> TestStream()
- .advance_watermark_to(10)
- .add_elements(['e']))
- side_stream = (p
- | 'side TestStream' >> TestStream()
- .add_elements([window.TimestampedValue(2, 2)])
- .add_elements([window.TimestampedValue(1, 1)])
- .add_elements([window.TimestampedValue(7, 7)])
- .add_elements([window.TimestampedValue(4, 4)])
- )
+ main_stream = (p
+ | 'main TestStream' >> TestStream()
+ .advance_watermark_to(10)
+ .add_elements(['e']))
+ side_stream = (p
+ | 'side TestStream' >> TestStream()
+ .add_elements([window.TimestampedValue(2, 2)])
+ .add_elements([window.TimestampedValue(1, 1)])
+ .add_elements([window.TimestampedValue(7, 7)])
+ .add_elements([window.TimestampedValue(4, 4)])
+ )
- class RecordFn(beam.DoFn):
- def process(self,
- elm=beam.DoFn.ElementParam,
- ts=beam.DoFn.TimestampParam,
- side=beam.DoFn.SideInputParam):
- yield (elm, ts, side)
+ class RecordFn(beam.DoFn):
+ def process(self,
+ elm=beam.DoFn.ElementParam,
+ ts=beam.DoFn.TimestampParam,
+ side=beam.DoFn.SideInputParam):
+ yield (elm, ts, side)
- records = (main_stream # pylint: disable=unused-variable
- | beam.ParDo(RecordFn(), beam.pvalue.AsList(side_stream)))
+ records = (main_stream # pylint: disable=unused-variable
+ | beam.ParDo(RecordFn(), beam.pvalue.AsList(side_stream)))
- assert_that(records, equal_to([('e', Timestamp(10), [2, 1, 7, 4])]))
-
- p.run()
+ assert_that(records, equal_to([('e', Timestamp(10), [2, 1, 7, 4])]))
def test_basic_execution_batch_sideinputs_fixed_windows(self):
options = PipelineOptions()
diff --git a/sdks/python/apache_beam/testing/test_utils.py b/sdks/python/apache_beam/testing/test_utils.py
index f9aa128..1c40324 100644
--- a/sdks/python/apache_beam/testing/test_utils.py
+++ b/sdks/python/apache_beam/testing/test_utils.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import hashlib
diff --git a/sdks/python/apache_beam/testing/test_utils_test.py b/sdks/python/apache_beam/testing/test_utils_test.py
index 2b16c30c..e68d26f 100644
--- a/sdks/python/apache_beam/testing/test_utils_test.py
+++ b/sdks/python/apache_beam/testing/test_utils_test.py
@@ -17,6 +17,8 @@
"""Unittest for testing utilities,"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/testing/util.py b/sdks/python/apache_beam/testing/util.py
index 5b6bc85..b41ae71 100644
--- a/sdks/python/apache_beam/testing/util.py
+++ b/sdks/python/apache_beam/testing/util.py
@@ -17,6 +17,8 @@
"""Utilities for testing Beam pipelines."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
@@ -174,15 +176,19 @@
# 2) As a fallback if we encounter a TypeError in python 3. this method
# works on collections that have different types.
except (BeamAssertException, TypeError):
+ unexpected = []
for element in actual:
try:
expected_list.remove(element)
except ValueError:
- raise BeamAssertException(
- 'Failed assert: %r == %r' % (expected, actual))
- if expected_list:
- raise BeamAssertException(
- 'Failed assert: %r == %r' % (expected, actual))
+ unexpected.append(element)
+ if unexpected or expected_list:
+ msg = 'Failed assert: %r == %r' % (expected, actual)
+ if unexpected:
+ msg = msg + ', unexpected elements %r' % unexpected
+ if expected_list:
+ msg = msg + ', missing elements %r' % expected_list
+ raise BeamAssertException(msg)
return _equal
diff --git a/sdks/python/apache_beam/testing/util_test.py b/sdks/python/apache_beam/testing/util_test.py
index 72c9205..c9385ee 100644
--- a/sdks/python/apache_beam/testing/util_test.py
+++ b/sdks/python/apache_beam/testing/util_test.py
@@ -17,6 +17,8 @@
"""Unit tests for testing utilities."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
@@ -42,6 +44,12 @@
class UtilTest(unittest.TestCase):
+ def setUp(self):
+ try: # Python 3
+ _ = self.assertRaisesRegex
+ except AttributeError: # Python 2
+ self.assertRaisesRegex = self.assertRaisesRegexp
+
def test_assert_that_passes(self):
with TestPipeline() as p:
assert_that(p | Create([1, 2, 3]), equal_to([1, 2, 3]))
@@ -67,6 +75,27 @@
with TestPipeline() as p:
assert_that(p | Create([1, 10, 100]), equal_to([1, 2, 3]))
+ def test_assert_missing(self):
+ with self.assertRaisesRegex(BeamAssertException,
+ r"missing elements \['c'\]"):
+ with TestPipeline() as p:
+ assert_that(p | Create(['a', 'b']), equal_to(['a', 'b', 'c']))
+
+ def test_assert_unexpected(self):
+ with self.assertRaisesRegex(BeamAssertException,
+ r"unexpected elements \['c', 'd'\]|"
+ r"unexpected elements \['d', 'c'\]"):
+ with TestPipeline() as p:
+ assert_that(p | Create(['a', 'b', 'c', 'd']), equal_to(['a', 'b']))
+
+ def test_assert_missing_and_unexpected(self):
+ with self.assertRaisesRegex(
+ BeamAssertException,
+ r"unexpected elements \['c'\].*missing elements \['d'\]"):
+ with TestPipeline() as p:
+ assert_that(p | Create(['a', 'b', 'c']),
+ equal_to(['a', 'b', 'd']))
+
def test_reified_value_passes(self):
expected = [TestWindowedValue(v, MIN_TIMESTAMP, [GlobalWindow()])
for v in [1, 2, 3]]
@@ -111,7 +140,7 @@
with TestPipeline() as p:
assert_that(p | Create([1, 2, 3]), is_not_empty())
- def test_assert_that_fails_on_empty_expected(self):
+ def test_assert_that_fails_on_is_not_empty_expected(self):
with self.assertRaises(BeamAssertException):
with TestPipeline() as p:
assert_that(p | Create([]), is_not_empty())
diff --git a/sdks/python/apache_beam/tools/coders_microbenchmark.py b/sdks/python/apache_beam/tools/coders_microbenchmark.py
index edaa3ea..707e447 100644
--- a/sdks/python/apache_beam/tools/coders_microbenchmark.py
+++ b/sdks/python/apache_beam/tools/coders_microbenchmark.py
@@ -28,6 +28,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/tools/distribution_counter_microbenchmark.py b/sdks/python/apache_beam/tools/distribution_counter_microbenchmark.py
index 06035d5..5889853 100644
--- a/sdks/python/apache_beam/tools/distribution_counter_microbenchmark.py
+++ b/sdks/python/apache_beam/tools/distribution_counter_microbenchmark.py
@@ -23,6 +23,8 @@
python -m apache_beam.tools.distribution_counter_microbenchmark
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
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 538f65f..d8506b2 100644
--- a/sdks/python/apache_beam/tools/fn_api_runner_microbenchmark.py
+++ b/sdks/python/apache_beam/tools/fn_api_runner_microbenchmark.py
@@ -54,6 +54,8 @@
R^2 0.95189
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/tools/map_fn_microbenchmark.py b/sdks/python/apache_beam/tools/map_fn_microbenchmark.py
index 6b4a143..191b335 100644
--- a/sdks/python/apache_beam/tools/map_fn_microbenchmark.py
+++ b/sdks/python/apache_beam/tools/map_fn_microbenchmark.py
@@ -30,6 +30,8 @@
python -m apache_beam.tools.map_fn_microbenchmark
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/tools/microbenchmarks_test.py b/sdks/python/apache_beam/tools/microbenchmarks_test.py
index 74949f6..850ef33 100644
--- a/sdks/python/apache_beam/tools/microbenchmarks_test.py
+++ b/sdks/python/apache_beam/tools/microbenchmarks_test.py
@@ -17,6 +17,8 @@
"""Unit tests for microbenchmarks code."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/tools/sideinput_microbenchmark.py b/sdks/python/apache_beam/tools/sideinput_microbenchmark.py
index 2a46aee..37cff3e 100644
--- a/sdks/python/apache_beam/tools/sideinput_microbenchmark.py
+++ b/sdks/python/apache_beam/tools/sideinput_microbenchmark.py
@@ -22,6 +22,8 @@
python -m apache_beam.tools.sideinput_microbenchmark
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/tools/utils.py b/sdks/python/apache_beam/tools/utils.py
index 41253a8..838e347 100644
--- a/sdks/python/apache_beam/tools/utils.py
+++ b/sdks/python/apache_beam/tools/utils.py
@@ -17,6 +17,8 @@
"""Utility functions for all microbenchmarks."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py
index 501e4d0..80bbdfc 100644
--- a/sdks/python/apache_beam/transforms/combiners.py
+++ b/sdks/python/apache_beam/transforms/combiners.py
@@ -17,6 +17,8 @@
"""A library of basic combiner PTransform subclasses."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -403,33 +405,44 @@
def process(self, key_and_bundles):
_, bundles = key_and_bundles
- heap = []
- for bundle in bundles:
- if not heap:
- if self._less_than or self._key:
- heap = [
+
+ def push(hp, e):
+ if len(hp) < self._n:
+ heapq.heappush(hp, e)
+ return False
+ elif e < hp[0]:
+ # Because _TopPerBundle returns sorted lists, all other elements
+ # will also be smaller.
+ return True
+ else:
+ heapq.heappushpop(hp, e)
+ return False
+
+ if self._less_than or self._key:
+ heapc = [] # type: List[cy_combiners.ComparableValue]
+ for bundle in bundles:
+ if not heapc:
+ heapc = [
cy_combiners.ComparableValue(element, self._less_than, self._key)
for element in bundle]
- else:
- heap = bundle
- continue
- for element in reversed(bundle):
- if self._less_than or self._key:
- element = cy_combiners.ComparableValue(
- element, self._less_than, self._key)
- if len(heap) < self._n:
- heapq.heappush(heap, element)
- elif element < heap[0]:
- # Because _TopPerBundle returns sorted lists, all other elements
- # will also be smaller.
- break
- else:
- heapq.heappushpop(heap, element)
+ continue
+ for element in reversed(bundle):
+ if push(heapc, cy_combiners.ComparableValue(
+ element, self._less_than, self._key)):
+ break
+ heapc.sort()
+ yield [wrapper.value for wrapper in reversed(heapc)]
- heap.sort()
- if self._less_than or self._key:
- yield [wrapper.value for wrapper in reversed(heap)]
else:
+ heap = [] # type: List[T]
+ for bundle in bundles:
+ if not heap:
+ heap = bundle
+ continue
+ for element in reversed(bundle):
+ if push(heap, element):
+ break
+ heap.sort()
yield heap[::-1]
diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py
index 975d3ae..412f079 100644
--- a/sdks/python/apache_beam/transforms/combiners_test.py
+++ b/sdks/python/apache_beam/transforms/combiners_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for our libraries of combine PTransforms."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -30,12 +32,16 @@
import apache_beam as beam
import apache_beam.transforms.combiners as combine
+from apache_beam.metrics import Metrics
+from apache_beam.metrics import MetricsFilter
from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_stream import TestStream
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.testing.util import equal_to_per_window
+from apache_beam.transforms import trigger
from apache_beam.transforms import window
from apache_beam.transforms.core import CombineGlobally
from apache_beam.transforms.core import Create
@@ -43,94 +49,129 @@
from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher
from apache_beam.transforms.ptransform import PTransform
+from apache_beam.transforms.trigger import AfterAll
+from apache_beam.transforms.trigger import AfterCount
+from apache_beam.transforms.trigger import AfterWatermark
+from apache_beam.transforms.window import GlobalWindows
from apache_beam.transforms.window import TimestampCombiner
from apache_beam.typehints import TypeCheckError
from apache_beam.utils.timestamp import Timestamp
+class SortedConcatWithCounters(beam.CombineFn):
+ """CombineFn for incrementing three different counters:
+ counter, distribution, gauge,
+ at the same time concatenating words."""
+
+ def __init__(self):
+ beam.CombineFn.__init__(self)
+ self.word_counter = Metrics.counter(self.__class__, 'word_counter')
+ self.word_lengths_counter = Metrics.counter(
+ self.__class__, 'word_lengths')
+ self.word_lengths_dist = Metrics.distribution(
+ self.__class__, 'word_len_dist')
+ self.last_word_len = Metrics.gauge(self.__class__, 'last_word_len')
+
+ def create_accumulator(self):
+ return ''
+
+ def add_input(self, acc, element):
+ self.word_counter.inc(1)
+ self.word_lengths_counter.inc(len(element))
+ self.word_lengths_dist.update(len(element))
+ self.last_word_len.set(len(element))
+
+ return acc + element
+
+ def merge_accumulators(self, accs):
+ return ''.join(accs)
+
+ def extract_output(self, acc):
+ # The sorted acc became a list of characters
+ # and has to be converted back to a string using join.
+ return ''.join(sorted(acc))
+
+
class CombineTest(unittest.TestCase):
def test_builtin_combines(self):
- pipeline = TestPipeline()
+ with TestPipeline() as pipeline:
- vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
- mean = sum(vals) / float(len(vals))
- size = len(vals)
+ vals = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
+ mean = sum(vals) / float(len(vals))
+ size = len(vals)
- # First for global combines.
- pcoll = pipeline | 'start' >> Create(vals)
- result_mean = pcoll | 'mean' >> combine.Mean.Globally()
- result_count = pcoll | 'count' >> combine.Count.Globally()
- assert_that(result_mean, equal_to([mean]), label='assert:mean')
- assert_that(result_count, equal_to([size]), label='assert:size')
+ # First for global combines.
+ pcoll = pipeline | 'start' >> Create(vals)
+ result_mean = pcoll | 'mean' >> combine.Mean.Globally()
+ result_count = pcoll | 'count' >> combine.Count.Globally()
+ assert_that(result_mean, equal_to([mean]), label='assert:mean')
+ assert_that(result_count, equal_to([size]), label='assert:size')
- # Again for per-key combines.
- pcoll = pipeline | 'start-perkey' >> Create([('a', x) for x in vals])
- result_key_mean = pcoll | 'mean-perkey' >> combine.Mean.PerKey()
- result_key_count = pcoll | 'count-perkey' >> combine.Count.PerKey()
- assert_that(result_key_mean, equal_to([('a', mean)]), label='key:mean')
- assert_that(result_key_count, equal_to([('a', size)]), label='key:size')
- pipeline.run()
+ # Again for per-key combines.
+ pcoll = pipeline | 'start-perkey' >> Create([('a', x) for x in vals])
+ result_key_mean = pcoll | 'mean-perkey' >> combine.Mean.PerKey()
+ result_key_count = pcoll | 'count-perkey' >> combine.Count.PerKey()
+ assert_that(result_key_mean, equal_to([('a', mean)]), label='key:mean')
+ assert_that(result_key_count, equal_to([('a', size)]), label='key:size')
def test_top(self):
- pipeline = TestPipeline()
+ with TestPipeline() as pipeline:
- # First for global combines.
- pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
- result_top = pcoll | 'top' >> combine.Top.Largest(5)
- result_bot = pcoll | 'bot' >> combine.Top.Smallest(4)
- assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top')
- assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot')
+ # First for global combines.
+ pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
+ result_top = pcoll | 'top' >> combine.Top.Largest(5)
+ result_bot = pcoll | 'bot' >> combine.Top.Smallest(4)
+ assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top')
+ assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot')
- # Again for per-key combines.
- pcoll = pipeline | 'start-perkey' >> Create(
- [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
- result_key_top = pcoll | 'top-perkey' >> combine.Top.LargestPerKey(5)
- result_key_bot = pcoll | 'bot-perkey' >> combine.Top.SmallestPerKey(4)
- assert_that(result_key_top, equal_to([('a', [9, 6, 6, 5, 3])]),
- label='key:top')
- assert_that(result_key_bot, equal_to([('a', [0, 1, 1, 1])]),
- label='key:bot')
- pipeline.run()
+ # Again for per-key combines.
+ pcoll = pipeline | 'start-perkey' >> Create(
+ [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
+ result_key_top = pcoll | 'top-perkey' >> combine.Top.LargestPerKey(5)
+ result_key_bot = pcoll | 'bot-perkey' >> combine.Top.SmallestPerKey(4)
+ assert_that(result_key_top, equal_to([('a', [9, 6, 6, 5, 3])]),
+ label='key:top')
+ assert_that(result_key_bot, equal_to([('a', [0, 1, 1, 1])]),
+ label='key:bot')
@unittest.skipIf(sys.version_info[0] > 2, 'deprecated comparator')
def test_top_py2(self):
- pipeline = TestPipeline()
+ with TestPipeline() as pipeline:
- # A parameter we'll be sharing with a custom comparator.
- names = {0: 'zo',
- 1: 'one',
- 2: 'twoo',
- 3: 'three',
- 5: 'fiiive',
- 6: 'sssssix',
- 9: 'nniiinne'}
+ # A parameter we'll be sharing with a custom comparator.
+ names = {0: 'zo',
+ 1: 'one',
+ 2: 'twoo',
+ 3: 'three',
+ 5: 'fiiive',
+ 6: 'sssssix',
+ 9: 'nniiinne'}
- # First for global combines.
- pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
+ # First for global combines.
+ pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
- result_cmp = pcoll | 'cmp' >> combine.Top.Of(
- 6,
- lambda a, b, names: len(names[a]) < len(names[b]),
- names) # Note parameter passed to comparator.
- result_cmp_rev = pcoll | 'cmp_rev' >> combine.Top.Of(
- 3,
- lambda a, b, names: len(names[a]) < len(names[b]),
- names, # Note parameter passed to comparator.
- reverse=True)
- assert_that(result_cmp, equal_to([[9, 6, 6, 5, 3, 2]]), label='assert:cmp')
- assert_that(result_cmp_rev, equal_to([[0, 1, 1]]), label='assert:cmp_rev')
+ result_cmp = pcoll | 'cmp' >> combine.Top.Of(
+ 6,
+ lambda a, b, names: len(names[a]) < len(names[b]),
+ names) # Note parameter passed to comparator.
+ result_cmp_rev = pcoll | 'cmp_rev' >> combine.Top.Of(
+ 3,
+ lambda a, b, names: len(names[a]) < len(names[b]),
+ names, # Note parameter passed to comparator.
+ reverse=True)
+ assert_that(result_cmp, equal_to([[9, 6, 6, 5, 3, 2]]), label='CheckCmp')
+ assert_that(result_cmp_rev, equal_to([[0, 1, 1]]), label='CheckCmpRev')
- # Again for per-key combines.
- pcoll = pipeline | 'start-perkye' >> Create(
- [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
- result_key_cmp = pcoll | 'cmp-perkey' >> combine.Top.PerKey(
- 6,
- lambda a, b, names: len(names[a]) < len(names[b]),
- names) # Note parameter passed to comparator.
- assert_that(result_key_cmp, equal_to([('a', [9, 6, 6, 5, 3, 2])]),
- label='key:cmp')
- pipeline.run()
+ # Again for per-key combines.
+ pcoll = pipeline | 'start-perkye' >> Create(
+ [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
+ result_key_cmp = pcoll | 'cmp-perkey' >> combine.Top.PerKey(
+ 6,
+ lambda a, b, names: len(names[a]) < len(names[b]),
+ names) # Note parameter passed to comparator.
+ assert_that(result_key_cmp, equal_to([('a', [9, 6, 6, 5, 3, 2])]),
+ label='key:cmp')
def test_empty_global_top(self):
with TestPipeline() as p:
@@ -141,12 +182,11 @@
elements = list(range(100))
random.shuffle(elements)
- pipeline = TestPipeline()
- shards = [pipeline | 'Shard%s' % shard >> beam.Create(elements[shard::7])
- for shard in range(7)]
- assert_that(shards | beam.Flatten() | combine.Top.Largest(10),
- equal_to([[99, 98, 97, 96, 95, 94, 93, 92, 91, 90]]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ shards = [pipeline | 'Shard%s' % shard >> beam.Create(elements[shard::7])
+ for shard in range(7)]
+ assert_that(shards | beam.Flatten() | combine.Top.Largest(10),
+ equal_to([[99, 98, 97, 96, 95, 94, 93, 92, 91, 90]]))
def test_top_key(self):
self.assertEqual(
@@ -228,22 +268,22 @@
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_top_shorthands(self):
- pipeline = TestPipeline()
+ with TestPipeline() as pipeline:
- pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
- result_top = pcoll | 'top' >> beam.CombineGlobally(combine.Largest(5))
- result_bot = pcoll | 'bot' >> beam.CombineGlobally(combine.Smallest(4))
- assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top')
- assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot')
+ pcoll = pipeline | 'start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
+ result_top = pcoll | 'top' >> beam.CombineGlobally(combine.Largest(5))
+ result_bot = pcoll | 'bot' >> beam.CombineGlobally(combine.Smallest(4))
+ assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='assert:top')
+ assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='assert:bot')
- pcoll = pipeline | 'start-perkey' >> Create(
- [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
- result_ktop = pcoll | 'top-perkey' >> beam.CombinePerKey(combine.Largest(5))
- result_kbot = pcoll | 'bot-perkey' >> beam.CombinePerKey(
- combine.Smallest(4))
- assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='k:top')
- assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='k:bot')
- pipeline.run()
+ pcoll = pipeline | 'start-perkey' >> Create(
+ [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
+ result_ktop = pcoll | 'top-perkey' >> beam.CombinePerKey(
+ combine.Largest(5))
+ result_kbot = pcoll | 'bot-perkey' >> beam.CombinePerKey(
+ combine.Smallest(4))
+ assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='ktop')
+ assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='kbot')
def test_top_no_compact(self):
@@ -252,24 +292,23 @@
def compact(self, accumulator):
return accumulator
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
- result_top = pcoll | 'Top' >> beam.CombineGlobally(
- TopCombineFnNoCompact(5, key=lambda x: x))
- result_bot = pcoll | 'Bot' >> beam.CombineGlobally(
- TopCombineFnNoCompact(4, reverse=True))
- assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='Assert:Top')
- assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='Assert:Bot')
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> Create([6, 3, 1, 1, 9, 1, 5, 2, 0, 6])
+ result_top = pcoll | 'Top' >> beam.CombineGlobally(
+ TopCombineFnNoCompact(5, key=lambda x: x))
+ result_bot = pcoll | 'Bot' >> beam.CombineGlobally(
+ TopCombineFnNoCompact(4, reverse=True))
+ assert_that(result_top, equal_to([[9, 6, 6, 5, 3]]), label='Assert:Top')
+ assert_that(result_bot, equal_to([[0, 1, 1, 1]]), label='Assert:Bot')
- pcoll = pipeline | 'Start-Perkey' >> Create(
- [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
- result_ktop = pcoll | 'Top-PerKey' >> beam.CombinePerKey(
- TopCombineFnNoCompact(5, key=lambda x: x))
- result_kbot = pcoll | 'Bot-PerKey' >> beam.CombinePerKey(
- TopCombineFnNoCompact(4, reverse=True))
- assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='K:Top')
- assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='K:Bot')
- pipeline.run()
+ pcoll = pipeline | 'Start-Perkey' >> Create(
+ [('a', x) for x in [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]])
+ result_ktop = pcoll | 'Top-PerKey' >> beam.CombinePerKey(
+ TopCombineFnNoCompact(5, key=lambda x: x))
+ result_kbot = pcoll | 'Bot-PerKey' >> beam.CombinePerKey(
+ TopCombineFnNoCompact(4, reverse=True))
+ assert_that(result_ktop, equal_to([('a', [9, 6, 6, 5, 3])]), label='KTop')
+ assert_that(result_kbot, equal_to([('a', [0, 1, 1, 1])]), label='KBot')
def test_global_sample(self):
def is_good_sample(actual):
@@ -285,21 +324,20 @@
label='check-%d' % ix)
def test_per_key_sample(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'start-perkey' >> Create(
- sum(([(i, 1), (i, 1), (i, 2), (i, 2)] for i in range(9)), []))
- result = pcoll | 'sample' >> combine.Sample.FixedSizePerKey(3)
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'start-perkey' >> Create(
+ sum(([(i, 1), (i, 1), (i, 2), (i, 2)] for i in range(9)), []))
+ result = pcoll | 'sample' >> combine.Sample.FixedSizePerKey(3)
- def matcher():
- def match(actual):
- for _, samples in actual:
- equal_to([3])([len(samples)])
- num_ones = sum(1 for x in samples if x == 1)
- num_twos = sum(1 for x in samples if x == 2)
- equal_to([1, 2])([num_ones, num_twos])
- return match
- assert_that(result, matcher())
- pipeline.run()
+ def matcher():
+ def match(actual):
+ for _, samples in actual:
+ equal_to([3])([len(samples)])
+ num_ones = sum(1 for x in samples if x == 1)
+ num_twos = sum(1 for x in samples if x == 2)
+ equal_to([1, 2])([num_ones, num_twos])
+ return match
+ assert_that(result, matcher())
def test_tuple_combine_fn(self):
with TestPipeline() as p:
@@ -321,30 +359,28 @@
assert_that(result, equal_to([(1, 7.0 / 4, 3)]))
def test_to_list_and_to_dict(self):
- pipeline = TestPipeline()
- the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
- pcoll = pipeline | 'start' >> Create(the_list)
- result = pcoll | 'to list' >> combine.ToList()
+ with TestPipeline() as pipeline:
+ the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
+ pcoll = pipeline | 'start' >> Create(the_list)
+ result = pcoll | 'to list' >> combine.ToList()
- def matcher(expected):
- def match(actual):
- equal_to(expected[0])(actual[0])
- return match
- assert_that(result, matcher([the_list]))
- pipeline.run()
+ def matcher(expected):
+ def match(actual):
+ equal_to(expected[0])(actual[0])
+ return match
+ assert_that(result, matcher([the_list]))
- pipeline = TestPipeline()
- pairs = [(1, 2), (3, 4), (5, 6)]
- pcoll = pipeline | 'start-pairs' >> Create(pairs)
- result = pcoll | 'to dict' >> combine.ToDict()
+ with TestPipeline() as pipeline:
+ pairs = [(1, 2), (3, 4), (5, 6)]
+ pcoll = pipeline | 'start-pairs' >> Create(pairs)
+ result = pcoll | 'to dict' >> combine.ToDict()
- def matcher():
- def match(actual):
- equal_to([1])([len(actual)])
- equal_to(pairs)(actual[0].items())
- return match
- assert_that(result, matcher())
- pipeline.run()
+ def matcher():
+ def match(actual):
+ equal_to([1])([len(actual)])
+ equal_to(pairs)(actual[0].items())
+ return match
+ assert_that(result, matcher())
def test_combine_globally_with_default(self):
with TestPipeline() as p:
@@ -399,6 +435,31 @@
| beam.CombineGlobally(combine.MeanCombineFn()).with_fanout(11))
assert_that(result, equal_to([49.5]))
+ def test_combining_with_accumulation_mode_and_fanout(self):
+ # PCollection will contain elements from 1 to 5.
+ elements = [i for i in range(1, 6)]
+
+ ts = TestStream().advance_watermark_to(0)
+ for i in elements:
+ ts.add_elements([i])
+ ts.advance_watermark_to_infinity()
+
+ options = PipelineOptions()
+ options.view_as(StandardOptions).streaming = True
+ with TestPipeline(options=options) as p:
+ result = (p
+ | ts
+ | beam.WindowInto(
+ GlobalWindows(),
+ accumulation_mode=trigger.AccumulationMode.ACCUMULATING,
+ trigger=AfterWatermark(early=AfterAll(AfterCount(1)))
+ )
+ | beam.CombineGlobally(sum).without_defaults().with_fanout(2))
+
+ # The frings for DISCARDING mode is [1, 2, 3, 4, 5, 0, 0].
+ firings = [1, 3, 6, 10, 15, 15, 15]
+ assert_that(result, equal_to(firings))
+
def test_MeanCombineFn_combine(self):
with TestPipeline() as p:
input = (p
@@ -483,6 +544,108 @@
equal_to([('c', 3), ('c', 10), ('d', 5), ('d', 17)]),
label='sum per key')
+ # Test that three different kinds of metrics work with a customized
+ # SortedConcatWithCounters CombineFn.
+ def test_custormized_counters_in_combine_fn(self):
+ p = TestPipeline()
+ input = (p
+ | beam.Create([('key1', 'a'),
+ ('key1', 'ab'),
+ ('key1', 'abc'),
+ ('key2', 'uvxy'),
+ ('key2', 'uvxyz')]))
+
+ # The result of concatenating all values regardless of key.
+ global_concat = (input
+ | beam.Values()
+ | beam.CombineGlobally(SortedConcatWithCounters()))
+
+ # The (key, concatenated_string) pairs for all keys.
+ concat_per_key = (input
+ | beam.CombinePerKey(SortedConcatWithCounters()))
+
+ # Verify the concatenated strings are correct.
+ expected_concat_per_key = [('key1', 'aaabbc'), ('key2', 'uuvvxxyyz')]
+ assert_that(global_concat, equal_to(['aaabbcuuvvxxyyz']),
+ label='global concat')
+ assert_that(concat_per_key, equal_to(expected_concat_per_key),
+ label='concat per key')
+
+ result = p.run()
+ result.wait_until_finish()
+
+ # Verify the values of metrics are correct.
+ word_counter_filter = MetricsFilter().with_name('word_counter')
+ query_result = result.metrics().query(word_counter_filter)
+ if query_result['counters']:
+ word_counter = query_result['counters'][0]
+ self.assertEqual(word_counter.result, 5)
+
+ word_lengths_filter = MetricsFilter().with_name('word_lengths')
+ query_result = result.metrics().query(word_lengths_filter)
+ if query_result['counters']:
+ word_lengths = query_result['counters'][0]
+ self.assertEqual(word_lengths.result, 15)
+
+ word_len_dist_filter = MetricsFilter().with_name('word_len_dist')
+ query_result = result.metrics().query(word_len_dist_filter)
+ if query_result['distributions']:
+ word_len_dist = query_result['distributions'][0]
+ self.assertEqual(word_len_dist.result.mean, 3)
+
+ last_word_len_filter = MetricsFilter().with_name('last_word_len')
+ query_result = result.metrics().query(last_word_len_filter)
+ if query_result['gauges']:
+ last_word_len = query_result['gauges'][0]
+ self.assertIn(last_word_len.result.value, [1, 2, 3, 4, 5])
+
+ # Test that three different kinds of metrics work with the customized
+ # SortedConcatWithCounters CombineFn when the PCollection is empty.
+ def test_custormized_counters_in_combine_fn_empty(self):
+ p = TestPipeline()
+ input = p | beam.Create([])
+
+ # The result of concatenating all values regardless of key.
+ global_concat = (input
+ | beam.Values()
+ | beam.CombineGlobally(SortedConcatWithCounters()))
+
+ # The (key, concatenated_string) pairs for all keys.
+ concat_per_key = (input | beam.CombinePerKey(
+ SortedConcatWithCounters()))
+
+ # Verify the concatenated strings are correct.
+ assert_that(global_concat, equal_to(['']), label='global concat')
+ assert_that(concat_per_key, equal_to([]), label='concat per key')
+
+ result = p.run()
+ result.wait_until_finish()
+
+ # Verify the values of metrics are correct.
+ word_counter_filter = MetricsFilter().with_name('word_counter')
+ query_result = result.metrics().query(word_counter_filter)
+ if query_result['counters']:
+ word_counter = query_result['counters'][0]
+ self.assertEqual(word_counter.result, 0)
+
+ word_lengths_filter = MetricsFilter().with_name('word_lengths')
+ query_result = result.metrics().query(word_lengths_filter)
+ if query_result['counters']:
+ word_lengths = query_result['counters'][0]
+ self.assertEqual(word_lengths.result, 0)
+
+ word_len_dist_filter = MetricsFilter().with_name('word_len_dist')
+ query_result = result.metrics().query(word_len_dist_filter)
+ if query_result['distributions']:
+ word_len_dist = query_result['distributions'][0]
+ self.assertEqual(word_len_dist.result.count, 0)
+
+ last_word_len_filter = MetricsFilter().with_name('last_word_len')
+ query_result = result.metrics().query(last_word_len_filter)
+
+ # No element has ever been recorded.
+ self.assertFalse(query_result['gauges'])
+
class LatestTest(unittest.TestCase):
diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py
index 25cc91f..1d13f26 100644
--- a/sdks/python/apache_beam/transforms/core.py
+++ b/sdks/python/apache_beam/transforms/core.py
@@ -17,13 +17,14 @@
"""Core PTransform subclasses, such as FlatMap, GroupByKey, and Map."""
+# pytype: skip-file
+
from __future__ import absolute_import
import copy
import inspect
import logging
import random
-import re
import types
import typing
from builtins import map
@@ -48,6 +49,7 @@
from apache_beam.transforms.display import HasDisplayData
from apache_beam.transforms.ptransform import PTransform
from apache_beam.transforms.ptransform import PTransformWithSideInputs
+from apache_beam.transforms.sideinputs import get_sideinput_index
from apache_beam.transforms.userstate import StateSpec
from apache_beam.transforms.userstate import TimerSpec
from apache_beam.transforms.window import GlobalWindows
@@ -419,9 +421,10 @@
TODO(BEAM-8537): Create WatermarkEstimatorProvider to support different types.
"""
def __init__(self):
- self._watermark = None
+ self._watermark = None # type: typing.Optional[timestamp.Timestamp]
def set_watermark(self, watermark):
+ # type: (timestamp.Timestamp) -> None
"""Update tracking output_watermark with latest output_watermark.
This function is called inside an SDF.Process() to track the watermark of
output element.
@@ -437,6 +440,7 @@
self._watermark = min(self._watermark, watermark)
def current_watermark(self):
+ # type: () -> typing.Optional[timestamp.Timestamp]
"""Get current output_watermark. This function is called by system."""
return self._watermark
@@ -1304,7 +1308,9 @@
return (
common_urns.primitives.PAR_DO.urn,
beam_runner_api_pb2.ParDoPayload(
- do_fn=beam_runner_api_pb2.FunctionSpec(urn=python_urns.PICKLED_DOFN_INFO, payload=picked_pardo_fn_data),
+ do_fn=beam_runner_api_pb2.FunctionSpec(
+ urn=python_urns.PICKLED_DOFN_INFO,
+ payload=picked_pardo_fn_data),
splittable=is_splittable,
restriction_coder_id=restriction_coder_id,
state_specs={spec.name: spec.to_runner_api(context)
@@ -1334,7 +1340,7 @@
# This is an ordered list stored as a dict (see the comments in
# to_runner_api_parameter above).
indexed_side_inputs = [
- (int(re.match('side([0-9]+)(-.*)?$', tag).group(1)),
+ (get_sideinput_index(tag),
pvalue.AsSideInput.from_runner_api(si, context))
for tag, si in pardo_payload.side_inputs.items()]
result.side_inputs = [si for _, si in sorted(indexed_side_inputs)]
@@ -2268,8 +2274,8 @@
def __init__(self,
windowfn, # type: WindowFn
triggerfn=None, # type: typing.Optional[TriggerFn]
- accumulation_mode=None, # typing.Optional[beam_runner_api_pb2.AccumulationMode]
- timestamp_combiner=None, # typing.Optional[beam_runner_api_pb2.OutputTime]
+ accumulation_mode=None, # type: typing.Optional[beam_runner_api_pb2.AccumulationMode]
+ timestamp_combiner=None, # type: typing.Optional[beam_runner_api_pb2.OutputTime]
allowed_lateness=0, # type: typing.Union[int, float]
):
"""Class representing the window strategy.
diff --git a/sdks/python/apache_beam/transforms/core_test.py b/sdks/python/apache_beam/transforms/core_test.py
index 1a27bd2..3791252 100644
--- a/sdks/python/apache_beam/transforms/core_test.py
+++ b/sdks/python/apache_beam/transforms/core_test.py
@@ -17,6 +17,8 @@
"""Unit tests for core module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/transforms/create_source.py b/sdks/python/apache_beam/transforms/create_source.py
index aa26ceb..b0188c6 100644
--- a/sdks/python/apache_beam/transforms/create_source.py
+++ b/sdks/python/apache_beam/transforms/create_source.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/transforms/create_test.py b/sdks/python/apache_beam/transforms/create_test.py
index 915056f..08c5c24 100644
--- a/sdks/python/apache_beam/transforms/create_test.py
+++ b/sdks/python/apache_beam/transforms/create_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the Create and _CreateSource classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/transforms/cy_combiners.py b/sdks/python/apache_beam/transforms/cy_combiners.py
index 139b8a3..33914f9 100644
--- a/sdks/python/apache_beam/transforms/cy_combiners.py
+++ b/sdks/python/apache_beam/transforms/cy_combiners.py
@@ -22,6 +22,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/transforms/dataflow_distribution_counter_test.py b/sdks/python/apache_beam/transforms/dataflow_distribution_counter_test.py
index bedad4b..f072cd1 100644
--- a/sdks/python/apache_beam/transforms/dataflow_distribution_counter_test.py
+++ b/sdks/python/apache_beam/transforms/dataflow_distribution_counter_test.py
@@ -14,6 +14,8 @@
otherwise, test on pure python module
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/transforms/display.py b/sdks/python/apache_beam/transforms/display.py
index 8a24fe6..21e9d32 100644
--- a/sdks/python/apache_beam/transforms/display.py
+++ b/sdks/python/apache_beam/transforms/display.py
@@ -36,6 +36,8 @@
and communicate it to the API.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import calendar
@@ -130,15 +132,18 @@
@classmethod
def create_from_options(cls, pipeline_options):
- """ Creates :class:`DisplayData` from a
+ """ Creates :class:`~apache_beam.transforms.display.DisplayData` from a
:class:`~apache_beam.options.pipeline_options.PipelineOptions` instance.
- When creating :class:`DisplayData`, this method will convert the value of
- any item of a non-supported type to its string representation.
+ When creating :class:`~apache_beam.transforms.display.DisplayData`, this
+ method will convert the value of any item of a non-supported type to its
+ string representation.
The normal :meth:`.create_from()` method rejects those items.
Returns:
- DisplayData: A :class:`DisplayData` instance with populated items.
+ ~apache_beam.transforms.display.DisplayData:
+ A :class:`~apache_beam.transforms.display.DisplayData` instance with
+ populated items.
Raises:
~exceptions.ValueError: If the **has_display_data** argument is
@@ -158,10 +163,13 @@
@classmethod
def create_from(cls, has_display_data):
- """ Creates :class:`DisplayData` from a :class:`HasDisplayData` instance.
+ """ Creates :class:`~apache_beam.transforms.display.DisplayData` from a
+ :class:`HasDisplayData` instance.
Returns:
- DisplayData: A :class:`DisplayData` instance with populated items.
+ ~apache_beam.transforms.display.DisplayData:
+ A :class:`~apache_beam.transforms.display.DisplayData` instance with
+ populated items.
Raises:
~exceptions.ValueError: If the **has_display_data** argument is
diff --git a/sdks/python/apache_beam/transforms/display_test.py b/sdks/python/apache_beam/transforms/display_test.py
index bdaade6..6ead789 100644
--- a/sdks/python/apache_beam/transforms/display_test.py
+++ b/sdks/python/apache_beam/transforms/display_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the DisplayData API."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/transforms/dofn_lifecycle_test.py b/sdks/python/apache_beam/transforms/dofn_lifecycle_test.py
index fd3eee6..bcba20d 100644
--- a/sdks/python/apache_beam/transforms/dofn_lifecycle_test.py
+++ b/sdks/python/apache_beam/transforms/dofn_lifecycle_test.py
@@ -16,6 +16,8 @@
#
"""UnitTests for DoFn lifecycle and bundle methods"""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
@@ -76,12 +78,10 @@
@attr('ValidatesRunner')
class DoFnLifecycleTest(unittest.TestCase):
def test_dofn_lifecycle(self):
- p = TestPipeline()
- _ = (p
- | 'Start' >> beam.Create([1, 2, 3])
- | 'Do' >> beam.ParDo(CallSequenceEnforcingDoFn()))
- result = p.run()
- result.wait_until_finish()
+ with TestPipeline() as p:
+ _ = (p
+ | 'Start' >> beam.Create([1, 2, 3])
+ | 'Do' >> beam.ParDo(CallSequenceEnforcingDoFn()))
# Assumes that the worker is run in the same process as the test.
diff --git a/sdks/python/apache_beam/transforms/environments.py b/sdks/python/apache_beam/transforms/environments.py
index dce3b13..9a0c7e2 100644
--- a/sdks/python/apache_beam/transforms/environments.py
+++ b/sdks/python/apache_beam/transforms/environments.py
@@ -19,11 +19,23 @@
For internal use only. No backwards compatibility guarantees."""
+# pytype: skip-file
+
from __future__ import absolute_import
import json
import logging
import sys
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Optional
+from typing import Tuple
+from typing import Type
+from typing import TypeVar
+from typing import Union
+from typing import overload
from google.protobuf import message
@@ -33,11 +45,20 @@
from apache_beam.portability.api import endpoints_pb2
from apache_beam.utils import proto_utils
+if TYPE_CHECKING:
+ from apache_beam.options.pipeline_options import PipelineOptions
+ from apache_beam.runners.pipeline_context import PipelineContext
+
__all__ = ['Environment',
'DockerEnvironment', 'ProcessEnvironment', 'ExternalEnvironment',
'EmbeddedPythonEnvironment', 'EmbeddedPythonGrpcEnvironment',
'SubprocessSDKEnvironment', 'RunnerAPIEnvironmentHolder']
+T = TypeVar('T')
+EnvironmentT = TypeVar('EnvironmentT', bound='Environment')
+ConstructorFn = Callable[
+ [Optional[Any], 'PipelineContext'],
+ Any]
def looks_like_json(s):
import re
@@ -53,12 +74,52 @@
For internal use only. No backwards compatibility guarantees.
"""
- _known_urns = {}
- _urn_to_env_cls = {}
+ _known_urns = {} # type: Dict[str, Tuple[Optional[type], ConstructorFn]]
+ _urn_to_env_cls = {} # type: Dict[str, type]
def to_runner_api_parameter(self, context):
+ # type: (PipelineContext) -> Tuple[str, Optional[Union[message.Message, bytes, str]]]
raise NotImplementedError
+
+ @classmethod
+ @overload
+ def register_urn(cls,
+ urn, # type: str
+ parameter_type, # type: Type[T]
+ ):
+ # type: (...) -> Callable[[Union[type, Callable[[T, PipelineContext], Any]]], Callable[[T, PipelineContext], Any]]
+ pass
+
+ @classmethod
+ @overload
+ def register_urn(cls,
+ urn, # type: str
+ parameter_type, # type: None
+ ):
+ # type: (...) -> Callable[[Union[type, Callable[[bytes, PipelineContext], Any]]], Callable[[bytes, PipelineContext], Any]]
+ pass
+
+ @classmethod
+ @overload
+ def register_urn(cls,
+ urn, # type: str
+ parameter_type, # type: Type[T]
+ constructor # type: Callable[[T, PipelineContext], Any]
+ ):
+ # type: (...) -> None
+ pass
+
+ @classmethod
+ @overload
+ def register_urn(cls,
+ urn, # type: str
+ parameter_type, # type: None
+ constructor # type: Callable[[bytes, PipelineContext], Any]
+ ):
+ # type: (...) -> None
+ pass
+
@classmethod
def register_urn(cls, urn, parameter_type, constructor=None):
@@ -86,6 +147,7 @@
return cls._urn_to_env_cls[urn]
def to_runner_api(self, context):
+ # type: (PipelineContext) -> beam_runner_api_pb2.Environment
urn, typed_param = self.to_runner_api_parameter(context)
return beam_runner_api_pb2.Environment(
urn=urn,
@@ -97,7 +159,11 @@
)
@classmethod
- def from_runner_api(cls, proto, context):
+ def from_runner_api(cls,
+ proto, # type: Optional[beam_runner_api_pb2.FunctionSpec]
+ context # type: PipelineContext
+ ):
+ # type: (...) -> Optional[Environment]
if proto is None or not proto.urn:
return None
parameter_type, constructor = cls._known_urns[proto.urn]
@@ -113,6 +179,7 @@
@classmethod
def from_options(cls, options):
+ # type: (Type[EnvironmentT], PipelineOptions) -> EnvironmentT
"""Creates an Environment object from PipelineOptions.
Args:
@@ -146,6 +213,7 @@
return 'DockerEnvironment(container_image=%s)' % self.container_image
def to_runner_api_parameter(self, context):
+ # type: (PipelineContext) -> Tuple[str, beam_runner_api_pb2.DockerPayload]
return (common_urns.environments.DOCKER.urn,
beam_runner_api_pb2.DockerPayload(
container_image=self.container_image))
@@ -156,6 +224,7 @@
@classmethod
def from_options(cls, options):
+ # type: (PipelineOptions) -> DockerEnvironment
return cls(container_image=options.environment_config)
@staticmethod
@@ -210,6 +279,7 @@
return 'ProcessEnvironment(%s)' % ','.join(repr_parts)
def to_runner_api_parameter(self, context):
+ # type: (PipelineContext) -> Tuple[str, beam_runner_api_pb2.ProcessPayload]
return (common_urns.environments.PROCESS.urn,
beam_runner_api_pb2.ProcessPayload(
os=self.os,
@@ -255,6 +325,7 @@
return 'ExternalEnvironment(url=%s,params=%s)' % (self.url, self.params)
def to_runner_api_parameter(self, context):
+ # type: (PipelineContext) -> Tuple[str, beam_runner_api_pb2.ExternalPayload]
return (common_urns.environments.EXTERNAL.urn,
beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(url=self.url),
@@ -295,6 +366,7 @@
return hash(self.__class__)
def to_runner_api_parameter(self, context):
+ # type: (PipelineContext) -> Tuple[str, None]
return python_urns.EMBEDDED_PYTHON, None
@staticmethod
@@ -336,6 +408,7 @@
return 'EmbeddedPythonGrpcEnvironment(%s)' % ','.join(repr_parts)
def to_runner_api_parameter(self, context):
+ # type: (PipelineContext) -> Tuple[str, bytes]
params = {}
if self.state_cache_size is not None:
params['state_cache_size'] = self.state_cache_size
@@ -399,9 +472,10 @@
return hash((self.__class__, self.command_string))
def __repr__(self):
- return 'SubprocessSDKEnvironment(command_string=%s)' % self.container_string
+ return 'SubprocessSDKEnvironment(command_string=%s)' % self.command_string
def to_runner_api_parameter(self, context):
+ # type: (PipelineContext) -> Tuple[str, bytes]
return python_urns.SUBPROCESS_SDK, self.command_string.encode('utf-8')
@staticmethod
diff --git a/sdks/python/apache_beam/transforms/environments_test.py b/sdks/python/apache_beam/transforms/environments_test.py
index b28af24..b79ca79 100644
--- a/sdks/python/apache_beam/transforms/environments_test.py
+++ b/sdks/python/apache_beam/transforms/environments_test.py
@@ -18,6 +18,8 @@
"""Unit tests for the transform.environments classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/transforms/external.py b/sdks/python/apache_beam/transforms/external.py
index bf45802..882d1ea 100644
--- a/sdks/python/apache_beam/transforms/external.py
+++ b/sdks/python/apache_beam/transforms/external.py
@@ -19,6 +19,8 @@
No backward compatibility guarantees. Everything in this module is experimental.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/transforms/external_test.py b/sdks/python/apache_beam/transforms/external_test.py
index ff34431..a2b7800 100644
--- a/sdks/python/apache_beam/transforms/external_test.py
+++ b/sdks/python/apache_beam/transforms/external_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the transform.external classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
import argparse
@@ -392,12 +394,8 @@
p
| beam.Create(list('aaabccxyyzzz'))
| beam.Map(unicode)
- # TODO(BEAM-6587): Use strings directly rather than ints.
- | beam.Map(lambda x: int(ord(x)))
| beam.ExternalTransform(TEST_FILTER_URN, b'middle', expansion_service)
| beam.ExternalTransform(TEST_COUNT_URN, None, expansion_service)
- # # TODO(BEAM-6587): Remove when above is removed.
- | beam.Map(lambda kv: (chr(kv[0]), kv[1]))
| beam.Map(lambda kv: '%s: %s' % kv))
assert_that(res, equal_to(['a: 3', 'b: 1', 'c: 2']))
diff --git a/sdks/python/apache_beam/transforms/external_test_it.py b/sdks/python/apache_beam/transforms/external_test_it.py
index 97f857c..d597886 100644
--- a/sdks/python/apache_beam/transforms/external_test_it.py
+++ b/sdks/python/apache_beam/transforms/external_test_it.py
@@ -17,6 +17,8 @@
"""Integration tests for cross-language transform expansion."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/transforms/external_test_py3.py b/sdks/python/apache_beam/transforms/external_test_py3.py
index c2e7f87..3f7b2a4 100644
--- a/sdks/python/apache_beam/transforms/external_test_py3.py
+++ b/sdks/python/apache_beam/transforms/external_test_py3.py
@@ -17,6 +17,8 @@
"""Unit tests for the transform.external classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
import typing
diff --git a/sdks/python/apache_beam/transforms/external_test_py37.py b/sdks/python/apache_beam/transforms/external_test_py37.py
index e01f532..8399b8b 100644
--- a/sdks/python/apache_beam/transforms/external_test_py37.py
+++ b/sdks/python/apache_beam/transforms/external_test_py37.py
@@ -17,6 +17,8 @@
"""Unit tests for the transform.external classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
import dataclasses
diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py
index dcc3d96..66ee5f3 100644
--- a/sdks/python/apache_beam/transforms/ptransform.py
+++ b/sdks/python/apache_beam/transforms/ptransform.py
@@ -34,6 +34,8 @@
FlatMap processing functions.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import contextlib
diff --git a/sdks/python/apache_beam/transforms/ptransform_test.py b/sdks/python/apache_beam/transforms/ptransform_test.py
index fac509e..ef3932e 100644
--- a/sdks/python/apache_beam/transforms/ptransform_test.py
+++ b/sdks/python/apache_beam/transforms/ptransform_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the PTransform and descendants."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -116,11 +118,10 @@
def process(self, element, addon):
return [element + addon]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
- result = pcoll | 'Do' >> beam.ParDo(AddNDoFn(), 10)
- assert_that(result, equal_to([11, 12, 13]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ result = pcoll | 'Do' >> beam.ParDo(AddNDoFn(), 10)
+ assert_that(result, equal_to([11, 12, 13]))
def test_do_with_unconstructed_do_fn(self):
class MyDoFn(beam.DoFn):
@@ -128,81 +129,74 @@
def process(self):
pass
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
with self.assertRaises(ValueError):
- pcoll | 'Do' >> beam.ParDo(MyDoFn) # Note the lack of ()'s
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ pcoll | 'Do' >> beam.ParDo(MyDoFn) # Note the lack of ()'s
def test_do_with_callable(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
- result = pcoll | 'Do' >> beam.FlatMap(lambda x, addon: [x + addon], 10)
- assert_that(result, equal_to([11, 12, 13]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ result = pcoll | 'Do' >> beam.FlatMap(lambda x, addon: [x + addon], 10)
+ assert_that(result, equal_to([11, 12, 13]))
def test_do_with_side_input_as_arg(self):
- pipeline = TestPipeline()
- side = pipeline | 'Side' >> beam.Create([10])
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
- result = pcoll | 'Do' >> beam.FlatMap(
- lambda x, addon: [x + addon], pvalue.AsSingleton(side))
- assert_that(result, equal_to([11, 12, 13]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ side = pipeline | 'Side' >> beam.Create([10])
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ result = pcoll | 'Do' >> beam.FlatMap(
+ lambda x, addon: [x + addon], pvalue.AsSingleton(side))
+ assert_that(result, equal_to([11, 12, 13]))
def test_do_with_side_input_as_keyword_arg(self):
- pipeline = TestPipeline()
- side = pipeline | 'Side' >> beam.Create([10])
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
- result = pcoll | 'Do' >> beam.FlatMap(
- lambda x, addon: [x + addon], addon=pvalue.AsSingleton(side))
- assert_that(result, equal_to([11, 12, 13]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ side = pipeline | 'Side' >> beam.Create([10])
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ result = pcoll | 'Do' >> beam.FlatMap(
+ lambda x, addon: [x + addon], addon=pvalue.AsSingleton(side))
+ assert_that(result, equal_to([11, 12, 13]))
def test_do_with_do_fn_returning_string_raises_warning(self):
- pipeline = TestPipeline()
- pipeline._options.view_as(TypeOptions).runtime_type_check = True
- pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])
- pcoll | 'Do' >> beam.FlatMap(lambda x: x + '1')
-
- # Since the DoFn directly returns a string we should get an error warning
- # us.
with self.assertRaises(typehints.TypeCheckError) as cm:
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pipeline._options.view_as(TypeOptions).runtime_type_check = True
+ pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])
+ pcoll | 'Do' >> beam.FlatMap(lambda x: x + '1')
+
+ # Since the DoFn directly returns a string we should get an
+ # error warning us when the pipeliene runs.
expected_error_prefix = ('Returning a str from a ParDo or FlatMap '
'is discouraged.')
self.assertStartswith(cm.exception.args[0], expected_error_prefix)
def test_do_with_do_fn_returning_dict_raises_warning(self):
- pipeline = TestPipeline()
- pipeline._options.view_as(TypeOptions).runtime_type_check = True
- pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])
- pcoll | 'Do' >> beam.FlatMap(lambda x: {x: '1'})
-
- # Since the DoFn directly returns a dict we should get an error warning
- # us.
with self.assertRaises(typehints.TypeCheckError) as cm:
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pipeline._options.view_as(TypeOptions).runtime_type_check = True
+ pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])
+ pcoll | 'Do' >> beam.FlatMap(lambda x: {x: '1'})
+
+ # Since the DoFn directly returns a dict we should get an error warning
+ # us when the pipeliene runs.
expected_error_prefix = ('Returning a dict from a ParDo or FlatMap '
'is discouraged.')
self.assertStartswith(cm.exception.args[0], expected_error_prefix)
def test_do_with_multiple_outputs_maintains_unique_name(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
- r1 = pcoll | 'A' >> beam.FlatMap(lambda x: [x + 1]).with_outputs(main='m')
- r2 = pcoll | 'B' >> beam.FlatMap(lambda x: [x + 2]).with_outputs(main='m')
- assert_that(r1.m, equal_to([2, 3, 4]), label='r1')
- assert_that(r2.m, equal_to([3, 4, 5]), label='r2')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ r1 = pcoll | 'A' >> beam.FlatMap(lambda x: [x + 1]).with_outputs(main='m')
+ r2 = pcoll | 'B' >> beam.FlatMap(lambda x: [x + 2]).with_outputs(main='m')
+ assert_that(r1.m, equal_to([2, 3, 4]), label='r1')
+ assert_that(r2.m, equal_to([3, 4, 5]), label='r2')
@attr('ValidatesRunner')
def test_impulse(self):
- pipeline = TestPipeline()
- result = pipeline | beam.Impulse() | beam.Map(lambda _: 0)
- assert_that(result, equal_to([0]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ result = pipeline | beam.Impulse() | beam.Map(lambda _: 0)
+ assert_that(result, equal_to([0]))
# TODO(BEAM-3544): Disable this test in streaming temporarily.
# Remove sickbay-streaming tag after it's resolved.
@@ -244,14 +238,13 @@
else:
yield pvalue.TaggedOutput('odd', element)
- pipeline = TestPipeline()
- nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
- results = nums | 'ClassifyNumbers' >> beam.ParDo(
- SomeDoFn()).with_outputs('odd', 'even', main='main')
- assert_that(results.main, equal_to([1, 2, 3, 4]))
- assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
- assert_that(results.even, equal_to([2, 4]), label='assert:even')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
+ results = nums | 'ClassifyNumbers' >> beam.ParDo(
+ SomeDoFn()).with_outputs('odd', 'even', main='main')
+ assert_that(results.main, equal_to([1, 2, 3, 4]))
+ assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
+ assert_that(results.even, equal_to([2, 4]), label='assert:even')
@attr('ValidatesRunner')
def test_par_do_with_multiple_outputs_and_using_return(self):
@@ -260,55 +253,51 @@
return [v, pvalue.TaggedOutput('even', v)]
return [v, pvalue.TaggedOutput('odd', v)]
- pipeline = TestPipeline()
- nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
- results = nums | 'ClassifyNumbers' >> beam.FlatMap(
- some_fn).with_outputs('odd', 'even', main='main')
- assert_that(results.main, equal_to([1, 2, 3, 4]))
- assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
- assert_that(results.even, equal_to([2, 4]), label='assert:even')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
+ results = nums | 'ClassifyNumbers' >> beam.FlatMap(
+ some_fn).with_outputs('odd', 'even', main='main')
+ assert_that(results.main, equal_to([1, 2, 3, 4]))
+ assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
+ assert_that(results.even, equal_to([2, 4]), label='assert:even')
@attr('ValidatesRunner')
def test_undeclared_outputs(self):
- pipeline = TestPipeline()
- nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
- results = nums | 'ClassifyNumbers' >> beam.FlatMap(
- lambda x: [x,
- pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x),
- pvalue.TaggedOutput('extra', x)]
- ).with_outputs()
- assert_that(results[None], equal_to([1, 2, 3, 4]))
- assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
- assert_that(results.even, equal_to([2, 4]), label='assert:even')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
+ results = nums | 'ClassifyNumbers' >> beam.FlatMap(
+ lambda x: [x,
+ pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x),
+ pvalue.TaggedOutput('extra', x)]
+ ).with_outputs()
+ assert_that(results[None], equal_to([1, 2, 3, 4]))
+ assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
+ assert_that(results.even, equal_to([2, 4]), label='assert:even')
@attr('ValidatesRunner')
def test_multiple_empty_outputs(self):
- pipeline = TestPipeline()
- nums = pipeline | 'Some Numbers' >> beam.Create([1, 3, 5])
- results = nums | 'ClassifyNumbers' >> beam.FlatMap(
- lambda x: [x,
- pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x)]
- ).with_outputs()
- assert_that(results[None], equal_to([1, 3, 5]))
- assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd')
- assert_that(results.even, equal_to([]), label='assert:even')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ nums = pipeline | 'Some Numbers' >> beam.Create([1, 3, 5])
+ results = nums | 'ClassifyNumbers' >> beam.FlatMap(
+ lambda x: [x,
+ pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x)]
+ ).with_outputs()
+ assert_that(results[None], equal_to([1, 3, 5]))
+ assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd')
+ assert_that(results.even, equal_to([]), label='assert:even')
def test_do_requires_do_fn_returning_iterable(self):
# This function is incorrect because it returns an object that isn't an
# iterable.
def incorrect_par_do_fn(x):
return x + 5
- pipeline = TestPipeline()
- pipeline._options.view_as(TypeOptions).runtime_type_check = True
- pcoll = pipeline | 'Start' >> beam.Create([2, 9, 3])
- pcoll | 'Do' >> beam.FlatMap(incorrect_par_do_fn)
- # It's a requirement that all user-defined functions to a ParDo return
- # an iterable.
with self.assertRaises(typehints.TypeCheckError) as cm:
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pipeline._options.view_as(TypeOptions).runtime_type_check = True
+ pcoll = pipeline | 'Start' >> beam.Create([2, 9, 3])
+ pcoll | 'Do' >> beam.FlatMap(incorrect_par_do_fn)
+ # It's a requirement that all user-defined functions to a ParDo return
+ # an iterable.
expected_error_prefix = 'FlatMap and ParDo must return an iterable.'
self.assertStartswith(cm.exception.args[0], expected_error_prefix)
@@ -321,19 +310,18 @@
def finish_bundle(self):
yield WindowedValue('finish', -1, [window.GlobalWindow()])
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
- result = pcoll | 'Do' >> beam.ParDo(MyDoFn())
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ result = pcoll | 'Do' >> beam.ParDo(MyDoFn())
- # May have many bundles, but each has a start and finish.
- def matcher():
- def match(actual):
- equal_to(['finish'])(list(set(actual)))
- equal_to([1])([actual.count('finish')])
- return match
+ # May have many bundles, but each has a start and finish.
+ def matcher():
+ def match(actual):
+ equal_to(['finish'])(list(set(actual)))
+ equal_to([1])([actual.count('finish')])
+ return match
- assert_that(result, matcher())
- pipeline.run()
+ assert_that(result, matcher())
def test_do_fn_with_windowing_in_finish_bundle(self):
windowfn = window.FixedWindows(2)
@@ -373,19 +361,18 @@
yield 'started'
self.state = 'process'
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
- result = pcoll | 'Do' >> beam.ParDo(MyDoFn())
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
+ result = pcoll | 'Do' >> beam.ParDo(MyDoFn())
- # May have many bundles, but each has a start and finish.
- def matcher():
- def match(actual):
- equal_to(['started'])(list(set(actual)))
- equal_to([1])([actual.count('started')])
- return match
+ # May have many bundles, but each has a start and finish.
+ def matcher():
+ def match(actual):
+ equal_to(['started'])(list(set(actual)))
+ equal_to([1])([actual.count('started')])
+ return match
- assert_that(result, matcher())
- pipeline.run()
+ assert_that(result, matcher())
def test_do_fn_with_start_error(self):
class MyDoFn(beam.DoFn):
@@ -395,17 +382,15 @@
def process(self, element):
pass
- pipeline = TestPipeline()
- pipeline | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(MyDoFn())
with self.assertRaises(RuntimeError):
- pipeline.run()
+ with TestPipeline() as p:
+ p | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(MyDoFn())
def test_filter(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3, 4])
- result = pcoll | 'Filter' >> beam.Filter(lambda x: x % 2 == 0)
- assert_that(result, equal_to([2, 4]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3, 4])
+ result = pcoll | 'Filter' >> beam.Filter(lambda x: x % 2 == 0)
+ assert_that(result, equal_to([2, 4]))
class _MeanCombineFn(beam.CombineFn):
@@ -428,68 +413,62 @@
def test_combine_with_combine_fn(self):
vals = [1, 2, 3, 4, 5, 6, 7]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(vals)
- result = pcoll | 'Mean' >> beam.CombineGlobally(self._MeanCombineFn())
- assert_that(result, equal_to([sum(vals) // len(vals)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(vals)
+ result = pcoll | 'Mean' >> beam.CombineGlobally(self._MeanCombineFn())
+ assert_that(result, equal_to([sum(vals) // len(vals)]))
def test_combine_with_callable(self):
vals = [1, 2, 3, 4, 5, 6, 7]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(vals)
- result = pcoll | beam.CombineGlobally(sum)
- assert_that(result, equal_to([sum(vals)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(vals)
+ result = pcoll | beam.CombineGlobally(sum)
+ assert_that(result, equal_to([sum(vals)]))
def test_combine_with_side_input_as_arg(self):
values = [1, 2, 3, 4, 5, 6, 7]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(values)
- divisor = pipeline | 'Divisor' >> beam.Create([2])
- result = pcoll | 'Max' >> beam.CombineGlobally(
- # Multiples of divisor only.
- lambda vals, d: max(v for v in vals if v % d == 0),
- pvalue.AsSingleton(divisor)).without_defaults()
- filt_vals = [v for v in values if v % 2 == 0]
- assert_that(result, equal_to([max(filt_vals)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(values)
+ divisor = pipeline | 'Divisor' >> beam.Create([2])
+ result = pcoll | 'Max' >> beam.CombineGlobally(
+ # Multiples of divisor only.
+ lambda vals, d: max(v for v in vals if v % d == 0),
+ pvalue.AsSingleton(divisor)).without_defaults()
+ filt_vals = [v for v in values if v % 2 == 0]
+ assert_that(result, equal_to([max(filt_vals)]))
def test_combine_per_key_with_combine_fn(self):
vals_1 = [1, 2, 3, 4, 5, 6, 7]
vals_2 = [2, 4, 6, 8, 10, 12, 14]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
- [('b', x) for x in vals_2]))
- result = pcoll | 'Mean' >> beam.CombinePerKey(self._MeanCombineFn())
- assert_that(result, equal_to([('a', sum(vals_1) // len(vals_1)),
- ('b', sum(vals_2) // len(vals_2))]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
+ [('b', x) for x in vals_2]))
+ result = pcoll | 'Mean' >> beam.CombinePerKey(self._MeanCombineFn())
+ assert_that(result, equal_to([('a', sum(vals_1) // len(vals_1)),
+ ('b', sum(vals_2) // len(vals_2))]))
def test_combine_per_key_with_callable(self):
vals_1 = [1, 2, 3, 4, 5, 6, 7]
vals_2 = [2, 4, 6, 8, 10, 12, 14]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
- [('b', x) for x in vals_2]))
- result = pcoll | beam.CombinePerKey(sum)
- assert_that(result, equal_to([('a', sum(vals_1)), ('b', sum(vals_2))]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
+ [('b', x) for x in vals_2]))
+ result = pcoll | beam.CombinePerKey(sum)
+ assert_that(result, equal_to([('a', sum(vals_1)), ('b', sum(vals_2))]))
def test_combine_per_key_with_side_input_as_arg(self):
vals_1 = [1, 2, 3, 4, 5, 6, 7]
vals_2 = [2, 4, 6, 8, 10, 12, 14]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
- [('b', x) for x in vals_2]))
- divisor = pipeline | 'Divisor' >> beam.Create([2])
- result = pcoll | beam.CombinePerKey(
- lambda vals, d: max(v for v in vals if v % d == 0),
- pvalue.AsSingleton(divisor)) # Multiples of divisor only.
- m_1 = max(v for v in vals_1 if v % 2 == 0)
- m_2 = max(v for v in vals_2 if v % 2 == 0)
- assert_that(result, equal_to([('a', m_1), ('b', m_2)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
+ [('b', x) for x in vals_2]))
+ divisor = pipeline | 'Divisor' >> beam.Create([2])
+ result = pcoll | beam.CombinePerKey(
+ lambda vals, d: max(v for v in vals if v % d == 0),
+ pvalue.AsSingleton(divisor)) # Multiples of divisor only.
+ m_1 = max(v for v in vals_1 if v % 2 == 0)
+ m_2 = max(v for v in vals_2 if v % 2 == 0)
+ assert_that(result, equal_to([('a', m_1), ('b', m_2)]))
def test_group_by_key(self):
pipeline = TestPipeline()
@@ -509,13 +488,12 @@
sum_val += sum(value_list)
return [(key, sum_val)]
- pipeline = TestPipeline()
- pcoll = pipeline | 'start' >> beam.Create(
- [(1, 1), (1, 2), (1, 3), (1, 4)])
- result = (pcoll | 'Group' >> beam.GroupByKey()
- | 'Reiteration-Sum' >> beam.ParDo(MyDoFn()))
- assert_that(result, equal_to([(1, 170)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'start' >> beam.Create(
+ [(1, 1), (1, 2), (1, 3), (1, 4)])
+ result = (pcoll | 'Group' >> beam.GroupByKey()
+ | 'Reiteration-Sum' >> beam.ParDo(MyDoFn()))
+ assert_that(result, equal_to([(1, 170)]))
def test_partition_with_partition_fn(self):
@@ -524,36 +502,33 @@
def partition_for(self, element, num_partitions, offset):
return (element % 3) + offset
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
- # Attempt nominal partition operation.
- partitions = pcoll | 'Part 1' >> beam.Partition(SomePartitionFn(), 4, 1)
- assert_that(partitions[0], equal_to([]))
- assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')
- assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')
- assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
+ # Attempt nominal partition operation.
+ partitions = pcoll | 'Part 1' >> beam.Partition(SomePartitionFn(), 4, 1)
+ assert_that(partitions[0], equal_to([]))
+ assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')
+ assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')
+ assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')
# Check that a bad partition label will yield an error. For the
# DirectRunner, this error manifests as an exception.
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
- partitions = pcoll | 'Part 2' >> beam.Partition(SomePartitionFn(), 4, 10000)
with self.assertRaises(ValueError):
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
+ partitions = pcoll | beam.Partition(SomePartitionFn(), 4, 10000)
def test_partition_with_callable(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
- partitions = (
- pcoll | 'part' >> beam.Partition(
- lambda e, n, offset: (e % 3) + offset, 4,
- 1))
- assert_that(partitions[0], equal_to([]))
- assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')
- assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')
- assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
+ partitions = (
+ pcoll | 'part' >> beam.Partition(
+ lambda e, n, offset: (e % 3) + offset, 4,
+ 1))
+ assert_that(partitions[0], equal_to([]))
+ assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')
+ assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')
+ assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')
def test_partition_followed_by_flatten_and_groupbykey(self):
"""Regression test for an issue with how partitions are handled."""
@@ -568,56 +543,50 @@
@attr('ValidatesRunner')
def test_flatten_pcollections(self):
- pipeline = TestPipeline()
- pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
- pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
- result = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten()
- assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
+ pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
+ result = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten()
+ assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))
def test_flatten_no_pcollections(self):
- pipeline = TestPipeline()
- with self.assertRaises(ValueError):
- () | 'PipelineArgMissing' >> beam.Flatten()
- result = () | 'Empty' >> beam.Flatten(pipeline=pipeline)
- assert_that(result, equal_to([]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ with self.assertRaises(ValueError):
+ () | 'PipelineArgMissing' >> beam.Flatten()
+ result = () | 'Empty' >> beam.Flatten(pipeline=pipeline)
+ assert_that(result, equal_to([]))
@attr('ValidatesRunner')
def test_flatten_one_single_pcollection(self):
- pipeline = TestPipeline()
- input = [0, 1, 2, 3]
- pcoll = pipeline | 'Input' >> beam.Create(input)
- result = (pcoll,)| 'Single Flatten' >> beam.Flatten()
- assert_that(result, equal_to(input))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ input = [0, 1, 2, 3]
+ pcoll = pipeline | 'Input' >> beam.Create(input)
+ result = (pcoll,)| 'Single Flatten' >> beam.Flatten()
+ assert_that(result, equal_to(input))
# TODO(BEAM-9002): Does not work in streaming mode on Dataflow.
@attr('ValidatesRunner', 'sickbay-streaming')
def test_flatten_same_pcollections(self):
- pipeline = TestPipeline()
- pc = pipeline | beam.Create(['a', 'b'])
- assert_that((pc, pc, pc) | beam.Flatten(), equal_to(['a', 'b'] * 3))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pc = pipeline | beam.Create(['a', 'b'])
+ assert_that((pc, pc, pc) | beam.Flatten(), equal_to(['a', 'b'] * 3))
def test_flatten_pcollections_in_iterable(self):
- pipeline = TestPipeline()
- pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
- pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
- result = [pcoll for pcoll in (pcoll_1, pcoll_2)] | beam.Flatten()
- assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
+ pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
+ result = [pcoll for pcoll in (pcoll_1, pcoll_2)] | beam.Flatten()
+ assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))
@attr('ValidatesRunner')
def test_flatten_a_flattened_pcollection(self):
- pipeline = TestPipeline()
- pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
- pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
- pcoll_3 = pipeline | 'Start 3' >> beam.Create([8, 9])
- pcoll_12 = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten()
- pcoll_123 = (pcoll_12, pcoll_3) | 'Flatten again' >> beam.Flatten()
- assert_that(pcoll_123, equal_to([x for x in range(10)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
+ pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
+ pcoll_3 = pipeline | 'Start 3' >> beam.Create([8, 9])
+ pcoll_12 = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten()
+ pcoll_123 = (pcoll_12, pcoll_3) | 'Flatten again' >> beam.Flatten()
+ assert_that(pcoll_123, equal_to([x for x in range(10)]))
def test_flatten_input_type_must_be_iterable(self):
# Inputs to flatten *must* be an iterable.
@@ -633,21 +602,20 @@
@attr('ValidatesRunner')
def test_flatten_multiple_pcollections_having_multiple_consumers(self):
- pipeline = TestPipeline()
- input = pipeline | 'Start' >> beam.Create(['AA', 'BBB', 'CC'])
+ with TestPipeline() as pipeline:
+ input = pipeline | 'Start' >> beam.Create(['AA', 'BBB', 'CC'])
- def split_even_odd(element):
- tag = 'even_length' if len(element) % 2 == 0 else 'odd_length'
- return pvalue.TaggedOutput(tag, element)
+ def split_even_odd(element):
+ tag = 'even_length' if len(element) % 2 == 0 else 'odd_length'
+ return pvalue.TaggedOutput(tag, element)
- even_length, odd_length = (input | beam.Map(split_even_odd)
- .with_outputs('even_length', 'odd_length'))
- merged = (even_length, odd_length) | 'Flatten' >> beam.Flatten()
+ even_length, odd_length = (input | beam.Map(split_even_odd)
+ .with_outputs('even_length', 'odd_length'))
+ merged = (even_length, odd_length) | 'Flatten' >> beam.Flatten()
- assert_that(merged, equal_to(['AA', 'BBB', 'CC']))
- assert_that(even_length, equal_to(['AA', 'CC']), label='assert:even')
- assert_that(odd_length, equal_to(['BBB']), label='assert:odd')
- pipeline.run()
+ assert_that(merged, equal_to(['AA', 'BBB', 'CC']))
+ assert_that(even_length, equal_to(['AA', 'CC']), label='assert:even')
+ assert_that(odd_length, equal_to(['BBB']), label='assert:odd')
def test_co_group_by_key_on_list(self):
pipeline = TestPipeline()
@@ -688,12 +656,10 @@
pipeline.run()
def test_group_by_key_input_must_be_kv_pairs(self):
- pipeline = TestPipeline()
- pcolls = pipeline | 'A' >> beam.Create([1, 2, 3, 4, 5])
-
with self.assertRaises(typehints.TypeCheckError) as e:
- pcolls | 'D' >> beam.GroupByKey()
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcolls = pipeline | 'A' >> beam.Create([1, 2, 3, 4, 5])
+ pcolls | 'D' >> beam.GroupByKey()
self.assertStartswith(
e.exception.args[0],
@@ -701,58 +667,52 @@
'Tuple[TypeVariable[K], TypeVariable[V]]')
def test_group_by_key_only_input_must_be_kv_pairs(self):
- pipeline = TestPipeline()
- pcolls = pipeline | 'A' >> beam.Create(['a', 'b', 'f'])
with self.assertRaises(typehints.TypeCheckError) as cm:
- pcolls | 'D' >> _GroupByKeyOnly()
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcolls = pipeline | 'A' >> beam.Create(['a', 'b', 'f'])
+ pcolls | 'D' >> _GroupByKeyOnly()
expected_error_prefix = ('Input type hint violation at D: expected '
'Tuple[TypeVariable[K], TypeVariable[V]]')
self.assertStartswith(cm.exception.args[0], expected_error_prefix)
def test_keys_and_values(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(
- [(3, 1), (2, 1), (1, 1), (3, 2), (2, 2), (3, 3)])
- keys = pcoll.apply(beam.Keys('keys'))
- vals = pcoll.apply(beam.Values('vals'))
- assert_that(keys, equal_to([1, 2, 2, 3, 3, 3]), label='assert:keys')
- assert_that(vals, equal_to([1, 1, 1, 2, 2, 3]), label='assert:vals')
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(
+ [(3, 1), (2, 1), (1, 1), (3, 2), (2, 2), (3, 3)])
+ keys = pcoll.apply(beam.Keys('keys'))
+ vals = pcoll.apply(beam.Values('vals'))
+ assert_that(keys, equal_to([1, 2, 2, 3, 3, 3]), label='assert:keys')
+ assert_that(vals, equal_to([1, 1, 1, 2, 2, 3]), label='assert:vals')
def test_kv_swap(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(
- [(6, 3), (1, 2), (7, 1), (5, 2), (3, 2)])
- result = pcoll.apply(beam.KvSwap(), label='swap')
- assert_that(result, equal_to([(1, 7), (2, 1), (2, 3), (2, 5), (3, 6)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(
+ [(6, 3), (1, 2), (7, 1), (5, 2), (3, 2)])
+ result = pcoll.apply(beam.KvSwap(), label='swap')
+ assert_that(result, equal_to([(1, 7), (2, 1), (2, 3), (2, 5), (3, 6)]))
def test_distinct(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(
- [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel'])
- result = pcoll.apply(beam.Distinct())
- assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel']))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(
+ [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel'])
+ result = pcoll.apply(beam.Distinct())
+ assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel']))
def test_remove_duplicates(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(
- [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel'])
- result = pcoll.apply(beam.RemoveDuplicates())
- assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel']))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(
+ [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel'])
+ result = pcoll.apply(beam.RemoveDuplicates())
+ assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel']))
def test_chained_ptransforms(self):
- pipeline = TestPipeline()
- t = (beam.Map(lambda x: (x, 1))
- | beam.GroupByKey()
- | beam.Map(lambda x_ones: (x_ones[0], sum(x_ones[1]))))
- result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t
- assert_that(result, equal_to([('a', 2), ('b', 1)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ t = (beam.Map(lambda x: (x, 1))
+ | beam.GroupByKey()
+ | beam.Map(lambda x_ones: (x_ones[0], sum(x_ones[1]))))
+ result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t
+ assert_that(result, equal_to([('a', 2), ('b', 1)]))
def test_apply_to_list(self):
self.assertCountEqual(
@@ -848,47 +808,43 @@
def test_chained_ptransforms(self):
"""Tests that chaining gets proper nesting."""
- pipeline = TestPipeline()
- map1 = 'Map1' >> beam.Map(lambda x: (x, 1))
- gbk = 'Gbk' >> beam.GroupByKey()
- map2 = 'Map2' >> beam.Map(lambda x_ones2: (x_ones2[0], sum(x_ones2[1])))
- t = (map1 | gbk | map2)
- result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t
- self.assertTrue('Map1|Gbk|Map2/Map1' in pipeline.applied_labels)
- self.assertTrue('Map1|Gbk|Map2/Gbk' in pipeline.applied_labels)
- self.assertTrue('Map1|Gbk|Map2/Map2' in pipeline.applied_labels)
- assert_that(result, equal_to([('a', 2), ('b', 1)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ map1 = 'Map1' >> beam.Map(lambda x: (x, 1))
+ gbk = 'Gbk' >> beam.GroupByKey()
+ map2 = 'Map2' >> beam.Map(lambda x_ones2: (x_ones2[0], sum(x_ones2[1])))
+ t = (map1 | gbk | map2)
+ result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t
+ self.assertTrue('Map1|Gbk|Map2/Map1' in pipeline.applied_labels)
+ self.assertTrue('Map1|Gbk|Map2/Gbk' in pipeline.applied_labels)
+ self.assertTrue('Map1|Gbk|Map2/Map2' in pipeline.applied_labels)
+ assert_that(result, equal_to([('a', 2), ('b', 1)]))
def test_apply_custom_transform_without_label(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3])
- custom = PTransformLabelsTest.CustomTransform()
- result = pipeline.apply(custom, pcoll)
- self.assertTrue('CustomTransform' in pipeline.applied_labels)
- self.assertTrue('CustomTransform/*Do*' in pipeline.applied_labels)
- assert_that(result, equal_to([2, 3, 4]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3])
+ custom = PTransformLabelsTest.CustomTransform()
+ result = pipeline.apply(custom, pcoll)
+ self.assertTrue('CustomTransform' in pipeline.applied_labels)
+ self.assertTrue('CustomTransform/*Do*' in pipeline.applied_labels)
+ assert_that(result, equal_to([2, 3, 4]))
def test_apply_custom_transform_with_label(self):
- pipeline = TestPipeline()
- pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3])
- custom = PTransformLabelsTest.CustomTransform('*Custom*')
- result = pipeline.apply(custom, pcoll)
- self.assertTrue('*Custom*' in pipeline.applied_labels)
- self.assertTrue('*Custom*/*Do*' in pipeline.applied_labels)
- assert_that(result, equal_to([2, 3, 4]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3])
+ custom = PTransformLabelsTest.CustomTransform('*Custom*')
+ result = pipeline.apply(custom, pcoll)
+ self.assertTrue('*Custom*' in pipeline.applied_labels)
+ self.assertTrue('*Custom*/*Do*' in pipeline.applied_labels)
+ assert_that(result, equal_to([2, 3, 4]))
def test_combine_without_label(self):
vals = [1, 2, 3, 4, 5, 6, 7]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(vals)
- combine = beam.CombineGlobally(sum)
- result = pcoll | combine
- self.assertTrue('CombineGlobally(sum)' in pipeline.applied_labels)
- assert_that(result, equal_to([sum(vals)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(vals)
+ combine = beam.CombineGlobally(sum)
+ result = pcoll | combine
+ self.assertTrue('CombineGlobally(sum)' in pipeline.applied_labels)
+ assert_that(result, equal_to([sum(vals)]))
def test_apply_ptransform_using_decorator(self):
pipeline = TestPipeline()
@@ -901,13 +857,12 @@
def test_combine_with_label(self):
vals = [1, 2, 3, 4, 5, 6, 7]
- pipeline = TestPipeline()
- pcoll = pipeline | 'Start' >> beam.Create(vals)
- combine = '*Sum*' >> beam.CombineGlobally(sum)
- result = pcoll | combine
- self.assertTrue('*Sum*' in pipeline.applied_labels)
- assert_that(result, equal_to([sum(vals)]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ pcoll = pipeline | 'Start' >> beam.Create(vals)
+ combine = '*Sum*' >> beam.CombineGlobally(sum)
+ result = pcoll | combine
+ self.assertTrue('*Sum*' in pipeline.applied_labels)
+ assert_that(result, equal_to([sum(vals)]))
def check_label(self, ptransform, expected_label):
pipeline = TestPipeline()
@@ -2224,11 +2179,10 @@
def MyTransform(pcoll):
return pcoll | beam.ParDo(lambda x: [x]).with_output_types(int)
- p = TestPipeline()
- _ = (p
- | beam.Create([1, 2])
- | MyTransform().with_output_types(int))
- p.run()
+ with TestPipeline() as p:
+ _ = (p
+ | beam.Create([1, 2])
+ | MyTransform().with_output_types(int))
def test_type_hints_arg(self):
# Tests passing type hints via the magic 'type_hints' argument name.
@@ -2239,11 +2193,10 @@
| beam.ParDo(lambda x: [x]).with_output_types(
type_hints.output_types[0][0]))
- p = TestPipeline()
- _ = (p
- | beam.Create([1, 2])
- | MyTransform('test').with_output_types(int))
- p.run()
+ with TestPipeline() as p:
+ _ = (p
+ | beam.Create([1, 2])
+ | MyTransform('test').with_output_types(int))
def _sort_lists(result):
diff --git a/sdks/python/apache_beam/transforms/py_dataflow_distribution_counter.py b/sdks/python/apache_beam/transforms/py_dataflow_distribution_counter.py
index 980abab..4905b0b 100644
--- a/sdks/python/apache_beam/transforms/py_dataflow_distribution_counter.py
+++ b/sdks/python/apache_beam/transforms/py_dataflow_distribution_counter.py
@@ -17,6 +17,8 @@
"""For internal use only; no backwards-compatibility guarantees."""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/transforms/sideinputs.py b/sdks/python/apache_beam/transforms/sideinputs.py
index cf63b72..79b8290 100644
--- a/sdks/python/apache_beam/transforms/sideinputs.py
+++ b/sdks/python/apache_beam/transforms/sideinputs.py
@@ -24,8 +24,11 @@
AsSingleton, AsIter, AsList and AsDict in apache_beam.pvalue.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
+import re
from builtins import object
from typing import TYPE_CHECKING
from typing import Any
@@ -58,6 +61,16 @@
return map_via_end
+def get_sideinput_index(tag):
+ # type: (str) -> int
+ match = re.match('side([0-9]+)(-.*)?$', tag,
+ re.DOTALL)
+ if match:
+ return int(match.group(1))
+ else:
+ raise RuntimeError("Invalid tag %r" % tag)
+
+
class SideInputMap(object):
"""Represents a mapping of windows to side input values."""
diff --git a/sdks/python/apache_beam/transforms/sideinputs_test.py b/sdks/python/apache_beam/transforms/sideinputs_test.py
index 4fe68ea..3aa87e8 100644
--- a/sdks/python/apache_beam/transforms/sideinputs_test.py
+++ b/sdks/python/apache_beam/transforms/sideinputs_test.py
@@ -17,6 +17,8 @@
"""Unit tests for side inputs."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/transforms/stats.py b/sdks/python/apache_beam/transforms/stats.py
index 5550e48..501d63c 100644
--- a/sdks/python/apache_beam/transforms/stats.py
+++ b/sdks/python/apache_beam/transforms/stats.py
@@ -17,6 +17,8 @@
"""This module has all statistic related transforms."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/transforms/stats_test.py b/sdks/python/apache_beam/transforms/stats_test.py
index 14027fd..b7f9604 100644
--- a/sdks/python/apache_beam/transforms/stats_test.py
+++ b/sdks/python/apache_beam/transforms/stats_test.py
@@ -16,6 +16,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -54,13 +56,12 @@
test_input = [random.randint(0, 1000) for _ in range(100)]
with self.assertRaises(ValueError) as e:
- pipeline = TestPipeline()
- _ = (pipeline
- | 'create'
- >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(size=sample_size))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ _ = (pipeline
+ | 'create'
+ >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(size=sample_size))
expected_msg = beam.ApproximateUnique._INPUT_SIZE_ERR_MSG % (sample_size)
@@ -73,12 +74,11 @@
test_input = [random.randint(0, 1000) for _ in range(100)]
with self.assertRaises(ValueError) as e:
- pipeline = TestPipeline()
- _ = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(size=sample_size))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ _ = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(size=sample_size))
expected_msg = beam.ApproximateUnique._INPUT_SIZE_ERR_MSG % (sample_size)
@@ -91,12 +91,11 @@
test_input = [random.randint(0, 1000) for _ in range(100)]
with self.assertRaises(ValueError) as e:
- pipeline = TestPipeline()
- _ = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(error=est_err))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ _ = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(error=est_err))
expected_msg = beam.ApproximateUnique._INPUT_ERROR_ERR_MSG % (est_err)
@@ -109,12 +108,11 @@
test_input = [random.randint(0, 1000) for _ in range(100)]
with self.assertRaises(ValueError) as e:
- pipeline = TestPipeline()
- _ = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(error=est_err))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ _ = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(error=est_err))
expected_msg = beam.ApproximateUnique._INPUT_ERROR_ERR_MSG % (est_err)
@@ -125,12 +123,11 @@
test_input = [random.randint(0, 1000) for _ in range(100)]
with self.assertRaises(ValueError) as e:
- pipeline = TestPipeline()
- _ = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally())
- pipeline.run()
+ with TestPipeline() as pipeline:
+ _ = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally())
expected_msg = beam.ApproximateUnique._NO_VALUE_ERR_MSG
assert e.exception.args[0] == expected_msg
@@ -142,12 +139,11 @@
sample_size = 30
with self.assertRaises(ValueError) as e:
- pipeline = TestPipeline()
- _ = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(size=sample_size, error=est_err))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ _ = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate' >> beam.ApproximateUnique.Globally(
+ size=sample_size, error=est_err))
expected_msg = beam.ApproximateUnique._MULTI_VALUE_ERR_MSG % (
sample_size, est_err)
@@ -176,18 +172,17 @@
actual_count = len(set(test_input))
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(size=sample_size)
- | 'compare'
- >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
- / actual_count <= max_err]))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(size=sample_size)
+ | 'compare'
+ >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
+ / actual_count <= max_err]))
- assert_that(result, equal_to([True]),
- label='assert:global_by_size')
- pipeline.run()
+ assert_that(result, equal_to([True]),
+ label='assert:global_by_size')
@retry(reraise=True, stop=stop_after_attempt(5))
def test_approximate_unique_global_by_sample_size_with_duplicates(self):
@@ -198,18 +193,17 @@
test_input = [10] * 50 + [20] * 50
actual_count = len(set(test_input))
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(size=sample_size)
- | 'compare'
- >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
- / actual_count <= max_err]))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(size=sample_size)
+ | 'compare'
+ >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
+ / actual_count <= max_err]))
- assert_that(result, equal_to([True]),
- label='assert:global_by_size_with_duplicates')
- pipeline.run()
+ assert_that(result, equal_to([True]),
+ label='assert:global_by_size_with_duplicates')
@retry(reraise=True, stop=stop_after_attempt(5))
def test_approximate_unique_global_by_sample_size_with_small_population(self):
@@ -221,15 +215,14 @@
221, 829, 965, 729, 35, 33, 115, 894, 827, 364]
actual_count = len(set(test_input))
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(size=sample_size))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(size=sample_size))
- assert_that(result, equal_to([actual_count]),
- label='assert:global_by_sample_size_with_small_population')
- pipeline.run()
+ assert_that(result, equal_to([actual_count]),
+ label='assert:global_by_sample_size_with_small_population')
@unittest.skip('Skip because hash function is not good enough. '
'TODO: BEAM-7654')
@@ -241,17 +234,16 @@
973, 386, 506, 546, 991, 450, 226, 889, 514, 693]
actual_count = len(set(test_input))
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(error=est_err)
- | 'compare'
- >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
- / actual_count <= est_err]))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(error=est_err)
+ | 'compare'
+ >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
+ / actual_count <= est_err]))
- assert_that(result, equal_to([True]), label='assert:global_by_error')
- pipeline.run()
+ assert_that(result, equal_to([True]), label='assert:global_by_error')
@retry(reraise=True, stop=stop_after_attempt(5))
def test_approximate_unique_global_by_error_with_small_population(self):
@@ -264,15 +256,14 @@
756, 755, 839, 79, 393]
actual_count = len(set(test_input))
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(error=est_err))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(error=est_err))
- assert_that(result, equal_to([actual_count]),
- label='assert:global_by_error_with_small_population')
- pipeline.run()
+ assert_that(result, equal_to([actual_count]),
+ label='assert:global_by_error_with_small_population')
@retry(reraise=True, stop=stop_after_attempt(5))
def test_approximate_unique_perkey_by_size(self):
@@ -290,20 +281,19 @@
for (x, y) in test_input:
actual_count_dict[x].add(y)
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.PerKey(size=sample_size)
- | 'compare'
- >> beam.FlatMap(lambda x: [abs(x[1]
- - len(actual_count_dict[x[0]]))
- * 1.0 / len(actual_count_dict[x[0]])
- <= max_err]))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.PerKey(size=sample_size)
+ | 'compare'
+ >> beam.FlatMap(lambda x: [abs(x[1]
+ - len(actual_count_dict[x[0]]))
+ * 1.0 / len(actual_count_dict[x[0]])
+ <= max_err]))
- assert_that(result, equal_to([True] * len(actual_count_dict)),
- label='assert:perkey_by_size')
- pipeline.run()
+ assert_that(result, equal_to([True] * len(actual_count_dict)),
+ label='assert:perkey_by_size')
@retry(reraise=True, stop=stop_after_attempt(5))
def test_approximate_unique_perkey_by_error(self):
@@ -316,20 +306,19 @@
for (x, y) in test_input:
actual_count_dict[x].add(y)
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.PerKey(error=est_err)
- | 'compare'
- >> beam.FlatMap(lambda x: [abs(x[1]
- - len(actual_count_dict[x[0]]))
- * 1.0 / len(actual_count_dict[x[0]])
- <= est_err]))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.PerKey(error=est_err)
+ | 'compare'
+ >> beam.FlatMap(lambda x: [abs(x[1]
+ - len(actual_count_dict[x[0]]))
+ * 1.0 / len(actual_count_dict[x[0]])
+ <= est_err]))
- assert_that(result, equal_to([True] * len(actual_count_dict)),
- label='assert:perkey_by_error')
- pipeline.run()
+ assert_that(result, equal_to([True] * len(actual_count_dict)),
+ label='assert:perkey_by_error')
@retry(reraise=True, stop=stop_after_attempt(5))
def test_approximate_unique_globally_by_error_with_skewed_data(self):
@@ -339,18 +328,17 @@
6, 55, 1, 13, 90, 4, 18, 52, 33, 0, 77, 21, 26, 5, 18]
actual_count = len(set(test_input))
- pipeline = TestPipeline()
- result = (pipeline
- | 'create' >> beam.Create(test_input)
- | 'get_estimate'
- >> beam.ApproximateUnique.Globally(error=est_err)
- | 'compare'
- >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
- / actual_count <= est_err]))
+ with TestPipeline() as pipeline:
+ result = (pipeline
+ | 'create' >> beam.Create(test_input)
+ | 'get_estimate'
+ >> beam.ApproximateUnique.Globally(error=est_err)
+ | 'compare'
+ >> beam.FlatMap(lambda x: [abs(x - actual_count) * 1.0
+ / actual_count <= est_err]))
- assert_that(result, equal_to([True]),
- label='assert:globally_by_error_with_skewed_data')
- pipeline.run()
+ assert_that(result, equal_to([True]),
+ label='assert:globally_by_error_with_skewed_data')
class ApproximateQuantilesTest(unittest.TestCase):
diff --git a/sdks/python/apache_beam/transforms/timeutil.py b/sdks/python/apache_beam/transforms/timeutil.py
index 2556571..aaa313d 100644
--- a/sdks/python/apache_beam/transforms/timeutil.py
+++ b/sdks/python/apache_beam/transforms/timeutil.py
@@ -17,6 +17,8 @@
"""Timestamp utilities."""
+# pytype: skip-file
+
from __future__ import absolute_import
from abc import ABCMeta
diff --git a/sdks/python/apache_beam/transforms/transforms_keyword_only_args_test_py3.py b/sdks/python/apache_beam/transforms/transforms_keyword_only_args_test_py3.py
index 6a3c311..b220373 100644
--- a/sdks/python/apache_beam/transforms/transforms_keyword_only_args_test_py3.py
+++ b/sdks/python/apache_beam/transforms/transforms_keyword_only_args_test_py3.py
@@ -17,6 +17,8 @@
"""Unit tests for side inputs."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
@@ -34,109 +36,106 @@
_multiprocess_can_split_ = True
def test_side_input_keyword_only_args(self):
- pipeline = TestPipeline()
+ with TestPipeline() as pipeline:
- def sort_with_side_inputs(x, *s, reverse=False):
- for y in s:
- yield sorted([x] + y, reverse=reverse)
+ def sort_with_side_inputs(x, *s, reverse=False):
+ for y in s:
+ yield sorted([x] + y, reverse=reverse)
- def sort_with_side_inputs_without_default_values(x, *s, reverse):
- for y in s:
- yield sorted([x] + y, reverse=reverse)
+ def sort_with_side_inputs_without_default_values(x, *s, reverse):
+ for y in s:
+ yield sorted([x] + y, reverse=reverse)
- pcol = pipeline | 'start' >> beam.Create([1, 2])
- side = pipeline | 'side' >> beam.Create([3, 4]) # 2 values in side input.
- result1 = pcol | 'compute1' >> beam.FlatMap(
- sort_with_side_inputs,
- beam.pvalue.AsList(side), reverse=True)
- assert_that(result1, equal_to([[4, 3, 1], [4, 3, 2]]), label='assert1')
+ pcol = pipeline | 'start' >> beam.Create([1, 2])
+ side = pipeline | 'side' >> beam.Create([3, 4]) # 2 values in side input.
+ result1 = pcol | 'compute1' >> beam.FlatMap(
+ sort_with_side_inputs,
+ beam.pvalue.AsList(side), reverse=True)
+ assert_that(result1, equal_to([[4, 3, 1], [4, 3, 2]]), label='assert1')
- result2 = pcol | 'compute2' >> beam.FlatMap(
- sort_with_side_inputs,
- beam.pvalue.AsList(side))
- assert_that(result2, equal_to([[1, 3, 4], [2, 3, 4]]), label='assert2')
+ result2 = pcol | 'compute2' >> beam.FlatMap(
+ sort_with_side_inputs,
+ beam.pvalue.AsList(side))
+ assert_that(result2, equal_to([[1, 3, 4], [2, 3, 4]]), label='assert2')
- result3 = pcol | 'compute3' >> beam.FlatMap(
- sort_with_side_inputs)
- assert_that(result3, equal_to([]), label='assert3')
+ result3 = pcol | 'compute3' >> beam.FlatMap(
+ sort_with_side_inputs)
+ assert_that(result3, equal_to([]), label='assert3')
- result4 = pcol | 'compute4' >> beam.FlatMap(
- sort_with_side_inputs, reverse=True)
- assert_that(result4, equal_to([]), label='assert4')
+ result4 = pcol | 'compute4' >> beam.FlatMap(
+ sort_with_side_inputs, reverse=True)
+ assert_that(result4, equal_to([]), label='assert4')
- result5 = pcol | 'compute5' >> beam.FlatMap(
- sort_with_side_inputs_without_default_values,
- beam.pvalue.AsList(side), reverse=True)
- assert_that(result5, equal_to([[4, 3, 1], [4, 3, 2]]), label='assert5')
+ result5 = pcol | 'compute5' >> beam.FlatMap(
+ sort_with_side_inputs_without_default_values,
+ beam.pvalue.AsList(side), reverse=True)
+ assert_that(result5, equal_to([[4, 3, 1], [4, 3, 2]]), label='assert5')
- result6 = pcol | 'compute6' >> beam.FlatMap(
- sort_with_side_inputs_without_default_values,
- beam.pvalue.AsList(side), reverse=False)
- assert_that(result6, equal_to([[1, 3, 4], [2, 3, 4]]), label='assert6')
+ result6 = pcol | 'compute6' >> beam.FlatMap(
+ sort_with_side_inputs_without_default_values,
+ beam.pvalue.AsList(side), reverse=False)
+ assert_that(result6, equal_to([[1, 3, 4], [2, 3, 4]]), label='assert6')
- result7 = pcol | 'compute7' >> beam.FlatMap(
- sort_with_side_inputs_without_default_values, reverse=False)
- assert_that(result7, equal_to([]), label='assert7')
+ result7 = pcol | 'compute7' >> beam.FlatMap(
+ sort_with_side_inputs_without_default_values, reverse=False)
+ assert_that(result7, equal_to([]), label='assert7')
- result8 = pcol | 'compute8' >> beam.FlatMap(
- sort_with_side_inputs_without_default_values, reverse=True)
- assert_that(result8, equal_to([]), label='assert8')
+ result8 = pcol | 'compute8' >> beam.FlatMap(
+ sort_with_side_inputs_without_default_values, reverse=True)
+ assert_that(result8, equal_to([]), label='assert8')
- pipeline.run()
def test_combine_keyword_only_args(self):
- pipeline = TestPipeline()
+ with TestPipeline() as pipeline:
- def bounded_sum(values, *s, bound=500):
- return min(sum(values) + sum(s), bound)
+ def bounded_sum(values, *s, bound=500):
+ return min(sum(values) + sum(s), bound)
- def bounded_sum_without_default_values(values, *s, bound):
- return min(sum(values) + sum(s), bound)
+ def bounded_sum_without_default_values(values, *s, bound):
+ return min(sum(values) + sum(s), bound)
- pcoll = pipeline | 'start' >> beam.Create([6, 3, 1])
- result1 = pcoll | 'sum1' >> beam.CombineGlobally(bounded_sum, 5, 8,
- bound=20)
- result2 = pcoll | 'sum2' >> beam.CombineGlobally(bounded_sum, 0, 0)
- result3 = pcoll | 'sum3' >> beam.CombineGlobally(bounded_sum)
- result4 = pcoll | 'sum4' >> beam.CombineGlobally(bounded_sum, bound=5)
- result5 = pcoll | 'sum5' >> beam.CombineGlobally(
- bounded_sum_without_default_values, 5, 8, bound=20)
- result6 = pcoll | 'sum6' >> beam.CombineGlobally(
- bounded_sum_without_default_values, 0, 0, bound=500)
- result7 = pcoll | 'sum7' >> beam.CombineGlobally(
- bounded_sum_without_default_values, bound=500)
- result8 = pcoll | 'sum8' >> beam.CombineGlobally(
- bounded_sum_without_default_values, bound=5)
+ pcoll = pipeline | 'start' >> beam.Create([6, 3, 1])
+ result1 = pcoll | 'sum1' >> beam.CombineGlobally(bounded_sum, 5, 8,
+ bound=20)
+ result2 = pcoll | 'sum2' >> beam.CombineGlobally(bounded_sum, 0, 0)
+ result3 = pcoll | 'sum3' >> beam.CombineGlobally(bounded_sum)
+ result4 = pcoll | 'sum4' >> beam.CombineGlobally(bounded_sum, bound=5)
+ result5 = pcoll | 'sum5' >> beam.CombineGlobally(
+ bounded_sum_without_default_values, 5, 8, bound=20)
+ result6 = pcoll | 'sum6' >> beam.CombineGlobally(
+ bounded_sum_without_default_values, 0, 0, bound=500)
+ result7 = pcoll | 'sum7' >> beam.CombineGlobally(
+ bounded_sum_without_default_values, bound=500)
+ result8 = pcoll | 'sum8' >> beam.CombineGlobally(
+ bounded_sum_without_default_values, bound=5)
- assert_that(result1, equal_to([20]), label='assert1')
- assert_that(result2, equal_to([10]), label='assert2')
- assert_that(result3, equal_to([10]), label='assert3')
- assert_that(result4, equal_to([5]), label='assert4')
- assert_that(result5, equal_to([20]), label='assert5')
- assert_that(result6, equal_to([10]), label='assert6')
- assert_that(result7, equal_to([10]), label='assert7')
- assert_that(result8, equal_to([5]), label='assert8')
+ assert_that(result1, equal_to([20]), label='assert1')
+ assert_that(result2, equal_to([10]), label='assert2')
+ assert_that(result3, equal_to([10]), label='assert3')
+ assert_that(result4, equal_to([5]), label='assert4')
+ assert_that(result5, equal_to([20]), label='assert5')
+ assert_that(result6, equal_to([10]), label='assert6')
+ assert_that(result7, equal_to([10]), label='assert7')
+ assert_that(result8, equal_to([5]), label='assert8')
- pipeline.run()
def test_do_fn_keyword_only_args(self):
- pipeline = TestPipeline()
+ with TestPipeline() as pipeline:
- class MyDoFn(beam.DoFn):
- def process(self, element, *s, bound=500):
- return [min(sum(s) + element, bound)]
+ class MyDoFn(beam.DoFn):
+ def process(self, element, *s, bound=500):
+ return [min(sum(s) + element, bound)]
- pcoll = pipeline | 'start' >> beam.Create([6, 3, 1])
- result1 = pcoll | 'sum1' >> beam.ParDo(MyDoFn(), 5, 8, bound=15)
- result2 = pcoll | 'sum2' >> beam.ParDo(MyDoFn(), 5, 8)
- result3 = pcoll | 'sum3' >> beam.ParDo(MyDoFn())
- result4 = pcoll | 'sum4' >> beam.ParDo(MyDoFn(), bound=5)
+ pcoll = pipeline | 'start' >> beam.Create([6, 3, 1])
+ result1 = pcoll | 'sum1' >> beam.ParDo(MyDoFn(), 5, 8, bound=15)
+ result2 = pcoll | 'sum2' >> beam.ParDo(MyDoFn(), 5, 8)
+ result3 = pcoll | 'sum3' >> beam.ParDo(MyDoFn())
+ result4 = pcoll | 'sum4' >> beam.ParDo(MyDoFn(), bound=5)
- assert_that(result1, equal_to([15, 15, 14]), label='assert1')
- assert_that(result2, equal_to([19, 16, 14]), label='assert2')
- assert_that(result3, equal_to([6, 3, 1]), label='assert3')
- assert_that(result4, equal_to([5, 3, 1]), label='assert4')
- pipeline.run()
+ assert_that(result1, equal_to([15, 15, 14]), label='assert1')
+ assert_that(result2, equal_to([19, 16, 14]), label='assert2')
+ assert_that(result3, equal_to([6, 3, 1]), label='assert3')
+ assert_that(result4, equal_to([5, 3, 1]), label='assert4')
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/transforms/trigger.py b/sdks/python/apache_beam/transforms/trigger.py
index 65bd4c7..d69f056 100644
--- a/sdks/python/apache_beam/transforms/trigger.py
+++ b/sdks/python/apache_beam/transforms/trigger.py
@@ -20,6 +20,8 @@
Triggers control when in processing time windows get emitted.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/transforms/trigger_test.py b/sdks/python/apache_beam/transforms/trigger_test.py
index bdc8e37..8899cf8 100644
--- a/sdks/python/apache_beam/transforms/trigger_test.py
+++ b/sdks/python/apache_beam/transforms/trigger_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the triggering classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/transforms/userstate.py b/sdks/python/apache_beam/transforms/userstate.py
index dbff82c..c9b2038 100644
--- a/sdks/python/apache_beam/transforms/userstate.py
+++ b/sdks/python/apache_beam/transforms/userstate.py
@@ -20,6 +20,8 @@
Experimental; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import types
@@ -48,8 +50,14 @@
class StateSpec(object):
"""Specification for a user DoFn state cell."""
- def __init__(self):
- raise NotImplementedError
+ def __init__(self, name, coder):
+ # type: (str, Coder) -> None
+ if not isinstance(name, str):
+ raise TypeError("name is not a string")
+ if not isinstance(coder, Coder):
+ raise TypeError("coder is not of type Coder")
+ self.name = name
+ self.coder = coder
def __repr__(self):
return '%s(%s)' % (self.__class__.__name__, self.name)
@@ -61,13 +69,6 @@
class BagStateSpec(StateSpec):
"""Specification for a user DoFn bag state cell."""
- def __init__(self, name, coder):
- # type: (str, Coder) -> None
- assert isinstance(name, str)
- assert isinstance(coder, Coder)
- self.name = name
- self.coder = coder
-
def to_runner_api(self, context):
# type: (PipelineContext) -> beam_runner_api_pb2.StateSpec
return beam_runner_api_pb2.StateSpec(
@@ -78,15 +79,6 @@
class SetStateSpec(StateSpec):
"""Specification for a user DoFn Set State cell"""
- def __init__(self, name, coder):
- # type: (str, Coder) -> None
- if not isinstance(name, str):
- raise TypeError("SetState name is not a string")
- if not isinstance(coder, Coder):
- raise TypeError("SetState coder is not of type Coder")
- self.name = name
- self.coder = coder
-
def to_runner_api(self, context):
return beam_runner_api_pb2.StateSpec(
set_spec=beam_runner_api_pb2.SetStateSpec(
@@ -126,14 +118,11 @@
else:
coder, combine_fn = None, coder
self.combine_fn = CombineFn.maybe_from_callable(combine_fn)
+ # The coder here should be for the accumulator type of the given CombineFn.
if coder is None:
coder = self.combine_fn.get_accumulator_coder()
- assert isinstance(name, str)
- assert isinstance(coder, Coder)
- self.name = name
- # The coder here should be for the accumulator type of the given CombineFn.
- self.coder = coder
+ super(CombiningValueStateSpec, self).__init__(name, coder)
def to_runner_api(self, context):
# type: (PipelineContext) -> beam_runner_api_pb2.StateSpec
diff --git a/sdks/python/apache_beam/transforms/userstate_test.py b/sdks/python/apache_beam/transforms/userstate_test.py
index 601a1d4..30aa88c 100644
--- a/sdks/python/apache_beam/transforms/userstate_test.py
+++ b/sdks/python/apache_beam/transforms/userstate_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the Beam State and Timer API interfaces."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
@@ -111,10 +113,10 @@
def test_spec_construction(self):
BagStateSpec('statename', VarIntCoder())
- with self.assertRaises(AssertionError):
+ with self.assertRaises(TypeError):
BagStateSpec(123, VarIntCoder())
CombiningValueStateSpec('statename', VarIntCoder(), TopCombineFn(10))
- with self.assertRaises(AssertionError):
+ with self.assertRaises(TypeError):
CombiningValueStateSpec(123, VarIntCoder(), TopCombineFn(10))
with self.assertRaises(TypeError):
CombiningValueStateSpec('statename', VarIntCoder(), object())
@@ -516,19 +518,16 @@
aggregated_value += saved_value
yield aggregated_value
- p = TestPipeline()
- values = p | beam.Create([('key', 1),
- ('key', 2),
- ('key', 3),
- ('key', 4),
- ('key', 3)], reshuffle=False)
- actual_values = (values
- | beam.ParDo(SetStatefulDoFn()))
+ with TestPipeline() as p:
+ values = p | beam.Create([('key', 1),
+ ('key', 2),
+ ('key', 3),
+ ('key', 4),
+ ('key', 3)], reshuffle=False)
+ actual_values = (values
+ | beam.ParDo(SetStatefulDoFn()))
+ assert_that(actual_values, equal_to([1, 3, 6, 10, 10]))
- assert_that(actual_values, equal_to([1, 3, 6, 10, 10]))
-
- result = p.run()
- result.wait_until_finish()
def test_stateful_set_state_clean_portably(self):
@@ -555,21 +554,19 @@
def emit_values(self, set_state=beam.DoFn.StateParam(SET_STATE)):
yield sorted(set_state.read())
- p = TestPipeline()
- values = p | beam.Create([('key', 1),
- ('key', 2),
- ('key', 3),
- ('key', 4),
- ('key', 5)])
- actual_values = (values
- | beam.Map(lambda t: window.TimestampedValue(t, 1))
- | beam.WindowInto(window.FixedWindows(1))
- | beam.ParDo(SetStateClearingStatefulDoFn()))
+ with TestPipeline() as p:
+ values = p | beam.Create([('key', 1),
+ ('key', 2),
+ ('key', 3),
+ ('key', 4),
+ ('key', 5)])
+ actual_values = (values
+ | beam.Map(lambda t: window.TimestampedValue(t, 1))
+ | beam.WindowInto(window.FixedWindows(1))
+ | beam.ParDo(SetStateClearingStatefulDoFn()))
- assert_that(actual_values, equal_to([[100]]))
+ assert_that(actual_values, equal_to([[100]]))
- result = p.run()
- result.wait_until_finish()
def test_stateful_dofn_nonkeyed_input(self):
p = TestPipeline()
diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py
index ddb5e50..a29ad68 100644
--- a/sdks/python/apache_beam/transforms/util.py
+++ b/sdks/python/apache_beam/transforms/util.py
@@ -18,6 +18,8 @@
"""Simple utility PTransforms.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py
index cf2e3b1..fc32874 100644
--- a/sdks/python/apache_beam/transforms/util_test.py
+++ b/sdks/python/apache_beam/transforms/util_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the transform.util classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
@@ -282,23 +284,22 @@
yield WindowedValue(
element, expected_timestamp, [expected_window])
- pipeline = TestPipeline()
- data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
- expected_windows = [
- TestWindowedValue(kv, expected_timestamp, [expected_window])
- for kv in data]
- before_identity = (pipeline
- | 'start' >> beam.Create(data)
- | 'add_windows' >> beam.ParDo(AddWindowDoFn()))
- assert_that(before_identity, equal_to(expected_windows),
- label='before_identity', reify_windows=True)
- after_identity = (before_identity
- | 'window' >> beam.WindowInto(
- beam.transforms.util._IdentityWindowFn(
- coders.IntervalWindowCoder())))
- assert_that(after_identity, equal_to(expected_windows),
- label='after_identity', reify_windows=True)
- pipeline.run()
+ with TestPipeline() as pipeline:
+ data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
+ expected_windows = [
+ TestWindowedValue(kv, expected_timestamp, [expected_window])
+ for kv in data]
+ before_identity = (pipeline
+ | 'start' >> beam.Create(data)
+ | 'add_windows' >> beam.ParDo(AddWindowDoFn()))
+ assert_that(before_identity, equal_to(expected_windows),
+ label='before_identity', reify_windows=True)
+ after_identity = (before_identity
+ | 'window' >> beam.WindowInto(
+ beam.transforms.util._IdentityWindowFn(
+ coders.IntervalWindowCoder())))
+ assert_that(after_identity, equal_to(expected_windows),
+ label='after_identity', reify_windows=True)
def test_no_window_context_fails(self):
expected_timestamp = timestamp.Timestamp(5)
@@ -309,40 +310,38 @@
def process(self, element):
yield window.TimestampedValue(element, expected_timestamp)
- pipeline = TestPipeline()
- data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
- expected_windows = [
- TestWindowedValue(kv, expected_timestamp, [expected_window])
- for kv in data]
- before_identity = (pipeline
- | 'start' >> beam.Create(data)
- | 'add_timestamps' >> beam.ParDo(AddTimestampDoFn()))
- assert_that(before_identity, equal_to(expected_windows),
- label='before_identity', reify_windows=True)
- after_identity = (before_identity
- | 'window' >> beam.WindowInto(
- beam.transforms.util._IdentityWindowFn(
- coders.GlobalWindowCoder()))
- # This DoFn will return TimestampedValues, making
- # WindowFn.AssignContext passed to IdentityWindowFn
- # contain a window of None. IdentityWindowFn should
- # raise an exception.
- | 'add_timestamps2' >> beam.ParDo(AddTimestampDoFn()))
- assert_that(after_identity, equal_to(expected_windows),
- label='after_identity', reify_windows=True)
with self.assertRaisesRegex(ValueError, r'window.*None.*add_timestamps2'):
- pipeline.run()
+ with TestPipeline() as pipeline:
+ data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
+ expected_windows = [
+ TestWindowedValue(kv, expected_timestamp, [expected_window])
+ for kv in data]
+ before_identity = (pipeline
+ | 'start' >> beam.Create(data)
+ | 'add_timestamps' >> beam.ParDo(AddTimestampDoFn()))
+ assert_that(before_identity, equal_to(expected_windows),
+ label='before_identity', reify_windows=True)
+ after_identity = (before_identity
+ | 'window' >> beam.WindowInto(
+ beam.transforms.util._IdentityWindowFn(
+ coders.GlobalWindowCoder()))
+ # This DoFn will return TimestampedValues, making
+ # WindowFn.AssignContext passed to IdentityWindowFn
+ # contain a window of None. IdentityWindowFn should
+ # raise an exception.
+ | 'add_timestamps2' >> beam.ParDo(AddTimestampDoFn()))
+ assert_that(after_identity, equal_to(expected_windows),
+ label='after_identity', reify_windows=True)
class ReshuffleTest(unittest.TestCase):
def test_reshuffle_contents_unchanged(self):
- pipeline = TestPipeline()
- data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]
- result = (pipeline
- | beam.Create(data)
- | beam.Reshuffle())
- assert_that(result, equal_to(data))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]
+ result = (pipeline
+ | beam.Create(data)
+ | beam.Reshuffle())
+ assert_that(result, equal_to(data))
def test_reshuffle_after_gbk_contents_unchanged(self):
pipeline = TestPipeline()
@@ -360,74 +359,72 @@
pipeline.run()
def test_reshuffle_timestamps_unchanged(self):
- pipeline = TestPipeline()
- timestamp = 5
- data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]
- expected_result = [TestWindowedValue(v, timestamp, [GlobalWindow()])
- for v in data]
- before_reshuffle = (pipeline
- | 'start' >> beam.Create(data)
- | 'add_timestamp' >> beam.Map(
- lambda v: beam.window.TimestampedValue(v,
- timestamp)))
- assert_that(before_reshuffle, equal_to(expected_result),
- label='before_reshuffle', reify_windows=True)
- after_reshuffle = before_reshuffle | beam.Reshuffle()
- assert_that(after_reshuffle, equal_to(expected_result),
- label='after_reshuffle', reify_windows=True)
- pipeline.run()
+ with TestPipeline() as pipeline:
+ timestamp = 5
+ data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)]
+ expected_result = [TestWindowedValue(v, timestamp, [GlobalWindow()])
+ for v in data]
+ before_reshuffle = (pipeline
+ | 'start' >> beam.Create(data)
+ | 'add_timestamp' >> beam.Map(
+ lambda v: beam.window.TimestampedValue(
+ v, timestamp)))
+ assert_that(before_reshuffle, equal_to(expected_result),
+ label='before_reshuffle', reify_windows=True)
+ after_reshuffle = before_reshuffle | beam.Reshuffle()
+ assert_that(after_reshuffle, equal_to(expected_result),
+ label='after_reshuffle', reify_windows=True)
def test_reshuffle_windows_unchanged(self):
- pipeline = TestPipeline()
- data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
- expected_data = [TestWindowedValue(v, t - .001, [w]) for (v, t, w) in [
- ((1, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
- ((2, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
- ((3, [1]), 3.0, IntervalWindow(1.0, 3.0)),
- ((1, [4]), 6.0, IntervalWindow(4.0, 6.0))]]
- before_reshuffle = (pipeline
- | 'start' >> beam.Create(data)
- | 'add_timestamp' >> beam.Map(
- lambda v: beam.window.TimestampedValue(v, v[1]))
- | 'window' >> beam.WindowInto(Sessions(gap_size=2))
- | 'group_by_key' >> beam.GroupByKey())
- assert_that(before_reshuffle, equal_to(expected_data),
- label='before_reshuffle', reify_windows=True)
- after_reshuffle = before_reshuffle | beam.Reshuffle()
- assert_that(after_reshuffle, equal_to(expected_data),
- label='after reshuffle', reify_windows=True)
- pipeline.run()
+ with TestPipeline() as pipeline:
+ data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
+ expected_data = [TestWindowedValue(v, t - .001, [w]) for (v, t, w) in [
+ ((1, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
+ ((2, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
+ ((3, [1]), 3.0, IntervalWindow(1.0, 3.0)),
+ ((1, [4]), 6.0, IntervalWindow(4.0, 6.0))]]
+ before_reshuffle = (pipeline
+ | 'start' >> beam.Create(data)
+ | 'add_timestamp' >> beam.Map(
+ lambda v: beam.window.TimestampedValue(v, v[1]))
+ | 'window' >> beam.WindowInto(Sessions(gap_size=2))
+ | 'group_by_key' >> beam.GroupByKey())
+ assert_that(before_reshuffle, equal_to(expected_data),
+ label='before_reshuffle', reify_windows=True)
+ after_reshuffle = before_reshuffle | beam.Reshuffle()
+ assert_that(after_reshuffle, equal_to(expected_data),
+ label='after reshuffle', reify_windows=True)
def test_reshuffle_window_fn_preserved(self):
- pipeline = TestPipeline()
- data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
- expected_windows = [TestWindowedValue(v, t, [w]) for (v, t, w) in [
- ((1, 1), 1.0, IntervalWindow(1.0, 3.0)),
- ((2, 1), 1.0, IntervalWindow(1.0, 3.0)),
- ((3, 1), 1.0, IntervalWindow(1.0, 3.0)),
- ((1, 2), 2.0, IntervalWindow(2.0, 4.0)),
- ((2, 2), 2.0, IntervalWindow(2.0, 4.0)),
- ((1, 4), 4.0, IntervalWindow(4.0, 6.0))]]
- expected_merged_windows = [
- TestWindowedValue(v, t - .001, [w]) for (v, t, w) in [
- ((1, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
- ((2, contains_in_any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
- ((3, [1]), 3.0, IntervalWindow(1.0, 3.0)),
- ((1, [4]), 6.0, IntervalWindow(4.0, 6.0))]]
- before_reshuffle = (pipeline
- | 'start' >> beam.Create(data)
- | 'add_timestamp' >> beam.Map(
- lambda v: TimestampedValue(v, v[1]))
- | 'window' >> beam.WindowInto(Sessions(gap_size=2)))
- assert_that(before_reshuffle, equal_to(expected_windows),
- label='before_reshuffle', reify_windows=True)
- after_reshuffle = before_reshuffle | beam.Reshuffle()
- assert_that(after_reshuffle, equal_to(expected_windows),
- label='after_reshuffle', reify_windows=True)
- after_group = after_reshuffle | beam.GroupByKey()
- assert_that(after_group, equal_to(expected_merged_windows),
- label='after_group', reify_windows=True)
- pipeline.run()
+ any_order = contains_in_any_order
+ with TestPipeline() as pipeline:
+ data = [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 4)]
+ expected_windows = [TestWindowedValue(v, t, [w]) for (v, t, w) in [
+ ((1, 1), 1.0, IntervalWindow(1.0, 3.0)),
+ ((2, 1), 1.0, IntervalWindow(1.0, 3.0)),
+ ((3, 1), 1.0, IntervalWindow(1.0, 3.0)),
+ ((1, 2), 2.0, IntervalWindow(2.0, 4.0)),
+ ((2, 2), 2.0, IntervalWindow(2.0, 4.0)),
+ ((1, 4), 4.0, IntervalWindow(4.0, 6.0))]]
+ expected_merged_windows = [
+ TestWindowedValue(v, t - .001, [w]) for (v, t, w) in [
+ ((1, any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
+ ((2, any_order([2, 1])), 4.0, IntervalWindow(1.0, 4.0)),
+ ((3, [1]), 3.0, IntervalWindow(1.0, 3.0)),
+ ((1, [4]), 6.0, IntervalWindow(4.0, 6.0))]]
+ before_reshuffle = (pipeline
+ | 'start' >> beam.Create(data)
+ | 'add_timestamp' >> beam.Map(
+ lambda v: TimestampedValue(v, v[1]))
+ | 'window' >> beam.WindowInto(Sessions(gap_size=2)))
+ assert_that(before_reshuffle, equal_to(expected_windows),
+ label='before_reshuffle', reify_windows=True)
+ after_reshuffle = before_reshuffle | beam.Reshuffle()
+ assert_that(after_reshuffle, equal_to(expected_windows),
+ label='after_reshuffle', reify_windows=True)
+ after_group = after_reshuffle | beam.GroupByKey()
+ assert_that(after_group, equal_to(expected_merged_windows),
+ label='after_group', reify_windows=True)
def test_reshuffle_global_window(self):
pipeline = TestPipeline()
@@ -582,16 +579,16 @@
return data
def test_in_global_window(self):
- pipeline = TestPipeline()
- collection = pipeline \
- | beam.Create(GroupIntoBatchesTest._create_test_data()) \
- | util.GroupIntoBatches(GroupIntoBatchesTest.BATCH_SIZE)
- num_batches = collection | beam.combiners.Count.Globally()
- assert_that(num_batches,
- equal_to([int(math.ceil(GroupIntoBatchesTest.NUM_ELEMENTS /
- GroupIntoBatchesTest.BATCH_SIZE))]))
- pipeline.run()
+ with TestPipeline() as pipeline:
+ collection = pipeline \
+ | beam.Create(GroupIntoBatchesTest._create_test_data()) \
+ | util.GroupIntoBatches(GroupIntoBatchesTest.BATCH_SIZE)
+ num_batches = collection | beam.combiners.Count.Globally()
+ assert_that(num_batches,
+ equal_to([int(math.ceil(GroupIntoBatchesTest.NUM_ELEMENTS /
+ GroupIntoBatchesTest.BATCH_SIZE))]))
+ @unittest.skip('BEAM-8748')
def test_in_streaming_mode(self):
timestamp_interval = 1
offset = itertools.count(0)
@@ -607,26 +604,23 @@
.advance_watermark_to(start_time +
GroupIntoBatchesTest.NUM_ELEMENTS)
.advance_watermark_to_infinity())
- pipeline = TestPipeline(options=StandardOptions(streaming=True))
- # window duration is 6 and batch size is 5, so output batch size should be
- # 5 (flush because of batchSize reached)
- expected_0 = 5
- # there is only one element left in the window so batch size should be 1
- # (flush because of end of window reached)
- expected_1 = 1
- # collection is 10 elements, there is only 4 left, so batch size should be
- # 4 (flush because end of collection reached)
- expected_2 = 4
+ with TestPipeline(options=StandardOptions(streaming=True)) as pipeline:
+ # window duration is 6 and batch size is 5, so output batch size
+ # should be 5 (flush because of batchSize reached)
+ expected_0 = 5
+ # there is only one element left in the window so batch size
+ # should be 1 (flush because of end of window reached)
+ expected_1 = 1
+ # collection is 10 elements, there is only 4 left, so batch size
+ # should be 4 (flush because end of collection reached)
+ expected_2 = 4
- collection = pipeline | test_stream \
- | WindowInto(FixedWindows(window_duration)) \
- | util.GroupIntoBatches(GroupIntoBatchesTest.BATCH_SIZE)
- num_elements_in_batches = collection | beam.Map(len)
-
- result = pipeline.run()
- result.wait_until_finish()
- assert_that(num_elements_in_batches,
- equal_to([expected_0, expected_1, expected_2]))
+ collection = pipeline | test_stream \
+ | WindowInto(FixedWindows(window_duration)) \
+ | util.GroupIntoBatches(GroupIntoBatchesTest.BATCH_SIZE)
+ num_elements_in_batches = collection | beam.Map(len)
+ assert_that(num_elements_in_batches,
+ equal_to([expected_0, expected_1, expected_2]))
class ToStringTest(unittest.TestCase):
diff --git a/sdks/python/apache_beam/transforms/window.py b/sdks/python/apache_beam/transforms/window.py
index b16d00c..b79ed20 100644
--- a/sdks/python/apache_beam/transforms/window.py
+++ b/sdks/python/apache_beam/transforms/window.py
@@ -47,6 +47,8 @@
WindowFn.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import abc
diff --git a/sdks/python/apache_beam/transforms/window_test.py b/sdks/python/apache_beam/transforms/window_test.py
index 30430cc..0c88e21 100644
--- a/sdks/python/apache_beam/transforms/window_test.py
+++ b/sdks/python/apache_beam/transforms/window_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the windowing classes."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/transforms/write_ptransform_test.py b/sdks/python/apache_beam/transforms/write_ptransform_test.py
index a8f56fd..5d3e7a9 100644
--- a/sdks/python/apache_beam/transforms/write_ptransform_test.py
+++ b/sdks/python/apache_beam/transforms/write_ptransform_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the write transform."""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/typehints/decorators.py b/sdks/python/apache_beam/typehints/decorators.py
index adc2173..4f53eb5 100644
--- a/sdks/python/apache_beam/typehints/decorators.py
+++ b/sdks/python/apache_beam/typehints/decorators.py
@@ -83,6 +83,8 @@
defined, or before importing a module containing type-hinted functions.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import inspect
@@ -309,6 +311,11 @@
Only affects instances with simple output types, otherwise is a no-op.
Does not modify self.
+ Designed to be used with type hints from callables of ParDo, FlatMap, DoFn.
+ Output type may be Optional[T], in which case the result of stripping T is
+ used as the output type.
+ Output type may be None/NoneType, in which case nothing is done.
+
Example: Generator[Tuple(int, int)] becomes Tuple(int, int)
Returns:
@@ -319,7 +326,20 @@
"""
if not self.has_simple_output_type():
return self
- yielded_type = typehints.get_yielded_type(self.output_types[0][0])
+ output_type = self.output_types[0][0]
+ if output_type is None or isinstance(output_type, type(None)):
+ return self
+ # If output_type == Optional[T]: output_type = T.
+ if isinstance(output_type, typehints.UnionConstraint):
+ types = list(output_type.union_types)
+ if len(types) == 2:
+ try:
+ types.remove(type(None))
+ output_type = types[0]
+ except ValueError:
+ pass
+
+ yielded_type = typehints.get_yielded_type(output_type)
res = self.copy()
res.output_types = ((yielded_type,), {})
return res
diff --git a/sdks/python/apache_beam/typehints/decorators_test.py b/sdks/python/apache_beam/typehints/decorators_test.py
index 20df674..f477f92 100644
--- a/sdks/python/apache_beam/typehints/decorators_test.py
+++ b/sdks/python/apache_beam/typehints/decorators_test.py
@@ -17,6 +17,8 @@
"""Tests for decorators module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
@@ -88,8 +90,7 @@
self._test_strip_iterable(before, None)
def test_strip_iterable(self):
- # TODO(BEAM-8492): Uncomment once #9895 is merged.
- # self._test_strip_iterable(None, None)
+ self._test_strip_iterable(None, None)
self._test_strip_iterable(typehints.Any, typehints.Any)
self._test_strip_iterable(typehints.Iterable[str], str)
self._test_strip_iterable(typehints.List[str], str)
diff --git a/sdks/python/apache_beam/typehints/decorators_test_py3.py b/sdks/python/apache_beam/typehints/decorators_test_py3.py
index e66b140..a9e2aad 100644
--- a/sdks/python/apache_beam/typehints/decorators_test_py3.py
+++ b/sdks/python/apache_beam/typehints/decorators_test_py3.py
@@ -17,6 +17,8 @@
"""Tests for decorators module with Python 3 syntax not supported by 2.7."""
+# pytype: skip-file
+
from __future__ import absolute_import
import typing
@@ -41,7 +43,7 @@
class IOTypeHintsTest(unittest.TestCase):
def test_from_callable(self):
- def fn(a: int, b: str = None, *args: Tuple[T], foo: List[int],
+ def fn(a: int, b: str = '', *args: Tuple[T], foo: List[int],
**kwargs: Dict[str, str]) -> Tuple[Any, ...]:
return a, b, args, foo, kwargs
th = decorators.IOTypeHints.from_callable(fn)
@@ -94,7 +96,7 @@
self.assertEqual(th.output_types, ((Tuple[Any, ...],), {}))
def test_getcallargs_forhints(self):
- def fn(a: int, b: str = None, *args: Tuple[T], foo: List[int],
+ def fn(a: int, b: str = '', *args: Tuple[T], foo: List[int],
**kwargs: Dict[str, str]) -> Tuple[Any, ...]:
return a, b, args, foo, kwargs
callargs = decorators.getcallargs_forhints(fn, float, foo=List[str])
diff --git a/sdks/python/apache_beam/typehints/native_type_compatibility.py b/sdks/python/apache_beam/typehints/native_type_compatibility.py
index f2802e0..d7e45a1 100644
--- a/sdks/python/apache_beam/typehints/native_type_compatibility.py
+++ b/sdks/python/apache_beam/typehints/native_type_compatibility.py
@@ -17,6 +17,8 @@
"""Module to convert Python's native typing types to Beam types."""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/typehints/native_type_compatibility_test.py b/sdks/python/apache_beam/typehints/native_type_compatibility_test.py
index 4bfb60b..969f9c7 100644
--- a/sdks/python/apache_beam/typehints/native_type_compatibility_test.py
+++ b/sdks/python/apache_beam/typehints/native_type_compatibility_test.py
@@ -17,6 +17,8 @@
"""Test for Beam type compatibility library."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
diff --git a/sdks/python/apache_beam/typehints/opcodes.py b/sdks/python/apache_beam/typehints/opcodes.py
index 6aa5de2..8034d0d 100644
--- a/sdks/python/apache_beam/typehints/opcodes.py
+++ b/sdks/python/apache_beam/typehints/opcodes.py
@@ -18,7 +18,8 @@
"""Defines the actions various bytecodes have on the frame.
Each function here corresponds to a bytecode documented in
-https://docs.python.org/2/library/dis.html. The first argument is a (mutable)
+https://docs.python.org/2/library/dis.html or
+https://docs.python.org/3/library/dis.html. The first argument is a (mutable)
FrameState object, the second the integer opcode argument.
Bytecodes with more complicated behavior (e.g. modifying the program counter)
@@ -26,6 +27,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import inspect
@@ -36,17 +39,17 @@
from past.builtins import unicode
-from . import typehints
-from .trivial_inference import BoundMethod
-from .trivial_inference import Const
-from .trivial_inference import element_type
-from .trivial_inference import union
-from .typehints import Any
-from .typehints import Dict
-from .typehints import Iterable
-from .typehints import List
-from .typehints import Tuple
-from .typehints import Union
+from apache_beam.typehints import typehints
+from apache_beam.typehints.trivial_inference import BoundMethod
+from apache_beam.typehints.trivial_inference import Const
+from apache_beam.typehints.trivial_inference import element_type
+from apache_beam.typehints.trivial_inference import union
+from apache_beam.typehints.typehints import Any
+from apache_beam.typehints.typehints import Dict
+from apache_beam.typehints.typehints import Iterable
+from apache_beam.typehints.typehints import List
+from apache_beam.typehints.typehints import Tuple
+from apache_beam.typehints.typehints import Union
def pop_one(state, unused_arg):
@@ -152,7 +155,7 @@
def binary_subscr(state, unused_arg):
index = state.stack.pop()
- base = state.stack.pop()
+ base = Const.unwrap(state.stack.pop())
if base in (str, unicode):
out = base
elif (isinstance(index, Const) and isinstance(index.value, int)
@@ -175,33 +178,16 @@
binary_xor = inplace_xor = symmetric_binary_op
binary_or = inpalce_or = symmetric_binary_op
-# As far as types are concerned.
-slice_0 = nop
-slice_1 = slice_2 = pop_top
-slice_3 = pop_two
-store_slice_0 = store_slice_1 = store_slice_2 = store_slice_3 = nop
-delete_slice_0 = delete_slice_1 = delete_slice_2 = delete_slice_3 = nop
-
def store_subscr(unused_state, unused_args):
# TODO(robertwb): Update element/value type of iterable/dict.
pass
-binary_divide = binary_floor_divide = binary_modulo = symmetric_binary_op
-binary_divide = binary_floor_divide = binary_modulo = symmetric_binary_op
-binary_divide = binary_floor_divide = binary_modulo = symmetric_binary_op
-
-# print_expr
print_item = pop_top
-# print_item_to
print_newline = nop
-# print_newline_to
-
-# break_loop
-# continue_loop
def list_append(state, arg):
new_element_type = Const.unwrap(state.stack.pop())
state.stack[-arg] = List[Union[element_type(state.stack[-arg]),
@@ -217,21 +203,10 @@
load_locals = push_value(Dict[str, Any])
-
-# return_value
-# yield_value
-# import_star
exec_stmt = pop_three
-# pop_block
-# end_finally
build_class = pop_three
-# setup_with
-# with_cleanup
-
-# store_name
-# delete_name
def unpack_sequence(state, arg):
t = state.stack.pop()
if isinstance(t, Const):
@@ -331,18 +306,11 @@
import_from = push_value(Any)
-# jump
-
-# for_iter
-
def load_global(state, arg):
state.stack.append(state.get_global(arg))
-# setup_loop
-# setup_except
-# setup_finally
store_map = pop_two
@@ -364,7 +332,6 @@
def load_deref(state, arg):
state.stack.append(state.closure_type(arg))
-# raise_varargs
def make_function(state, arg):
diff --git a/sdks/python/apache_beam/typehints/schemas.py b/sdks/python/apache_beam/typehints/schemas.py
index 812cbe1..eae83cd 100644
--- a/sdks/python/apache_beam/typehints/schemas.py
+++ b/sdks/python/apache_beam/typehints/schemas.py
@@ -41,6 +41,8 @@
ByteString <-----> BYTES
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
diff --git a/sdks/python/apache_beam/typehints/schemas_test.py b/sdks/python/apache_beam/typehints/schemas_test.py
index 9dd1bc2..b8916d8 100644
--- a/sdks/python/apache_beam/typehints/schemas_test.py
+++ b/sdks/python/apache_beam/typehints/schemas_test.py
@@ -16,6 +16,8 @@
#
"""Tests for schemas."""
+# pytype: skip-file
+
from __future__ import absolute_import
import itertools
diff --git a/sdks/python/apache_beam/typehints/trivial_inference.py b/sdks/python/apache_beam/typehints/trivial_inference.py
index f120dfb..8e4b14e 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference.py
@@ -19,6 +19,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
@@ -65,18 +67,27 @@
elif t == tuple:
return typehints.Tuple[[instance_to_type(item) for item in o]]
elif t == list:
- return typehints.List[
- typehints.Union[[instance_to_type(item) for item in o]]
- ]
+ if len(o) > 0:
+ return typehints.List[
+ typehints.Union[[instance_to_type(item) for item in o]]
+ ]
+ else:
+ return typehints.List[typehints.Any]
elif t == set:
- return typehints.Set[
- typehints.Union[[instance_to_type(item) for item in o]]
- ]
+ if len(o) > 0:
+ return typehints.Set[
+ typehints.Union[[instance_to_type(item) for item in o]]
+ ]
+ else:
+ return typehints.Set[typehints.Any]
elif t == dict:
- return typehints.Dict[
- typehints.Union[[instance_to_type(k) for k, v in o.items()]],
- typehints.Union[[instance_to_type(v) for k, v in o.items()]],
- ]
+ if len(o) > 0:
+ return typehints.Dict[
+ typehints.Union[[instance_to_type(k) for k, v in o.items()]],
+ typehints.Union[[instance_to_type(v) for k, v in o.items()]],
+ ]
+ else:
+ return typehints.Dict[typehints.Any, typehints.Any]
else:
raise TypeInferenceError('Unknown forbidden type: %s' % t)
diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test.py b/sdks/python/apache_beam/typehints/trivial_inference_test.py
index ff0949b..cd2ce29 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference_test.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference_test.py
@@ -17,9 +17,12 @@
"""Tests for apache_beam.typehints.trivial_inference."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
+import types
import unittest
from apache_beam.typehints import trivial_inference
@@ -69,6 +72,11 @@
self.assertReturnType(str, lambda v: v[::-1], [str])
self.assertReturnType(typehints.Any, lambda v: v[::-1], [typehints.Any])
self.assertReturnType(typehints.Any, lambda v: v[::-1], [object])
+ if sys.version_info >= (3,):
+ # Test binary_subscr on a slice of a Const. On Py2.7 this will use the
+ # unsupported opcode SLICE+0.
+ test_list = ['a', 'b']
+ self.assertReturnType(typehints.List[str], lambda: test_list[:], [])
def testUnpack(self):
def reverse(a_b):
@@ -275,6 +283,35 @@
lambda x1, x2, _dict: fn(x1, x2, **_dict),
[str, float, typehints.List[int]])
+ def testInstanceToType(self):
+ class MyClass(object):
+ def method(self):
+ pass
+
+ test_cases = [
+ (typehints.Dict[str, int], {'a': 1}),
+ (typehints.Dict[str, typehints.Union[str, int]], {'a': 1, 'b': 'c'}),
+ (typehints.Dict[typehints.Any, typehints.Any], {}),
+ (typehints.Set[str], {'a'}),
+ (typehints.Set[typehints.Union[str, float]], {'a', 0.4}),
+ (typehints.Set[typehints.Any], set()),
+ (typehints.Tuple[int], (1, )),
+ (typehints.Tuple[int, int, str], (1, 2, '3')),
+ (typehints.Tuple[()], ()),
+ (typehints.List[int], [1]),
+ (typehints.List[typehints.Union[int, str]], [1, 'a']),
+ (typehints.List[typehints.Any], []),
+ (type(None), None),
+ (type(MyClass), MyClass),
+ (MyClass, MyClass()),
+ (type(MyClass.method), MyClass.method),
+ (types.MethodType, MyClass().method),
+ ]
+ for expected_type, instance in test_cases:
+ self.assertEqual(expected_type,
+ trivial_inference.instance_to_type(instance),
+ msg=instance)
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test_py3.py b/sdks/python/apache_beam/typehints/trivial_inference_test_py3.py
index 291e52e..9c368a7 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference_test_py3.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference_test_py3.py
@@ -18,6 +18,8 @@
"""Tests for apache_beam.typehints.trivial_inference that use Python 3 syntax.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/typehints/typecheck.py b/sdks/python/apache_beam/typehints/typecheck.py
index e9187f0..09b0cb2 100644
--- a/sdks/python/apache_beam/typehints/typecheck.py
+++ b/sdks/python/apache_beam/typehints/typecheck.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/typehints/typed_pipeline_test.py b/sdks/python/apache_beam/typehints/typed_pipeline_test.py
index 9726064..52cd4ee 100644
--- a/sdks/python/apache_beam/typehints/typed_pipeline_test.py
+++ b/sdks/python/apache_beam/typehints/typed_pipeline_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the type-hint objects and decorators."""
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
@@ -135,19 +137,18 @@
self.assertEqual([1, 3], [1, 2, 3] | beam.Filter(filter_fn))
def test_partition(self):
- p = TestPipeline()
- even, odd = (p
- | beam.Create([1, 2, 3])
- | 'even_odd' >> beam.Partition(lambda e, _: e % 2, 2))
- self.assertIsNotNone(even.element_type)
- self.assertIsNotNone(odd.element_type)
- res_even = (even
- | 'id_even' >> beam.ParDo(lambda e: [e]).with_input_types(int))
- res_odd = (odd
- | 'id_odd' >> beam.ParDo(lambda e: [e]).with_input_types(int))
- assert_that(res_even, equal_to([2]), label='even_check')
- assert_that(res_odd, equal_to([1, 3]), label='odd_check')
- p.run()
+ with TestPipeline() as p:
+ even, odd = (p
+ | beam.Create([1, 2, 3])
+ | 'even_odd' >> beam.Partition(lambda e, _: e % 2, 2))
+ self.assertIsNotNone(even.element_type)
+ self.assertIsNotNone(odd.element_type)
+ res_even = (even
+ | 'IdEven' >> beam.ParDo(lambda e: [e]).with_input_types(int))
+ res_odd = (odd
+ | 'IdOdd' >> beam.ParDo(lambda e: [e]).with_input_types(int))
+ assert_that(res_even, equal_to([2]), label='even_check')
+ assert_that(res_odd, equal_to([1, 3]), label='odd_check')
def test_typed_dofn_multi_output(self):
class MyDoFn(beam.DoFn):
diff --git a/sdks/python/apache_beam/typehints/typed_pipeline_test_py3.py b/sdks/python/apache_beam/typehints/typed_pipeline_test_py3.py
index a718c8b..7f85384 100644
--- a/sdks/python/apache_beam/typehints/typed_pipeline_test_py3.py
+++ b/sdks/python/apache_beam/typehints/typed_pipeline_test_py3.py
@@ -18,6 +18,8 @@
"""Unit tests for type-hint objects and decorators - Python 3 syntax specific.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
@@ -124,6 +126,31 @@
with self.assertRaisesRegex(ValueError, r'str.*is not iterable'):
_ = [1, 2, 3] | beam.ParDo(MyDoFn())
+ def test_typed_dofn_method_return_none(self):
+ class MyDoFn(beam.DoFn):
+ def process(self, unused_element: int) -> None:
+ pass
+
+ result = [1, 2, 3] | beam.ParDo(MyDoFn())
+ self.assertListEqual([], result)
+
+ def test_typed_dofn_method_return_optional(self):
+ class MyDoFn(beam.DoFn):
+ def process(self, unused_element: int) -> typehints.Optional[
+ typehints.Iterable[int]]:
+ pass
+
+ result = [1, 2, 3] | beam.ParDo(MyDoFn())
+ self.assertListEqual([], result)
+
+ def test_typed_dofn_method_return_optional_not_iterable(self):
+ class MyDoFn(beam.DoFn):
+ def process(self, unused_element: int) -> typehints.Optional[int]:
+ pass
+
+ with self.assertRaisesRegex(ValueError, r'int.*is not iterable'):
+ _ = [1, 2, 3] | beam.ParDo(MyDoFn())
+
def test_typed_callable_not_iterable(self):
def do_fn(element: int) -> int:
return [element] # Return a list to not fail the pipeline.
@@ -176,6 +203,57 @@
result = [1, 2] | beam.ParDo(MyDoFn())
self.assertEqual([['1', '1'], ['2', '2']], sorted(result))
+ def test_typed_map(self):
+ def fn(element: int) -> int:
+ return element * 2
+
+ result = [1, 2, 3] | beam.Map(fn)
+ self.assertEqual([2, 4, 6], sorted(result))
+
+ def test_typed_map_return_optional(self):
+ # None is a valid element value for Map.
+ def fn(element: int) -> typehints.Optional[int]:
+ if element > 1:
+ return element
+
+ result = [1, 2, 3] | beam.Map(fn)
+ self.assertCountEqual([None, 2, 3], result)
+
+ def test_typed_flatmap(self):
+ def fn(element: int) -> typehints.Iterable[int]:
+ yield element * 2
+
+ result = [1, 2, 3] | beam.FlatMap(fn)
+ self.assertCountEqual([2, 4, 6], result)
+
+ def test_typed_flatmap_output_hint_not_iterable(self):
+ def fn(element: int) -> int:
+ return element * 2
+
+ # TODO(BEAM-8466): This case currently only generates a warning instead of a
+ # typehints.TypeCheckError.
+ with self.assertRaisesRegex(TypeError, r'int.*is not iterable'):
+ _ = [1, 2, 3] | beam.FlatMap(fn)
+
+ def test_typed_flatmap_output_value_not_iterable(self):
+ def fn(element: int) -> typehints.Iterable[int]:
+ return element * 2
+
+ with self.assertRaisesRegex(TypeError, r'int.*is not iterable'):
+ _ = [1, 2, 3] | beam.FlatMap(fn)
+
+ def test_typed_flatmap_optional(self):
+ def fn(element: int) -> typehints.Optional[typehints.Iterable[int]]:
+ if element > 1:
+ yield element * 2
+
+ # Verify that the output type of fn is int and not Optional[int].
+ def fn2(element: int) -> int:
+ return element
+
+ result = [1, 2, 3] | beam.FlatMap(fn) | beam.Map(fn2)
+ self.assertCountEqual([4, 6], result)
+
class AnnotationsTest(unittest.TestCase):
@@ -230,6 +308,15 @@
self.assertEqual(th.input_types, ((int,), {}))
self.assertEqual(th.output_types, ((int,), {}))
+ def test_flat_map_wrapper_optional_output(self):
+ # Optional should not affect output type (Nones are ignored).
+ def map_fn(element: int) -> typehints.Optional[typehints.Iterable[int]]:
+ return [element, element + 1]
+
+ th = beam.FlatMap(map_fn).get_type_hints()
+ self.assertEqual(th.input_types, ((int,), {}))
+ self.assertEqual(th.output_types, ((int,), {}))
+
@unittest.skip('BEAM-8662: Py3 annotations not yet supported for MapTuple')
def test_flat_map_tuple_wrapper(self):
# TODO(BEAM-8662): Also test with a fn that accepts default arguments.
@@ -248,6 +335,15 @@
self.assertEqual(th.input_types, ((int,), {}))
self.assertEqual(th.output_types, ((int,), {}))
+ def test_map_wrapper_optional_output(self):
+ # Optional does affect output type (Nones are NOT ignored).
+ def map_fn(unused_element: int) -> typehints.Optional[int]:
+ return 1
+
+ th = beam.Map(map_fn).get_type_hints()
+ self.assertEqual(th.input_types, ((int,), {}))
+ self.assertEqual(th.output_types, ((typehints.Optional[int],), {}))
+
@unittest.skip('BEAM-8662: Py3 annotations not yet supported for MapTuple')
def test_map_tuple(self):
# TODO(BEAM-8662): Also test with a fn that accepts default arguments.
diff --git a/sdks/python/apache_beam/typehints/typehints.py b/sdks/python/apache_beam/typehints/typehints.py
index 052beb3..78bff46 100644
--- a/sdks/python/apache_beam/typehints/typehints.py
+++ b/sdks/python/apache_beam/typehints/typehints.py
@@ -63,6 +63,8 @@
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import collections
diff --git a/sdks/python/apache_beam/typehints/typehints_test.py b/sdks/python/apache_beam/typehints/typehints_test.py
index 9b73a7f..38c0118 100644
--- a/sdks/python/apache_beam/typehints/typehints_test.py
+++ b/sdks/python/apache_beam/typehints/typehints_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the type-hint objects and decorators."""
+# pytype: skip-file
+
from __future__ import absolute_import
import functools
diff --git a/sdks/python/apache_beam/typehints/typehints_test_py3.py b/sdks/python/apache_beam/typehints/typehints_test_py3.py
index 01df57c..56042ff 100644
--- a/sdks/python/apache_beam/typehints/typehints_test_py3.py
+++ b/sdks/python/apache_beam/typehints/typehints_test_py3.py
@@ -18,6 +18,8 @@
"""Unit tests for the type-hint objects and decorators with Python 3 syntax not
supported by 2.7."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import print_function
diff --git a/sdks/python/apache_beam/utils/annotations.py b/sdks/python/apache_beam/utils/annotations.py
index af44fac..b60fb18 100644
--- a/sdks/python/apache_beam/utils/annotations.py
+++ b/sdks/python/apache_beam/utils/annotations.py
@@ -81,6 +81,8 @@
print(exp_multiply(5,6))
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import warnings
diff --git a/sdks/python/apache_beam/utils/annotations_test.py b/sdks/python/apache_beam/utils/annotations_test.py
index b9bd774..58d0f40 100644
--- a/sdks/python/apache_beam/utils/annotations_test.py
+++ b/sdks/python/apache_beam/utils/annotations_test.py
@@ -14,6 +14,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/utils/counters.py b/sdks/python/apache_beam/utils/counters.py
index 94a19d1..51d1ab0 100644
--- a/sdks/python/apache_beam/utils/counters.py
+++ b/sdks/python/apache_beam/utils/counters.py
@@ -23,6 +23,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import threading
diff --git a/sdks/python/apache_beam/utils/counters_test.py b/sdks/python/apache_beam/utils/counters_test.py
index d868861..be1ecf9 100644
--- a/sdks/python/apache_beam/utils/counters_test.py
+++ b/sdks/python/apache_beam/utils/counters_test.py
@@ -17,6 +17,8 @@
"""Unit tests for counters and counter names."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/utils/interactive_utils.py b/sdks/python/apache_beam/utils/interactive_utils.py
index ac4e5d8..2163e94 100644
--- a/sdks/python/apache_beam/utils/interactive_utils.py
+++ b/sdks/python/apache_beam/utils/interactive_utils.py
@@ -19,6 +19,8 @@
For experimental usage only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/utils/plugin.py b/sdks/python/apache_beam/utils/plugin.py
index 1425874..7a92489 100644
--- a/sdks/python/apache_beam/utils/plugin.py
+++ b/sdks/python/apache_beam/utils/plugin.py
@@ -20,6 +20,8 @@
For experimental usage only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/utils/processes.py b/sdks/python/apache_beam/utils/processes.py
index cfd82bd..d1214d3 100644
--- a/sdks/python/apache_beam/utils/processes.py
+++ b/sdks/python/apache_beam/utils/processes.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import platform
diff --git a/sdks/python/apache_beam/utils/processes_test.py b/sdks/python/apache_beam/utils/processes_test.py
index 354a076..491bfeb 100644
--- a/sdks/python/apache_beam/utils/processes_test.py
+++ b/sdks/python/apache_beam/utils/processes_test.py
@@ -16,6 +16,8 @@
#
"""Unit tests for the processes module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import subprocess
diff --git a/sdks/python/apache_beam/utils/profiler.py b/sdks/python/apache_beam/utils/profiler.py
index b3761d1..ab1da2f 100644
--- a/sdks/python/apache_beam/utils/profiler.py
+++ b/sdks/python/apache_beam/utils/profiler.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import cProfile # pylint: disable=bad-python3-import
@@ -105,6 +107,7 @@
if random.random() < options.profile_sample_rate:
return Profile(profile_id, options.profile_location, **kwargs)
return create_profiler
+ return None
class MemoryReporter(object):
diff --git a/sdks/python/apache_beam/utils/proto_utils.py b/sdks/python/apache_beam/utils/proto_utils.py
index 2dfefe5..da9ca58 100644
--- a/sdks/python/apache_beam/utils/proto_utils.py
+++ b/sdks/python/apache_beam/utils/proto_utils.py
@@ -17,6 +17,8 @@
"""For internal use only; no backwards-compatibility guarantees."""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/utils/retry.py b/sdks/python/apache_beam/utils/retry.py
index c291393..d62c3cf 100644
--- a/sdks/python/apache_beam/utils/retry.py
+++ b/sdks/python/apache_beam/utils/retry.py
@@ -25,6 +25,8 @@
needed right now use a @retry.no_retries decorator.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
import functools
diff --git a/sdks/python/apache_beam/utils/retry_test.py b/sdks/python/apache_beam/utils/retry_test.py
index 7f471fb..8053b1d 100644
--- a/sdks/python/apache_beam/utils/retry_test.py
+++ b/sdks/python/apache_beam/utils/retry_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the retry module."""
+# pytype: skip-file
+
from __future__ import absolute_import
import unittest
diff --git a/sdks/python/apache_beam/utils/subprocess_server.py b/sdks/python/apache_beam/utils/subprocess_server.py
index fd55f18..b7c8030 100644
--- a/sdks/python/apache_beam/utils/subprocess_server.py
+++ b/sdks/python/apache_beam/utils/subprocess_server.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import logging
diff --git a/sdks/python/apache_beam/utils/thread_pool_executor.py b/sdks/python/apache_beam/utils/thread_pool_executor.py
index aba8f5ad..a71a113 100644
--- a/sdks/python/apache_beam/utils/thread_pool_executor.py
+++ b/sdks/python/apache_beam/utils/thread_pool_executor.py
@@ -15,6 +15,8 @@
# limitations under the License.
#
+# pytype: skip-file
+
from __future__ import absolute_import
import sys
@@ -25,7 +27,7 @@
try: # Python3
import queue
except Exception: # Python2
- import Queue as queue
+ import Queue as queue # type: ignore[no-redef]
class _WorkItem(object):
diff --git a/sdks/python/apache_beam/utils/thread_pool_executor_test.py b/sdks/python/apache_beam/utils/thread_pool_executor_test.py
index c82d0f9..d9bbae4 100644
--- a/sdks/python/apache_beam/utils/thread_pool_executor_test.py
+++ b/sdks/python/apache_beam/utils/thread_pool_executor_test.py
@@ -17,6 +17,8 @@
"""Unit tests for UnboundedThreadPoolExecutor."""
+# pytype: skip-file
+
from __future__ import absolute_import
import itertools
diff --git a/sdks/python/apache_beam/utils/timestamp.py b/sdks/python/apache_beam/utils/timestamp.py
index b9a683e..babe5dc 100644
--- a/sdks/python/apache_beam/utils/timestamp.py
+++ b/sdks/python/apache_beam/utils/timestamp.py
@@ -20,6 +20,8 @@
For internal use only; no backwards-compatibility guarantees.
"""
+# pytype: skip-file
+
from __future__ import absolute_import
from __future__ import division
diff --git a/sdks/python/apache_beam/utils/timestamp_test.py b/sdks/python/apache_beam/utils/timestamp_test.py
index c728b10..18e31f7 100644
--- a/sdks/python/apache_beam/utils/timestamp_test.py
+++ b/sdks/python/apache_beam/utils/timestamp_test.py
@@ -17,6 +17,8 @@
"""Unit tests for time utilities."""
+# pytype: skip-file
+
from __future__ import absolute_import
import datetime
diff --git a/sdks/python/apache_beam/utils/urns.py b/sdks/python/apache_beam/utils/urns.py
index 47b8584..1bd31e9 100644
--- a/sdks/python/apache_beam/utils/urns.py
+++ b/sdks/python/apache_beam/utils/urns.py
@@ -17,6 +17,8 @@
"""For internal use only; no backwards-compatibility guarantees."""
+# pytype: skip-file
+
from __future__ import absolute_import
import abc
diff --git a/sdks/python/apache_beam/utils/windowed_value.py b/sdks/python/apache_beam/utils/windowed_value.py
index 786484f..dc16a58 100644
--- a/sdks/python/apache_beam/utils/windowed_value.py
+++ b/sdks/python/apache_beam/utils/windowed_value.py
@@ -27,6 +27,8 @@
#cython: profile=True
+# pytype: skip-file
+
from __future__ import absolute_import
from builtins import object
diff --git a/sdks/python/apache_beam/utils/windowed_value_test.py b/sdks/python/apache_beam/utils/windowed_value_test.py
index 5549aee..27512d3 100644
--- a/sdks/python/apache_beam/utils/windowed_value_test.py
+++ b/sdks/python/apache_beam/utils/windowed_value_test.py
@@ -17,6 +17,8 @@
"""Unit tests for the windowed_value."""
+# pytype: skip-file
+
from __future__ import absolute_import
import copy
diff --git a/sdks/python/apache_beam/version.py b/sdks/python/apache_beam/version.py
index ba28fb7..be5d97f 100644
--- a/sdks/python/apache_beam/version.py
+++ b/sdks/python/apache_beam/version.py
@@ -18,4 +18,4 @@
"""Apache Beam SDK version information and utilities."""
-__version__ = '2.19.0.dev'
+__version__ = '2.20.0.dev'
diff --git a/sdks/python/build.gradle b/sdks/python/build.gradle
index d3f65a9..ca9cae8 100644
--- a/sdks/python/build.gradle
+++ b/sdks/python/build.gradle
@@ -48,7 +48,7 @@
args '-c', ". ${envdir}/bin/activate && python setup.py -q sdist --formats zip,gztar --dist-dir ${buildDir}"
}
- def collection = fileTree(buildDir){ include "**/*${project['python_sdk_version']}*.tar.gz" exclude 'srcs/**'}
+ def collection = fileTree(buildDir){ include "**/*${project.sdk_version}*.tar.gz" exclude 'srcs/**'}
// we need a fixed name for the artifact
copy { from collection.singleFile; into buildDir; rename { tarball } }
diff --git a/sdks/python/container/base_image_requirements.txt b/sdks/python/container/base_image_requirements.txt
index 359d6e5..47c5dd5 100644
--- a/sdks/python/container/base_image_requirements.txt
+++ b/sdks/python/container/base_image_requirements.txt
@@ -67,7 +67,7 @@
protorpc==0.12.0
python-gflags==3.0.6
-tensorflow==1.14.0
+tensorflow==2.1.0
pymongo==3.8.0
# Packages needed for testing.
diff --git a/sdks/python/container/py2/build.gradle b/sdks/python/container/py2/build.gradle
index 64f39f0..4e29de0 100644
--- a/sdks/python/container/py2/build.gradle
+++ b/sdks/python/container/py2/build.gradle
@@ -57,7 +57,7 @@
root: project.rootProject.hasProperty(["docker-repository-root"]) ?
project.rootProject["docker-repository-root"] : "apachebeam",
tag: project.rootProject.hasProperty(["docker-tag"]) ?
- project.rootProject["docker-tag"] : project['python_sdk_version'])
+ project.rootProject["docker-tag"] : project.sdk_version)
files "../Dockerfile", "./build"
buildArgs(['py_version': "2.7"])
}
diff --git a/sdks/python/container/py35/build.gradle b/sdks/python/container/py35/build.gradle
index 024847b..f8cd8c2 100644
--- a/sdks/python/container/py35/build.gradle
+++ b/sdks/python/container/py35/build.gradle
@@ -57,7 +57,7 @@
root: project.rootProject.hasProperty(["docker-repository-root"]) ?
project.rootProject["docker-repository-root"] : "apachebeam",
tag: project.rootProject.hasProperty(["docker-tag"]) ?
- project.rootProject["docker-tag"] : project['python_sdk_version'])
+ project.rootProject["docker-tag"] : project.sdk_version)
files "../Dockerfile", "./build"
buildArgs(['py_version': "3.5"])
}
diff --git a/sdks/python/container/py36/build.gradle b/sdks/python/container/py36/build.gradle
index f81f6ec..b7ced3d 100644
--- a/sdks/python/container/py36/build.gradle
+++ b/sdks/python/container/py36/build.gradle
@@ -57,7 +57,7 @@
root: project.rootProject.hasProperty(["docker-repository-root"]) ?
project.rootProject["docker-repository-root"] : "apachebeam",
tag: project.rootProject.hasProperty(["docker-tag"]) ?
- project.rootProject["docker-tag"] : project['python_sdk_version'])
+ project.rootProject["docker-tag"] : project.sdk_version)
files "../Dockerfile", "./build"
buildArgs(['py_version': "3.6"])
}
diff --git a/sdks/python/container/py37/build.gradle b/sdks/python/container/py37/build.gradle
index a7f10c4..d9e9b5d 100644
--- a/sdks/python/container/py37/build.gradle
+++ b/sdks/python/container/py37/build.gradle
@@ -57,7 +57,7 @@
root: project.rootProject.hasProperty(["docker-repository-root"]) ?
project.rootProject["docker-repository-root"] : "apachebeam",
tag: project.rootProject.hasProperty(["docker-tag"]) ?
- project.rootProject["docker-tag"] : project['python_sdk_version'])
+ project.rootProject["docker-tag"] : project.sdk_version)
files "../Dockerfile", "./build"
buildArgs(['py_version': "3.7"])
}
diff --git a/sdks/python/gen_protos.py b/sdks/python/gen_protos.py
index 5105ad5..3a7b23b 100644
--- a/sdks/python/gen_protos.py
+++ b/sdks/python/gen_protos.py
@@ -19,11 +19,14 @@
from __future__ import absolute_import
from __future__ import print_function
+import contextlib
import glob
+import inspect
import logging
import multiprocessing
import os
import platform
+import re
import shutil
import subprocess
import sys
@@ -47,6 +50,172 @@
]
+def generate_urn_files(log, out_dir):
+ """
+ Create python files with statically defined URN constants.
+
+ Creates a <proto>_pb2_urn.py file for each <proto>_pb2.py file that contains
+ an enum type.
+
+ This works by importing each api.<proto>_pb2 module created by `protoc`,
+ inspecting the module's contents, and generating a new side-car urn module.
+ This is executed at build time rather than dynamically on import to ensure
+ that it is compatible with static type checkers like mypy.
+ """
+ import google.protobuf.message as message
+ import google.protobuf.pyext._message as pyext_message
+
+ class Context(object):
+ INDENT = ' '
+ CAP_SPLIT = re.compile('([A-Z][^A-Z]*|^[a-z]+)')
+
+ def __init__(self, indent=0):
+ self.lines = []
+ self.imports = set()
+ self.empty_types = set()
+ self._indent = indent
+
+ @contextlib.contextmanager
+ def indent(self):
+ self._indent += 1
+ yield
+ self._indent -= 1
+
+ def prepend(self, s):
+ if s:
+ self.lines.insert(0, (self.INDENT * self._indent) + s + '\n')
+ else:
+ self.lines.insert(0, '\n')
+
+ def line(self, s):
+ if s:
+ self.lines.append((self.INDENT * self._indent) + s + '\n')
+ else:
+ self.lines.append('\n')
+
+ def import_type(self, typ):
+ modname = typ.__module__
+ if modname in ('__builtin__', 'builtin'):
+ return typ.__name__
+ else:
+ self.imports.add(modname)
+ return modname + '.' + typ.__name__
+
+ @staticmethod
+ def is_message_type(obj):
+ return isinstance(obj, type) and \
+ issubclass(obj, message.Message)
+
+ @staticmethod
+ def is_enum_type(obj):
+ return type(obj).__name__ == 'EnumTypeWrapper'
+
+ def python_repr(self, obj):
+ if isinstance(obj, message.Message):
+ return self.message_repr(obj)
+ elif isinstance(obj, (list,
+ pyext_message.RepeatedCompositeContainer, # pylint: disable=c-extension-no-member
+ pyext_message.RepeatedScalarContainer)): # pylint: disable=c-extension-no-member
+ return '[%s]' % ', '.join(self.python_repr(x) for x in obj)
+ else:
+ return repr(obj)
+
+ def empty_type(self, typ):
+ name = ('EMPTY_' +
+ '_'.join(x.upper()
+ for x in self.CAP_SPLIT.findall(typ.__name__)))
+ self.empty_types.add('%s = %s()' % (name, self.import_type(typ)))
+ return name
+
+ def message_repr(self, msg):
+ parts = []
+ for field, value in msg.ListFields():
+ parts.append('%s=%s' % (field.name, self.python_repr(value)))
+ if parts:
+ return '%s(%s)' % (self.import_type(type(msg)), ', '.join(parts))
+ else:
+ return self.empty_type(type(msg))
+
+ def write_enum(self, enum_name, enum, indent):
+ ctx = Context(indent=indent)
+
+ with ctx.indent():
+ for v in enum.DESCRIPTOR.values:
+ extensions = v.GetOptions().Extensions
+
+ prop = (
+ extensions[beam_runner_api_pb2.beam_urn],
+ extensions[beam_runner_api_pb2.beam_constant],
+ extensions[metrics_pb2.monitoring_info_spec],
+ extensions[metrics_pb2.label_props],
+ )
+ reprs = [self.python_repr(x) for x in prop]
+ if all(x == "''" or x.startswith('EMPTY_') for x in reprs):
+ continue
+ ctx.line('%s = PropertiesFromEnumValue(%s)' %
+ (v.name, ', '.join(self.python_repr(x) for x in prop)))
+
+ if ctx.lines:
+ ctx.prepend('class %s(object):' % enum_name)
+ ctx.prepend('')
+ ctx.line('')
+ return ctx.lines
+
+ def write_message(self, message_name, message, indent=0):
+ ctx = Context(indent=indent)
+
+ with ctx.indent():
+ for obj_name, obj in inspect.getmembers(message):
+ if self.is_message_type(obj):
+ ctx.lines += self.write_message(obj_name, obj, ctx._indent)
+ elif self.is_enum_type(obj):
+ ctx.lines += self.write_enum(obj_name, obj, ctx._indent)
+
+ if ctx.lines:
+ ctx.prepend('class %s(object):' % message_name)
+ ctx.prepend('')
+ return ctx.lines
+
+ pb2_files = [path for path in glob.glob(os.path.join(out_dir, '*_pb2.py'))]
+ api_path = os.path.dirname(pb2_files[0])
+ sys.path.insert(0, os.path.dirname(api_path))
+
+ def _import(m):
+ # TODO: replace with importlib when we drop support for python2.
+ return __import__('api.%s' % m, fromlist=[None])
+
+ try:
+ beam_runner_api_pb2 = _import('beam_runner_api_pb2')
+ metrics_pb2 = _import('metrics_pb2')
+
+ for pb2_file in pb2_files:
+ modname = os.path.splitext(pb2_file)[0]
+ out_file = modname + '_urns.py'
+ modname = os.path.basename(modname)
+ mod = _import(modname)
+
+ ctx = Context()
+ for obj_name, obj in inspect.getmembers(mod):
+ if ctx.is_message_type(obj):
+ ctx.lines += ctx.write_message(obj_name, obj)
+
+ if ctx.lines:
+ for line in reversed(sorted(ctx.empty_types)):
+ ctx.prepend(line)
+
+ for modname in reversed(sorted(ctx.imports)):
+ ctx.prepend('from . import %s' % modname)
+
+ ctx.prepend('from ..utils import PropertiesFromEnumValue')
+
+ log.info("Writing urn stubs: %s" % out_file)
+ with open(out_file, 'w') as f:
+ f.writelines(ctx.lines)
+
+ finally:
+ sys.path.pop(0)
+
+
def generate_proto_files(force=False, log=None):
try:
@@ -114,7 +283,8 @@
# Note that this requires a separate module from setup.py for Windows:
# https://docs.python.org/2/library/multiprocessing.html#windows
p = multiprocessing.Process(
- target=_install_grpcio_tools_and_generate_proto_files)
+ target=_install_grpcio_tools_and_generate_proto_files,
+ kwargs={'force': force})
p.start()
p.join()
if p.exitcode:
@@ -151,6 +321,11 @@
raise RuntimeError(
'Error applying futurize to generated protobuf python files.')
+ generate_urn_files(log, out_dir)
+
+ else:
+ log.info('Skipping proto regeneration: all files up to date')
+
# Though wheels are available for grpcio-tools, setup_requires uses
# easy_install which doesn't understand them. This means that it is
@@ -158,7 +333,7 @@
# protoc compiler). Instead, we attempt to install a wheel in a temporary
# directory and add it to the path as needed.
# See https://github.com/pypa/setuptools/issues/377
-def _install_grpcio_tools_and_generate_proto_files():
+def _install_grpcio_tools_and_generate_proto_files(force=False):
py_sdk_root = os.path.dirname(os.path.abspath(__file__))
install_path = os.path.join(py_sdk_root, '.eggs', 'grpcio-wheels')
build_path = install_path + '-build'
@@ -179,10 +354,11 @@
shutil.rmtree(build_path, ignore_errors=True)
sys.path.append(install_path)
try:
- generate_proto_files()
+ generate_proto_files(force=force)
finally:
sys.stderr.flush()
if __name__ == '__main__':
+ logging.getLogger().setLevel(logging.INFO)
generate_proto_files(force=True)
diff --git a/sdks/python/scripts/run_integration_test.sh b/sdks/python/scripts/run_integration_test.sh
index 1133147..0804ace 100755
--- a/sdks/python/scripts/run_integration_test.sh
+++ b/sdks/python/scripts/run_integration_test.sh
@@ -194,8 +194,10 @@
fi
# Install test dependencies for ValidatesRunner tests.
- echo "pyhamcrest" > postcommit_requirements.txt
- echo "mock" >> postcommit_requirements.txt
+ # pyhamcrest==1.10.0 doesn't work on Py2.
+ # See: https://github.com/hamcrest/PyHamcrest/issues/131.
+ echo "pyhamcrest!=1.10.0,<2.0.0" > postcommit_requirements.txt
+ echo "mock<3.0.0" >> postcommit_requirements.txt
# Options used to run testing pipeline on Cloud Dataflow Service. Also used for
# running on DirectRunner (some options ignored).
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index fcb4fd4..a396d17 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -138,7 +138,10 @@
'avro>=1.8.1,<2.0.0; python_version < "3.0"',
'avro-python3>=1.8.1,<2.0.0; python_version >= "3.0"',
'crcmod>=1.7,<2.0',
- # Dill doesn't guarantee compatibility between releases within minor version.
+ # Dill doesn't have forwards-compatibility guarantees within minor version.
+ # Pickles created with a new version of dill may not unpickle using older
+ # version of dill. It is best to use the same version of dill on client and
+ # server, therefore list of allowed versions is very narrow.
# See: https://github.com/uqfoundation/dill/issues/341.
'dill>=0.3.1.1,<0.3.2',
'fastavro>=0.21.4,<0.22',
@@ -178,7 +181,9 @@
'nose_xunitmp>=0.4.1',
'pandas>=0.23.4,<0.25',
'parameterized>=0.6.0,<0.8.0',
- 'pyhamcrest>=1.9,<2.0',
+ # pyhamcrest==1.10.0 doesn't work on Py2. Beam still supports Py2.
+ # See: https://github.com/hamcrest/PyHamcrest/issues/131.
+ 'pyhamcrest>=1.9,!=1.10.0,<2.0.0',
'pyyaml>=3.12,<6.0.0',
'requests_mock>=1.7,<2.0',
'tenacity>=5.0.2,<6.0',
@@ -199,6 +204,8 @@
'google-cloud-bigtable>=0.31.1,<1.1.0',
# [BEAM-4543] googledatastore is not supported in Python 3.
'proto-google-cloud-datastore-v1>=0.90.0,<=0.90.4; python_version < "3.0"',
+ 'google-cloud-spanner>=1.7.1<1.8.0',
+ 'grpcio-gcp>=0.2.2,<1',
]
INTERACTIVE_BEAM = [
diff --git a/sdks/python/test-suites/dataflow/py2/build.gradle b/sdks/python/test-suites/dataflow/py2/build.gradle
index 098df88..248f71d 100644
--- a/sdks/python/test-suites/dataflow/py2/build.gradle
+++ b/sdks/python/test-suites/dataflow/py2/build.gradle
@@ -146,18 +146,19 @@
}
}
-task dataflowChicagoTaxiExample {
+task chicagoTaxiExample {
dependsOn 'installChicagoTaxiExampleRequirements'
dependsOn ':sdks:python:sdist'
def gcsRoot = findProperty('gcsRoot')
- def cliArgs = "${gcsRoot} DataflowRunner ${files(configurations.distTarBall.files).singleFile}"
+ def pipelineOptions = findProperty('pipelineOptions') ?: ""
+ pipelineOptions += " --sdk_location=\"${files(configurations.distTarBall.files).singleFile}\""
doLast {
exec {
workingDir "$rootProject.projectDir/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/"
executable 'sh'
- args '-c', ". ${envdir}/bin/activate && ./run_chicago.sh ${cliArgs}"
+ args '-c', ". ${envdir}/bin/activate && ./run_chicago.sh ${gcsRoot} DataflowRunner ${pipelineOptions}"
}
}
}
diff --git a/sdks/python/test-suites/direct/py2/build.gradle b/sdks/python/test-suites/direct/py2/build.gradle
index 436cd7e..acb7fee 100644
--- a/sdks/python/test-suites/direct/py2/build.gradle
+++ b/sdks/python/test-suites/direct/py2/build.gradle
@@ -77,11 +77,10 @@
}
task hdfsIntegrationTest {
- dependsOn 'installGcpTest'
doLast {
exec {
executable 'sh'
- args '-c', ". ${envdir}/bin/activate && ${rootDir}/sdks/python/apache_beam/io/hdfs_integration_test/hdfs_integration_test.sh python:2"
+ args '-c', "${rootDir}/sdks/python/apache_beam/io/hdfs_integration_test/hdfs_integration_test.sh python:2"
}
}
}
diff --git a/sdks/python/test-suites/direct/py37/build.gradle b/sdks/python/test-suites/direct/py37/build.gradle
index c5c30fb..1f54f90 100644
--- a/sdks/python/test-suites/direct/py37/build.gradle
+++ b/sdks/python/test-suites/direct/py37/build.gradle
@@ -56,11 +56,10 @@
}
task hdfsIntegrationTest {
- dependsOn 'installGcpTest'
doLast {
exec {
executable 'sh'
- args '-c', ". ${envdir}/bin/activate && ${pythonDir}/apache_beam/io/hdfs_integration_test/hdfs_integration_test.sh python:3.7"
+ args '-c', "${pythonDir}/apache_beam/io/hdfs_integration_test/hdfs_integration_test.sh python:3.7"
}
}
}
diff --git a/sdks/python/test-suites/portable/py2/build.gradle b/sdks/python/test-suites/portable/py2/build.gradle
index f9fb5f4..128a5f6 100644
--- a/sdks/python/test-suites/portable/py2/build.gradle
+++ b/sdks/python/test-suites/portable/py2/build.gradle
@@ -175,6 +175,33 @@
dependsOn "crossLanguagePortableWordCount"
}
+task installChicagoTaxiExampleRequirements {
+ dependsOn 'installGcpTest'
+
+ doLast {
+ exec {
+ workingDir "$rootProject.projectDir/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/"
+ executable 'sh'
+ args '-c', ". ${envdir}/bin/activate && pip install -r requirements.txt"
+ }
+ }
+}
+
+task chicagoTaxiExample {
+ dependsOn 'installChicagoTaxiExampleRequirements'
+
+ def gcsRoot = findProperty('gcsRoot')
+ def pipelineOptions = findProperty('pipelineOptions') ?: ""
+
+ doLast {
+ exec {
+ workingDir "$rootProject.projectDir/sdks/python/apache_beam/testing/benchmarks/chicago_taxi/"
+ executable 'sh'
+ args '-c', ". ${envdir}/bin/activate && ./run_chicago.sh ${gcsRoot} PortableRunner ${pipelineOptions}"
+ }
+ }
+}
+
/*************************************************************************************************/
task createProcessWorker {
diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini
index c78bddf..e6908d9 100644
--- a/sdks/python/tox.ini
+++ b/sdks/python/tox.ini
@@ -315,3 +315,30 @@
coverage report --skip-covered
# Generate report in xml format
coverage xml
+
+[testenv:hdfs_integration_test]
+# Used by hdfs_integration_test.sh. Do not run this directly, as it depends on
+# nodes defined in hdfs_integration_test/docker-compose.yml.
+deps =
+ -r build-requirements.txt
+ gsutil==4.47
+ holdup==1.8.0
+extras =
+ gcp
+whitelist_externals =
+ echo
+ sleep
+passenv = HDFSCLI_CONFIG
+commands =
+ holdup -t 45 http://namenode:50070 http://datanode:50075
+ echo "Waiting for safe mode to end."
+ sleep 45
+ gsutil cp gs://dataflow-samples/shakespeare/kinglear.txt .
+ hdfscli -v -v -v upload -f kinglear.txt /
+ python -m apache_beam.examples.wordcount \
+ --input hdfs://kinglear* \
+ --output hdfs://py-wordcount-integration \
+ --hdfs_host namenode --hdfs_port 50070 --hdfs_user root
+# Disable pip check. TODO: remove this once gsutil does not conflict with
+# apache_beam (oauth2client).
+commands_pre =
diff --git a/settings.gradle b/settings.gradle
index 2952e09..c7ca4db 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -74,6 +74,7 @@
include ":sdks:java:extensions:sketching"
include ":sdks:java:extensions:sorter"
include ":sdks:java:extensions:sql"
+include ":sdks:java:extensions:sql:perf-tests"
include ":sdks:java:extensions:sql:jdbc"
include ":sdks:java:extensions:sql:shell"
include ":sdks:java:extensions:sql:hcatalog"
@@ -147,7 +148,7 @@
include ":sdks:python:test-suites:tox:py35"
include ":sdks:python:test-suites:tox:py36"
include ":sdks:python:test-suites:tox:py37"
-include ":vendor:grpc-1_21_0"
+include ":vendor:grpc-1_26_0"
include ":vendor:bytebuddy-1_9_3"
include ":vendor:calcite-1_20_0"
include ":vendor:guava-26_0-jre"
diff --git a/vendor/README.md b/vendor/README.md
new file mode 100644
index 0000000..c4b38cd
--- /dev/null
+++ b/vendor/README.md
@@ -0,0 +1,38 @@
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+-->
+
+# Vendored Dependencies Release
+
+The upgrading of the vendored dependencies should be performed in two steps:
+- Firstly, we need to perform a formal release of the vendored dependency.
+ The [release process](http://s.apache.org/beam-release-vendored-artifacts) of the vendored dependency
+ is separate from the release of Apache Beam.
+- When the release of the vendored dependency is out, we can migrate Apache Beam to use the newly released
+ vendored dependency.
+
+# How to validate the vendored dependencies
+
+The [linkage tool](https://lists.apache.org/thread.html/eb5d95b9a33d7e32dc9bcd0f7d48ba8711d42bd7ed03b9cf0f1103f1%40%3Cdev.beam.apache.org%3E)
+is useful for the vendored dependency upgrades. It reports the linkage errors across multiple Apache Beam artifact ids.
+
+For example, when we upgrade the version of gRPC to 1.26.0 and the version of the vendored gRPC is 0.1-SNAPSHOT,
+we could run the linkage tool as following:
+```
+./gradlew -PvendoredDependenciesOnly -Ppublishing -PjavaLinkageArtifactIds=beam-vendor-grpc-1_26_0:0.1-SNAPSHOT :checkJavaLinkage
+```
diff --git a/vendor/grpc-1_21_0/build.gradle b/vendor/grpc-1_26_0/build.gradle
similarity index 72%
rename from vendor/grpc-1_21_0/build.gradle
rename to vendor/grpc-1_26_0/build.gradle
index 16ded33..85f68bb 100644
--- a/vendor/grpc-1_21_0/build.gradle
+++ b/vendor/grpc-1_26_0/build.gradle
@@ -16,21 +16,22 @@
* limitations under the License.
*/
-import org.apache.beam.gradle.GrpcVendoring
+import org.apache.beam.gradle.GrpcVendoring_1_26_0
plugins { id 'org.apache.beam.vendor-java' }
-description = "Apache Beam :: Vendored Dependencies :: gRPC :: 1.21.0"
+description = "Apache Beam :: Vendored Dependencies :: gRPC :: 1.26.0"
group = "org.apache.beam"
version = "0.1"
vendorJava(
- dependencies: GrpcVendoring.dependencies(),
- runtimeDependencies: GrpcVendoring.runtimeDependencies(),
- relocations: GrpcVendoring.relocations(),
- exclusions: GrpcVendoring.exclusions(),
- artifactId: "beam-vendor-grpc-1_21_0",
+ dependencies: GrpcVendoring_1_26_0.dependencies(),
+ runtimeDependencies: GrpcVendoring_1_26_0.runtimeDependencies(),
+ testDependencies: GrpcVendoring_1_26_0.testDependencies(),
+ relocations: GrpcVendoring_1_26_0.relocations(),
+ exclusions: GrpcVendoring_1_26_0.exclusions(),
+ artifactId: "beam-vendor-grpc-1_26_0",
groupId: group,
version: version,
)
diff --git a/vendor/sdks-java-extensions-protobuf/build.gradle b/vendor/sdks-java-extensions-protobuf/build.gradle
index e3f0c949..a174db6 100644
--- a/vendor/sdks-java-extensions-protobuf/build.gradle
+++ b/vendor/sdks-java-extensions-protobuf/build.gradle
@@ -16,14 +16,16 @@
* limitations under the License.
*/
+import org.apache.beam.gradle.GrpcVendoring_1_26_0
+
plugins { id 'org.apache.beam.module' }
applyJavaNature(
automaticModuleName: 'org.apache.beam.vendor.sdks.java.extensions.protobuf',
exportJavadoc: false,
shadowClosure: {
dependencies {
- include(dependency('com.google.guava:guava:26.0-jre'))
- include(dependency('com.google.protobuf:protobuf-java:3.7.1'))
+ include(dependency("com.google.guava:guava:${GrpcVendoring_1_26_0.guava_version}"))
+ include(dependency("com.google.protobuf:protobuf-java:${GrpcVendoring_1_26_0.protobuf_version}"))
}
// We specifically relocate beam-sdks-extensions-protobuf under a vendored namespace
// but also vendor guava and protobuf to the same vendored namespace as the model/*
@@ -32,10 +34,10 @@
relocate "org.apache.beam.sdk.extensions.protobuf", "org.apache.beam.vendor.sdk.v2.sdk.extensions.protobuf"
// guava uses the com.google.common and com.google.thirdparty package namespaces
- relocate "com.google.common", "org.apache.beam.vendor.grpc.v1p21p0.com.google.common"
- relocate "com.google.thirdparty", "org.apache.beam.vendor.grpc.v1p21p0.com.google.thirdparty"
+ relocate "com.google.common", "org.apache.beam.vendor.grpc.v1p26p0.com.google.common"
+ relocate "com.google.thirdparty", "org.apache.beam.vendor.grpc.v1p26p0.com.google.thirdparty"
- relocate "com.google.protobuf", "org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf"
+ relocate "com.google.protobuf", "org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf"
}
)
@@ -54,7 +56,7 @@
}
dependencies {
- compile 'com.google.guava:guava:26.0-jre'
- compile 'com.google.protobuf:protobuf-java:3.7.1'
+ compile "com.google.guava:guava:${GrpcVendoring_1_26_0.guava_version}"
+ compile "com.google.protobuf:protobuf-java:${GrpcVendoring_1_26_0.protobuf_version}"
shadow project(path: ":sdks:java:core", configuration: "shadow")
}
diff --git a/website/_config.yml b/website/_config.yml
index ff65ef8..75c534a 100644
--- a/website/_config.yml
+++ b/website/_config.yml
@@ -62,7 +62,7 @@
toc_levels: 2..6
# The most recent release of Beam.
-release_latest: 2.17.0
+release_latest: 2.18.0
# Plugins are configured in the Gemfile.
diff --git a/website/src/.htaccess b/website/src/.htaccess
index 3f3eb2b..b620b2e 100644
--- a/website/src/.htaccess
+++ b/website/src/.htaccess
@@ -21,4 +21,4 @@
# The following redirect maintains the previously supported URLs.
RedirectMatch permanent "/documentation/sdks/(javadoc|pydoc)(.*)" "https://beam.apache.org/releases/$1$2"
# Keep this updated to point to the current release.
-RedirectMatch "/releases/([^/]+)/current(.*)" "https://beam.apache.org/releases/$1/2.17.0$2"
+RedirectMatch "/releases/([^/]+)/current(.*)" "https://beam.apache.org/releases/$1/2.18.0$2"
diff --git a/website/src/_data/authors.yml b/website/src/_data/authors.yml
index a8e874d..5812439 100644
--- a/website/src/_data/authors.yml
+++ b/website/src/_data/authors.yml
@@ -135,3 +135,7 @@
name: Tanay Tummalapalli
email: ttanay100@gmail.com
twitter: ttanay100
+udim:
+ name: Udi Meiri
+ email: udim@apache.org
+ twitter: udim
diff --git a/website/src/_includes/section-menu/get-started.html b/website/src/_includes/section-menu/get-started.html
index 61270a4..96facb3 100644
--- a/website/src/_includes/section-menu/get-started.html
+++ b/website/src/_includes/section-menu/get-started.html
@@ -30,4 +30,4 @@
</ul>
</li>
<li><a href="{{ site.baseurl }}/get-started/downloads">Downloads</a></li>
-
+<li><a href="{{ site.baseurl }}/security">Security</a></li>
diff --git a/website/src/_posts/2017-02-13-stateful-processing.md b/website/src/_posts/2017-02-13-stateful-processing.md
index a778dde..1105b7c 100644
--- a/website/src/_posts/2017-02-13-stateful-processing.md
+++ b/website/src/_posts/2017-02-13-stateful-processing.md
@@ -230,7 +230,7 @@
This presents a good opportunity to talk about big data and parallelism,
because the algorithm in those bullet points is not parallelizable at all! If
you wanted to apply this logic over an entire `PCollection`, you would have to
-process each element of the `PCollection` one-at-a-time... this is obvious a
+process each element of the `PCollection` one-at-a-time... this is obviously a
bad idea. State in Beam is tightly scoped so that most of the time a stateful
`ParDo` transform should still be possible for a runner to execute in parallel,
though you still have to be thoughtful about it.
diff --git a/website/src/_posts_2019-12-16-beam-2.17.0.md b/website/src/_posts/2020-01-06-beam-2.17.0.md
similarity index 94%
rename from website/src/_posts_2019-12-16-beam-2.17.0.md
rename to website/src/_posts/2020-01-06-beam-2.17.0.md
index 7997f9e..5c7caeb 100644
--- a/website/src/_posts_2019-12-16-beam-2.17.0.md
+++ b/website/src/_posts/2020-01-06-beam-2.17.0.md
@@ -3,7 +3,7 @@
title: "Apache Beam 2.17.0"
date: 2020-01-06 00:00:01 -0800
# Date above corrected but keep the old URL:
-permalink: /blog/2019/12/16/beam-2.17.0.html
+permalink: /blog/2020/01/06/beam-2.17.0.html
excerpt_separator: <!--more-->
categories: blog
authors:
@@ -25,7 +25,9 @@
-->
We are happy to present the new 2.17.0 release of Beam. This release includes both improvements and new functionality.
-See the [download page]({{ site.baseurl }}/get-started/downloads/#2170-2019-12-16) for this release.<!--more-->
+Users of the MongoDbIO connector are encouraged to upgrade to this release to address a [security vulnerability]({{ site.baseurl }}/security/CVE-2020-1929/).
+
+See the [download page]({{ site.baseurl }}/get-started/downloads/#2170-2020-01-06) for this release.<!--more-->
For more information on changes in 2.17.0, check out the
[detailed release notes](https://issues.apache.org/jira/secure/ReleaseNote.jspa?version=12345970&projectId=12319527).
@@ -40,7 +42,7 @@
### New Features / Improvements
* [BEAM-7730](https://issues.apache.org/jira/browse/BEAM-7730) - Add Flink 1.9 build target and Make FlinkRunner compatible with Flink 1.9.
-* [BEAM-7990](https://issues.apache.org/jira/browse/BEAM-7990) - Add ability to read parquet files into PCollection<pyarrow.Table>.
+* [BEAM-7990](https://issues.apache.org/jira/browse/BEAM-7990) - Add ability to read parquet files into PCollection of pyarrow.Table.
* [BEAM-8355](https://issues.apache.org/jira/browse/BEAM-8355) - Make BooleanCoder a standard coder.
* [BEAM-8394](https://issues.apache.org/jira/browse/BEAM-8394) - Add withDataSourceConfiguration() method in JdbcIO.ReadRows class.
* [BEAM-5428](https://issues.apache.org/jira/browse/BEAM-5428) - Implement cross-bundle state caching.
diff --git a/website/src/_posts/2020-01-13-beam-2.18.0.md b/website/src/_posts/2020-01-13-beam-2.18.0.md
new file mode 100644
index 0000000..c434645
--- /dev/null
+++ b/website/src/_posts/2020-01-13-beam-2.18.0.md
@@ -0,0 +1,108 @@
+---
+layout: post
+title: "Apache Beam 2.18.0"
+date: 2020-01-23 00:00:01 -0800
+# Date above corrected but keep the old URL:
+permalink: /blog/2020/01/13/beam-2.18.0.html
+excerpt_separator: <!--more-->
+categories: blog
+authors:
+ - udim
+ - altay
+
+---
+<!--
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+We are happy to present the new 2.18.0 release of Beam. This release includes both improvements and new functionality.
+See the [download page]({{ site.baseurl }}/get-started/downloads/#2180-2020-01-23) for this release.<!--more-->
+For more information on changes in 2.18.0, check out the
+[detailed release notes](https://issues.apache.org/jira/secure/ReleaseNote.jspa?version=12346383&projectId=12319527).
+
+## Highlights
+
+ * [BEAM-8470](https://issues.apache.org/jira/browse/BEAM-8470) - Create a new Spark runner based on Spark Structured streaming framework
+
+### I/Os
+* [BEAM-7636](https://issues.apache.org/jira/browse/BEAM-7636) - Added SqsIO v2 support.
+* [BEAM-8513](https://issues.apache.org/jira/browse/BEAM-8513) - RabbitMqIO: Allow reads from exchange-bound queue without declaring the exchange.
+* [BEAM-8540](https://issues.apache.org/jira/browse/BEAM-8540) - Fix CSVSink example in FileIO docs
+
+### New Features / Improvements
+
+* [BEAM-5878](https://issues.apache.org/jira/browse/BEAM-5878) - Added support DoFns with Keyword-only arguments in Python 3.
+* [BEAM-6756](https://issues.apache.org/jira/browse/BEAM-6756) - Improved support for lazy iterables in schemas (Java).
+* [BEAM-4776](https://issues.apache.org/jira/browse/BEAM-4776) AND [BEAM-4777](https://issues.apache.org/jira/browse/BEAM-4777) - Added metrics supports to portable runners.
+* Various improvements to Interactive Beam: [BEAM-7760](https://issues.apache.org/jira/browse/BEAM-7760), [BEAM-8379](https://issues.apache.org/jira/browse/BEAM-8379), [BEAM-8016](https://issues.apache.org/jira/browse/BEAM-8016), [BEAM-8016](https://issues.apache.org/jira/browse/BEAM-8016).
+* [BEAM-8658](https://issues.apache.org/jira/browse/BEAM-8658) - Optionally set artifact staging port in FlinkUberJarJobServer.
+* [BEAM-8660](https://issues.apache.org/jira/browse/BEAM-8660) - Override returned artifact staging endpoint
+
+### SQL
+* [BEAM-8343](https://issues.apache.org/jira/browse/BEAM-8343) - [SQL] Add means for IO APIs to support predicate and/or project push-down when running SQL pipelines. And [BEAM-8468](https://issues.apache.org/jira/browse/BEAM-8468), [BEAM-8365](https://issues.apache.org/jira/browse/BEAM-8365), [BEAM-8508](https://issues.apache.org/jira/browse/BEAM-8508).
+* [BEAM-8427](https://issues.apache.org/jira/browse/BEAM-8427) - [SQL] Add support for MongoDB source.
+* [BEAM-8456](https://issues.apache.org/jira/browse/BEAM-8456) - Add pipeline option to control truncate of BigQuery data processed by Beam SQL.
+
+### Breaking Changes
+
+* [BEAM-8814](https://issues.apache.org/jira/browse/BEAM-8814) - --no_auth flag changed to boolean type.
+
+
+### Deprecations
+
+* [BEAM-8252](https://issues.apache.org/jira/browse/BEAM-8252) AND [BEAM-8254](https://issues.apache.org/jira/browse/BEAM-8254) Add worker_region and worker_zone options. Deprecated --zone flag and --worker_region experiment argument.
+
+### Dependency Changes
+* [BEAM-7078](https://issues.apache.org/jira/browse/BEAM-7078) - com.amazonaws:amazon-kinesis-client updated to 1.13.0.
+* [BEAM-8822](https://issues.apache.org/jira/browse/BEAM-8822) - Upgrade Hadoop dependencies to version 2.8.
+
+### Bugfixes
+
+* [BEAM-7917](https://issues.apache.org/jira/browse/BEAM-7917) - Python datastore v1new fails on retry.
+* [BEAM-7981](https://issues.apache.org/jira/browse/BEAM-7981) - ParDo function wrapper doesn't support Iterable output types.
+* [BEAM-8146](https://issues.apache.org/jira/browse/BEAM-8146) - SchemaCoder/RowCoder have no equals() function.
+* [BEAM-8347](https://issues.apache.org/jira/browse/BEAM-8347) - UnboundedRabbitMqReader can fail to advance watermark if no new data comes in.
+* [BEAM-8352](https://issues.apache.org/jira/browse/BEAM-8352) - Reading records in background may lead to OOM errors
+* [BEAM-8480](https://issues.apache.org/jira/browse/BEAM-8480) - Explicitly set restriction coder for bounded reader wrapper SDF.
+* [BEAM-8515](https://issues.apache.org/jira/browse/BEAM-8515) - Ensure that ValueProvider types have equals/hashCode implemented for comparison reasons.
+* [BEAM-8579](https://issues.apache.org/jira/browse/BEAM-8579) - Strip UTF-8 BOM bytes (if present) in TextSource.
+* [BEAM-8657](https://issues.apache.org/jira/browse/BEAM-8657) - Not doing Combiner lifting for data-driven triggers.
+* [BEAM-8663](https://issues.apache.org/jira/browse/BEAM-8663) - BundleBasedRunner Stacked Bundles don't respect PaneInfo.
+* [BEAM-8667](https://issues.apache.org/jira/browse/BEAM-8667) - Data channel should to avoid unlimited buffering in Python SDK.
+* [BEAM-8802](https://issues.apache.org/jira/browse/BEAM-8802) - Timestamp combiner not respected across bundles in streaming mode.
+* [BEAM-8803](https://issues.apache.org/jira/browse/BEAM-8803) - Default behaviour for Python BQ Streaming inserts sink should be to retry always.
+* [BEAM-8825](https://issues.apache.org/jira/browse/BEAM-8825) - OOM when writing large numbers of 'narrow' rows.
+* [BEAM-8835](https://issues.apache.org/jira/browse/BEAM-8835) - Artifact retrieval fails with FlinkUberJarJobServer
+* [BEAM-8836](https://issues.apache.org/jira/browse/BEAM-8836) - ExternalTransform is not providing a unique name
+* [BEAM-8884](https://issues.apache.org/jira/browse/BEAM-8884) - Python MongoDBIO TypeError when splitting.
+* [BEAM-9041](https://issues.apache.org/jira/browse/BEAM-9041) - SchemaCoder equals should not rely on from/toRowFunction equality.
+* [BEAM-9042](https://issues.apache.org/jira/browse/BEAM-9042) - AvroUtils.schemaCoder(schema) produces a not serializable SchemaCoder.
+* [BEAM-9065](https://issues.apache.org/jira/browse/BEAM-9065) - Spark runner accumulates metrics (incorrectly) between runs.
+* [BEAM-6303](https://issues.apache.org/jira/browse/BEAM-6303) - Add .parquet extension to files in ParquetIO.
+* Various bug fixes and performance improvements.
+
+### Known Issues
+
+* [BEAM-8882](https://issues.apache.org/jira/browse/BEAM-8882) - Python: `beam.Create` no longer preserves order unless `reshuffle=False` is passed in as an argument.
+
+ You may encounter this issue when using DirectRunner.
+* [BEAM-9065](https://issues.apache.org/jira/browse/BEAM-9065) - Spark runner accumulates metrics (incorrectly) between runs
+* [BEAM-9123](https://issues.apache.org/jira/browse/BEAM-9123) - HadoopResourceId returns wrong directory name
+* See a full list of open [issues that affect](https://issues.apache.org/jira/issues/?jql=project%20%3D%20BEAM%20AND%20affectedVersion%20%3D%202.18.0%20ORDER%20BY%20priority%20DESC%2C%20updated%20DESC) this version.
+
+
+## List of Contributors
+
+According to git shortlog, the following people contributed to the 2.18.0 release. Thank you to all contributors!
+
+Ahmet Altay, Aizhamal Nurmamat kyzy, Alan Myrvold, Alexey Romanenko, Alex Van Boxel, Andre Araujo, Andrew Crites, Andrew Pilloud, Aryan Naraghi, Boyuan Zhang, Brian Hulette, bumblebee-coming, Cerny Ondrej, Chad Dombrova, Chamikara Jayalath, Changming Ma, Chun Yang, cmachgodaddy, Colm O hEigeartaigh, Craig Chambers, Daniel Oliveira, Daniel Robert, David Cavazos, David Moravek, David Song, dependabot[bot], Derek, Dmytro Sadovnychyi, Elliotte Rusty Harold, Etienne Chauchot, Hai Lu, Henry Suryawirawan, Ismaël Mejía, Jack Whelpton, Jan Lukavský, Jean-Baptiste Onofré, Jeff Klukas, Jincheng Sun, Jing, Jing Chen, Joe Tsai, Jonathan Alvarez-Gutierrez, Kamil Wasilewski, KangZhiDong, Kasia Kucharczyk, Kenneth Knowles, kirillkozlov, Kirill Kozlov, Kyle Weaver, liumomo315, lostluck, Łukasz Gajowy, Luke Cwik, Mark Liu, Maximilian Michels, Michal Walenia, Mikhail Gryzykhin, Niel Markwick, Ning Kang, nlofeudo, pabloem, Pablo Estrada, Pankaj Gudlani, Piotr Szczepanik, Primevenn, Reuven Lax, Robert Bradshaw, Robert Burke, Rui Wang, Ruoyun Huang, RusOr10n, Ryan Skraba, Saikat Maitra, sambvfx, Sam Rohde, Samuel Husso, Stefano, Steve Koonce, Steve Niemitz, sunjincheng121, Thomas Weise, Tianyang Hu, Tim Robertson, Tomo Suzuki, tvalentyn, Udi Meiri, Valentyn Tymofieiev, Viola Lyu, Wenjia Liu, Yichi Zhang, Yifan Zou, yoshiki.obata, Yueyang Qiu, ziel, 康智冬
diff --git a/website/src/contribute/index.md b/website/src/contribute/index.md
index 9d84167..984d522 100644
--- a/website/src/contribute/index.md
+++ b/website/src/contribute/index.md
@@ -38,9 +38,9 @@
- review proposed design ideas on [dev@beam.apache.org]({{ site.baseurl
}}/community/contact-us/)
- improve the documentation
- - contribute [bug reports](https://issues.apache.org/jira/projects/BEAM/issues)
- - contribute by testing releases
- - contribute by reviewing [changes](https://github.com/apache/beam/pulls)
+ - file [bug reports](https://issues.apache.org/jira/projects/BEAM/issues)
+ - test releases
+ - review [changes](https://github.com/apache/beam/pulls)
- write new examples
- improve your favorite language SDK (Java, Python, Go, etc)
- improve specific runners (Apache Apex, Apache Flink, Apache Spark, Google
@@ -58,7 +58,7 @@
## Contributing code
-Below is a tutorial for contributing [code to Beam](https://github.com/apache/beam), covering our tools and typical process in
+Below is a tutorial for contributing code to Beam, covering our tools and typical process in
detail.
### Prerequisites
@@ -67,14 +67,38 @@
- a GitHub account
- a Linux, macOS, or Microsoft Windows development environment with Java JDK 8 installed
- - [Docker](https://www.docker.com/) installed for some tasks including building worker containers and testing website
+ - [Docker](https://www.docker.com/) installed for some tasks including building worker containers and testing this website
changes locally
- - [Go](https://golang.org) 1.10 or later installed for Go SDK development
- - Python, virtualenv, and tox installed for Python SDK development
+ - [Go](https://golang.org) 1.12 or later installed for Go SDK development
+ - Python 2.7, 3.5, 3.6, and 3.7. Yes, you need all four versions installed.
+ - pip, setuptools, virtualenv, and tox installed for Python development
- for large contributions, a signed [Individual Contributor License
Agreement](https://www.apache.org/licenses/icla.pdf) (ICLA) to the Apache
Software Foundation (ASF).
+To install these in a Debian-based distribution:
+
+```
+sudo apt-get install \
+ openjdk-8-jdk \
+ python-setuptools \
+ python-pip \
+ virtualenv \
+ tox \
+ docker-ce
+```
+
+You also need to [install Go](https://golang.org/doc/install]).
+
+Once Go is installed, install goavro:
+
+```
+$ export GOPATH=`pwd`/sdks/go/examples/.gogradle/project_gopath
+$ go get github.com/linkedin/goavro
+```
+
+gLinux users should configure their machines for sudoless Docker.
+
### Connect With the Beam community
1. Consider subscribing to the [dev@ mailing list]({{ site.baseurl}}/community/contact-us/), especially
@@ -182,7 +206,7 @@
1. If you don't get any response in 3 business days, email the [dev@ mailing list]({{ site.baseurl }}/community/contact-us) to ask for someone to look at your pull
request.
-### Make reviewer's job easier
+### Make the reviewer's job easier
1. Provide context for your changes in the associated JIRA issue and/or PR description.
diff --git a/website/src/contribute/release-guide.md b/website/src/contribute/release-guide.md
index 2a27bde..c110dce 100644
--- a/website/src/contribute/release-guide.md
+++ b/website/src/contribute/release-guide.md
@@ -125,7 +125,7 @@
* Determine your Apache GPG Key and Key ID, as follows:
- gpg --list-keys
+ gpg --list-sigs --keyid-format LONG
This will list your GPG keys. One of these should reflect your Apache account, for example:
@@ -377,12 +377,15 @@
After the release branch is cut you need to make sure it builds and has no significant issues that would block the creation of the release candidate.
There are 2 ways to perform this verification, either running automation script(recommended), or running all commands manually.
+! Dataflow tests will fail if Dataflow worker container is not created and published by this time. (Should be done by Google)
+
#### Run automation script (verify_release_build.sh)
* Script: [verify_release_build.sh](https://github.com/apache/beam/blob/master/release/src/main/scripts/verify_release_build.sh)
* Usage
1. Create a personal access token from your Github account. See instruction [here](https://help.github.com/en/articles/creating-a-personal-access-token-for-the-command-line).
It'll be used by the script for accessing Github API.
+ You don't have to add any permissions to this token.
1. Update required configurations listed in `RELEASE_BUILD_CONFIGS` in [script.config](https://github.com/apache/beam/blob/master/release/src/main/scripts/script.config)
1. Then run
```
@@ -397,7 +400,9 @@
1. Create a test PR against release branch;
Jenkins job `beam_Release_Gradle_Build` basically run `./gradlew build -PisRelease`.
-This only verifies that everything builds with unit tests passing.
+This only verifies that everything builds with unit tests passing.
+
+You can refer to [this script](https://gist.github.com/Ardagan/13e6031e8d1c9ebbd3029bf365c1a517) to mass-comment on PR.
#### Verify the build succeeds
@@ -579,9 +584,6 @@
* The script will:
1. Run gradle release to create rc tag and push source release into github repo.
1. Run gradle publish to push java artifacts into Maven staging repo.
-
- __NOTE__: In order to public staging artifacts, you need to goto the [staging repo](https://repository.apache.org/#stagingRepositories) to close the staging repository on Apache Nexus.
- When prompted for a description, enter “Apache Beam, version X, release candidate Y”.
1. Stage source release into dist.apache.org dev [repo](https://dist.apache.org/repos/dist/dev/beam/).
1. Stage,sign and hash python binaries into dist.apache.ord dev repo python dir
1. Stage SDK docker images to [https://hub.docker.com/u/apachebeam](https://hub.docker.com/u/apachebeam).
@@ -595,6 +597,9 @@
1. Update last release download links in `website/src/get-started/downloads.md`.
1. Update `website/src/.htaccess` to redirect to the new version.
1. Build and stage python wheels.
+ 1. Publish staging artifacts
+ 1. Go to the staging repo to close the staging repository on [Apache Nexus](https://repository.apache.org/#stagingRepositories).
+ 1. When prompted for a description, enter “Apache Beam, version X, release candidate Y”.
### (Alternative) Run all steps manually
@@ -706,9 +711,9 @@
* Build Flink job server images and push to DockerHub.
```
-FLINK_VER=($(ls -1 runners/flink | awk '/^[0-9]+\.[0-9]+$/{print}'))
+FLINK_VER=("1.7" "1.8" "1.9")
for ver in "${FLINK_VER[@]}"; do
- ./gradlew ":runners:flink:${ver}:job-server-container:dockerPush" -Pdocker-tag="${RELEASE}_rc${RC_NUM}"
+ ./gradlew ":runners:flink:${ver}:job-server-container:dockerPush" -Pdocker-tag="${RELEASE}_rc${RC_NUM}"
done
```
@@ -867,7 +872,7 @@
* {$KNOWN_ISSUE_1}
* {$KNOWN_ISSUE_2}
- * See a full list of open [issues that affects](https://issues.apache.org/jira/browse/BEAM-8989?jql=project = BEAM AND affectedVersion = 2.16.0 ORDER BY priority DESC, updated DESC) this version.
+ * See a full list of open [issues that affect](https://issues.apache.org/jira/issues/?jql=project%20%3D%20BEAM%20AND%20affectedVersion%20%3D%20{$RELEASE}%20ORDER%20BY%20priority%20DESC%2C%20updated%20DESC) this version.
## List of Contributors
@@ -926,7 +931,7 @@
* Java artifacts were built with Maven MAVEN_VERSION and OpenJDK/Oracle JDK JDK_VERSION.
* Python artifacts are deployed along with the source release to the dist.apache.org [2].
* Validation sheet with a tab for 1.2.3 release to help with validation [9].
- * Docker images puhlished to Docker Hub [10].
+ * Docker images published to Docker Hub [10].
The vote will be open for at least 72 hours. It is adopted by majority approval, with at least 3 PMC affirmative votes.
@@ -946,7 +951,7 @@
If there are any issues found in the release candidate, reply on the vote thread to cancel the vote. There’s no need to wait 72 hours. Proceed to the `Fix Issues` step below and address the problem. However, some issues don’t require cancellation. For example, if an issue is found in the website pull request, just correct it on the spot and the vote can continue as-is.
-If there are no issues, reply on the vote thread to close the voting. Then, tally the votes in a separate email. Here’s an email template; please adjust as you see fit.
+If there are no issues, reply on the vote thread to close the voting. Then, tally the votes in a separate email thread. Here’s an email template; please adjust as you see fit.
From: Release Manager
To: dev@beam.apache.org
@@ -1237,7 +1242,7 @@
### Deploy artifacts to Maven Central Repository
-Use the Apache Nexus repository to release the staged binary artifacts to the Maven Central repository. In the `Staging Repositories` section, find the relevant release candidate `orgapachebeam-XXX` entry and click `Release`. Drop all other release candidates that are not being released.
+Use the [Apache Nexus repository manager](https://repository.apache.org/#stagingRepositories) to release the staged binary artifacts to the Maven Central repository. In the `Staging Repositories` section, find the relevant release candidate `orgapachebeam-XXX` entry and click `Release`. Drop all other release candidates that are not being released.
__NOTE__: If you are using [GitHub two-factor authentication](https://help.github.com/articles/securing-your-account-with-two-factor-authentication-2fa/) and haven't configure HTTPS access,
please follow [the guide](https://help.github.com/articles/creating-a-personal-access-token-for-the-command-line/) to configure command line access.
@@ -1248,7 +1253,10 @@
delete the `.asc`, `.sha512`;
3. Upload the new release `twine upload *` from the directory with the `.zip` and `.whl` files;
-#### Deploy source release to dist.apache.org
+[Installing twine](https://packaging.python.org/tutorials/packaging-projects/#uploading-the-distribution-archives): `pip install twine`. You can install twine under [virtualenv](https://virtualenv.pypa.io/en/latest/) if preferred.
+
+
+### Deploy source release to dist.apache.org
Copy the source release from the `dev` repository to the `release` repository at `dist.apache.org` using Subversion.
@@ -1259,9 +1267,7 @@
Make sure the download address for last release version is upldaed, [example PR](https://github.com/apache/beam-site/pull/478).
### Deploy SDK docker images to DockerHub
-TODO(hannahjiang): change link to master branch after #9560 is merged.
-
-* Script: [publish_docker_images.sh](https://github.com/Hannah-Jiang/beam/blob/release_script_for_containers/release/src/main/scripts/publish_docker_images.sh)
+* Script: [publish_docker_images.sh](https://github.com/apache/beam/blob/master/release/src/main/scripts/publish_docker_images.sh)
* Usage
```
./beam/release/src/main/scripts/publish_docker_images.sh
@@ -1281,7 +1287,7 @@
### Merge website pull request
-Merge the website pull request to [list the release]({{ site.baseurl }}/get-started/downloads/), publish the [Python API reference manual](https://beam.apache.org/releases/pydoc/), and the [Java API reference manual](https://beam.apache.org/releases/javadoc/) created earlier.
+Merge the website pull request to [list the release]({{ site.baseurl }}/get-started/downloads/), publish the [Python API reference manual](https://beam.apache.org/releases/pydoc/), the [Java API reference manual](https://beam.apache.org/releases/javadoc/) and Blogpost created earlier.
### Mark the version as released in JIRA
diff --git a/website/src/documentation/runners/direct.md b/website/src/documentation/runners/direct.md
index 2e763bf..acfc6bb 100644
--- a/website/src/documentation/runners/direct.md
+++ b/website/src/documentation/runners/direct.md
@@ -86,7 +86,7 @@
Python [FnApiRunner](https://beam.apache.org/contribute/runner-guide/#the-fn-api) supports multi-threading and multi-processing mode.
-#### Setting parallelism
+<strong>Setting parallelism</strong>
Number of threads or subprocesses is defined by setting the `direct_num_workers` option. There are several ways to set this option.
@@ -108,6 +108,23 @@
pipeline_options.view_as(DirectOptions).direct_num_workers = 2
```
+
+
+<strong>Setting running mode</strong>
+
+From 2.19, a new option was added to set running mode. We can use `direct_running_mode` option to set the running mode.
+`direct_running_mode` can be one of [`'in_memory'`, `'multi_threading'`, `'multi_processing'`].
+
+<b>in_memory</b>: Runner and workers' communication happens in memory (not through gRPC). This is a default mode.
+
+<b>multi_threading</b>: Runner and workers communicate through gRPC and each worker runs in a thread.
+
+<b>multi_processing</b>: Runner and workers communicate through gRPC and each worker runs in a subprocess.
+
+Same as other options, `direct_running_mode` can be passed through CLI or set with `PipelineOptions`.
+
+For the versions before 2.19.0, the running mode should be set with `FnApiRunner()`. Please refer following examples.
+
#### Running with multi-threading mode
```
diff --git a/website/src/documentation/sdks/python-dependencies.md b/website/src/documentation/sdks/python-dependencies.md
index 70da2bd..a983bd6 100644
--- a/website/src/documentation/sdks/python-dependencies.md
+++ b/website/src/documentation/sdks/python-dependencies.md
@@ -29,6 +29,46 @@
<p>To see the compile and runtime dependencies for your Beam SDK version, expand
the relevant section below.</p>
+<details><summary markdown="span"><b>2.17.0</b></summary>
+
+<p>Beam SDK for Python 2.17.0 has the following compile and runtime dependencies.</p>
+
+<table class="table-bordered table-striped">
+<tr><th>Package</th><th>Version</th></tr>
+ <tr><td>avro-python3</td><td>>=1.8.1,<2.0.0; python_version >= "3.0"</td></tr>
+ <tr><td>avro</td><td>>=1.8.1,<2.0.0; python_version < "3.0"</td></tr>
+ <tr><td>cachetools</td><td>>=3.1.0,<4</td></tr>
+ <tr><td>crcmod</td><td>>=1.7,<2.0</td></tr>
+ <tr><td>dill</td><td>>=0.3.0,<0.3.1</td></tr>
+ <tr><td>fastavro</td><td>>=0.21.4,<0.22</td></tr>
+ <tr><td>funcsigs</td><td>>=1.0.2,<2; python_version < "3.0"</td></tr>
+ <tr><td>future</td><td>>=0.16.0,<1.0.0</td></tr>
+ <tr><td>futures</td><td>>=3.2.0,<4.0.0; python_version < "3.0"</td></tr>
+ <tr><td>google-apitools</td><td>>=0.5.28,<0.5.29</td></tr>
+ <tr><td>google-cloud-bigquery</td><td>>=1.6.0,<1.18.0</td></tr>
+ <tr><td>google-cloud-bigtable</td><td>>=0.31.1,<1.1.0</td></tr>
+ <tr><td>google-cloud-core</td><td>>=0.28.1,<2</td></tr>
+ <tr><td>google-cloud-datastore</td><td>>=1.7.1,<1.8.0</td></tr>
+ <tr><td>google-cloud-pubsub</td><td>>=0.39.0,<1.1.0</td></tr>
+ <tr><td>googledatastore</td><td>>=7.0.1,<7.1; python_version < "3.0"</td></tr>
+ <tr><td>grpcio</td><td>>=1.12.1,<2</td></tr>
+ <tr><td>hdfs</td><td>>=2.1.0,<3.0.0</td></tr>
+ <tr><td>httplib2</td><td>>=0.8,<=0.12.0</td></tr>
+ <tr><td>mock</td><td>>=1.0.1,<3.0.0</td></tr>
+ <tr><td>oauth2client</td><td>>=2.0.1,<4</td></tr>
+ <tr><td>proto-google-cloud-datastore-v1</td><td>>=0.90.0,<=0.90.4; python_version < "3.0"</td></tr>
+ <tr><td>protobuf</td><td>>=3.5.0.post1,<4</td></tr>
+ <tr><td>pyarrow</td><td>>=0.15.1,<0.16.0; python_version >= "3.0" or platform_system != "Windows"</td></tr>
+ <tr><td>pydot</td><td>>=1.2.0,<2</td></tr>
+ <tr><td>pymongo</td><td>>=3.8.0,<4.0.0</td></tr>
+ <tr><td>python-dateutil</td><td>>=2.8.0,<3</td></tr>
+ <tr><td>pytz</td><td>>=2018.3</td></tr>
+ <tr><td>pyvcf</td><td>>=0.6.8,<0.7.0; python_version < "3.0"</td></tr>
+ <tr><td>typing</td><td>>=3.6.0,<3.7.0; python_version < "3.5.0"</td></tr>
+</table>
+
+</details>
+
<details><summary markdown="span"><b>2.16.0</b></summary>
<p>Beam SDK for Python 2.16.0 has the following compile and
@@ -443,4 +483,3 @@
</table>
</details>
-
diff --git a/website/src/get-started/downloads.md b/website/src/get-started/downloads.md
index d70f7d7..eb5175a 100644
--- a/website/src/get-started/downloads.md
+++ b/website/src/get-started/downloads.md
@@ -90,21 +90,28 @@
## Releases
-## 2.17.0 (2020-01-06)
+### 2.18.0 (2020-01-23)
+Official [source code download](http://www.apache.org/dyn/closer.cgi/beam/2.18.0/apache-beam-2.18.0-source-release.zip).
+[SHA-512](https://www.apache.org/dist/beam/2.18.0/apache-beam-2.18.0-source-release.zip.sha512).
+[signature](https://www.apache.org/dist/beam/2.18.0/apache-beam-2.18.0-source-release.zip.asc).
+
+[Release notes](https://issues.apache.org/jira/secure/ReleaseNote.jspa?version=12346383&projectId=12319527).
+
+### 2.17.0 (2020-01-06)
Official [source code download](http://www.apache.org/dyn/closer.cgi/beam/2.17.0/apache-beam-2.17.0-source-release.zip).
[SHA-512](https://www.apache.org/dist/beam/2.17.0/apache-beam-2.17.0-source-release.zip.sha512).
[signature](https://www.apache.org/dist/beam/2.17.0/apache-beam-2.17.0-source-release.zip.asc).
[Release notes](https://issues.apache.org/jira/secure/ReleaseNote.jspa?projectId=12319527&version=12345970).
-## 2.16.0 (2019-10-07)
+### 2.16.0 (2019-10-07)
Official [source code download](http://www.apache.org/dyn/closer.cgi/beam/2.16.0/apache-beam-2.16.0-source-release.zip).
[SHA-512](https://www.apache.org/dist/beam/2.16.0/apache-beam-2.16.0-source-release.zip.sha512).
[signature](https://www.apache.org/dist/beam/2.16.0/apache-beam-2.16.0-source-release.zip.asc).
[Release notes](https://issues.apache.org/jira/secure/ReleaseNote.jspa?projectId=12319527&version=12345494).
-## 2.15.0 (2019-08-22)
+### 2.15.0 (2019-08-22)
Official [source code download](http://www.apache.org/dyn/closer.cgi/beam/2.15.0/apache-beam-2.15.0-source-release.zip).
[SHA-512](https://www.apache.org/dist/beam/2.15.0/apache-beam-2.15.0-source-release.zip.sha512).
[signature](https://www.apache.org/dist/beam/2.15.0/apache-beam-2.15.0-source-release.zip.asc).
diff --git a/website/src/security/CVE-2020-1929.md b/website/src/security/CVE-2020-1929.md
new file mode 100644
index 0000000..27facc4
--- /dev/null
+++ b/website/src/security/CVE-2020-1929.md
@@ -0,0 +1,17 @@
+---
+permalink: /security/CVE-2020-1929/
+redirect_to: /security/index.html#cve-2020-1929
+---
+<!--
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
diff --git a/website/src/security/index.md b/website/src/security/index.md
new file mode 100644
index 0000000..c8db8e1
--- /dev/null
+++ b/website/src/security/index.md
@@ -0,0 +1,56 @@
+---
+layout: section
+title: "Beam Security"
+permalink: security/
+section_menu: section-menu/get-started.html
+---
+<!--
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+# Reporting Security Issues
+
+Apache Beam uses the standard process outlined by the [Apache Security
+Team](https://www.apache.org/security/) for reporting vulnerabilities. Note
+that vulnerabilities should not be publicly disclosed until the project has
+responded.
+
+To report a possible security vulnerability, please email
+`security@apache.org` and `pmc@beam.apache.org`. This is a non-public list
+that will reach the Beam PMC.
+
+# Known Security Issues
+
+## CVE-2020-1929
+
+[CVE-2020-1929] Apache Beam MongoDB IO connector disables certificate trust verification
+
+Severity: Major
+Vendor: The Apache Software Foundation
+
+Versions Affected:
+Apache Beam 2.10.0 to 2.16.0
+
+Description:
+The Apache Beam MongoDB connector in versions 2.10.0 to 2.16.0 has an option to
+disable SSL trust verification. However this configuration is not respected and
+the certificate verification disables trust verification in every case. This
+exclusion also gets registered globally which disables trust checking for any
+code running in the same JVM.
+
+Mitigation:
+Users of the affected versions should apply one of the following mitigations:
+- Upgrade to Apache Beam 2.17.0 or later
+
+Acknowledgements:
+This issue was reported (and fixed) by Colm Ó hÉigeartaigh.
