Merge pull request #12412 from [BEAM-10603] Add ElementLimiters to all Cache Managers.
[BEAM-10603] Add ElementLimiters to all Cache Managers.
diff --git a/.github/workflows/build_wheels.yml b/.github/workflows/build_wheels.yml
index cd16b1c..319cb44 100644
--- a/.github/workflows/build_wheels.yml
+++ b/.github/workflows/build_wheels.yml
@@ -36,6 +36,21 @@
jobs:
+ check_gcp_variables:
+ timeout-minutes: 5
+ name: "Check are GCP variables set"
+ runs-on: ubuntu-latest
+ outputs:
+ run-tests: ${{ steps.check_gcp_variables.outputs.are-gcp-variables-set }}
+ steps:
+ - uses: actions/checkout@v2
+ - name: "Check are GCP variables set"
+ run: "./scripts/ci/ci_check_are_gcp_variables_set.sh"
+ id: check_gcp_variables
+ env:
+ GCP_SA_EMAIL: ${{ secrets.GCP_SA_EMAIL }}
+ GCP_SA_KEY: ${{ secrets.GCP_SA_KEY }}
+
build_source:
runs-on: ubuntu-latest
name: Build python source distribution
@@ -78,9 +93,11 @@
prepare_gcs:
name: Prepare GCS
- needs: build_source
+ needs:
+ - build_source
+ - check_gcp_variables
runs-on: ubuntu-latest
- if: github.repository_owner == 'apache' && github.event_name != 'pull_request'
+ if: needs.check_gcp_variables.outputs.are-gcp-variables-set == 'true' && github.event_name != 'pull_request'
steps:
- name: Authenticate on GCP
uses: GoogleCloudPlatform/github-actions/setup-gcloud@master
@@ -92,9 +109,11 @@
upload_source_to_gcs:
name: Upload python source distribution to GCS bucket
- needs: prepare_gcs
+ needs:
+ - prepare_gcs
+ - check_gcp_variables
runs-on: ubuntu-latest
- if: github.repository_owner == 'apache'
+ if: needs.check_gcp_variables.outputs.are-gcp-variables-set == 'true'
steps:
- name: Download compressed sources from artifacts
uses: actions/download-artifact@v2
@@ -156,10 +175,12 @@
path: apache-beam-source/wheelhouse/
upload_wheels_to_gcs:
- name: Upload python wheels to GCS bucket
- needs: build_wheels
+ name: Upload wheels to GCS bucket
+ needs:
+ - build_wheels
+ - check_gcp_variables
runs-on: ubuntu-latest
- if: github.repository_owner == 'apache' && github.event_name != 'pull_request'
+ if: needs.check_gcp_variables.outputs.are-gcp-variables-set == 'true' && github.event_name != 'pull_request'
strategy:
matrix:
os : [ubuntu-latest, macos-latest, windows-latest]
@@ -199,9 +220,11 @@
list_files_on_gcs:
name: List files on Google Cloud Storage Bucket
- needs: upload_wheels_to_gcs
+ needs:
+ - upload_wheels_to_gcs
+ - check_gcp_variables
runs-on: ubuntu-latest
- if: github.repository_owner == 'apache' && github.event_name != 'pull_request'
+ if: needs.check_gcp_variables.outputs.are-gcp-variables-set == 'true' && github.event_name != 'pull_request'
steps:
- name: Authenticate on GCP
uses: GoogleCloudPlatform/github-actions/setup-gcloud@master
diff --git a/CHANGES.md b/CHANGES.md
index 04f3094..c3b7959 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -64,11 +64,13 @@
* Add cross-language support to Java's JdbcIO, now available in the Python module `apache_beam.io.external.jdbc` ([BEAM-10135](https://issues.apache.org/jira/browse/BEAM-10135), [BEAM-10136](https://issues.apache.org/jira/browse/BEAM-10136)).
* Add support of AWS SDK v2 for KinesisIO.Read (Java) ([BEAM-9702](https://issues.apache.org/jira/browse/BEAM-9702)).
* Support for X source added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+* Add streaming support to SnowflakeIO in Java SDK ([BEAM-9896](https://issues.apache.org/jira/browse/BEAM-9896 ))
## New Features / Improvements
* Shared library for simplifying management of large shared objects added to Python SDK. Example use case is sharing a large TF model object across threads ([BEAM-10417](https://issues.apache.org/jira/browse/BEAM-10417)).
* X feature added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+* Dataflow streaming timers are not strictly time ordered when set earlier mid-bundle ([BEAM-8543](https://issues.apache.org/jira/browse/BEAM-8543)).
* OnTimerContext should not create a new one when processing each element/timer in FnApiDoFnRunner ([BEAM-9839](https://issues.apache.org/jira/browse/BEAM-9839))
## Breaking Changes
@@ -100,6 +102,7 @@
reading data by exporting to JSON files. This has small differences in behavior for Time and Date-related fields. See
Pydoc for more information.
* Add dispositions for SnowflakeIO.write ([BEAM-10343](https://issues.apache.org/jira/browse/BEAM-10343))
+* Add cross-language support to SnowflakeIO.Read([BEAM-9897](https://issues.apache.org/jira/browse/BEAM-9897)).
## New Features / Improvements
diff --git a/build.gradle b/build.gradle
index 93ca237..7f93fe2 100644
--- a/build.gradle
+++ b/build.gradle
@@ -247,7 +247,6 @@
}
task python2PostCommit() {
- dependsOn ":sdks:python:test-suites:portable:py2:crossLanguagePythonJavaKafkaIOFlink"
dependsOn ":sdks:python:test-suites:portable:py2:crossLanguageTests"
dependsOn ":sdks:python:test-suites:dataflow:py2:postCommitIT"
dependsOn ":sdks:python:test-suites:direct:py2:directRunnerIT"
@@ -275,7 +274,6 @@
}
task python38PostCommit() {
- dependsOn ":sdks:python:test-suites:portable:py38:crossLanguagePythonJavaKafkaIOFlink"
dependsOn ":sdks:python:test-suites:dataflow:py38:postCommitIT"
dependsOn ":sdks:python:test-suites:direct:py38:postCommitIT"
dependsOn ":sdks:python:test-suites:direct:py38:hdfsIntegrationTest"
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 3622764..82541ef 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -391,7 +391,7 @@
def gax_version = "1.54.0"
def generated_grpc_ga_version = "1.85.1"
def google_auth_version = "0.19.0"
- def google_clients_version = "1.30.9"
+ def google_clients_version = "1.30.10"
def google_cloud_bigdataoss_version = "2.1.3"
def google_cloud_core_version = "1.92.2"
def google_cloud_pubsublite_version = "0.1.6"
@@ -469,13 +469,13 @@
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",
google_api_common : "com.google.api:api-common:1.8.1",
- google_api_services_bigquery : "com.google.apis:google-api-services-bigquery:v2-rev20191211-$google_clients_version",
- google_api_services_clouddebugger : "com.google.apis:google-api-services-clouddebugger:v2-rev20200313-$google_clients_version",
- google_api_services_cloudresourcemanager : "com.google.apis:google-api-services-cloudresourcemanager:v1-rev20200311-$google_clients_version",
- google_api_services_dataflow : "com.google.apis:google-api-services-dataflow:v1b3-rev20200305-$google_clients_version",
- google_api_services_healthcare : "com.google.apis:google-api-services-healthcare:v1beta1-rev20200525-$google_clients_version",
- google_api_services_pubsub : "com.google.apis:google-api-services-pubsub:v1-rev20200312-$google_clients_version",
- google_api_services_storage : "com.google.apis:google-api-services-storage:v1-rev20200226-$google_clients_version",
+ google_api_services_bigquery : "com.google.apis:google-api-services-bigquery:v2-rev20200719-$google_clients_version",
+ google_api_services_clouddebugger : "com.google.apis:google-api-services-clouddebugger:v2-rev20200501-$google_clients_version",
+ google_api_services_cloudresourcemanager : "com.google.apis:google-api-services-cloudresourcemanager:v1-rev20200720-$google_clients_version",
+ google_api_services_dataflow : "com.google.apis:google-api-services-dataflow:v1b3-rev20200713-$google_clients_version",
+ google_api_services_healthcare : "com.google.apis:google-api-services-healthcare:v1beta1-rev20200713-$google_clients_version",
+ google_api_services_pubsub : "com.google.apis:google-api-services-pubsub:v1-rev20200713-$google_clients_version",
+ google_api_services_storage : "com.google.apis:google-api-services-storage:v1-rev20200611-$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",
google_cloud_bigquery : "com.google.cloud:google-cloud-bigquery:1.108.0",
diff --git a/examples/java/build.gradle b/examples/java/build.gradle
index 0c91573..eaad35f 100644
--- a/examples/java/build.gradle
+++ b/examples/java/build.gradle
@@ -52,6 +52,7 @@
compile project(path: ":sdks:java:core", configuration: "shadow")
compile project(":sdks:java:extensions:google-cloud-platform-core")
compile project(":sdks:java:io:google-cloud-platform")
+ compile project(":sdks:java:extensions:ml")
compile library.java.avro
compile library.java.bigdataoss_util
compile library.java.google_api_client
@@ -60,6 +61,7 @@
compile library.java.google_auth_library_credentials
compile library.java.google_auth_library_oauth2_http
compile library.java.google_cloud_datastore_v1_proto_client
+ compile library.java.google_code_gson
compile library.java.google_http_client
compile library.java.joda_time
compile library.java.proto_google_cloud_datastore_v1
@@ -67,6 +69,7 @@
compile library.java.slf4j_jdk14
runtime project(path: ":runners:direct-java", configuration: "shadow")
testCompile project(":sdks:java:io:google-cloud-platform")
+ testCompile project(":sdks:java:extensions:ml")
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
testCompile library.java.junit
diff --git a/examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java b/examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java
index 3393e6d..cb55475 100644
--- a/examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java
+++ b/examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java
@@ -22,6 +22,14 @@
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import com.google.api.services.bigquery.model.TimePartitioning;
+import com.google.cloud.language.v1.AnnotateTextRequest;
+import com.google.cloud.language.v1.AnnotateTextResponse;
+import com.google.cloud.language.v1.Document;
+import com.google.cloud.language.v1.Entity;
+import com.google.cloud.language.v1.Sentence;
+import com.google.cloud.language.v1.Token;
+import com.google.gson.Gson;
+import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
@@ -30,12 +38,14 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.stream.Collectors;
import org.apache.avro.generic.GenericRecord;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.DefaultCoder;
import org.apache.beam.sdk.coders.DoubleCoder;
+import org.apache.beam.sdk.extensions.ml.AnnotateText;
import org.apache.beam.sdk.io.Compression;
import org.apache.beam.sdk.io.FileIO;
import org.apache.beam.sdk.io.GenerateSequence;
@@ -63,6 +73,7 @@
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.PeriodicImpulse;
+import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.Sum;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.Watch;
@@ -984,4 +995,167 @@
return result;
}
}
+
+ public static class NaturalLanguageIntegration {
+ private static final SerializableFunction<AnnotateTextResponse, List<Map<String, List<String>>>>
+ // [START NlpAnalyzeDependencyTree]
+ analyzeDependencyTree =
+ (SerializableFunction<AnnotateTextResponse, List<Map<String, List<String>>>>)
+ response -> {
+ List<Map<String, List<String>>> adjacencyLists = new ArrayList<>();
+ int index = 0;
+ for (Sentence s : response.getSentencesList()) {
+ Map<String, List<String>> adjacencyMap = new HashMap<>();
+ int sentenceBegin = s.getText().getBeginOffset();
+ int sentenceEnd = sentenceBegin + s.getText().getContent().length() - 1;
+ while (index < response.getTokensCount()
+ && response.getTokens(index).getText().getBeginOffset() <= sentenceEnd) {
+ Token token = response.getTokensList().get(index);
+ int headTokenIndex = token.getDependencyEdge().getHeadTokenIndex();
+ String headTokenContent =
+ response.getTokens(headTokenIndex).getText().getContent();
+ List<String> adjacencyList =
+ adjacencyMap.getOrDefault(headTokenContent, new ArrayList<>());
+ adjacencyList.add(token.getText().getContent());
+ adjacencyMap.put(headTokenContent, adjacencyList);
+ index++;
+ }
+ adjacencyLists.add(adjacencyMap);
+ }
+ return adjacencyLists;
+ };
+ // [END NlpAnalyzeDependencyTree]
+
+ private static final SerializableFunction<? super AnnotateTextResponse, TextSentiments>
+ // [START NlpExtractSentiments]
+ extractSentiments =
+ (SerializableFunction<AnnotateTextResponse, TextSentiments>)
+ annotateTextResponse -> {
+ TextSentiments sentiments = new TextSentiments();
+ sentiments.setDocumentSentiment(
+ annotateTextResponse.getDocumentSentiment().getMagnitude());
+ Map<String, Float> sentenceSentimentsMap =
+ annotateTextResponse.getSentencesList().stream()
+ .collect(
+ Collectors.toMap(
+ (Sentence s) -> s.getText().getContent(),
+ (Sentence s) -> s.getSentiment().getMagnitude()));
+ sentiments.setSentenceSentiments(sentenceSentimentsMap);
+ return sentiments;
+ };
+ // [END NlpExtractSentiments]
+
+ private static final SerializableFunction<? super AnnotateTextResponse, Map<String, String>>
+ // [START NlpExtractEntities]
+ extractEntities =
+ (SerializableFunction<AnnotateTextResponse, Map<String, String>>)
+ annotateTextResponse ->
+ annotateTextResponse.getEntitiesList().stream()
+ .collect(
+ Collectors.toMap(Entity::getName, (Entity e) -> e.getType().toString()));
+ // [END NlpExtractEntities]
+
+ private static final SerializableFunction<? super Map<String, String>, String>
+ mapEntitiesToJson =
+ (SerializableFunction<Map<String, String>, String>)
+ item -> {
+ StringBuilder builder = new StringBuilder("[");
+ builder.append(
+ item.entrySet().stream()
+ .map(
+ entry -> "{\"" + entry.getKey() + "\": \"" + entry.getValue() + "\"}")
+ .collect(Collectors.joining(",")));
+ builder.append("]");
+ return builder.toString();
+ };
+
+ private static final SerializableFunction<List<Map<String, List<String>>>, String>
+ mapDependencyTreesToJson =
+ (SerializableFunction<List<Map<String, List<String>>>, String>)
+ tree -> {
+ Gson gson = new Gson();
+ return gson.toJson(tree);
+ };
+
+ public static void main(Pipeline p) {
+ // [START NlpAnalyzeText]
+ AnnotateTextRequest.Features features =
+ AnnotateTextRequest.Features.newBuilder()
+ .setExtractEntities(true)
+ .setExtractDocumentSentiment(true)
+ .setExtractEntitySentiment(true)
+ .setExtractSyntax(true)
+ .build();
+ AnnotateText annotateText = AnnotateText.newBuilder().setFeatures(features).build();
+
+ PCollection<AnnotateTextResponse> responses =
+ p.apply(
+ Create.of(
+ "My experience so far has been fantastic, "
+ + "I\'d really recommend this product."))
+ .apply(
+ MapElements.into(TypeDescriptor.of(Document.class))
+ .via(
+ (SerializableFunction<String, Document>)
+ input ->
+ Document.newBuilder()
+ .setContent(input)
+ .setType(Document.Type.PLAIN_TEXT)
+ .build()))
+ .apply(annotateText);
+
+ responses
+ .apply(MapElements.into(TypeDescriptor.of(TextSentiments.class)).via(extractSentiments))
+ .apply(
+ MapElements.into(TypeDescriptors.strings())
+ .via((SerializableFunction<TextSentiments, String>) TextSentiments::toJson))
+ .apply(TextIO.write().to("sentiments.txt"));
+
+ responses
+ .apply(
+ MapElements.into(
+ TypeDescriptors.maps(TypeDescriptors.strings(), TypeDescriptors.strings()))
+ .via(extractEntities))
+ .apply(MapElements.into(TypeDescriptors.strings()).via(mapEntitiesToJson))
+ .apply(TextIO.write().to("entities.txt"));
+
+ responses
+ .apply(
+ MapElements.into(
+ TypeDescriptors.lists(
+ TypeDescriptors.maps(
+ TypeDescriptors.strings(),
+ TypeDescriptors.lists(TypeDescriptors.strings()))))
+ .via(analyzeDependencyTree))
+ .apply(MapElements.into(TypeDescriptors.strings()).via(mapDependencyTreesToJson))
+ .apply(TextIO.write().to("adjacency_list.txt"));
+ // [END NlpAnalyzeText]
+ }
+
+ private static class TextSentiments implements Serializable {
+ private Float documentSentiment;
+ private Map<String, Float> sentenceSentiments;
+
+ public void setSentenceSentiments(Map<String, Float> sentenceSentiments) {
+ this.sentenceSentiments = sentenceSentiments;
+ }
+
+ public Float getDocumentSentiment() {
+ return documentSentiment;
+ }
+
+ public void setDocumentSentiment(Float documentSentiment) {
+ this.documentSentiment = documentSentiment;
+ }
+
+ public Map<String, Float> getSentenceSentiments() {
+ return sentenceSentiments;
+ }
+
+ public String toJson() {
+ Gson gson = new Gson();
+ return gson.toJson(this);
+ }
+ }
+ }
}
diff --git a/learning/katas/go/core_transforms/section-info.yaml b/learning/katas/go/core_transforms/section-info.yaml
index 32ae3aa..bcdab5d 100644
--- a/learning/katas/go/core_transforms/section-info.yaml
+++ b/learning/katas/go/core_transforms/section-info.yaml
@@ -28,3 +28,4 @@
- additional_outputs
- branching
- composite
+- windowing
diff --git a/learning/katas/go/core_transforms/windowing/lesson-info.yaml b/learning/katas/go/core_transforms/windowing/lesson-info.yaml
new file mode 100644
index 0000000..2315b4c
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/lesson-info.yaml
@@ -0,0 +1,21 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+content:
+- windowing
diff --git a/learning/katas/go/core_transforms/windowing/lesson-remote-info.yaml b/learning/katas/go/core_transforms/windowing/lesson-remote-info.yaml
new file mode 100644
index 0000000..7a20a3a
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/lesson-remote-info.yaml
@@ -0,0 +1,3 @@
+id: 387853
+update_date: Thu, 06 Aug 2020 17:53:20 UTC
+unit: 377026
diff --git a/learning/katas/go/core_transforms/windowing/windowing/cmd/main.go b/learning/katas/go/core_transforms/windowing/windowing/cmd/main.go
new file mode 100644
index 0000000..e1cfcf0
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/windowing/cmd/main.go
@@ -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 main
+
+import (
+ "beam.apache.org/learning/katas/core_transforms/windowing/windowing/pkg/common"
+ "beam.apache.org/learning/katas/core_transforms/windowing/windowing/pkg/task"
+ "context"
+ "github.com/apache/beam/sdks/go/pkg/beam"
+ "github.com/apache/beam/sdks/go/pkg/beam/log"
+ "github.com/apache/beam/sdks/go/pkg/beam/x/beamx"
+ "github.com/apache/beam/sdks/go/pkg/beam/x/debug"
+)
+
+func main() {
+ ctx := context.Background()
+
+ p, s := beam.NewPipelineWithRoot()
+
+ input := common.CreateLines(s)
+
+ result := task.ApplyTransform(s, input)
+
+ output := beam.ParDo(s, func(commit task.Commit) string {
+ return commit.String()
+ }, result)
+
+ debug.Print(s, output)
+
+ err := beamx.Run(ctx, p)
+
+ if err != nil {
+ log.Exitf(context.Background(), "Failed to execute job: %v", err)
+ }
+}
diff --git a/learning/katas/go/core_transforms/windowing/windowing/pkg/common/input.go b/learning/katas/go/core_transforms/windowing/windowing/pkg/common/input.go
new file mode 100644
index 0000000..3353e2b
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/windowing/pkg/common/input.go
@@ -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 common
+
+import (
+ "github.com/apache/beam/sdks/go/pkg/beam"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/graph/mtime"
+ "time"
+)
+
+var (
+ lines = map[time.Time]string{
+ time.Date(2020, 7, 31, 15, 52, 5, 0, time.UTC): "3c6c45924a Remove trailing whitespace from README",
+ time.Date(2020, 7, 31, 15, 59, 40, 0, time.UTC): "a52be99b62 Merge pull request #12443 from KevinGG/whitespace",
+ time.Date(2020, 7, 31, 16, 7, 36, 0, time.UTC): "7c1772d13f Merge pull request #12439 from ibzib/beam-9199-1",
+ time.Date(2020, 7, 31, 16, 35, 41, 0, time.UTC): "d971ba13b8 Widen ranges for GCP libraries (#12198)",
+ time.Date(2020, 8, 1, 0, 7, 25, 0, time.UTC): "875620111b Enable all Jenkins jobs triggering for committers (#12407)",
+ }
+)
+
+func CreateLines(s beam.Scope) beam.PCollection {
+ return beam.ParDo(s, timestampFn, beam.Impulse(s))
+}
+
+func timestampFn(_ []byte, emit func(beam.EventTime, string)) {
+ for timestamp, line := range lines {
+ emit(mtime.FromTime(timestamp), line)
+ }
+}
diff --git a/learning/katas/go/core_transforms/windowing/windowing/pkg/task/task.go b/learning/katas/go/core_transforms/windowing/windowing/pkg/task/task.go
new file mode 100644
index 0000000..1838578
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/windowing/pkg/task/task.go
@@ -0,0 +1,54 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 task
+
+import (
+ "fmt"
+ "github.com/apache/beam/sdks/go/pkg/beam"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/graph/window"
+ "time"
+)
+
+func ApplyTransform(s beam.Scope, input beam.PCollection) beam.PCollection {
+ windowed := beam.WindowInto(s, window.NewFixedWindows(time.Hour), input)
+ return beam.ParDo(s, timestampFn, windowed)
+}
+
+func timestampFn(iw beam.Window, et beam.EventTime, line string) Commit {
+ return Commit{
+ MaxTimestampWindow: toTime(iw.MaxTimestamp()),
+ EventTimestamp: toTime(et),
+ Line: line,
+ }
+}
+
+func toTime(et beam.EventTime) time.Time {
+ return time.Unix(0, et.Milliseconds() * int64(time.Millisecond))
+}
+
+type Commit struct {
+ MaxTimestampWindow time.Time
+ EventTimestamp time.Time
+ Line string
+}
+
+func (c Commit) String() string {
+ return fmt.Sprintf("Window ending at: %v contains timestamp: %v for commit: \"%s\"",
+ c.MaxTimestampWindow.Format(time.Kitchen),
+ c.EventTimestamp.Format(time.Kitchen),
+ c.Line)
+}
+
diff --git a/learning/katas/go/core_transforms/windowing/windowing/task-info.yaml b/learning/katas/go/core_transforms/windowing/windowing/task-info.yaml
new file mode 100644
index 0000000..b0e38f7
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/windowing/task-info.yaml
@@ -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.
+#
+
+type: edu
+files:
+- name: test/task_test.go
+ visible: false
+- name: cmd/main.go
+ visible: true
+- name: pkg/common/input.go
+ visible: true
+- name: pkg/task/task.go
+ visible: true
+ placeholders:
+ - offset: 1030
+ length: 60
+ placeholder_text: TODO()
+ - offset: 1099
+ length: 36
+ placeholder_text: TODO()
+ - offset: 1156
+ length: 46
+ placeholder_text: TODO()
+ - offset: 1221
+ length: 121
+ placeholder_text: TODO()
diff --git a/learning/katas/go/core_transforms/windowing/windowing/task-remote-info.yaml b/learning/katas/go/core_transforms/windowing/windowing/task-remote-info.yaml
new file mode 100644
index 0000000..235e663
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/windowing/task-remote-info.yaml
@@ -0,0 +1,2 @@
+id: 1464828
+update_date: Thu, 06 Aug 2020 17:53:26 UTC
diff --git a/learning/katas/go/core_transforms/windowing/windowing/task.md b/learning/katas/go/core_transforms/windowing/windowing/task.md
new file mode 100644
index 0000000..22444a6
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/windowing/task.md
@@ -0,0 +1,77 @@
+<!--
+ ~ Licensed to the Apache Software Foundation (ASF) under one
+ ~ or more contributor license agreements. See the NOTICE file
+ ~ distributed with this work for additional information
+ ~ regarding copyright ownership. The ASF licenses this file
+ ~ to you under the Apache License, Version 2.0 (the
+ ~ "License"); you may not use this file except in compliance
+ ~ with the License. You may obtain a copy of the License at
+ ~
+ ~ http://www.apache.org/licenses/LICENSE-2.0
+ ~
+ ~ Unless required by applicable law or agreed to in writing, software
+ ~ distributed under the License is distributed on an "AS IS" BASIS,
+ ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ ~ See the License for the specific language governing permissions and
+ ~ limitations under the License.
+ -->
+
+# Windowing
+
+This lesson introduces the concept of windowed PCollection elements. A window is a view into a fixed beginning and
+fixed end to a set of data. In the beam model, windowing subdivides a PCollection according to the
+timestamps of its individual elements. An element can be a part of one or more windows.
+
+A DoFn can request timestamp and windowing information about the element it is processing. All the previous lessons
+had this information available as well. This lesson makes use of these parameters. The simple dataset
+has five git commit messages and their timestamps from the
+[Apache Beam public repository](https://github.com/apache/beam). Timestamps have been applied to this PCollection
+input according to the date and time of these messages.
+
+**Kata:** This lesson challenges you to apply an hourly fixed window to a PCollection. You are then to
+apply a ParDo to that hourly fixed windowed PCollection to produce a PCollection of a Commit struct. The
+Commit struct is provided for you. You are encouraged to run the pipeline at cmd/main.go of this task
+to visualize the windowing and timestamps.
+
+<div class="hint">
+ Use <a href="https://godoc.org/github.com/apache/beam/sdks/go/pkg/beam#ParDo">
+ beam.ParDo</a>
+ with a <a href="https://godoc.org/github.com/apache/beam/sdks/go/pkg/beam#hdr-DoFns">
+ DoFn</a> to accomplish this lesson.
+</div>
+
+<div class="hint">
+ Use <a href="https://godoc.org/github.com/apache/beam/sdks/go/pkg/beam#WindowInto">
+ beam.WindowInto</a>
+ with <a href="https://godoc.org/github.com/apache/beam/sdks/go/pkg/beam/core/graph/window#NewFixedWindows">
+ window.NewFixedWindows(time.Hour)</a>
+ on your PCollection input to apply an hourly windowing strategy to each element.
+</div>
+
+<div class="hint">
+ To access <a href="https://godoc.org/github.com/apache/beam/sdks/go/pkg/beam#Window">
+ beam.Window</a>
+ and <a href="https://godoc.org/github.com/apache/beam/sdks/go/pkg/beam#EventTime">
+ beam.EventTime</a> in your DoFn, add the parameters in the set order.
+
+```
+func doFn(iw beam.Window, et beam.EventTime, element X) Y {
+ // do something with iw, et and element to return Y
+}
+```
+</div>
+
+<div class="hint">
+ The Commit struct provided for you has a MaxTimestampWindow property that can be set from
+ <a href="https://godoc.org/github.com/apache/beam/sdks/go/pkg/beam#Window">
+ beam.Window</a>'s MaxTimestamp().
+</div>
+
+<div class="hint">
+ Refer to the Beam Programming Guide for additional information about
+ <a href="https://beam.apache.org/documentation/programming-guide/#other-dofn-parameters">
+ additional DoFn parameters</a> and
+ <a href="https://beam.apache.org/documentation/programming-guide/#windowing">
+ windowing</a>.
+</div>
+
diff --git a/learning/katas/go/core_transforms/windowing/windowing/test/task_test.go b/learning/katas/go/core_transforms/windowing/windowing/test/task_test.go
new file mode 100644
index 0000000..8d4a2e2
--- /dev/null
+++ b/learning/katas/go/core_transforms/windowing/windowing/test/task_test.go
@@ -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.
+
+package test
+
+import (
+ "beam.apache.org/learning/katas/core_transforms/windowing/windowing/pkg/common"
+ "beam.apache.org/learning/katas/core_transforms/windowing/windowing/pkg/task"
+ "github.com/apache/beam/sdks/go/pkg/beam"
+ "github.com/apache/beam/sdks/go/pkg/beam/testing/ptest"
+ "github.com/google/go-cmp/cmp"
+ "testing"
+ "time"
+)
+
+func TestApplyTransform(t *testing.T) {
+ p, s := beam.NewPipelineWithRoot()
+ tests := []struct {
+ input beam.PCollection
+ want []interface{}
+ }{
+ {
+ input: common.CreateLines(s),
+ want: []interface{}{
+ task.Commit{
+ MaxTimestampWindow: time.Unix(1596211199, 0),
+ EventTimestamp: time.Unix(1596210725, 0),
+ Line: "3c6c45924a Remove trailing whitespace from README",
+ },
+ task.Commit{
+ MaxTimestampWindow: time.Unix(1596211199, 0),
+ EventTimestamp: time.Unix(1596211180, 0),
+ Line: "a52be99b62 Merge pull request #12443 from KevinGG/whitespace",
+ },
+ task.Commit{
+ MaxTimestampWindow: time.Unix(1596214799, 0),
+ EventTimestamp: time.Unix(1596211656, 0),
+ Line: "7c1772d13f Merge pull request #12439 from ibzib/beam-9199-1",
+ },
+ task.Commit{
+ MaxTimestampWindow: time.Unix(1596214799, 0),
+ EventTimestamp: time.Unix(1596213341, 0),
+ Line: "d971ba13b8 Widen ranges for GCP libraries (#12198)",
+ },
+ task.Commit{
+ MaxTimestampWindow: time.Unix(1596243599, 0),
+ EventTimestamp: time.Unix(1596240445, 0),
+ Line: "875620111b Enable all Jenkins jobs triggering for committers (#12407)",
+ },
+ },
+ },
+ }
+ for _, tt := range tests {
+ got := task.ApplyTransform(s, tt.input)
+ cmp.Equal(got, tt.want)
+ if err := ptest.Run(p); err != nil {
+ t.Error(err)
+ }
+ }
+}
diff --git a/learning/katas/go/course-remote-info.yaml b/learning/katas/go/course-remote-info.yaml
index 90e7821..e944389 100644
--- a/learning/katas/go/course-remote-info.yaml
+++ b/learning/katas/go/course-remote-info.yaml
@@ -1,2 +1,2 @@
id: 70387
-update_date: Mon, 27 Jul 2020 20:44:48 UTC
+update_date: Wed, 29 Jul 2020 20:42:36 UTC
diff --git a/model/fn-execution/src/main/resources/org/apache/beam/model/fnexecution/v1/standard_coders.yaml b/model/fn-execution/src/main/resources/org/apache/beam/model/fnexecution/v1/standard_coders.yaml
index 80e4eb8..08b823d 100644
--- a/model/fn-execution/src/main/resources/org/apache/beam/model/fnexecution/v1/standard_coders.yaml
+++ b/model/fn-execution/src/main/resources/org/apache/beam/model/fnexecution/v1/standard_coders.yaml
@@ -347,7 +347,6 @@
urn: "beam:coder:row:v1"
# str: string, i32: int32, f64: float64, arr: array[string]
payload: "\n\t\n\x03str\x1a\x02\x10\x07\n\t\n\x03i32\x1a\x02\x10\x03\n\t\n\x03f64\x1a\x02\x10\x06\n\r\n\x03arr\x1a\x06\x1a\x04\n\x02\x10\x07\x12$4e5e554c-d4c1-4a5d-b5e1-f3293a6b9f05"
-nested: false
examples:
"\u0004\u0000\u0003foo\u00a9\u0046\u003f\u00b9\u0099\u0099\u0099\u0099\u0099\u009a\0\0\0\u0003\u0003foo\u0003bar\u0003baz": {str: "foo", i32: 9001, f64: "0.1", arr: ["foo", "bar", "baz"]}
@@ -357,7 +356,6 @@
urn: "beam:coder:row:v1"
# str: nullable string, i32: nullable int32, f64: nullable float64
payload: "\n\x0b\n\x03str\x1a\x04\x08\x01\x10\x07\n\x0b\n\x03i32\x1a\x04\x08\x01\x10\x03\n\x0b\n\x03f64\x1a\x04\x08\x01\x10\x06\x12$b20c6545-57af-4bc8-b2a9-51ace21c7393"
-nested: false
examples:
"\u0003\u0001\u0007": {str: null, i32: null, f64: null}
"\u0003\u0001\u0004\u0003foo\u00a9\u0046": {str: "foo", i32: 9001, f64: null}
@@ -380,7 +378,33 @@
urn: "beam:coder:row:v1"
# f_bool: boolean, f_bytes: nullable bytes
payload: "\n\x0c\n\x06f_bool\x1a\x02\x10\x08\n\x0f\n\x07f_bytes\x1a\x04\x08\x01\x10\t\x12$eea1b747-7571-43d3-aafa-9255afdceafb"
-nested: false
examples:
"\x02\x01\x02\x01": {f_bool: True, f_bytes: null}
"\x02\x00\x00\x04ab\x00c": {f_bool: False, f_bytes: "ab\0c"}
+
+---
+
+# Binary data generated with the python SDK:
+#
+# import typing
+# import apache_beam as beam
+# class Test(typing.NamedTuple):
+# f_map: typing.Mapping[str,int]
+# schema = beam.typehints.schemas.named_tuple_to_schema(Test)
+# coder = beam.coders.row_coder.RowCoder(schema)
+# print("payload = %s" % schema.SerializeToString())
+# examples = (Test(f_map={}),
+# Test(f_map={"foo": 9001, "bar": 9223372036854775807}),
+# Test(f_map={"everything": None, "is": None, "null!": None, "¯\_(ツ)_/¯": None}))
+# for example in examples:
+# print("example = %s" % coder.encode(example))
+coder:
+ urn: "beam:coder:row:v1"
+ # f_map: map<str, nullable int64>
+ payload: "\n\x15\n\x05f_map\x1a\x0c*\n\n\x02\x10\x07\x12\x04\x08\x01\x10\x04\x12$d8c8f969-14e6-457f-a8b5-62a1aec7f1cd"
+ # map ordering is non-deterministic
+ non_deterministic: True
+examples:
+ "\x01\x00\x00\x00\x00\x00": {f_map: {}}
+ "\x01\x00\x00\x00\x00\x02\x03foo\x01\xa9F\x03bar\x01\xff\xff\xff\xff\xff\xff\xff\xff\x7f": {f_map: {"foo": 9001, "bar": 9223372036854775807}}
+ "\x01\x00\x00\x00\x00\x04\neverything\x00\x02is\x00\x05null!\x00\r\xc2\xaf\\_(\xe3\x83\x84)_/\xc2\xaf\x00": {f_map: {"everything":null, "is": null, "null!": null, "¯\\_(ツ)_/¯": null}}
diff --git a/model/pipeline/src/main/proto/beam_runner_api.proto b/model/pipeline/src/main/proto/beam_runner_api.proto
index 3623790..8b1ce6b 100644
--- a/model/pipeline/src/main/proto/beam_runner_api.proto
+++ b/model/pipeline/src/main/proto/beam_runner_api.proto
@@ -855,10 +855,21 @@
// BOOLEAN: beam:coder:bool:v1
// BYTES: beam:coder:bytes:v1
// ArrayType: beam:coder:iterable:v1 (always has a known length)
- // MapType: not yet a standard coder (BEAM-7996)
+ // MapType: not a standard coder, specification defined below.
// RowType: beam:coder:row:v1
// LogicalType: Uses the coder for its representation.
//
+ // The MapType is encoded by:
+ // - An INT32 representing the size of the map (N)
+ // - Followed by N interleaved keys and values, encoded with their
+ // corresponding coder.
+ //
+ // Nullable types in container types (ArrayType, MapType) are encoded by:
+ // - A one byte null indicator, 0x00 for null values, or 0x01 for present
+ // values.
+ // - For present values the null indicator is followed by the value
+ // encoded with it's corresponding coder.
+ //
// The payload for RowCoder is an instance of Schema.
// Components: None
// Experimental.
diff --git a/model/pipeline/src/main/proto/schema.proto b/model/pipeline/src/main/proto/schema.proto
index dcf75ca..bffa5f1 100644
--- a/model/pipeline/src/main/proto/schema.proto
+++ b/model/pipeline/src/main/proto/schema.proto
@@ -19,6 +19,9 @@
// ** Experimental **
// Protocol Buffers describing Beam Schemas, a portable representation for
// complex types.
+//
+// The primary application of Schema is as the payload for the standard coder
+// "beam:coder:row:v1", defined in beam_runner_api.proto
syntax = "proto3";
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 10221b8..2c4090e 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
@@ -21,7 +21,6 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects.firstNonNull;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList;
-import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap.toImmutableMap;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasItem;
import static org.hamcrest.Matchers.instanceOf;
@@ -43,6 +42,7 @@
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
+import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
@@ -340,6 +340,10 @@
}
private static Object parseField(Object value, Schema.FieldType fieldType) {
+ if (value == null) {
+ return null;
+ }
+
switch (fieldType.getTypeName()) {
case BYTE:
return ((Number) value).byteValue();
@@ -366,14 +370,18 @@
.map((element) -> parseField(element, fieldType.getCollectionElementType()))
.collect(toImmutableList());
case MAP:
- Map<Object, Object> kvMap = (Map<Object, Object>) value;
- return kvMap.entrySet().stream()
- .collect(
- toImmutableMap(
- (pair) -> parseField(pair.getKey(), fieldType.getMapKeyType()),
- (pair) -> parseField(pair.getValue(), fieldType.getMapValueType())));
+ Map<Object, Object> kvMap = new HashMap<>();
+ ((Map<Object, Object>) value)
+ .entrySet().stream()
+ .forEach(
+ (entry) ->
+ kvMap.put(
+ parseField(entry.getKey(), fieldType.getMapKeyType()),
+ parseField(entry.getValue(), fieldType.getMapValueType())));
+ return kvMap;
case ROW:
- Map<String, Object> rowMap = (Map<String, Object>) value;
+ // Clone map so we don't mutate the underlying value
+ Map<String, Object> rowMap = new HashMap<>((Map<String, Object>) value);
Schema schema = fieldType.getRowSchema();
Row.Builder row = Row.withSchema(schema);
for (Schema.Field field : schema.getFields()) {
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ValidateRunnerXlangTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ValidateRunnerXlangTest.java
index 3a0f60b..b96c2c3 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ValidateRunnerXlangTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ValidateRunnerXlangTest.java
@@ -54,7 +54,30 @@
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
-/** Test External transforms. */
+/**
+ * Runner Validation Test Suite for Cross-language Transforms.
+ *
+ * <p>As per Beams's Portability Framework design, Cross-language transforms should work out of the
+ * box. In spite of this, there always exists a possibility of rough edges existing. It could be
+ * caused due to unpolished implementation of any part of the execution code path, for example: –>
+ * Transform expansion [SDK] –> Pipeline construction [SDK] –> Cross-language artifact staging
+ * [Runner] –> Language specific serialization/deserialization of PCollection (and other data types)
+ * [Runner/SDK]
+ *
+ * <p>In an effort to improve developer visibility into potential problems, this test suite
+ * validates correct execution of 5 Core Beam transforms when used as cross-language transforms
+ * within the Java SDK from any foreign SDK: –> ParDo
+ * (https://beam.apache.org/documentation/programming-guide/#pardo) –> GroupByKey
+ * (https://beam.apache.org/documentation/programming-guide/#groupbykey) –> CoGroupByKey
+ * (https://beam.apache.org/documentation/programming-guide/#cogroupbykey) –> Combine
+ * (https://beam.apache.org/documentation/programming-guide/#combine) –> Flatten
+ * (https://beam.apache.org/documentation/programming-guide/#flatten) –> Partition
+ * (https://beam.apache.org/documentation/programming-guide/#partition)
+ *
+ * <p>See Runner Validation Test Plan for Cross-language transforms
+ * (https://docs.google.com/document/d/1xQp0ElIV84b8OCVz8CD2hvbiWdR8w4BvWxPTZJZA6NA") for further
+ * details.
+ */
@RunWith(JUnit4.class)
public class ValidateRunnerXlangTest implements Serializable {
@Rule public transient TestPipeline testPipeline = TestPipeline.create();
@@ -110,6 +133,14 @@
}
}
+ /**
+ * Motivation behind singleInputOutputTest.
+ *
+ * <p>Target transform – ParDo (https://beam.apache.org/documentation/programming-guide/#pardo)
+ * Test scenario – Mapping elements from a single input collection to a single output collection
+ * Boundary conditions checked – –> PCollection<?> to external transforms –> PCollection<?> from
+ * external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void singleInputOutputTest() throws IOException {
@@ -120,6 +151,14 @@
PAssert.that(col).containsInAnyOrder("01", "02", "03");
}
+ /**
+ * Motivation behind multiInputOutputWithSideInputTest.
+ *
+ * <p>Target transform – ParDo (https://beam.apache.org/documentation/programming-guide/#pardo)
+ * Test scenario – Mapping elements from multiple input collections (main and side) to multiple
+ * output collections (main and side) Boundary conditions checked – –> PCollectionTuple to
+ * external transforms –> PCollectionTuple from external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void multiInputOutputWithSideInputTest() {
@@ -135,6 +174,15 @@
PAssert.that(pTuple.get("side")).containsInAnyOrder("ss");
}
+ /**
+ * Motivation behind groupByKeyTest.
+ *
+ * <p>Target transform – GroupByKey
+ * (https://beam.apache.org/documentation/programming-guide/#groupbykey) Test scenario – Grouping
+ * a collection of KV<K,V> to a collection of KV<K, Iterable<V>> by key Boundary conditions
+ * checked – –> PCollection<KV<?, ?>> to external transforms –> PCollection<KV<?, Iterable<?>>>
+ * from external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void groupByKeyTest() {
@@ -154,6 +202,15 @@
PAssert.that(col).containsInAnyOrder("0:1,2", "1:3");
}
+ /**
+ * Motivation behind coGroupByKeyTest.
+ *
+ * <p>Target transform – CoGroupByKey
+ * (https://beam.apache.org/documentation/programming-guide/#cogroupbykey) Test scenario –
+ * Grouping multiple input collections with keys to a collection of KV<K, CoGbkResult> by key
+ * Boundary conditions checked – –> KeyedPCollectionTuple<?> to external transforms –>
+ * PCollection<KV<?, Iterable<?>>> from external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void coGroupByKeyTest() {
@@ -177,6 +234,14 @@
PAssert.that(col).containsInAnyOrder("0:1,2,4", "1:3,5,6");
}
+ /**
+ * Motivation behind combineGloballyTest.
+ *
+ * <p>Target transform – Combine
+ * (https://beam.apache.org/documentation/programming-guide/#combine) Test scenario – Combining
+ * elements globally with a predefined simple CombineFn Boundary conditions checked – –>
+ * PCollection<?> to external transforms –> PCollection<?> from external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void combineGloballyTest() {
@@ -187,6 +252,14 @@
PAssert.that(col).containsInAnyOrder(6L);
}
+ /**
+ * Motivation behind combinePerKeyTest.
+ *
+ * <p>Target transform – Combine
+ * (https://beam.apache.org/documentation/programming-guide/#combine) Test scenario – Combining
+ * elements per key with a predefined simple merging function Boundary conditions checked – –>
+ * PCollection<?> to external transforms –> PCollection<?> from external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void combinePerKeyTest() {
@@ -197,6 +270,14 @@
PAssert.that(col).containsInAnyOrder(KV.of("a", 3L), KV.of("b", 3L));
}
+ /**
+ * Motivation behind flattenTest.
+ *
+ * <p>Target transform – Flatten
+ * (https://beam.apache.org/documentation/programming-guide/#flatten) Test scenario – Merging
+ * multiple collections into a single collection Boundary conditions checked – –>
+ * PCollectionList<?> to external transforms –> PCollection<?> from external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void flattenTest() {
@@ -209,6 +290,15 @@
PAssert.that(col).containsInAnyOrder(1L, 2L, 3L, 4L, 5L, 6L);
}
+ /**
+ * Motivation behind partitionTest.
+ *
+ * <p>Target transform – Partition
+ * (https://beam.apache.org/documentation/programming-guide/#partition) Test scenario – Splitting
+ * a single collection into multiple collections with a predefined simple PartitionFn Boundary
+ * conditions checked – –> PCollection<?> to external transforms –> PCollectionList<?> from
+ * external transforms
+ */
@Test
@Category({ValidatesRunner.class, UsesCrossLanguageTransforms.class})
public void partitionTest() {
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java
index 988921e..f7c74c0 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java
@@ -337,26 +337,30 @@
*/
@Override
public State waitUntilFinish(Duration duration) {
- State startState = this.state;
- if (!startState.isTerminal()) {
- try {
- state = executor.waitUntilFinish(duration);
- } catch (UserCodeException uce) {
- // Emulates the behavior of Pipeline#run(), where a stack trace caused by a
- // UserCodeException is truncated and replaced with the stack starting at the call to
- // waitToFinish
- throw new Pipeline.PipelineExecutionException(uce.getCause());
- } catch (Exception e) {
- if (e instanceof InterruptedException) {
- Thread.currentThread().interrupt();
- }
- if (e instanceof RuntimeException) {
- throw (RuntimeException) e;
- }
- throw new RuntimeException(e);
- }
+ if (this.state.isTerminal()) {
+ return this.state;
}
- return this.state;
+ final State endState;
+ try {
+ endState = executor.waitUntilFinish(duration);
+ } catch (UserCodeException uce) {
+ // Emulates the behavior of Pipeline#run(), where a stack trace caused by a
+ // UserCodeException is truncated and replaced with the stack starting at the call to
+ // waitToFinish
+ throw new Pipeline.PipelineExecutionException(uce.getCause());
+ } catch (Exception e) {
+ if (e instanceof InterruptedException) {
+ Thread.currentThread().interrupt();
+ }
+ if (e instanceof RuntimeException) {
+ throw (RuntimeException) e;
+ }
+ throw new RuntimeException(e);
+ }
+ if (endState != null) {
+ this.state = endState;
+ }
+ return endState;
}
}
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java
index 8054a07..fbcf0c0 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java
@@ -328,8 +328,9 @@
// The pipeline should never complete;
assertThat(result.getState(), is(State.RUNNING));
// Must time out, otherwise this test will never complete
- result.waitUntilFinish(Duration.millis(1L));
- assertEquals(null, result.getState());
+ assertEquals(null, result.waitUntilFinish(Duration.millis(1L)));
+ // Ensure multiple calls complete
+ assertEquals(null, result.waitUntilFinish(Duration.millis(1L)));
}
private static final AtomicLong TEARDOWN_CALL = new AtomicLong(-1);
diff --git a/scripts/ci/ci_check_are_gcp_variables_set.sh b/scripts/ci/ci_check_are_gcp_variables_set.sh
new file mode 100755
index 0000000..28631b1
--- /dev/null
+++ b/scripts/ci/ci_check_are_gcp_variables_set.sh
@@ -0,0 +1,43 @@
+#!/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.
+
+set -e
+
+echo "This script checks of presence of variables required to perform operations on Google Cloud Platform. They should be stored as secrets."
+echo "More detailed information about Google Cloud Platform Credentials can be found in CI.md"
+
+function check_vars() {
+ ret=true
+ for var in "$@"; do
+ if [ -n "${!var}" ]; then
+ echo "$var is set"
+ else
+ echo >&2 "$var is not set"
+ ret=false
+ fi
+ done
+ $ret
+}
+
+if ! check_vars "GCP_SA_EMAIL" "GCP_SA_KEY"; then
+ echo "::set-output name=are-gcp-variables-set::false"
+ echo >&2 "!!! WARNING !!!"
+ echo >&2 "Not all GCP variables are set. Jobs which require them will be skipped."
+else
+ echo "::set-output name=are-gcp-variables-set::true"
+fi
diff --git a/sdks/go/pkg/beam/core/graph/coder/map.go b/sdks/go/pkg/beam/core/graph/coder/map.go
new file mode 100644
index 0000000..4e5dc2c
--- /dev/null
+++ b/sdks/go/pkg/beam/core/graph/coder/map.go
@@ -0,0 +1,102 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package coder
+
+import (
+ "io"
+ "reflect"
+)
+
+// TODO(lostluck): 2020.08.04 export these for use for others?
+
+// mapDecoder produces a decoder for the beam schema map encoding.
+func mapDecoder(rt reflect.Type, decodeToKey, decodeToElem func(reflect.Value, io.Reader) error) func(reflect.Value, io.Reader) error {
+ return func(ret reflect.Value, r io.Reader) error {
+ // (1) Read count prefixed encoded data
+ size, err := DecodeInt32(r)
+ if err != nil {
+ return err
+ }
+ n := int(size)
+ ret.Set(reflect.MakeMapWithSize(rt, n))
+ for i := 0; i < n; i++ {
+ rvk := reflect.New(rt.Key()).Elem()
+ if err := decodeToKey(rvk, r); err != nil {
+ return err
+ }
+ rvv := reflect.New(rt.Elem()).Elem()
+ if err := decodeToElem(rvv, r); err != nil {
+ return err
+ }
+ ret.SetMapIndex(rvk, rvv)
+ }
+ return nil
+ }
+}
+
+// containerNilDecoder handles when a value is nillable for map or iterable components.
+// Nillable types have an extra byte prefixing them indicating nil status.
+func containerNilDecoder(decodeToElem func(reflect.Value, io.Reader) error) func(reflect.Value, io.Reader) error {
+ return func(ret reflect.Value, r io.Reader) error {
+ hasValue, err := DecodeBool(r)
+ if err != nil {
+ return err
+ }
+ if !hasValue {
+ return nil
+ }
+ rv := reflect.New(ret.Type().Elem())
+ if err := decodeToElem(rv.Elem(), r); err != nil {
+ return err
+ }
+ ret.Set(rv)
+ return nil
+ }
+}
+
+// mapEncoder reflectively encodes a map or array type using the beam map encoding.
+func mapEncoder(rt reflect.Type, encodeKey, encodeValue func(reflect.Value, io.Writer) error) func(reflect.Value, io.Writer) error {
+ return func(rv reflect.Value, w io.Writer) error {
+ size := rv.Len()
+ if err := EncodeInt32((int32)(size), w); err != nil {
+ return err
+ }
+ iter := rv.MapRange()
+ for iter.Next() {
+ if err := encodeKey(iter.Key(), w); err != nil {
+ return err
+ }
+ if err := encodeValue(iter.Value(), w); err != nil {
+ return err
+ }
+ }
+ return nil
+ }
+}
+
+// containerNilEncoder handles when a value is nillable for map or iterable components.
+// Nillable types have an extra byte prefixing them indicating nil status.
+func containerNilEncoder(encodeElem func(reflect.Value, io.Writer) error) func(reflect.Value, io.Writer) error {
+ return func(rv reflect.Value, w io.Writer) error {
+ if rv.IsNil() {
+ return EncodeBool(false, w)
+ }
+ if err := EncodeBool(true, w); err != nil {
+ return err
+ }
+ return encodeElem(rv.Elem(), w)
+ }
+}
diff --git a/sdks/go/pkg/beam/core/graph/coder/map_test.go b/sdks/go/pkg/beam/core/graph/coder/map_test.go
new file mode 100644
index 0000000..0b825c2
--- /dev/null
+++ b/sdks/go/pkg/beam/core/graph/coder/map_test.go
@@ -0,0 +1,98 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package coder
+
+import (
+ "bytes"
+ "fmt"
+ "io"
+ "reflect"
+ "testing"
+
+ "github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
+ "github.com/google/go-cmp/cmp"
+)
+
+func TestEncodeDecodeMap(t *testing.T) {
+ byteEnc := containerEncoderForType(reflectx.Uint8)
+ byteDec := containerDecoderForType(reflectx.Uint8)
+ bytePtrEnc := containerEncoderForType(reflect.PtrTo(reflectx.Uint8))
+ bytePtrDec := containerDecoderForType(reflect.PtrTo(reflectx.Uint8))
+
+ ptrByte := byte(42)
+
+ tests := []struct {
+ v interface{}
+ encK, encV func(reflect.Value, io.Writer) error
+ decK, decV func(reflect.Value, io.Reader) error
+ encoded []byte
+ decodeOnly bool
+ }{
+ {
+ v: map[byte]byte{10: 42},
+ encK: byteEnc,
+ encV: byteEnc,
+ decK: byteDec,
+ decV: byteDec,
+ encoded: []byte{0, 0, 0, 1, 10, 42},
+ }, {
+ v: map[byte]*byte{10: &ptrByte},
+ encK: byteEnc,
+ encV: bytePtrEnc,
+ decK: byteDec,
+ decV: bytePtrDec,
+ encoded: []byte{0, 0, 0, 1, 10, 1, 42},
+ }, {
+ v: map[byte]*byte{10: &ptrByte, 23: nil, 53: nil},
+ encK: byteEnc,
+ encV: bytePtrEnc,
+ decK: byteDec,
+ decV: bytePtrDec,
+ encoded: []byte{0, 0, 0, 3, 10, 1, 42, 23, 0, 53, 0},
+ decodeOnly: true,
+ },
+ }
+ for _, test := range tests {
+ test := test
+ if !test.decodeOnly {
+ t.Run(fmt.Sprintf("encode %q", test.v), func(t *testing.T) {
+ var buf bytes.Buffer
+ err := mapEncoder(reflect.TypeOf(test.v), test.encK, test.encV)(reflect.ValueOf(test.v), &buf)
+ if err != nil {
+ t.Fatalf("mapEncoder(%q) = %v", test.v, err)
+ }
+ if d := cmp.Diff(test.encoded, buf.Bytes()); d != "" {
+ t.Errorf("mapEncoder(%q) = %v, want %v diff(-want,+got):\n %v", test.v, buf.Bytes(), test.encoded, d)
+ }
+ })
+ }
+ t.Run(fmt.Sprintf("decode %v", test.v), func(t *testing.T) {
+ buf := bytes.NewBuffer(test.encoded)
+ rt := reflect.TypeOf(test.v)
+ var dec func(reflect.Value, io.Reader) error
+ dec = mapDecoder(rt, test.decK, test.decV)
+ rv := reflect.New(rt).Elem()
+ err := dec(rv, buf)
+ if err != nil {
+ t.Fatalf("mapDecoder(%q) = %v", test.encoded, err)
+ }
+ got := rv.Interface()
+ if d := cmp.Diff(test.v, got); d != "" {
+ t.Errorf("mapDecoder(%q) = %q, want %v diff(-want,+got):\n %v", test.encoded, got, test.v, d)
+ }
+ })
+ }
+}
diff --git a/sdks/go/pkg/beam/core/graph/coder/row.go b/sdks/go/pkg/beam/core/graph/coder/row.go
index aac34ac..00b4c26 100644
--- a/sdks/go/pkg/beam/core/graph/coder/row.go
+++ b/sdks/go/pkg/beam/core/graph/coder/row.go
@@ -152,21 +152,26 @@
return nil
}
}
- decf := decoderForSingleTypeReflect(t.Elem())
- sdec := iterableDecoderForSlice(t, decf)
- return func(rv reflect.Value, r io.Reader) error {
- return sdec(rv, r)
- }
+ decf := containerDecoderForType(t.Elem())
+ return iterableDecoderForSlice(t, decf)
case reflect.Array:
- decf := decoderForSingleTypeReflect(t.Elem())
- sdec := iterableDecoderForArray(t, decf)
- return func(rv reflect.Value, r io.Reader) error {
- return sdec(rv, r)
- }
+ decf := containerDecoderForType(t.Elem())
+ return iterableDecoderForArray(t, decf)
+ case reflect.Map:
+ decK := containerDecoderForType(t.Key())
+ decV := containerDecoderForType(t.Elem())
+ return mapDecoder(t, decK, decV)
}
panic(fmt.Sprintf("unimplemented type to decode: %v", t))
}
+func containerDecoderForType(t reflect.Type) func(reflect.Value, io.Reader) error {
+ if t.Kind() == reflect.Ptr {
+ return containerNilDecoder(decoderForSingleTypeReflect(t.Elem()))
+ }
+ return decoderForSingleTypeReflect(t)
+}
+
type typeDecoderReflect struct {
typ reflect.Type
fields []func(reflect.Value, io.Reader) error
@@ -270,15 +275,26 @@
return EncodeBytes(rv.Bytes(), w)
}
}
- encf := encoderForSingleTypeReflect(t.Elem())
+ encf := containerEncoderForType(t.Elem())
return iterableEncoder(t, encf)
case reflect.Array:
- encf := encoderForSingleTypeReflect(t.Elem())
+ encf := containerEncoderForType(t.Elem())
return iterableEncoder(t, encf)
+ case reflect.Map:
+ encK := containerEncoderForType(t.Key())
+ encV := containerEncoderForType(t.Elem())
+ return mapEncoder(t, encK, encV)
}
panic(fmt.Sprintf("unimplemented type to encode: %v", t))
}
+func containerEncoderForType(t reflect.Type) func(reflect.Value, io.Writer) error {
+ if t.Kind() == reflect.Ptr {
+ return containerNilEncoder(encoderForSingleTypeReflect(t.Elem()))
+ }
+ return encoderForSingleTypeReflect(t)
+}
+
type typeEncoderReflect struct {
fields []func(reflect.Value, io.Writer) error
}
diff --git a/sdks/go/pkg/beam/core/graph/coder/row_test.go b/sdks/go/pkg/beam/core/graph/coder/row_test.go
index f1089b8..38b7c5d 100644
--- a/sdks/go/pkg/beam/core/graph/coder/row_test.go
+++ b/sdks/go/pkg/beam/core/graph/coder/row_test.go
@@ -78,16 +78,18 @@
V12 [0]int
V13 [2]int
V14 []int
- // V15 map[string]int // not yet a standard coder (BEAM-7996)
+ V15 map[string]int
V16 float32
V17 float64
V18 []byte
+ V19 [2]*int
+ V20 map[*string]*int
}{},
}, {
want: struct {
V00 bool
- V01 byte
- V02 uint8
+ V01 byte // unsupported by spec (same as uint8)
+ V02 uint8 // unsupported by spec
V03 int16
// V04 uint16 // unsupported by spec
V05 int32
@@ -100,10 +102,13 @@
V12 [0]int
V13 [2]int
V14 []int
- // V15 map[string]int // not yet a standard coder (BEAM-7996) (encoding unspecified)
+ V15 map[string]int
V16 float32
V17 float64
V18 []byte
+ V19 [2]*int
+ V20 map[string]*int
+ V21 []*int
}{
V00: true,
V01: 1,
@@ -117,9 +122,16 @@
V12: [0]int{},
V13: [2]int{72, 908},
V14: []int{12, 9326, 641346, 6},
+ V15: map[string]int{"pants": 42},
V16: 3.14169,
V17: 2.6e100,
V18: []byte{21, 17, 65, 255, 0, 16},
+ V19: [2]*int{nil, &num},
+ V20: map[string]*int{
+ "notnil": &num,
+ "nil": nil,
+ },
+ V21: []*int{nil, &num, nil},
},
// TODO add custom types such as protocol buffers.
},
diff --git a/sdks/go/pkg/beam/core/runtime/exec/datasource.go b/sdks/go/pkg/beam/core/runtime/exec/datasource.go
index 054e34f..ad7b601 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/datasource.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/datasource.go
@@ -16,6 +16,7 @@
package exec
import (
+ "bytes"
"context"
"fmt"
"io"
@@ -25,7 +26,6 @@
"time"
"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
- "github.com/apache/beam/sdks/go/pkg/beam/core/sdf"
"github.com/apache/beam/sdks/go/pkg/beam/core/util/ioutilx"
"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
"github.com/apache/beam/sdks/go/pkg/beam/log"
@@ -47,20 +47,24 @@
splitIdx int64
start time.Time
- // rt is non-nil if this DataSource feeds directly to a splittable unit,
- // and receives the current restriction tracker being processed.
- rt chan sdf.RTracker
+ // su is non-nil if this DataSource feeds directly to a splittable unit,
+ // and receives that splittable unit when it is available for splitting.
+ // While the splittable unit is received, it is blocked from processing
+ // new elements, so it must be sent back through the channel once the
+ // DataSource is finished using it.
+ su chan SplittableUnit
mu sync.Mutex
}
-// Initializes the rt channel from the following unit when applicable.
+// InitSplittable initializes the SplittableUnit channel from the output unit,
+// if it provides one.
func (n *DataSource) InitSplittable() {
if n.Out == nil {
return
}
if u, ok := n.Out.(*ProcessSizedElementsAndRestrictions); ok == true {
- n.rt = u.Rt
+ n.su = u.SU
}
}
@@ -284,16 +288,21 @@
// Split takes a sorted set of potential split indices and a fraction of the
// remainder to split at, selects and actuates a split on an appropriate split
-// index, and returns the selected split index if successful or an error when
-// unsuccessful.
+// index, and returns the selected split index in a SplitResult if successful or
+// an error when unsuccessful.
+//
+// If the following transform is splittable, and the split indices and fraction
+// allow for splitting on the currently processing element, then a sub-element
+// split is performed, and the appropriate information is returned in the
+// SplitResult.
//
// The bufSize param specifies the estimated number of elements that will be
// sent to this DataSource, and is used to be able to perform accurate splits
// even if the DataSource has not yet received all its elements. A bufSize of
// 0 or less indicates that its unknown, and so uses the current known size.
-func (n *DataSource) Split(splits []int64, frac float64, bufSize int64) (int64, error) {
+func (n *DataSource) Split(splits []int64, frac float64, bufSize int64) (SplitResult, error) {
if n == nil {
- return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource not initialized", splits)
+ return SplitResult{}, fmt.Errorf("failed to split at requested splits: {%v}, DataSource not initialized", splits)
}
if frac > 1.0 {
frac = 1.0
@@ -302,31 +311,83 @@
}
n.mu.Lock()
+ defer n.mu.Unlock()
+
var currProg float64 // Current element progress.
- if n.index < 0 { // Progress is at the end of the non-existant -1st element.
+ var su SplittableUnit
+ if n.index < 0 { // Progress is at the end of the non-existant -1st element.
currProg = 1.0
- } else if n.rt == nil { // If this isn't sub-element splittable, estimate some progress.
+ } else if n.su == nil { // If this isn't sub-element splittable, estimate some progress.
currProg = 0.5
} else { // If this is sub-element splittable, get progress of the current element.
- rt := <-n.rt
- d, r := rt.GetProgress()
- currProg = d / (d + r)
- n.rt <- rt
+
+ select {
+ case su = <-n.su:
+ // If an element is processing, we'll get a splittable unit.
+ if su == nil {
+ return SplitResult{}, fmt.Errorf("failed to split: splittable unit was nil")
+ }
+ defer func() {
+ n.su <- su
+ }()
+ currProg = su.GetProgress()
+ case <-time.After(500 * time.Millisecond):
+ // Otherwise, the current element hasn't started processing yet
+ // or has already finished. By adding a short timeout, we avoid
+ // the first possibility, and can assume progress is at max.
+ currProg = 1.0
+ }
}
// Size to split within is the minimum of bufSize or splitIdx so we avoid
// including elements we already know won't be processed.
if bufSize <= 0 || n.splitIdx < bufSize {
bufSize = n.splitIdx
}
- s, _, err := splitHelper(n.index, bufSize, currProg, splits, frac, false)
+ s, f, err := splitHelper(n.index, bufSize, currProg, splits, frac, su != nil)
if err != nil {
- n.mu.Unlock()
- return 0, err
+ return SplitResult{}, err
}
- n.splitIdx = s
- fs := n.splitIdx
- n.mu.Unlock()
- return fs, nil
+
+ // No fraction returned, perform channel split.
+ if f < 0 {
+ n.splitIdx = s
+ return SplitResult{PI: s - 1, RI: s}, nil
+ }
+ // Otherwise, perform a sub-element split.
+ fr := f / (1.0 - currProg)
+ p, r, err := su.Split(fr)
+ if err != nil {
+ return SplitResult{}, err
+ }
+
+ if p == nil || r == nil { // Unsuccessful split.
+ // Fallback to channel split, so split at next elm, not current.
+ n.splitIdx = s + 1
+ return SplitResult{PI: s, RI: s + 1}, nil
+ }
+
+ // TODO(BEAM-10579) Eventually encode elements with the splittable
+ // unit's input coder instead of the DataSource's coder.
+ wc := MakeWindowEncoder(n.Coder.Window)
+ ec := MakeElementEncoder(coder.SkipW(n.Coder))
+ pEnc, err := encodeElm(p, wc, ec)
+ if err != nil {
+ return SplitResult{}, err
+ }
+ rEnc, err := encodeElm(r, wc, ec)
+ if err != nil {
+ return SplitResult{}, err
+ }
+ n.splitIdx = s + 1 // In a sub-element split, s is currIdx.
+ res := SplitResult{
+ PI: s - 1,
+ RI: s + 1,
+ PS: pEnc,
+ RS: rEnc,
+ TId: su.GetTransformId(),
+ InId: su.GetInputId(),
+ }
+ return res, nil
}
// splitHelper is a helper function that finds a split point in a range.
@@ -348,7 +409,8 @@
// Returns the element index to split at (first element of residual), and the
// fraction within that element to split, iff the split point is the current
// element, the splittable param is set to true, and both the element being
-// split and the following element are valid split points.
+// split and the following element are valid split points. If there is no
+// fraction, returns -1.
func splitHelper(
currIdx, endIdx int64,
currProg float64,
@@ -377,7 +439,7 @@
if splitIdx < safeStart {
splitIdx = safeStart
}
- return splitIdx, 0.0, nil
+ return splitIdx, -1.0, nil
}
// Cases where we have to find a valid split point.
@@ -419,10 +481,21 @@
}
}
if bestS != -1 {
- return bestS, 0.0, nil
+ return bestS, -1.0, nil
}
- return -1, 0.0, fmt.Errorf("failed to split DataSource (at index: %v) at requested splits: {%v}", currIdx, splits)
+ return -1, -1.0, fmt.Errorf("failed to split DataSource (at index: %v) at requested splits: {%v}", currIdx, splits)
+}
+
+func encodeElm(elm *FullValue, wc WindowEncoder, ec ElementEncoder) ([]byte, error) {
+ var b bytes.Buffer
+ if err := EncodeWindowedValueHeader(wc, elm.Windows, elm.Timestamp, &b); err != nil {
+ return nil, err
+ }
+ if err := ec.Encode(elm, &b); err != nil {
+ return nil, err
+ }
+ return b.Bytes(), nil
}
type concatReStream struct {
diff --git a/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go b/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go
index ff6126e..4089a1c 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go
@@ -25,6 +25,7 @@
"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
"github.com/apache/beam/sdks/go/pkg/beam/core/graph/mtime"
"github.com/apache/beam/sdks/go/pkg/beam/core/graph/window"
+ "github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
)
func TestDataSource_PerElement(t *testing.T) {
@@ -278,11 +279,17 @@
runOnRoots(ctx, t, p, "StartBundle", func(root Root, ctx context.Context) error { return root.StartBundle(ctx, "1", dc) })
// SDK never splits on 0, so check that every test.
- if splitIdx, err := p.Split(SplitPoints{Splits: []int64{0, test.splitIdx}}); err != nil {
+ splitRes, err := p.Split(SplitPoints{Splits: []int64{0, test.splitIdx}})
+ if err != nil {
t.Fatalf("error in Split: %v", err)
- } else if got, want := splitIdx, test.splitIdx; got != want {
+ }
+ if got, want := splitRes.RI, test.splitIdx; got != want {
t.Fatalf("error in Split: got splitIdx = %v, want %v ", got, want)
}
+ if got, want := splitRes.PI, test.splitIdx-1; got != want {
+ t.Fatalf("error in Split: got primary index = %v, want %v ", got, want)
+ }
+
runOnRoots(ctx, t, p, "Process", Root.Process)
runOnRoots(ctx, t, p, "FinishBundle", Root.FinishBundle)
@@ -342,10 +349,15 @@
<-blockedCh
// Validate that we do not split on the element we're blocking on index.
// The first valid split is at test.splitIdx.
- if splitIdx, err := source.Split([]int64{0, 1, 2, 3, 4, 5}, -1, 0); err != nil {
+ if splitRes, err := source.Split([]int64{0, 1, 2, 3, 4, 5}, -1, 0); err != nil {
t.Errorf("error in Split: %v", err)
- } else if got, want := splitIdx, test.splitIdx; got != want {
- t.Errorf("error in Split: got splitIdx = %v, want %v ", got, want)
+ } else {
+ if got, want := splitRes.RI, test.splitIdx; got != want {
+ t.Errorf("error in Split: got splitIdx = %v, want %v ", got, want)
+ }
+ if got, want := splitRes.PI, test.splitIdx-1; got != want {
+ t.Errorf("error in Split: got primary index = %v, want %v ", got, want)
+ }
}
// Validate that our progress is where we expect it to be. (test.splitIdx - 1)
if got, want := source.Progress().Count, test.splitIdx-1; got != want {
@@ -412,17 +424,120 @@
// SDK never splits on 0, so check that every test.
sp := SplitPoints{Splits: test.splitPts, Frac: test.frac, BufSize: test.bufSize}
- if splitIdx, err := p.Split(sp); err != nil {
+ splitRes, err := p.Split(sp)
+ if err != nil {
t.Fatalf("error in Split: %v", err)
- } else if got, want := splitIdx, test.splitIdx; got != want {
+ }
+ if got, want := splitRes.RI, test.splitIdx; got != want {
t.Fatalf("error in Split: got splitIdx = %v, want %v ", got, want)
}
+ if got, want := splitRes.PI, test.splitIdx-1; got != want {
+ t.Fatalf("error in Split: got primary index = %v, want %v ", got, want)
+ }
runOnRoots(ctx, t, p, "Process", Root.Process)
runOnRoots(ctx, t, p, "FinishBundle", Root.FinishBundle)
validateSource(t, out, source, makeValues(test.expected...))
})
+ // Test splitting on sub-elements works when available.
+ t.Run("subElement", func(t *testing.T) {
+ // Each test will process up to an element, then split at different
+ // fractions and check that a sub-element split either was, or was not
+ // performed.
+ const blockOn int64 = 3 // Should leave 2 elements unprocessed, including blocked element.
+ numElms := int64(len(elements))
+ tests := []struct {
+ fraction float64
+ splitIdx int64
+ isSubElm bool
+ }{
+ {fraction: 0.0, splitIdx: blockOn + 1, isSubElm: true},
+ {fraction: 0.01, splitIdx: blockOn + 1, isSubElm: true},
+ {fraction: 0.49, splitIdx: blockOn + 1, isSubElm: true}, // Should be just within current element.
+ {fraction: 0.51, splitIdx: blockOn + 1, isSubElm: false}, // Should be just past current element.
+ {fraction: 0.99, splitIdx: numElms, isSubElm: false},
+ }
+ for _, test := range tests {
+ test := test
+ name := fmt.Sprintf("withFraction_%v", test.fraction)
+ t.Run(name, func(t *testing.T) {
+ source, out, pr := initSourceTest(name)
+ unblockCh, blockedCh := make(chan struct{}), make(chan struct{}, 1)
+ // Block on the one less than the desired split,
+ // so the desired split is the first valid split.
+ blocker := &BlockingNode{
+ UID: 3,
+ Block: func(elm *FullValue) bool {
+ if source.index == blockOn {
+ // Signal to call Split
+ blockedCh <- struct{}{}
+ return true
+ }
+ return false
+ },
+ Unblock: unblockCh,
+ Out: out,
+ }
+ source.Out = blocker
+
+ splittableCh := make(chan SplittableUnit, 1)
+ source.su = splittableCh
+ splittableCh <- &TestSplittableUnit{elm: elements[blockOn]}
+
+ go func() {
+ // Wait to call Split until the DoFn is blocked at the desired element.
+ <-blockedCh
+ // Validate that we either do or do not perform a sub-element split with the
+ // given fraction.
+ if splitRes, err := source.Split([]int64{0, 1, 2, 3, 4, 5}, test.fraction, int64(len(elements))); err != nil {
+ t.Errorf("error in Split: %v", err)
+ } else {
+ // For sub-element splits, check sub-element split only results.
+ isSubElm := splitRes.RS != nil && splitRes.PS != nil
+ if isSubElm != test.isSubElm {
+ t.Errorf("error in Split: got sub-element split = %t, want %t", isSubElm, test.isSubElm)
+ }
+ if isSubElm {
+ if got, want := splitRes.TId, testTransformId; got != want {
+ t.Errorf("error in Split: got incorrect Transform Id = %v, want %v", got, want)
+ }
+ if got, want := splitRes.InId, testInputId; got != want {
+ t.Errorf("error in Split: got incorrect Input Id = %v, want %v", got, want)
+ }
+ }
+
+ // Check that split indices are correct, for both sub-element and channel splits.
+ var wantPI, wantRI = test.splitIdx - 1, test.splitIdx
+ if isSubElm {
+ // In sub-element splits, primary index is expected to be one element
+ // before the current (split) element.
+ wantPI--
+ }
+ if splitRes.PI != wantPI || splitRes.RI != wantRI {
+ t.Errorf("error in Split: got split indices of (primary, residual) = (%d, %d), want (%d, %d)",
+ splitRes.PI, splitRes.RI, wantPI, wantRI)
+ }
+ }
+ // Validate that our progress is where we expect it to be. (blockOn)
+ if got, want := source.Progress().Count, blockOn; got != want {
+ t.Errorf("error in Progress: got finished processing Count = %v, want %v ", got, want)
+ }
+ unblockCh <- struct{}{}
+ }()
+
+ constructAndExecutePlanWithContext(t, []Unit{out, blocker, source}, DataContext{
+ Data: &TestDataManager{R: pr},
+ })
+
+ validateSource(t, out, source, makeValues(elements[:test.splitIdx]...))
+ if got, want := source.Progress().Count, test.splitIdx; got != want {
+ t.Fatalf("progress didn't match split: got %v, want %v", got, want)
+ }
+ })
+ }
+ })
+
// Test expects splitting errors, but for processing to be successful.
t.Run("errors", func(t *testing.T) {
source, out, pr := initSourceTest("noSplitsUntilStarted")
@@ -471,6 +586,39 @@
})
}
+const testTransformId = "transform_id"
+const testInputId = "input_id"
+
+// TestSplittableUnit is an implementation of the SplittableUnit interface
+// for DataSource tests.
+type TestSplittableUnit struct {
+ elm interface{} // The element to split.
+}
+
+// Split checks the input fraction for correctness, but otherwise always returns
+// a successful split. The split elements are just copies of the original.
+func (n *TestSplittableUnit) Split(f float64) (*FullValue, *FullValue, error) {
+ if f > 1.0 || f < 0.0 {
+ return nil, nil, errors.Errorf("Error")
+ }
+ return &FullValue{Elm: n.elm}, &FullValue{Elm: n.elm}, nil
+}
+
+// GetProgress always returns 0, to keep tests consistent.
+func (n *TestSplittableUnit) GetProgress() float64 {
+ return 0
+}
+
+// GetTransformId returns a constant transform ID that can be tested for.
+func (n *TestSplittableUnit) GetTransformId() string {
+ return testTransformId
+}
+
+// GetInputId returns a constant input ID that can be tested for.
+func (n *TestSplittableUnit) GetInputId() string {
+ return testInputId
+}
+
func floatEquals(a, b, epsilon float64) bool {
return math.Abs(a-b) < epsilon
}
@@ -501,7 +649,7 @@
for _, test := range tests {
test := test
t.Run(fmt.Sprintf("(%v of [%v, %v])", test.frac, test.curr, test.size), func(t *testing.T) {
- wantFrac := 0.0
+ wantFrac := -1.0
got, gotFrac, err := splitHelper(test.curr, test.size, 0.0, nil, test.frac, false)
if err != nil {
t.Fatalf("error in splitHelper: %v", err)
@@ -533,7 +681,7 @@
for _, test := range tests {
test := test
t.Run(fmt.Sprintf("(%v of [%v, %v])", test.frac, float64(test.curr)+test.currProg, test.size), func(t *testing.T) {
- wantFrac := 0.0
+ wantFrac := -1.0
got, gotFrac, err := splitHelper(test.curr, test.size, test.currProg, nil, test.frac, false)
if err != nil {
t.Fatalf("error in splitHelper: %v", err)
@@ -572,7 +720,7 @@
for _, test := range tests {
test := test
t.Run(fmt.Sprintf("(%v of [%v, %v], splits = %v)", test.frac, test.curr, test.size, test.splits), func(t *testing.T) {
- wantFrac := 0.0
+ wantFrac := -1.0
got, gotFrac, err := splitHelper(test.curr, test.size, 0.0, test.splits, test.frac, false)
if test.err {
if err == nil {
@@ -602,10 +750,10 @@
wantFrac float64
}{
// Split between future elements at element boundaries.
- {curr: 0, currProg: 0, size: 4, frac: 0.51, want: 2},
- {curr: 0, currProg: 0, size: 4, frac: 0.49, want: 2},
- {curr: 0, currProg: 0, size: 4, frac: 0.26, want: 1},
- {curr: 0, currProg: 0, size: 4, frac: 0.25, want: 1},
+ {curr: 0, currProg: 0, size: 4, frac: 0.51, want: 2, wantFrac: -1.0},
+ {curr: 0, currProg: 0, size: 4, frac: 0.49, want: 2, wantFrac: -1.0},
+ {curr: 0, currProg: 0, size: 4, frac: 0.26, want: 1, wantFrac: -1.0},
+ {curr: 0, currProg: 0, size: 4, frac: 0.25, want: 1, wantFrac: -1.0},
// If the split falls inside the first, splittable element, split there.
{curr: 0, currProg: 0, size: 4, frac: 0.20, want: 0, wantFrac: 0.8},
@@ -613,11 +761,11 @@
{curr: 0, currProg: 0, size: 4, frac: 0.125, want: 0, wantFrac: 0.5},
// Here we are far enough into the first element that splitting at 0.2 of the
// remainder falls outside the first element.
- {curr: 0, currProg: 0.5, size: 4, frac: 0.2, want: 1},
+ {curr: 0, currProg: 0.5, size: 4, frac: 0.2, want: 1, wantFrac: -1.0},
// Verify the above logic when we are partially through the stream.
- {curr: 2, currProg: 0, size: 4, frac: 0.6, want: 3},
- {curr: 2, currProg: 0.9, size: 4, frac: 0.6, want: 4},
+ {curr: 2, currProg: 0, size: 4, frac: 0.6, want: 3, wantFrac: -1.0},
+ {curr: 2, currProg: 0.9, size: 4, frac: 0.6, want: 4, wantFrac: -1.0},
{curr: 2, currProg: 0.5, size: 4, frac: 0.2, want: 2, wantFrac: 0.8},
}
for _, test := range tests {
@@ -649,12 +797,12 @@
// This is where we would like to split, when all split points are available.
{curr: 2, currProg: 0, size: 5, frac: 0.2, splits: []int64{1, 2, 3, 4, 5}, want: 2, wantFrac: 0.6},
// We can't split element at index 2, because 3 is not a split point.
- {curr: 2, currProg: 0, size: 5, frac: 0.2, splits: []int64{1, 2, 4, 5}, want: 4},
+ {curr: 2, currProg: 0, size: 5, frac: 0.2, splits: []int64{1, 2, 4, 5}, want: 4, wantFrac: -1.0},
// We can't even split element at index 4 as above, because 4 is also not a
// split point.
- {curr: 2, currProg: 0, size: 5, frac: 0.2, splits: []int64{1, 2, 5}, want: 5},
+ {curr: 2, currProg: 0, size: 5, frac: 0.2, splits: []int64{1, 2, 5}, want: 5, wantFrac: -1.0},
// We can't split element at index 2, because 2 is not a split point.
- {curr: 2, currProg: 0, size: 5, frac: 0.2, splits: []int64{1, 3, 4, 5}, want: 3},
+ {curr: 2, currProg: 0, size: 5, frac: 0.2, splits: []int64{1, 3, 4, 5}, want: 3, wantFrac: -1.0},
}
for _, test := range tests {
test := test
diff --git a/sdks/go/pkg/beam/core/runtime/exec/dynsplit_test.go b/sdks/go/pkg/beam/core/runtime/exec/dynsplit_test.go
new file mode 100644
index 0000000..3b17a0a
--- /dev/null
+++ b/sdks/go/pkg/beam/core/runtime/exec/dynsplit_test.go
@@ -0,0 +1,456 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 exec
+
+import (
+ "bytes"
+ "context"
+ "io"
+ "reflect"
+ "sync"
+ "testing"
+
+ "github.com/apache/beam/sdks/go/pkg/beam/core/graph"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/graph/mtime"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/graph/window"
+ "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"
+ "github.com/apache/beam/sdks/go/pkg/beam/io/rtrackers/offsetrange"
+ "github.com/google/go-cmp/cmp"
+ "github.com/google/go-cmp/cmp/cmpopts"
+)
+
+// TestDynamicSplit tests that a dynamic split of an in-progress SDF succeeds
+// with valid input. It coordinates the two threads (processing and splitting)
+// to test what happens if operations happen in various orders. The test then
+// validates that the output of the SDF is correct according to the split.
+func TestDynamicSplit(t *testing.T) {
+ tests := []struct {
+ name string
+ // driver is a function determining how the processing and splitting
+ // threads are created and coordinated.
+ driver func(*Plan, DataContext, *splitTestSdf) (error, splitResult)
+ }{
+ {
+ // Complete a split before beginning processing.
+ name: "Simple",
+ driver: nonBlockingDriver,
+ },
+ {
+ // Try claiming while blocked on a split.
+ name: "BlockOnSplit",
+ driver: splitBlockingDriver,
+ },
+ {
+ // Try splitting while blocked on a claim.
+ name: "BlockOnClaim",
+ driver: claimBlockingDriver,
+ },
+ }
+ for _, test := range tests {
+ test := test
+ t.Run(test.name, func(t *testing.T) {
+ // Create pipeline.
+ sdf := newSplitTestSdf()
+ dfn, err := graph.NewDoFn(sdf, graph.NumMainInputs(graph.MainSingle))
+ if err != nil {
+ t.Fatalf("invalid function: %v", err)
+ }
+ cdr := createSplitTestInCoder()
+ plan, out := createSdfPlan(t, t.Name(), dfn, cdr)
+
+ // Create thread to send element to pipeline.
+ pr, pw := io.Pipe()
+ elm := createElm()
+ go writeElm(elm, cdr, pw)
+ dc := DataContext{Data: &TestDataManager{R: pr}}
+
+ // Call driver to coordinate processing & splitting threads.
+ procRes, splitRes := test.driver(plan, dc, sdf)
+
+ // Validate we get a valid split result, aside from split elements.
+ if splitRes.err != nil {
+ t.Fatalf("Plan.Split failed: %v", splitRes.err)
+ }
+ wantSplit := SplitResult{
+ PI: -1,
+ RI: 1,
+ PS: nil,
+ RS: nil,
+ TId: testTransformId,
+ InId: indexToInputId(0),
+ }
+ if diff := cmp.Diff(splitRes.split, wantSplit, cmpopts.IgnoreFields(SplitResult{}, "PS", "RS")); diff != "" {
+ t.Errorf("Incorrect split result (ignoring split elements): %v", diff)
+ }
+
+ // Validate split elements are encoded correctly by decoding them
+ // with the input coder to the path.
+ // TODO(BEAM-10579) Switch to using splittable unit's input coder
+ // once that is implemented.
+ p, err := decodeDynSplitElm(splitRes.split.PS, cdr)
+ if err != nil {
+ t.Errorf("Failed decoding primary element split: %v", err)
+ }
+ _, err = decodeDynSplitElm(splitRes.split.RS, cdr)
+ if err != nil {
+ t.Errorf("Failed decoding residual element split: %v", err)
+ }
+
+ // Validate SDF output. Make sure each restriction matches the split result.
+ if err := procRes; err != nil {
+ t.Fatal(err)
+ }
+ pRest := p.Elm.(*FullValue).Elm2.(offsetrange.Restriction)
+ if got, want := len(out.Elements), int(pRest.End-pRest.Start); got != want {
+ t.Errorf("Unexpected number of elements: got: %v, want: %v", got, want)
+ }
+ for i, fv := range out.Elements {
+ rest := fv.Elm.(offsetrange.Restriction)
+ if got, want := rest, pRest; !cmp.Equal(got, want) {
+ t.Errorf("Output element %v had incorrect restriction: got: %v, want: %v", i, got, want)
+ }
+ }
+ })
+ }
+}
+
+// nonBlockingDriver performs a split before starting processing, so no thread
+// is forced to wait on a mutex.
+func nonBlockingDriver(plan *Plan, dc DataContext, sdf *splitTestSdf) (procRes error, splitRes splitResult) {
+ // Begin processing pipeline.
+ procResCh := make(chan error)
+ go processPlan(plan, dc, procResCh)
+ rt := <-sdf.rt // Tracker is created first, retrieve that.
+
+ // Complete a split before unblocking processing.
+ splitResCh := make(chan splitResult)
+ go splitPlan(plan, splitResCh)
+ <-rt.split
+ <-rt.blockSplit
+ splitRes = <-splitResCh
+
+ // Unblock and finishing processing.
+ <-sdf.proc
+ <-rt.claim
+ <-rt.blockClaim
+ <-rt.endClaim
+ procRes = <-procResCh
+
+ return procRes, splitRes
+}
+
+// splitBlockingDriver blocks on a split request so that the SDF attempts to
+// claim while the split is occurring.
+func splitBlockingDriver(plan *Plan, dc DataContext, sdf *splitTestSdf) (procRes error, splitRes splitResult) {
+ // Begin processing pipeline.
+ procResCh := make(chan error)
+ go processPlan(plan, dc, procResCh)
+ rt := <-sdf.rt // Tracker is created first, retrieve that.
+
+ // Start a split, but block on it so it holds the mutex.
+ splitResCh := make(chan splitResult)
+ go splitPlan(plan, splitResCh)
+ <-rt.split
+
+ // Start processing and start a claim, that'll be waiting for the mutex.
+ <-sdf.proc
+ <-rt.claim
+
+ // Unblock and finish splitting and free the mutex.
+ <-rt.blockSplit
+ splitRes = <-splitResCh
+
+ // Unblock and finish claiming and processing.
+ <-rt.blockClaim
+ <-rt.endClaim
+ procRes = <-procResCh
+
+ return procRes, splitRes
+}
+
+// claimBlockingDriver blocks on a claim request so that the SDF attempts to
+// split while the claim is occurring.
+func claimBlockingDriver(plan *Plan, dc DataContext, sdf *splitTestSdf) (procRes error, splitRes splitResult) {
+ // Begin processing pipeline.
+ procResCh := make(chan error)
+ go processPlan(plan, dc, procResCh)
+ rt := <-sdf.rt // Tracker is created first, retrieve that.
+
+ // Start a claim, but block on it so it holds the mutex.
+ <-sdf.proc
+ <-rt.claim
+
+ // Start a split that'll be waiting for the mutex.
+ splitResCh := make(chan splitResult)
+ go splitPlan(plan, splitResCh)
+ <-rt.split
+
+ // Unblock the claim, freeing the mutex (but not finishing processing yet).
+ <-rt.blockClaim
+
+ // Finish splitting, allowing processing to finish.
+ <-rt.blockSplit
+ splitRes = <-splitResCh
+ <-rt.endClaim // Delay the claim end so we don't process too much before splitting.
+ procRes = <-procResCh
+
+ return procRes, splitRes
+}
+
+// createElm creates the element for our test pipeline.
+func createElm() *FullValue {
+ return &FullValue{
+ Elm: &FullValue{
+ Elm: 20,
+ Elm2: offsetrange.Restriction{Start: 0, End: 20},
+ },
+ Elm2: float64(20),
+ }
+}
+
+// createSplitTestInCoder outputs the coder for inputs to our test pipeline,
+// (in particular, the DataSource transform of the pipeline). For the specific
+// element this is a coder for, see createElm.
+func createSplitTestInCoder() *coder.Coder {
+ restT := reflect.TypeOf((*offsetrange.Restriction)(nil)).Elem()
+ restCdr := coder.LookupCustomCoder(restT)
+
+ cdr := coder.NewW(
+ coder.NewKV([]*coder.Coder{
+ coder.NewKV([]*coder.Coder{
+ intCoder(reflectx.Int),
+ {Kind: coder.Custom, T: typex.New(restT), Custom: restCdr},
+ }),
+ coder.NewDouble(),
+ }),
+ coder.NewGlobalWindow())
+ return cdr
+}
+
+// createSdfPlan creates a plan containing the test pipeline, which is
+// DataSource -> SDF.ProcessSizedElementsAndRestrictions -> CaptureNode.
+func createSdfPlan(t *testing.T, name string, fn *graph.DoFn, cdr *coder.Coder) (*Plan, *CaptureNode) {
+ out := &CaptureNode{UID: 0}
+ n := &ParDo{UID: 1, Fn: fn, Out: []Node{out}}
+ sdf := &ProcessSizedElementsAndRestrictions{PDo: n, TfId: testTransformId}
+ ds := &DataSource{
+ UID: 2,
+ SID: StreamID{PtransformID: "DataSource"},
+ Name: "name",
+ Coder: cdr,
+ Out: sdf,
+ }
+ units := []Unit{ds, sdf, out}
+
+ p, err := NewPlan(name+"_plan", units)
+ if err != nil {
+ t.Fatalf("NewPlan failed: %v", err)
+ }
+ return p, out
+}
+
+// writeElm is meant to be the goroutine for feeding an element to the
+// DataSourc of the test pipeline.
+func writeElm(elm *FullValue, cdr *coder.Coder, pw *io.PipeWriter) {
+ wc := MakeWindowEncoder(cdr.Window)
+ ec := MakeElementEncoder(coder.SkipW(cdr))
+ if err := EncodeWindowedValueHeader(wc, window.SingleGlobalWindow, mtime.ZeroTimestamp, pw); err != nil {
+ panic("err")
+ }
+ if err := ec.Encode(elm, pw); err != nil {
+ panic("err")
+ }
+ if err := pw.Close(); err != nil {
+ panic("err")
+ }
+}
+
+func decodeDynSplitElm(elm []byte, cdr *coder.Coder) (*FullValue, error) {
+ wd := MakeWindowDecoder(cdr.Window)
+ ed := MakeElementDecoder(coder.SkipW(cdr))
+ b := bytes.NewBuffer(elm)
+ w, t, err := DecodeWindowedValueHeader(wd, b)
+ if err != nil {
+ return nil, err
+ }
+ e, err := ed.Decode(b)
+ if err != nil {
+ return nil, err
+ }
+ e.Windows = w
+ e.Timestamp = t
+ return e, nil
+}
+
+// processPlan is meant to be the goroutine representing the thread processing
+// the SDF.
+func processPlan(plan *Plan, dc DataContext, result chan error) {
+ if err := plan.Execute(context.Background(), plan.ID()+"_execute", dc); err != nil {
+ result <- errors.Wrap(err, "Plan.Execute failed")
+ }
+ if err := plan.Down(context.Background()); err != nil {
+ result <- errors.Wrap(err, "Plan.Down failed")
+ }
+ result <- nil
+}
+
+type splitResult struct {
+ split SplitResult
+ err error
+}
+
+// splitPlan is meant to be the goroutine representing the thread handling a
+// split request for the SDF.
+func splitPlan(plan *Plan, result chan splitResult) {
+ split, err := plan.Split(SplitPoints{Frac: 0.5, BufSize: 1})
+ result <- splitResult{split: split, err: err}
+}
+
+// splitTestRTracker adds signals needed to coordinate splitting and claiming
+// over multiple threads for this test. Semantically, this tracker is an
+// offset range tracker representing a range of integers to output.
+type splitTestRTracker struct {
+ mu sync.Mutex // Lock on accessing underlying tracker.
+ rt *offsetrange.Tracker
+
+ // Send signals when starting a claim, blocking a claim, and ending a claim.
+ claim chan struct{}
+ blockClaim chan struct{}
+ endClaim chan struct{}
+ blockInd int64 // Only send signals when claiming a specific position.
+
+ // Send signals when starting a split, and blocking a split. Important note:
+ // the spot to use these in this test is dependent on the first operation
+ // taking place on a split, which may not necessarily be TrySplit.
+ split chan struct{}
+ blockSplit chan struct{}
+}
+
+func newSplitTestRTracker(rest offsetrange.Restriction) *splitTestRTracker {
+ return &splitTestRTracker{
+ rt: offsetrange.NewTracker(rest),
+ claim: make(chan struct{}, 1),
+ blockClaim: make(chan struct{}),
+ endClaim: make(chan struct{}),
+ blockInd: rest.Start,
+ split: make(chan struct{}, 1),
+ blockSplit: make(chan struct{}),
+ }
+}
+
+func (rt *splitTestRTracker) TryClaim(pos interface{}) bool {
+ i := pos.(int64)
+ if i == rt.blockInd {
+ rt.claim <- struct{}{}
+ }
+
+ rt.mu.Lock()
+ if i == rt.blockInd {
+ rt.blockClaim <- struct{}{}
+ }
+ result := rt.rt.TryClaim(pos)
+ rt.mu.Unlock()
+
+ if i == rt.blockInd {
+ rt.endClaim <- struct{}{}
+ }
+ return result
+}
+
+func (rt *splitTestRTracker) GetError() error {
+ rt.mu.Lock()
+ defer rt.mu.Unlock()
+ return rt.rt.GetError()
+}
+
+func (rt *splitTestRTracker) TrySplit(fraction float64) (interface{}, interface{}, error) {
+ rt.mu.Lock()
+ defer rt.mu.Unlock()
+ rt.blockSplit <- struct{}{}
+ return rt.rt.TrySplit(fraction)
+}
+
+func (rt *splitTestRTracker) GetProgress() (float64, float64) {
+ // Note: Currently, GetProgress is called first in a split and blocks if
+ // TryClaim is being called.
+ rt.split <- struct{}{}
+
+ rt.mu.Lock()
+ defer rt.mu.Unlock()
+ return rt.rt.GetProgress()
+}
+
+func (rt *splitTestRTracker) IsDone() bool {
+ rt.mu.Lock()
+ defer rt.mu.Unlock()
+ return rt.rt.IsDone()
+}
+
+func (rt *splitTestRTracker) GetRestriction() interface{} {
+ rt.mu.Lock()
+ defer rt.mu.Unlock()
+ return rt.rt.GetRestriction()
+}
+
+// splitTestSdf has signals needed to control processing behavior over multiple
+// threads. The actual behavior is to accept an integer N as the element and
+// output each element in the range of [0, N).
+type splitTestSdf struct {
+ proc chan struct{}
+ rt chan *splitTestRTracker // Used to provide created trackers to the test code.
+}
+
+func newSplitTestSdf() *splitTestSdf {
+ return &splitTestSdf{
+ proc: make(chan struct{}),
+ rt: make(chan *splitTestRTracker),
+ }
+}
+
+func (fn *splitTestSdf) ProcessElement(rt *splitTestRTracker, _ int, emit func(offsetrange.Restriction, int)) {
+ i := rt.GetRestriction().(offsetrange.Restriction).Start
+ fn.proc <- struct{}{}
+
+ for rt.TryClaim(i) == true {
+ rest := rt.GetRestriction().(offsetrange.Restriction)
+ emit(rest, int(i))
+ i++
+ }
+}
+
+func (fn *splitTestSdf) CreateInitialRestriction(i int) offsetrange.Restriction {
+ return offsetrange.Restriction{
+ Start: 0,
+ End: int64(i),
+ }
+}
+
+func (fn *splitTestSdf) SplitRestriction(_ int, rest offsetrange.Restriction) []offsetrange.Restriction {
+ return []offsetrange.Restriction{rest}
+}
+
+func (fn *splitTestSdf) RestrictionSize(_ int, rest offsetrange.Restriction) float64 {
+ return rest.Size()
+}
+
+func (fn *splitTestSdf) CreateTracker(rest offsetrange.Restriction) *splitTestRTracker {
+ rt := newSplitTestRTracker(rest)
+ fn.rt <- rt
+ return rt
+}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/plan.go b/sdks/go/pkg/beam/core/runtime/exec/plan.go
index 53891d6..23ce1c7 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/plan.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/plan.go
@@ -207,13 +207,28 @@
BufSize int64
}
+// SplitResult contains the result of performing a split on a Plan.
+type SplitResult struct {
+ // Indices are always included, for both channel and sub-element splits.
+ PI int64 // Primary index, last element of the primary.
+ RI int64 // Residual index, first element of the residual.
+
+ // Extra information included for sub-element splits. If PS and RS are
+ // present then a sub-element split occurred.
+ PS []byte // Primary split. If an element is split, this is the encoded primary.
+ RS []byte // Residual split. If an element is split, this is the encoded residual.
+ TId string // Transform ID of the transform receiving the split elements.
+ InId string // Input ID of the input the split elements are received from.
+}
+
// Split takes a set of potential split indexes, and if successful returns
-// the split index of the first element of the residual, on which processing
-// will be halted.
+// the split result.
// Returns an error when unable to split.
-func (p *Plan) Split(s SplitPoints) (int64, error) {
+func (p *Plan) Split(s SplitPoints) (SplitResult, error) {
+ // TODO: When bundles with multiple sources, are supported, perform splits
+ // on all sources.
if p.source != nil {
return p.source.Split(s.Splits, s.Frac, s.BufSize)
}
- return 0, fmt.Errorf("failed to split at requested splits: {%v}, Source not initialized", s)
+ return SplitResult{}, fmt.Errorf("failed to split at requested splits: {%v}, Source not initialized", s)
}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf.go b/sdks/go/pkg/beam/core/runtime/exec/sdf.go
index c24e34a..29b40c8 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf.go
@@ -43,7 +43,7 @@
}
// Up performs one-time setup for this executor.
-func (n *PairWithRestriction) Up(ctx context.Context) error {
+func (n *PairWithRestriction) Up(_ context.Context) error {
fn := (*graph.SplittableDoFn)(n.Fn).CreateInitialRestrictionFn()
var err error
if n.inv, err = newCreateInitialRestrictionInvoker(fn); err != nil {
@@ -88,7 +88,7 @@
}
// Down currently does nothing.
-func (n *PairWithRestriction) Down(ctx context.Context) error {
+func (n *PairWithRestriction) Down(_ context.Context) error {
return nil
}
@@ -117,7 +117,7 @@
}
// Up performs one-time setup for this executor.
-func (n *SplitAndSizeRestrictions) Up(ctx context.Context) error {
+func (n *SplitAndSizeRestrictions) Up(_ context.Context) error {
fn := (*graph.SplittableDoFn)(n.Fn).SplitRestrictionFn()
var err error
if n.splitInv, err = newSplitRestrictionInvoker(fn); err != nil {
@@ -200,7 +200,7 @@
}
// Down currently does nothing.
-func (n *SplitAndSizeRestrictions) Down(ctx context.Context) error {
+func (n *SplitAndSizeRestrictions) Down(_ context.Context) error {
return nil
}
@@ -215,14 +215,27 @@
// changes to support the SDF's method signatures and the expected structure
// of the FullValue being received.
type ProcessSizedElementsAndRestrictions struct {
- PDo *ParDo
+ PDo *ParDo
+ TfId string // Transform ID. Needed for splitting.
+ ctInv *ctInvoker
+ sizeInv *rsInvoker
- inv *ctInvoker
+ // SU is a buffered channel for indicating when this unit is splittable.
+ // When this unit is processing an element, it sends a SplittableUnit
+ // interface through the channel. That interface can be received on other
+ // threads and used to perform splitting or other related operation.
+ //
+ // This channel should be received on in a non-blocking manner, to avoid
+ // hanging if no element is processing.
+ //
+ // Receiving the SplittableUnit prevents the current element from finishing
+ // processing, so the element does not unexpectedly change during a split.
+ // Therefore, receivers of the SplittableUnit must send it back through the
+ // channel once finished with it, or it will block indefinitely.
+ SU chan SplittableUnit
- // Rt allows this unit to send out restriction trackers being processed.
- // Receivers of the tracker do not own it, and must send it back through the
- // same channel once finished with it.
- Rt chan sdf.RTracker
+ elm *FullValue // Currently processing element.
+ rt sdf.RTracker // Currently processing element's restriction tracker.
}
// ID calls the ParDo's ID method.
@@ -234,10 +247,14 @@
func (n *ProcessSizedElementsAndRestrictions) Up(ctx context.Context) error {
fn := (*graph.SplittableDoFn)(n.PDo.Fn).CreateTrackerFn()
var err error
- if n.inv, err = newCreateTrackerInvoker(fn); err != nil {
+ if n.ctInv, err = newCreateTrackerInvoker(fn); err != nil {
return errors.WithContextf(err, "%v", n)
}
- n.Rt = make(chan sdf.RTracker, 1)
+ fn = (*graph.SplittableDoFn)(n.PDo.Fn).RestrictionSizeFn()
+ if n.sizeInv, err = newRestrictionSizeInvoker(fn); err != nil {
+ return errors.WithContextf(err, "%v", n)
+ }
+ n.SU = make(chan SplittableUnit, 1)
return n.PDo.Up(ctx)
}
@@ -268,15 +285,22 @@
// and processes each element using the underlying ParDo and adding the
// restriction tracker to the normal invocation. Sizing information is present
// but currently ignored. Output is forwarded to the underlying ParDo's outputs.
-func (n *ProcessSizedElementsAndRestrictions) ProcessElement(ctx context.Context, elm *FullValue, values ...ReStream) error {
+func (n *ProcessSizedElementsAndRestrictions) ProcessElement(_ context.Context, elm *FullValue, values ...ReStream) error {
if n.PDo.status != Active {
err := errors.Errorf("invalid status %v, want Active", n.PDo.status)
return errors.WithContextf(err, "%v", n)
}
rest := elm.Elm.(*FullValue).Elm2
- rt := n.inv.Invoke(rest)
- n.Rt <- rt
+ rt := n.ctInv.Invoke(rest)
+
+ n.rt = rt
+ n.elm = elm
+ n.SU <- n
+ defer func() {
+ <-n.SU
+ }()
+
mainIn := &MainInput{
Values: values,
RTracker: rt,
@@ -303,20 +327,18 @@
}
}
- err := n.PDo.processMainInput(mainIn)
- <-n.Rt
- return err
+ return n.PDo.processMainInput(mainIn)
}
// FinishBundle resets the invokers and then calls the ParDo's FinishBundle method.
func (n *ProcessSizedElementsAndRestrictions) FinishBundle(ctx context.Context) error {
- n.inv.Reset()
+ n.ctInv.Reset()
+ n.sizeInv.Reset()
return n.PDo.FinishBundle(ctx)
}
-// Down closes open channels and calls the ParDo's Down method.
+// Down calls the ParDo's Down method.
func (n *ProcessSizedElementsAndRestrictions) Down(ctx context.Context) error {
- close(n.Rt)
return n.PDo.Down(ctx)
}
@@ -325,6 +347,104 @@
return fmt.Sprintf("SDF.ProcessSizedElementsAndRestrictions[%v] UID:%v Out:%v", path.Base(n.PDo.Fn.Name()), n.PDo.ID(), IDs(n.PDo.Out...))
}
+// SplittableUnit is an interface that defines sub-element splitting operations
+// for a unit, and provides access to them on other threads.
+type SplittableUnit interface {
+ // Split performs a split on a fraction of a currently processing element
+ // and returns the primary and residual elements resulting from it, or an
+ // error if the split failed.
+ Split(fraction float64) (primary, residual *FullValue, err error)
+
+ // GetProgress returns the fraction of progress the current element has
+ // made in processing. (ex. 0.0 means no progress, and 1.0 means fully
+ // processed.)
+ GetProgress() float64
+
+ // GetTransformId returns the transform ID of the splittable unit.
+ GetTransformId() string
+
+ // GetInputId returns the local input ID of the input that the element being
+ // split was received from.
+ GetInputId() string
+}
+
+// Split splits the currently processing element using its restriction tracker.
+// Then it returns an element for primary and residual, following the expected
+// input structure to this unit, including updating the size of the split
+// elements.
+func (n *ProcessSizedElementsAndRestrictions) Split(f float64) (*FullValue, *FullValue, error) {
+ addContext := func(err error) error {
+ return errors.WithContext(err, "Attempting split in ProcessSizedElementsAndRestrictions")
+ }
+
+ // Check that the restriction tracker is in a state where it can be split.
+ if n.rt == nil {
+ return nil, nil, addContext(errors.New("Restriction tracker missing."))
+ }
+ if err := n.rt.GetError(); err != nil {
+ return nil, nil, addContext(err)
+ }
+ if n.rt.IsDone() { // Not an error, but not splittable.
+ return nil, nil, nil
+ }
+
+ p, r, err := n.rt.TrySplit(f)
+ if err != nil {
+ return nil, nil, addContext(err)
+ }
+ if r == nil { // If r is nil then the split failed/returned an empty residual.
+ return nil, nil, nil
+ }
+
+ var pfv, rfv *FullValue
+ var pSize, rSize float64
+ elm := n.elm.Elm.(*FullValue).Elm
+ if fv, ok := elm.(*FullValue); ok {
+ pSize = n.sizeInv.Invoke(fv, p)
+ rSize = n.sizeInv.Invoke(fv, r)
+ } else {
+ fv := &FullValue{Elm: elm}
+ pSize = n.sizeInv.Invoke(fv, p)
+ rSize = n.sizeInv.Invoke(fv, r)
+ }
+ pfv = &FullValue{
+ Elm: &FullValue{
+ Elm: elm,
+ Elm2: p,
+ },
+ Elm2: pSize,
+ Timestamp: n.elm.Timestamp,
+ Windows: n.elm.Windows,
+ }
+ rfv = &FullValue{
+ Elm: &FullValue{
+ Elm: elm,
+ Elm2: r,
+ },
+ Elm2: rSize,
+ Timestamp: n.elm.Timestamp,
+ Windows: n.elm.Windows,
+ }
+ return pfv, rfv, nil
+}
+
+// GetProgress returns the current restriction tracker's progress as a fraction.
+func (n *ProcessSizedElementsAndRestrictions) GetProgress() float64 {
+ d, r := n.rt.GetProgress()
+ return d / (d + r)
+}
+
+// GetTransformId returns this transform's transform ID.
+func (n *ProcessSizedElementsAndRestrictions) GetTransformId() string {
+ return n.TfId
+}
+
+// GetInputId returns the main input ID, since main input elements are being
+// split.
+func (n *ProcessSizedElementsAndRestrictions) GetInputId() string {
+ return indexToInputId(0)
+}
+
// SdfFallback is an executor used when an SDF isn't expanded into steps by the
// runner, indicating that the runner doesn't support splitting. It executes all
// the SDF steps together in one unit.
@@ -370,7 +490,7 @@
// restrictions, and then creating restriction trackers and processing each
// restriction with the underlying ParDo. This executor skips the sizing step
// because sizing information is unnecessary for unexpanded SDFs.
-func (n *SdfFallback) ProcessElement(ctx context.Context, elm *FullValue, values ...ReStream) error {
+func (n *SdfFallback) ProcessElement(_ context.Context, elm *FullValue, values ...ReStream) error {
if n.PDo.status != Active {
err := errors.Errorf("invalid status %v, want Active", n.PDo.status)
return errors.WithContextf(err, "%v", n)
diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go b/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go
index 5e3abd6..7dbe3e8 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go
@@ -16,9 +16,10 @@
package exec
import (
+ "testing"
+
"github.com/apache/beam/sdks/go/pkg/beam/core/graph"
"github.com/google/go-cmp/cmp"
- "testing"
)
// TestInvokes runs tests on each SDF method invoker, using the SDFs defined
diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go b/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go
index e746f85..fb85620 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go
@@ -16,11 +16,13 @@
package exec
import (
+ "context"
+ "testing"
+
"github.com/apache/beam/sdks/go/pkg/beam/core/graph"
"github.com/apache/beam/sdks/go/pkg/beam/core/graph/window"
"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
"github.com/google/go-cmp/cmp"
- "testing"
)
// testTimestamp is a constant used to check that timestamps are retained.
@@ -386,3 +388,145 @@
}
})
}
+
+// TestAsSplittableUnit tests ProcessSizedElementsAndRestrictions' implementation
+// of the SplittableUnit interface.
+func TestAsSplittableUnit(t *testing.T) {
+ dfn, err := graph.NewDoFn(&VetSdf{}, graph.NumMainInputs(graph.MainSingle))
+ if err != nil {
+ t.Fatalf("invalid function: %v", err)
+ }
+ kvdfn, err := graph.NewDoFn(&VetKvSdf{}, graph.NumMainInputs(graph.MainKv))
+ if err != nil {
+ t.Fatalf("invalid function: %v", err)
+ }
+
+ // Test that Split returns properly structured results and calls Split on
+ // the restriction tracker.
+ t.Run("Split", func(t *testing.T) {
+ tests := []struct {
+ name string
+ fn *graph.DoFn
+ in FullValue
+ wantPrimary FullValue
+ wantResidual FullValue
+ }{
+ {
+ name: "SingleElem",
+ fn: dfn,
+ in: FullValue{
+ Elm: &FullValue{
+ Elm: 1,
+ Elm2: &VetRestriction{ID: "Sdf"},
+ },
+ Elm2: 1.0,
+ Timestamp: testTimestamp,
+ Windows: testWindows,
+ },
+ wantPrimary: FullValue{
+ Elm: &FullValue{
+ Elm: 1,
+ Elm2: &VetRestriction{ID: "Sdf.1", RestSize: true, Val: 1},
+ },
+ Elm2: 1.0,
+ Timestamp: testTimestamp,
+ Windows: testWindows,
+ },
+ wantResidual: FullValue{
+ Elm: &FullValue{
+ Elm: 1,
+ Elm2: &VetRestriction{ID: "Sdf.2", RestSize: true, Val: 1},
+ },
+ Elm2: 1.0,
+ Timestamp: testTimestamp,
+ Windows: testWindows,
+ },
+ },
+ {
+ name: "KvElem",
+ fn: kvdfn,
+ in: FullValue{
+ Elm: &FullValue{
+ Elm: &FullValue{
+ Elm: 1,
+ Elm2: 2,
+ },
+ Elm2: &VetRestriction{ID: "KvSdf"},
+ },
+ Elm2: 3.0,
+ Timestamp: testTimestamp,
+ Windows: testWindows,
+ },
+ wantPrimary: FullValue{
+ Elm: &FullValue{
+ Elm: &FullValue{
+ Elm: 1,
+ Elm2: 2,
+ },
+ Elm2: &VetRestriction{ID: "KvSdf.1", RestSize: true, Key: 1, Val: 2},
+ },
+ Elm2: 3.0,
+ Timestamp: testTimestamp,
+ Windows: testWindows,
+ },
+ wantResidual: FullValue{
+ Elm: &FullValue{
+ Elm: &FullValue{
+ Elm: 1,
+ Elm2: 2,
+ },
+ Elm2: &VetRestriction{ID: "KvSdf.2", RestSize: true, Key: 1, Val: 2},
+ },
+ Elm2: 3.0,
+ Timestamp: testTimestamp,
+ Windows: testWindows,
+ },
+ },
+ }
+ for _, test := range tests {
+ test := test
+ t.Run(test.name, func(t *testing.T) {
+ // Setup, create transforms, inputs, and desired outputs.
+ n := &ParDo{UID: 1, Fn: test.fn, Out: []Node{}}
+ node := &ProcessSizedElementsAndRestrictions{PDo: n}
+ node.rt = &SplittableUnitRTracker{
+ VetRTracker: VetRTracker{Rest: test.in.Elm.(*FullValue).Elm2.(*VetRestriction)},
+ }
+ node.elm = &test.in
+
+ // Call from SplittableUnit and check results.
+ su := SplittableUnit(node)
+ frac := 0.5
+ if err := node.Up(context.Background()); err != nil {
+ t.Fatalf("ProcessSizedElementsAndRestrictions.Up() failed: %v", err)
+ }
+ gotPrimary, gotResidual, err := su.Split(frac)
+ if err != nil {
+ t.Fatalf("SplittableUnit.Split(%v) failed: %v", frac, err)
+ }
+ if diff := cmp.Diff(gotPrimary, &test.wantPrimary); diff != "" {
+ t.Errorf("SplittableUnit.Split(%v) has incorrect primary: %v", frac, diff)
+ }
+ if diff := cmp.Diff(gotResidual, &test.wantResidual); diff != "" {
+ t.Errorf("SplittableUnit.Split(%v) has incorrect residual: %v", frac, diff)
+ }
+ })
+ }
+ })
+}
+
+// SplittableUnitRTracker is a VetRTracker with some added behavior needed for
+// TestAsSplittableUnit.
+type SplittableUnitRTracker struct {
+ VetRTracker
+}
+
+func (rt *SplittableUnitRTracker) IsDone() bool { return false }
+
+func (rt *SplittableUnitRTracker) TrySplit(_ float64) (interface{}, interface{}, error) {
+ rest1 := rt.Rest.copy()
+ rest1.ID += ".1"
+ rest2 := rt.Rest.copy()
+ rest2.ID += ".2"
+ return &rest1, &rest2, nil
+}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/translate.go b/sdks/go/pkg/beam/core/runtime/exec/translate.go
index caa5a22..fbbdab3 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/translate.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/translate.go
@@ -418,7 +418,7 @@
}
u = n
if urn == urnProcessSizedElementsAndRestrictions {
- u = &ProcessSizedElementsAndRestrictions{PDo: n}
+ u = &ProcessSizedElementsAndRestrictions{PDo: n, TfId: id.to}
} else if dofn.IsSplittable() {
u = &SdfFallback{PDo: n}
}
@@ -580,7 +580,7 @@
var unordered []string
for key := range m {
- if i, err := strconv.Atoi(strings.TrimPrefix(key, "i")); strings.HasPrefix(key, "i") && err == nil {
+ if i, err := inputIdToIndex(key); err == nil {
if i < len(m) {
ordered[i] = key
continue
@@ -605,6 +605,24 @@
return ret
}
+// inputIdToIndex converts a local input ID for a transform into an index. Use
+// this to avoid relying on format details for input IDs.
+//
+// Currently, expects IDs in the format "iN" where N is the index. If the ID is
+// in an invalid form, returns an error.
+func inputIdToIndex(id string) (int, error) {
+ if !strings.HasPrefix(id, "i") {
+ return 0, errors.New("invalid input ID format")
+ }
+ return strconv.Atoi(strings.TrimPrefix(id, "i"))
+}
+
+// inputIdToIndex converts an index into a local input ID for a transform. Use
+// this to avoid relying on format details for input IDs.
+func indexToInputId(i int) string {
+ return "i" + strconv.Itoa(i)
+}
+
func unmarshalPort(data []byte) (Port, string, error) {
var port fnpb.RemoteGrpcPort
if err := proto.Unmarshal(data, &port); err != nil {
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go
new file mode 100644
index 0000000..ac724b3
--- /dev/null
+++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go
@@ -0,0 +1,269 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 schema contains utility functions for relating Go types and Beam Schemas.
+//
+// Not all Go types can be converted to schemas. This is Go is more expressive than
+// Beam schemas. Just as not all Go types can be serialized, similarly,
+// not all Beam Schemas will have a conversion to Go types, until the correct
+// mechanism exists in the SDK to handle them.
+//
+// While efforts will be made to have conversions be reversable, this will not
+// be possible in all instances. Eg. Go arrays as fields will be converted to
+// Beam Arrays, but a Beam Array type will map by default to a Go slice.
+package schema
+
+import (
+ "fmt"
+ "reflect"
+ "strings"
+
+ "github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
+ "github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
+ pipepb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
+)
+
+// FromType returns a Beam Schema of the passed in type.
+// Returns an error if the type cannot be converted to a Schema.
+func FromType(ot reflect.Type) (*pipepb.Schema, error) {
+ t := ot // keep the original type for errors.
+ // The top level schema for a pointer to struct and the struct is the same.
+ if t.Kind() == reflect.Ptr {
+ t = t.Elem()
+ }
+ if t.Kind() != reflect.Struct {
+ return nil, errors.Errorf("cannot convert %v to schema. FromType only converts structs to schemas", ot)
+ }
+ return structToSchema(t)
+}
+
+func structToSchema(t reflect.Type) (*pipepb.Schema, error) {
+ fields := make([]*pipepb.Field, 0, t.NumField())
+ for i := 0; i < t.NumField(); i++ {
+ f, err := structFieldToField(t.Field(i))
+ if err != nil {
+ return nil, errors.Wrapf(err, "cannot convert field %v to schema", t.Field(i).Name)
+ }
+ fields = append(fields, f)
+ }
+ return &pipepb.Schema{
+ Fields: fields,
+ }, nil
+}
+
+func structFieldToField(sf reflect.StructField) (*pipepb.Field, error) {
+ name := sf.Name
+ if tag := sf.Tag.Get("beam"); tag != "" {
+ name, _ = parseTag(tag)
+ }
+ ftype, err := reflectTypeToFieldType(sf.Type)
+ if err != nil {
+ return nil, err
+ }
+ return &pipepb.Field{
+ Name: name,
+ Type: ftype,
+ }, nil
+}
+
+func reflectTypeToFieldType(ot reflect.Type) (*pipepb.FieldType, error) {
+ var isPtr bool
+ t := ot
+ if t.Kind() == reflect.Ptr {
+ isPtr = true
+ t = t.Elem()
+ }
+ switch t.Kind() {
+ case reflect.Map:
+ kt, err := reflectTypeToFieldType(t.Key())
+ if err != nil {
+ return nil, errors.Wrapf(err, "unable to convert key of %v to schema field", ot)
+ }
+ vt, err := reflectTypeToFieldType(t.Elem())
+ if err != nil {
+ return nil, errors.Wrapf(err, "unable to convert value of %v to schema field", ot)
+ }
+ return &pipepb.FieldType{
+ Nullable: isPtr,
+ TypeInfo: &pipepb.FieldType_MapType{
+ MapType: &pipepb.MapType{
+ KeyType: kt,
+ ValueType: vt,
+ },
+ },
+ }, nil
+ case reflect.Struct:
+ sch, err := structToSchema(t)
+ if err != nil {
+ return nil, errors.Wrapf(err, "unable to convert %v to schema field", ot)
+ }
+ return &pipepb.FieldType{
+ Nullable: isPtr,
+ TypeInfo: &pipepb.FieldType_RowType{
+ RowType: &pipepb.RowType{
+ Schema: sch,
+ },
+ },
+ }, nil
+ case reflect.Slice, reflect.Array:
+ // Special handling for []byte
+ if t == reflectx.ByteSlice {
+ return &pipepb.FieldType{
+ Nullable: isPtr,
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_BYTES,
+ },
+ }, nil
+ }
+ vt, err := reflectTypeToFieldType(t.Elem())
+ if err != nil {
+ return nil, errors.Wrapf(err, "unable to convert element type of %v to schema field", ot)
+ }
+ return &pipepb.FieldType{
+ Nullable: isPtr,
+ TypeInfo: &pipepb.FieldType_ArrayType{
+ ArrayType: &pipepb.ArrayType{
+ ElementType: vt,
+ },
+ },
+ }, nil
+ case reflect.Interface, reflect.Chan, reflect.UnsafePointer, reflect.Complex128, reflect.Complex64:
+ return nil, errors.Errorf("unable to convert unsupported type %v to schema", ot)
+ default: // must be an atomic type
+ if enum, ok := reflectTypeToAtomicTypeMap[t.Kind()]; ok {
+ return &pipepb.FieldType{
+ Nullable: isPtr,
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: enum,
+ },
+ }, nil
+ }
+ return nil, errors.Errorf("unable to map %v to pipepb.AtomicType", t)
+ }
+}
+
+var reflectTypeToAtomicTypeMap = map[reflect.Kind]pipepb.AtomicType{
+ reflect.Uint8: pipepb.AtomicType_BYTE,
+ reflect.Int16: pipepb.AtomicType_INT16,
+ reflect.Int32: pipepb.AtomicType_INT32,
+ reflect.Int64: pipepb.AtomicType_INT64,
+ reflect.Int: pipepb.AtomicType_INT64,
+ reflect.Float32: pipepb.AtomicType_FLOAT,
+ reflect.Float64: pipepb.AtomicType_DOUBLE,
+ reflect.String: pipepb.AtomicType_STRING,
+ reflect.Bool: pipepb.AtomicType_BOOLEAN,
+}
+
+// ToType returns a Go type of the passed in Schema.
+// Types returned by ToType are always of Struct kind.
+// Returns an error if the Schema cannot be converted to a type.
+func ToType(s *pipepb.Schema) (reflect.Type, error) {
+ fields := make([]reflect.StructField, 0, len(s.GetFields()))
+ for _, sf := range s.GetFields() {
+ rf, err := fieldToStructField(sf)
+ if err != nil {
+ return nil, errors.Wrapf(err, "cannot convert schema field %v to field", sf.GetName())
+ }
+ fields = append(fields, rf)
+ }
+ return reflect.StructOf(fields), nil
+}
+
+func fieldToStructField(sf *pipepb.Field) (reflect.StructField, error) {
+ name := sf.GetName()
+ rt, err := fieldTypeToReflectType(sf.GetType())
+ if err != nil {
+ return reflect.StructField{}, err
+ }
+ return reflect.StructField{
+ Name: strings.ToUpper(name[:1]) + name[1:], // Go field name must be capitalized for export and encoding.
+ Type: rt,
+ Tag: reflect.StructTag(fmt.Sprintf("beam:\"%s\"", name)),
+ }, nil
+}
+
+var atomicTypeToReflectType = map[pipepb.AtomicType]reflect.Type{
+ pipepb.AtomicType_BYTE: reflectx.Uint8,
+ pipepb.AtomicType_INT16: reflectx.Int16,
+ pipepb.AtomicType_INT32: reflectx.Int32,
+ pipepb.AtomicType_INT64: reflectx.Int64,
+ pipepb.AtomicType_FLOAT: reflectx.Float32,
+ pipepb.AtomicType_DOUBLE: reflectx.Float64,
+ pipepb.AtomicType_STRING: reflectx.String,
+ pipepb.AtomicType_BOOLEAN: reflectx.Bool,
+ pipepb.AtomicType_BYTES: reflectx.ByteSlice,
+}
+
+func fieldTypeToReflectType(sft *pipepb.FieldType) (reflect.Type, error) {
+ var t reflect.Type
+ switch sft.GetTypeInfo().(type) {
+ case *pipepb.FieldType_AtomicType:
+ var ok bool
+ if t, ok = atomicTypeToReflectType[sft.GetAtomicType()]; !ok {
+ return nil, errors.Errorf("unknown atomic type: %v", sft.GetAtomicType())
+ }
+ case *pipepb.FieldType_ArrayType:
+ rt, err := fieldTypeToReflectType(sft.GetArrayType().GetElementType())
+ if err != nil {
+ return nil, errors.Wrap(err, "unable to convert array element type")
+ }
+ t = reflect.SliceOf(rt)
+ case *pipepb.FieldType_MapType:
+ kt, err := fieldTypeToReflectType(sft.GetMapType().GetKeyType())
+ if err != nil {
+ return nil, errors.Wrap(err, "unable to convert map key type")
+ }
+ vt, err := fieldTypeToReflectType(sft.GetMapType().GetValueType())
+ if err != nil {
+ return nil, errors.Wrap(err, "unable to convert map value type")
+ }
+ t = reflect.MapOf(kt, vt) // Panics for invalid map keys (slices/iterables)
+ case *pipepb.FieldType_RowType:
+ rt, err := ToType(sft.GetRowType().GetSchema())
+ if err != nil {
+ return nil, errors.Wrapf(err, "unable to convert row type: %v", sft.GetRowType().GetSchema().GetId())
+ }
+ t = rt
+ // case *pipepb.FieldType_IterableType:
+ // TODO(BEAM-9615): handle IterableTypes.
+
+ // case *pipepb.FieldType_LogicalType:
+ // TODO(BEAM-9615): handle LogicalTypes types.
+
+ // Logical Types are for things that have more specialized user representation already, or
+ // things like Time or protocol buffers.
+ // They would be encoded with the schema encoding.
+
+ default:
+ return nil, errors.Errorf("unknown fieldtype: %T", sft.GetTypeInfo())
+ }
+ if sft.GetNullable() {
+ return reflect.PtrTo(t), nil
+ }
+ return t, nil
+}
+
+// parseTag splits a struct field's beam tag into its name and
+// comma-separated options.
+func parseTag(tag string) (string, options) {
+ if idx := strings.Index(tag, ","); idx != -1 {
+ return tag[:idx], options(tag[idx+1:])
+ }
+ return tag, options("")
+}
+
+type options string
+
+// TODO(BEAM-9615): implement looking up specific options from the tags.
diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go
new file mode 100644
index 0000000..f6b5ff7
--- /dev/null
+++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go
@@ -0,0 +1,182 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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 schema
+
+import (
+ "fmt"
+ "reflect"
+ "testing"
+
+ pipepb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
+ "github.com/google/go-cmp/cmp"
+)
+
+func TestSchemaConversion(t *testing.T) {
+ tests := []struct {
+ st *pipepb.Schema
+ rt reflect.Type
+ }{
+ {
+ st: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ &pipepb.Field{
+ Name: "firstField",
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_INT32,
+ },
+ },
+ },
+ },
+ },
+ rt: reflect.TypeOf(struct {
+ FirstField int32 `beam:"firstField"`
+ }{}),
+ }, {
+ st: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ &pipepb.Field{
+ Name: "stringField",
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_STRING,
+ },
+ },
+ },
+ &pipepb.Field{
+ Name: "intPtrField",
+ Type: &pipepb.FieldType{
+ Nullable: true,
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_INT32,
+ },
+ },
+ },
+ },
+ },
+ rt: reflect.TypeOf(struct {
+ StringField string `beam:"stringField"`
+ IntPtrField *int32 `beam:"intPtrField"`
+ }{}),
+ }, {
+ st: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ &pipepb.Field{
+ Name: "cypher",
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_MapType{
+ MapType: &pipepb.MapType{
+ KeyType: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_BOOLEAN,
+ },
+ },
+ ValueType: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_FLOAT,
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ rt: reflect.TypeOf(struct {
+ Cypher map[bool]float32 `beam:"cypher"`
+ }{}),
+ }, {
+ st: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ &pipepb.Field{
+ Name: "wrapper",
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_RowType{
+ RowType: &pipepb.RowType{
+ Schema: &pipepb.Schema{
+ Fields: []*pipepb.Field{{
+ Name: "threshold",
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_DOUBLE,
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ rt: reflect.TypeOf(struct {
+ Wrapper struct {
+ Threshold float64 `beam:"threshold"`
+ } `beam:"wrapper"`
+ }{}),
+ }, {
+ st: &pipepb.Schema{
+ Fields: []*pipepb.Field{
+ &pipepb.Field{
+ Name: "payloads",
+ Type: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_ArrayType{
+ ArrayType: &pipepb.ArrayType{
+ ElementType: &pipepb.FieldType{
+ TypeInfo: &pipepb.FieldType_AtomicType{
+ AtomicType: pipepb.AtomicType_BYTES,
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ rt: reflect.TypeOf(struct {
+ Payloads [][]byte `beam:"payloads"`
+ }{}),
+ },
+ }
+
+ for _, test := range tests {
+ test := test
+ t.Run(fmt.Sprintf("%v", test.rt), func(t *testing.T) {
+ {
+ got, err := ToType(test.st)
+ if err != nil {
+ t.Fatalf("error ToType(%v) = %v", test.st, err)
+ }
+
+ if d := cmp.Diff(reflect.New(test.rt).Elem().Interface(), reflect.New(got).Elem().Interface()); d != "" {
+ t.Errorf("diff (-want, +got): %v", d)
+ }
+ }
+ {
+ got, err := FromType(test.rt)
+ if err != nil {
+ t.Fatalf("error FromType(%v) = %v", test.rt, err)
+ }
+
+ if d := cmp.Diff(test.st, got); d != "" {
+ t.Errorf("diff (-want, +got): %v", d)
+ }
+
+ }
+ })
+ }
+}
diff --git a/sdks/go/pkg/beam/core/runtime/harness/harness.go b/sdks/go/pkg/beam/core/runtime/harness/harness.go
index 4437529..3aa8d43 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/harness.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/harness.go
@@ -315,7 +315,7 @@
if ds == nil {
return fail(ctx, instID, "failed to split: desired splits for root of %v was empty.", ref)
}
- split, err := plan.Split(exec.SplitPoints{
+ sr, err := plan.Split(exec.SplitPoints{
Splits: ds.GetAllowedSplitPoints(),
Frac: ds.GetFractionOfRemainder(),
BufSize: ds.GetEstimatedInputElements(),
@@ -329,12 +329,22 @@
InstructionId: string(instID),
Response: &fnpb.InstructionResponse_ProcessBundleSplit{
ProcessBundleSplit: &fnpb.ProcessBundleSplitResponse{
- ChannelSplits: []*fnpb.ProcessBundleSplitResponse_ChannelSplit{
- &fnpb.ProcessBundleSplitResponse_ChannelSplit{
- LastPrimaryElement: split - 1,
- FirstResidualElement: split,
+ PrimaryRoots: []*fnpb.BundleApplication{{
+ TransformId: sr.TId,
+ InputId: sr.InId,
+ Element: sr.PS,
+ }},
+ ResidualRoots: []*fnpb.DelayedBundleApplication{{
+ Application: &fnpb.BundleApplication{
+ TransformId: sr.TId,
+ InputId: sr.InId,
+ Element: sr.RS,
},
- },
+ }},
+ ChannelSplits: []*fnpb.ProcessBundleSplitResponse_ChannelSplit{{
+ LastPrimaryElement: sr.PI,
+ FirstResidualElement: sr.RI,
+ }},
},
},
}
diff --git a/sdks/go/pkg/beam/core/sdf/sdf.go b/sdks/go/pkg/beam/core/sdf/sdf.go
index 3d663d5..1de3f5c 100644
--- a/sdks/go/pkg/beam/core/sdf/sdf.go
+++ b/sdks/go/pkg/beam/core/sdf/sdf.go
@@ -64,8 +64,9 @@
//
// This method modifies the underlying restriction in the RTracker to reflect the primary. It
// then returns a copy of the newly modified restriction as a primary, and returns a new
- // restriction for the residual. If the split would produce an empty residual (i.e. the only
- // split point is the end of the restriction), then the returned residual is nil.
+ // restriction for the residual. If the split would produce an empty residual (either because
+ // the only split point is the end of the restriction, or the split failed for some recoverable
+ // reason), then this function returns nil as the residual.
//
// If an error is returned, some catastrophic failure occurred and the entire bundle will fail.
TrySplit(fraction float64) (primary, residual interface{}, err error)
@@ -77,8 +78,8 @@
// IsDone returns a boolean indicating whether all blocks inside the restriction have been
// claimed. This method is called by the SDK Harness to validate that a splittable DoFn has
- // correctly processed all work in a restriction before finishing. If this method returns false
- // then GetError is expected to return a non-nil error.
+ // correctly processed all work in a restriction before finishing. If this method still returns
+ // false after processing, then GetError is expected to return a non-nil error.
IsDone() bool
// GetRestriction returns the restriction this tracker is tracking, or nil if the restriction
diff --git a/sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange.go b/sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange.go
index f01ba50..e140a7e 100644
--- a/sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange.go
+++ b/sdks/go/pkg/beam/io/rtrackers/offsetrange/offsetrange.go
@@ -20,16 +20,39 @@
package offsetrange
import (
+ "bytes"
+ "encoding/binary"
"errors"
"math"
"reflect"
- "github.com/apache/beam/sdks/go/pkg/beam"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
+ "github.com/apache/beam/sdks/go/pkg/beam/core/runtime"
)
func init() {
- beam.RegisterType(reflect.TypeOf((*Tracker)(nil)))
- beam.RegisterType(reflect.TypeOf((*Restriction)(nil)))
+ runtime.RegisterType(reflect.TypeOf((*Tracker)(nil)))
+ runtime.RegisterType(reflect.TypeOf((*Restriction)(nil)).Elem())
+ runtime.RegisterFunction(restEnc)
+ runtime.RegisterFunction(restDec)
+ coder.RegisterCoder(reflect.TypeOf((*Restriction)(nil)).Elem(), restEnc, restDec)
+}
+
+func restEnc(in Restriction) ([]byte, error) {
+ buf := new(bytes.Buffer)
+ if err := binary.Write(buf, binary.BigEndian, in); err != nil {
+ return nil, err
+ }
+ return buf.Bytes(), nil
+}
+
+func restDec(in []byte) (Restriction, error) {
+ buf := bytes.NewBuffer(in)
+ rest := Restriction{}
+ if err := binary.Read(buf, binary.BigEndian, &rest); err != nil {
+ return rest, err
+ }
+ return rest, nil
}
// Restriction is an offset range restriction, which represents a range of
@@ -44,10 +67,10 @@
//
// Num should be greater than 0. Otherwise there is no way to split the
// restriction and this function will return the original restriction.
-func (r *Restriction) EvenSplits(num int64) (splits []Restriction) {
+func (r Restriction) EvenSplits(num int64) (splits []Restriction) {
if num <= 1 {
// Don't split, just return original restriction.
- return append(splits, *r)
+ return append(splits, r)
}
offset := r.Start
@@ -67,7 +90,7 @@
}
// Size returns the restriction's size as the difference between Start and End.
-func (r *Restriction) Size() float64 {
+func (r Restriction) Size() float64 {
return float64(r.End - r.Start)
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/expansion/ExternalTransformRegistrar.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/expansion/ExternalTransformRegistrar.java
index bbdd3a5..aa5288c 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/expansion/ExternalTransformRegistrar.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/expansion/ExternalTransformRegistrar.java
@@ -17,10 +17,14 @@
*/
package org.apache.beam.sdk.expansion;
+import java.lang.reflect.Constructor;
import java.util.Map;
+import java.util.Map.Entry;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
+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;
/**
* A registrar which contains a mapping from URNs to available {@link ExternalTransformBuilder}s.
@@ -29,6 +33,42 @@
@Experimental(Kind.PORTABILITY)
public interface ExternalTransformRegistrar {
- /** A mapping from URN to an {@link ExternalTransformBuilder} class. */
- Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders();
+ /**
+ * A mapping from URN to an {@link ExternalTransformBuilder} class.
+ *
+ * @deprecated Prefer implementing 'knownBuilderInstances'. This method will be removed in a
+ * future version of Beam.
+ */
+ @Deprecated
+ default Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders() {
+ return ImmutableMap.<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>>builder()
+ .build();
+ }
+
+ /** A mapping from URN to an {@link ExternalTransformBuilder} instance. */
+ default Map<String, ExternalTransformBuilder<?, ?, ?>> knownBuilderInstances() {
+ ImmutableMap.Builder builder = ImmutableMap.<String, ExternalTransformBuilder>builder();
+ Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders = knownBuilders();
+ for (Entry<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilder :
+ knownBuilders.entrySet()) {
+ Preconditions.checkState(
+ ExternalTransformBuilder.class.isAssignableFrom(knownBuilder.getValue()),
+ "Provided identifier %s is not an ExternalTransformBuilder.",
+ knownBuilder.getValue().getName());
+ try {
+ Constructor<? extends ExternalTransformBuilder> constructor =
+ knownBuilder.getValue().getDeclaredConstructor();
+
+ constructor.setAccessible(true);
+ builder.put(knownBuilder.getKey(), constructor.newInstance());
+
+ } catch (RuntimeException e) {
+ throw e;
+ } catch (Exception e) {
+ throw new RuntimeException(
+ "Unable to instantiate ExternalTransformBuilder from constructor.");
+ }
+ }
+ return builder.build();
+ }
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/logicaltypes/DateTime.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/logicaltypes/DateTime.java
new file mode 100644
index 0000000..148c726
--- /dev/null
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/logicaltypes/DateTime.java
@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.schemas.logicaltypes;
+
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.values.Row;
+
+/**
+ * A datetime without a time-zone.
+ *
+ * <p>It cannot represent an instant on the time-line without additional information such as an
+ * offset or time-zone.
+ *
+ * <p>Its input type is a {@link LocalDateTime}, and base type is a {@link Row} containing Date
+ * field and Time field. Date field is the same as the base type of {@link Date}, which is a Long
+ * that represents incrementing count of days where day 0 is 1970-01-01 (ISO). Time field is the
+ * same as the base type of {@link Time}, which is a Long that represents a count of time in
+ * nanoseconds.
+ */
+public class DateTime implements Schema.LogicalType<LocalDateTime, Row> {
+ public static final String DATE_FIELD_NAME = "Date";
+ public static final String TIME_FIELD_NAME = "Time";
+ public static final Schema DATETIME_SCHEMA =
+ Schema.builder().addInt64Field(DATE_FIELD_NAME).addInt64Field(TIME_FIELD_NAME).build();
+
+ @Override
+ public String getIdentifier() {
+ return "beam:logical_type:datetime:v1";
+ }
+
+ // unused
+ @Override
+ public Schema.FieldType getArgumentType() {
+ return Schema.FieldType.STRING;
+ }
+
+ // unused
+ @Override
+ public String getArgument() {
+ return "";
+ }
+
+ @Override
+ public Schema.FieldType getBaseType() {
+ return Schema.FieldType.row(DATETIME_SCHEMA);
+ }
+
+ @Override
+ public Row toBaseType(LocalDateTime input) {
+ return input == null
+ ? null
+ : Row.withSchema(DATETIME_SCHEMA)
+ .addValues(input.toLocalDate().toEpochDay(), input.toLocalTime().toNanoOfDay())
+ .build();
+ }
+
+ @Override
+ public LocalDateTime toInputType(Row base) {
+ return base == null
+ ? null
+ : LocalDateTime.of(
+ LocalDate.ofEpochDay(base.getInt64(DATE_FIELD_NAME)),
+ LocalTime.ofNanoOfDay(base.getInt64(TIME_FIELD_NAME)));
+ }
+}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/logicaltypes/SqlTypes.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/logicaltypes/SqlTypes.java
index d22b77a..ef6a68a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/logicaltypes/SqlTypes.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/logicaltypes/SqlTypes.java
@@ -18,8 +18,10 @@
package org.apache.beam.sdk.schemas.logicaltypes;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import org.apache.beam.sdk.schemas.Schema.LogicalType;
+import org.apache.beam.sdk.values.Row;
/** Beam {@link org.apache.beam.sdk.schemas.Schema.LogicalType}s corresponding to SQL data types. */
public class SqlTypes {
@@ -31,4 +33,7 @@
/** Beam LogicalType corresponding to ZetaSQL/CalciteSQL TIME type. */
public static final LogicalType<LocalTime, Long> TIME = new Time();
+
+ /** Beam LogicalType corresponding to ZetaSQL DATETIME type. */
+ public static final LogicalType<LocalDateTime, Row> DATETIME = new DateTime();
}
diff --git a/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java b/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java
index c20e8d2..09a42e3 100644
--- a/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java
+++ b/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java
@@ -22,7 +22,6 @@
import com.google.auto.service.AutoService;
import java.io.IOException;
-import java.lang.reflect.Constructor;
import java.lang.reflect.Method;
import java.util.ArrayDeque;
import java.util.Collections;
@@ -107,17 +106,21 @@
ImmutableMap.builder();
for (ExternalTransformRegistrar registrar :
ServiceLoader.load(ExternalTransformRegistrar.class)) {
- for (Map.Entry<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> entry :
- registrar.knownBuilders().entrySet()) {
+ for (Map.Entry<String, ExternalTransformBuilder<?, ?, ?>> entry :
+ registrar.knownBuilderInstances().entrySet()) {
String urn = entry.getKey();
- Class<? extends ExternalTransformBuilder<?, ?, ?>> builderClass = entry.getValue();
+ ExternalTransformBuilder builderInstance = entry.getValue();
builder.put(
urn,
spec -> {
try {
ExternalTransforms.ExternalConfigurationPayload payload =
ExternalTransforms.ExternalConfigurationPayload.parseFrom(spec.getPayload());
- return translate(payload, builderClass);
+ return builderInstance.buildExternal(
+ payloadToConfig(
+ payload,
+ (Class<? extends ExternalTransformBuilder<?, ?, ?>>)
+ builderInstance.getClass()));
} catch (Exception e) {
throw new RuntimeException(
String.format("Failed to build transform %s from spec %s", urn, spec), e);
@@ -125,21 +128,17 @@
});
}
}
+
return builder.build();
}
- private static PTransform<?, ?> translate(
+ Object payloadToConfig(
ExternalTransforms.ExternalConfigurationPayload payload,
Class<? extends ExternalTransformBuilder<?, ?, ?>> builderClass)
throws Exception {
- Preconditions.checkState(
- ExternalTransformBuilder.class.isAssignableFrom(builderClass),
- "Provided identifier %s is not an ExternalTransformBuilder.",
- builderClass.getName());
-
Object configObject = initConfiguration(builderClass);
populateConfiguration(configObject, payload);
- return buildTransform(builderClass, configObject);
+ return configObject;
}
private static Object initConfiguration(
@@ -239,28 +238,6 @@
return coderBuilder.build();
}
-
- private static PTransform<?, ?> buildTransform(
- Class<? extends ExternalTransformBuilder<?, ?, ?>> builderClass, Object configObject)
- throws Exception {
- Constructor<? extends ExternalTransformBuilder<?, ?, ?>> constructor =
- builderClass.getDeclaredConstructor();
- constructor.setAccessible(true);
- ExternalTransformBuilder<?, ?, ?> externalTransformBuilder = constructor.newInstance();
- Method buildMethod = builderClass.getMethod("buildExternal", configObject.getClass());
- buildMethod.setAccessible(true);
-
- PTransform<?, ?> transform =
- (PTransform<?, ?>)
- checkArgumentNotNull(
- buildMethod.invoke(externalTransformBuilder, configObject),
- "Invoking %s.%s(%s) returned null, violating its type.",
- builderClass.getCanonicalName(),
- "buildExternal",
- configObject);
-
- return transform;
- }
}
/**
diff --git a/sdks/java/extensions/sql/datacatalog/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datacatalog/SchemaUtils.java b/sdks/java/extensions/sql/datacatalog/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datacatalog/SchemaUtils.java
index 4604a00..42efbbd 100644
--- a/sdks/java/extensions/sql/datacatalog/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datacatalog/SchemaUtils.java
+++ b/sdks/java/extensions/sql/datacatalog/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datacatalog/SchemaUtils.java
@@ -39,7 +39,7 @@
.put("BOOL", FieldType.BOOLEAN)
.put("BYTES", FieldType.BYTES)
.put("DATE", FieldType.logicalType(SqlTypes.DATE))
- .put("DATETIME", FieldType.DATETIME)
+ .put("DATETIME", FieldType.logicalType(SqlTypes.DATETIME))
.put("DOUBLE", FieldType.DOUBLE)
.put("FLOAT", FieldType.DOUBLE)
.put("FLOAT64", FieldType.DOUBLE)
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/ScalarFunctionImpl.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/ScalarFunctionImpl.java
index 3ef4d9f..b4a9d7e 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/ScalarFunctionImpl.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/ScalarFunctionImpl.java
@@ -124,7 +124,7 @@
@Override
public RelDataType getReturnType(RelDataTypeFactory typeFactory) {
- return CalciteUtils.sqlTypeWithAutoCast(typeFactory, method.getReturnType());
+ return CalciteUtils.sqlTypeWithAutoCast(typeFactory, method.getGenericReturnType());
}
@Override
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRelDataTypeSystem.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRelDataTypeSystem.java
index 1d234b0..4452422 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRelDataTypeSystem.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRelDataTypeSystem.java
@@ -48,6 +48,8 @@
switch (typeName) {
case TIME:
return 6; // support microsecond time precision
+ case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
+ return 6; // support microsecond datetime precision
default:
return super.getMaxPrecision(typeName);
}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java
index b6d606b..d9ad401 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java
@@ -27,6 +27,7 @@
import java.lang.reflect.Type;
import java.math.BigDecimal;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.AbstractList;
import java.util.AbstractMap;
@@ -35,13 +36,14 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.TimeZone;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlPipelineOptions;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamJavaTypeFactory;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.CharType;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.TimeWithLocalTzType;
-import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.TimestampWithLocalTzType;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.logicaltypes.DateTime;
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
@@ -92,6 +94,7 @@
public class BeamCalcRel extends AbstractBeamCalcRel {
private static final long NANOS_PER_MILLISECOND = 1000000L;
+ private static final long MILLIS_PER_DAY = 86400000L;
private static final ParameterExpression outputSchemaParam =
Expressions.parameter(Schema.class, "outputSchema");
@@ -344,6 +347,18 @@
valueDateTime = Expressions.unbox(valueDateTime);
}
valueDateTime = Expressions.call(LocalDate.class, "ofEpochDay", valueDateTime);
+ } else if (CalciteUtils.TIMESTAMP_WITH_LOCAL_TZ.typesEqual(toType)
+ || CalciteUtils.NULLABLE_TIMESTAMP_WITH_LOCAL_TZ.typesEqual(toType)) {
+ // Convert TimeStamp_With_Local_TimeZone to LocalDateTime
+ Expression dateValue =
+ Expressions.divide(valueDateTime, Expressions.constant(MILLIS_PER_DAY));
+ Expression date = Expressions.call(LocalDate.class, "ofEpochDay", dateValue);
+ Expression timeValue =
+ Expressions.multiply(
+ Expressions.modulo(valueDateTime, Expressions.constant(MILLIS_PER_DAY)),
+ Expressions.constant(NANOS_PER_MILLISECOND));
+ Expression time = Expressions.call(LocalTime.class, "ofNanoOfDay", timeValue);
+ valueDateTime = Expressions.call(LocalDateTime.class, "of", date, time);
} else {
throw new UnsupportedOperationException("Unknown DateTime type " + toType);
}
@@ -385,7 +400,7 @@
.put(SqlTypes.DATE.getIdentifier(), Long.class)
.put(SqlTypes.TIME.getIdentifier(), Long.class)
.put(TimeWithLocalTzType.IDENTIFIER, ReadableInstant.class)
- .put(TimestampWithLocalTzType.IDENTIFIER, ReadableInstant.class)
+ .put(SqlTypes.DATETIME.getIdentifier(), Row.class)
.put(CharType.IDENTIFIER, String.class)
.build();
@@ -442,6 +457,16 @@
value, Expressions.divide(value, Expressions.constant(NANOS_PER_MILLISECOND)));
} else if (SqlTypes.DATE.getIdentifier().equals(logicalId)) {
return value;
+ } else if (SqlTypes.DATETIME.getIdentifier().equals(logicalId)) {
+ Expression dateValue =
+ Expressions.call(value, "getInt64", Expressions.constant(DateTime.DATE_FIELD_NAME));
+ Expression timeValue =
+ Expressions.call(value, "getInt64", Expressions.constant(DateTime.TIME_FIELD_NAME));
+ Expression returnValue =
+ Expressions.add(
+ Expressions.multiply(dateValue, Expressions.constant(MILLIS_PER_DAY)),
+ Expressions.divide(timeValue, Expressions.constant(NANOS_PER_MILLISECOND)));
+ return nullOr(value, returnValue);
} else if (!CharType.IDENTIFIER.equals(logicalId)) {
throw new UnsupportedOperationException(
"Unknown LogicalType " + type.getLogicalType().getIdentifier());
@@ -563,6 +588,8 @@
|| name.equals(DataContext.Variable.CURRENT_TIMESTAMP.camelName)
|| name.equals(DataContext.Variable.LOCAL_TIMESTAMP.camelName)) {
return System.currentTimeMillis();
+ } else if (name.equals(Variable.TIME_ZONE.camelName)) {
+ return TimeZone.getDefault();
}
return null;
}
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/utils/CalciteUtils.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/utils/CalciteUtils.java
index acb4ee1..6ded492 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/utils/CalciteUtils.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/utils/CalciteUtils.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.extensions.sql.impl.utils;
+import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.Date;
import java.util.Map;
@@ -53,15 +54,6 @@
}
}
- /** A LogicalType corresponding to TIMESTAMP_WITH_LOCAL_TIME_ZONE. */
- public static class TimestampWithLocalTzType extends PassThroughLogicalType<Instant> {
- public static final String IDENTIFIER = "SqlTimestampWithLocalTzType";
-
- public TimestampWithLocalTzType() {
- super(IDENTIFIER, FieldType.STRING, "", FieldType.DATETIME);
- }
- }
-
/** A LogicalType corresponding to CHAR. */
public static class CharType extends PassThroughLogicalType<String> {
public static final String IDENTIFIER = "SqlCharType";
@@ -82,7 +74,7 @@
return logicalId.equals(SqlTypes.DATE.getIdentifier())
|| logicalId.equals(SqlTypes.TIME.getIdentifier())
|| logicalId.equals(TimeWithLocalTzType.IDENTIFIER)
- || logicalId.equals(TimestampWithLocalTzType.IDENTIFIER);
+ || logicalId.equals(SqlTypes.DATETIME.getIdentifier());
}
return false;
}
@@ -121,8 +113,9 @@
FieldType.logicalType(new TimeWithLocalTzType());
public static final FieldType TIMESTAMP = FieldType.DATETIME;
public static final FieldType NULLABLE_TIMESTAMP = FieldType.DATETIME.withNullable(true);
- public static final FieldType TIMESTAMP_WITH_LOCAL_TZ =
- FieldType.logicalType(new TimestampWithLocalTzType());
+ public static final FieldType TIMESTAMP_WITH_LOCAL_TZ = FieldType.logicalType(SqlTypes.DATETIME);
+ public static final FieldType NULLABLE_TIMESTAMP_WITH_LOCAL_TZ =
+ FieldType.logicalType(SqlTypes.DATETIME).withNullable(true);
private static final BiMap<FieldType, SqlTypeName> BEAM_TO_CALCITE_TYPE_MAPPING =
ImmutableBiMap.<FieldType, SqlTypeName>builder()
@@ -283,18 +276,26 @@
/**
* SQL-Java type mapping, with specified Beam rules: <br>
- * 1. redirect {@link AbstractInstant} to {@link Date} so Calcite can recognize it.
+ * 1. redirect {@link AbstractInstant} to {@link Date} so Calcite can recognize it. <br>
+ * 2. For a list, the component type is needed to create a Sql array type.
*
- * @param rawType
- * @return
+ * @param type
+ * @return Calcite RelDataType
*/
- public static RelDataType sqlTypeWithAutoCast(RelDataTypeFactory typeFactory, Type rawType) {
+ public static RelDataType sqlTypeWithAutoCast(RelDataTypeFactory typeFactory, Type type) {
// For Joda time types, return SQL type for java.util.Date.
- if (rawType instanceof Class && AbstractInstant.class.isAssignableFrom((Class<?>) rawType)) {
+ if (type instanceof Class && AbstractInstant.class.isAssignableFrom((Class<?>) type)) {
return typeFactory.createJavaType(Date.class);
- } else if (rawType instanceof Class && ByteString.class.isAssignableFrom((Class<?>) rawType)) {
+ } else if (type instanceof Class && ByteString.class.isAssignableFrom((Class<?>) type)) {
return typeFactory.createJavaType(byte[].class);
+ } else if (type instanceof ParameterizedType
+ && java.util.List.class.isAssignableFrom(
+ (Class<?>) ((ParameterizedType) type).getRawType())) {
+ ParameterizedType parameterizedType = (ParameterizedType) type;
+ Class<?> genericType = (Class<?>) parameterizedType.getActualTypeArguments()[0];
+ RelDataType collectionElementType = typeFactory.createJavaType(genericType);
+ return typeFactory.createArrayType(collectionElementType, UNLIMITED_ARRAY_SIZE);
}
- return typeFactory.createJavaType((Class) rawType);
+ return typeFactory.createJavaType((Class) type);
}
}
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 7aa5032..f107dc3 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
@@ -38,6 +38,7 @@
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.calcite.util.TimestampString;
import org.checkerframework.checker.nullness.qual.Nullable;
public class BeamSqlUnparseContext extends SqlImplementor.SimpleContext {
@@ -69,8 +70,12 @@
public SqlNode toSql(RexProgram program, RexNode rex) {
if (rex.getKind().equals(SqlKind.LITERAL)) {
final RexLiteral literal = (RexLiteral) rex;
- SqlTypeFamily family = literal.getTypeName().getFamily();
- if (SqlTypeFamily.BINARY.equals(family)) {
+ SqlTypeName name = literal.getTypeName();
+ SqlTypeFamily family = name.getFamily();
+ if (SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE.equals(name)) {
+ TimestampString timestampString = literal.getValueAs(TimestampString.class);
+ return new SqlDateTimeLiteral(timestampString, POS);
+ } else if (SqlTypeFamily.BINARY.equals(family)) {
ByteString byteString = literal.getValueAs(ByteString.class);
BitString bitString = BitString.createFromHexString(byteString.toString(16));
return new SqlByteStringLiteral(bitString, POS);
@@ -92,6 +97,21 @@
return super.toSql(program, rex);
}
+ private static class SqlDateTimeLiteral extends SqlLiteral {
+
+ private final TimestampString timestampString;
+
+ SqlDateTimeLiteral(TimestampString timestampString, SqlParserPos pos) {
+ super(timestampString, SqlTypeName.TIMESTAMP, pos);
+ this.timestampString = timestampString;
+ }
+
+ @Override
+ public void unparse(SqlWriter writer, int leftPrec, int rightPrec) {
+ writer.literal("DATETIME '" + timestampString.toString() + "'");
+ }
+ }
+
private static class SqlByteStringLiteral extends SqlLiteral {
SqlByteStringLiteral(BitString bytes, SqlParserPos pos) {
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamComplexTypeTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamComplexTypeTest.java
index 1b1641f..0519798 100644
--- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamComplexTypeTest.java
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamComplexTypeTest.java
@@ -18,6 +18,7 @@
package org.apache.beam.sdk.extensions.sql;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.Arrays;
import java.util.HashMap;
@@ -368,22 +369,135 @@
}
@Test
- public void testNullDatetimeFields() {
+ public void testDatetimeFields() {
Instant current = new Instant(1561671380000L); // Long value corresponds to 27/06/2019
Schema dateTimeFieldSchema =
Schema.builder()
.addField("dateTimeField", FieldType.DATETIME)
.addNullableField("nullableDateTimeField", FieldType.DATETIME)
- .addField("timeTypeField", FieldType.logicalType(SqlTypes.TIME))
- .addNullableField("nullableTimeTypeField", FieldType.logicalType(SqlTypes.TIME))
+ .build();
+
+ Row dateTimeRow = Row.withSchema(dateTimeFieldSchema).addValues(current, null).build();
+
+ PCollection<Row> outputRow =
+ pipeline
+ .apply(Create.of(dateTimeRow))
+ .setRowSchema(dateTimeFieldSchema)
+ .apply(
+ SqlTransform.query(
+ "select EXTRACT(YEAR from dateTimeField) as yyyy, "
+ + " EXTRACT(YEAR from nullableDateTimeField) as year_with_null, "
+ + " EXTRACT(MONTH from dateTimeField) as mm, "
+ + " EXTRACT(MONTH from nullableDateTimeField) as month_with_null "
+ + " from PCOLLECTION"));
+
+ Schema outputRowSchema =
+ Schema.builder()
+ .addField("yyyy", FieldType.INT64)
+ .addNullableField("year_with_null", FieldType.INT64)
+ .addField("mm", FieldType.INT64)
+ .addNullableField("month_with_null", FieldType.INT64)
+ .build();
+
+ PAssert.that(outputRow)
+ .containsInAnyOrder(
+ Row.withSchema(outputRowSchema).addValues(2019L, null, 06L, null).build());
+
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
+ }
+
+ @Test
+ public void testSqlLogicalTypeDateFields() {
+ Schema dateTimeFieldSchema =
+ Schema.builder()
.addField("dateTypeField", FieldType.logicalType(SqlTypes.DATE))
.addNullableField("nullableDateTypeField", FieldType.logicalType(SqlTypes.DATE))
.build();
+ Row dateRow =
+ Row.withSchema(dateTimeFieldSchema).addValues(LocalDate.of(2019, 6, 27), null).build();
+
+ PCollection<Row> outputRow =
+ pipeline
+ .apply(Create.of(dateRow))
+ .setRowSchema(dateTimeFieldSchema)
+ .apply(
+ SqlTransform.query(
+ "select EXTRACT(DAY from dateTypeField) as dd, "
+ + " EXTRACT(DAY from nullableDateTypeField) as day_with_null, "
+ + " dateTypeField + interval '1' day as date_with_day_added, "
+ + " nullableDateTypeField + interval '1' day as day_added_with_null "
+ + " from PCOLLECTION"));
+
+ Schema outputRowSchema =
+ Schema.builder()
+ .addField("dd", FieldType.INT64)
+ .addNullableField("day_with_null", FieldType.INT64)
+ .addField("date_with_day_added", FieldType.logicalType(SqlTypes.DATE))
+ .addNullableField("day_added_with_null", FieldType.logicalType(SqlTypes.DATE))
+ .build();
+
+ PAssert.that(outputRow)
+ .containsInAnyOrder(
+ Row.withSchema(outputRowSchema)
+ .addValues(27L, null, LocalDate.of(2019, 6, 28), null)
+ .build());
+
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
+ }
+
+ @Test
+ public void testSqlLogicalTypeTimeFields() {
+ Schema dateTimeFieldSchema =
+ Schema.builder()
+ .addField("timeTypeField", FieldType.logicalType(SqlTypes.TIME))
+ .addNullableField("nullableTimeTypeField", FieldType.logicalType(SqlTypes.TIME))
+ .build();
+
+ Row timeRow =
+ Row.withSchema(dateTimeFieldSchema).addValues(LocalTime.of(1, 0, 0), null).build();
+
+ PCollection<Row> outputRow =
+ pipeline
+ .apply(Create.of(timeRow))
+ .setRowSchema(dateTimeFieldSchema)
+ .apply(
+ SqlTransform.query(
+ "select timeTypeField + interval '1' hour as time_with_hour_added, "
+ + " nullableTimeTypeField + interval '1' hour as hour_added_with_null, "
+ + " timeTypeField - INTERVAL '60' SECOND as time_with_seconds_added, "
+ + " nullableTimeTypeField - INTERVAL '60' SECOND as seconds_added_with_null "
+ + " from PCOLLECTION"));
+
+ Schema outputRowSchema =
+ Schema.builder()
+ .addField("time_with_hour_added", FieldType.logicalType(SqlTypes.TIME))
+ .addNullableField("hour_added_with_null", FieldType.logicalType(SqlTypes.TIME))
+ .addField("time_with_seconds_added", FieldType.logicalType(SqlTypes.TIME))
+ .addNullableField("seconds_added_with_null", FieldType.logicalType(SqlTypes.TIME))
+ .build();
+
+ PAssert.that(outputRow)
+ .containsInAnyOrder(
+ Row.withSchema(outputRowSchema)
+ .addValues(LocalTime.of(2, 0, 0), null, LocalTime.of(0, 59, 0), null)
+ .build());
+
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
+ }
+
+ @Test
+ public void testSqlLogicalTypeDatetimeFields() {
+ Schema dateTimeFieldSchema =
+ Schema.builder()
+ .addField("dateTimeField", FieldType.logicalType(SqlTypes.DATETIME))
+ .addNullableField("nullableDateTimeField", FieldType.logicalType(SqlTypes.DATETIME))
+ .build();
+
Row dateTimeRow =
Row.withSchema(dateTimeFieldSchema)
- .addValues(current, null, LocalTime.of(1, 0, 0), null, LocalDate.of(2019, 6, 27), null)
+ .addValues(LocalDateTime.of(2008, 12, 25, 15, 30, 0), null)
.build();
PCollection<Row> outputRow =
@@ -396,14 +510,14 @@
+ " EXTRACT(YEAR from nullableDateTimeField) as year_with_null, "
+ " EXTRACT(MONTH from dateTimeField) as mm, "
+ " EXTRACT(MONTH from nullableDateTimeField) as month_with_null, "
- + " timeTypeField + interval '1' hour as time_with_hour_added, "
- + " nullableTimeTypeField + interval '1' hour as hour_added_with_null, "
- + " timeTypeField - INTERVAL '60' SECOND as time_with_seconds_added, "
- + " nullableTimeTypeField - INTERVAL '60' SECOND as seconds_added_with_null, "
- + " EXTRACT(DAY from dateTypeField) as dd, "
- + " EXTRACT(DAY from nullableDateTypeField) as day_with_null, "
- + " dateTypeField + interval '1' day as date_with_day_added, "
- + " nullableDateTypeField + interval '1' day as day_added_with_null "
+ + " dateTimeField + interval '1' hour as time_with_hour_added, "
+ + " nullableDateTimeField + interval '1' hour as hour_added_with_null, "
+ + " dateTimeField - INTERVAL '60' SECOND as time_with_seconds_added, "
+ + " nullableDateTimeField - INTERVAL '60' SECOND as seconds_added_with_null, "
+ + " EXTRACT(DAY from dateTimeField) as dd, "
+ + " EXTRACT(DAY from nullableDateTimeField) as day_with_null, "
+ + " dateTimeField + interval '1' day as date_with_day_added, "
+ + " nullableDateTimeField + interval '1' day as day_added_with_null "
+ " from PCOLLECTION"));
Schema outputRowSchema =
@@ -412,31 +526,31 @@
.addNullableField("year_with_null", FieldType.INT64)
.addField("mm", FieldType.INT64)
.addNullableField("month_with_null", FieldType.INT64)
- .addField("time_with_hour_added", FieldType.logicalType(SqlTypes.TIME))
- .addNullableField("hour_added_with_null", FieldType.logicalType(SqlTypes.TIME))
- .addField("time_with_seconds_added", FieldType.logicalType(SqlTypes.TIME))
- .addNullableField("seconds_added_with_null", FieldType.logicalType(SqlTypes.TIME))
+ .addField("time_with_hour_added", FieldType.logicalType(SqlTypes.DATETIME))
+ .addNullableField("hour_added_with_null", FieldType.logicalType(SqlTypes.DATETIME))
+ .addField("time_with_seconds_added", FieldType.logicalType(SqlTypes.DATETIME))
+ .addNullableField("seconds_added_with_null", FieldType.logicalType(SqlTypes.DATETIME))
.addField("dd", FieldType.INT64)
.addNullableField("day_with_null", FieldType.INT64)
- .addField("date_with_day_added", FieldType.logicalType(SqlTypes.DATE))
- .addNullableField("day_added_with_null", FieldType.logicalType(SqlTypes.DATE))
+ .addField("date_with_day_added", FieldType.logicalType(SqlTypes.DATETIME))
+ .addNullableField("day_added_with_null", FieldType.logicalType(SqlTypes.DATETIME))
.build();
PAssert.that(outputRow)
.containsInAnyOrder(
Row.withSchema(outputRowSchema)
.addValues(
- 2019L,
+ 2008L,
null,
- 06L,
+ 12L,
null,
- LocalTime.of(2, 0, 0),
+ LocalDateTime.of(2008, 12, 25, 16, 30, 0),
null,
- LocalTime.of(0, 59, 0),
+ LocalDateTime.of(2008, 12, 25, 15, 29, 0),
null,
- 27L,
+ 25L,
null,
- LocalDate.of(2019, 6, 28),
+ LocalDateTime.of(2008, 12, 26, 15, 30, 0),
null)
.build());
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUdfUdafTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUdfUdafTest.java
index 75e8a08..c2afc5d 100644
--- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUdfUdafTest.java
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUdfUdafTest.java
@@ -23,6 +23,7 @@
import com.google.auto.service.AutoService;
import java.sql.Timestamp;
+import java.util.Arrays;
import java.util.Map;
import java.util.stream.IntStream;
import org.apache.beam.sdk.extensions.sql.impl.BeamCalciteTable;
@@ -30,6 +31,7 @@
import org.apache.beam.sdk.extensions.sql.meta.provider.UdfUdafProvider;
import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestBoundedTable;
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.transforms.Combine.CombineFn;
import org.apache.beam.sdk.transforms.SerializableFunction;
@@ -101,6 +103,29 @@
pipeline.run().waitUntilFinish();
}
+ @Test
+ public void testListUdf() throws Exception {
+ Schema resultType1 = Schema.builder().addArrayField("array_field", FieldType.INT64).build();
+ Row row1 = Row.withSchema(resultType1).addValue(Arrays.asList(1L)).build();
+ String sql1 = "SELECT test_array(1)";
+ PCollection<Row> result1 =
+ boundedInput1.apply(
+ "testArrayUdf",
+ SqlTransform.query(sql1).registerUdf("test_array", TestReturnTypeList.class));
+ PAssert.that(result1).containsInAnyOrder(row1);
+
+ Schema resultType2 = Schema.builder().addInt32Field("int_field").build();
+ Row row2 = Row.withSchema(resultType2).addValue(3).build();
+ String sql2 = "select array_length(ARRAY[1, 2, 3])";
+ PCollection<Row> result2 =
+ boundedInput1.apply(
+ "testArrayUdf2",
+ SqlTransform.query(sql2).registerUdf("array_length", TestListLength.class));
+ PAssert.that(result2).containsInAnyOrder(row2);
+
+ pipeline.run().waitUntilFinish();
+ }
+
/** Test that an indirect subclass of a {@link CombineFn} works as a UDAF. BEAM-3777 */
@Test
public void testUdafMultiLevelDescendent() {
@@ -347,6 +372,20 @@
}
}
+ /** A UDF to test support of array as return type. */
+ public static final class TestReturnTypeList implements BeamSqlUdf {
+ public static java.util.List<Long> eval(Long i) {
+ return Arrays.asList(i);
+ }
+ }
+
+ /** A UDF to test support of array as argument type. */
+ public static final class TestListLength implements BeamSqlUdf {
+ public static Integer eval(java.util.List<Long> i) {
+ return i.size();
+ }
+ }
+
/**
* UDF to test support for {@link
* org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.TableMacro}.
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamSqlRowCoderTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamSqlRowCoderTest.java
index f0854bc..d4819fc 100644
--- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamSqlRowCoderTest.java
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamSqlRowCoderTest.java
@@ -19,6 +19,7 @@
import java.math.BigDecimal;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
@@ -51,6 +52,7 @@
.add("col_string_varchar", SqlTypeName.VARCHAR)
.add("col_time", SqlTypeName.TIME)
.add("col_date", SqlTypeName.DATE)
+ .add("col_timestamp_with_local_time_zone", SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE)
.add("col_timestamp", SqlTypeName.TIMESTAMP)
.add("col_boolean", SqlTypeName.BOOLEAN)
.build();
@@ -70,6 +72,7 @@
"hello",
LocalTime.now(),
LocalDate.now(),
+ LocalDateTime.now(),
DateTime.now().toInstant(),
true)
.build();
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/DateTimeUtils.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/DateTimeUtils.java
index cc25380..5186099 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/DateTimeUtils.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/DateTimeUtils.java
@@ -20,11 +20,13 @@
import com.google.zetasql.CivilTimeEncoder;
import com.google.zetasql.Value;
import io.grpc.Status;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.List;
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.util.DateString;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.TimeString;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.TimestampString;
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.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
@@ -161,10 +163,17 @@
}
}
- public static TimeString convertTimeValueToTimeString(Value value) {
- LocalTime localTime = CivilTimeEncoder.decodePacked64TimeNanosAsJavaTime(value.getTimeValue());
- return new TimeString(localTime.getHour(), localTime.getMinute(), localTime.getSecond())
- .withNanos(localTime.getNano());
+ public static TimestampString convertDateTimeValueToTimeStampString(Value value) {
+ LocalDateTime dateTime =
+ CivilTimeEncoder.decodePacked96DatetimeNanosAsJavaTime(value.getDatetimeValue());
+ return new TimestampString(
+ dateTime.getYear(),
+ dateTime.getMonthValue(),
+ dateTime.getDayOfMonth(),
+ dateTime.getHour(),
+ dateTime.getMinute(),
+ dateTime.getSecond())
+ .withNanos(dateTime.getNano());
}
// dates are represented as an int32 value, indicating the offset
@@ -174,6 +183,12 @@
return DateString.fromDaysSinceEpoch(value.getDateValue());
}
+ public static TimeString convertTimeValueToTimeString(Value value) {
+ LocalTime localTime = CivilTimeEncoder.decodePacked64TimeNanosAsJavaTime(value.getTimeValue());
+ return new TimeString(localTime.getHour(), localTime.getMinute(), localTime.getSecond())
+ .withNanos(localTime.getNano());
+ }
+
public static Value parseDateToValue(String dateString) {
DateTime dateTime = parseDate(dateString);
return Value.createDateValue((int) (dateTime.getMillis() / MILLIS_PER_DAY));
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SupportedZetaSqlBuiltinFunctions.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SupportedZetaSqlBuiltinFunctions.java
index 6ccfb26..eb8d51f 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SupportedZetaSqlBuiltinFunctions.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/SupportedZetaSqlBuiltinFunctions.java
@@ -197,23 +197,23 @@
// Time functions
FunctionSignatureId.FN_CURRENT_DATE, // current_date
- // FunctionSignatureId.FN_CURRENT_DATETIME, // current_datetime
+ FunctionSignatureId.FN_CURRENT_DATETIME, // current_datetime
FunctionSignatureId.FN_CURRENT_TIME, // current_time
FunctionSignatureId.FN_CURRENT_TIMESTAMP, // current_timestamp
FunctionSignatureId.FN_DATE_ADD_DATE, // date_add
- // FunctionSignatureId.FN_DATETIME_ADD, // datetime_add
+ FunctionSignatureId.FN_DATETIME_ADD, // datetime_add
FunctionSignatureId.FN_TIME_ADD, // time_add
FunctionSignatureId.FN_TIMESTAMP_ADD, // timestamp_add
FunctionSignatureId.FN_DATE_DIFF_DATE, // date_diff
- // FunctionSignatureId.FN_DATETIME_DIFF, // datetime_diff
+ FunctionSignatureId.FN_DATETIME_DIFF, // datetime_diff
FunctionSignatureId.FN_TIME_DIFF, // time_diff
FunctionSignatureId.FN_TIMESTAMP_DIFF, // timestamp_diff
FunctionSignatureId.FN_DATE_SUB_DATE, // date_sub
- // FunctionSignatureId.FN_DATETIME_SUB, // datetime_sub
+ FunctionSignatureId.FN_DATETIME_SUB, // datetime_sub
FunctionSignatureId.FN_TIME_SUB, // time_sub
FunctionSignatureId.FN_TIMESTAMP_SUB, // timestamp_sub
FunctionSignatureId.FN_DATE_TRUNC_DATE, // date_trunc
- // FunctionSignatureId.FN_DATETIME_TRUNC, // datetime_trunc
+ FunctionSignatureId.FN_DATETIME_TRUNC, // datetime_trunc
FunctionSignatureId.FN_TIME_TRUNC, // time_trunc
FunctionSignatureId.FN_TIMESTAMP_TRUNC, // timestamp_trunc
FunctionSignatureId.FN_DATE_FROM_UNIX_DATE, // date_from_unix_date
@@ -234,19 +234,18 @@
FunctionSignatureId.FN_UNIX_MILLIS_FROM_TIMESTAMP,
// FunctionSignatureId.FN_UNIX_MICROS_FROM_TIMESTAMP,
FunctionSignatureId.FN_DATE_FROM_TIMESTAMP, // date
- // FunctionSignatureId.FN_DATE_FROM_DATETIME, // date
+ FunctionSignatureId.FN_DATE_FROM_DATETIME, // date
FunctionSignatureId.FN_DATE_FROM_YEAR_MONTH_DAY, // date
FunctionSignatureId.FN_TIMESTAMP_FROM_STRING, // timestamp
FunctionSignatureId.FN_TIMESTAMP_FROM_DATE, // timestamp
- // FunctionSignatureId.FN_TIMESTAMP_FROM_DATETIME, // timestamp
+ FunctionSignatureId.FN_TIMESTAMP_FROM_DATETIME, // timestamp
FunctionSignatureId.FN_TIME_FROM_HOUR_MINUTE_SECOND, // time
FunctionSignatureId.FN_TIME_FROM_TIMESTAMP, // time
- // FunctionSignatureId.FN_TIME_FROM_DATETIME, // time
- // FunctionSignatureId.FN_DATETIME_FROM_DATE_AND_TIME, // datetime
- // FunctionSignatureId.FN_DATETIME_FROM_YEAR_MONTH_DAY_HOUR_MINUTE_SECOND, // datetime
- // FunctionSignatureId.FN_DATETIME_FROM_TIMESTAMP, // datetime
- // FunctionSignatureId.FN_DATETIME_FROM_DATE, // datetime
-
+ FunctionSignatureId.FN_TIME_FROM_DATETIME, // time
+ FunctionSignatureId.FN_DATETIME_FROM_DATE_AND_TIME, // datetime
+ FunctionSignatureId.FN_DATETIME_FROM_YEAR_MONTH_DAY_HOUR_MINUTE_SECOND, // datetime
+ FunctionSignatureId.FN_DATETIME_FROM_TIMESTAMP, // datetime
+ FunctionSignatureId.FN_DATETIME_FROM_DATE, // datetime
FunctionSignatureId.FN_STRING_FROM_TIMESTAMP, // string
// Signatures for extracting date parts, taking a date/timestamp
@@ -258,23 +257,24 @@
// Signatures specific to extracting the DATE date part from a DATETIME or a
// TIMESTAMP.
- // FunctionSignatureId.FN_EXTRACT_DATE_FROM_DATETIME, // $extract_date
+ FunctionSignatureId.FN_EXTRACT_DATE_FROM_DATETIME, // $extract_date
FunctionSignatureId.FN_EXTRACT_DATE_FROM_TIMESTAMP, // $extract_date
// Signatures specific to extracting the TIME date part from a DATETIME or a
// TIMESTAMP.
- // FunctionSignatureId.FN_EXTRACT_TIME_FROM_DATETIME, // $extract_time
+ FunctionSignatureId.FN_EXTRACT_TIME_FROM_DATETIME, // $extract_time
FunctionSignatureId.FN_EXTRACT_TIME_FROM_TIMESTAMP, // $extract_time
// Signature specific to extracting the DATETIME date part from a TIMESTAMP.
- // FunctionSignatureId.FN_EXTRACT_DATETIME_FROM_TIMESTAMP, // $extract_datetime
+ FunctionSignatureId.FN_EXTRACT_DATETIME_FROM_TIMESTAMP, // $extract_datetime
+ // Signature for formatting and parsing
FunctionSignatureId.FN_FORMAT_DATE, // format_date
- // FunctionSignatureId.FN_FORMAT_DATETIME, // format_datetime
+ FunctionSignatureId.FN_FORMAT_DATETIME, // format_datetime
FunctionSignatureId.FN_FORMAT_TIME, // format_time
FunctionSignatureId.FN_FORMAT_TIMESTAMP, // format_timestamp
FunctionSignatureId.FN_PARSE_DATE, // parse_date
- // FunctionSignatureId.FN_PARSE_DATETIME, // parse_datetime
+ FunctionSignatureId.FN_PARSE_DATETIME, // parse_datetime
FunctionSignatureId.FN_PARSE_TIME, // parse_time
FunctionSignatureId.FN_PARSE_TIMESTAMP, // parse_timestamp
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtils.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtils.java
index 073aa41..dbab34a 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtils.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtils.java
@@ -28,6 +28,7 @@
import com.google.zetasql.ZetaSQLType.TypeKind;
import java.math.BigDecimal;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.ArrayList;
import java.util.List;
@@ -36,6 +37,7 @@
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.Field;
import org.apache.beam.sdk.schemas.Schema.FieldType;
+import org.apache.beam.sdk.schemas.logicaltypes.DateTime;
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.LongMath;
@@ -45,7 +47,6 @@
* Utility methods for ZetaSQL <=> Beam translation.
*
* <p>Unsupported ZetaSQL types: INT32, UINT32, UINT64, FLOAT, ENUM, PROTO, GEOGRAPHY
- * TODO[BEAM-10238]: support ZetaSQL types: TIME, DATETIME, NUMERIC
*/
@Internal
public final class ZetaSqlBeamTranslationUtils {
@@ -106,6 +107,9 @@
} else if (SqlTypes.TIME.getIdentifier().equals(identifier)) {
// Time type
return TypeFactory.createSimpleType(TypeKind.TYPE_TIME);
+ } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) {
+ // DateTime type
+ return TypeFactory.createSimpleType(TypeKind.TYPE_DATETIME);
} else {
throw new UnsupportedOperationException("Unknown Beam logical type: " + identifier);
}
@@ -184,6 +188,20 @@
} else { // input type
return Value.createTimeValue(CivilTimeEncoder.encodePacked64TimeNanos((LocalTime) object));
}
+ } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) {
+ // DateTime value
+ LocalDateTime datetime;
+ if (object instanceof Row) { // base type
+ datetime =
+ LocalDateTime.of(
+ LocalDate.ofEpochDay(((Row) object).getInt64(DateTime.DATE_FIELD_NAME)),
+ LocalTime.ofNanoOfDay(((Row) object).getInt64(DateTime.TIME_FIELD_NAME)));
+ } else { // input type
+ datetime = (LocalDateTime) object;
+ }
+ // TODO[BEAM-10611]: Create ZetaSQL Value.createDatetimeValue(LocalDateTime) function
+ return Value.createDatetimeValue(
+ CivilTimeEncoder.encodePacked64DatetimeSeconds(datetime), datetime.getNano());
} else {
throw new UnsupportedOperationException("Unknown Beam logical type: " + identifier);
}
@@ -208,6 +226,8 @@
return FieldType.logicalType(SqlTypes.DATE).withNullable(true);
case TYPE_TIME:
return FieldType.logicalType(SqlTypes.TIME).withNullable(true);
+ case TYPE_DATETIME:
+ return FieldType.logicalType(SqlTypes.DATETIME).withNullable(true);
case TYPE_TIMESTAMP:
return FieldType.DATETIME.withNullable(true);
case TYPE_ARRAY:
@@ -314,6 +334,9 @@
} else if (SqlTypes.TIME.getIdentifier().equals(identifier)) {
// Time value
return CivilTimeEncoder.decodePacked64TimeNanosAsJavaTime(value.getTimeValue());
+ } else if (SqlTypes.DATETIME.getIdentifier().equals(identifier)) {
+ // DateTime value
+ return CivilTimeEncoder.decodePacked96DatetimeNanosAsJavaTime(value.getDatetimeValue());
} else {
throw new UnsupportedOperationException("Unknown Beam logical type: " + identifier);
}
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlCalciteTranslationUtils.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlCalciteTranslationUtils.java
index d8394ab..81bc142 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlCalciteTranslationUtils.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlCalciteTranslationUtils.java
@@ -20,6 +20,7 @@
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_BOOL;
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_BYTES;
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_DATE;
+import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_DATETIME;
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_DOUBLE;
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_INT64;
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_NUMERIC;
@@ -46,7 +47,6 @@
* Utility methods for ZetaSQL <=> Calcite translation.
*
* <p>Unsupported ZetaSQL types: INT32, UINT32, UINT64, FLOAT, ENUM, PROTO, GEOGRAPHY
- * TODO[BEAM-10238]: support ZetaSQL types: TIME, DATETIME, NUMERIC
*/
@Internal
public final class ZetaSqlCalciteTranslationUtils {
@@ -72,6 +72,8 @@
return TypeFactory.createSimpleType(TYPE_DATE);
case TIME:
return TypeFactory.createSimpleType(TYPE_TIME);
+ case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
+ return TypeFactory.createSimpleType(TYPE_DATETIME);
case TIMESTAMP:
return TypeFactory.createSimpleType(TYPE_TIMESTAMP);
case ARRAY:
@@ -107,6 +109,8 @@
return SqlTypeName.DATE;
case TYPE_TIME:
return SqlTypeName.TIME;
+ case TYPE_DATETIME:
+ return SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE;
case TYPE_TIMESTAMP:
// TODO: handle timestamp with time zone.
return SqlTypeName.TIMESTAMP;
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 19f18cd..fd5651f 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
@@ -24,6 +24,7 @@
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_INT64;
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_STRING;
import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_TIMESTAMP;
+import static org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.convertDateTimeValueToTimeStampString;
import static org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.convertDateValueToDateString;
import static org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.convertTimeValueToTimeString;
import static org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.safeMicrosToMillis;
@@ -545,7 +546,7 @@
case TYPE_TIMESTAMP:
case TYPE_DATE:
case TYPE_TIME:
- // case TYPE_DATETIME:
+ case TYPE_DATETIME:
case TYPE_BYTES:
case TYPE_ARRAY:
case TYPE_STRUCT:
@@ -709,7 +710,7 @@
case TYPE_TIMESTAMP:
case TYPE_DATE:
case TYPE_TIME:
- // case TYPE_DATETIME:
+ case TYPE_DATETIME:
case TYPE_BYTES:
ret = convertSimpleValueToRexNode(type.getKind(), value);
break;
@@ -792,9 +793,7 @@
rexBuilder()
.makeCall(
SqlOperators.createZetaSqlFunction(wrapperFun, returnType.getSqlTypeName()),
- ImmutableList.of(
- rexBuilder()
- .makeApproxLiteral(new BigDecimal(Math.random()), returnType)));
+ rexBuilder().makeApproxLiteral(new BigDecimal(Math.random()), returnType));
;
} else {
ret =
@@ -823,12 +822,11 @@
SqlOperators.createZetaSqlFunction(
BeamBigQuerySqlDialect.NUMERIC_LITERAL_FUNCTION,
ZetaSqlCalciteTranslationUtils.toCalciteTypeName(kind)),
- ImmutableList.of(
- rexBuilder()
- .makeExactLiteral(
- value.getNumericValue(),
- ZetaSqlCalciteTranslationUtils.toSimpleRelDataType(
- kind, rexBuilder()))));
+ rexBuilder()
+ .makeExactLiteral(
+ value.getNumericValue(),
+ ZetaSqlCalciteTranslationUtils.toSimpleRelDataType(
+ kind, rexBuilder())));
break;
case TYPE_TIMESTAMP:
ret =
@@ -850,6 +848,15 @@
// TODO: Doing micro to mills truncation, need to throw exception.
ret = rexBuilder().makeLiteral(convertTimeValueToTimeString(value), timeType, false);
break;
+ case TYPE_DATETIME:
+ ret =
+ rexBuilder()
+ .makeTimestampWithLocalTimeZoneLiteral(
+ convertDateTimeValueToTimeStampString(value),
+ typeFactory()
+ .getTypeSystem()
+ .getMaxPrecision(SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE));
+ break;
case TYPE_BYTES:
ret = rexBuilder().makeBinaryLiteral(new ByteString(value.getBytesValue().toByteArray()));
break;
diff --git a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/TableScanConverter.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/TableScanConverter.java
index 25f40ed..9137b94 100644
--- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/TableScanConverter.java
+++ b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/TableScanConverter.java
@@ -17,17 +17,13 @@
*/
package org.apache.beam.sdk.extensions.sql.zetasql.translation;
-import static com.google.zetasql.ZetaSQLType.TypeKind.TYPE_DATETIME;
import static org.apache.beam.vendor.calcite.v1_20_0.com.google.common.base.Preconditions.checkNotNull;
-import com.google.zetasql.ZetaSQLType.TypeKind;
-import com.google.zetasql.resolvedast.ResolvedColumn;
import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedTableScan;
import java.util.List;
import java.util.Properties;
import org.apache.beam.sdk.extensions.sql.zetasql.TableResolution;
import org.apache.beam.vendor.calcite.v1_20_0.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.calcite.v1_20_0.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.config.CalciteConnectionConfigImpl;
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.RelOptCluster;
@@ -44,16 +40,12 @@
/** Converts table scan. */
class TableScanConverter extends RelConverter<ResolvedTableScan> {
- private static final ImmutableSet<TypeKind> UNSUPPORTED_DATA_TYPES =
- ImmutableSet.of(TYPE_DATETIME);
-
TableScanConverter(ConversionContext context) {
super(context);
}
@Override
public RelNode convert(ResolvedTableScan zetaNode, List<RelNode> inputs) {
- checkTableScanSchema(zetaNode.getColumnList());
List<String> tablePath = getTablePath(zetaNode.getTable());
@@ -115,15 +107,4 @@
}
};
}
-
- private void checkTableScanSchema(List<ResolvedColumn> columnList) {
- if (columnList != null) {
- for (ResolvedColumn resolvedColumn : columnList) {
- if (UNSUPPORTED_DATA_TYPES.contains(resolvedColumn.getType().getKind())) {
- throw new UnsupportedOperationException(
- "Does not support " + UNSUPPORTED_DATA_TYPES + " types in source tables");
- }
- }
- }
- }
}
diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/TestInput.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/TestInput.java
index a761d9f..2edb4d0 100644
--- a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/TestInput.java
+++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/TestInput.java
@@ -21,6 +21,7 @@
import java.nio.charset.StandardCharsets;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.Arrays;
import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestBoundedTable;
@@ -225,47 +226,59 @@
public static final TestBoundedTable TABLE_EMPTY =
TestBoundedTable.of(Schema.builder().addInt64Field("ColId").addStringField("Value").build());
- private static final Schema TABLE_WTH_MAP_SCHEMA =
+ private static final Schema TABLE_WITH_MAP_SCHEMA =
Schema.builder()
.addMapField("map_field", FieldType.STRING, FieldType.STRING)
.addRowField("row_field", structSchema)
.build();
public static final TestBoundedTable TABLE_WITH_MAP =
- TestBoundedTable.of(TABLE_WTH_MAP_SCHEMA)
+ TestBoundedTable.of(TABLE_WITH_MAP_SCHEMA)
.addRows(
ImmutableMap.of("MAP_KEY_1", "MAP_VALUE_1"),
Row.withSchema(structSchema).addValues(1L, "data1").build());
- private static final Schema TABLE_WTH_DATE_SCHEMA =
+ private static final Schema TABLE_WITH_DATE_SCHEMA =
Schema.builder()
.addLogicalTypeField("date_field", SqlTypes.DATE)
.addStringField("str_field")
.build();
public static final TestBoundedTable TABLE_WITH_DATE =
- TestBoundedTable.of(TABLE_WTH_DATE_SCHEMA)
+ TestBoundedTable.of(TABLE_WITH_DATE_SCHEMA)
.addRows(LocalDate.of(2008, 12, 25), "s")
.addRows(LocalDate.of(2020, 4, 7), "s");
- private static final Schema TABLE_WTH_TIME_SCHEMA =
+ private static final Schema TABLE_WITH_TIME_SCHEMA =
Schema.builder()
.addLogicalTypeField("time_field", SqlTypes.TIME)
.addStringField("str_field")
.build();
public static final TestBoundedTable TABLE_WITH_TIME =
- TestBoundedTable.of(TABLE_WTH_TIME_SCHEMA)
+ TestBoundedTable.of(TABLE_WITH_TIME_SCHEMA)
.addRows(LocalTime.of(15, 30, 0), "s")
.addRows(LocalTime.of(23, 35, 59), "s");
- private static final Schema TABLE_WTH_NUMERIC_SCHEMA =
+ private static final Schema TABLE_WITH_NUMERIC_SCHEMA =
Schema.builder().addDecimalField("numeric_field").addStringField("str_field").build();
+
public static final TestBoundedTable TABLE_WITH_NUMERIC =
- TestBoundedTable.of(TABLE_WTH_NUMERIC_SCHEMA)
+ TestBoundedTable.of(TABLE_WITH_NUMERIC_SCHEMA)
.addRows(ZetaSqlTypesUtils.bigDecimalAsNumeric("123.4567"), "str1")
.addRows(ZetaSqlTypesUtils.bigDecimalAsNumeric("765.4321"), "str2")
.addRows(ZetaSqlTypesUtils.bigDecimalAsNumeric("-555.5555"), "str3");
+ private static final Schema TABLE_WITH_DATETIME_SCHEMA =
+ Schema.builder()
+ .addLogicalTypeField("datetime_field", SqlTypes.DATETIME)
+ .addStringField("str_field")
+ .build();
+
+ public static final TestBoundedTable TABLE_WITH_DATETIME =
+ TestBoundedTable.of(TABLE_WITH_DATETIME_SCHEMA)
+ .addRows(LocalDateTime.of(2008, 12, 25, 15, 30, 0).withNano(123456000), "s")
+ .addRows(LocalDateTime.of(2012, 10, 6, 11, 45, 0).withNano(987654000), "s");
+
private static byte[] stringToBytes(String s) {
return s.getBytes(StandardCharsets.UTF_8);
}
diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtilsTest.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtilsTest.java
index 0510590..7b450fb 100644
--- a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtilsTest.java
+++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlBeamTranslationUtilsTest.java
@@ -28,6 +28,7 @@
import com.google.zetasql.Value;
import com.google.zetasql.ZetaSQLType.TypeKind;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.Arrays;
import org.apache.beam.sdk.schemas.Schema;
@@ -54,12 +55,12 @@
.addField("f_string", FieldType.STRING)
.addField("f_bytes", FieldType.BYTES)
.addLogicalTypeField("f_date", SqlTypes.DATE)
- // .addLogicalTypeField("f_datetime", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_datetime", SqlTypes.DATETIME)
.addLogicalTypeField("f_time", SqlTypes.TIME)
.addField("f_timestamp", FieldType.DATETIME)
.addArrayField("f_array", FieldType.DOUBLE)
.addRowField("f_struct", TEST_INNER_SCHEMA)
- // .addLogicalTypeField("f_numeric", SqlTypes.NUMERIC)
+ .addField("f_numeric", FieldType.DECIMAL)
.addNullableField("f_null", FieldType.INT64)
.build();
@@ -83,10 +84,12 @@
new StructField("f_string", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)),
new StructField("f_bytes", TypeFactory.createSimpleType(TypeKind.TYPE_BYTES)),
new StructField("f_date", TypeFactory.createSimpleType(TypeKind.TYPE_DATE)),
+ new StructField("f_datetime", TypeFactory.createSimpleType(TypeKind.TYPE_DATETIME)),
new StructField("f_time", TypeFactory.createSimpleType(TypeKind.TYPE_TIME)),
new StructField("f_timestamp", TypeFactory.createSimpleType(TypeKind.TYPE_TIMESTAMP)),
new StructField("f_array", TEST_INNER_ARRAY_TYPE),
new StructField("f_struct", TEST_INNER_STRUCT_TYPE),
+ new StructField("f_numeric", TypeFactory.createSimpleType(TypeKind.TYPE_NUMERIC)),
new StructField("f_null", TypeFactory.createSimpleType(TypeKind.TYPE_INT64))));
private static final Row TEST_ROW =
@@ -97,10 +100,12 @@
.addValue("Hello")
.addValue(new byte[] {0x11, 0x22})
.addValue(LocalDate.of(2020, 6, 4))
+ .addValue(LocalDateTime.of(2008, 12, 25, 15, 30, 0))
.addValue(LocalTime.of(15, 30, 45))
.addValue(Instant.ofEpochMilli(12345678L))
.addArray(3.0, 6.5)
.addValue(Row.withSchema(TEST_INNER_SCHEMA).addValues(0L, "world").build())
+ .addValue(ZetaSqlTypesUtils.bigDecimalAsNumeric("12346"))
.addValue(null)
.build();
@@ -114,6 +119,10 @@
Value.createStringValue("Hello"),
Value.createBytesValue(ByteString.copyFrom(new byte[] {0x11, 0x22})),
Value.createDateValue((int) LocalDate.of(2020, 6, 4).toEpochDay()),
+ Value.createDatetimeValue(
+ CivilTimeEncoder.encodePacked64DatetimeSeconds(
+ LocalDateTime.of(2008, 12, 25, 15, 30, 0)),
+ LocalDateTime.of(2008, 12, 25, 15, 30, 0).getNano()),
Value.createTimeValue(
CivilTimeEncoder.encodePacked64TimeNanos(LocalTime.of(15, 30, 45))),
Value.createTimestampValueFromUnixMicros(12345678000L),
@@ -123,6 +132,7 @@
Value.createStructValue(
TEST_INNER_STRUCT_TYPE,
Arrays.asList(Value.createInt64Value(0L), Value.createStringValue("world"))),
+ Value.createNumericValue(ZetaSqlTypesUtils.bigDecimalAsNumeric("12346")),
Value.createNullValue(TypeFactory.createSimpleType(TypeKind.TYPE_INT64))));
@Test
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 e9c51c7..d148012 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
@@ -2942,7 +2942,7 @@
}
@Test
- @Ignore("BEAM-9515")
+ @Ignore("[BEAM-9515] ArrayScanToUncollectConverter Unnest does not support sub-queries")
public void testUNNESTExpression() {
String sql = "SELECT * FROM UNNEST(ARRAY(SELECT Value FROM KeyValue));";
ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTestBase.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTestBase.java
index 6d9ba67..483a9c1 100644
--- a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTestBase.java
+++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTestBase.java
@@ -61,6 +61,7 @@
testBoundedTableMap.put("table_with_date", TestInput.TABLE_WITH_DATE);
testBoundedTableMap.put("table_with_time", TestInput.TABLE_WITH_TIME);
testBoundedTableMap.put("table_with_numeric", TestInput.TABLE_WITH_NUMERIC);
+ testBoundedTableMap.put("table_with_datetime", TestInput.TABLE_WITH_DATETIME);
testBoundedTableMap.put(
"table_with_struct_ts_string", TestInput.TABLE_WITH_STRUCT_TIMESTAMP_STRING);
testBoundedTableMap.put("streaming_sql_test_table_a", TestInput.STREAMING_SQL_TABLE_A);
diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTimeFunctionsTest.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTimeFunctionsTest.java
index 6789d63..109ca1e 100644
--- a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTimeFunctionsTest.java
+++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSqlTimeFunctionsTest.java
@@ -23,9 +23,11 @@
import static org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.parseTimestampWithTimeZone;
import static org.apache.beam.sdk.extensions.sql.zetasql.DateTimeUtils.parseTimestampWithUTCTimeZone;
+import com.google.zetasql.CivilTimeEncoder;
import com.google.zetasql.Value;
import com.google.zetasql.ZetaSQLType.TypeKind;
import java.time.LocalDate;
+import java.time.LocalDateTime;
import java.time.LocalTime;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlPipelineOptions;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
@@ -40,7 +42,6 @@
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.joda.time.Duration;
import org.junit.Before;
-import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
@@ -148,7 +149,7 @@
+ " EXTRACT(ISOYEAR FROM date) AS isoyear,\n"
+ " EXTRACT(YEAR FROM date) AS year,\n"
+ " EXTRACT(ISOWEEK FROM date) AS isoweek,\n"
- // TODO[BEAM-9178]: Add tests for DATE_TRUNC and EXTRACT with "week with weekday" date
+ // TODO[BEAM-10606]: Add tests for DATE_TRUNC and EXTRACT with "week with weekday" date
// parts once they are supported
// + " EXTRACT(WEEK FROM date) AS week,\n"
+ " EXTRACT(MONTH FROM date) AS month,\n"
@@ -219,6 +220,22 @@
}
@Test
+ public void testDateFromDateTime() {
+ String sql = "SELECT DATE(DATETIME '2008-12-25 15:30:00.123456')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addLogicalTypeField("f_date", SqlTypes.DATE).build())
+ .addValues(LocalDate.of(2008, 12, 25))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testDateAdd() {
String sql =
"SELECT "
@@ -580,6 +597,22 @@
}
@Test
+ public void testTimeFromDateTime() {
+ String sql = "SELECT TIME(DATETIME '2008-12-25 15:30:00.123456')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addLogicalTypeField("f_time", SqlTypes.TIME).build())
+ .addValues(LocalTime.of(15, 30, 0, 123456000))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testTimeAdd() {
String sql =
"SELECT "
@@ -753,13 +786,420 @@
/////////////////////////////////////////////////////////////////////////////
@Test
- @Ignore("Does not support Datetime literal.")
- public void testDatetimeLiteral() {
- String sql = "SELECT DATETIME '2018-01-01 05:30:00.334'";
+ public void testDateTimeLiteral() {
+ String sql = "SELECT DATETIME '2008-12-25 15:30:00.123456'";
+
ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
- thrown.expect(RuntimeException.class);
- thrown.expectMessage("Unsupported ResolvedLiteral type: DATETIME");
- zetaSQLQueryPlanner.convertToBeamRel(sql);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("f_datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2008, 12, 25, 15, 30, 0).withNano(123456000))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeColumn() {
+ String sql = "SELECT FORMAT_DATETIME('%D %T %E6S', datetime_field) FROM table_with_datetime";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addStringField("f_datetime_str").build())
+ .addValues("12/25/08 15:30:00 00.123456")
+ .build(),
+ Row.withSchema(Schema.builder().addStringField("f_datetime_str").build())
+ .addValues("10/06/12 11:45:00 00.987654")
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testGroupByDateTime() {
+ String sql = "SELECT datetime_field, COUNT(*) FROM table_with_datetime GROUP BY datetime_field";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ final Schema schema =
+ Schema.builder()
+ .addLogicalTypeField("datetime_field", SqlTypes.DATETIME)
+ .addInt64Field("count")
+ .build();
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(schema)
+ .addValues(LocalDateTime.of(2008, 12, 25, 15, 30, 0).withNano(123456000), 1L)
+ .build(),
+ Row.withSchema(schema)
+ .addValues(LocalDateTime.of(2012, 10, 6, 11, 45, 0).withNano(987654000), 1L)
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testAggregateOnDateTime() {
+ String sql = "SELECT MAX(datetime_field) FROM table_with_datetime GROUP BY str_field";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder()
+ .addLogicalTypeField("datetime_field", SqlTypes.DATETIME)
+ .build())
+ .addValues(LocalDateTime.of(2012, 10, 6, 11, 45, 0).withNano(987654000))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ // TODO[BEAM-9166]: Add a test for CURRENT_DATETIME function ("SELECT CURRENT_DATETIME()")
+
+ @Test
+ public void testExtractFromDateTime() {
+ String sql =
+ "SELECT "
+ + "EXTRACT(YEAR FROM DATETIME '2008-12-25 15:30:00') as year, "
+ + "EXTRACT(QUARTER FROM DATETIME '2008-12-25 15:30:00') as quarter, "
+ + "EXTRACT(MONTH FROM DATETIME '2008-12-25 15:30:00') as month, "
+ // TODO[BEAM-10606]: Add tests for DATETIME_TRUNC and EXTRACT with "week with weekday"
+ // date parts once they are supported
+ // + "EXTRACT(WEEK FROM DATETIME '2008-12-25 15:30:00') as week, "
+ + "EXTRACT(DAY FROM DATETIME '2008-12-25 15:30:00') as day, "
+ + "EXTRACT(DAYOFWEEK FROM DATETIME '2008-12-25 15:30:00') as dayofweek, "
+ + "EXTRACT(DAYOFYEAR FROM DATETIME '2008-12-25 15:30:00') as dayofyear, "
+ + "EXTRACT(HOUR FROM DATETIME '2008-12-25 15:30:00.123456') as hour, "
+ + "EXTRACT(MINUTE FROM DATETIME '2008-12-25 15:30:00.123456') as minute, "
+ + "EXTRACT(SECOND FROM DATETIME '2008-12-25 15:30:00.123456') as second, "
+ + "EXTRACT(MILLISECOND FROM DATETIME '2008-12-25 15:30:00.123456') as millisecond, "
+ + "EXTRACT(MICROSECOND FROM DATETIME '2008-12-25 15:30:00.123456') as microsecond, "
+ + "EXTRACT(DATE FROM DATETIME '2008-12-25 15:30:00.123456') as date, "
+ + "EXTRACT(TIME FROM DATETIME '2008-12-25 15:30:00.123456') as time ";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ final Schema schema =
+ Schema.builder()
+ .addInt64Field("year")
+ .addInt64Field("quarter")
+ .addInt64Field("month")
+ // .addInt64Field("week")
+ .addInt64Field("day")
+ .addInt64Field("dayofweek")
+ .addInt64Field("dayofyear")
+ .addInt64Field("hour")
+ .addInt64Field("minute")
+ .addInt64Field("second")
+ .addInt64Field("millisecond")
+ .addInt64Field("microsecond")
+ .addLogicalTypeField("date", SqlTypes.DATE)
+ .addLogicalTypeField("time", SqlTypes.TIME)
+ .build();
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(schema)
+ .addValues(
+ 2008L,
+ 4L,
+ 12L,
+ // 52L,
+ 25L,
+ 5L,
+ 360L,
+ 15L,
+ 30L,
+ 0L,
+ 123L,
+ 123456L,
+ LocalDate.of(2008, 12, 25),
+ LocalTime.of(15, 30, 0, 123456000))
+ .build());
+
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeFromDateAndTime() {
+ String sql = "SELECT DATETIME(DATE '2008-12-25', TIME '15:30:00.123456')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("f_datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2008, 12, 25, 15, 30, 0).withNano(123456000))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeFromDate() {
+ String sql = "SELECT DATETIME(DATE '2008-12-25')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("f_datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2008, 12, 25, 0, 0, 0))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeFromYearMonthDayHourMinuteSecond() {
+ String sql = "SELECT DATETIME(2008, 12, 25, 15, 30, 0)";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("f_datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2008, 12, 25, 15, 30, 0))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeFromTimestamp() {
+ String sql = "SELECT DATETIME(TIMESTAMP '2008-12-25 15:30:00+08', 'America/Los_Angeles')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("f_datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2008, 12, 24, 23, 30, 0))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeAdd() {
+ String sql =
+ "SELECT "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MICROSECOND), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MILLISECOND), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 SECOND), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MINUTE), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 HOUR), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 DAY), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MONTH), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 QUARTER), "
+ + "DATETIME_ADD(DATETIME '2008-12-25 15:30:00', INTERVAL 10 YEAR) ";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder()
+ .addLogicalTypeField("f_time1", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time2", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time3", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time4", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time5", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time6", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time7", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time8", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time9", SqlTypes.DATETIME)
+ .build())
+ .addValues(
+ LocalDateTime.of(2008, 12, 25, 15, 30, 0).withNano(10000),
+ LocalDateTime.of(2008, 12, 25, 15, 30, 0).withNano(10000000),
+ LocalDateTime.of(2008, 12, 25, 15, 30, 10),
+ LocalDateTime.of(2008, 12, 25, 15, 40, 0),
+ LocalDateTime.of(2008, 12, 26, 1, 30, 0),
+ LocalDateTime.of(2009, 1, 4, 15, 30, 0),
+ LocalDateTime.of(2009, 10, 25, 15, 30, 0),
+ LocalDateTime.of(2011, 6, 25, 15, 30, 0),
+ LocalDateTime.of(2018, 12, 25, 15, 30, 0))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeAddWithParameter() {
+ String sql = "SELECT DATETIME_ADD(@p0, INTERVAL @p1 HOUR)";
+
+ LocalDateTime datetime = LocalDateTime.of(2008, 12, 25, 15, 30, 00).withNano(123456000);
+ ImmutableMap<String, Value> params =
+ ImmutableMap.of(
+ "p0",
+ Value.createDatetimeValue(
+ CivilTimeEncoder.encodePacked64DatetimeSeconds(datetime), datetime.getNano()),
+ "p1", Value.createInt64Value(3L));
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql, params);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("f_datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2008, 12, 25, 18, 30, 00).withNano(123456000))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeSub() {
+ String sql =
+ "SELECT "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MICROSECOND), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MILLISECOND), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 SECOND), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MINUTE), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 HOUR), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 DAY), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 MONTH), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 QUARTER), "
+ + "DATETIME_SUB(DATETIME '2008-12-25 15:30:00', INTERVAL 10 YEAR) ";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder()
+ .addLogicalTypeField("f_time1", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time2", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time3", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time4", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time5", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time6", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time7", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time8", SqlTypes.DATETIME)
+ .addLogicalTypeField("f_time9", SqlTypes.DATETIME)
+ .build())
+ .addValues(
+ LocalDateTime.of(2008, 12, 25, 15, 29, 59).withNano(999990000),
+ LocalDateTime.of(2008, 12, 25, 15, 29, 59).withNano(990000000),
+ LocalDateTime.of(2008, 12, 25, 15, 29, 50),
+ LocalDateTime.of(2008, 12, 25, 15, 20, 0),
+ LocalDateTime.of(2008, 12, 25, 5, 30, 0),
+ LocalDateTime.of(2008, 12, 15, 15, 30, 0),
+ LocalDateTime.of(2008, 2, 25, 15, 30, 0),
+ LocalDateTime.of(2006, 6, 25, 15, 30, 0),
+ LocalDateTime.of(1998, 12, 25, 15, 30, 0))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeDiff() {
+ String sql =
+ "SELECT DATETIME_DIFF(DATETIME '2008-12-25 15:30:00', DATETIME '2008-10-25 15:30:00', DAY)";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addInt64Field("f_datetime_diff").build())
+ .addValues(61L)
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeDiffNegativeResult() {
+ String sql =
+ "SELECT DATETIME_DIFF(DATETIME '2008-10-25 15:30:00', DATETIME '2008-12-25 15:30:00', DAY)";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addInt64Field("f_datetime_diff").build())
+ .addValues(-61L)
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testDateTimeTrunc() {
+ String sql = "SELECT DATETIME_TRUNC(DATETIME '2008-12-25 15:30:00', HOUR)";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder()
+ .addLogicalTypeField("f_datetime_trunc", SqlTypes.DATETIME)
+ .build())
+ .addValues(LocalDateTime.of(2008, 12, 25, 15, 0, 0))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testFormatDateTime() {
+ String sql = "SELECT FORMAT_DATETIME('%D %T %E6S', DATETIME '2008-12-25 15:30:00.123456')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addStringField("f_datetime_str").build())
+ .addValues("12/25/08 15:30:00 00.123456")
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ public void testParseDateTime() {
+ String sql = "SELECT PARSE_DATETIME('%Y-%m-%d %H:%M:%E6S', '2008-12-25 15:30:00.123456')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("f_datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2008, 12, 25, 15, 30, 0).withNano(123456000))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
}
/////////////////////////////////////////////////////////////////////////////
@@ -846,7 +1286,7 @@
+ " EXTRACT(ISOYEAR FROM timestamp) AS isoyear,\n"
+ " EXTRACT(YEAR FROM timestamp) AS year,\n"
+ " EXTRACT(ISOWEEK FROM timestamp) AS isoweek,\n"
- // TODO[BEAM-9178]: Add tests for TIMESTAMP_TRUNC and EXTRACT with "week with weekday"
+ // TODO[BEAM-10606]: Add tests for TIMESTAMP_TRUNC and EXTRACT with "week with weekday"
// date parts once they are supported
// + " EXTRACT(WEEK FROM timestamp) AS week,\n"
+ " EXTRACT(MONTH FROM timestamp) AS month,\n"
@@ -926,6 +1366,23 @@
}
@Test
+ public void testExtractDateTimeFromTimestamp() {
+ String sql = "SELECT EXTRACT(DATETIME FROM TIMESTAMP '2017-05-26 12:34:56')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(
+ Schema.builder().addLogicalTypeField("datetime", SqlTypes.DATETIME).build())
+ .addValues(LocalDateTime.of(2017, 5, 26, 12, 34, 56))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testExtractFromTimestampAtTimeZone() {
String sql =
"WITH Timestamps AS (\n"
@@ -1028,6 +1485,45 @@
}
@Test
+ public void testTimestampFromDateTime() {
+ String sql = "SELECT TIMESTAMP(DATETIME '2008-12-25 15:30:00')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addDateTimeField("f_timestamp").build())
+ .addValues(parseTimestampWithTimeZone("2008-12-25 15:30:00+00"))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
+ // test default timezone works properly in query execution stage
+ public void testTimestampFromDateTimeWithDefaultTimezoneSet() {
+ String sql = "SELECT TIMESTAMP(DATETIME '2008-12-25 15:30:00')";
+
+ ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+ zetaSQLQueryPlanner.setDefaultTimezone("Asia/Shanghai");
+ pipeline
+ .getOptions()
+ .as(BeamSqlPipelineOptions.class)
+ .setZetaSqlDefaultTimezone("Asia/Shanghai");
+
+ BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+ PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+ PAssert.that(stream)
+ .containsInAnyOrder(
+ Row.withSchema(Schema.builder().addDateTimeField("f_timestamp").build())
+ .addValues(parseTimestampWithTimeZone("2008-12-25 15:30:00+08"))
+ .build());
+ pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+ }
+
+ @Test
public void testTimestampAdd() {
String sql =
"SELECT "
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 9521443..3f449d2 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
@@ -129,7 +129,7 @@
new PubsubMessage().encodeData(outgoingMessage.message().getData().toByteArray());
pubsubMessage.setAttributes(getMessageAttributes(outgoingMessage));
if (!outgoingMessage.message().getOrderingKey().isEmpty()) {
- pubsubMessage.put("orderingKey", outgoingMessage.message().getOrderingKey());
+ pubsubMessage.setOrderingKey(outgoingMessage.message().getOrderingKey());
}
pubsubMessages.add(pubsubMessage);
}
@@ -156,6 +156,7 @@
}
@Override
+ @SuppressWarnings("ProtoFieldNullComparison")
public List<IncomingMessage> pull(
long requestTimeMsSinceEpoch,
SubscriptionPath subscription,
@@ -207,8 +208,12 @@
com.google.pubsub.v1.PubsubMessage.newBuilder();
protoMessage.setData(ByteString.copyFrom(elementBytes));
protoMessage.putAllAttributes(attributes);
- protoMessage.setOrderingKey(
- (String) pubsubMessage.getUnknownKeys().getOrDefault("orderingKey", ""));
+ // PubsubMessage uses `null` to represent no ordering key where we want a default of "".
+ if (pubsubMessage.getOrderingKey() != null) {
+ protoMessage.setOrderingKey(pubsubMessage.getOrderingKey());
+ } else {
+ protoMessage.setOrderingKey("");
+ }
incomingMessages.add(
IncomingMessage.of(
protoMessage.build(),
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 aad9729..22c1cb1 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
@@ -101,7 +101,7 @@
.setAttributes(
ImmutableMap.of(
TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME), ID_ATTRIBUTE, RECORD_ID))
- .set("orderingKey", ORDERING_KEY);
+ .setOrderingKey(ORDERING_KEY);
ReceivedMessage expectedReceivedMessage =
new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
PullResponse expectedResponse =
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/CrossLanguageConfiguration.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/CrossLanguageConfiguration.java
deleted file mode 100644
index df0ad3e..0000000
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/CrossLanguageConfiguration.java
+++ /dev/null
@@ -1,76 +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.sdk.io.jdbc;
-
-import java.util.List;
-import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
-
-/** Abstract Parameters class to expose the Jdbc transforms to an external SDK. */
-public abstract class CrossLanguageConfiguration {
- String driverClassName;
- String jdbcUrl;
- String username;
- String password;
- String connectionProperties;
- Iterable<String> connectionInitSqls;
-
- public void setDriverClassName(String driverClassName) {
- this.driverClassName = driverClassName;
- }
-
- public void setJdbcUrl(String jdbcUrl) {
- this.jdbcUrl = jdbcUrl;
- }
-
- public void setUsername(String username) {
- this.username = username;
- }
-
- public void setPassword(String password) {
- this.password = password;
- }
-
- public void setConnectionProperties(String connectionProperties) {
- this.connectionProperties = connectionProperties;
- }
-
- public void setConnectionInitSqls(Iterable<String> connectionInitSqls) {
- this.connectionInitSqls = connectionInitSqls;
- }
-
- protected JdbcIO.DataSourceConfiguration getDataSourceConfiguration() {
- JdbcIO.DataSourceConfiguration dataSourceConfiguration =
- JdbcIO.DataSourceConfiguration.create(driverClassName, jdbcUrl)
- .withUsername(username)
- .withPassword(password);
-
- if (connectionProperties != null) {
- dataSourceConfiguration =
- dataSourceConfiguration.withConnectionProperties(connectionProperties);
- }
-
- if (connectionInitSqls != null) {
- List<String> initSqls =
- StreamSupport.stream(connectionInitSqls.spliterator(), false)
- .collect(Collectors.toList());
- dataSourceConfiguration = dataSourceConfiguration.withConnectionInitSqls(initSqls);
- }
- return dataSourceConfiguration;
- }
-}
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcReadRowsRegistrar.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcReadRowsRegistrar.java
deleted file mode 100644
index 72c8b05..0000000
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcReadRowsRegistrar.java
+++ /dev/null
@@ -1,84 +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.sdk.io.jdbc;
-
-import com.google.auto.service.AutoService;
-import java.util.Map;
-import org.apache.beam.sdk.annotations.Experimental;
-import org.apache.beam.sdk.annotations.Experimental.Kind;
-import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
-import org.apache.beam.sdk.io.jdbc.JdbcIO.DataSourceConfiguration;
-import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.values.PBegin;
-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.ImmutableMap;
-
-/** Exposes {@link JdbcIO.ReadRows} as an external transform for cross-language usage. */
-@Experimental(Kind.PORTABILITY)
-@AutoService(ExternalTransformRegistrar.class)
-public class JdbcReadRowsRegistrar implements ExternalTransformRegistrar {
-
- public static final String URN = "beam:external:java:jdbc:read_rows:v1";
-
- @Override
- public Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders() {
- return ImmutableMap.of(URN, JdbcReadRowsRegistrar.Builder.class);
- }
-
- /** Parameters class to expose the Read transform to an external SDK. */
- public static class ReadConfiguration extends CrossLanguageConfiguration {
- private String query;
- private Integer fetchSize;
- private Boolean outputParallelization;
-
- public void setOutputParallelization(Boolean outputParallelization) {
- this.outputParallelization = outputParallelization;
- }
-
- public void setFetchSize(Integer fetchSize) {
- this.fetchSize = fetchSize;
- }
-
- public void setQuery(String query) {
- this.query = query;
- }
- }
-
- public static class Builder
- implements ExternalTransformBuilder<ReadConfiguration, PBegin, PCollection<Row>> {
- @Override
- public PTransform<PBegin, PCollection<Row>> buildExternal(ReadConfiguration configuration) {
- DataSourceConfiguration dataSourceConfiguration = configuration.getDataSourceConfiguration();
-
- JdbcIO.ReadRows readRows =
- JdbcIO.readRows()
- .withDataSourceConfiguration(dataSourceConfiguration)
- .withQuery(configuration.query);
-
- if (configuration.fetchSize != null) {
- readRows = readRows.withFetchSize(configuration.fetchSize);
- }
- if (configuration.outputParallelization != null) {
- readRows = readRows.withOutputParallelization(configuration.outputParallelization);
- }
- return readRows;
- }
- }
-}
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcTransformsRegistrar.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcTransformsRegistrar.java
new file mode 100644
index 0000000..cfa22d5
--- /dev/null
+++ b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcTransformsRegistrar.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.sdk.io.jdbc;
+
+import com.google.auto.service.AutoService;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
+import org.apache.beam.sdk.io.jdbc.JdbcIO.DataSourceConfiguration;
+import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+
+/**
+ * Exposes {@link JdbcIO.Write} and {@link JdbcIO.ReadRows} as the external transforms for
+ * cross-language usage.
+ */
+@Experimental(Kind.PORTABILITY)
+@AutoService(ExternalTransformRegistrar.class)
+public class JdbcTransformsRegistrar implements ExternalTransformRegistrar {
+
+ public static final String READ_ROWS_URN = "beam:external:java:jdbc:read_rows:v1";
+ public static final String WRITE_URN = "beam:external:java:jdbc:write:v1";
+
+ @Override
+ public Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders() {
+ return ImmutableMap.of(READ_ROWS_URN, ReadRowsBuilder.class, WRITE_URN, WriteBuilder.class);
+ }
+
+ private static class CrossLanguageConfiguration {
+ String driverClassName;
+ String jdbcUrl;
+ String username;
+ String password;
+ String connectionProperties;
+ Iterable<String> connectionInitSqls;
+
+ public void setDriverClassName(String driverClassName) {
+ this.driverClassName = driverClassName;
+ }
+
+ public void setJdbcUrl(String jdbcUrl) {
+ this.jdbcUrl = jdbcUrl;
+ }
+
+ public void setUsername(String username) {
+ this.username = username;
+ }
+
+ public void setPassword(String password) {
+ this.password = password;
+ }
+
+ public void setConnectionProperties(String connectionProperties) {
+ this.connectionProperties = connectionProperties;
+ }
+
+ public void setConnectionInitSqls(Iterable<String> connectionInitSqls) {
+ this.connectionInitSqls = connectionInitSqls;
+ }
+
+ protected JdbcIO.DataSourceConfiguration getDataSourceConfiguration() {
+ JdbcIO.DataSourceConfiguration dataSourceConfiguration =
+ JdbcIO.DataSourceConfiguration.create(driverClassName, jdbcUrl)
+ .withUsername(username)
+ .withPassword(password);
+
+ if (connectionProperties != null) {
+ dataSourceConfiguration =
+ dataSourceConfiguration.withConnectionProperties(connectionProperties);
+ }
+
+ if (connectionInitSqls != null) {
+ List<String> initSqls =
+ StreamSupport.stream(connectionInitSqls.spliterator(), false)
+ .collect(Collectors.toList());
+ dataSourceConfiguration = dataSourceConfiguration.withConnectionInitSqls(initSqls);
+ }
+ return dataSourceConfiguration;
+ }
+ }
+
+ public static class ReadRowsBuilder
+ implements ExternalTransformBuilder<ReadRowsBuilder.Configuration, PBegin, PCollection<Row>> {
+
+ public static class Configuration extends CrossLanguageConfiguration {
+ private String query;
+ private Integer fetchSize;
+ private Boolean outputParallelization;
+
+ public void setOutputParallelization(Boolean outputParallelization) {
+ this.outputParallelization = outputParallelization;
+ }
+
+ public void setFetchSize(Integer fetchSize) {
+ this.fetchSize = fetchSize;
+ }
+
+ public void setQuery(String query) {
+ this.query = query;
+ }
+ }
+
+ @Override
+ public PTransform<PBegin, PCollection<Row>> buildExternal(Configuration configuration) {
+ JdbcIO.ReadRows readRows =
+ JdbcIO.readRows()
+ .withDataSourceConfiguration(configuration.getDataSourceConfiguration())
+ .withQuery(configuration.query);
+ if (configuration.fetchSize != null) {
+ readRows = readRows.withFetchSize(configuration.fetchSize);
+ }
+ if (configuration.outputParallelization != null) {
+ readRows = readRows.withOutputParallelization(configuration.outputParallelization);
+ }
+ return readRows;
+ }
+ }
+
+ public static class WriteBuilder
+ implements ExternalTransformBuilder<WriteBuilder.Configuration, PCollection<Row>, PDone> {
+
+ public static class Configuration extends CrossLanguageConfiguration {
+ private String statement;
+
+ public void setStatement(String statement) {
+ this.statement = statement;
+ }
+ }
+
+ @Override
+ public PTransform<PCollection<Row>, PDone> buildExternal(Configuration configuration) {
+ DataSourceConfiguration dataSourceConfiguration = configuration.getDataSourceConfiguration();
+
+ // TODO: BEAM-10396 use writeRows() when it's available
+ return JdbcIO.<Row>write()
+ .withDataSourceConfiguration(dataSourceConfiguration)
+ .withStatement(configuration.statement)
+ .withPreparedStatementSetter(new JdbcUtil.BeamRowPreparedStatementSetter());
+ }
+ }
+}
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcWriteRegistrar.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcWriteRegistrar.java
deleted file mode 100644
index 39997e0..0000000
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcWriteRegistrar.java
+++ /dev/null
@@ -1,67 +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.sdk.io.jdbc;
-
-import com.google.auto.service.AutoService;
-import java.util.Map;
-import org.apache.beam.sdk.annotations.Experimental;
-import org.apache.beam.sdk.annotations.Experimental.Kind;
-import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
-import org.apache.beam.sdk.io.jdbc.JdbcIO.DataSourceConfiguration;
-import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PDone;
-import org.apache.beam.sdk.values.Row;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
-
-/** Exposes {@link JdbcIO.Write} as an external transform for cross-language usage. */
-@Experimental(Kind.PORTABILITY)
-@AutoService(ExternalTransformRegistrar.class)
-public class JdbcWriteRegistrar implements ExternalTransformRegistrar {
-
- public static final String URN = "beam:external:java:jdbc:write:v1";
-
- @Override
- public Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders() {
- return ImmutableMap.of(URN, JdbcWriteRegistrar.Builder.class);
- }
-
- /** Parameters class to expose the Write transform to an external SDK. */
- public static class WriteConfiguration extends CrossLanguageConfiguration {
- private String statement;
-
- public void setStatement(String statement) {
- this.statement = statement;
- }
- }
-
- public static class Builder
- implements ExternalTransformBuilder<WriteConfiguration, PCollection<Row>, PDone> {
- @Override
- public PTransform<PCollection<Row>, PDone> buildExternal(WriteConfiguration configuration) {
- DataSourceConfiguration dataSourceConfiguration = configuration.getDataSourceConfiguration();
-
- // TODO: BEAM-10396 use writeRows() when it's available
- return JdbcIO.<Row>write()
- .withDataSourceConfiguration(dataSourceConfiguration)
- .withStatement(configuration.statement)
- .withPreparedStatementSetter(new JdbcUtil.BeamRowPreparedStatementSetter());
- }
- }
-}
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 13aabc8..08847f6 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -452,9 +452,15 @@
// Set required defaults
setTopicPartitions(Collections.emptyList());
setConsumerFactoryFn(Read.KAFKA_CONSUMER_FACTORY_FN);
- setMaxNumRecords(Long.MAX_VALUE);
+ if (config.maxReadTime != null) {
+ setMaxReadTime(Duration.standardSeconds(config.maxReadTime));
+ }
+ setMaxNumRecords(config.maxNumRecords == null ? Long.MAX_VALUE : config.maxNumRecords);
setCommitOffsetsInFinalizeEnabled(false);
setTimestampPolicyFactory(TimestampPolicyFactory.withProcessingTime());
+ if (config.startReadTime != null) {
+ setStartReadTime(Instant.ofEpochMilli(config.startReadTime));
+ }
// We do not include Metadata until we can encode KafkaRecords cross-language
return build().withoutMetadata();
}
@@ -507,6 +513,9 @@
private Iterable<String> topics;
private String keyDeserializer;
private String valueDeserializer;
+ private Long startReadTime;
+ private Long maxNumRecords;
+ private Long maxReadTime;
public void setConsumerConfig(Iterable<KV<String, String>> consumerConfig) {
this.consumerConfig = consumerConfig;
@@ -523,6 +532,18 @@
public void setValueDeserializer(String valueDeserializer) {
this.valueDeserializer = valueDeserializer;
}
+
+ public void setStartReadTime(Long startReadTime) {
+ this.startReadTime = startReadTime;
+ }
+
+ public void setMaxNumRecords(Long maxNumRecords) {
+ this.maxNumRecords = maxNumRecords;
+ }
+
+ public void setMaxReadTime(Long maxReadTime) {
+ this.maxReadTime = maxReadTime;
+ }
}
}
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 d157c16..3e44c17 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
@@ -34,6 +34,7 @@
import org.apache.beam.sdk.coders.IterableCoder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.expansion.service.ExpansionService;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Impulse;
@@ -67,6 +68,7 @@
.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer)
.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer)
.build();
+ Long startReadTime = 100L;
ExternalTransforms.ExternalConfigurationPayload payload =
ExternalTransforms.ExternalConfigurationPayload.newBuilder()
@@ -98,6 +100,12 @@
.addCoderUrn("beam:coder:string_utf8:v1")
.setPayload(ByteString.copyFrom(encodeString(valueDeserializer)))
.build())
+ .putConfiguration(
+ "start_read_time",
+ ExternalTransforms.ConfigValue.newBuilder()
+ .addCoderUrn("beam:coder:varint:v1")
+ .setPayload(ByteString.copyFrom(encodeLong(startReadTime)))
+ .build())
.build();
RunnerApi.Components defaultInstance = RunnerApi.Components.getDefaultInstance();
@@ -280,6 +288,12 @@
return baos.toByteArray();
}
+ private static byte[] encodeLong(Long str) throws IOException {
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ VarLongCoder.of().encode(str, baos);
+ return baos.toByteArray();
+ }
+
private static class TestStreamObserver<T> implements StreamObserver<T> {
private T result;
diff --git a/sdks/java/io/snowflake/build.gradle b/sdks/java/io/snowflake/build.gradle
index 32ad7af..ea8204d 100644
--- a/sdks/java/io/snowflake/build.gradle
+++ b/sdks/java/io/snowflake/build.gradle
@@ -22,6 +22,7 @@
automaticModuleName: 'org.apache.beam.sdk.io.snowflake')
provideIntegrationTestingDependencies()
enableJavaPerformanceTesting()
+
description = "Apache Beam :: SDKs :: Java :: IO :: Snowflake"
ext.summary = "IO to read and write on Snowflake."
dependencies {
@@ -31,6 +32,7 @@
compile library.java.slf4j_api
compile group: 'net.snowflake', name: 'snowflake-jdbc', version: '3.12.7'
compile group: 'com.opencsv', name: 'opencsv', version: '5.0'
+ compile 'net.snowflake:snowflake-ingest-sdk:0.9.9'
testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
testCompile project(path: ":sdks:java:io:common", configuration: "testRuntime")
testCompile project(path: ":sdks:java:testing:test-utils", configuration: "testRuntime")
diff --git a/sdks/java/io/snowflake/expansion-service/build.gradle b/sdks/java/io/snowflake/expansion-service/build.gradle
new file mode 100644
index 0000000..8a6ea6c
--- /dev/null
+++ b/sdks/java/io/snowflake/expansion-service/build.gradle
@@ -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.
+ */
+
+apply plugin: 'org.apache.beam.module'
+apply plugin: 'application'
+mainClassName = "org.apache.beam.sdk.expansion.service.ExpansionService"
+
+applyJavaNature(enableChecker:false,
+ automaticModuleName: 'org.apache.beam.sdk.io.expansion.service',
+ exportJavadoc: false,
+ validateShadowJar: false,
+ shadowClosure: {},
+)
+
+description = "Apache Beam :: SDKs :: Java :: IO :: Snowflake ::Expansion Service"
+ ext.summary = "Expansion service serving Snowflake IO"
+
+dependencies {
+ compile project(":sdks:java:expansion-service")
+ compile project(":sdks:java:io:snowflake")
+ runtime library.java.slf4j_jdk14
+}
+
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java
index 845f137..a2c0a7e 100644
--- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java
@@ -37,6 +37,8 @@
import java.util.stream.Collectors;
import javax.sql.DataSource;
import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
+import net.snowflake.ingest.SimpleIngestManager;
+import net.snowflake.ingest.connection.HistoryResponse;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.ListCoder;
@@ -51,13 +53,18 @@
import org.apache.beam.sdk.io.snowflake.credentials.KeyPairSnowflakeCredentials;
import org.apache.beam.sdk.io.snowflake.credentials.OAuthTokenSnowflakeCredentials;
import org.apache.beam.sdk.io.snowflake.credentials.SnowflakeCredentials;
+import org.apache.beam.sdk.io.snowflake.credentials.SnowflakeCredentialsFactory;
import org.apache.beam.sdk.io.snowflake.credentials.UsernamePasswordSnowflakeCredentials;
import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
+import org.apache.beam.sdk.io.snowflake.enums.StreamingLogLevel;
import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeBatchServiceConfig;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeBatchServiceImpl;
import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
-import org.apache.beam.sdk.io.snowflake.services.SnowflakeServiceConfig;
-import org.apache.beam.sdk.io.snowflake.services.SnowflakeServiceImpl;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeStreamingServiceConfig;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeStreamingServiceImpl;
+import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
@@ -73,12 +80,20 @@
import org.apache.beam.sdk.transforms.Wait;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.transforms.windowing.AfterFirst;
+import org.apache.beam.sdk.transforms.windowing.AfterPane;
+import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
+import org.apache.beam.sdk.transforms.windowing.Repeatedly;
+import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Joiner;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Splitter;
import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -92,9 +107,9 @@
*
* <p>To configure SnowflakeIO to read/write from your Snowflake instance, you have to provide a
* {@link DataSourceConfiguration} using {@link
- * DataSourceConfiguration#create(SnowflakeCredentials)}, where {@link SnowflakeCredentials might be
- * created using {@link org.apache.beam.sdk.io.snowflake.credentials.SnowflakeCredentialsFactory}}.
- * Additionally one of {@link DataSourceConfiguration#withServerName(String)} or {@link
+ * DataSourceConfiguration#create(SnowflakeCredentials)}, where {@link SnowflakeCredentials} might
+ * be created using {@link SnowflakeCredentialsFactory }. Additionally one of {@link
+ * DataSourceConfiguration#withServerName(String)} or {@link
* DataSourceConfiguration#withUrl(String)} must be used to tell SnowflakeIO which instance to use.
* <br>
* There are also other options available to configure connection to Snowflake:
@@ -147,8 +162,8 @@
*
* <h3>Writing to Snowflake</h3>
*
- * <p>SnowflakeIO.Write supports writing records into a database. It writes a {@link PCollection<T>}
- * to the database by converting each T into a {@link Object[]} via a user-provided {@link
+ * <p>SnowflakeIO.Write supports writing records into a database. It writes a {@link PCollection} to
+ * the database by converting each T into a {@link Object[]} via a user-provided {@link
* UserDataMapper}.
*
* <p>For example
@@ -171,7 +186,13 @@
private static final Logger LOG = LoggerFactory.getLogger(SnowflakeIO.class);
private static final String CSV_QUOTE_CHAR = "'";
- private static final String WRITE_TMP_PATH = "data";
+
+ static final int DEFAULT_FLUSH_ROW_LIMIT = 10000;
+ static final int DEFAULT_STREAMING_SHARDS_NUMBER = 1;
+ static final int DEFAULT_BATCH_SHARDS_NUMBER = 0;
+ static final Duration DEFAULT_FLUSH_TIME_LIMIT = Duration.millis(30000); // 30 seconds
+ static final Duration DEFAULT_STREAMING_LOGS_MAX_SLEEP = Duration.standardMinutes(2);
+ static final Duration DEFAULT_SLEEP_STREAMING_LOGS = Duration.standardSeconds(5000);
/**
* Read data from Snowflake.
@@ -182,6 +203,7 @@
public static <T> Read<T> read(SnowflakeService snowflakeService) {
return new AutoValue_SnowflakeIO_Read.Builder<T>()
.setSnowflakeService(snowflakeService)
+ .setQuotationMark(CSV_QUOTE_CHAR)
.build();
}
@@ -191,7 +213,7 @@
* @param <T> Type of the data to be read.
*/
public static <T> Read<T> read() {
- return read(new SnowflakeServiceImpl());
+ return read(new SnowflakeBatchServiceImpl());
}
/**
@@ -226,6 +248,10 @@
.setFileNameTemplate("output")
.setCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
.setWriteDisposition(WriteDisposition.APPEND)
+ .setFlushTimeLimit(DEFAULT_FLUSH_TIME_LIMIT)
+ .setShardsNumber(DEFAULT_BATCH_SHARDS_NUMBER)
+ .setFlushRowLimit(DEFAULT_FLUSH_ROW_LIMIT)
+ .setQuotationMark(CSV_QUOTE_CHAR)
.build();
}
@@ -249,6 +275,9 @@
abstract @Nullable SnowflakeService getSnowflakeService();
+ @Nullable
+ abstract String getQuotationMark();
+
abstract Builder<T> toBuilder();
@AutoValue.Builder
@@ -270,13 +299,15 @@
abstract Builder<T> setSnowflakeService(SnowflakeService snowflakeService);
+ abstract Builder<T> setQuotationMark(String quotationMark);
+
abstract Read<T> build();
}
/**
* Setting information about Snowflake server.
*
- * @param config - An instance of {@link DataSourceConfiguration}.
+ * @param config An instance of {@link DataSourceConfiguration}.
*/
public Read<T> withDataSourceConfiguration(final DataSourceConfiguration config) {
return withDataSourceProviderFn(new DataSourceProviderFromDataSourceConfiguration(config));
@@ -295,7 +326,7 @@
/**
* A query to be executed in Snowflake.
*
- * @param query - String with query.
+ * @param query String with query.
*/
public Read<T> fromQuery(String query) {
return toBuilder().setQuery(query).build();
@@ -304,7 +335,7 @@
/**
* A table name to be read in Snowflake.
*
- * @param table - String with the name of the table.
+ * @param table String with the name of the table.
*/
public Read<T> fromTable(String table) {
return toBuilder().setTable(table).build();
@@ -313,9 +344,12 @@
/**
* Name of the cloud bucket (GCS by now) to use as tmp location of CSVs during COPY statement.
*
- * @param stagingBucketName - String with the name of the bucket.
+ * @param stagingBucketName String with the name of the bucket.
*/
public Read<T> withStagingBucketName(String stagingBucketName) {
+ checkArgument(
+ stagingBucketName.endsWith("/"),
+ "stagingBucketName must be a cloud storage path ending with /");
return toBuilder().setStagingBucketName(stagingBucketName).build();
}
@@ -324,7 +358,7 @@
* https://docs.snowflake.com/en/sql-reference/sql/create-storage-integration.html for
* reference.
*
- * @param integrationName - String with the name of the Storage Integration.
+ * @param integrationName String with the name of the Storage Integration.
*/
public Read<T> withStorageIntegrationName(String integrationName) {
return toBuilder().setStorageIntegrationName(integrationName).build();
@@ -333,7 +367,7 @@
/**
* User-defined function mapping CSV lines into user data.
*
- * @param csvMapper - an instance of {@link CsvMapper}.
+ * @param csvMapper an instance of {@link CsvMapper}.
*/
public Read<T> withCsvMapper(CsvMapper<T> csvMapper) {
return toBuilder().setCsvMapper(csvMapper).build();
@@ -342,21 +376,29 @@
/**
* A Coder to be used by the output PCollection generated by the source.
*
- * @param coder - an instance of {@link Coder}.
+ * @param coder an instance of {@link Coder}.
*/
public Read<T> withCoder(Coder<T> coder) {
return toBuilder().setCoder(coder).build();
}
+ /**
+ * Sets Snowflake-specific quotations around strings.
+ *
+ * @param quotationMark with possible single quote {@code '}, double quote {@code "} or nothing.
+ * Default value is single quotation {@code '}.
+ * @return
+ */
+ public Read<T> withQuotationMark(String quotationMark) {
+ return toBuilder().setQuotationMark(quotationMark).build();
+ }
+
@Override
public PCollection<T> expand(PBegin input) {
checkArguments();
- String tmpDirName = makeTmpDirName();
- String stagingBucketDir = String.format("%s/%s/", getStagingBucketName(), tmpDirName);
-
PCollection<Void> emptyCollection = input.apply(Create.of((Void) null));
-
+ String tmpDirName = makeTmpDirName();
PCollection<T> output =
emptyCollection
.apply(
@@ -366,20 +408,22 @@
getQuery(),
getTable(),
getStorageIntegrationName(),
- stagingBucketDir,
- getSnowflakeService())))
+ getStagingBucketName(),
+ tmpDirName,
+ getSnowflakeService(),
+ getQuotationMark())))
.apply(Reshuffle.viaRandomKey())
.apply(FileIO.matchAll())
.apply(FileIO.readMatches())
.apply(readFiles())
- .apply(ParDo.of(new MapCsvToStringArrayFn()))
+ .apply(ParDo.of(new MapCsvToStringArrayFn(getQuotationMark())))
.apply(ParDo.of(new MapStringArrayToUserDataFn<>(getCsvMapper())));
output.setCoder(getCoder());
emptyCollection
.apply(Wait.on(output))
- .apply(ParDo.of(new CleanTmpFilesFromGcsFn(stagingBucketDir)));
+ .apply(ParDo.of(new CleanTmpFilesFromGcsFn(getStagingBucketName(), tmpDirName)));
return output;
}
@@ -387,8 +431,9 @@
// Either table or query is required. If query is present, it's being used, table is used
// otherwise
- checkArgument(getStorageIntegrationName() != null, "withStorageIntegrationName is required");
- checkArgument(getStagingBucketName() != null, "withStagingBucketName is required");
+ checkArgument(
+ getStorageIntegrationName() != null, "withStorageIntegrationName() is required");
+ checkArgument(getStagingBucketName() != null, "withStagingBucketName() is required");
checkArgument(
getQuery() != null || getTable() != null, "fromTable() or fromQuery() is required");
@@ -414,9 +459,13 @@
private final SerializableFunction<Void, DataSource> dataSourceProviderFn;
private final String query;
private final String table;
+ private final String database;
+ private final String schema;
+ private final String tmpDirName;
private final String storageIntegrationName;
private final String stagingBucketDir;
private final SnowflakeService snowflakeService;
+ private final String quotationMark;
private CopyIntoStageFn(
SerializableFunction<Void, DataSource> dataSourceProviderFn,
@@ -424,23 +473,44 @@
String table,
String storageIntegrationName,
String stagingBucketDir,
- SnowflakeService snowflakeService) {
+ String tmpDirName,
+ SnowflakeService snowflakeService,
+ String quotationMark) {
this.dataSourceProviderFn = dataSourceProviderFn;
this.query = query;
this.table = table;
this.storageIntegrationName = storageIntegrationName;
- this.stagingBucketDir =
- String.format(
- "%s/run_%s/", stagingBucketDir, UUID.randomUUID().toString().subSequence(0, 8));
this.snowflakeService = snowflakeService;
+ this.quotationMark = quotationMark;
+ this.stagingBucketDir = stagingBucketDir;
+ this.tmpDirName = tmpDirName;
+ DataSourceProviderFromDataSourceConfiguration
+ dataSourceProviderFromDataSourceConfiguration =
+ (DataSourceProviderFromDataSourceConfiguration) this.dataSourceProviderFn;
+ DataSourceConfiguration config = dataSourceProviderFromDataSourceConfiguration.getConfig();
+
+ this.database = config.getDatabase();
+ this.schema = config.getSchema();
}
@ProcessElement
public void processElement(ProcessContext context) throws Exception {
- SnowflakeServiceConfig config =
- new SnowflakeServiceConfig(
- dataSourceProviderFn, table, query, storageIntegrationName, stagingBucketDir);
+ String stagingBucketRunDir =
+ String.format(
+ "%s/%s/run_%s/",
+ stagingBucketDir, tmpDirName, UUID.randomUUID().toString().subSequence(0, 8));
+
+ SnowflakeBatchServiceConfig config =
+ new SnowflakeBatchServiceConfig(
+ dataSourceProviderFn,
+ database,
+ schema,
+ table,
+ query,
+ storageIntegrationName,
+ stagingBucketRunDir,
+ quotationMark);
String output = snowflakeService.read(config);
@@ -448,11 +518,21 @@
}
}
+ /**
+ * Parses {@code String} from incoming data in {@link PCollection} to have proper format for CSV
+ * files.
+ */
public static class MapCsvToStringArrayFn extends DoFn<String, String[]> {
+ private String quoteChar;
+
+ public MapCsvToStringArrayFn(String quoteChar) {
+ this.quoteChar = quoteChar;
+ }
+
@ProcessElement
public void processElement(ProcessContext c) throws IOException {
String csvLine = c.element();
- CSVParser parser = new CSVParserBuilder().withQuoteChar(CSV_QUOTE_CHAR.charAt(0)).build();
+ CSVParser parser = new CSVParserBuilder().withQuoteChar(quoteChar.charAt(0)).build();
String[] parts = parser.parseLine(csvLine);
c.output(parts);
}
@@ -471,16 +551,25 @@
}
}
+ /** Removes temporary staged files after reading. */
public static class CleanTmpFilesFromGcsFn extends DoFn<Object, Object> {
private final String stagingBucketDir;
+ private final String tmpDirName;
- public CleanTmpFilesFromGcsFn(String stagingBucketDir) {
+ /**
+ * Created object that will remove temp files from stage.
+ *
+ * @param stagingBucketDir bucket and directory where temporary files are saved
+ * @param tmpDirName temporary directory created on bucket where files were saved
+ */
+ public CleanTmpFilesFromGcsFn(String stagingBucketDir, String tmpDirName) {
this.stagingBucketDir = stagingBucketDir;
+ this.tmpDirName = tmpDirName;
}
@ProcessElement
public void processElement(ProcessContext c) throws IOException {
- String combinedPath = stagingBucketDir + "/**";
+ String combinedPath = String.format("%s/%s/**", stagingBucketDir, tmpDirName);
List<ResourceId> paths =
FileSystems.match(combinedPath).metadata().stream()
.map(metadata -> metadata.resourceId())
@@ -521,10 +610,18 @@
abstract @Nullable String getStagingBucketName();
- abstract @Nullable String getQuery();
+ abstract @Nullable ValueProvider<String> getSnowPipe();
+
+ abstract @Nullable Integer getFlushRowLimit();
+
+ abstract @Nullable Integer getShardsNumber();
+
+ abstract @Nullable Duration getFlushTimeLimit();
abstract @Nullable String getFileNameTemplate();
+ abstract @Nullable String getQuery();
+
abstract @Nullable WriteDisposition getWriteDisposition();
abstract @Nullable CreateDisposition getCreateDisposition();
@@ -535,6 +632,12 @@
abstract @Nullable SnowflakeService getSnowflakeService();
+ @Nullable
+ abstract String getQuotationMark();
+
+ @Nullable
+ abstract StreamingLogLevel getDebugMode();
+
abstract Builder<T> toBuilder();
@AutoValue.Builder
@@ -550,6 +653,14 @@
abstract Builder<T> setQuery(String query);
+ abstract Builder<T> setSnowPipe(ValueProvider<String> snowPipe);
+
+ abstract Builder<T> setFlushRowLimit(Integer rowsCount);
+
+ abstract Builder<T> setShardsNumber(Integer shardsNumber);
+
+ abstract Builder<T> setFlushTimeLimit(Duration triggeringFrequency);
+
abstract Builder<T> setFileNameTemplate(String fileNameTemplate);
abstract Builder<T> setUserDataMapper(UserDataMapper userDataMapper);
@@ -562,13 +673,17 @@
abstract Builder<T> setSnowflakeService(SnowflakeService snowflakeService);
+ abstract Builder<T> setQuotationMark(String quotationMark);
+
+ abstract Builder<T> setDebugMode(StreamingLogLevel debugLevel);
+
abstract Write<T> build();
}
/**
* Setting information about Snowflake server.
*
- * @param config - An instance of {@link DataSourceConfiguration}.
+ * @param config An instance of {@link DataSourceConfiguration}.
*/
public Write<T> withDataSourceConfiguration(final DataSourceConfiguration config) {
return withDataSourceProviderFn(new DataSourceProviderFromDataSourceConfiguration(config));
@@ -587,7 +702,7 @@
/**
* A table name to be written in Snowflake.
*
- * @param table - String with the name of the table.
+ * @param table String with the name of the table.
*/
public Write<T> to(String table) {
return toBuilder().setTable(table).build();
@@ -596,9 +711,12 @@
/**
* Name of the cloud bucket (GCS by now) to use as tmp location of CSVs during COPY statement.
*
- * @param stagingBucketName - String with the name of the bucket.
+ * @param stagingBucketName String with the name of the bucket.
*/
public Write<T> withStagingBucketName(String stagingBucketName) {
+ checkArgument(
+ stagingBucketName.endsWith("/"),
+ "stagingBucketName must be a cloud storage path ending with /");
return toBuilder().setStagingBucketName(stagingBucketName).build();
}
@@ -607,7 +725,7 @@
* https://docs.snowflake.com/en/sql-reference/sql/create-storage-integration.html for
* reference.
*
- * @param integrationName - String with the name of the Storage Integration.
+ * @param integrationName String with the name of the Storage Integration.
*/
public Write<T> withStorageIntegrationName(String integrationName) {
return toBuilder().setStorageIntegrationName(integrationName).build();
@@ -616,7 +734,7 @@
/**
* A query to be executed in Snowflake.
*
- * @param query - String with query.
+ * @param query String with query.
*/
public Write<T> withQueryTransformation(String query) {
return toBuilder().setQuery(query).build();
@@ -625,7 +743,7 @@
/**
* A template name for files saved to GCP.
*
- * @param fileNameTemplate - String with template name for files.
+ * @param fileNameTemplate String with template name for files.
*/
public Write<T> withFileNameTemplate(String fileNameTemplate) {
return toBuilder().setFileNameTemplate(fileNameTemplate).build();
@@ -634,16 +752,96 @@
/**
* User-defined function mapping user data into CSV lines.
*
- * @param userDataMapper - an instance of {@link UserDataMapper}.
+ * @param userDataMapper an instance of {@link UserDataMapper}.
*/
public Write<T> withUserDataMapper(UserDataMapper userDataMapper) {
return toBuilder().setUserDataMapper(userDataMapper).build();
}
/**
+ * Sets duration how often staged files will be created and then how often ingested by Snowflake
+ * during streaming.
+ *
+ * @param triggeringFrequency time for triggering frequency in {@link Duration} type.
+ * @return
+ */
+ public Write<T> withFlushTimeLimit(Duration triggeringFrequency) {
+ return toBuilder().setFlushTimeLimit(triggeringFrequency).build();
+ }
+
+ /**
+ * Sets name of <a
+ * href="https://docs.snowflake.com/en/user-guide/data-load-snowpipe-intro.html">SnowPipe</a>
+ * which can be created in Snowflake dashboard or cli:
+ *
+ * <pre>{@code
+ * CREATE snowPipeName AS COPY INTO your_table from @yourstage;
+ * }</pre>
+ *
+ * <p>The stage in <a
+ * href="https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html">COPY</a>
+ * statement should be pointing to the cloud <a
+ * href="https://docs.snowflake.com/en/sql-reference/sql/create-storage-integration.html">integration</a>
+ * with the valid bucket url, ex. for GCS:
+ *
+ * <pre>{@code
+ * CREATE STAGE yourstage
+ * URL = 'gcs://yourbucket/path/'
+ * STORAGE_INTEGRATION = your_integration;
+ * }</pre>
+ *
+ * <pre>{@code
+ * CREATE STORAGE INTEGRATION your_integration
+ * TYPE = EXTERNAL_STAGE
+ * STORAGE_PROVIDER = GCS
+ * ENABLED = TRUE
+ * STORAGE_ALLOWED_LOCATIONS = ('gcs://yourbucket/path/')
+ * }</pre>
+ *
+ * @param snowPipe name of created SnowPipe in Snowflake dashboard.
+ * @return
+ */
+ public Write<T> withSnowPipe(String snowPipe) {
+ return toBuilder().setSnowPipe(ValueProvider.StaticValueProvider.of(snowPipe)).build();
+ }
+
+ /**
+ * Same as {@code withSnowPipe(String}, but with a {@link ValueProvider}.
+ *
+ * @param snowPipe name of created SnowPipe in Snowflake dashboard.
+ * @return
+ */
+ public Write<T> withSnowPipe(ValueProvider<String> snowPipe) {
+ return toBuilder().setSnowPipe(snowPipe).build();
+ }
+
+ /**
+ * Number of shards that are created per window.
+ *
+ * @param shardsNumber defined number of shards or 1 by default.
+ * @return
+ */
+ public Write<T> withShardsNumber(Integer shardsNumber) {
+ return toBuilder().setShardsNumber(shardsNumber).build();
+ }
+
+ /**
+ * Sets number of row limit that will be saved to the staged file and then loaded to Snowflake.
+ * If the number of rows will be lower than the limit it will be loaded with current number of
+ * rows after certain time specified by setting {@code withFlushTimeLimit(Duration
+ * triggeringFrequency)}
+ *
+ * @param rowsCount Number of rows that will be in one file staged for loading. Default: 10000.
+ * @return
+ */
+ public Write<T> withFlushRowLimit(Integer rowsCount) {
+ return toBuilder().setFlushRowLimit(rowsCount).build();
+ }
+
+ /**
* A disposition to be used during writing to table phase.
*
- * @param writeDisposition - an instance of {@link WriteDisposition}.
+ * @param writeDisposition an instance of {@link WriteDisposition}.
*/
public Write<T> withWriteDisposition(WriteDisposition writeDisposition) {
return toBuilder().setWriteDisposition(writeDisposition).build();
@@ -668,28 +866,56 @@
}
/**
- * A snowflake service which is supposed to be used. Note: Currently we have {@link
- * SnowflakeServiceImpl} with corresponding {@link FakeSnowflakeServiceImpl} used for testing.
+ * A snowflake service {@link SnowflakeService} implementation which is supposed to be used.
*
- * @param snowflakeService - an instance of {@link SnowflakeService}.
+ * @param snowflakeService an instance of {@link SnowflakeService}.
*/
public Write<T> withSnowflakeService(SnowflakeService snowflakeService) {
return toBuilder().setSnowflakeService(snowflakeService).build();
}
+ /**
+ * Sets Snowflake-specific quotations around strings.
+ *
+ * @param quotationMark with possible single quote {@code '}, double quote {@code "} or nothing.
+ * Default value is single quotation {@code '}.
+ * @return
+ */
+ public Write<T> withQuotationMark(String quotationMark) {
+ return toBuilder().setQuotationMark(quotationMark).build();
+ }
+
+ /**
+ * The option to verbose info (or only errors) of loaded files while streaming. It is not set by
+ * default because it may influence performance. For details: <a
+ * href="https://docs.snowflake.com/en/user-guide/data-load-snowpipe-rest-apis.html#endpoint-insertreport">insert
+ * report REST API.</a>
+ *
+ * @param debugLevel error or info debug level from enum {@link StreamingLogLevel}
+ * @return
+ */
+ public Write<T> withDebugMode(StreamingLogLevel debugLevel) {
+ return toBuilder().setDebugMode(debugLevel).build();
+ }
+
@Override
public PDone expand(PCollection<T> input) {
- checkArguments();
+ checkArguments(input);
- String stagingBucketDir = String.format("%s/%s/", getStagingBucketName(), WRITE_TMP_PATH);
+ PCollection out;
- PCollection<String> out = write(input, stagingBucketDir);
+ if (getSnowPipe() != null) {
+ out = writeStream(input, getStagingBucketName());
+ } else {
+ out = writeBatch(input, getStagingBucketName());
+ }
+
out.setCoder(StringUtf8Coder.of());
return PDone.in(out.getPipeline());
}
- private void checkArguments() {
+ private void checkArguments(PCollection<T> input) {
checkArgument(getStagingBucketName() != null, "withStagingBucketName is required");
checkArgument(getUserDataMapper() != null, "withUserDataMapper() is required");
@@ -698,14 +924,65 @@
(getDataSourceProviderFn() != null),
"withDataSourceConfiguration() or withDataSourceProviderFn() is required");
- checkArgument(getTable() != null, "to() is required");
+ if (input.isBounded() == PCollection.IsBounded.UNBOUNDED) {
+ checkArgument(
+ getSnowPipe() != null,
+ "in streaming (unbounded) write it is required to specify SnowPipe name via withSnowPipe() method.");
+ } else {
+ checkArgument(
+ getTable() != null,
+ "in batch writing it is required to specify destination table name via to() method.");
+ }
}
- private PCollection<String> write(PCollection<T> input, String stagingBucketDir) {
+ private PCollection<T> writeStream(PCollection<T> input, String stagingBucketDir) {
SnowflakeService snowflakeService =
- getSnowflakeService() != null ? getSnowflakeService() : new SnowflakeServiceImpl();
+ getSnowflakeService() != null
+ ? getSnowflakeService()
+ : new SnowflakeStreamingServiceImpl();
- PCollection<String> files = writeFiles(input, stagingBucketDir);
+ /* Ensure that files will be created after specific record count or duration specified */
+ PCollection<T> inputInGlobalWindow =
+ input.apply(
+ "Rewindow Into Global",
+ Window.<T>into(new GlobalWindows())
+ .triggering(
+ Repeatedly.forever(
+ AfterFirst.of(
+ AfterProcessingTime.pastFirstElementInPane()
+ .plusDelayOf(getFlushTimeLimit()),
+ AfterPane.elementCountAtLeast(getFlushRowLimit()))))
+ .discardingFiredPanes());
+
+ int shards = (getShardsNumber() > 0) ? getShardsNumber() : DEFAULT_STREAMING_SHARDS_NUMBER;
+ PCollection files = writeFiles(inputInGlobalWindow, stagingBucketDir, shards);
+
+ /* Ensuring that files will be ingested after flush time */
+ files =
+ (PCollection)
+ files.apply(
+ "Apply User Trigger",
+ Window.<T>into(new GlobalWindows())
+ .triggering(
+ Repeatedly.forever(
+ AfterProcessingTime.pastFirstElementInPane()
+ .plusDelayOf(getFlushTimeLimit())))
+ .discardingFiredPanes());
+ files =
+ (PCollection)
+ files.apply(
+ "Create list of files for loading via SnowPipe",
+ Combine.globally(new Concatenate()).withoutDefaults());
+
+ return (PCollection)
+ files.apply("Stream files to table", streamToTable(snowflakeService, stagingBucketDir));
+ }
+
+ private PCollection writeBatch(PCollection input, String stagingBucketDir) {
+ SnowflakeService snowflakeService =
+ getSnowflakeService() != null ? getSnowflakeService() : new SnowflakeBatchServiceImpl();
+
+ PCollection<String> files = writeBatchFiles(input, stagingBucketDir);
// Combining PCollection of files as a side input into one list of files
ListCoder<String> coder = ListCoder.of(StringUtf8Coder.of());
@@ -721,7 +998,12 @@
files.apply("Copy files to table", copyToTable(snowflakeService, stagingBucketDir));
}
- private PCollection<String> writeFiles(PCollection<T> input, String stagingBucketDir) {
+ private PCollection writeBatchFiles(PCollection<T> input, String outputDirectory) {
+ return writeFiles(input, outputDirectory, DEFAULT_BATCH_SHARDS_NUMBER);
+ }
+
+ private PCollection<String> writeFiles(
+ PCollection<T> input, String stagingBucketDir, int numShards) {
PCollection<String> mappedUserData =
input
@@ -733,7 +1015,9 @@
return getUserDataMapper().mapRow(element);
}
}))
- .apply("Map Objects array to CSV lines", ParDo.of(new MapObjectsArrayToCsvFn()))
+ .apply(
+ "Map Objects array to CSV lines",
+ ParDo.of(new MapObjectsArrayToCsvFn(getQuotationMark())))
.setCoder(StringUtf8Coder.of());
WriteFilesResult filesResult =
@@ -742,8 +1026,9 @@
FileIO.<String>write()
.via(TextIO.sink())
.to(stagingBucketDir)
- .withPrefix(getFileNameTemplate())
+ .withPrefix(UUID.randomUUID().toString().subSequence(0, 8).toString())
.withSuffix(".csv")
+ .withNumShards(numShards)
.withCompression(Compression.GZIP));
return (PCollection)
@@ -764,10 +1049,25 @@
getCreateDisposition(),
getWriteDisposition(),
getTableSchema(),
+ snowflakeService,
+ getQuotationMark()));
+ }
+
+ protected PTransform streamToTable(SnowflakeService snowflakeService, String stagingBucketDir) {
+ return ParDo.of(
+ new StreamToTableFn(
+ getDataSourceProviderFn(),
+ getSnowPipe(),
+ stagingBucketDir,
+ getDebugMode(),
snowflakeService));
}
}
+ /**
+ * Combines list of {@code String} to provide one {@code String} with paths where files were
+ * staged for write.
+ */
public static class Concatenate extends Combine.CombineFn<String, List<String>, List<String>> {
@Override
public List<String> createAccumulator() {
@@ -801,6 +1101,11 @@
* <p>Adds Snowflake-specific quotations around strings.
*/
private static class MapObjectsArrayToCsvFn extends DoFn<Object[], String> {
+ private String quotationMark;
+
+ public MapObjectsArrayToCsvFn(String quotationMark) {
+ this.quotationMark = quotationMark;
+ }
@ProcessElement
public void processElement(ProcessContext context) {
@@ -820,7 +1125,7 @@
}
private String quoteField(String field) {
- return quoteField(field, CSV_QUOTE_CHAR);
+ return quoteField(field, this.quotationMark);
}
private String quoteField(String field, String quotation) {
@@ -830,11 +1135,14 @@
private static class CopyToTableFn<ParameterT, OutputT> extends DoFn<ParameterT, OutputT> {
private final SerializableFunction<Void, DataSource> dataSourceProviderFn;
+ private final String database;
+ private final String schema;
private final String table;
private final String query;
private final SnowflakeTableSchema tableSchema;
private final String stagingBucketDir;
private final String storageIntegrationName;
+ private final String quotationMark;
private final WriteDisposition writeDisposition;
private final CreateDisposition createDisposition;
private final SnowflakeService snowflakeService;
@@ -848,35 +1156,156 @@
CreateDisposition createDisposition,
WriteDisposition writeDisposition,
SnowflakeTableSchema tableSchema,
- SnowflakeService snowflakeService) {
+ SnowflakeService snowflakeService,
+ String quotationMark) {
this.dataSourceProviderFn = dataSourceProviderFn;
- this.table = table;
this.query = query;
+ this.table = table;
this.stagingBucketDir = stagingBucketDir;
this.storageIntegrationName = storageIntegrationName;
this.writeDisposition = writeDisposition;
this.createDisposition = createDisposition;
this.tableSchema = tableSchema;
this.snowflakeService = snowflakeService;
+ this.quotationMark = quotationMark;
+
+ DataSourceProviderFromDataSourceConfiguration dataSourceProviderFromDataSourceConfiguration =
+ (DataSourceProviderFromDataSourceConfiguration) this.dataSourceProviderFn;
+ DataSourceConfiguration config = dataSourceProviderFromDataSourceConfiguration.getConfig();
+
+ this.database = config.getDatabase();
+ this.schema = config.getSchema();
}
@ProcessElement
public void processElement(ProcessContext context) throws Exception {
- SnowflakeServiceConfig config =
- new SnowflakeServiceConfig(
+ SnowflakeBatchServiceConfig config =
+ new SnowflakeBatchServiceConfig(
dataSourceProviderFn,
(List<String>) context.element(),
+ database,
+ schema,
table,
query,
tableSchema,
createDisposition,
writeDisposition,
storageIntegrationName,
- stagingBucketDir);
+ stagingBucketDir,
+ quotationMark);
snowflakeService.write(config);
}
}
+ /** Custom DoFn that streams data to Snowflake table. */
+ private static class StreamToTableFn<ParameterT, OutputT> extends DoFn<ParameterT, OutputT> {
+ private final SerializableFunction<Void, DataSource> dataSourceProviderFn;
+ private final String stagingBucketDir;
+ private final ValueProvider<String> snowPipe;
+ private final StreamingLogLevel debugMode;
+ private final SnowflakeService snowflakeService;
+ private transient SimpleIngestManager ingestManager;
+
+ private transient DataSource dataSource;
+ ArrayList<String> trackedFilesNames;
+
+ StreamToTableFn(
+ SerializableFunction<Void, DataSource> dataSourceProviderFn,
+ ValueProvider<String> snowPipe,
+ String stagingBucketDir,
+ StreamingLogLevel debugMode,
+ SnowflakeService snowflakeService) {
+ this.dataSourceProviderFn = dataSourceProviderFn;
+ this.stagingBucketDir = stagingBucketDir;
+ this.snowPipe = snowPipe;
+ this.debugMode = debugMode;
+ this.snowflakeService = snowflakeService;
+ trackedFilesNames = new ArrayList<>();
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ dataSource = dataSourceProviderFn.apply(null);
+
+ DataSourceProviderFromDataSourceConfiguration dataSourceProviderFromDataSourceConfiguration =
+ (DataSourceProviderFromDataSourceConfiguration) this.dataSourceProviderFn;
+ DataSourceConfiguration config = dataSourceProviderFromDataSourceConfiguration.getConfig();
+
+ checkArgument(config.getPrivateKey() != null, "KeyPair is required for authentication");
+
+ String hostName = config.getServerName();
+ List<String> path = Splitter.on('.').splitToList(hostName);
+ String account = path.get(0);
+ String username = config.getUsername();
+ PrivateKey privateKey = config.getPrivateKey();
+ String schema = config.getSchema();
+ String database = config.getDatabase();
+ String snowPipeName = String.format("%s.%s.%s", database, schema, snowPipe.get());
+
+ this.ingestManager =
+ new SimpleIngestManager(
+ account, username, snowPipeName, privateKey, "https", hostName, 443);
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext context) throws Exception {
+ List<String> filesList = (List<String>) context.element();
+
+ if (debugMode != null) {
+ trackedFilesNames.addAll(filesList);
+ }
+ SnowflakeStreamingServiceConfig config =
+ new SnowflakeStreamingServiceConfig(filesList, this.stagingBucketDir, this.ingestManager);
+ snowflakeService.write(config);
+ }
+
+ @FinishBundle
+ public void finishBundle() throws Exception {
+ if (debugMode != null) {
+ String beginMark = null;
+ Duration currentSleep = Duration.ZERO;
+
+ while (currentSleep.isShorterThan(DEFAULT_STREAMING_LOGS_MAX_SLEEP)
+ && trackedFilesNames.size() > 0) {
+ Thread.sleep(DEFAULT_SLEEP_STREAMING_LOGS.getMillis());
+ currentSleep = currentSleep.plus(DEFAULT_SLEEP_STREAMING_LOGS);
+ HistoryResponse response = ingestManager.getHistory(null, null, beginMark);
+
+ if (response != null && response.getNextBeginMark() != null) {
+ beginMark = response.getNextBeginMark();
+ }
+ if (response != null && response.files != null) {
+ response.files.forEach(
+ entry -> {
+ if (entry.getPath() != null && entry.isComplete()) {
+ String responseFileName =
+ String.format("'%s%s'", entry.getStageLocation(), entry.getPath())
+ .toLowerCase()
+ .replace("gcs://", "gs://");
+ if (trackedFilesNames.contains(responseFileName)) {
+ trackedFilesNames.remove(responseFileName);
+
+ if (entry.getErrorsSeen() > 0) {
+ LOG.error(String.format("Snowflake SnowPipe ERROR: %s", entry.toString()));
+ } else if (entry.getErrorsSeen() == 0
+ && debugMode.equals(StreamingLogLevel.INFO)) {
+ LOG.info(String.format("Snowflake SnowPipe INFO: %s", entry.toString()));
+ }
+ }
+ }
+ });
+ }
+ }
+ trackedFilesNames.forEach(
+ file -> LOG.info(String.format("File %s was not found in ingest history", file)));
+ }
+ }
+ }
+
+ private static String getValueOrNull(ValueProvider<String> valueProvider) {
+ return valueProvider != null ? valueProvider.get() : null;
+ }
+
/**
* A POJO describing a {@link DataSource}, providing all properties allowing to create a {@link
* DataSource}.
@@ -954,7 +1383,7 @@
/**
* Creates {@link DataSourceConfiguration} from existing instance of {@link DataSource}.
*
- * @param dataSource - an instance of {@link DataSource}.
+ * @param dataSource an instance of {@link DataSource}.
*/
public static DataSourceConfiguration create(DataSource dataSource) {
checkArgument(dataSource instanceof Serializable, "dataSource must be Serializable");
@@ -967,7 +1396,7 @@
/**
* Creates {@link DataSourceConfiguration} from instance of {@link SnowflakeCredentials}.
*
- * @param credentials - an instance of {@link SnowflakeCredentials}.
+ * @param credentials an instance of {@link SnowflakeCredentials}.
*/
public static DataSourceConfiguration create(SnowflakeCredentials credentials) {
if (credentials instanceof UsernamePasswordSnowflakeCredentials) {
@@ -998,7 +1427,7 @@
*
* <p>Either withUrl or withServerName is required.
*
- * @param url - String with URL of the Snowflake server.
+ * @param url String with URL of the Snowflake server.
*/
public DataSourceConfiguration withUrl(String url) {
checkArgument(
@@ -1013,7 +1442,7 @@
/**
* Sets database to use.
*
- * @param database - String with database name.
+ * @param database String with database name.
*/
public DataSourceConfiguration withDatabase(String database) {
return builder().setDatabase(database).build();
@@ -1022,7 +1451,7 @@
/**
* Sets Snowflake Warehouse to use.
*
- * @param warehouse - String with warehouse name.
+ * @param warehouse String with warehouse name.
*/
public DataSourceConfiguration withWarehouse(String warehouse) {
return builder().setWarehouse(warehouse).build();
@@ -1031,7 +1460,7 @@
/**
* Sets schema to use when connecting to Snowflake.
*
- * @param schema - String with schema name.
+ * @param schema String with schema name.
*/
public DataSourceConfiguration withSchema(String schema) {
return builder().setSchema(schema).build();
@@ -1043,7 +1472,7 @@
*
* <p>Either withServerName or withUrl is required.
*
- * @param serverName - String with server name.
+ * @param serverName String with server name.
*/
public DataSourceConfiguration withServerName(String serverName) {
checkArgument(
@@ -1055,7 +1484,7 @@
/**
* Sets port number to use to connect to Snowflake.
*
- * @param portNumber - Integer with port number.
+ * @param portNumber Integer with port number.
*/
public DataSourceConfiguration withPortNumber(Integer portNumber) {
return builder().setPortNumber(portNumber).build();
@@ -1064,16 +1493,16 @@
/**
* Sets user's role to be used when running queries on Snowflake.
*
- * @param role - String with role name.
+ * @param role String with role name.
*/
public DataSourceConfiguration withRole(String role) {
return builder().setRole(role).build();
}
/**
- * Sets loginTimeout that will be used in {@link SnowflakeBasicDataSource:setLoginTimeout}.
+ * Sets loginTimeout that will be used in {@link SnowflakeBasicDataSource#setLoginTimeout}.
*
- * @param loginTimeout - Integer with timeout value.
+ * @param loginTimeout Integer with timeout value.
*/
public DataSourceConfiguration withLoginTimeout(Integer loginTimeout) {
return builder().setLoginTimeout(loginTimeout).build();
@@ -1156,6 +1585,7 @@
}
}
+ /** Wraps {@link DataSourceConfiguration} to provide DataSource. */
public static class DataSourceProviderFromDataSourceConfiguration
implements SerializableFunction<Void, DataSource>, HasDisplayData {
private static final ConcurrentHashMap<DataSourceConfiguration, DataSource> instances =
@@ -1188,5 +1618,9 @@
public void populateDisplayData(DisplayData.Builder builder) {
config.populateDisplayData(builder);
}
+
+ public DataSourceConfiguration getConfig() {
+ return this.config;
+ }
}
}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakePipelineOptions.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakePipelineOptions.java
index bf91e0c..201f71c 100644
--- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakePipelineOptions.java
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakePipelineOptions.java
@@ -20,9 +20,11 @@
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.options.Validation;
+import org.apache.beam.sdk.options.ValueProvider;
-public interface SnowflakePipelineOptions extends PipelineOptions {
+public interface SnowflakePipelineOptions extends PipelineOptions, StreamingOptions {
String BASIC_CONNECTION_INFO_VALIDATION_GROUP = "BASIC_CONNECTION_INFO_GROUP";
String AUTH_VALIDATION_GROUP = "AUTH_VALIDATION_GROUP";
@@ -120,4 +122,9 @@
String getStorageIntegrationName();
void setStorageIntegrationName(String storageIntegrationName);
+
+ @Description("SnowPipe name. Optional.")
+ ValueProvider<String> getSnowPipe();
+
+ void setSnowPipe(ValueProvider<String> snowPipe);
}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/credentials/SnowflakeCredentialsFactory.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/credentials/SnowflakeCredentialsFactory.java
index 3876c2f..2b45dc1 100644
--- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/credentials/SnowflakeCredentialsFactory.java
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/credentials/SnowflakeCredentialsFactory.java
@@ -18,38 +18,52 @@
package org.apache.beam.sdk.io.snowflake.credentials;
import org.apache.beam.sdk.io.snowflake.SnowflakePipelineOptions;
+import org.apache.beam.sdk.io.snowflake.crosslanguage.SnowflakeReadRegistrar;
/**
* Factory class for creating implementations of {@link SnowflakeCredentials} from {@link
* SnowflakePipelineOptions}.
*/
public class SnowflakeCredentialsFactory {
- public static SnowflakeCredentials of(SnowflakePipelineOptions options) {
- if (oauthOptionsAvailable(options)) {
- return new OAuthTokenSnowflakeCredentials(options.getOauthToken());
- } else if (usernamePasswordOptionsAvailable(options)) {
- return new UsernamePasswordSnowflakeCredentials(options.getUsername(), options.getPassword());
- } else if (keyPairOptionsAvailable(options)) {
+ public static SnowflakeCredentials of(SnowflakePipelineOptions o) {
+ if (oauthOptionsAvailable(o.getOauthToken())) {
+ return new OAuthTokenSnowflakeCredentials(o.getOauthToken());
+ } else if (usernamePasswordOptionsAvailable(o.getUsername(), o.getPassword())) {
+ return new UsernamePasswordSnowflakeCredentials(o.getUsername(), o.getPassword());
+ } else if (keyPairOptionsAvailable(
+ o.getUsername(), o.getPrivateKeyPath(), o.getPrivateKeyPassphrase())) {
return new KeyPairSnowflakeCredentials(
- options.getUsername(), options.getPrivateKeyPath(), options.getPrivateKeyPassphrase());
+ o.getUsername(), o.getPrivateKeyPath(), o.getPrivateKeyPassphrase());
}
throw new RuntimeException("Can't get credentials from Options");
}
- private static boolean oauthOptionsAvailable(SnowflakePipelineOptions options) {
- return options.getOauthToken() != null && !options.getOauthToken().isEmpty();
+ public static SnowflakeCredentials of(SnowflakeReadRegistrar.ReadConfiguration c) {
+ if (oauthOptionsAvailable(c.getOAuthToken())) {
+ return new OAuthTokenSnowflakeCredentials(c.getOAuthToken());
+ } else if (usernamePasswordOptionsAvailable(c.getUsername(), c.getPassword())) {
+ return new UsernamePasswordSnowflakeCredentials(c.getUsername(), c.getPassword());
+ } else if (keyPairOptionsAvailable(
+ c.getUsername(), c.getPrivateKeyPath(), c.getPrivateKeyPassphrase())) {
+ return new KeyPairSnowflakeCredentials(
+ c.getUsername(), c.getPrivateKeyPath(), c.getPrivateKeyPassphrase());
+ }
+ throw new RuntimeException("Can't get credentials from Options");
}
- private static boolean usernamePasswordOptionsAvailable(SnowflakePipelineOptions options) {
- return options.getUsername() != null
- && !options.getUsername().isEmpty()
- && !options.getPassword().isEmpty();
+ private static boolean oauthOptionsAvailable(String token) {
+ return token != null && !token.isEmpty();
}
- private static boolean keyPairOptionsAvailable(SnowflakePipelineOptions options) {
- return options.getUsername() != null
- && !options.getUsername().isEmpty()
- && !options.getPrivateKeyPath().isEmpty()
- && !options.getPrivateKeyPassphrase().isEmpty();
+ private static boolean usernamePasswordOptionsAvailable(String username, String password) {
+ return username != null && !username.isEmpty() && !password.isEmpty();
+ }
+
+ private static boolean keyPairOptionsAvailable(
+ String username, String privateKeyPath, String privateKeyPassphrase) {
+ return username != null
+ && !username.isEmpty()
+ && !privateKeyPath.isEmpty()
+ && !privateKeyPassphrase.isEmpty();
}
}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/Configuration.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/Configuration.java
new file mode 100644
index 0000000..38162ae
--- /dev/null
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/Configuration.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.io.snowflake.crosslanguage;
+
+/** Parameters abstract class to expose the transforms to an external SDK. */
+public abstract class Configuration {
+ private String serverName;
+ private String username;
+ private String password;
+ private String privateKeyPath;
+ private String privateKeyPassphrase;
+ private String oAuthToken;
+ private String database;
+ private String schema;
+ private String table;
+ private String query;
+ private String stagingBucketName;
+ private String storageIntegrationName;
+
+ public String getServerName() {
+ return serverName;
+ }
+
+ public void setServerName(String serverName) {
+ this.serverName = serverName;
+ }
+
+ public String getUsername() {
+ return username;
+ }
+
+ public void setUsername(String username) {
+ this.username = username;
+ }
+
+ public String getPassword() {
+ return password;
+ }
+
+ public void setPassword(String password) {
+ this.password = password;
+ }
+
+ public String getPrivateKeyPath() {
+ return privateKeyPath;
+ }
+
+ public void setPrivateKeyPath(String privateKeyPath) {
+ this.privateKeyPath = privateKeyPath;
+ }
+
+ public String getPrivateKeyPassphrase() {
+ return privateKeyPassphrase;
+ }
+
+ public void setPrivateKeyPassphrase(String privateKeyPassphrase) {
+ this.privateKeyPassphrase = privateKeyPassphrase;
+ }
+
+ public String getOAuthToken() {
+ return oAuthToken;
+ }
+
+ public void setOAuthToken(String oAuthToken) {
+ this.oAuthToken = oAuthToken;
+ }
+
+ public String getDatabase() {
+ return database;
+ }
+
+ public void setDatabase(String database) {
+ this.database = database;
+ }
+
+ public String getSchema() {
+ return schema;
+ }
+
+ public void setSchema(String schema) {
+ this.schema = schema;
+ }
+
+ public String getTable() {
+ return table;
+ }
+
+ public void setTable(String table) {
+ this.table = table;
+ }
+
+ public String getQuery() {
+ return query;
+ }
+
+ public void setQuery(String query) {
+ this.query = query;
+ }
+
+ public String getStagingBucketName() {
+ return stagingBucketName;
+ }
+
+ public void setStagingBucketName(String stagingBucketName) {
+ this.stagingBucketName = stagingBucketName;
+ }
+
+ public String getStorageIntegrationName() {
+ return storageIntegrationName;
+ }
+
+ public void setStorageIntegrationName(String storageIntegrationName) {
+ this.storageIntegrationName = storageIntegrationName;
+ }
+}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/SnowflakeReadRegistrar.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/SnowflakeReadRegistrar.java
new file mode 100644
index 0000000..1e7be0f
--- /dev/null
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/SnowflakeReadRegistrar.java
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.io.snowflake.crosslanguage;
+
+import com.google.auto.service.AutoService;
+import java.io.Serializable;
+import java.nio.charset.Charset;
+import java.util.Map;
+import javax.sql.DataSource;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
+import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
+import org.apache.beam.sdk.io.snowflake.credentials.SnowflakeCredentials;
+import org.apache.beam.sdk.io.snowflake.credentials.SnowflakeCredentialsFactory;
+import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
+
+/** Exposes {@link SnowflakeIO.Read} as an external transform for cross-language usage. */
+@Experimental
+@AutoService(ExternalTransformRegistrar.class)
+public final class SnowflakeReadRegistrar implements ExternalTransformRegistrar {
+
+ public static final String URN = "beam:external:java:snowflake:read:v1";
+
+ @Override
+ public Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders() {
+ return ImmutableMap.of(URN, ReadBuilder.class);
+ }
+
+ /** Parameters class to expose the transform to an external SDK. */
+ public static class ReadConfiguration extends Configuration {}
+
+ public static class ReadBuilder
+ implements ExternalTransformBuilder<ReadConfiguration, PBegin, PCollection<byte[]>> {
+ public ReadBuilder() {}
+
+ @Override
+ public PTransform<PBegin, PCollection<byte[]>> buildExternal(ReadConfiguration c) {
+ SnowflakeCredentials credentials = SnowflakeCredentialsFactory.of(c);
+
+ SerializableFunction<Void, DataSource> dataSourceSerializableFunction =
+ SnowflakeIO.DataSourceProviderFromDataSourceConfiguration.of(
+ SnowflakeIO.DataSourceConfiguration.create(credentials)
+ .withServerName(c.getServerName())
+ .withDatabase(c.getDatabase())
+ .withSchema(c.getSchema()));
+
+ return SnowflakeIO.<byte[]>read()
+ .withStorageIntegrationName(c.getStorageIntegrationName())
+ .withStagingBucketName(c.getStagingBucketName())
+ .withDataSourceProviderFn(dataSourceSerializableFunction)
+ .withCsvMapper(CsvMapper.getCsvMapper())
+ .withCoder(ByteArrayCoder.of())
+ .fromTable(c.getTable())
+ .fromQuery(c.getQuery());
+ }
+ }
+
+ private static class CsvMapper implements Serializable {
+
+ public static SnowflakeIO.CsvMapper getCsvMapper() {
+ return (SnowflakeIO.CsvMapper<byte[]>)
+ parts -> {
+ String partsCSV = String.join(",", parts);
+
+ return partsCSV.getBytes(Charset.defaultCharset());
+ };
+ }
+ }
+}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/package-info.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/package-info.java
new file mode 100644
index 0000000..7e24ee9
--- /dev/null
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/crosslanguage/package-info.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/** Cross-language for SnowflakeIO. */
+@Experimental(Kind.PORTABILITY)
+@DefaultAnnotation(NonNull.class)
+package org.apache.beam.sdk.io.snowflake.crosslanguage;
+
+import edu.umd.cs.findbugs.annotations.DefaultAnnotation;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.annotations.Experimental.Kind;
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/enums/StreamingLogLevel.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/enums/StreamingLogLevel.java
new file mode 100644
index 0000000..f547b07
--- /dev/null
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/enums/StreamingLogLevel.java
@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.snowflake.enums;
+
+public enum StreamingLogLevel {
+ INFO,
+ ERROR
+}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/ServiceConfig.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/ServiceConfig.java
index 09e1368..1826ce9 100644
--- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/ServiceConfig.java
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/ServiceConfig.java
@@ -17,4 +17,8 @@
*/
package org.apache.beam.sdk.io.snowflake.services;
+/**
+ * Configuration abstract class for {@link SnowflakeService} that gives parameters for write and
+ * read (batch and streaming).
+ */
public abstract class ServiceConfig {}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeServiceConfig.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeBatchServiceConfig.java
similarity index 60%
rename from sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeServiceConfig.java
rename to sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeBatchServiceConfig.java
index fc68a00..726e9d7 100644
--- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeServiceConfig.java
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeBatchServiceConfig.java
@@ -24,81 +24,120 @@
import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;
import org.apache.beam.sdk.transforms.SerializableFunction;
-public class SnowflakeServiceConfig extends ServiceConfig {
- private SerializableFunction<Void, DataSource> dataSourceProviderFn;
+/** Class for preparing configuration for batch write and read. */
+public class SnowflakeBatchServiceConfig extends ServiceConfig {
+ private final SerializableFunction<Void, DataSource> dataSourceProviderFn;
- private String table;
- private String query;
- private String storageIntegrationName;
+ private final String database;
+ private final String schema;
+ private final String table;
+ private final String query;
+ private final String storageIntegrationName;
private List<String> filesList;
-
private WriteDisposition writeDisposition;
private CreateDisposition createDisposition;
private SnowflakeTableSchema tableSchema;
- private String stagingBucketDir;
+ private final String stagingBucketDir;
+ private final String quotationMark;
- public SnowflakeServiceConfig(
+ /** Creating a batch configuration for reading. */
+ public SnowflakeBatchServiceConfig(
SerializableFunction<Void, DataSource> dataSourceProviderFn,
+ String database,
+ String schema,
String table,
String query,
- String storageIntegration,
- String stagingBucketDir) {
+ String storageIntegrationName,
+ String stagingBucketDir,
+ String quotationMark) {
this.dataSourceProviderFn = dataSourceProviderFn;
+ this.database = database;
+ this.schema = schema;
this.table = table;
this.query = query;
- this.storageIntegrationName = storageIntegration;
+ this.storageIntegrationName = storageIntegrationName;
this.stagingBucketDir = stagingBucketDir;
+ this.quotationMark = quotationMark;
}
- public SnowflakeServiceConfig(
+ /** Creating a batch configuration for writing. */
+ public SnowflakeBatchServiceConfig(
SerializableFunction<Void, DataSource> dataSourceProviderFn,
List<String> filesList,
+ String database,
+ String schema,
String table,
String query,
SnowflakeTableSchema tableSchema,
CreateDisposition createDisposition,
WriteDisposition writeDisposition,
String storageIntegrationName,
- String stagingBucketDir) {
+ String stagingBucketDir,
+ String quotationMark) {
this.dataSourceProviderFn = dataSourceProviderFn;
this.filesList = filesList;
+ this.database = database;
+ this.schema = schema;
this.table = table;
this.query = query;
- this.tableSchema = tableSchema;
this.writeDisposition = writeDisposition;
this.createDisposition = createDisposition;
+ this.tableSchema = tableSchema;
this.storageIntegrationName = storageIntegrationName;
this.stagingBucketDir = stagingBucketDir;
+ this.quotationMark = quotationMark;
}
+ /** Getting a DataSource provider function for connection credentials. */
public SerializableFunction<Void, DataSource> getDataSourceProviderFn() {
return dataSourceProviderFn;
}
+ /** Getting a table as a source of reading or destination to write. */
public String getTable() {
return table;
}
+ /** Getting a query which can be source for reading. */
public String getQuery() {
return query;
}
+ /** Getting Snowflake integration which is used in COPY statement. */
public String getStorageIntegrationName() {
return storageIntegrationName;
}
+ /** Getting directory where files are staged. */
public String getStagingBucketDir() {
return stagingBucketDir;
}
+ /** Getting list of names of staged files. */
public List<String> getFilesList() {
return filesList;
}
+ /** Getting disposition how write data to table, see: {@link WriteDisposition}. */
public WriteDisposition getWriteDisposition() {
return writeDisposition;
}
+ /** Getting a character that will surround {@code String} in staged CSV files. */
+ public String getQuotationMark() {
+ return quotationMark;
+ }
+
+ /** Getting a Snowflake database. */
+ public String getDatabase() {
+ return database;
+ }
+
+ /** Getting a schema of a Snowflake table. */
+ public String getSchema() {
+ return schema;
+ }
+
public CreateDisposition getCreateDisposition() {
return createDisposition;
}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeServiceImpl.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeBatchServiceImpl.java
similarity index 79%
rename from sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeServiceImpl.java
rename to sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeBatchServiceImpl.java
index 36e9f3b..2b73bd5 100644
--- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeServiceImpl.java
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeBatchServiceImpl.java
@@ -19,6 +19,8 @@
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import java.math.BigInteger;
+import java.nio.charset.Charset;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
@@ -35,23 +37,36 @@
import org.slf4j.LoggerFactory;
/** Implemenation of {@link SnowflakeService} used in production. */
-public class SnowflakeServiceImpl implements SnowflakeService<SnowflakeServiceConfig> {
- private static final Logger LOG = LoggerFactory.getLogger(SnowflakeServiceImpl.class);
+public class SnowflakeBatchServiceImpl implements SnowflakeService<SnowflakeBatchServiceConfig> {
+ private static final Logger LOG = LoggerFactory.getLogger(SnowflakeBatchServiceImpl.class);
private static final String SNOWFLAKE_GCS_PREFIX = "gcs://";
private static final String GCS_PREFIX = "gs://";
+ /** Writing data to Snowflake in batch mode. */
@Override
- public void write(SnowflakeServiceConfig config) throws Exception {
+ public void write(SnowflakeBatchServiceConfig config) throws Exception {
copyToTable(config);
}
+ /** Reading data from Snowflake tables in batch processing. */
@Override
- public String read(SnowflakeServiceConfig config) throws Exception {
+ public String read(SnowflakeBatchServiceConfig config) throws Exception {
return copyIntoStage(config);
}
- public String copyIntoStage(SnowflakeServiceConfig config) throws SQLException {
+ /**
+ * Copies data from specified table to stage (bucket and directory). Uses Snowflake's <a
+ * href="https://docs.snowflake.com/en/sql-reference/sql/copy-into-location.html">COPY method</a>.
+ * All the details needed for COPY are inside passed configuration.
+ *
+ * @param config object with configuration to perform COPY query.
+ * @return destination where files were copied into
+ * @throws SQLException in case COPY query failed.
+ */
+ private String copyIntoStage(SnowflakeBatchServiceConfig config) throws SQLException {
SerializableFunction<Void, DataSource> dataSourceProviderFn = config.getDataSourceProviderFn();
+ String database = config.getDatabase();
+ String schema = config.getSchema();
String table = config.getTable();
String query = config.getQuery();
String storageIntegrationName = config.getStorageIntegrationName();
@@ -62,7 +77,7 @@
// Query must be surrounded with brackets
source = String.format("(%s)", query);
} else {
- source = table;
+ source = getTablePath(database, schema, table);
}
String copyQuery =
@@ -71,17 +86,31 @@
getProperBucketDir(stagingBucketDir),
source,
storageIntegrationName,
- CSV_QUOTE_CHAR_FOR_COPY);
+ getASCIICharRepresentation(config.getQuotationMark()));
runStatement(copyQuery, getConnection(dataSourceProviderFn), null);
return stagingBucketDir.concat("*");
}
- public void copyToTable(SnowflakeServiceConfig config) throws SQLException {
+ private String getASCIICharRepresentation(String input) {
+ return String.format("0x%x", new BigInteger(1, input.getBytes(Charset.defaultCharset())));
+ }
+
+ /**
+ * Copies staged data from bucket directory to table. Uses Snowflake's <a
+ * href="https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html">COPY method</a>.
+ * All the details needed for COPY are inside passed configuration.
+ *
+ * @param config object with configuration to perform COPY query.
+ * @throws SQLException
+ */
+ private void copyToTable(SnowflakeBatchServiceConfig config) throws SQLException {
SerializableFunction<Void, DataSource> dataSourceProviderFn = config.getDataSourceProviderFn();
List<String> filesList = config.getFilesList();
+ String database = config.getDatabase();
+ String schema = config.getSchema();
String table = config.getTable();
String query = config.getQuery();
SnowflakeTableSchema tableSchema = config.getTableSchema();
@@ -110,28 +139,29 @@
query =
String.format(
"COPY INTO %s FROM %s FILES=(%s) FILE_FORMAT=(TYPE=CSV FIELD_OPTIONALLY_ENCLOSED_BY='%s' COMPRESSION=GZIP) STORAGE_INTEGRATION=%s;",
- table,
+ getTablePath(database, schema, table),
getProperBucketDir(source),
files,
- CSV_QUOTE_CHAR_FOR_COPY,
+ getASCIICharRepresentation(config.getQuotationMark()),
storageIntegrationName);
} else {
query =
String.format(
"COPY INTO %s FROM %s FILES=(%s) FILE_FORMAT=(TYPE=CSV FIELD_OPTIONALLY_ENCLOSED_BY='%s' COMPRESSION=GZIP);",
- table, source, files, CSV_QUOTE_CHAR_FOR_COPY);
+ table, source, files, getASCIICharRepresentation(config.getQuotationMark()));
}
runStatement(query, dataSource.getConnection(), null);
}
- private void truncateTable(DataSource dataSource, String table) throws SQLException {
- String query = String.format("TRUNCATE %s;", table);
+ private void truncateTable(DataSource dataSource, String tablePath) throws SQLException {
+ String query = String.format("TRUNCATE %s;", tablePath);
runConnectionWithStatement(dataSource, query, null);
}
- private static void checkIfTableIsEmpty(DataSource dataSource, String table) throws SQLException {
- String selectQuery = String.format("SELECT count(*) FROM %s LIMIT 1;", table);
+ private static void checkIfTableIsEmpty(DataSource dataSource, String tablePath)
+ throws SQLException {
+ String selectQuery = String.format("SELECT count(*) FROM %s LIMIT 1;", tablePath);
runConnectionWithStatement(
dataSource,
selectQuery,
@@ -275,4 +305,8 @@
}
return bucketDir;
}
+
+ private String getTablePath(String database, String schema, String table) {
+ return String.format("%s.%s.%s", database, schema, table);
+ }
}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeService.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeService.java
index 16cd3c6..192856e 100644
--- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeService.java
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeService.java
@@ -21,8 +21,6 @@
/** Interface which defines common methods for interacting with Snowflake. */
public interface SnowflakeService<T extends ServiceConfig> extends Serializable {
- String CSV_QUOTE_CHAR_FOR_COPY = "''";
-
String read(T config) throws Exception;
void write(T config) throws Exception;
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeStreamingServiceConfig.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeStreamingServiceConfig.java
new file mode 100644
index 0000000..7039c89
--- /dev/null
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeStreamingServiceConfig.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.io.snowflake.services;
+
+import java.util.List;
+import net.snowflake.ingest.SimpleIngestManager;
+
+/** Class for preparing configuration for streaming write. */
+public class SnowflakeStreamingServiceConfig extends ServiceConfig {
+ private final SimpleIngestManager ingestManager;
+ private final List<String> filesList;
+ private final String stagingBucketDir;
+
+ /**
+ * Constructor to create configuration for streaming write.
+ *
+ * @param filesList list of strings of staged files' names.
+ * @param stagingBucketDir name of a bucket and directory inside where files are staged and awaits
+ * for being loaded to Snowflake.
+ * @param ingestManager instance of {@link SimpleIngestManager}.
+ */
+ public SnowflakeStreamingServiceConfig(
+ List<String> filesList, String stagingBucketDir, SimpleIngestManager ingestManager) {
+ this.filesList = filesList;
+ this.stagingBucketDir = stagingBucketDir;
+ this.ingestManager = ingestManager;
+ }
+
+ /**
+ * Getter for ingest manager which serves API to load data in streaming mode and retrieve a report
+ * about loaded data.
+ *
+ * @return instance of {@link SimpleIngestManager}.
+ */
+ public SimpleIngestManager getIngestManager() {
+ return ingestManager;
+ }
+
+ /**
+ * Getter for a list of staged files which are will be loaded to Snowflake.
+ *
+ * @return list of strings of staged files' names.
+ */
+ public List<String> getFilesList() {
+ return filesList;
+ }
+
+ /**
+ * Getter for a bucket name with directory where files were staged and waiting for loading.
+ *
+ * @return name of a bucket and directory inside in form {@code gs://mybucket/dir/}
+ */
+ public String getStagingBucketDir() {
+ return stagingBucketDir;
+ }
+}
diff --git a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeStreamingServiceImpl.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeStreamingServiceImpl.java
new file mode 100644
index 0000000..8d555ed
--- /dev/null
+++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/services/SnowflakeStreamingServiceImpl.java
@@ -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.
+ */
+package org.apache.beam.sdk.io.snowflake.services;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import net.snowflake.ingest.SimpleIngestManager;
+import net.snowflake.ingest.connection.IngestResponseException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Implemenation of {@link SnowflakeService} used in production. */
+public class SnowflakeStreamingServiceImpl
+ implements SnowflakeService<SnowflakeStreamingServiceConfig> {
+ private static final Logger LOG = LoggerFactory.getLogger(SnowflakeStreamingServiceImpl.class);
+ private transient SimpleIngestManager ingestManager;
+
+ /** Writing data to Snowflake in streaming mode. */
+ @Override
+ public void write(SnowflakeStreamingServiceConfig config) throws Exception {
+ ingest(config);
+ }
+
+ /** Reading data from Snowflake in streaming mode is not supported. */
+ @Override
+ public String read(SnowflakeStreamingServiceConfig config) throws Exception {
+ throw new UnsupportedOperationException("Not supported by SnowflakeIO.");
+ }
+
+ /**
+ * SnowPipe is processing files from stage in streaming mode.
+ *
+ * @param config configuration object containing parameters for writing files to Snowflake
+ * @throws IngestResponseException REST API response error
+ * @throws IOException Snowflake problem while streaming
+ * @throws URISyntaxException creating request error
+ */
+ private void ingest(SnowflakeStreamingServiceConfig config)
+ throws IngestResponseException, IOException, URISyntaxException {
+ List<String> filesList = config.getFilesList();
+ String stagingBucketDir = config.getStagingBucketDir();
+ ingestManager = config.getIngestManager();
+
+ Set<String> files =
+ filesList.stream()
+ .map(e -> e.replaceAll(String.valueOf(stagingBucketDir), ""))
+ .map(e -> e.replaceAll("'", ""))
+ .collect(Collectors.toSet());
+
+ if (!files.isEmpty()) {
+ this.ingestManager.ingestFiles(SimpleIngestManager.wrapFilepaths(files), null);
+ }
+ }
+}
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeServiceImpl.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeBatchServiceImpl.java
similarity index 89%
rename from sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeServiceImpl.java
rename to sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeBatchServiceImpl.java
index 7ff097d..90ee4b9 100644
--- a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeServiceImpl.java
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeBatchServiceImpl.java
@@ -28,23 +28,24 @@
import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeBatchServiceConfig;
import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
-import org.apache.beam.sdk.io.snowflake.services.SnowflakeServiceConfig;
/** Fake implementation of {@link SnowflakeService} used in tests. */
-public class FakeSnowflakeServiceImpl implements SnowflakeService<SnowflakeServiceConfig> {
+public class FakeSnowflakeBatchServiceImpl
+ implements SnowflakeService<SnowflakeBatchServiceConfig> {
@Override
- public void write(SnowflakeServiceConfig config) throws Exception {
+ public void write(SnowflakeBatchServiceConfig config) throws Exception {
copyToTable(config);
}
@Override
- public String read(SnowflakeServiceConfig config) throws Exception {
+ public String read(SnowflakeBatchServiceConfig config) throws Exception {
return copyIntoStage(config);
}
- public String copyIntoStage(SnowflakeServiceConfig config) throws SQLException {
+ public String copyIntoStage(SnowflakeBatchServiceConfig config) throws SQLException {
String table = config.getTable();
String query = config.getQuery();
@@ -60,7 +61,7 @@
return String.format("./%s/*", stagingBucketDir);
}
- public void copyToTable(SnowflakeServiceConfig config) throws SQLException {
+ public void copyToTable(SnowflakeBatchServiceConfig config) throws SQLException {
List<String> filesList = config.getFilesList();
String table = config.getTable();
SnowflakeTableSchema tableSchema = config.getTableSchema();
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeIngestManager.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeIngestManager.java
new file mode 100644
index 0000000..e144ae7
--- /dev/null
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeIngestManager.java
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.io.snowflake.test;
+
+import java.util.List;
+
+public class FakeSnowflakeIngestManager {
+ // Only for testing purposes
+ private String table = "TEST_TABLE";
+
+ public FakeSnowflakeIngestManager() {}
+
+ public void ingestFiles(List<String> rows) {
+ FakeSnowflakeDatabase.createTableWithElements(this.table, rows);
+ }
+
+ public String getTable() {
+ return this.table;
+ }
+}
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeStreamingServiceImpl.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeStreamingServiceImpl.java
new file mode 100644
index 0000000..362eb5e
--- /dev/null
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/FakeSnowflakeStreamingServiceImpl.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.sdk.io.snowflake.test;
+
+import java.io.BufferedReader;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.nio.charset.Charset;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.zip.GZIPInputStream;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeStreamingServiceConfig;
+
+/** Fake implementation of {@link SnowflakeService} used in tests. */
+public class FakeSnowflakeStreamingServiceImpl
+ implements SnowflakeService<SnowflakeStreamingServiceConfig> {
+ private FakeSnowflakeIngestManager snowflakeIngestManager;
+
+ @Override
+ public void write(SnowflakeStreamingServiceConfig config) throws Exception {
+ snowflakeIngestManager = new FakeSnowflakeIngestManager();
+ ingest(config);
+ }
+
+ @Override
+ public String read(SnowflakeStreamingServiceConfig config) throws Exception {
+ throw new UnsupportedOperationException("Streaming read is not supported in SnowflakeIO.");
+ }
+
+ public void ingest(SnowflakeStreamingServiceConfig config) {
+ List<String> rows = new ArrayList<>();
+ List<String> filesList = config.getFilesList();
+ for (String file : filesList) {
+ rows.addAll(readGZIPFile(file.replace("'", "")));
+ }
+
+ snowflakeIngestManager.ingestFiles(rows);
+ }
+
+ private List<String> readGZIPFile(String file) {
+ List<String> lines = new ArrayList<>();
+ try {
+ GZIPInputStream gzip = new GZIPInputStream(new FileInputStream(file));
+ BufferedReader br = new BufferedReader(new InputStreamReader(gzip, Charset.defaultCharset()));
+
+ String line;
+ while ((line = br.readLine()) != null) {
+ lines.add(line);
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to read file", e);
+ }
+
+ return lines;
+ }
+}
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/TestUtils.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/TestUtils.java
index 05c80c5..41eac70 100644
--- a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/TestUtils.java
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/TestUtils.java
@@ -138,6 +138,10 @@
return (SnowflakeIO.UserDataMapper<String[]>) recordLine -> recordLine;
}
+ public static SnowflakeIO.UserDataMapper<String> getStringCsvMapper() {
+ return (SnowflakeIO.UserDataMapper<String>) recordLine -> new String[] {recordLine};
+ }
+
public static class ParseToKv extends DoFn<Long, KV<String, Long>> {
@ProcessElement
public void processElement(ProcessContext c) {
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/read/SnowflakeIOReadTest.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/read/SnowflakeIOReadTest.java
index 6016a66..b844760 100644
--- a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/read/SnowflakeIOReadTest.java
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/read/SnowflakeIOReadTest.java
@@ -28,13 +28,15 @@
import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBasicDataSource;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBatchServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeDatabase;
-import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeServiceImpl;
+import org.apache.beam.sdk.io.snowflake.test.TestUtils;
import org.apache.beam.sdk.io.snowflake.test.unit.TestPipelineOptions;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
@@ -46,6 +48,7 @@
public class SnowflakeIOReadTest implements Serializable {
public static final String FAKE_TABLE = "FAKE_TABLE";
public static final String FAKE_QUERY = "SELECT * FROM FAKE_TABLE";
+ public static final String BUCKET_NAME = "BUCKET/";
private static final TestPipelineOptions options =
TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class);;
@@ -71,19 +74,24 @@
options.setServerName("NULL.snowflakecomputing.com");
options.setStorageIntegrationName("STORAGE_INTEGRATION");
- options.setStagingBucketName("BUCKET");
+ options.setStagingBucketName(BUCKET_NAME);
dataSourceConfiguration =
SnowflakeIO.DataSourceConfiguration.create(new FakeSnowflakeBasicDataSource())
.withServerName(options.getServerName());
- snowflakeService = new FakeSnowflakeServiceImpl();
+ snowflakeService = new FakeSnowflakeBatchServiceImpl();
+ }
+
+ @AfterClass
+ public static void tearDown() {
+ TestUtils.removeTempDir(BUCKET_NAME);
}
@Test
public void testConfigIsMissingStagingBucketName() {
thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("withStagingBucketName is required");
+ thrown.expectMessage("withStagingBucketName() is required");
pipeline.apply(
SnowflakeIO.<GenericRecord>read(snowflakeService)
@@ -99,7 +107,7 @@
@Test
public void testConfigIsMissingStorageIntegration() {
thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("withStorageIntegrationName is required");
+ thrown.expectMessage("withStorageIntegrationName() is required");
pipeline.apply(
SnowflakeIO.<GenericRecord>read(snowflakeService)
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/CreateDispositionTest.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/CreateDispositionTest.java
index 4b0f728..2dcd88b 100644
--- a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/CreateDispositionTest.java
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/CreateDispositionTest.java
@@ -32,8 +32,8 @@
import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBasicDataSource;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBatchServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeDatabase;
-import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.TestUtils;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -73,7 +73,7 @@
stagingBucketName = options.getStagingBucketName();
storageIntegrationName = options.getStorageIntegrationName();
- snowflakeService = new FakeSnowflakeServiceImpl();
+ snowflakeService = new FakeSnowflakeBatchServiceImpl();
testData = LongStream.range(0, 100).boxed().collect(Collectors.toList());
dc =
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/QueryDispositionLocationTest.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/QueryDispositionLocationTest.java
index 2825c25..5c06dfe 100644
--- a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/QueryDispositionLocationTest.java
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/QueryDispositionLocationTest.java
@@ -29,8 +29,8 @@
import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;
import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBasicDataSource;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBatchServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeDatabase;
-import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.TestUtils;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -47,7 +47,7 @@
@RunWith(JUnit4.class)
public class QueryDispositionLocationTest {
private static final String FAKE_TABLE = "FAKE_TABLE";
- private static final String BUCKET_NAME = "BUCKET";
+ private static final String BUCKET_NAME = "BUCKET/";
@Rule public final transient TestPipeline pipeline = TestPipeline.create();
@Rule public ExpectedException exceptionRule = ExpectedException.none();
@@ -63,7 +63,7 @@
PipelineOptionsFactory.register(SnowflakePipelineOptions.class);
options = TestPipeline.testingPipelineOptions().as(SnowflakePipelineOptions.class);
- snowflakeService = new FakeSnowflakeServiceImpl();
+ snowflakeService = new FakeSnowflakeBatchServiceImpl();
testData = LongStream.range(0, 100).boxed().collect(Collectors.toList());
}
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SchemaDispositionTest.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SchemaDispositionTest.java
index ac0af04..fe8e98e 100644
--- a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SchemaDispositionTest.java
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SchemaDispositionTest.java
@@ -38,8 +38,8 @@
import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBasicDataSource;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBatchServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeDatabase;
-import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.TestUtils;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -78,7 +78,7 @@
stagingBucketName = options.getStagingBucketName();
storageIntegrationName = options.getStorageIntegrationName();
- snowflakeService = new FakeSnowflakeServiceImpl();
+ snowflakeService = new FakeSnowflakeBatchServiceImpl();
dc =
SnowflakeIO.DataSourceConfiguration.create(new FakeSnowflakeBasicDataSource())
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SnowflakeIOWriteTest.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SnowflakeIOWriteTest.java
index e73760b..95798c5 100644
--- a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SnowflakeIOWriteTest.java
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/SnowflakeIOWriteTest.java
@@ -20,6 +20,7 @@
import static org.junit.Assert.assertTrue;
import java.sql.SQLException;
+import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.LongStream;
@@ -28,8 +29,8 @@
import org.apache.beam.sdk.io.snowflake.SnowflakePipelineOptions;
import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBasicDataSource;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBatchServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeDatabase;
-import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeServiceImpl;
import org.apache.beam.sdk.io.snowflake.test.TestUtils;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -48,7 +49,7 @@
@RunWith(JUnit4.class)
public class SnowflakeIOWriteTest {
private static final String FAKE_TABLE = "FAKE_TABLE";
- private static final String BUCKET_NAME = "BUCKET";
+ private static final String BUCKET_NAME = "BUCKET/";
@Rule public final transient TestPipeline pipeline = TestPipeline.create();
@@ -59,11 +60,21 @@
private static SnowflakeService snowflakeService;
private static List<Long> testData;
+ private static List<String> testDataInStrings;
@BeforeClass
public static void setupAll() {
- snowflakeService = new FakeSnowflakeServiceImpl();
+ snowflakeService = new FakeSnowflakeBatchServiceImpl();
testData = LongStream.range(0, 100).boxed().collect(Collectors.toList());
+
+ testDataInStrings = new ArrayList<>();
+ testDataInStrings.add("First row");
+ testDataInStrings.add("Second row with 'single' quotation");
+ testDataInStrings.add("Second row with single one ' quotation");
+ testDataInStrings.add("Second row with single twice '' quotation");
+ testDataInStrings.add("Third row with \"double\" quotation");
+ testDataInStrings.add("Third row with double one \" quotation");
+ testDataInStrings.add("Third row with double twice \"\" quotation");
}
@Before
@@ -147,7 +158,7 @@
List<String> actualData = FakeSnowflakeDatabase.getElements(FAKE_TABLE);
List<String> testDataInStrings =
- testData.stream().map(e -> e.toString()).collect(Collectors.toList());
+ testData.stream().map(Object::toString).collect(Collectors.toList());
assertTrue(TestUtils.areListsEqual(testDataInStrings, actualData));
}
@@ -174,4 +185,55 @@
assertTrue(TestUtils.areListsEqual(testData, actualData));
}
+
+ @Test
+ public void writeToExternalWithDoubleQuotation() throws SnowflakeSQLException {
+
+ pipeline
+ .apply(Create.of(testDataInStrings))
+ .apply(
+ "Write SnowflakeIO",
+ SnowflakeIO.<String>write()
+ .withDataSourceConfiguration(dc)
+ .withUserDataMapper(TestUtils.getStringCsvMapper())
+ .to(FAKE_TABLE)
+ .withStagingBucketName(options.getStagingBucketName())
+ .withStorageIntegrationName(options.getStorageIntegrationName())
+ .withSnowflakeService(snowflakeService)
+ .withQuotationMark("\""));
+
+ pipeline.run(options).waitUntilFinish();
+
+ List<String> actualData = FakeSnowflakeDatabase.getElements(FAKE_TABLE);
+ List<String> escapedTestData =
+ testDataInStrings.stream()
+ .map(e -> e.replace("'", "''"))
+ .map(e -> String.format("\"%s\"", e))
+ .collect(Collectors.toList());
+ assertTrue(TestUtils.areListsEqual(escapedTestData, actualData));
+ }
+
+ @Test
+ public void writeToExternalWithBlankQuotation() throws SnowflakeSQLException {
+ pipeline
+ .apply(Create.of(testDataInStrings))
+ .apply(
+ "Write SnowflakeIO",
+ SnowflakeIO.<String>write()
+ .withDataSourceConfiguration(dc)
+ .withUserDataMapper(TestUtils.getStringCsvMapper())
+ .to(FAKE_TABLE)
+ .withStagingBucketName(options.getStagingBucketName())
+ .withStorageIntegrationName(options.getStorageIntegrationName())
+ .withSnowflakeService(snowflakeService)
+ .withQuotationMark(""));
+
+ pipeline.run(options).waitUntilFinish();
+
+ List<String> actualData = FakeSnowflakeDatabase.getElements(FAKE_TABLE);
+
+ List<String> escapedTestData =
+ testDataInStrings.stream().map(e -> e.replace("'", "''")).collect(Collectors.toList());
+ assertTrue(TestUtils.areListsEqual(escapedTestData, actualData));
+ }
}
diff --git a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/StreamingWriteTest.java b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/StreamingWriteTest.java
new file mode 100644
index 0000000..3f8ffe5
--- /dev/null
+++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/test/unit/write/StreamingWriteTest.java
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.io.snowflake.test.unit.write;
+
+import static org.hamcrest.CoreMatchers.equalTo;
+
+import java.io.IOException;
+import java.nio.file.DirectoryStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.LongStream;
+import net.snowflake.client.jdbc.SnowflakeSQLException;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.snowflake.SnowflakeIO;
+import org.apache.beam.sdk.io.snowflake.SnowflakePipelineOptions;
+import org.apache.beam.sdk.io.snowflake.credentials.SnowflakeCredentialsFactory;
+import org.apache.beam.sdk.io.snowflake.services.SnowflakeService;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeBasicDataSource;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeDatabase;
+import org.apache.beam.sdk.io.snowflake.test.FakeSnowflakeStreamingServiceImpl;
+import org.apache.beam.sdk.io.snowflake.test.TestUtils;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.testing.TestStream;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.values.TimestampedValue;
+import org.hamcrest.MatcherAssert;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.After;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@RunWith(JUnit4.class)
+public class StreamingWriteTest {
+ private static final Logger LOG = LoggerFactory.getLogger(StreamingWriteTest.class);
+ private static final String FAKE_TABLE = "TEST_TABLE";
+ private static final String STAGING_BUCKET_NAME = "BUCKET/";
+ private static final String STORAGE_INTEGRATION_NAME = "STORAGE_INTEGRATION";
+ private static final String SNOW_PIPE = "Snowpipe";
+ private static final Instant START_TIME = new Instant(0);
+
+ @Rule public final transient TestPipeline pipeline = TestPipeline.create();
+
+ @Rule public ExpectedException exceptionRule = ExpectedException.none();
+ private static SnowflakeIO.DataSourceConfiguration dataSourceConfiguration;
+ private static SnowflakeService snowflakeService;
+ private static SnowflakePipelineOptions options;
+ private static List<Long> testData;
+
+ private static final List<String> SENTENCES =
+ Arrays.asList(
+ "Snowflake window 1 1",
+ "Snowflake window 1 2",
+ "Snowflake window 1 3",
+ "Snowflake window 1 4",
+ "Snowflake window 2 1",
+ "Snowflake window 2 2");
+
+ private static final List<String> FIRST_WIN_WORDS = SENTENCES.subList(0, 4);
+ private static final List<String> SECOND_WIN_WORDS = SENTENCES.subList(4, 6);
+ private static final Duration WINDOW_DURATION = Duration.standardMinutes(1);
+
+ @BeforeClass
+ public static void setup() {
+ snowflakeService = new FakeSnowflakeStreamingServiceImpl();
+
+ PipelineOptionsFactory.register(SnowflakePipelineOptions.class);
+ options = TestPipeline.testingPipelineOptions().as(SnowflakePipelineOptions.class);
+ options.setUsername("username");
+
+ options.setServerName("NULL.snowflakecomputing.com");
+
+ testData = LongStream.range(0, 100).boxed().collect(Collectors.toList());
+
+ FakeSnowflakeDatabase.createTable(FAKE_TABLE);
+ dataSourceConfiguration =
+ SnowflakeIO.DataSourceConfiguration.create(new FakeSnowflakeBasicDataSource())
+ .withServerName(options.getServerName())
+ .withoutValidation()
+ .withSchema("PUBLIC")
+ .withDatabase("DATABASE")
+ .withWarehouse("WAREHOUSE");
+ }
+
+ @After
+ public void tearDown() {
+ TestUtils.removeTempDir(STAGING_BUCKET_NAME);
+ }
+
+ @Test
+ public void streamWriteWithOAuthFails() {
+ options.setOauthToken("token");
+ dataSourceConfiguration =
+ SnowflakeIO.DataSourceConfiguration.create(SnowflakeCredentialsFactory.of(options))
+ .withoutValidation()
+ .withServerName(options.getServerName())
+ .withSchema("PUBLIC")
+ .withDatabase("DATABASE")
+ .withWarehouse("WAREHOUSE");
+
+ exceptionRule.expectMessage("KeyPair is required for authentication");
+
+ pipeline
+ .apply(Create.of(testData))
+ .apply(
+ SnowflakeIO.<Long>write()
+ .withDataSourceConfiguration(dataSourceConfiguration)
+ .to(FAKE_TABLE)
+ .withStagingBucketName(STAGING_BUCKET_NAME)
+ .withStorageIntegrationName(STORAGE_INTEGRATION_NAME)
+ .withSnowPipe(SNOW_PIPE)
+ .withUserDataMapper(TestUtils.getLongCsvMapper())
+ .withSnowflakeService(snowflakeService));
+
+ pipeline.run(options);
+ }
+
+ @Test
+ public void streamWriteWithUserPasswordFails() {
+ options.setPassword("password");
+ dataSourceConfiguration =
+ SnowflakeIO.DataSourceConfiguration.create(SnowflakeCredentialsFactory.of(options))
+ .withoutValidation()
+ .withServerName(options.getServerName())
+ .withSchema("PUBLIC")
+ .withDatabase("DATABASE")
+ .withWarehouse("WAREHOUSE");
+
+ exceptionRule.expectMessage("KeyPair is required for authentication");
+
+ pipeline
+ .apply(Create.of(testData))
+ .apply(
+ SnowflakeIO.<Long>write()
+ .withDataSourceConfiguration(dataSourceConfiguration)
+ .to(FAKE_TABLE)
+ .withStagingBucketName(STAGING_BUCKET_NAME)
+ .withStorageIntegrationName(STORAGE_INTEGRATION_NAME)
+ .withSnowPipe(SNOW_PIPE)
+ .withUserDataMapper(TestUtils.getLongCsvMapper())
+ .withSnowflakeService(snowflakeService));
+
+ pipeline.run(options);
+ }
+
+ @Test
+ public void streamWriteWithKey() throws SnowflakeSQLException {
+ String quotationMark = "'";
+ options.setPrivateKeyPath(TestUtils.getPrivateKeyPath(getClass()));
+ options.setPrivateKeyPassphrase(TestUtils.getPrivateKeyPassphrase());
+
+ TestStream<String> stringsStream =
+ TestStream.create(StringUtf8Coder.of())
+ .advanceWatermarkTo(START_TIME)
+ .addElements(event(FIRST_WIN_WORDS.get(0), 2L))
+ .advanceWatermarkTo(START_TIME.plus(Duration.standardSeconds(27L)))
+ .addElements(
+ event(FIRST_WIN_WORDS.get(1), 25L),
+ event(FIRST_WIN_WORDS.get(2), 18L),
+ event(FIRST_WIN_WORDS.get(3), 26L))
+ .advanceWatermarkTo(START_TIME.plus(Duration.standardSeconds(65L)))
+ // This are late elements after window ends so they should not be saved
+ .addElements(event(SECOND_WIN_WORDS.get(0), 67L), event(SECOND_WIN_WORDS.get(1), 68L))
+ .advanceWatermarkToInfinity();
+
+ dataSourceConfiguration =
+ SnowflakeIO.DataSourceConfiguration.create(SnowflakeCredentialsFactory.of(options))
+ .withServerName(options.getServerName())
+ .withoutValidation()
+ .withSchema("PUBLIC")
+ .withDatabase("DATABASE")
+ .withWarehouse("WAREHOUSE");
+
+ pipeline
+ .apply(stringsStream)
+ .apply(Window.into(FixedWindows.of(WINDOW_DURATION)))
+ .apply(
+ SnowflakeIO.<String>write()
+ .withDataSourceConfiguration(dataSourceConfiguration)
+ .withStagingBucketName(STAGING_BUCKET_NAME)
+ .withStorageIntegrationName(STORAGE_INTEGRATION_NAME)
+ .withSnowPipe(SNOW_PIPE)
+ .withFlushRowLimit(4)
+ .withFlushTimeLimit(WINDOW_DURATION)
+ .withUserDataMapper(TestUtils.getStringCsvMapper())
+ .withSnowflakeService(snowflakeService));
+
+ pipeline.run(options).waitUntilFinish();
+
+ List<String> actualDataFirstWin =
+ parseResults(FakeSnowflakeDatabase.getElements(String.format(FAKE_TABLE)), quotationMark);
+
+ Map<String, List<String>> mapOfResults = getMapOfFilesAndResults();
+
+ String firstFileKey = "0";
+ List<String> filesResult = parseResults(mapOfResults.get(firstFileKey), quotationMark);
+
+ int amountOfCreatedFiles = 2;
+ MatcherAssert.assertThat(mapOfResults.size(), equalTo(amountOfCreatedFiles));
+ MatcherAssert.assertThat(filesResult, equalTo(FIRST_WIN_WORDS));
+ MatcherAssert.assertThat(actualDataFirstWin, equalTo(SENTENCES));
+ }
+
+ @Test
+ public void streamWriteWithDoubleQuotation() throws SnowflakeSQLException {
+ String quotationMark = "\"";
+ options.setPrivateKeyPath(TestUtils.getPrivateKeyPath(getClass()));
+ options.setPrivateKeyPassphrase(TestUtils.getPrivateKeyPassphrase());
+
+ TestStream<String> stringsStream =
+ TestStream.create(StringUtf8Coder.of())
+ .advanceWatermarkTo(START_TIME)
+ .addElements(event(FIRST_WIN_WORDS.get(0), 2L))
+ .advanceWatermarkTo(START_TIME.plus(Duration.standardSeconds(27L)))
+ .addElements(
+ event(FIRST_WIN_WORDS.get(1), 25L),
+ event(FIRST_WIN_WORDS.get(2), 18L),
+ event(FIRST_WIN_WORDS.get(3), 26L))
+ .advanceWatermarkTo(START_TIME.plus(Duration.standardSeconds(65L)))
+ // This are late elements after window ends so they should not be saved
+ .addElements(event(SECOND_WIN_WORDS.get(0), 67L), event(SECOND_WIN_WORDS.get(1), 68L))
+ .advanceWatermarkToInfinity();
+
+ dataSourceConfiguration =
+ SnowflakeIO.DataSourceConfiguration.create(SnowflakeCredentialsFactory.of(options))
+ .withServerName(options.getServerName())
+ .withoutValidation()
+ .withSchema("PUBLIC")
+ .withDatabase("DATABASE")
+ .withWarehouse("WAREHOUSE");
+
+ pipeline
+ .apply(stringsStream)
+ .apply(Window.into(FixedWindows.of(WINDOW_DURATION)))
+ .apply(
+ SnowflakeIO.<String>write()
+ .withDataSourceConfiguration(dataSourceConfiguration)
+ .withStagingBucketName(STAGING_BUCKET_NAME)
+ .withStorageIntegrationName(STORAGE_INTEGRATION_NAME)
+ .withSnowPipe(SNOW_PIPE)
+ .withFlushRowLimit(4)
+ .withQuotationMark(quotationMark)
+ .withFlushTimeLimit(WINDOW_DURATION)
+ .withUserDataMapper(TestUtils.getStringCsvMapper())
+ .withSnowflakeService(snowflakeService));
+
+ pipeline.run(options).waitUntilFinish();
+
+ List<String> actualDataFirstWin =
+ parseResults(FakeSnowflakeDatabase.getElements(String.format(FAKE_TABLE)), quotationMark);
+
+ Map<String, List<String>> mapOfResults = getMapOfFilesAndResults();
+
+ String firstFileKey = "0";
+ List<String> filesResult = parseResults(mapOfResults.get(firstFileKey), quotationMark);
+
+ int amountOfCreatedFiles = 2;
+ MatcherAssert.assertThat(mapOfResults.size(), equalTo(amountOfCreatedFiles));
+ MatcherAssert.assertThat(filesResult, equalTo(FIRST_WIN_WORDS));
+ MatcherAssert.assertThat(actualDataFirstWin, equalTo(SENTENCES));
+ }
+
+ private List<String> parseResults(List<String> resultsList, String quotationMark) {
+ return resultsList.stream()
+ .map(s -> s.replaceAll(quotationMark, ""))
+ .collect(Collectors.toList());
+ }
+
+ private Map<String, List<String>> getMapOfFilesAndResults() {
+ return new HashMap<>(getFiles(Paths.get(STAGING_BUCKET_NAME)));
+ }
+
+ private Map<String, List<String>> getFiles(Path file) {
+ Map<String, List<String>> fileNames = new HashMap<>();
+ try (DirectoryStream<Path> paths = Files.newDirectoryStream(file, "*.gz")) {
+
+ paths.forEach(
+ path -> {
+ String key = path.getFileName().toString().split("-", -1)[1];
+ fileNames.put(key, TestUtils.readGZIPFile(path.toString()));
+ });
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to retrieve files", e);
+ }
+ return fileNames;
+ }
+
+ private TimestampedValue<String> event(String word, Long timestamp) {
+ return TimestampedValue.of(word, START_TIME.plus(new Duration(timestamp)));
+ }
+}
diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py
index 8fdc4a4..54c63e7 100644
--- a/sdks/python/apache_beam/coders/coder_impl.py
+++ b/sdks/python/apache_beam/coders/coder_impl.py
@@ -530,6 +530,92 @@
return 1
+class MapCoderImpl(StreamCoderImpl):
+ """For internal use only; no backwards-compatibility guarantees.
+
+ Note this implementation always uses nested context when encoding keys
+ and values. This differs from Java's MapCoder, which uses
+ nested=False if possible for the last value encoded.
+
+ This difference is acceptable because MapCoder is not standard. It is only
+ used in a standard context by RowCoder which always uses nested context for
+ attribute values.
+
+ A coder for typing.Mapping objects."""
+ def __init__(
+ self,
+ key_coder, # type: CoderImpl
+ value_coder # type: CoderImpl
+ ):
+ self._key_coder = key_coder
+ self._value_coder = value_coder
+
+ def encode_to_stream(self, dict_value, out, nested):
+ out.write_bigendian_int32(len(dict_value))
+ for key, value in dict_value.items():
+ # Note this implementation always uses nested context when encoding keys
+ # and values which differs from Java. See note in docstring.
+ self._key_coder.encode_to_stream(key, out, True)
+ self._value_coder.encode_to_stream(value, out, True)
+
+ def decode_from_stream(self, in_stream, nested):
+ size = in_stream.read_bigendian_int32()
+ result = {}
+ for _ in range(size):
+ # Note this implementation always uses nested context when encoding keys
+ # and values which differs from Java. See note in docstring.
+ key = self._key_coder.decode_from_stream(in_stream, True)
+ value = self._value_coder.decode_from_stream(in_stream, True)
+ result[key] = value
+
+ return result
+
+ def estimate_size(self, unused_value, nested=False):
+ estimate = 4 # 4 bytes for int32 size prefix
+ for key, value in unused_value.items():
+ estimate += self._key_coder.estimate_size(key, True)
+ estimate += self._value_coder.estimate_size(value, True)
+ return estimate
+
+
+class NullableCoderImpl(StreamCoderImpl):
+ """For internal use only; no backwards-compatibility guarantees.
+
+ A coder for typing.Optional objects."""
+
+ ENCODE_NULL = 0
+ ENCODE_PRESENT = 1
+
+ def __init__(
+ self,
+ value_coder # type: CoderImpl
+ ):
+ self._value_coder = value_coder
+
+ def encode_to_stream(self, value, out, nested):
+ if value is None:
+ out.write_byte(self.ENCODE_NULL)
+ else:
+ out.write_byte(self.ENCODE_PRESENT)
+ self._value_coder.encode_to_stream(value, out, nested)
+
+ def decode_from_stream(self, in_stream, nested):
+ null_indicator = in_stream.read_byte()
+ if null_indicator == self.ENCODE_NULL:
+ return None
+ elif null_indicator == self.ENCODE_PRESENT:
+ return self._value_coder.decode_from_stream(in_stream, nested)
+ else:
+ raise ValueError(
+ "Encountered unexpected value for null indicator: '%s'" %
+ null_indicator)
+
+ def estimate_size(self, unused_value, nested=False):
+ return 1 + (
+ self._value_coder.estimate_size(unused_value)
+ if unused_value is not None else 0)
+
+
class FloatCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
def encode_to_stream(self, value, out, nested):
diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py
index 399a46d..b6aca0a 100644
--- a/sdks/python/apache_beam/coders/coders.py
+++ b/sdks/python/apache_beam/coders/coders.py
@@ -83,6 +83,8 @@
'FastPrimitivesCoder',
'FloatCoder',
'IterableCoder',
+ 'MapCoder',
+ 'NullableCoder',
'PickleCoder',
'ProtoCoder',
'SingletonCoder',
@@ -520,6 +522,57 @@
Coder.register_structured_urn(common_urns.coders.BOOL.urn, BooleanCoder)
+class MapCoder(FastCoder):
+ def __init__(self, key_coder, value_coder):
+ # type: (Coder, Coder) -> None
+ self._key_coder = key_coder
+ self._value_coder = value_coder
+
+ def _create_impl(self):
+ return coder_impl.MapCoderImpl(
+ self._key_coder.get_impl(), self._value_coder.get_impl())
+
+ def to_type_hint(self):
+ return typehints.Dict[self._key_coder.to_type_hint(),
+ self._value_coder.to_type_hint()]
+
+ def is_deterministic(self):
+ # () -> bool
+ # Map ordering is non-deterministic
+ return False
+
+ def __eq__(self, other):
+ return (
+ type(self) == type(other) and self._key_coder == other._key_coder and
+ self._value_coder == other._value_coder)
+
+ def __hash__(self):
+ return hash(type(self)) + hash(self._key_coder) + hash(self._value_coder)
+
+
+class NullableCoder(FastCoder):
+ def __init__(self, value_coder):
+ # type: (Coder) -> None
+ self._value_coder = value_coder
+
+ def _create_impl(self):
+ return coder_impl.NullableCoderImpl(self._value_coder.get_impl())
+
+ def to_type_hint(self):
+ return typehints.Optional[self._value_coder.to_type_hint()]
+
+ def is_deterministic(self):
+ # () -> bool
+ return self._value_coder.is_deterministic()
+
+ def __eq__(self, other):
+ return (
+ type(self) == type(other) and self._value_coder == other._value_coder)
+
+ def __hash__(self):
+ return hash(type(self)) + hash(self._value_coder)
+
+
class VarIntCoder(FastCoder):
"""Variable-length integer coder."""
def _create_impl(self):
diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py
index f1a771a..e1ce23b 100644
--- a/sdks/python/apache_beam/coders/coders_test_common.py
+++ b/sdks/python/apache_beam/coders/coders_test_common.py
@@ -82,8 +82,8 @@
coders.ToBytesCoder
])
cls.seen_nested -= set([coders.ProtoCoder, CustomCoder])
- assert not standard - cls.seen
- assert not cls.seen_nested - standard
+ assert not standard - cls.seen, str(standard - cls.seen)
+ assert not cls.seen_nested - standard, str(cls.seen_nested - standard)
@classmethod
def _observe(cls, coder):
@@ -560,6 +560,16 @@
context=context,
test_size_estimation=False)
+ def test_nullable_coder(self):
+ self.check_coder(coders.NullableCoder(coders.VarIntCoder()), None, 2 * 64)
+
+ def test_map_coder(self):
+ self.check_coder(
+ coders.MapCoder(coders.VarIntCoder(), coders.StrUtf8Coder()), {
+ 1: "one", 300: "three hundred"
+ }, {}, {i: str(i)
+ for i in range(5000)})
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/coders/row_coder.py b/sdks/python/apache_beam/coders/row_coder.py
index 3ad880f..02c7f06 100644
--- a/sdks/python/apache_beam/coders/row_coder.py
+++ b/sdks/python/apache_beam/coders/row_coder.py
@@ -30,6 +30,8 @@
from apache_beam.coders.coders import FastCoder
from apache_beam.coders.coders import FloatCoder
from apache_beam.coders.coders import IterableCoder
+from apache_beam.coders.coders import MapCoder
+from apache_beam.coders.coders import NullableCoder
from apache_beam.coders.coders import StrUtf8Coder
from apache_beam.coders.coders import TupleCoder
from apache_beam.coders.coders import VarIntCoder
@@ -58,8 +60,9 @@
to encode/decode.
"""
self.schema = schema
+ # Use non-null coders because null values are represented separately
self.components = [
- RowCoder.coder_from_type(field.type) for field in self.schema.fields
+ _nonnull_coder_from_type(field.type) for field in self.schema.fields
]
def _create_impl(self):
@@ -102,32 +105,6 @@
# type: (bytes) -> RowCoder
return RowCoder(proto_utils.parse_Bytes(payload, schema_pb2.Schema))
- @staticmethod
- def coder_from_type(field_type):
- type_info = field_type.WhichOneof("type_info")
- if type_info == "atomic_type":
- if field_type.atomic_type in (schema_pb2.INT32, schema_pb2.INT64):
- return VarIntCoder()
- elif field_type.atomic_type == schema_pb2.DOUBLE:
- return FloatCoder()
- elif field_type.atomic_type == schema_pb2.STRING:
- return StrUtf8Coder()
- elif field_type.atomic_type == schema_pb2.BOOLEAN:
- return BooleanCoder()
- elif field_type.atomic_type == schema_pb2.BYTES:
- return BytesCoder()
- elif type_info == "array_type":
- return IterableCoder(
- RowCoder.coder_from_type(field_type.array_type.element_type))
- elif type_info == "row_type":
- return RowCoder(field_type.row_type.schema)
-
- # The Java SDK supports several more types, but the coders are not yet
- # standard, and are not implemented in Python.
- raise ValueError(
- "Encountered a type that is not currently supported by RowCoder: %s" %
- field_type)
-
def __reduce__(self):
# when pickling, use bytes representation of the schema. schema_pb2.Schema
# objects cannot be pickled.
@@ -137,6 +114,43 @@
typecoders.registry.register_coder(row_type.RowTypeConstraint, RowCoder)
+def _coder_from_type(field_type):
+ coder = _nonnull_coder_from_type(field_type)
+ if field_type.nullable:
+ return NullableCoder(coder)
+ else:
+ return coder
+
+
+def _nonnull_coder_from_type(field_type):
+ type_info = field_type.WhichOneof("type_info")
+ if type_info == "atomic_type":
+ if field_type.atomic_type in (schema_pb2.INT32, schema_pb2.INT64):
+ return VarIntCoder()
+ elif field_type.atomic_type == schema_pb2.DOUBLE:
+ return FloatCoder()
+ elif field_type.atomic_type == schema_pb2.STRING:
+ return StrUtf8Coder()
+ elif field_type.atomic_type == schema_pb2.BOOLEAN:
+ return BooleanCoder()
+ elif field_type.atomic_type == schema_pb2.BYTES:
+ return BytesCoder()
+ elif type_info == "array_type":
+ return IterableCoder(_coder_from_type(field_type.array_type.element_type))
+ elif type_info == "map_type":
+ return MapCoder(
+ _coder_from_type(field_type.map_type.key_type),
+ _coder_from_type(field_type.map_type.value_type))
+ elif type_info == "row_type":
+ return RowCoder(field_type.row_type.schema)
+
+ # The Java SDK supports several more types, but the coders are not yet
+ # standard, and are not implemented in Python.
+ raise ValueError(
+ "Encountered a type that is not currently supported by RowCoder: %s" %
+ field_type)
+
+
class RowCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
SIZE_CODER = VarIntCoder().get_impl()
diff --git a/sdks/python/apache_beam/coders/row_coder_test.py b/sdks/python/apache_beam/coders/row_coder_test.py
index 65b1024..4277e58 100644
--- a/sdks/python/apache_beam/coders/row_coder_test.py
+++ b/sdks/python/apache_beam/coders/row_coder_test.py
@@ -45,14 +45,15 @@
("aliases", typing.List[unicode]),
("knows_javascript", bool),
# TODO(BEAM-7372): Use bytes instead of ByteString
- ("payload", typing.Optional[typing.ByteString])
+ ("payload", typing.Optional[typing.ByteString]),
+ ("custom_metadata", typing.Mapping[unicode, int])
])
coders_registry.register_coder(Person, RowCoder)
class RowCoderTest(unittest.TestCase):
- JON_SNOW = Person("Jon Snow", 23, None, ["crow", "wildling"], False, None)
+ JON_SNOW = Person("Jon Snow", 23, None, ["crow", "wildling"], False, None, {})
PEOPLE = [
JON_SNOW,
Person(
@@ -60,8 +61,9 @@
25,
"Westeros", ["Mother of Dragons"],
False,
- None),
- Person("Michael Bluth", 30, None, [], True, b"I've made a huge mistake")
+ None, {"dragons": 3}),
+ Person(
+ "Michael Bluth", 30, None, [], True, b"I've made a huge mistake", {})
]
def test_create_row_coder_from_named_tuple(self):
@@ -102,6 +104,15 @@
name="payload",
type=schema_pb2.FieldType(
atomic_type=schema_pb2.BYTES, nullable=True)),
+ schema_pb2.Field(
+ name="custom_metadata",
+ type=schema_pb2.FieldType(
+ map_type=schema_pb2.MapType(
+ key_type=schema_pb2.FieldType(
+ atomic_type=schema_pb2.STRING),
+ value_type=schema_pb2.FieldType(
+ atomic_type=schema_pb2.INT64),
+ ))),
])
coder = RowCoder(schema)
diff --git a/sdks/python/apache_beam/coders/standard_coders_test.py b/sdks/python/apache_beam/coders/standard_coders_test.py
index 4a13895..df88978 100644
--- a/sdks/python/apache_beam/coders/standard_coders_test.py
+++ b/sdks/python/apache_beam/coders/standard_coders_test.py
@@ -29,6 +29,7 @@
import sys
import unittest
from builtins import map
+from copy import deepcopy
from typing import Dict
from typing import Tuple
@@ -78,6 +79,13 @@
def value_parser_from_schema(schema):
def attribute_parser_from_type(type_):
+ parser = nonnull_attribute_parser_from_type(type_)
+ if type_.nullable:
+ return lambda x: None if x is None else parser(x)
+ else:
+ return parser
+
+ def nonnull_attribute_parser_from_type(type_):
# TODO: This should be exhaustive
type_info = type_.WhichOneof("type_info")
if type_info == "atomic_type":
@@ -89,8 +97,8 @@
element_parser = attribute_parser_from_type(type_.array_type.element_type)
return lambda x: list(map(element_parser, x))
elif type_info == "map_type":
- key_parser = attribute_parser_from_type(type_.array_type.key_type)
- value_parser = attribute_parser_from_type(type_.array_type.value_type)
+ key_parser = attribute_parser_from_type(type_.map_type.key_type)
+ value_parser = attribute_parser_from_type(type_.map_type.value_type)
return lambda x: dict(
(key_parser(k), value_parser(v)) for k, v in x.items())
@@ -101,6 +109,7 @@
def value_parser(x):
result = []
+ x = deepcopy(x)
for name, parser in parsers:
value = x.pop(name)
result.append(None if value is None else parser(value))
diff --git a/sdks/python/apache_beam/dataframe/__init__.py b/sdks/python/apache_beam/dataframe/__init__.py
index 427fee1..9071a88 100644
--- a/sdks/python/apache_beam/dataframe/__init__.py
+++ b/sdks/python/apache_beam/dataframe/__init__.py
@@ -13,3 +13,7 @@
# WITHOUT 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 apache_beam.dataframe.expressions import allow_non_parallel_operations
diff --git a/sdks/python/apache_beam/dataframe/doctests.py b/sdks/python/apache_beam/dataframe/doctests.py
index 2d135ea..93f33d1 100644
--- a/sdks/python/apache_beam/dataframe/doctests.py
+++ b/sdks/python/apache_beam/dataframe/doctests.py
@@ -361,6 +361,8 @@
original_doc_test_runner = doctest.DocTestRunner
doctest.DocTestRunner = lambda **kwargs: BeamDataframeDoctestRunner(
env, use_beam=use_beam, skip=skip, **kwargs)
- return func(*args, extraglobs=extraglobs, optionflags=optionflags, **kwargs)
+ with expressions.allow_non_parallel_operations():
+ return func(
+ *args, extraglobs=extraglobs, optionflags=optionflags, **kwargs)
finally:
doctest.DocTestRunner = original_doc_test_runner
diff --git a/sdks/python/apache_beam/dataframe/expressions.py b/sdks/python/apache_beam/dataframe/expressions.py
index 376efb8..34e01cca 100644
--- a/sdks/python/apache_beam/dataframe/expressions.py
+++ b/sdks/python/apache_beam/dataframe/expressions.py
@@ -16,6 +16,9 @@
from __future__ import absolute_import
+import contextlib
+import threading
+
from typing import Any
from typing import Callable
from typing import Iterable
@@ -203,6 +206,11 @@
be partitioned by index whenever all of its inputs are partitioned by
index.
"""
+ if (not _get_allow_non_parallel() and
+ requires_partition_by == partitionings.Singleton()):
+ raise NonParallelOperation(
+ "Using non-parallel form of %s "
+ "outside of allow_non_parallel_operations block." % name)
args = tuple(args)
if proxy is None:
proxy = func(*(arg.proxy() for arg in args))
@@ -236,3 +244,22 @@
args,
requires_partition_by=partitionings.Nothing(),
preserves_partition_by=partitionings.Singleton())
+
+
+_ALLOW_NON_PARALLEL = threading.local()
+_ALLOW_NON_PARALLEL.value = False
+
+
+def _get_allow_non_parallel():
+ return _ALLOW_NON_PARALLEL.value
+
+
+@contextlib.contextmanager
+def allow_non_parallel_operations(allow=True):
+ old_value, _ALLOW_NON_PARALLEL.value = _ALLOW_NON_PARALLEL.value, allow
+ yield
+ _ALLOW_NON_PARALLEL.value = old_value
+
+
+class NonParallelOperation(Exception):
+ pass
diff --git a/sdks/python/apache_beam/dataframe/frame_base.py b/sdks/python/apache_beam/dataframe/frame_base.py
index 7780cd3..0b1c296 100644
--- a/sdks/python/apache_beam/dataframe/frame_base.py
+++ b/sdks/python/apache_beam/dataframe/frame_base.py
@@ -16,7 +16,9 @@
from __future__ import absolute_import
+import functools
import inspect
+import sys
from typing import Any
from typing import Callable
from typing import Dict
@@ -29,6 +31,18 @@
from apache_beam.dataframe import expressions
from apache_beam.dataframe import partitionings
+# pylint: disable=deprecated-method
+if sys.version_info < (3, ):
+ _getargspec = inspect.getargspec
+
+ def _unwrap(func):
+ while hasattr(func, '__wrapped__'):
+ func = func.__wrapped__
+ return func
+else:
+ _getargspec = inspect.getfullargspec
+ _unwrap = inspect.unwrap
+
class DeferredBase(object):
@@ -146,8 +160,7 @@
value = kwargs[key]
else:
try:
- # pylint: disable=deprecated-method
- ix = inspect.getargspec(func).args.index(key)
+ ix = _getargspec(func).args.index(key)
except ValueError:
# TODO: fix for delegation?
continue
@@ -226,6 +239,68 @@
return wrapper
+def maybe_inplace(func):
+ @functools.wraps(func)
+ def wrapper(self, inplace=False, **kwargs):
+ result = func(self, **kwargs)
+ if inplace:
+ self._expr = result._expr
+ else:
+ return result
+
+ return wrapper
+
+
+def args_to_kwargs(base_type):
+ def wrap(func):
+ arg_names = _getargspec(_unwrap(getattr(base_type, func.__name__))).args
+
+ @functools.wraps(func)
+ def wrapper(*args, **kwargs):
+ for name, value in zip(arg_names, args):
+ if name in kwargs:
+ raise TypeError(
+ "%s() got multiple values for argument '%s'" %
+ (func.__name__, name))
+ kwargs[name] = value
+ return func(**kwargs)
+
+ return wrapper
+
+ return wrap
+
+
+def populate_defaults(base_type):
+ def wrap(func):
+ base_argspec = _getargspec(_unwrap(getattr(base_type, func.__name__)))
+ if not base_argspec.defaults:
+ return func
+
+ arg_to_default = dict(
+ zip(
+ base_argspec.args[-len(base_argspec.defaults):],
+ base_argspec.defaults))
+
+ unwrapped_func = _unwrap(func)
+ # args that do not have defaults in func, but do have defaults in base
+ func_argspec = _getargspec(unwrapped_func)
+ num_non_defaults = len(func_argspec.args) - len(func_argspec.defaults or ())
+ defaults_to_populate = set(
+ func_argspec.args[:num_non_defaults]).intersection(
+ arg_to_default.keys())
+
+ @functools.wraps(func)
+ def wrapper(**kwargs):
+ for name in defaults_to_populate:
+ if name not in kwargs:
+ kwargs[name] = arg_to_default[name]
+ return func(**kwargs)
+
+ return wrapper
+
+ return wrap
+
+
class WontImplementError(NotImplementedError):
"""An subclass of NotImplementedError to raise indicating that implementing
the given method is infeasible.
diff --git a/sdks/python/apache_beam/dataframe/frame_base_test.py b/sdks/python/apache_beam/dataframe/frame_base_test.py
index 392272c..b527da0 100644
--- a/sdks/python/apache_beam/dataframe/frame_base_test.py
+++ b/sdks/python/apache_beam/dataframe/frame_base_test.py
@@ -41,6 +41,59 @@
self.assertTrue(sub(x, b)._expr.evaluate_at(session).equals(a - b))
self.assertTrue(sub(a, y)._expr.evaluate_at(session).equals(a - b))
+ def test_maybe_inplace(self):
+ @frame_base.maybe_inplace
+ def add_one(frame):
+ return frame + 1
+
+ frames.DeferredSeries.add_one = add_one
+ original_expr = expressions.PlaceholderExpression(pd.Series([1, 2, 3]))
+ x = frames.DeferredSeries(original_expr)
+ x.add_one()
+ self.assertIs(x._expr, original_expr)
+ x.add_one(inplace=False)
+ self.assertIs(x._expr, original_expr)
+ x.add_one(inplace=True)
+ self.assertIsNot(x._expr, original_expr)
+
+ def test_args_to_kwargs(self):
+ class Base(object):
+ def func(self, a=1, b=2, c=3):
+ pass
+
+ class Proxy(object):
+ @frame_base.args_to_kwargs(Base)
+ def func(self, **kwargs):
+ return kwargs
+
+ proxy = Proxy()
+ # pylint: disable=too-many-function-args
+ self.assertEqual(proxy.func(), {})
+ self.assertEqual(proxy.func(100), {'a': 100})
+ self.assertEqual(proxy.func(2, 4, 6), {'a': 2, 'b': 4, 'c': 6})
+ self.assertEqual(proxy.func(2, c=6), {'a': 2, 'c': 6})
+ self.assertEqual(proxy.func(c=6, a=2), {'a': 2, 'c': 6})
+
+ def test_args_to_kwargs_populates_defaults(self):
+ class Base(object):
+ def func(self, a=1, b=2, c=3):
+ pass
+
+ class Proxy(object):
+ @frame_base.args_to_kwargs(Base)
+ @frame_base.populate_defaults(Base)
+ def func(self, a, c=1000, **kwargs):
+ return dict(kwargs, a=a, c=c)
+
+ proxy = Proxy()
+ # pylint: disable=too-many-function-args
+ self.assertEqual(proxy.func(), {'a': 1, 'c': 1000})
+ self.assertEqual(proxy.func(100), {'a': 100, 'c': 1000})
+ self.assertEqual(proxy.func(2, 4, 6), {'a': 2, 'b': 4, 'c': 6})
+ self.assertEqual(proxy.func(2, c=6), {'a': 2, 'c': 6})
+ self.assertEqual(proxy.func(c=6, a=2), {'a': 2, 'c': 6})
+ self.assertEqual(proxy.func(c=6), {'a': 1, 'c': 6})
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py
index 89e9154..9e2e97a 100644
--- a/sdks/python/apache_beam/dataframe/frames.py
+++ b/sdks/python/apache_beam/dataframe/frames.py
@@ -54,29 +54,20 @@
'order-sensitive')
diff = frame_base.wont_implement_method('order-sensitive')
- def replace(
- self,
- to_replace=None,
- value=None,
- inplace=False,
- limit=None,
- *args,
- **kwargs):
+ @frame_base.args_to_kwargs(pd.Series)
+ @frame_base.populate_defaults(pd.Series)
+ @frame_base.maybe_inplace
+ def replace(self, limit, **kwargs):
if limit is None:
requires_partition_by = partitionings.Nothing()
else:
requires_partition_by = partitionings.Singleton()
- result = frame_base.DeferredFrame.wrap(
+ return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'replace',
- lambda df: df.replace(
- to_replace, value, False, limit, *args, **kwargs), [self._expr],
+ lambda df: df.replace(limit=limit, **kwargs), [self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=requires_partition_by))
- if inplace:
- self._expr = result._expr
- else:
- return result
def unstack(self, *args, **kwargs):
raise frame_base.WontImplementError('non-deferred column values')
@@ -159,14 +150,15 @@
def loc(self):
return _DeferredLoc(self)
- def aggregate(self, *args, **kwargs):
- if 'axis' in kwargs and kwargs['axis'] is None:
- return self.agg(*args, **dict(kwargs, axis=1)).agg(
- *args, **dict(kwargs, axis=0))
+ @frame_base.args_to_kwargs(pd.DataFrame)
+ @frame_base.populate_defaults(pd.DataFrame)
+ def aggregate(self, axis, **kwargs):
+ if axis is None:
+ return self.agg(axis=1, **kwargs).agg(axis=0, **kwargs)
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'aggregate',
- lambda df: df.agg(*args, **kwargs),
+ lambda df: df.agg(axis=axis, **kwargs),
[self._expr],
# TODO(robertwb): Sub-aggregate when possible.
requires_partition_by=partitionings.Singleton()))
@@ -188,32 +180,22 @@
min = frame_base._associative_agg_method('min')
mode = frame_base._agg_method('mode')
- def dropna(
- self,
- axis=0,
- how='any',
- thresh=None,
- subset=None,
- inplace=False,
- *args,
- **kwargs):
+ @frame_base.args_to_kwargs(pd.DataFrame)
+ @frame_base.populate_defaults(pd.DataFrame)
+ @frame_base.maybe_inplace
+ def dropna(self, axis, **kwargs):
# TODO(robertwb): This is a common pattern. Generalize?
if axis == 1 or axis == 'columns':
requires_partition_by = partitionings.Singleton()
else:
requires_partition_by = partitionings.Nothing()
- result = frame_base.DeferredFrame.wrap(
+ return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'dropna',
- lambda df: df.dropna(
- axis, how, thresh, subset, False, *args, **kwargs),
+ lambda df: df.dropna(axis=axis, **kwargs),
[self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=requires_partition_by))
- if inplace:
- self._expr = result._expr
- else:
- return result
items = itertuples = iterrows = iteritems = frame_base.wont_implement_method(
'non-lazy')
@@ -223,13 +205,15 @@
prod = product = frame_base._associative_agg_method('prod')
- def quantile(self, q=0.5, axis=0, *args, **kwargs):
- if axis != 0:
+ @frame_base.args_to_kwargs(pd.DataFrame)
+ @frame_base.populate_defaults(pd.DataFrame)
+ def quantile(self, axis, **kwargs):
+ if axis == 1 or axis == 'columns':
raise frame_base.WontImplementError('non-deferred column values')
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'quantile',
- lambda df: df.quantile(q, axis, *args, **kwargs),
+ lambda df: df.quantile(axis=axis, **kwargs),
[self._expr],
#TODO(robertwb): Approximate quantiles?
requires_partition_by=partitionings.Singleton(),
@@ -237,28 +221,26 @@
query = frame_base._elementwise_method('query')
- def replace(self, to_replace=None,
- value=None,
- inplace=False,
- limit=None, *args, **kwargs):
+ @frame_base.args_to_kwargs(pd.DataFrame)
+ @frame_base.populate_defaults(pd.DataFrame)
+ @frame_base.maybe_inplace
+ def replace(self, limit, **kwargs):
if limit is None:
requires_partition_by = partitionings.Nothing()
else:
requires_partition_by = partitionings.Singleton()
- result = frame_base.DeferredFrame.wrap(
+ return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'replace',
- lambda df: df.replace(
- to_replace, value, False, limit, *args, **kwargs),
+ lambda df: df.replace(limit=limit, **kwargs),
[self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=requires_partition_by))
- if inplace:
- self._expr = result._expr
- else:
- return result
- def reset_index(self, level=None, drop=False, inplace=False, *args, **kwargs):
+ @frame_base.args_to_kwargs(pd.DataFrame)
+ @frame_base.populate_defaults(pd.DataFrame)
+ @frame_base.maybe_inplace
+ def reset_index(self, level, **kwargs):
if level is not None and not isinstance(level, (tuple, list)):
level = [level]
if level is None or len(level) == len(self._expr.proxy().index.levels):
@@ -266,22 +248,20 @@
requires_partition_by = partitionings.Singleton()
else:
requires_partition_by = partitionings.Nothing()
- result = frame_base.DeferredFrame.wrap(
+ return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'reset_index',
- lambda df: df.reset_index(level, drop, False, *args, **kwargs),
+ lambda df: df.reset_index(level=level, **kwargs),
[self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=requires_partition_by))
- if inplace:
- self._expr = result._expr
- else:
- return result
round = frame_base._elementwise_method('round')
select_dtypes = frame_base._elementwise_method('select_dtypes')
- def shift(self, periods=1, freq=None, axis=0, *args, **kwargs):
+ @frame_base.args_to_kwargs(pd.DataFrame)
+ @frame_base.populate_defaults(pd.DataFrame)
+ def shift(self, axis, **kwargs):
if axis == 1 or axis == 'columns':
requires_partition_by = partitionings.Nothing()
else:
@@ -289,7 +269,7 @@
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'shift',
- lambda df: df.shift(periods, freq, axis, *args, **kwargs),
+ lambda df: df.shift(axis=axis, **kwargs),
[self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=requires_partition_by))
@@ -298,24 +278,21 @@
def shape(self):
raise frame_base.WontImplementError('scalar value')
- def sort_values(
- self, by, axis=0, ascending=True, inplace=False, *args, **kwargs):
+ @frame_base.args_to_kwargs(pd.DataFrame)
+ @frame_base.populate_defaults(pd.DataFrame)
+ @frame_base.maybe_inplace
+ def sort_values(self, axis, **kwargs):
if axis == 1 or axis == 'columns':
requires_partition_by = partitionings.Nothing()
else:
requires_partition_by = partitionings.Singleton()
- result = frame_base.DeferredFrame.wrap(
+ return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'sort_values',
- lambda df: df.sort_values(
- by, axis, ascending, False, *args, **kwargs),
+ lambda df: df.sort_values(axis=axis, **kwargs),
[self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=requires_partition_by))
- if inplace:
- self._expr = result._expr
- else:
- return result
stack = frame_base._elementwise_method('stack')
diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py
index 1383a52..773b3ba 100644
--- a/sdks/python/apache_beam/dataframe/frames_test.py
+++ b/sdks/python/apache_beam/dataframe/frames_test.py
@@ -21,6 +21,7 @@
import numpy as np
import pandas as pd
+import apache_beam as beam
from apache_beam.dataframe import expressions
from apache_beam.dataframe import frame_base
from apache_beam.dataframe import frames # pylint: disable=unused-import
@@ -81,5 +82,35 @@
self._run_test(lambda df: df.loc[lambda df: df.A > 10], df)
+class AllowNonParallelTest(unittest.TestCase):
+ def _use_non_parallel_operation(self):
+ _ = frame_base.DeferredFrame.wrap(
+ expressions.PlaceholderExpression(pd.Series([1, 2, 3]))).replace(
+ 'a', 'b', limit=1)
+
+ def test_disallow_non_parallel(self):
+ with self.assertRaises(expressions.NonParallelOperation):
+ self._use_non_parallel_operation()
+
+ def test_allow_non_parallel_in_context(self):
+ with beam.dataframe.allow_non_parallel_operations():
+ self._use_non_parallel_operation()
+
+ def test_allow_non_parallel_nesting(self):
+ # disallowed
+ with beam.dataframe.allow_non_parallel_operations():
+ # allowed
+ self._use_non_parallel_operation()
+ with beam.dataframe.allow_non_parallel_operations(False):
+ # disallowed again
+ with self.assertRaises(expressions.NonParallelOperation):
+ self._use_non_parallel_operation()
+ # allowed
+ self._use_non_parallel_operation()
+ # disallowed
+ with self.assertRaises(expressions.NonParallelOperation):
+ self._use_non_parallel_operation()
+
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/dataframe/transforms_test.py b/sdks/python/apache_beam/dataframe/transforms_test.py
index 81917cf..e010b71 100644
--- a/sdks/python/apache_beam/dataframe/transforms_test.py
+++ b/sdks/python/apache_beam/dataframe/transforms_test.py
@@ -81,8 +81,9 @@
'Animal': ['Falcon', 'Falcon', 'Parrot', 'Parrot'],
'Speed': [380., 370., 24., 26.]
})
- self.run_scenario(df, lambda df: df.groupby('Animal').sum())
- self.run_scenario(df, lambda df: df.groupby('Animal').mean())
+ with expressions.allow_non_parallel_operations():
+ self.run_scenario(df, lambda df: df.groupby('Animal').sum())
+ self.run_scenario(df, lambda df: df.groupby('Animal').mean())
def test_filter(self):
df = pd.DataFrame({
@@ -95,19 +96,21 @@
df, lambda df: df.set_index('Animal').filter(regex='F.*', axis='index'))
def test_aggregate(self):
- a = pd.DataFrame({'col': [1, 2, 3]})
- self.run_scenario(a, lambda a: a.agg(sum))
- self.run_scenario(a, lambda a: a.agg(['mean', 'min', 'max']))
+ with expressions.allow_non_parallel_operations():
+ a = pd.DataFrame({'col': [1, 2, 3]})
+ self.run_scenario(a, lambda a: a.agg(sum))
+ self.run_scenario(a, lambda a: a.agg(['mean', 'min', 'max']))
def test_scalar(self):
- a = pd.Series([1, 2, 6])
- self.run_scenario(a, lambda a: a.agg(sum))
- self.run_scenario(a, lambda a: a / a.agg(sum))
+ with expressions.allow_non_parallel_operations():
+ a = pd.Series([1, 2, 6])
+ self.run_scenario(a, lambda a: a.agg(sum))
+ self.run_scenario(a, lambda a: a / a.agg(sum))
- # Tests scalar being used as an input to a downstream stage.
- df = pd.DataFrame({'key': ['a', 'a', 'b'], 'val': [1, 2, 6]})
- self.run_scenario(
- df, lambda df: df.groupby('key').sum().val / df.val.agg(sum))
+ # Tests scalar being used as an input to a downstream stage.
+ df = pd.DataFrame({'key': ['a', 'a', 'b'], 'val': [1, 2, 6]})
+ self.run_scenario(
+ df, lambda df: df.groupby('key').sum().val / df.val.agg(sum))
def test_input_output_polymorphism(self):
one_series = pd.Series([1])
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test_py3.py b/sdks/python/apache_beam/examples/snippets/snippets_test_py3.py
index 0f0b668..5eb1d4b 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_test_py3.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_test_py3.py
@@ -96,6 +96,18 @@
ids = numbers | 'to_id' >> beam.Map(my_fn)
# [END type_hints_map_annotations]
+ # Example using an annotated PTransform.
+ with self.assertRaises(typehints.TypeCheckError):
+ # [START type_hints_ptransforms]
+ from apache_beam.pvalue import PCollection
+
+ class IntToStr(beam.PTransform):
+ def expand(self, pcoll: PCollection[int]) -> PCollection[str]:
+ return pcoll | beam.Map(lambda elem: str(elem))
+
+ ids = numbers | 'convert to str' >> IntToStr()
+ # [END type_hints_ptransforms]
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/examples/wordcount.py b/sdks/python/apache_beam/examples/wordcount.py
index 6732568..aa07802 100644
--- a/sdks/python/apache_beam/examples/wordcount.py
+++ b/sdks/python/apache_beam/examples/wordcount.py
@@ -30,24 +30,12 @@
import apache_beam as beam
from apache_beam.io import ReadFromText
from apache_beam.io import WriteToText
-from apache_beam.metrics import Metrics
-from apache_beam.metrics.metric import MetricsFilter
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions
class WordExtractingDoFn(beam.DoFn):
"""Parse each line of input text into words."""
- def __init__(self):
- # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
- # super(WordExtractingDoFn, self).__init__()
- beam.DoFn.__init__(self)
- self.words_counter = Metrics.counter(self.__class__, 'words')
- self.word_lengths_counter = Metrics.counter(self.__class__, 'word_lengths')
- self.word_lengths_dist = Metrics.distribution(
- self.__class__, 'word_len_dist')
- self.empty_line_counter = Metrics.counter(self.__class__, 'empty_lines')
-
def process(self, element):
"""Returns an iterator over the words of this element.
@@ -59,15 +47,7 @@
Returns:
The processed element.
"""
- text_line = element.strip()
- if not text_line:
- self.empty_line_counter.inc(1)
- words = re.findall(r'[\w\']+', text_line, re.UNICODE)
- for w in words:
- self.words_counter.inc()
- self.word_lengths_counter.inc(len(w))
- self.word_lengths_dist.update(len(w))
- return words
+ return re.findall(r'[\w\']+', element, re.UNICODE)
def run(argv=None, save_main_session=True):
@@ -89,52 +69,29 @@
# workflow rely on global context (e.g., a module imported at module level).
pipeline_options = PipelineOptions(pipeline_args)
pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
- p = beam.Pipeline(options=pipeline_options)
- # Read the text file[pattern] into a PCollection.
- lines = p | 'read' >> ReadFromText(known_args.input)
+ # The pipeline will be run on exiting the with block.
+ with beam.Pipeline(options=pipeline_options) as p:
- # Count the occurrences of each word.
- def count_ones(word_ones):
- (word, ones) = word_ones
- return (word, sum(ones))
+ # Read the text file[pattern] into a PCollection.
+ lines = p | 'Read' >> ReadFromText(known_args.input)
- counts = (
- lines
- | 'split' >>
- (beam.ParDo(WordExtractingDoFn()).with_output_types(unicode))
- | 'pair_with_one' >> beam.Map(lambda x: (x, 1))
- | 'group' >> beam.GroupByKey()
- | 'count' >> beam.Map(count_ones))
+ counts = (
+ lines
+ | 'Split' >>
+ (beam.ParDo(WordExtractingDoFn()).with_output_types(unicode))
+ | 'PairWIthOne' >> beam.Map(lambda x: (x, 1))
+ | 'GroupAndSum' >> beam.CombinePerKey(sum))
- # Format the counts into a PCollection of strings.
- def format_result(word_count):
- (word, count) = word_count
- return '%s: %d' % (word, count)
+ # Format the counts into a PCollection of strings.
+ def format_result(word, count):
+ return '%s: %d' % (word, count)
- output = counts | 'format' >> beam.Map(format_result)
+ output = counts | 'Format' >> beam.MapTuple(format_result)
- # Write the output using a "Write" transform that has side effects.
- # pylint: disable=expression-not-assigned
- output | 'write' >> WriteToText(known_args.output)
-
- result = p.run()
- result.wait_until_finish()
-
- # Do not query metrics when creating a template which doesn't run
- if (not hasattr(result, 'has_job') # direct runner
- or result.has_job): # not just a template creation
- empty_lines_filter = MetricsFilter().with_name('empty_lines')
- query_result = result.metrics().query(empty_lines_filter)
- if query_result['counters']:
- empty_lines_counter = query_result['counters'][0]
- logging.info('number of empty lines: %d', empty_lines_counter.result)
-
- word_lengths_filter = MetricsFilter().with_name('word_len_dist')
- query_result = result.metrics().query(word_lengths_filter)
- if query_result['distributions']:
- word_lengths_dist = query_result['distributions'][0]
- logging.info('average word length: %d', word_lengths_dist.result.mean)
+ # Write the output using a "Write" transform that has side effects.
+ # pylint: disable=expression-not-assigned
+ output | 'Write' >> WriteToText(known_args.output)
if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/examples/wordcount_with_metrics.py b/sdks/python/apache_beam/examples/wordcount_with_metrics.py
new file mode 100644
index 0000000..6732568
--- /dev/null
+++ b/sdks/python/apache_beam/examples/wordcount_with_metrics.py
@@ -0,0 +1,142 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""A word-counting workflow."""
+
+# pytype: skip-file
+
+from __future__ import absolute_import
+
+import argparse
+import logging
+import re
+
+from past.builtins import unicode
+
+import apache_beam as beam
+from apache_beam.io import ReadFromText
+from apache_beam.io import WriteToText
+from apache_beam.metrics import Metrics
+from apache_beam.metrics.metric import MetricsFilter
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import SetupOptions
+
+
+class WordExtractingDoFn(beam.DoFn):
+ """Parse each line of input text into words."""
+ def __init__(self):
+ # TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
+ # super(WordExtractingDoFn, self).__init__()
+ beam.DoFn.__init__(self)
+ self.words_counter = Metrics.counter(self.__class__, 'words')
+ self.word_lengths_counter = Metrics.counter(self.__class__, 'word_lengths')
+ self.word_lengths_dist = Metrics.distribution(
+ self.__class__, 'word_len_dist')
+ self.empty_line_counter = Metrics.counter(self.__class__, 'empty_lines')
+
+ def process(self, element):
+ """Returns an iterator over the words of this element.
+
+ The element is a line of text. If the line is blank, note that, too.
+
+ Args:
+ element: the element being processed
+
+ Returns:
+ The processed element.
+ """
+ text_line = element.strip()
+ if not text_line:
+ self.empty_line_counter.inc(1)
+ words = re.findall(r'[\w\']+', text_line, re.UNICODE)
+ for w in words:
+ self.words_counter.inc()
+ self.word_lengths_counter.inc(len(w))
+ self.word_lengths_dist.update(len(w))
+ return words
+
+
+def run(argv=None, save_main_session=True):
+ """Main entry point; defines and runs the wordcount pipeline."""
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ '--input',
+ dest='input',
+ default='gs://dataflow-samples/shakespeare/kinglear.txt',
+ help='Input file to process.')
+ parser.add_argument(
+ '--output',
+ dest='output',
+ required=True,
+ help='Output file to write results to.')
+ known_args, pipeline_args = parser.parse_known_args(argv)
+
+ # We use the save_main_session option because one or more DoFn's in this
+ # workflow rely on global context (e.g., a module imported at module level).
+ pipeline_options = PipelineOptions(pipeline_args)
+ pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
+ p = beam.Pipeline(options=pipeline_options)
+
+ # Read the text file[pattern] into a PCollection.
+ lines = p | 'read' >> ReadFromText(known_args.input)
+
+ # Count the occurrences of each word.
+ def count_ones(word_ones):
+ (word, ones) = word_ones
+ return (word, sum(ones))
+
+ counts = (
+ lines
+ | 'split' >>
+ (beam.ParDo(WordExtractingDoFn()).with_output_types(unicode))
+ | 'pair_with_one' >> beam.Map(lambda x: (x, 1))
+ | 'group' >> beam.GroupByKey()
+ | 'count' >> beam.Map(count_ones))
+
+ # Format the counts into a PCollection of strings.
+ def format_result(word_count):
+ (word, count) = word_count
+ return '%s: %d' % (word, count)
+
+ output = counts | 'format' >> beam.Map(format_result)
+
+ # Write the output using a "Write" transform that has side effects.
+ # pylint: disable=expression-not-assigned
+ output | 'write' >> WriteToText(known_args.output)
+
+ result = p.run()
+ result.wait_until_finish()
+
+ # Do not query metrics when creating a template which doesn't run
+ if (not hasattr(result, 'has_job') # direct runner
+ or result.has_job): # not just a template creation
+ empty_lines_filter = MetricsFilter().with_name('empty_lines')
+ query_result = result.metrics().query(empty_lines_filter)
+ if query_result['counters']:
+ empty_lines_counter = query_result['counters'][0]
+ logging.info('number of empty lines: %d', empty_lines_counter.result)
+
+ word_lengths_filter = MetricsFilter().with_name('word_len_dist')
+ query_result = result.metrics().query(word_lengths_filter)
+ if query_result['distributions']:
+ word_lengths_dist = query_result['distributions'][0]
+ logging.info('average word length: %d', word_lengths_dist.result.mean)
+
+
+if __name__ == '__main__':
+ logging.getLogger().setLevel(logging.INFO)
+ run()
diff --git a/sdks/python/apache_beam/io/external/snowflake.py b/sdks/python/apache_beam/io/external/snowflake.py
new file mode 100644
index 0000000..e7ffa6a
--- /dev/null
+++ b/sdks/python/apache_beam/io/external/snowflake.py
@@ -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.
+#
+
+"""Snowflake transforms tested against Flink portable runner.
+
+ **Setup**
+
+ Transforms provided in this module are cross-language transforms
+ implemented in the Beam Java SDK. During the pipeline construction, Python SDK
+ will connect to a Java expansion service to expand these transforms.
+ To facilitate this, a small amount of setup is needed before using these
+ transforms in a Beam Python pipeline.
+
+ There are several ways to setup cross-language Snowflake transforms.
+
+ * Option 1: use the default expansion service
+ * Option 2: specify a custom expansion service
+
+ See below for details regarding each of these options.
+
+ *Option 1: Use the default expansion service*
+
+ This is the recommended and easiest setup option for using Python Snowflake
+ transforms.This option requires following pre-requisites
+ before running the Beam pipeline.
+
+ * Install Java runtime in the computer from where the pipeline is constructed
+ and make sure that 'java' command is available.
+
+ In this option, Python SDK will either download (for released Beam version) or
+ build (when running from a Beam Git clone) a expansion service jar and use
+ that to expand transforms. Currently Snowflake transforms use the
+ 'beam-sdks-java-io-expansion-service' jar for this purpose.
+
+ *Option 2: specify a custom expansion service*
+
+ In this option, you startup your own expansion service and provide that as
+ a parameter when using the transforms provided in this module.
+
+ This option requires following pre-requisites before running the Beam
+ pipeline.
+
+ * Startup your own expansion service.
+ * Update your pipeline to provide the expansion service address when
+ initiating Snowflake transforms provided in this module.
+
+ Flink Users can use the built-in Expansion Service of the Flink Runner's
+ Job Server. If you start Flink's Job Server, the expansion service will be
+ started on port 8097. For a different address, please set the
+ expansion_service parameter.
+
+ **More information**
+
+ For more information regarding cross-language transforms see:
+ - https://beam.apache.org/roadmap/portability/
+
+ For more information specific to Flink runner see:
+ - https://beam.apache.org/documentation/runners/flink/
+"""
+
+# pytype: skip-file
+
+from __future__ import absolute_import
+
+import typing
+
+from past.builtins import unicode
+
+import apache_beam as beam
+from apache_beam.transforms.external import BeamJarExpansionService
+from apache_beam.transforms.external import ExternalTransform
+from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder
+
+ReadFromSnowflakeSchema = typing.NamedTuple(
+ 'ReadFromSnowflakeSchema',
+ [
+ ('server_name', unicode),
+ ('schema', unicode),
+ ('database', unicode),
+ ('staging_bucket_name', unicode),
+ ('storage_integration_name', unicode),
+ ('username', typing.Optional[unicode]),
+ ('password', typing.Optional[unicode]),
+ ('private_key_path', typing.Optional[unicode]),
+ ('private_key_passphrase', typing.Optional[unicode]),
+ ('o_auth_token', typing.Optional[unicode]),
+ ('table', typing.Optional[unicode]),
+ ('query', typing.Optional[unicode]),
+ ])
+
+
+def default_io_expansion_service():
+ return BeamJarExpansionService(
+ 'sdks:java:io:snowflake:expansion-service:shadowJar')
+
+
+class ReadFromSnowflake(beam.PTransform):
+ """
+ An external PTransform which reads from Snowflake.
+ """
+
+ URN = 'beam:external:java:snowflake:read:v1'
+
+ def __init__(
+ self,
+ server_name,
+ schema,
+ database,
+ staging_bucket_name,
+ storage_integration_name,
+ csv_mapper,
+ username=None,
+ password=None,
+ private_key_path=None,
+ private_key_passphrase=None,
+ o_auth_token=None,
+ table=None,
+ query=None,
+ expansion_service=None):
+ """
+ Initializes a read operation from Snowflake.
+
+ Required parameters:
+
+ :param server_name: full Snowflake server name with the following format
+ account.region.gcp.snowflakecomputing.com.
+ :param schema: name of the Snowflake schema in the database to use.
+ :param database: name of the Snowflake database to use.
+ :param staging_bucket_name: name of the Google Cloud Storage bucket.::
+ Bucket will be used as a temporary location for storing CSV files.
+ Those temporary directories will be named
+ 'sf_copy_csv_DATE_TIME_RANDOMSUFFIX'
+ and they will be removed automatically once Read operation finishes.
+ :param storage_integration_name: is the name of storage integration
+ object created according to Snowflake documentation.
+ :param csv_mapper: specifies a function which must translate
+ user-defined object to array of strings.
+ SnowflakeIO uses a COPY INTO <location> statement to move data from
+ a Snowflake table to Google Cloud Storage as CSV files.These files
+ are then downloaded via FileIO and processed line by line.
+ Each line is split into an array of Strings using the OpenCSV
+ The csv_mapper function job is to give the user the possibility to
+ convert the array of Strings to a user-defined type,
+ ie. GenericRecord for Avro or Parquet files, or custom objects.
+ Example:
+ def csv_mapper(strings_array)
+ return User(strings_array[0], int(strings_array[1])))
+ :param table: specifies a Snowflake table name.
+ :param query: specifies a Snowflake custom SQL query.
+ :param expansion_service: specifies URL of expansion service.
+
+ Authentication parameters:
+
+ :param username: specifies username for
+ username/password authentication method.
+ :param password: specifies password for
+ username/password authentication method.
+ :param private_key_path: specifies a private key file for
+ key/ pair authentication method.
+ :param private_key_passphrase: specifies password for
+ key/ pair authentication method.
+ :param o_auth_token: specifies access token for
+ OAuth authentication method.
+ """
+ self.params = ReadFromSnowflakeSchema(
+ server_name=server_name,
+ schema=schema,
+ database=database,
+ staging_bucket_name=staging_bucket_name,
+ storage_integration_name=storage_integration_name,
+ username=username,
+ password=password,
+ private_key_path=private_key_path,
+ private_key_passphrase=private_key_passphrase,
+ o_auth_token=o_auth_token,
+ table=table,
+ query=query)
+ self.csv_mapper = csv_mapper
+ self.expansion_service = expansion_service or default_io_expansion_service()
+
+ def expand(self, pbegin):
+ return (
+ pbegin
+ | ExternalTransform(
+ self.URN,
+ NamedTupleBasedPayloadBuilder(self.params),
+ self.expansion_service,
+ )
+ | 'CSV to array mapper' >> beam.Map(lambda csv: csv.split(b','))
+ | 'CSV mapper' >> beam.Map(self.csv_mapper))
diff --git a/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py b/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
index 0dad234..cec1d9b 100644
--- a/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
@@ -28,12 +28,17 @@
import time
import typing
import unittest
+import uuid
import apache_beam as beam
from apache_beam.io.kafka import ReadFromKafka
from apache_beam.io.kafka import WriteToKafka
from apache_beam.metrics import Metrics
from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+
+NUM_RECORDS = 1000
class CrossLanguageKafkaIO(object):
@@ -47,7 +52,7 @@
_ = (
pipeline
| 'Impulse' >> beam.Impulse()
- | 'Generate' >> beam.FlatMap(lambda x: range(1000)) # pylint: disable=range-builtin-not-iterating
+ | 'Generate' >> beam.FlatMap(lambda x: range(NUM_RECORDS)) # pylint: disable=range-builtin-not-iterating
| 'Reshuffle' >> beam.Reshuffle()
| 'MakeKV' >> beam.Map(lambda x:
(b'', str(x).encode())).with_output_types(
@@ -57,8 +62,8 @@
topic=self.topic,
expansion_service=self.expansion_service))
- def build_read_pipeline(self, pipeline):
- _ = (
+ def build_read_pipeline(self, pipeline, max_num_records=None):
+ kafka_records = (
pipeline
| 'ReadFromKafka' >> ReadFromKafka(
consumer_config={
@@ -66,7 +71,14 @@
'auto.offset.reset': 'earliest'
},
topics=[self.topic],
- expansion_service=self.expansion_service)
+ max_num_records=max_num_records,
+ expansion_service=self.expansion_service))
+
+ if max_num_records:
+ return kafka_records
+
+ return (
+ kafka_records
| 'Windowing' >> beam.WindowInto(
beam.window.FixedWindows(300),
trigger=beam.transforms.trigger.AfterProcessingTime(60),
@@ -86,6 +98,30 @@
os.environ.get('LOCAL_KAFKA_JAR'),
"LOCAL_KAFKA_JAR environment var is not provided.")
class CrossLanguageKafkaIOTest(unittest.TestCase):
+ def test_kafkaio(self):
+ kafka_topic = 'xlang_kafkaio_test_{}'.format(uuid.uuid4())
+ local_kafka_jar = os.environ.get('LOCAL_KAFKA_JAR')
+ with self.local_kafka_service(local_kafka_jar) as kafka_port:
+ bootstrap_servers = '{}:{}'.format(
+ self.get_platform_localhost(), kafka_port)
+ pipeline_creator = CrossLanguageKafkaIO(bootstrap_servers, kafka_topic)
+
+ self.run_kafka_write(pipeline_creator)
+ self.run_kafka_read(pipeline_creator)
+
+ def run_kafka_write(self, pipeline_creator):
+ with TestPipeline() as pipeline:
+ pipeline.not_use_test_runner_api = True
+ pipeline_creator.build_write_pipeline(pipeline)
+
+ def run_kafka_read(self, pipeline_creator):
+ with TestPipeline() as pipeline:
+ pipeline.not_use_test_runner_api = True
+ result = pipeline_creator.build_read_pipeline(pipeline, NUM_RECORDS)
+ assert_that(
+ result,
+ equal_to([(b'', str(i).encode()) for i in range(NUM_RECORDS)]))
+
def get_platform_localhost(self):
if sys.platform == 'darwin':
return 'host.docker.internal'
@@ -119,18 +155,6 @@
if kafka_server:
kafka_server.kill()
- def test_kafkaio_write(self):
- local_kafka_jar = os.environ.get('LOCAL_KAFKA_JAR')
- with self.local_kafka_service(local_kafka_jar) as kafka_port:
- p = TestPipeline()
- p.not_use_test_runner_api = True
- xlang_kafkaio = CrossLanguageKafkaIO(
- '%s:%s' % (self.get_platform_localhost(), kafka_port),
- 'xlang_kafkaio_test')
- xlang_kafkaio.build_write_pipeline(p)
- job = p.run()
- job.wait_until_finish()
-
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index 1230d1d..826af40 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -63,7 +63,7 @@
# Protect against environments where bigquery library is not available.
# pylint: disable=wrong-import-order, wrong-import-position
try:
- from apitools.base.py.exceptions import HttpError
+ from apitools.base.py.exceptions import HttpError, HttpForbiddenError
except ImportError:
pass
@@ -290,9 +290,10 @@
"""
Get the location of tables referenced in a query.
- This method returns the location of the first referenced table in the query
- and depends on the BigQuery service to provide error handling for
- queries that reference tables in multiple locations.
+ This method returns the location of the first available referenced
+ table for user in the query and depends on the BigQuery service to
+ provide error handling for queries that reference tables in multiple
+ locations.
"""
reference = bigquery.JobReference(
jobId=uuid.uuid4().hex, projectId=project_id)
@@ -318,17 +319,25 @@
referenced_tables = response.statistics.query.referencedTables
if referenced_tables: # Guards against both non-empty and non-None
- table = referenced_tables[0]
- location = self.get_table_location(
- table.projectId, table.datasetId, table.tableId)
- _LOGGER.info(
- "Using location %r from table %r referenced by query %s",
- location,
- table,
- query)
- return location
+ for table in referenced_tables:
+ try:
+ location = self.get_table_location(
+ table.projectId, table.datasetId, table.tableId)
+ except HttpForbiddenError:
+ # Permission access for table (i.e. from authorized_view),
+ # try next one
+ continue
+ _LOGGER.info(
+ "Using location %r from table %r referenced by query %s",
+ location,
+ table,
+ query)
+ return location
- _LOGGER.debug("Query %s does not reference any tables.", query)
+ _LOGGER.debug(
+ "Query %s does not reference any tables or "
+ "you don't have permission to inspect them.",
+ query)
return None
@retry.with_exponential_backoff(
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 90716d0..985af6b 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -39,7 +39,6 @@
import apache_beam as beam
from apache_beam.internal.gcp.json_value import to_json_value
from apache_beam.io.gcp.bigquery import TableRowJsonCoder
-from apache_beam.io.gcp.bigquery_test import HttpError
from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR
from apache_beam.io.gcp.bigquery_tools import AvroRowWriter
from apache_beam.io.gcp.bigquery_tools import BigQueryJobTypes
@@ -50,6 +49,15 @@
from apache_beam.io.gcp.internal.clients import bigquery
from apache_beam.options.pipeline_options import PipelineOptions
+# Protect against environments where bigquery library is not available.
+# pylint: disable=wrong-import-order, wrong-import-position
+try:
+ from apitools.base.py.exceptions import HttpError, HttpForbiddenError
+except ImportError:
+ HttpError = None
+ HttpForbiddenError = None
+# pylint: enable=wrong-import-order, wrong-import-position
+
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestTableSchemaParser(unittest.TestCase):
@@ -267,6 +275,37 @@
'The maximum number of retries has been reached',
str(context.exception))
+ def test_get_query_location(self):
+ client = mock.Mock()
+ query = """
+ SELECT
+ av.column1, table.column1
+ FROM `dataset.authorized_view` as av
+ JOIN `dataset.table` as table ON av.column2 = table.column2
+ """
+ job = mock.MagicMock(spec=bigquery.Job)
+ job.statistics.query.referencedTables = [
+ bigquery.TableReference(
+ projectId="first_project_id",
+ datasetId="first_dataset",
+ tableId="table_used_by_authorized_view"),
+ bigquery.TableReference(
+ projectId="second_project_id",
+ datasetId="second_dataset",
+ tableId="table"),
+ ]
+ client.jobs.Insert.return_value = job
+
+ wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client)
+ wrapper.get_table_location = mock.Mock(
+ side_effect=[
+ HttpForbiddenError(response={'status': '404'}, url='', content=''),
+ "US"
+ ])
+ location = wrapper.get_query_location(
+ project_id="second_project_id", query=query, use_legacy_sql=False)
+ self.assertEqual("US", location)
+
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestBigQueryReader(unittest.TestCase):
diff --git a/sdks/python/apache_beam/io/kafka.py b/sdks/python/apache_beam/io/kafka.py
index 4336bec..dc75b73 100644
--- a/sdks/python/apache_beam/io/kafka.py
+++ b/sdks/python/apache_beam/io/kafka.py
@@ -97,6 +97,9 @@
('topics', typing.List[unicode]),
('key_deserializer', unicode),
('value_deserializer', unicode),
+ ('start_read_time', typing.Optional[int]),
+ ('max_num_records', typing.Optional[int]),
+ ('max_read_time', typing.Optional[int]),
])
@@ -125,24 +128,30 @@
topics,
key_deserializer=byte_array_deserializer,
value_deserializer=byte_array_deserializer,
- expansion_service=None):
+ start_read_time=None,
+ max_num_records=None,
+ max_read_time=None,
+ expansion_service=None,
+ ):
"""
Initializes a read operation from Kafka.
:param consumer_config: A dictionary containing the consumer configuration.
:param topics: A list of topic strings.
:param key_deserializer: A fully-qualified Java class name of a Kafka
- Deserializer for the topic's key, e.g.
- 'org.apache.kafka.common.
- serialization.LongDeserializer'.
- Default: 'org.apache.kafka.common.
- serialization.ByteArrayDeserializer'.
+ Deserializer for the topic's key, e.g.
+ 'org.apache.kafka.common.serialization.LongDeserializer'.
+ Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
:param value_deserializer: A fully-qualified Java class name of a Kafka
- Deserializer for the topic's value, e.g.
- 'org.apache.kafka.common.
- serialization.LongDeserializer'.
- Default: 'org.apache.kafka.common.
- serialization.ByteArrayDeserializer'.
+ Deserializer for the topic's value, e.g.
+ 'org.apache.kafka.common.serialization.LongDeserializer'.
+ Default: 'org.apache.kafka.common.serialization.ByteArrayDeserializer'.
+ :param start_read_time: Use timestamp to set up start offset in milliseconds
+ epoch.
+ :param max_num_records: Maximum amount of records to be read. Mainly used
+ for tests and demo applications.
+ :param max_read_time: Maximum amount of time in seconds the transform
+ executes. Mainly used for tests and demo applications.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
super(ReadFromKafka, self).__init__(
@@ -153,6 +162,9 @@
topics=topics,
key_deserializer=key_deserializer,
value_deserializer=value_deserializer,
+ max_num_records=max_num_records,
+ max_read_time=max_read_time,
+ start_read_time=start_read_time,
)),
expansion_service or default_io_expansion_service())
@@ -195,17 +207,13 @@
:param producer_config: A dictionary containing the producer configuration.
:param topic: A Kafka topic name.
:param key_deserializer: A fully-qualified Java class name of a Kafka
- Serializer for the topic's key, e.g.
- 'org.apache.kafka.common.
- serialization.LongSerializer'.
- Default: 'org.apache.kafka.common.
- serialization.ByteArraySerializer'.
+ Serializer for the topic's key, e.g.
+ 'org.apache.kafka.common.serialization.LongSerializer'.
+ Default: 'org.apache.kafka.common.serialization.ByteArraySerializer'.
:param value_deserializer: A fully-qualified Java class name of a Kafka
- Serializer for the topic's value, e.g.
- 'org.apache.kafka.common.
- serialization.LongSerializer'.
- Default: 'org.apache.kafka.common.
- serialization.ByteArraySerializer'.
+ Serializer for the topic's value, e.g.
+ 'org.apache.kafka.common.serialization.LongSerializer'.
+ Default: 'org.apache.kafka.common.serialization.ByteArraySerializer'.
:param expansion_service: The address (host:port) of the ExpansionService.
"""
super(WriteToKafka, self).__init__(
diff --git a/sdks/python/apache_beam/testing/test_stream_service.py b/sdks/python/apache_beam/testing/test_stream_service.py
index d69465f..618e208 100644
--- a/sdks/python/apache_beam/testing/test_stream_service.py
+++ b/sdks/python/apache_beam/testing/test_stream_service.py
@@ -34,6 +34,8 @@
"""
def __init__(self, reader, endpoint=None, exception_handler=None):
self._server = grpc.server(ThreadPoolExecutor(max_workers=10))
+ self._server_started = False
+ self._server_stopped = False
if endpoint:
self.endpoint = endpoint
@@ -50,9 +52,18 @@
self._exception_handler = lambda _: False
def start(self):
+ # A server can only be started if never started and never stopped before.
+ if self._server_started or self._server_stopped:
+ return
+ self._server_started = True
self._server.start()
def stop(self):
+ # A server can only be stopped if already started and never stopped before.
+ if not self._server_started or self._server_stopped:
+ return
+ self._server_started = False
+ self._server_stopped = True
self._server.stop(0)
# This was introduced in grpcio 1.24 and might be gone in the future. Keep
# this check in case the runtime is on a older, current or future grpcio.
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 01b16a1..7a5b403 100644
--- a/sdks/python/apache_beam/testing/test_stream_service_test.py
+++ b/sdks/python/apache_beam/testing/test_stream_service_test.py
@@ -19,6 +19,7 @@
from __future__ import absolute_import
+import sys
import unittest
import grpc
@@ -30,6 +31,13 @@
from apache_beam.portability.api.beam_runner_api_pb2 import TestStreamPayload
from apache_beam.testing.test_stream_service import TestStreamServiceController
+# TODO(BEAM-8288): clean up the work-around of nose tests using Python2 without
+# unittest.mock module.
+try:
+ from unittest.mock import patch
+except ImportError:
+ from mock import patch # type: ignore[misc]
+
# Nose automatically detects tests if they match a regex. Here, it mistakens
# these protos as tests. For more info see the Nose docs at:
# https://nose.readthedocs.io/en/latest/writing_tests.html
@@ -116,5 +124,78 @@
self.assertEqual(events_b, expected_events)
+@unittest.skipIf(
+ sys.version_info < (3, 6), 'The tests require at least Python 3.6 to work.')
+class TestStreamServiceStartStopTest(unittest.TestCase):
+
+ # Weak internal use needs to be explicitly imported.
+ from grpc import _server
+
+ def setUp(self):
+ self.controller = TestStreamServiceController(
+ EventsReader(expected_key=[('full', EXPECTED_KEY)]))
+ self.assertFalse(self.controller._server_started)
+ self.assertFalse(self.controller._server_stopped)
+
+ def tearDown(self):
+ self.controller.stop()
+
+ def test_start_when_never_started(self):
+ with patch.object(self._server._Server,
+ 'start',
+ wraps=self.controller._server.start) as mock_start:
+ self.controller.start()
+ mock_start.assert_called_once()
+ self.assertTrue(self.controller._server_started)
+ self.assertFalse(self.controller._server_stopped)
+
+ def test_start_noop_when_already_started(self):
+ with patch.object(self._server._Server,
+ 'start',
+ wraps=self.controller._server.start) as mock_start:
+ self.controller.start()
+ mock_start.assert_called_once()
+ self.controller.start()
+ mock_start.assert_called_once()
+
+ def test_start_noop_when_already_stopped(self):
+ with patch.object(self._server._Server,
+ 'start',
+ wraps=self.controller._server.start) as mock_start:
+ self.controller.start()
+ self.controller.stop()
+ mock_start.assert_called_once()
+ self.controller.start()
+ mock_start.assert_called_once()
+
+ def test_stop_noop_when_not_started(self):
+ with patch.object(self._server._Server,
+ 'stop',
+ wraps=self.controller._server.stop) as mock_stop:
+ self.controller.stop()
+ mock_stop.assert_not_called()
+
+ def test_stop_when_already_started(self):
+ with patch.object(self._server._Server,
+ 'stop',
+ wraps=self.controller._server.stop) as mock_stop:
+ self.controller.start()
+ mock_stop.assert_not_called()
+ self.controller.stop()
+ mock_stop.assert_called_once()
+ self.assertFalse(self.controller._server_started)
+ self.assertTrue(self.controller._server_stopped)
+
+ def test_stop_noop_when_already_stopped(self):
+ with patch.object(self._server._Server,
+ 'stop',
+ wraps=self.controller._server.stop) as mock_stop:
+ self.controller.start()
+ self.controller.stop()
+ mock_stop.assert_called_once()
+ self.controller.stop()
+ mock_stop.assert_called_once()
+
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py
index 2ab121e..4a0e80f 100644
--- a/sdks/python/apache_beam/transforms/ptransform.py
+++ b/sdks/python/apache_beam/transforms/ptransform.py
@@ -68,13 +68,16 @@
from apache_beam.internal import pickler
from apache_beam.internal import util
from apache_beam.portability import python_urns
+from apache_beam.pvalue import DoOutputsTuple
from apache_beam.transforms.display import DisplayDataItem
from apache_beam.transforms.display import HasDisplayData
from apache_beam.typehints import native_type_compatibility
from apache_beam.typehints import typehints
+from apache_beam.typehints.decorators import IOTypeHints
from apache_beam.typehints.decorators import TypeCheckError
from apache_beam.typehints.decorators import WithTypeHints
from apache_beam.typehints.decorators import get_signature
+from apache_beam.typehints.decorators import get_type_hints
from apache_beam.typehints.decorators import getcallargs_forhints
from apache_beam.typehints.trivial_inference import instance_to_type
from apache_beam.typehints.typehints import validate_composite_type_param
@@ -350,6 +353,14 @@
# type: () -> str
return self.__class__.__name__
+ def default_type_hints(self):
+ fn_type_hints = IOTypeHints.from_callable(self.expand)
+ if fn_type_hints is not None:
+ fn_type_hints = fn_type_hints.strip_pcoll()
+
+ # Prefer class decorator type hints for backwards compatibility.
+ return get_type_hints(self.__class__).with_defaults(fn_type_hints)
+
def with_input_types(self, input_type_hint):
"""Annotates the input type of a :class:`PTransform` with a type-hint.
@@ -419,6 +430,8 @@
root_hint = (
arg_hints[0] if len(arg_hints) == 1 else arg_hints or kwarg_hints)
for context, pvalue_, hint in _ZipPValues().visit(pvalueish, root_hint):
+ if isinstance(pvalue_, DoOutputsTuple):
+ continue
if pvalue_.element_type is None:
# TODO(robertwb): It's a bug that we ever get here. (typecheck)
continue
diff --git a/sdks/python/apache_beam/transforms/validate_runner_xlang_test.py b/sdks/python/apache_beam/transforms/validate_runner_xlang_test.py
index 4d6b56b..d6ff006 100644
--- a/sdks/python/apache_beam/transforms/validate_runner_xlang_test.py
+++ b/sdks/python/apache_beam/transforms/validate_runner_xlang_test.py
@@ -14,6 +14,42 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+
+"""
+###########################################################
+Runner Validation Test Suite for Cross-language Transforms
+###########################################################
+ As per Beams's Portability Framework design, Cross-language transforms
+ should work out of the box. In spite of this, there always exists a
+ possibility of rough edges existing. It could be caused due to unpolished
+ implementation of any part of the execution code path, for example:
+ - Transform expansion [SDK]
+ - Pipeline construction [SDK]
+ - Cross-language artifact staging [Runner]
+ - Language specific serialization/deserialization of PCollection (and
+ other data types) [Runner/SDK]
+
+ In an effort to improve developer visibility into potential problems,
+ this test suite validates correct execution of 5 Core Beam transforms when
+ used as cross-language transforms within the Python SDK from any foreign SDK:
+ - ParDo
+ (https://beam.apache.org/documentation/programming-guide/#pardo)
+ - GroupByKey
+ (https://beam.apache.org/documentation/programming-guide/#groupbykey)
+ - CoGroupByKey
+ (https://beam.apache.org/documentation/programming-guide/#cogroupbykey)
+ - Combine
+ (https://beam.apache.org/documentation/programming-guide/#combine)
+ - Flatten
+ (https://beam.apache.org/documentation/programming-guide/#flatten)
+ - Partition
+ (https://beam.apache.org/documentation/programming-guide/#partition)
+
+ See Runner Validation Test Plan for Cross-language transforms at
+https://docs.google.com/document/d/1xQp0ElIV84b8OCVz8CD2hvbiWdR8w4BvWxPTZJZA6NA
+ for further details.
+"""
+
from __future__ import absolute_import
import logging
@@ -46,6 +82,15 @@
'localhost:%s' % os.environ.get('EXPANSION_PORT'))
def run_prefix(self, pipeline):
+ """
+ Target transform - ParDo
+ (https://beam.apache.org/documentation/programming-guide/#pardo)
+ Test scenario - Mapping elements from a single input collection to a
+ single output collection
+ Boundary conditions checked -
+ - PCollection<?> to external transforms
+ - PCollection<?> from external transforms
+ """
with pipeline as p:
res = (
p
@@ -57,6 +102,15 @@
assert_that(res, equal_to(['0a', '0b']))
def run_multi_input_output_with_sideinput(self, pipeline):
+ """
+ Target transform - ParDo
+ (https://beam.apache.org/documentation/programming-guide/#pardo)
+ Test scenario - Mapping elements from multiple input collections (main
+ and side) to multiple output collections (main and side)
+ Boundary conditions checked -
+ - PCollectionTuple to external transforms
+ - PCollectionTuple from external transforms
+ """
with pipeline as p:
main1 = p | 'Main1' >> beam.Create(
['a', 'bb'], reshuffle=False).with_output_types(unicode)
@@ -70,6 +124,15 @@
assert_that(res['side'], equal_to(['ss']), label='CheckSide')
def run_group_by_key(self, pipeline):
+ """
+ Target transform - GroupByKey
+ (https://beam.apache.org/documentation/programming-guide/#groupbykey)
+ Test scenario - Grouping a collection of KV<K,V> to a collection of
+ KV<K, Iterable<V>> by key
+ Boundary conditions checked -
+ - PCollection<KV<?, ?>> to external transforms
+ - PCollection<KV<?, Iterable<?>>> from external transforms
+ """
with pipeline as p:
res = (
p
@@ -81,6 +144,15 @@
assert_that(res, equal_to(['0:1,2', '1:3']))
def run_cogroup_by_key(self, pipeline):
+ """
+ Target transform - CoGroupByKey
+ (https://beam.apache.org/documentation/programming-guide/#cogroupbykey)
+ Test scenario - Grouping multiple input collections with keys to a
+ collection of KV<K, CoGbkResult> by key
+ Boundary conditions checked -
+ - KeyedPCollectionTuple<?> to external transforms
+ - PCollection<KV<?, Iterable<?>>> from external transforms
+ """
with pipeline as p:
col1 = p | 'create_col1' >> beam.Create(
[(0, "1"), (0, "2"), (1, "3")], reshuffle=False).with_output_types(
@@ -95,6 +167,15 @@
assert_that(res, equal_to(['0:1,2,4', '1:3,5,6']))
def run_combine_globally(self, pipeline):
+ """
+ Target transform - Combine
+ (https://beam.apache.org/documentation/programming-guide/#combine)
+ Test scenario - Combining elements globally with a predefined simple
+ CombineFn
+ Boundary conditions checked -
+ - PCollection<?> to external transforms
+ - PCollection<?> from external transforms
+ """
with pipeline as p:
res = (
p
@@ -104,6 +185,15 @@
assert_that(res, equal_to([6]))
def run_combine_per_key(self, pipeline):
+ """
+ Target transform - Combine
+ (https://beam.apache.org/documentation/programming-guide/#combine)
+ Test scenario - Combining elements per key with a predefined simple
+ merging function
+ Boundary conditions checked -
+ - PCollection<?> to external transforms
+ - PCollection<?> from external transforms
+ """
with pipeline as p:
res = (
p
@@ -114,6 +204,14 @@
assert_that(res, equal_to([('a', 3), ('b', 3)]))
def run_flatten(self, pipeline):
+ """
+ Target transform - Flatten
+ (https://beam.apache.org/documentation/programming-guide/#flatten)
+ Test scenario - Merging multiple collections into a single collection
+ Boundary conditions checked -
+ - PCollectionList<?> to external transforms
+ - PCollection<?> from external transforms
+ """
with pipeline as p:
col1 = p | 'col1' >> beam.Create([1, 2, 3]).with_output_types(int)
col2 = p | 'col2' >> beam.Create([4, 5, 6]).with_output_types(int)
@@ -123,6 +221,15 @@
assert_that(res, equal_to([1, 2, 3, 4, 5, 6]))
def run_partition(self, pipeline):
+ """
+ Target transform - Partition
+ (https://beam.apache.org/documentation/programming-guide/#partition)
+ Test scenario - Splitting a single collection into multiple collections
+ with a predefined simple PartitionFn
+ Boundary conditions checked -
+ - PCollection<?> to external transforms
+ - PCollectionList<?> from external transforms
+ """
with pipeline as p:
res = (
p
diff --git a/sdks/python/apache_beam/typehints/decorators.py b/sdks/python/apache_beam/typehints/decorators.py
index dad0a31..4cd7681 100644
--- a/sdks/python/apache_beam/typehints/decorators.py
+++ b/sdks/python/apache_beam/typehints/decorators.py
@@ -105,6 +105,7 @@
from typing import Optional
from typing import Tuple
from typing import TypeVar
+from typing import Union
from apache_beam.typehints import native_type_compatibility
from apache_beam.typehints import typehints
@@ -378,6 +379,75 @@
self.output_types and len(self.output_types[0]) == 1 and
not self.output_types[1])
+ def strip_pcoll(self):
+ from apache_beam.pipeline import Pipeline
+ from apache_beam.pvalue import PBegin
+ from apache_beam.pvalue import PDone
+
+ return self.strip_pcoll_helper(self.input_types,
+ self._has_input_types,
+ 'input_types',
+ [Pipeline, PBegin],
+ 'This input type hint will be ignored '
+ 'and not used for type-checking purposes. '
+ 'Typically, input type hints for a '
+ 'PTransform are single (or nested) types '
+ 'wrapped by a PCollection, or PBegin.',
+ 'strip_pcoll_input()').\
+ strip_pcoll_helper(self.output_types,
+ self.has_simple_output_type,
+ 'output_types',
+ [PDone, None],
+ 'This output type hint will be ignored '
+ 'and not used for type-checking purposes. '
+ 'Typically, output type hints for a '
+ 'PTransform are single (or nested) types '
+ 'wrapped by a PCollection, PDone, or None.',
+ 'strip_pcoll_output()')
+
+ def strip_pcoll_helper(
+ self,
+ my_type, # type: any
+ has_my_type, # type: Callable[[], bool]
+ my_key, # type: str
+ special_containers, # type: List[Union[PBegin, PDone, PCollection]]
+ error_str, # type: str
+ source_str # type: str
+ ):
+ # type: (...) -> IOTypeHints
+ from apache_beam.pvalue import PCollection
+
+ if not has_my_type() or not my_type or len(my_type[0]) != 1:
+ return self
+
+ my_type = my_type[0][0]
+
+ if isinstance(my_type, typehints.AnyTypeConstraint):
+ return self
+
+ special_containers += [PCollection]
+ kwarg_dict = {}
+
+ if (my_type not in special_containers and
+ getattr(my_type, '__origin__', None) != PCollection):
+ logging.warning(error_str + ' Got: %s instead.' % my_type)
+ kwarg_dict[my_key] = None
+ return self._replace(
+ origin=self._make_origin([self], tb=False, msg=[source_str]),
+ **kwarg_dict)
+
+ if (getattr(my_type, '__args__', -1) in [-1, None] or
+ len(my_type.__args__) == 0):
+ # e.g. PCollection (or PBegin/PDone)
+ kwarg_dict[my_key] = ((typehints.Any, ), {})
+ else:
+ # e.g. PCollection[type]
+ kwarg_dict[my_key] = ((convert_to_beam_type(my_type.__args__[0]), ), {})
+
+ return self._replace(
+ origin=self._make_origin([self], tb=False, msg=[source_str]),
+ **kwarg_dict)
+
def strip_iterable(self):
# type: () -> IOTypeHints
diff --git a/sdks/python/apache_beam/typehints/schemas.py b/sdks/python/apache_beam/typehints/schemas.py
index 0cc513f..cb4cf01 100644
--- a/sdks/python/apache_beam/typehints/schemas.py
+++ b/sdks/python/apache_beam/typehints/schemas.py
@@ -178,6 +178,11 @@
return schema_pb2.FieldType(
array_type=schema_pb2.ArrayType(element_type=element_type))
+ elif _safe_issubclass(type_, Mapping):
+ key_type, value_type = map(typing_to_runner_api, _get_args(type_))
+ return schema_pb2.FieldType(
+ map_type=schema_pb2.MapType(key_type=key_type, value_type=value_type))
+
raise ValueError("Unsupported type: %s" % type_)
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 2016871..e12930d 100644
--- a/sdks/python/apache_beam/typehints/typed_pipeline_test_py3.py
+++ b/sdks/python/apache_beam/typehints/typed_pipeline_test_py3.py
@@ -22,6 +22,7 @@
from __future__ import absolute_import
+import typing
import unittest
import apache_beam as beam
@@ -257,6 +258,135 @@
result = [1, 2, 3] | beam.FlatMap(fn) | beam.Map(fn2)
self.assertCountEqual([4, 6], result)
+ def test_typed_ptransform_with_no_error(self):
+ class StrToInt(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[int]:
+ return pcoll | beam.Map(lambda x: int(x))
+
+ class IntToStr(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]:
+ return pcoll | beam.Map(lambda x: str(x))
+
+ _ = ['1', '2', '3'] | StrToInt() | IntToStr()
+
+ def test_typed_ptransform_with_bad_typehints(self):
+ class StrToInt(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[int]:
+ return pcoll | beam.Map(lambda x: int(x))
+
+ class IntToStr(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[str]:
+ return pcoll | beam.Map(lambda x: str(x))
+
+ with self.assertRaisesRegex(typehints.TypeCheckError,
+ "Input type hint violation at IntToStr: "
+ "expected <class 'str'>, got <class 'int'>"):
+ _ = ['1', '2', '3'] | StrToInt() | IntToStr()
+
+ def test_typed_ptransform_with_bad_input(self):
+ class StrToInt(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[str]) -> beam.pvalue.PCollection[int]:
+ return pcoll | beam.Map(lambda x: int(x))
+
+ class IntToStr(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]:
+ return pcoll | beam.Map(lambda x: str(x))
+
+ with self.assertRaisesRegex(typehints.TypeCheckError,
+ "Input type hint violation at StrToInt: "
+ "expected <class 'str'>, got <class 'int'>"):
+ # Feed integers to a PTransform that expects strings
+ _ = [1, 2, 3] | StrToInt() | IntToStr()
+
+ def test_typed_ptransform_with_partial_typehints(self):
+ class StrToInt(beam.PTransform):
+ def expand(self, pcoll) -> beam.pvalue.PCollection[int]:
+ return pcoll | beam.Map(lambda x: int(x))
+
+ class IntToStr(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]:
+ return pcoll | beam.Map(lambda x: str(x))
+
+ # Feed integers to a PTransform that should expect strings
+ # but has no typehints so it expects any
+ _ = [1, 2, 3] | StrToInt() | IntToStr()
+
+ def test_typed_ptransform_with_bare_wrappers(self):
+ class StrToInt(beam.PTransform):
+ def expand(
+ self, pcoll: beam.pvalue.PCollection) -> beam.pvalue.PCollection:
+ return pcoll | beam.Map(lambda x: int(x))
+
+ class IntToStr(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]:
+ return pcoll | beam.Map(lambda x: str(x))
+
+ _ = [1, 2, 3] | StrToInt() | IntToStr()
+
+ def test_typed_ptransform_with_no_typehints(self):
+ class StrToInt(beam.PTransform):
+ def expand(self, pcoll):
+ return pcoll | beam.Map(lambda x: int(x))
+
+ class IntToStr(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[int]) -> beam.pvalue.PCollection[str]:
+ return pcoll | beam.Map(lambda x: str(x))
+
+ # Feed integers to a PTransform that should expect strings
+ # but has no typehints so it expects any
+ _ = [1, 2, 3] | StrToInt() | IntToStr()
+
+ def test_typed_ptransform_with_generic_annotations(self):
+ T = typing.TypeVar('T')
+
+ class IntToInt(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[T]) -> beam.pvalue.PCollection[T]:
+ return pcoll | beam.Map(lambda x: x)
+
+ class IntToStr(beam.PTransform):
+ def expand(
+ self,
+ pcoll: beam.pvalue.PCollection[T]) -> beam.pvalue.PCollection[str]:
+ return pcoll | beam.Map(lambda x: str(x))
+
+ _ = [1, 2, 3] | IntToInt() | IntToStr()
+
+ def test_typed_ptransform_with_do_outputs_tuple_compiles(self):
+ class MyDoFn(beam.DoFn):
+ def process(self, element: int, *args, **kwargs):
+ if element % 2:
+ yield beam.pvalue.TaggedOutput('odd', 1)
+ else:
+ yield beam.pvalue.TaggedOutput('even', 1)
+
+ class MyPTransform(beam.PTransform):
+ def expand(self, pcoll: beam.pvalue.PCollection[int]):
+ return pcoll | beam.ParDo(MyDoFn()).with_outputs('odd', 'even')
+
+ # This test fails if you remove the following line from ptransform.py
+ # if isinstance(pvalue_, DoOutputsTuple): continue
+ _ = [1, 2, 3] | MyPTransform()
+
class AnnotationsTest(unittest.TestCase):
def test_pardo_dofn(self):
diff --git a/sdks/python/apache_beam/typehints/typehints_test_py3.py b/sdks/python/apache_beam/typehints/typehints_test_py3.py
index a7c23f0..5a36330 100644
--- a/sdks/python/apache_beam/typehints/typehints_test_py3.py
+++ b/sdks/python/apache_beam/typehints/typehints_test_py3.py
@@ -23,11 +23,19 @@
from __future__ import absolute_import
from __future__ import print_function
+import typing
import unittest
+import apache_beam.typehints.typehints as typehints
+from apache_beam import Map
+from apache_beam import PTransform
+from apache_beam.pvalue import PBegin
+from apache_beam.pvalue import PCollection
+from apache_beam.pvalue import PDone
from apache_beam.transforms.core import DoFn
from apache_beam.typehints import KV
from apache_beam.typehints import Iterable
+from apache_beam.typehints.typehints import Any
class TestParDoAnnotations(unittest.TestCase):
@@ -46,11 +54,221 @@
def process(self, element: int) -> Iterable[str]:
pass
- print(MyDoFn().get_type_hints())
th = MyDoFn().get_type_hints()
self.assertEqual(th.input_types, ((int, ), {}))
self.assertEqual(th.output_types, ((str, ), {}))
+class TestPTransformAnnotations(unittest.TestCase):
+ def test_pep484_annotations(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: PCollection[int]) -> PCollection[str]:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((str, ), {}))
+
+ def test_annotations_without_input_pcollection_wrapper(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: int) -> PCollection[str]:
+ return pcoll | Map(lambda num: str(num))
+
+ error_str = (
+ r'This input type hint will be ignored and not used for '
+ r'type-checking purposes. Typically, input type hints for a '
+ r'PTransform are single (or nested) types wrapped by a '
+ r'PCollection, or PBegin. Got: {} instead.'.format(int))
+
+ with self.assertLogs(level='WARN') as log:
+ MyPTransform().get_type_hints()
+ self.assertIn(error_str, log.output[0])
+
+ def test_annotations_without_output_pcollection_wrapper(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: PCollection[int]) -> str:
+ return pcoll | Map(lambda num: str(num))
+
+ error_str = (
+ r'This output type hint will be ignored and not used for '
+ r'type-checking purposes. Typically, output type hints for a '
+ r'PTransform are single (or nested) types wrapped by a '
+ r'PCollection, PDone, or None. Got: {} instead.'.format(str))
+
+ with self.assertLogs(level='WARN') as log:
+ th = MyPTransform().get_type_hints()
+ self.assertIn(error_str, log.output[0])
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, None)
+
+ def test_annotations_without_input_internal_type(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: PCollection) -> PCollection[str]:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, ((str, ), {}))
+
+ def test_annotations_without_output_internal_type(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: PCollection[int]) -> PCollection:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((Any, ), {}))
+
+ def test_annotations_without_any_internal_type(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: PCollection) -> PCollection:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, ((Any, ), {}))
+
+ def test_annotations_without_input_typehint(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll) -> PCollection[str]:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, ((str, ), {}))
+
+ def test_annotations_without_output_typehint(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: PCollection[int]):
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((Any, ), {}))
+
+ def test_annotations_without_any_typehints(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll):
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, None)
+ self.assertEqual(th.output_types, None)
+
+ def test_annotations_with_pbegin(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: PBegin):
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, ((Any, ), {}))
+
+ def test_annotations_with_pdone(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll) -> PDone:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, ((Any, ), {}))
+
+ def test_annotations_with_none_input(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: None) -> PCollection[str]:
+ return pcoll | Map(lambda num: str(num))
+
+ error_str = (
+ r'This input type hint will be ignored and not used for '
+ r'type-checking purposes. Typically, input type hints for a '
+ r'PTransform are single (or nested) types wrapped by a '
+ r'PCollection, or PBegin. Got: {} instead.'.format(None))
+
+ with self.assertLogs(level='WARN') as log:
+ th = MyPTransform().get_type_hints()
+ self.assertIn(error_str, log.output[0])
+ self.assertEqual(th.input_types, None)
+ self.assertEqual(th.output_types, ((str, ), {}))
+
+ def test_annotations_with_none_output(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll) -> None:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, ((Any, ), {}))
+
+ def test_annotations_with_arbitrary_output(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll) -> str:
+ return pcoll | Map(lambda num: str(num))
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, None)
+
+ def test_annotations_with_arbitrary_input_and_output(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: int) -> str:
+ return pcoll | Map(lambda num: str(num))
+
+ input_error_str = (
+ r'This input type hint will be ignored and not used for '
+ r'type-checking purposes. Typically, input type hints for a '
+ r'PTransform are single (or nested) types wrapped by a '
+ r'PCollection, or PBegin. Got: {} instead.'.format(int))
+
+ output_error_str = (
+ r'This output type hint will be ignored and not used for '
+ r'type-checking purposes. Typically, output type hints for a '
+ r'PTransform are single (or nested) types wrapped by a '
+ r'PCollection, PDone, or None. Got: {} instead.'.format(str))
+
+ with self.assertLogs(level='WARN') as log:
+ th = MyPTransform().get_type_hints()
+ self.assertIn(input_error_str, log.output[0])
+ self.assertIn(output_error_str, log.output[1])
+ self.assertEqual(th.input_types, None)
+ self.assertEqual(th.output_types, None)
+
+ def test_typing_module_annotations_are_converted_to_beam_annotations(self):
+ class MyPTransform(PTransform):
+ def expand(
+ self, pcoll: PCollection[typing.Dict[str, str]]
+ ) -> PCollection[typing.Dict[str, str]]:
+ return pcoll
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((typehints.Dict[str, str], ), {}))
+ self.assertEqual(th.input_types, ((typehints.Dict[str, str], ), {}))
+
+ def test_nested_typing_annotations_are_converted_to_beam_annotations(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll:
+ PCollection[typing.Union[int, typing.Any, typing.Dict[str, float]]]) \
+ -> PCollection[typing.Union[int, typing.Any, typing.Dict[str, float]]]:
+ return pcoll
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(
+ th.input_types,
+ ((typehints.Union[int, typehints.Any, typehints.Dict[str,
+ float]], ), {}))
+ self.assertEqual(
+ th.input_types,
+ ((typehints.Union[int, typehints.Any, typehints.Dict[str,
+ float]], ), {}))
+
+ def test_mixed_annotations_are_converted_to_beam_annotations(self):
+ class MyPTransform(PTransform):
+ def expand(self, pcoll: typing.Any) -> typehints.Any:
+ return pcoll
+
+ th = MyPTransform().get_type_hints()
+ self.assertEqual(th.input_types, ((typehints.Any, ), {}))
+ self.assertEqual(th.input_types, ((typehints.Any, ), {}))
+
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/utils/shared.py b/sdks/python/apache_beam/utils/shared.py
index 31d4ed6..9f0ac94 100644
--- a/sdks/python/apache_beam/utils/shared.py
+++ b/sdks/python/apache_beam/utils/shared.py
@@ -102,10 +102,12 @@
def __init__(self):
self._lock = threading.Lock()
self._ref = None
+ self._tag = None
def acquire(
self,
- constructor_fn # type: Callable[[], Any]
+ constructor_fn, # type: Callable[[], Any]
+ tag=None # type: Any
):
# type: (...) -> Any
@@ -116,6 +118,9 @@
present in the cache. This function should take no arguments. It should
return an initialised object, or None if the object could not be
initialised / constructed.
+ tag: an optional indentifier to store with the cached object. If
+ subsequent calls to acquire use different tags, the object will be
+ reloaded rather than returned from cache.
Returns:
An initialised object, either from a previous initialisation, or
@@ -124,7 +129,8 @@
with self._lock:
# self._ref is None if this is a new control block.
# self._ref() is None if the weak reference was GCed.
- if self._ref is None or self._ref() is None:
+ # self._tag != tag if user specifies a new identifier
+ if self._ref is None or self._ref() is None or self._tag != tag:
result = constructor_fn()
if result is None:
return None
@@ -205,6 +211,7 @@
self,
key, # type: Text
constructor_fn, # type: Callable[[], Any]
+ tag=None # type: Any
):
# type: (...) -> Any
@@ -216,6 +223,9 @@
present in the cache. This function should take no arguments. It should
return an initialised object, or None if the object could not be
initialised / constructed.
+ tag: an optional indentifier to store with the cached object. If
+ subsequent calls to acquire use different tags, the object will be
+ reloaded rather than returned from cache.
Returns:
A reference to the initialised object, either from the cache, or
@@ -227,7 +237,7 @@
control_block = _SharedControlBlock()
self._cache_map[key] = control_block
- result = control_block.acquire(constructor_fn)
+ result = control_block.acquire(constructor_fn, tag)
# Because we release the lock in between, if we acquire multiple Shareds
# in a short time, there's no guarantee as to which one will be kept alive.
@@ -266,7 +276,8 @@
def acquire(
self,
- constructor_fn # type: Callable[[], Any]
+ constructor_fn, # type: Callable[[], Any]
+ tag=None # type: Any
):
# type: (...) -> Any
@@ -277,9 +288,12 @@
present in the cache. This function should take no arguments. It should
return an initialised object, or None if the object could not be
initialised / constructed.
+ tag: an optional indentifier to store with the cached object. If
+ subsequent calls to acquire use different tags, the object will be
+ reloaded rather than returned from cache.
Returns:
A reference to an initialised object, either from the cache, or
newly-constructed.
"""
- return _shared_map.acquire(self._key, constructor_fn)
+ return _shared_map.acquire(self._key, constructor_fn, tag)
diff --git a/sdks/python/apache_beam/utils/shared_test.py b/sdks/python/apache_beam/utils/shared_test.py
index 35cdb26..28ba7c8 100644
--- a/sdks/python/apache_beam/utils/shared_test.py
+++ b/sdks/python/apache_beam/utils/shared_test.py
@@ -236,6 +236,28 @@
self.assertEqual('sequence3', f3.get_name())
self.assertEqual('sequence4', s3.get_name())
+ def testTagCacheEviction(self):
+ shared1 = shared.Shared()
+ shared2 = shared.Shared()
+
+ def acquire_fn_1():
+ return NamedObject('obj_1')
+
+ def acquire_fn_2():
+ return NamedObject('obj_2')
+
+ # with no tag, shared handle does not know when to evict objects
+ p1 = shared1.acquire(acquire_fn_1)
+ assert p1.get_name() == 'obj_1'
+ p2 = shared1.acquire(acquire_fn_2)
+ assert p2.get_name() == 'obj_1'
+
+ # cache eviction can be forced by specifying different tags
+ p1 = shared2.acquire(acquire_fn_1, tag='1')
+ assert p1.get_name() == 'obj_1'
+ p2 = shared2.acquire(acquire_fn_2, tag='2')
+ assert p2.get_name() == 'obj_2'
+
if __name__ == '__main__':
unittest.main()
diff --git a/sdks/python/apache_beam/utils/subprocess_server.py b/sdks/python/apache_beam/utils/subprocess_server.py
index 0d40181..018a837 100644
--- a/sdks/python/apache_beam/utils/subprocess_server.py
+++ b/sdks/python/apache_beam/utils/subprocess_server.py
@@ -291,11 +291,20 @@
if port:
return port
else:
- s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ try:
+ s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ except OSError as e:
+ # [Errno 97] Address family not supported by protocol
+ # Likely indicates we are in an IPv6-only environment (BEAM-10618). Try
+ # again with AF_INET6.
+ if e.errno == 97:
+ s = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
+ else:
+ raise e
+
sockets.append(s)
s.bind(('localhost', 0))
- _, free_port = s.getsockname()
- return free_port
+ return s.getsockname()[1]
ports = list(map(find_free_port, ports))
# Close sockets only now to avoid the same port to be chosen twice
diff --git a/sdks/python/container/base_image_requirements.txt b/sdks/python/container/base_image_requirements.txt
index fd9163b..3484d38 100644
--- a/sdks/python/container/base_image_requirements.txt
+++ b/sdks/python/container/base_image_requirements.txt
@@ -45,7 +45,7 @@
typing-extensions==3.7.4.1
# GCP extra features
-google-api-core==1.20.0
+google-api-core==1.21.0
google-apitools==0.5.28
google-cloud-pubsub==1.0.2
google-cloud-bigquery==1.26.1
diff --git a/sdks/python/test-suites/portable/common.gradle b/sdks/python/test-suites/portable/common.gradle
index 2e60afa..48312a6 100644
--- a/sdks/python/test-suites/portable/common.gradle
+++ b/sdks/python/test-suites/portable/common.gradle
@@ -101,38 +101,6 @@
}
}
-task crossLanguagePythonJavaKafkaIOFlink {
- dependsOn 'setupVirtualenv'
- dependsOn ':runners:flink:1.10:job-server:shadowJar'
- dependsOn ":sdks:python:container:py${pythonVersionSuffix}:docker"
- dependsOn ':sdks:java:container:docker'
- dependsOn ':sdks:java:io:expansion-service:shadowJar'
- dependsOn ':sdks:java:testing:kafka-service:buildTestKafkaServiceJar'
-
- doLast {
- def kafkaJar = project(":sdks:java:testing:kafka-service:").buildTestKafkaServiceJar.archivePath
- def options = [
- "--runner=FlinkRunner",
- "--parallelism=2",
- "--environment_type=DOCKER",
- "--environment_cache_millis=10000",
- "--experiment=beam_fn_api",
- ]
- exec {
- environment "LOCAL_KAFKA_JAR", kafkaJar
- executable 'sh'
- args '-c', """
- . ${envdir}/bin/activate \\
- && cd ${pythonRootDir} \\
- && pip install -e .[test] \\
- && python setup.py nosetests \\
- --tests apache_beam.io.external.xlang_kafkaio_it_test:CrossLanguageKafkaIOTest \\
- --test-pipeline-options='${options.join(' ')}'
- """
- }
- }
-}
-
task createProcessWorker {
dependsOn ':sdks:python:container:build'
dependsOn 'setupVirtualenv'
@@ -223,12 +191,14 @@
':runners:flink:1.10:job-server:shadowJar',
':sdks:java:container:docker',
':sdks:java:io:expansion-service:shadowJar',
+ ':sdks:java:testing:kafka-service:buildTestKafkaServiceJar'
]
doLast {
def tests = [
"apache_beam.io.gcp.bigquery_read_it_test",
"apache_beam.io.external.xlang_jdbcio_it_test",
+ "apache_beam.io.external.xlang_kafkaio_it_test",
]
def testOpts = ["--tests=${tests.join(',')}"]
def cmdArgs = mapToArgString([
@@ -236,7 +206,9 @@
"suite": "postCommitIT-flink-py${pythonVersionSuffix}",
"pipeline_opts": "--runner=FlinkRunner --project=apache-beam-testing --environment_type=LOOPBACK --temp_location=gs://temp-storage-for-end-to-end-tests/temp-it",
])
+ def kafkaJar = project(":sdks:java:testing:kafka-service:").buildTestKafkaServiceJar.archivePath
exec {
+ environment "LOCAL_KAFKA_JAR", kafkaJar
executable 'sh'
args '-c', ". ${envdir}/bin/activate && ${pythonRootDir}/scripts/run_integration_test.sh $cmdArgs"
}
diff --git a/settings.gradle b/settings.gradle
index 0220857..638c216 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -120,6 +120,7 @@
include ":sdks:java:io:redis"
include ":sdks:java:io:solr"
include ":sdks:java:io:snowflake"
+include ":sdks:java:io:snowflake:expansion-service"
include ":sdks:java:io:splunk"
include ":sdks:java:io:thrift"
include ":sdks:java:io:tika"
diff --git a/website/www/site/content/en/contribute/release-guide.md b/website/www/site/content/en/contribute/release-guide.md
index 87eb732..8599b41 100644
--- a/website/www/site/content/en/contribute/release-guide.md
+++ b/website/www/site/content/en/contribute/release-guide.md
@@ -37,6 +37,9 @@
1. Decide to release
1. Prepare for the release
+1. Investigate performance regressions
+1. Create a release branch
+1. Verify release branch
1. Build a release candidate
1. Vote on the release candidate
1. During vote process, run validation tests
@@ -45,7 +48,7 @@
1. Promote the release
-### Decide to release
+## 1. Decide to release
Deciding to release and selecting a Release Manager is the first step of the release process. This is a consensus-based decision of the entire community.
@@ -60,7 +63,7 @@
**********
-## 1. Prepare for the release
+## 2. Prepare for the release
Before your first release, you should perform one-time configuration steps. This will set up your security keys for signing the release and access to various release repositories.
@@ -244,7 +247,23 @@
**********
-## 2. Create a release branch in apache/beam repository
+## 3. Investigate performance regressions
+
+Check the Beam load tests for possible performance regressions. Measurements are available on [metrics.beam.apache.org](http://metrics.beam.apache.org).
+
+All Runners which publish data should be checked for the following, in both *batch* and *streaming* mode:
+
+- [ParDo](http://metrics.beam.apache.org/d/MOi-kf3Zk/pardo-load-tests) and [GBK](http://metrics.beam.apache.org/d/UYZ-oJ3Zk/gbk-load-test): Runtime, latency, checkpoint duration
+- [Nexmark](http://metrics.beam.apache.org/d/ahuaA_zGz/nexmark): Query runtime for all queries
+- [IO](http://metrics.beam.apache.org/d/bnlHKP3Wz/java-io-it-tests-dataflow): Runtime
+
+If regressions are found, the release branch can still be created, but the regressions should be investigated and fixed as part of the release process.
+The role of the release manager is to file JIRA issues for each regression with the 'Fix Version' set to the to-be-released version. The release manager
+oversees these just like any other JIRA issue marked with the 'Fix Version' of the release.
+
+The mailing list should be informed to allow fixing the regressions in the course of the release.
+
+## 4. Create a release branch in apache/beam repository
Attention: Only committer has permission to create release branch in apache/beam.
@@ -367,7 +386,7 @@
**********
-## 3. Verify release branch
+## 5. Verify release branch
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.
@@ -477,11 +496,16 @@
* Description: Description of failure
+#### Inform the mailing list
+
+The dev@beam.apache.org mailing list should be informed about the release branch being cut. Alongside with this note,
+a list of pending issues and to-be-trigated issues should be included. Afterwards, this list can be refined and updated
+by the release manager and the Beam community.
**********
-## 4. Triage release-blocking issues in JIRA
+## 6. Triage release-blocking issues in JIRA
There could be outstanding release-blocking issues, which should be triaged before proceeding to build a release candidate. We track them by assigning a specific `Fix version` field even before the issue resolved.
@@ -528,7 +552,7 @@
**********
-## 5. Build a release candidate
+## 7. Build a release candidate
### Checklist before proceeding
@@ -599,7 +623,7 @@
**********
-## 6. Prepare documents
+## 8. Prepare documents
### Update and Verify Javadoc
@@ -748,7 +772,7 @@
**********
-## 7. Vote and validate release candidate
+## 9. Vote and validate release candidate
Once you have built and individually reviewed the release candidate, please share it for the community-wide review. Please review foundation-wide [voting guidelines](http://www.apache.org/foundation/voting.html) for more information.
@@ -1074,7 +1098,7 @@
**********
-## 8. Finalize the release
+## 10. Finalize the release
Once the release candidate has been reviewed and approved by the community, the release should be finalized. This involves the final deployment of the release candidate to the release repositories, merging of the website changes, etc.
@@ -1163,7 +1187,7 @@
**********
-## 9. Promote the release
+## 11. Promote the release
Once the release has been finalized, the last step of the process is to promote the release within the project and beyond.
diff --git a/website/www/site/content/en/documentation/io/built-in/snowflake.md b/website/www/site/content/en/documentation/io/built-in/snowflake.md
new file mode 100644
index 0000000..078a373
--- /dev/null
+++ b/website/www/site/content/en/documentation/io/built-in/snowflake.md
@@ -0,0 +1,364 @@
+---
+title: "Apache Snowflake I/O connector"
+---
+<!--
+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.
+-->
+
+[Built-in I/O Transforms](/documentation/io/built-in/)
+
+# Snowflake I/O
+Pipeline options and general information about using and running Snowflake IO.
+
+## Authentication
+All authentication methods available for the Snowflake JDBC Driver are possible to use with the IO transforms:
+
+- Username and password
+- Key pair
+- OAuth token
+
+Passing credentials is done via Pipeline options.
+
+Passing credentials is done via Pipeline options used to instantiate `SnowflakeIO.DataSourceConfiguration`:
+{{< highlight java >}}
+SnowflakePipelineOptions options = PipelineOptionsFactory
+ .fromArgs(args)
+ .withValidation()
+ .as(SnowflakePipelineOptions.class);
+SnowflakeCredentials credentials = SnowflakeCredentialsFactory.of(options);
+
+SnowflakeIO.DataSourceConfiguration.create(credentials)
+ .(other DataSourceConfiguration options)
+{{< /highlight >}}
+### Username and password
+To use username/password authentication in SnowflakeIO, invoke your pipeline with the following Pipeline options:
+{{< highlight >}}
+--username=<USERNAME> --password=<PASSWORD>
+{{< /highlight >}}
+### Key pair
+To use this authentication method, you must first generate a key pair and associate the public key with the Snowflake user that will connect using the IO transform. For instructions, see the [Snowflake documentation](https://docs.snowflake.com/en/user-guide/jdbc-configure.html).
+
+To use key pair authentication with SnowflakeIO, invoke your pipeline with following Pipeline options:
+{{< highlight >}}
+--username=<USERNAME> --privateKeyPath=<PATH_TO_P8_FILE> --privateKeyPassphrase=<PASSWORD_FOR_KEY>
+{{< /highlight >}}
+
+### OAuth token
+SnowflakeIO also supports OAuth token.
+
+**IMPORTANT**: SnowflakeIO requires a valid OAuth access token. It will neither be able to refresh the token nor obtain it using a web-based flow. For information on configuring an OAuth integration and obtaining the token, see the [Snowflake documentation](https://docs.snowflake.com/en/user-guide/oauth-intro.html).
+
+Once you have the token, invoke your pipeline with following Pipeline Options:
+{{< highlight >}}
+--oauthToken=<TOKEN>
+{{< /highlight >}}
+## DataSource Configuration
+DataSource configuration is required in both read and write object for configuring Snowflake connection properties for IO purposes.
+### General usage
+Create the DataSource configuration:
+{{< highlight java >}}
+ SnowflakeIO.DataSourceConfiguration
+ .create(SnowflakeCredentialsFactory.of(options))
+ .withUrl(options.getUrl())
+ .withServerName(options.getServerName())
+ .withDatabase(options.getDatabase())
+ .withWarehouse(options.getWarehouse())
+ .withSchema(options.getSchema());
+{{< /highlight >}}
+Where parameters can be:
+
+- ` .withUrl(...)`
+ - JDBC-like URL for your Snowflake account, including account name and region, without any parameters.
+ - Example: `.withUrl("jdbc:snowflake://account.snowflakecomputing.com")`
+- `.withServerName(...)`
+ - Server Name - full server name with account, zone and domain.
+ - Example: `.withServerName("account.snowflakecomputing.com")`
+- `.withDatabase(...)`
+ - Name of the Snowflake database to use.
+ - Example: `.withDatabase("MY_DATABASE")`
+- `.withWarehouse(...)`
+ - Name of the Snowflake warehouse to use. This parameter is optional. If no warehouse name is specified, the default warehouse for the user is used.
+ - Example: `.withWarehouse("MY_WAREHOUSE")`
+- `.withSchema(...)`
+ - Name of the schema in the database to use. This parameter is optional.
+ - Example: `.withSchema("PUBLIC")`
+
+
+**Note** - either `.withUrl(...)` or `.withServerName(...)` **is required**.
+
+## Pipeline options
+Use Beam’s [Pipeline options](https://beam.apache.org/releases/javadoc/2.17.0/org/apache/beam/sdk/options/PipelineOptions.html) to set options via the command line.
+### Snowflake Pipeline options
+Snowflake IO library supports following options that can be passed via the [command line](https://beam.apache.org/documentation/io/built-in/snowflake/#running-main-command-with-pipeline-options) by default when a Pipeline uses them:
+
+`--url` Snowflake's JDBC-like url including account name and region without any parameters.
+
+`--serverName` Full server name with account, zone and domain.
+
+`--username` Required for username/password and Private Key authentication.
+
+`--oauthToken` Required for OAuth authentication only.
+
+`--password` Required for username/password authentication only.
+
+`--privateKeyPath` Path to Private Key file. Required for Private Key authentication only.
+
+`--privateKeyPassphrase` Private Key's passphrase. Required for Private Key authentication only.
+
+`--stagingBucketName` External bucket path ending with `/`. I.e. `gs://bucket/`. Sub-directories are allowed.
+
+`--storageIntegrationName` Storage integration name
+
+`--warehouse` Warehouse to use. Optional.
+
+`--database` Database name to connect to. Optional.
+
+`--schema` Schema to use. Optional.
+
+`--table` Table to use. Optional.
+
+`--query` Query to use. Optional.
+
+`--role` Role to use. Optional.
+
+`--authenticator` Authenticator to use. Optional.
+
+`--portNumber` Port number. Optional.
+
+`--loginTimeout` Login timeout. Optional.
+
+## Running pipelines on Dataflow
+By default, pipelines are run on [Direct Runner](https://beam.apache.org/documentation/runners/direct/) on your local machine. To run a pipeline on [Google Dataflow](https://cloud.google.com/dataflow/), you must provide the following Pipeline options:
+
+- `--runner=DataflowRunner`
+ - The Dataflow’s specific runner.
+
+- `--project=<GCS PROJECT>`
+ - Name of the Google Cloud Platform project.
+
+- `--stagingBucketName=<GCS BUCKET NAME>`
+ - Google Cloud Services bucket where the Beam files will be staged.
+
+- `--maxNumWorkers=5`
+ - (optional) Maximum number of workers.
+
+- `--appName=<JOB NAME>`
+ - (optional) Prefix for the job name in the Dataflow Dashboard.
+
+More pipeline options for Dataflow can be found [here](https://beam.apache.org/releases/javadoc/2.17.0/org/apache/beam/runners/dataflow/options/DataflowPipelineOptions.html).
+
+**Note**: To properly authenticate with Google Cloud, please use [gcloud](https://cloud.google.com/sdk/gcloud/) or follow the [Google Cloud documentation](https://cloud.google.com/docs/authentication/).
+
+**Important**: Please acknowledge [Google Dataflow pricing](Important: Please acknowledge Google Dataflow pricing).
+
+## Writing to Snowflake tables
+One of the functions of SnowflakeIO is writing to Snowflake tables. This transformation enables you to finish the Beam pipeline with an output operation that sends the user's [PCollection](https://beam.apache.org/releases/javadoc/2.17.0/org/apache/beam/sdk/values/PCollection.html) to your Snowflake database.
+### Batch write (from a bounded source)
+The basic .`write()` operation usage is as follows:
+{{< highlight java >}}
+data.apply(
+ SnowflakeIO.<type>write()
+ .withDataSourceConfiguration(dc)
+ .to("MY_TABLE")
+ .withStagingBucketName("BUCKET NAME")
+ .withStorageIntegrationName("STORAGE INTEGRATION NAME")
+ .withUserDataMapper(mapper)
+)
+{{< /highlight >}}
+Replace type with the data type of the PCollection object to write; for example, SnowflakeIO.<String> for an input PCollection of Strings.
+
+All the below parameters are required:
+
+- `.withDataSourceConfiguration()` Accepts a DatasourceConfiguration object.
+
+- `.to()` Accepts the target Snowflake table name.
+
+- `.withStagingBucketName()` Accepts a cloud bucket path ended with slash.
+ -Example: `.withStagingBucketName("gs://mybucket/my/dir/")`
+
+- `.withStorageIntegrationName()` Accepts a name of a Snowflake storage integration object created according to Snowflake documentationt. Example:
+{{< highlight >}}
+CREATE OR REPLACE STORAGE INTEGRATION test_integration
+TYPE = EXTERNAL_STAGE
+STORAGE_PROVIDER = GCS
+ENABLED = TRUE
+STORAGE_ALLOWED_LOCATIONS = ('gcs://bucket/');
+{{< /highlight >}}
+Then:
+{{< highlight >}}
+.withStorageIntegrationName(test_integration)
+{{< /highlight >}}
+
+- `.withUserDataMapper()` Accepts the UserDataMapper function that will map a user's PCollection to an array of String values `(String[])`.
+
+**Note**:
+SnowflakeIO uses COPY statements behind the scenes to write (using [COPY to table](https://docs.snowflake.net/manuals/sql-reference/sql/copy-into-table.html)). StagingBucketName will be used to save CSV files which will end up in Snowflake. Those CSV files will be saved under the “stagingBucketName” path.
+
+### UserDataMapper function
+The UserDataMapper function is required to map data from a PCollection to an array of String values before the `write()` operation saves the data to temporary .csv files. For example:
+{{< highlight java >}}
+public static SnowflakeIO.UserDataMapper<Long> getCsvMapper() {
+ return (SnowflakeIO.UserDataMapper<Long>) recordLine -> new String[] {recordLine.toString()};
+}
+{{< /highlight >}}
+
+### Additional write options
+#### Transformation query
+The `.withQueryTransformation()` option for the `write()` operation accepts a SQL query as a String value, which will be performed while transfering data staged in CSV files directly to the target Snowflake table. For information about the transformation SQL syntax, see the [Snowflake Documentation](https://docs.snowflake.net/manuals/sql-reference/sql/copy-into-table.html#transformation-parameters).
+
+Usage:
+{{< highlight java >}}
+String query = "SELECT t.$1 from YOUR_TABLE;";
+data.apply(
+ SnowflakeIO.<~>write()
+ .withDataSourceConfiguration(dc)
+ .to("MY_TABLE")
+ .withStagingBucketName("BUCKET NAME")
+ .withStorageIntegrationName("STORAGE INTEGRATION NAME")
+ .withUserDataMapper(mapper)
+ .withQueryTransformation(query)
+)
+{{< /highlight >}}
+
+#### Write disposition
+Define the write behaviour based on the table where data will be written to by specifying the `.withWriteDisposition(...)` option for the `write()` operation. The following values are supported:
+
+- APPEND - Default behaviour. Written data is added to the existing rows in the table,
+
+- EMPTY - The target table must be empty; otherwise, the write operation fails,
+
+- TRUNCATE - The write operation deletes all rows from the target table before writing to it.
+
+Example of usage:
+{{< highlight java >}}
+data.apply(
+ SnowflakeIO.<~>write()
+ .withDataSourceConfiguration(dc)
+ .to("MY_TABLE")
+ .withStagingBucketName("BUCKET NAME")
+ .withStorageIntegrationName("STORAGE INTEGRATION NAME")
+ .withUserDataMapper(mapper)
+ .withWriteDisposition(TRUNCATE)
+)
+{{< /highlight >}}
+
+#### Create disposition
+The `.withCreateDisposition()` option defines the behavior of the write operation if the target table does not exist . The following values are supported:
+
+- CREATE_IF_NEEDED - default behaviour. The write operation checks whether the specified target table exists; if it does not, the write operation attempts to create the table Specify the schema for the target table using the `.withTableSchema()` option.
+
+- CREATE_NEVER - The write operation fails if the target table does not exist.
+
+Usage:
+{{< highlight java >}}
+data.apply(
+ SnowflakeIO.<~>write()
+ .withDataSourceConfiguration(dc)
+ .to("MY_TABLE")
+ .withStagingBucketName("BUCKET NAME")
+ .withStorageIntegrationName("STORAGE INTEGRATION NAME")
+ .withUserDataMapper(mapper)
+ .withCreateDisposition(CREATE_NEVER)
+)
+{{< /highlight >}}
+
+#### Table schema disposition
+When the `.withCreateDisposition()` .option is set to `CREATE_IF_NEEDED`, the `.withTableSchema()` option enables specifying the schema for the created target table.
+A table schema is a list of `SFColumn` objects with name and type corresponding to column type for each column in the table.
+
+Usage:
+{{< highlight java >}}
+SFTableSchema tableSchema =
+ new SFTableSchema(
+ SFColumn.of("my_date", new SFDate(), true),
+ new SFColumn("id", new SFNumber()),
+ SFColumn.of("name", new SFText(), true));
+
+data.apply(
+ SnowflakeIO.<~>write()
+ .withDataSourceConfiguration(dc)
+ .to("MY_TABLE")
+ .withStagingBucketName("BUCKET NAME")
+ .withStorageIntegrationName("STORAGE INTEGRATION NAME")
+ .withUserDataMapper(mapper)
+ .withTableSchema(tableSchema)
+)
+{{< /highlight >}}
+## Reading from Snowflake
+One of the functions of SnowflakeIO is reading Snowflake tables - either full tables via table name or custom data via query. Output of the read transform is a [PCollection](https://beam.apache.org/releases/javadoc/2.17.0/org/apache/beam/sdk/values/PCollection.html) of user-defined data type.
+
+### General usage
+
+The basic `.read()` operation usage:
+{{< highlight java >}}
+PCollection<USER_DATA_TYPE> items = pipeline.apply(
+ SnowflakeIO.<USER_DATA_TYPE>read()
+ .withDataSourceConfiguration(dc)
+ .fromTable("MY_TABLE") // or .fromQuery("QUERY")
+ .withStagingBucketName("BUCKET NAME")
+ .withStorageIntegrationName("STORAGE INTEGRATION NAME")
+ .withCsvMapper(mapper)
+ .withCoder(coder));
+)
+{{< /highlight >}}
+Where all below parameters are required:
+
+- `.withDataSourceConfiguration(...)`
+ - Accepts a DataSourceConfiguration object.
+
+- `.fromTable(...) or .fromQuery(...)`
+ - Specifies a Snowflake table name or custom SQL query.
+
+- `.withStagingBucketName()`
+ - Accepts a cloud bucket name.
+
+- `.withStorageIntegrationName()`
+ - Accepts a name of a Snowflake storage integration object created according to Snowflake documentation. Example:
+{{< highlight >}}
+CREATE OR REPLACE STORAGE INTEGRATION test_integration
+TYPE = EXTERNAL_STAGE
+STORAGE_PROVIDER = GCS
+ENABLED = TRUE
+STORAGE_ALLOWED_LOCATIONS = ('gcs://bucket/');
+{{< /highlight >}}
+Then:
+{{< highlight >}}
+.withStorageIntegrationName(test_integration)
+{{< /highlight >}}
+
+- `.withCsvMapper(mapper)`
+ - Accepts a [CSVMapper](https://beam.apache.org/documentation/io/built-in/snowflake/#csvmapper) instance for mapping String[] to USER_DATA_TYPE.
+- `.withCoder(coder)`
+ - Accepts the [Coder](https://beam.apache.org/releases/javadoc/2.0.0/org/apache/beam/sdk/coders/Coder.html) for USER_DATA_TYPE.
+
+**Note**:
+SnowflakeIO uses COPY statements behind the scenes to read (using [COPY to location](https://docs.snowflake.net/manuals/sql-reference/sql/copy-into-location.html)) files staged in cloud storage.StagingBucketName will be used as a temporary location for storing CSV files. Those temporary directories will be named `sf_copy_csv_DATE_TIME_RANDOMSUFFIX` and they will be removed automatically once Read operation finishes.
+
+### CSVMapper
+SnowflakeIO uses a [COPY INTO <location>](https://docs.snowflake.net/manuals/sql-reference/sql/copy-into-location.html) statement to move data from a Snowflake table to Google Cloud Storage as CSV files. These files are then downloaded via [FileIO](https://beam.apache.org/releases/javadoc/2.3.0/index.html?org/apache/beam/sdk/io/FileIO.html) and processed line by line. Each line is split into an array of Strings using the [OpenCSV](http://opencsv.sourceforge.net/) library.
+
+The CSVMapper’s job is to give the user the possibility to convert the array of Strings to a user-defined type, ie. GenericRecord for Avro or Parquet files, or custom POJO.
+
+Example implementation of CsvMapper for GenericRecord:
+{{< highlight java >}}
+static SnowflakeIO.CsvMapper<GenericRecord> getCsvMapper() {
+ return (SnowflakeIO.CsvMapper<GenericRecord>)
+ parts -> {
+ return new GenericRecordBuilder(PARQUET_SCHEMA)
+ .set("ID", Long.valueOf(parts[0]))
+ .set("NAME", parts[1])
+ [...]
+ .build();
+ };
+}
+{{< /highlight >}}
\ No newline at end of file
diff --git a/website/www/site/content/en/documentation/patterns/ai-platform.md b/website/www/site/content/en/documentation/patterns/ai-platform.md
index 905f895..b2a7b10 100644
--- a/website/www/site/content/en/documentation/patterns/ai-platform.md
+++ b/website/www/site/content/en/documentation/patterns/ai-platform.md
@@ -35,7 +35,7 @@
{{< /highlight >}}
{{< highlight java >}}
-// Java examples will be available on Beam 2.23 release.
+{{< code_sample "examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java" NlpAnalyzeText >}}
{{< /highlight >}}
@@ -79,7 +79,7 @@
{{< /highlight >}}
{{< highlight java >}}
-// Java examples will be available on Beam 2.23 release.
+{{< code_sample "examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java" NlpExtractSentiments >}}
{{< /highlight >}}
The snippet loops over `sentences` and, for each sentence, extracts the sentiment score.
@@ -99,7 +99,7 @@
{{< /highlight >}}
{{< highlight java >}}
-// Java examples will be available on Beam 2.23 release.
+{{< code_sample "examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java" NlpExtractEntities >}}
{{< /highlight >}}
Entities can be found in `entities` attribute. Just like before, `entities` is a sequence, that's why list comprehension is a viable choice. The most tricky part is interpreting the types of entities. Natural Language API defines entity types as enum. In a response object, entity types are returned as integers. That's why a user has to instantiate `naturallanguageml.enums.Entity.Type` to access a human-readable name.
@@ -119,7 +119,7 @@
{{< /highlight >}}
{{< highlight java >}}
-// Java examples will be available on Beam 2.23 release.
+{{< code_sample "examples/java/src/main/java/org/apache/beam/examples/snippets/Snippets.java" NlpAnalyzeDependencyTree >}}
{{< /highlight >}}
The output is below. For better readability, indexes are replaced by text which they refer to:
diff --git a/website/www/site/content/en/documentation/sdks/python-type-safety.md b/website/www/site/content/en/documentation/sdks/python-type-safety.md
index 074460c..755f795 100644
--- a/website/www/site/content/en/documentation/sdks/python-type-safety.md
+++ b/website/www/site/content/en/documentation/sdks/python-type-safety.md
@@ -71,7 +71,7 @@
Using Annotations has the added benefit of allowing use of a static type checker (such as mypy) to additionally type check your code.
If you already use a type checker, using annotations instead of decorators reduces code duplication.
However, annotations do not cover all the use cases that decorators and inline declarations do.
-Two such are the `expand` of a composite transform and lambda functions.
+For instance, they do not work for lambda functions.
### Declaring Type Hints Using Type Annotations
@@ -82,6 +82,7 @@
Annotations are currently supported on:
- `process()` methods on `DoFn` subclasses.
+ - `expand()` methods on `PTransform` subclasses.
- Functions passed to: `ParDo`, `Map`, `FlatMap`, `Filter`.
The following code declares an `int` input and a `str` output type hint on the `to_id` transform, using annotations on `my_fn`.
@@ -90,6 +91,15 @@
{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test_py3.py" type_hints_map_annotations >}}
{{< /highlight >}}
+The following code demonstrates how to use annotations on `PTransform` subclasses.
+A valid annotation is a `PCollection` that wraps an internal (nested) type, `PBegin`, `PDone`, or `None`.
+The following code declares typehints on a custom PTransform, that takes a `PCollection[int]` input
+and outputs a `PCollection[str]`, using annotations.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/snippets_test_py3.py" type_hints_ptransforms >}}
+{{< /highlight >}}
+
The following code declares `int` input and output type hints on `filter_evens`, using annotations on `FilterEvensDoFn.process`.
Since `process` returns a generator, the output type for a DoFn producing a `PCollection[int]` is annotated as `Iterable[int]` (`Generator[int, None, None]` would also work here).
Beam will remove the outer iterable of the return type on the `DoFn.process` method and functions passed to `FlatMap` to deduce the element type of resulting PCollection .
@@ -182,6 +192,7 @@
* `Iterable[T]`
* `Iterator[T]`
* `Generator[T]`
+* `PCollection[T]`
**Note:** The `Tuple[T, U]` type hint is a tuple with a fixed number of heterogeneously typed elements, while the `Tuple[T, ...]` type hint is a tuple with a variable of homogeneously typed elements.
diff --git a/website/www/site/content/en/documentation/transforms/python/aggregation/sum.md b/website/www/site/content/en/documentation/transforms/python/aggregation/sum.md
index 8d49f5b..a9aef38 100644
--- a/website/www/site/content/en/documentation/transforms/python/aggregation/sum.md
+++ b/website/www/site/content/en/documentation/transforms/python/aggregation/sum.md
@@ -17,7 +17,61 @@
# Sum
-## Examples
-See [BEAM-7390](https://issues.apache.org/jira/browse/BEAM-7390) for updates.
+{{< localstorage language language-py >}}
-## Related transforms
\ No newline at end of file
+{{< button-pydoc path="apache_beam.transforms.core" class="CombineGlobally" >}}
+
+Sums all the elements within each aggregation.
+
+## Examples
+
+In the following example, we create a pipeline with a `PCollection`.
+Then, we get the sum of all the element values in different ways.
+
+### Example 1: Sum of the elements in a PCollection
+
+We use `Combine.Globally()` to get sum of all the element values from the *entire* `PCollection`.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum.py" sum_globally >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py" total >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum.py" >}}
+
+### Example 2: Sum of the elements for each key
+
+We use `Combine.PerKey()` to get the sum of all the element values for each unique key in a `PCollection` of key-values.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum.py" sum_per_key >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py" totals_per_key >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum.py" >}}
+
+## Related transforms
+
+* [CombineGlobally](/documentation/transforms/python/aggregation/combineglobally)
+* [CombinePerKey](/documentation/transforms/python/aggregation/combineperkey)
+* [Max](/documentation/transforms/python/aggregation/max)
+* [Mean](/documentation/transforms/python/aggregation/mean)
+* [Min](/documentation/transforms/python/aggregation/min)
+
+{{< button-pydoc path="apache_beam.transforms.core" class="CombineGlobally" >}}
diff --git a/website/www/site/content/en/documentation/transforms/python/aggregation/top.md b/website/www/site/content/en/documentation/transforms/python/aggregation/top.md
index b1d5eba..419c0f8 100644
--- a/website/www/site/content/en/documentation/transforms/python/aggregation/top.md
+++ b/website/www/site/content/en/documentation/transforms/python/aggregation/top.md
@@ -14,24 +14,144 @@
See the License for the specific language governing permissions and
limitations under the License.
-->
-# Top
-<table align="left">
- <a target="_blank" class="button"
- href="https://beam.apache.org/releases/pydoc/2.12.0/apache_beam.transforms.combiners.html?#apache_beam.transforms.combiners.Top">
- <img src="https://beam.apache.org/images/logos/sdks/python.png" width="20px" height="20px"
- alt="Pydoc" />
- Pydoc
- </a>
-</table>
-<br><br>
+# Top
+
+{{< localstorage language language-py >}}
+
+{{< button-pydoc path="apache_beam.transforms.combiners" class="Top" >}}
Transforms for finding the largest (or smallest) set of elements in
a collection, or the largest (or smallest) set of values associated
with each key in a collection of key-value pairs.
## Examples
-See [BEAM-7390](https://issues.apache.org/jira/browse/BEAM-7390) for updates.
-## Related transforms
-* [Sample](/documentation/transforms/python/aggregation/sample) to combine elements. takes samples of the elements in a collection.
\ No newline at end of file
+In the following example, we create a pipeline with a `PCollection`.
+Then, we get the largest or smallest elements in different ways.
+
+### Example 1: Largest elements from a PCollection
+
+We use `Top.Largest()` to get the largest elements from the *entire* `PCollection`.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" top_largest >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py" largest_elements >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" >}}
+
+### Example 2: Largest elements for each key
+
+We use `Top.LargestPerKey()` to get the largest elements for each unique key in a `PCollection` of key-values.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" top_largest_per_key >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py" largest_elements_per_key >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" >}}
+
+### Example 3: Smallest elements from a PCollection
+
+We use `Top.Smallest()` to get the smallest elements from the *entire* `PCollection`.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" top_smallest >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py" smallest_elements >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" >}}
+
+### Example 4: Smallest elements for each key
+
+We use `Top.SmallestPerKey()` to get the smallest elements for each unique key in a `PCollection` of key-values.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" top_smallest_per_key >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py" smallest_elements_per_key >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" >}}
+
+### Example 5: Custom elements from a PCollection
+
+We use `Top.Of()` to get elements with customized rules from the *entire* `PCollection`.
+
+You can change how the elements are compared with `key`.
+By default you get the largest elements, but you can get the smallest by setting `reverse=True`.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" top_of >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py" shortest_elements >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" >}}
+
+### Example 6: Custom elements for each key
+
+We use `Top.PerKey()` to get elements with customized rules for each unique key in a `PCollection` of key-values.
+
+You can change how the elements are compared with `key`.
+By default you get the largest elements, but you can get the smallest by setting `reverse=True`.
+
+{{< highlight py >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" top_per_key >}}
+{{< /highlight >}}
+
+{{< paragraph class="notebook-skip" >}}
+Output:
+{{< /paragraph >}}
+
+{{< highlight class="notebook-skip" >}}
+{{< code_sample "sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py" shortest_elements_per_key >}}
+{{< /highlight >}}
+
+{{< buttons-code-snippet
+ py="sdks/python/apache_beam/examples/snippets/transforms/aggregation/top.py" >}}
+
+## Related transforms
+
+* [Sample](/documentation/transforms/python/aggregation/sample) to combine elements. Takes samples of the elements in a collection.
+
+{{< button-pydoc path="apache_beam.transforms.combiners" class="Top" >}}
diff --git a/website/www/site/content/en/documentation/transforms/python/overview.md b/website/www/site/content/en/documentation/transforms/python/overview.md
index d648ac7..059abcb 100644
--- a/website/www/site/content/en/documentation/transforms/python/overview.md
+++ b/website/www/site/content/en/documentation/transforms/python/overview.md
@@ -60,7 +60,7 @@
<tr><td><a href="/documentation/transforms/python/aggregation/mean">Mean</a></td><td>Computes the average within each aggregation.</td></tr>
<tr><td><a href="/documentation/transforms/python/aggregation/min">Min</a></td><td>Gets the element with the minimum value within each aggregation.</td></tr>
<tr><td><a href="/documentation/transforms/python/aggregation/sample">Sample</a></td><td>Randomly select some number of elements from each aggregation.</td></tr>
- <tr><td>Sum</td><td>Not available.</td></tr>
+ <tr><td><a href="/documentation/transforms/python/aggregation/sum">Sum</a></td><td>Sums all the elements within each aggregation.</td></tr>
<tr><td><a href="/documentation/transforms/python/aggregation/top">Top</a></td><td>Compute the largest element(s) in each aggregation.</td></tr>
</table>
diff --git a/website/www/site/content/en/roadmap/connectors-multi-sdk.md b/website/www/site/content/en/roadmap/connectors-multi-sdk.md
index 464ad83..3b13f59 100644
--- a/website/www/site/content/en/roadmap/connectors-multi-sdk.md
+++ b/website/www/site/content/en/roadmap/connectors-multi-sdk.md
@@ -80,6 +80,7 @@
* Java KafkaIO - completed - [BEAM-7029](https://issues.apache.org/jira/browse/BEAM-7029)
* Java KinesisIO - In progress - [BEAM-10137](https://issues.apache.org/jira/browse/BEAM-10137), [BEAM-10138](https://issues.apache.org/jira/browse/BEAM-10138)
* Java PubSubIO - In progress - [BEAM-7738](https://issues.apache.org/jira/browse/BEAM-7738)
+* Java SnowflakeIO - In progress - [BEAM-9897](https://issues.apache.org/jira/browse/BEAM-9897), [BEAM-9898](https://issues.apache.org/jira/browse/BEAM-9898)
* Java SpannerIO - In progress - [BEAM-10139](https://issues.apache.org/jira/browse/BEAM-10139), [BEAM-10140](https://issues.apache.org/jira/browse/BEAM-10140)
* Java SQL - completed - [BEAM-8603](https://issues.apache.org/jira/browse/BEAM-8603)
diff --git a/website/www/site/data/io_matrix.yaml b/website/www/site/data/io_matrix.yaml
index 5923a86..84637c6 100644
--- a/website/www/site/data/io_matrix.yaml
+++ b/website/www/site/data/io_matrix.yaml
@@ -305,6 +305,7 @@
url: https://beam.apache.org/releases/pydoc/current/apache_beam.io.gcp.datastore.v1new.datastoreio.html
- transform: SnowflakeIO
description: Experimental Transforms for reading from and writing to [Snowflake](https://www.snowflake.com/).
+ docs: /documentation/io/built-in/snowflake
implementations:
- language: java
name: org.apache.beam.sdk.io.snowflake.SnowflakeIO
diff --git a/website/www/site/layouts/partials/section-menu/en/documentation.html b/website/www/site/layouts/partials/section-menu/en/documentation.html
index c8d0da9..f715f0c 100644
--- a/website/www/site/layouts/partials/section-menu/en/documentation.html
+++ b/website/www/site/layouts/partials/section-menu/en/documentation.html
@@ -74,6 +74,7 @@
<li><a href="/documentation/io/built-in/hadoop/">Hadoop Input/Output Format IO</a></li>
<li><a href="/documentation/io/built-in/hcatalog/">HCatalog IO</a></li>
<li><a href="/documentation/io/built-in/google-bigquery/">Google BigQuery I/O connector</a></li>
+ <li><a href="/documentation/io/built-in/snowflake/">Snowflake I/O connector</a></li>
</ul>
</li>
@@ -204,6 +205,7 @@
<li><a href="/documentation/transforms/python/aggregation/mean/">Mean</a></li>
<li><a href="/documentation/transforms/python/aggregation/min/">Min</a></li>
<li><a href="/documentation/transforms/python/aggregation/sample/">Sample</a></li>
+ <li><a href="/documentation/transforms/python/aggregation/sum/">Sum</a></li>
<li><a href="/documentation/transforms/python/aggregation/top/">Top</a></li>
</ul>
</li>
