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