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:
[email protected]
