You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by GitBox <gi...@apache.org> on 2020/12/07 19:11:39 UTC
[GitHub] [beam] nehsyc opened a new pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
nehsyc opened a new pull request #13496:
URL: https://github.com/apache/beam/pull/13496
Use `GroupIntoBatches.WithShardedKey` API to group and batch write before streaming to BigQuery service. Currently batching is done best-effort on bundle finalization.
This PR
- adds an option to `BigQueryOptions` to toggle between the existing and new implementation;
- extracts the shared code between the old and new implementation to a class `BatchedStreamingWrite` and provides an option to choose the implementation.
------------------------
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
- [ ] [**Choose reviewer(s)**](https://beam.apache.org/contribute/#make-your-change) and mention them in a comment (`R: @username`).
- [ ] Format the pull request title like `[BEAM-XXX] Fixes bug in ApproximateQuantiles`, where you replace `BEAM-XXX` with the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.
- [ ] Update `CHANGES.md` with noteworthy changes.
- [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://beam.apache.org/contribute/#make-reviewers-job-easier).
Post-Commit Tests Status (on master branch)
------------------------------------------------------------------------------------------------
Lang | SDK | Dataflow | Flink | Samza | Spark | Twister2
--- | --- | --- | --- | --- | --- | ---
Go | [](https://ci-beam.apache.org/job/beam_PostCommit_Go/lastCompletedBuild/) | --- | [](https://ci-beam.apache.org/job/beam_PostCommit_Go_VR_Flink/lastCompletedBuild/) | --- | [](https://ci-beam.apache.org/job/beam_PostCommit_Go_VR_Spark/lastCompletedBuild/) | ---
Java | [](https://ci-beam.apache.org/job/beam_PostCommit_Java/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Java_VR_Dataflow_V2/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Dataflow_Java11/lastCompletedBuild/) | [](https://ci-beam
.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Flink_Java11/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Java_PVR_Flink_Batch/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Java_PVR_Flink_Streaming/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Samza/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Spark/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Java_PVR_Spark_Batch/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PostCommit_Java_ValidatesRunner_Twister2/lastCompletedBuild/)
Python | [](https://ci-beam.apache.org/job/beam_PostCommit_Python36/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Python37/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Python38/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PostCommit_Py_VR_Dataflow/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Py_VR_Dataflow_V2/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Py_ValCont/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PreCommit_Python_PVR_Flink_Cron/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PostCommit_Python_VR_Flink/lastCompletedBuild/) | --- | [](https://ci-beam.apache.org/job/beam_PostCommit_Python_VR_Spark/lastCompletedBuild/) | ---
XLang | [](https://ci-beam.apache.org/job/beam_PostCommit_XVR_Direct/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PostCommit_XVR_Dataflow/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PostCommit_XVR_Flink/lastCompletedBuild/) | --- | [](https://ci-beam.apache.org/job/beam_PostCommit_XVR_Spark/lastCompletedBuild/) | ---
Pre-Commit Tests Status (on master branch)
------------------------------------------------------------------------------------------------
--- |Java | Python | Go | Website | Whitespace | Typescript
--- | --- | --- | --- | --- | --- | ---
Non-portable | [](https://ci-beam.apache.org/job/beam_PreCommit_Java_Cron/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PreCommit_Python_Cron/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PreCommit_PythonLint_Cron/lastCompletedBuild/)<br>[](https://ci-beam.apache.org/job/beam_PreCommit_PythonDocker_Cron/lastCompletedBuild/) <br>[](https://ci-beam.apache.org/job/beam_PreCommit_PythonDocs_Cron/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PreCommit_Go_Cron/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PreCommit_Website_Cron/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PreCommit_Whitespace_Cron/lastCompletedBuild/) | [](https://ci-beam.apache.org/job/beam_PreCommit_Typescript_Cron/lastCompletedBuild/)
Portable | --- | [](https://ci-beam.apache.org/job/beam_PreCommit_Portable_Python_Cron/lastCompletedBuild/) | --- | --- | --- | ---
See [.test-infra/jenkins/README](https://github.com/apache/beam/blob/master/.test-infra/jenkins/README.md) for trigger phrase, status and link of all Jenkins jobs.
GitHub Actions Tests Status (on master branch)
------------------------------------------------------------------------------------------------
[](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768627401
Run Java_Examples_Dataflow_Java11 PreCommit
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768741698
Run Java PreCommit
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r554123442
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
##########
@@ -78,4 +79,11 @@
Integer getLatencyLoggingFrequency();
void setLatencyLoggingFrequency(Integer value);
+
+ @Experimental
+ @Description("Whether dynamic sharding is enabled for writing to BigQuery in streaming.")
+ @Default.Boolean(false)
+ Boolean getEnableStreamingAutoSharding();
Review comment:
I was thinking that adding the option to BigQueryOptions might be easier for the users to switch on/off the feature. If it is an option of BigQueryIO users would need to update their code (vs. adding a flag) right? It might not be a concern...
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768664872
Run Java PostCommit
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r553774115
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.standardSeconds(10);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
Review comment:
Good point! I added a global window right before `GroupIntoBatches`.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
##########
@@ -243,61 +243,108 @@ public WriteResult expand(PCollection<KV<TableDestination, ElementT>> input) {
AtomicCoder<T> coder,
ErrorContainer<T> errorContainer) {
BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
- int numShards = options.getNumStreamingKeys();
// A naive implementation would be to simply stream data directly to BigQuery.
// However, this could occasionally lead to duplicated data, e.g., when
// a VM that runs this code is restarted and the code is re-run.
// The above risk is mitigated in this implementation by relying on
// BigQuery built-in best effort de-dup mechanism.
-
// To use this mechanism, each input TableRow is tagged with a generated
- // unique id, which is then passed to BigQuery and used to ignore duplicates
- // We create 50 keys per BigQuery table to generate output on. This is few enough that we
- // get good batching into BigQuery's insert calls, and enough that we can max out the
- // streaming insert quota.
- PCollection<KV<ShardedKey<String>, TableRowInfo<ElementT>>> tagged =
- input
- .apply("ShardTableWrites", ParDo.of(new GenerateShardedTable<>(numShards)))
- .setCoder(KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), elementCoder))
- .apply("TagWithUniqueIds", ParDo.of(new TagWithUniqueIds<>()))
- .setCoder(
- KvCoder.of(
- ShardedKeyCoder.of(StringUtf8Coder.of()), TableRowInfoCoder.of(elementCoder)));
-
- TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+ // unique id, which is then passed to BigQuery and used to ignore duplicates.
// To prevent having the same TableRow processed more than once with regenerated
// different unique ids, this implementation relies on "checkpointing", which is
- // achieved as a side effect of having StreamingWriteFn immediately follow a GBK,
- // performed by Reshuffle.
- PCollectionTuple tuple =
- tagged
- .apply(Reshuffle.of())
- // Put in the global window to ensure that DynamicDestinations side inputs are accessed
- // correctly.
- .apply(
- "GlobalWindow",
- Window.<KV<ShardedKey<String>, TableRowInfo<ElementT>>>into(new GlobalWindows())
- .triggering(DefaultTrigger.of())
- .discardingFiredPanes())
- .apply(
- "StreamingWrite",
- ParDo.of(
- new StreamingWriteFn<>(
- bigQueryServices,
- retryPolicy,
- failedInsertsTag,
- errorContainer,
- skipInvalidRows,
- ignoreUnknownValues,
- ignoreInsertIds,
- toTableRow,
- toFailsafeTableRow))
- .withOutputTags(mainOutputTag, TupleTagList.of(failedInsertsTag)));
- PCollection<T> failedInserts = tuple.get(failedInsertsTag);
- failedInserts.setCoder(coder);
- return failedInserts;
+ // achieved as a side effect of having BigQuery insertion immediately follow a GBK.
+
+ if (options.getEnableStreamingAutoSharding()) {
Review comment:
Done.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
##########
@@ -78,4 +79,11 @@
Integer getLatencyLoggingFrequency();
void setLatencyLoggingFrequency(Integer value);
+
+ @Experimental
+ @Description("Whether dynamic sharding is enabled for writing to BigQuery in streaming.")
+ @Default.Boolean(false)
+ Boolean getEnableStreamingAutoSharding();
Review comment:
Done.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r565688130
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.millis(200);
Review comment:
Yeah sounds reasonable to proceed with this for now.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r553774322
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.millis(200);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
+ // Group and batch table rows such that each batch has no more than
+ // getMaxStreamingRowsToBatch rows. Also set a buffering time limit to avoid being
+ // stuck at a partial batch forever, especially in a global window.
+ .apply(
+ GroupIntoBatches.<String, TableRowInfo<ElementT>>ofSize(
+ options.getMaxStreamingRowsToBatch())
+ .withMaxBufferingDuration(BATCH_MAX_BUFFERING_DURATION)
+ .withShardedKey())
+ .setCoder(
+ KvCoder.of(
+ ShardedKey.Coder.of(StringUtf8Coder.of()), IterableCoder.of(valueCoder)))
+ .apply(
+ ParDo.of(new InsertBatchedElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ private class InsertBatchedElements
Review comment:
There is a transform override in Dataflow to add a preceding `Reshuffle` to DoFns marked with `RequiresStableInput`.
https://github.com/apache/beam/blob/74ec6093255d12ad03f17717c26c2392e3991a1e/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RequiresStableInputParDoOverrides.java#L73
The override is disabled though. So I guess currently Dataflow does nothing for this tag.
https://github.com/apache/beam/blob/df74d74ef295aa972b0d1ddef2429f427b988a25/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java#L601
Also my understanding is that dding a `Reshuffle` before `GroupIntoBatches` will introduce an extra shuffle as `Reshuffle` is essentially a GBK + value expansion in Dataflow.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.millis(200);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
+ // Group and batch table rows such that each batch has no more than
+ // getMaxStreamingRowsToBatch rows. Also set a buffering time limit to avoid being
+ // stuck at a partial batch forever, especially in a global window.
+ .apply(
+ GroupIntoBatches.<String, TableRowInfo<ElementT>>ofSize(
+ options.getMaxStreamingRowsToBatch())
+ .withMaxBufferingDuration(BATCH_MAX_BUFFERING_DURATION)
+ .withShardedKey())
+ .setCoder(
+ KvCoder.of(
+ ShardedKey.Coder.of(StringUtf8Coder.of()), IterableCoder.of(valueCoder)))
+ .apply(
+ ParDo.of(new InsertBatchedElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ private class InsertBatchedElements
+ extends DoFn<KV<ShardedKey<String>, Iterable<TableRowInfo<ElementT>>>, Void> {
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV<ShardedKey<String>, Iterable<TableRowInfo<ElementT>>> input,
+ BoundedWindow window,
+ ProcessContext context)
+ throws InterruptedException {
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> tableRows = new ArrayList<>();
+ List<String> uniqueIds = new ArrayList<>();
+ for (TableRowInfo<ElementT> row : input.getValue()) {
+ TableRow tableRow = toTableRow.apply(row.tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(row.tableRow);
+ tableRows.add(
+ FailsafeValueInSingleWindow.of(
+ tableRow, context.timestamp(), window, context.pane(), failsafeTableRow));
+ uniqueIds.add(row.uniqueId);
+ }
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(input.getKey().getKey());
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ flushRows(tableReference, tableRows, uniqueIds, options, failedInserts);
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue());
Review comment:
Done.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r549453687
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.millis(200);
Review comment:
Not sure about this limit. What would be a proper value? Should we make it configurable?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768741780
Postcommit from previous commit: https://ci-beam.apache.org/job/beam_PostCommit_Java_PR/558/
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r565646491
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.millis(200);
Review comment:
to be hones, I am not sure what's a good duration either. I think this is acceptable for now, until we find out more. Thoughts?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] reuvenlax commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
reuvenlax commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r537890083
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
##########
@@ -243,61 +243,108 @@ public WriteResult expand(PCollection<KV<TableDestination, ElementT>> input) {
AtomicCoder<T> coder,
ErrorContainer<T> errorContainer) {
BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
- int numShards = options.getNumStreamingKeys();
// A naive implementation would be to simply stream data directly to BigQuery.
// However, this could occasionally lead to duplicated data, e.g., when
// a VM that runs this code is restarted and the code is re-run.
// The above risk is mitigated in this implementation by relying on
// BigQuery built-in best effort de-dup mechanism.
-
// To use this mechanism, each input TableRow is tagged with a generated
- // unique id, which is then passed to BigQuery and used to ignore duplicates
- // We create 50 keys per BigQuery table to generate output on. This is few enough that we
- // get good batching into BigQuery's insert calls, and enough that we can max out the
- // streaming insert quota.
- PCollection<KV<ShardedKey<String>, TableRowInfo<ElementT>>> tagged =
- input
- .apply("ShardTableWrites", ParDo.of(new GenerateShardedTable<>(numShards)))
- .setCoder(KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), elementCoder))
- .apply("TagWithUniqueIds", ParDo.of(new TagWithUniqueIds<>()))
- .setCoder(
- KvCoder.of(
- ShardedKeyCoder.of(StringUtf8Coder.of()), TableRowInfoCoder.of(elementCoder)));
-
- TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+ // unique id, which is then passed to BigQuery and used to ignore duplicates.
// To prevent having the same TableRow processed more than once with regenerated
// different unique ids, this implementation relies on "checkpointing", which is
- // achieved as a side effect of having StreamingWriteFn immediately follow a GBK,
- // performed by Reshuffle.
- PCollectionTuple tuple =
- tagged
- .apply(Reshuffle.of())
- // Put in the global window to ensure that DynamicDestinations side inputs are accessed
- // correctly.
- .apply(
- "GlobalWindow",
- Window.<KV<ShardedKey<String>, TableRowInfo<ElementT>>>into(new GlobalWindows())
- .triggering(DefaultTrigger.of())
- .discardingFiredPanes())
- .apply(
- "StreamingWrite",
- ParDo.of(
- new StreamingWriteFn<>(
- bigQueryServices,
- retryPolicy,
- failedInsertsTag,
- errorContainer,
- skipInvalidRows,
- ignoreUnknownValues,
- ignoreInsertIds,
- toTableRow,
- toFailsafeTableRow))
- .withOutputTags(mainOutputTag, TupleTagList.of(failedInsertsTag)));
- PCollection<T> failedInserts = tuple.get(failedInsertsTag);
- failedInserts.setCoder(coder);
- return failedInserts;
+ // achieved as a side effect of having BigQuery insertion immediately follow a GBK.
+
+ if (options.getEnableStreamingAutoSharding()) {
Review comment:
I would make this specified in the BigQueryIO builder instead.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
##########
@@ -78,4 +79,11 @@
Integer getLatencyLoggingFrequency();
void setLatencyLoggingFrequency(Integer value);
+
+ @Experimental
+ @Description("Whether dynamic sharding is enabled for writing to BigQuery in streaming.")
+ @Default.Boolean(false)
+ Boolean getEnableStreamingAutoSharding();
Review comment:
I think this should be an option on BigQueryIO.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
Review comment:
add space
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.millis(200);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
+ // Group and batch table rows such that each batch has no more than
+ // getMaxStreamingRowsToBatch rows. Also set a buffering time limit to avoid being
+ // stuck at a partial batch forever, especially in a global window.
+ .apply(
+ GroupIntoBatches.<String, TableRowInfo<ElementT>>ofSize(
+ options.getMaxStreamingRowsToBatch())
+ .withMaxBufferingDuration(BATCH_MAX_BUFFERING_DURATION)
+ .withShardedKey())
+ .setCoder(
+ KvCoder.of(
+ ShardedKey.Coder.of(StringUtf8Coder.of()), IterableCoder.of(valueCoder)))
+ .apply(
+ ParDo.of(new InsertBatchedElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ private class InsertBatchedElements
+ extends DoFn<KV<ShardedKey<String>, Iterable<TableRowInfo<ElementT>>>, Void> {
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV<ShardedKey<String>, Iterable<TableRowInfo<ElementT>>> input,
+ BoundedWindow window,
+ ProcessContext context)
+ throws InterruptedException {
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> tableRows = new ArrayList<>();
+ List<String> uniqueIds = new ArrayList<>();
+ for (TableRowInfo<ElementT> row : input.getValue()) {
+ TableRow tableRow = toTableRow.apply(row.tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(row.tableRow);
+ tableRows.add(
+ FailsafeValueInSingleWindow.of(
+ tableRow, context.timestamp(), window, context.pane(), failsafeTableRow));
+ uniqueIds.add(row.uniqueId);
+ }
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(input.getKey().getKey());
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ flushRows(tableReference, tableRows, uniqueIds, options, failedInserts);
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue());
Review comment:
use a MultiOutputReceiver instead
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
Review comment:
seems more direct to just call tableRows.computeIfAbsent(tableSpec, ...)
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.millis(200);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
+ // Group and batch table rows such that each batch has no more than
+ // getMaxStreamingRowsToBatch rows. Also set a buffering time limit to avoid being
+ // stuck at a partial batch forever, especially in a global window.
+ .apply(
+ GroupIntoBatches.<String, TableRowInfo<ElementT>>ofSize(
+ options.getMaxStreamingRowsToBatch())
+ .withMaxBufferingDuration(BATCH_MAX_BUFFERING_DURATION)
+ .withShardedKey())
+ .setCoder(
+ KvCoder.of(
+ ShardedKey.Coder.of(StringUtf8Coder.of()), IterableCoder.of(valueCoder)))
+ .apply(
+ ParDo.of(new InsertBatchedElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ private class InsertBatchedElements
Review comment:
We are relying on the fact that the GroupIntoBatches produces stable output. Really we should tag this with RequiresStableInput. Can you find out if this is safe to do in Dataflow?
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.standardSeconds(10);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
Review comment:
Hmm good question. I wonder if someone refactored the code at some point to change things? I'm not entirely sure about the global window here.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
Review comment:
Why is this "stateful"?
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.standardSeconds(10);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
Review comment:
However we probably want to put in global window here anyway as the GroupIntoBatches should logically be in hte global window, right? If the user had tiny windows, we don't want that to result in tiny grouping.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
Review comment:
instead add an OutputReceiver parameter to finishBundle.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768623857
Run Java PostCommit
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768615727
There are a couple checkstyle warnings on Precommit:
```
[ant:checkstyle] [ERROR] /home/jenkins/jenkins-slave/workspace/beam_PreCommit_Java_Phrase/src/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java:85:24: Name 'LOG' must match pattern '^[a-z][a-zA-Z0-9]*$'. [MemberName] | Â
 | [ant:checkstyle] [ERROR] /home/jenkins/jenkins-slave/workspace/beam_PreCommit_Java_Phrase/src/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java:279:28: Name 'BATCH_MAX_BUFFERING_DURATION' must match pattern '^[a-z][a-zA-Z0-9]*$'. [MemberName]
```
These mean that the variables should be static to have CAPITALIZED_NAMES, or should be named with camelCase if they are not static. Can you fix that?
Also, it seems there are some merge conflicts. Can you fix that as well?
The last thing to figure out is how to address Reuven's comment regarding stable input
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768628744
> There are a couple checkstyle warnings on Precommit:
>
> ```
> [ant:checkstyle] [ERROR] /home/jenkins/jenkins-slave/workspace/beam_PreCommit_Java_Phrase/src/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java:85:24: Name 'LOG' must match pattern '^[a-z][a-zA-Z0-9]*$'. [MemberName] | Â
> Â | [ant:checkstyle] [ERROR] /home/jenkins/jenkins-slave/workspace/beam_PreCommit_Java_Phrase/src/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java:279:28: Name 'BATCH_MAX_BUFFERING_DURATION' must match pattern '^[a-z][a-zA-Z0-9]*$'. [MemberName]
> ```
>
> These mean that the variables should be static to have CAPITALIZED_NAMES, or should be named with camelCase if they are not static. Can you fix that?
Thanks for pointing out! Pushed a commit to fix those.
>
> Also, it seems there are some merge conflicts. Can you fix that as well?
Done.
>
> The last thing to figure out is how to address Reuven's comment regarding stable input
Based on my understanding Dataflow currently doesn't do anything for @RequiresStableInputs so it may be considered as safe for now but if we add naive support like a Reshuffle it would be adding duplicated shuffles. How about adding a TODO here so we don't forget?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768567644
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768769653
Thanks @nehsyc !
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r553774215
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
Review comment:
Done.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
Review comment:
Done.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
Review comment:
`finishBundle` doesn't seems to accept parameters other than `FinishBundleContext`.
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
Review comment:
Because input is batched through a stateful DoFn, `GroupIntoBatches`. Any suggestion on the name?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-740200014
R: @reuvenlax @chamikaramj
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-767207889
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-767207889
R: @pabloem
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-767209597
> R: @pabloem
Hey Pablo, let me know if it is ok for you to include the changes in FILE_LOADS in this PR as well. If so, I will push a new commit.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r537860362
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchedStreamingWrite.java
##########
@@ -0,0 +1,402 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT 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.gcp.bigquery;
+
+import com.google.api.services.bigquery.model.TableReference;
+import com.google.api.services.bigquery.model.TableRow;
+import java.io.IOException;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.AtomicCoder;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.SinkMetrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.Histogram;
+import org.apache.beam.sdk.util.ShardedKey;
+import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** PTransform to perform batched streaming BigQuery write. */
+@SuppressWarnings({
+ "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
+})
+class BatchedStreamingWrite<ErrorT, ElementT>
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private static final TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
+
+ private final BigQueryServices bqServices;
+ private final InsertRetryPolicy retryPolicy;
+ private final TupleTag<ErrorT> failedOutputTag;
+ private final AtomicCoder<ErrorT> failedOutputCoder;
+ private final ErrorContainer<ErrorT> errorContainer;
+ private final boolean skipInvalidRows;
+ private final boolean ignoreUnknownValues;
+ private final boolean ignoreInsertIds;
+ private final SerializableFunction<ElementT, TableRow> toTableRow;
+ private final SerializableFunction<ElementT, TableRow> toFailsafeTableRow;
+
+ /** Tracks histogram of bytes written. Reset at the start of every bundle. */
+ private transient Histogram histogram = Histogram.linear(0, 20, 3000);
+
+ private transient Long lastReportedSystemClockMillis = System.currentTimeMillis();
+
+ private final Logger LOG = LoggerFactory.getLogger(BatchedStreamingWrite.class);
+
+ /** Tracks bytes written, exposed as "ByteCount" Counter. */
+ private Counter byteCounter = SinkMetrics.bytesWritten();
+
+ /** Switches the method of batching. */
+ private final boolean batchViaStateful;
+
+ public BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = false;
+ }
+
+ private BatchedStreamingWrite(
+ BigQueryServices bqServices,
+ InsertRetryPolicy retryPolicy,
+ TupleTag<ErrorT> failedOutputTag,
+ AtomicCoder<ErrorT> failedOutputCoder,
+ ErrorContainer<ErrorT> errorContainer,
+ boolean skipInvalidRows,
+ boolean ignoreUnknownValues,
+ boolean ignoreInsertIds,
+ SerializableFunction<ElementT, TableRow> toTableRow,
+ SerializableFunction<ElementT, TableRow> toFailsafeTableRow,
+ boolean batchViaStateful) {
+ this.bqServices = bqServices;
+ this.retryPolicy = retryPolicy;
+ this.failedOutputTag = failedOutputTag;
+ this.failedOutputCoder = failedOutputCoder;
+ this.errorContainer = errorContainer;
+ this.skipInvalidRows = skipInvalidRows;
+ this.ignoreUnknownValues = ignoreUnknownValues;
+ this.ignoreInsertIds = ignoreInsertIds;
+ this.toTableRow = toTableRow;
+ this.toFailsafeTableRow = toFailsafeTableRow;
+ this.batchViaStateful = batchViaStateful;
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are batched and
+ * flushed upon bundle finalization.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaDoFnFinalization() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ false);
+ }
+
+ /**
+ * A transform that performs batched streaming BigQuery write; input elements are grouped on table
+ * destinations and batched via a stateful DoFn. This also enables dynamic sharding during
+ * grouping to parallelize writes.
+ */
+ public BatchedStreamingWrite<ErrorT, ElementT> viaStateful() {
+ return new BatchedStreamingWrite<>(
+ bqServices,
+ retryPolicy,
+ failedOutputTag,
+ failedOutputCoder,
+ errorContainer,
+ skipInvalidRows,
+ ignoreUnknownValues,
+ ignoreInsertIds,
+ toTableRow,
+ toFailsafeTableRow,
+ true);
+ }
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ return batchViaStateful
+ ? input.apply(new ViaStateful())
+ : input.apply(new ViaBundleFinalization());
+ }
+
+ private class ViaBundleFinalization
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ PCollectionTuple result =
+ input.apply(
+ ParDo.of(new BatchAndInsertElements())
+ .withOutputTags(mainOutputTag, TupleTagList.of(failedOutputTag)));
+ PCollection<ErrorT> failedInserts = result.get(failedOutputTag);
+ failedInserts.setCoder(failedOutputCoder);
+ return failedInserts;
+ }
+ }
+
+ @VisibleForTesting
+ private class BatchAndInsertElements extends DoFn<KV<String, TableRowInfo<ElementT>>, Void> {
+
+ /** JsonTableRows to accumulate BigQuery rows in order to batch writes. */
+ private transient Map<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> tableRows;
+
+ /** The list of unique ids for each BigQuery table row. */
+ private transient Map<String, List<String>> uniqueIdsForTableRows;
+
+ @Setup
+ public void setup() {
+ // record latency upto 60 seconds in the resolution of 20ms
+ histogram = Histogram.linear(0, 20, 3000);
+ lastReportedSystemClockMillis = System.currentTimeMillis();
+ }
+
+ @Teardown
+ public void teardown() {
+ if (histogram.getTotalCount() > 0) {
+ logPercentiles();
+ histogram.clear();
+ }
+ }
+
+ /** Prepares a target BigQuery table. */
+ @StartBundle
+ public void startBundle() {
+ tableRows = new HashMap<>();
+ uniqueIdsForTableRows = new HashMap<>();
+ }
+
+ /** Accumulates the input into JsonTableRows and uniqueIdsForTableRows. */
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, TableRowInfo<ElementT>> element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane) {
+ String tableSpec = element.getKey();
+ List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows =
+ BigQueryHelpers.getOrCreateMapListValue(tableRows, tableSpec);
+ List<String> uniqueIds =
+ BigQueryHelpers.getOrCreateMapListValue(uniqueIdsForTableRows, tableSpec);
+
+ TableRow tableRow = toTableRow.apply(element.getValue().tableRow);
+ TableRow failsafeTableRow = toFailsafeTableRow.apply(element.getValue().tableRow);
+ rows.add(FailsafeValueInSingleWindow.of(tableRow, timestamp, window, pane, failsafeTableRow));
+ uniqueIds.add(element.getValue().uniqueId);
+ }
+
+ /** Writes the accumulated rows into BigQuery with streaming API. */
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) throws Exception {
+ List<ValueInSingleWindow<ErrorT>> failedInserts = Lists.newArrayList();
+ BigQueryOptions options = context.getPipelineOptions().as(BigQueryOptions.class);
+ for (Map.Entry<String, List<FailsafeValueInSingleWindow<TableRow, TableRow>>> entry :
+ tableRows.entrySet()) {
+ TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
+ flushRows(
+ tableReference,
+ entry.getValue(),
+ uniqueIdsForTableRows.get(entry.getKey()),
+ options,
+ failedInserts);
+ }
+ tableRows.clear();
+ uniqueIdsForTableRows.clear();
+
+ for (ValueInSingleWindow<ErrorT> row : failedInserts) {
+ context.output(failedOutputTag, row.getValue(), row.getTimestamp(), row.getWindow());
+ }
+
+ updateAndLogHistogram(options);
+ }
+ }
+
+ private class ViaStateful
+ extends PTransform<PCollection<KV<String, TableRowInfo<ElementT>>>, PCollection<ErrorT>> {
+ private final Duration BATCH_MAX_BUFFERING_DURATION = Duration.standardSeconds(10);
+
+ @Override
+ public PCollection<ErrorT> expand(PCollection<KV<String, TableRowInfo<ElementT>>> input) {
+ BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+ KvCoder<String, TableRowInfo<ElementT>> inputCoder = (KvCoder) input.getCoder();
+ TableRowInfoCoder<ElementT> valueCoder =
+ (TableRowInfoCoder) inputCoder.getCoderArguments().get(1);
+ PCollectionTuple result =
+ input
Review comment:
Didn't add a global window around here as the existing implementation:
https://github.com/apache/beam/blob/2462fe996626874a6a11f35c00405acc831c8dfd/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java#L281
According to the documentation the global window is for correct access to side inputs but it seems to me that up to this point the dynamic destination has been read. @reuvenlax I might be missing something but let me know if I should keep the global window (and if so where?).
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on a change in pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on a change in pull request #13496:
URL: https://github.com/apache/beam/pull/13496#discussion_r554123442
##########
File path: sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java
##########
@@ -78,4 +79,11 @@
Integer getLatencyLoggingFrequency();
void setLatencyLoggingFrequency(Integer value);
+
+ @Experimental
+ @Description("Whether dynamic sharding is enabled for writing to BigQuery in streaming.")
+ @Default.Boolean(false)
+ Boolean getEnableStreamingAutoSharding();
Review comment:
I was thinking that adding the option to BigQueryOptions might be easier for the users to switch on/off the feature. If it is an option of BigQueryIO users would need to update their code (vs. adding a flag) right? It might not be a concern...
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-768629703
There seems to be an issue with a try/catch: https://ci-beam.apache.org/job/beam_PreCommit_Java_Examples_Dataflow_Commit/12361/console
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] pabloem merged pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
pabloem merged pull request #13496:
URL: https://github.com/apache/beam/pull/13496
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [beam] nehsyc commented on pull request #13496: [BEAM-11408] Integrate BigQuery sink streaming inserts with GroupIntoBatches
Posted by GitBox <gi...@apache.org>.
nehsyc commented on pull request #13496:
URL: https://github.com/apache/beam/pull/13496#issuecomment-761198472
@reuvenlax I also have changes for FILE_LOADS ready in my local branch. If it is ok for you, I can merge those into this PR. Otherwise I will send a follow-up PR.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org