Merge pull request #25723: #25722 Add option to propagate successful storage-api writes

Author: reuvenlax

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java

@@ -329,7 +329,7 @@ public long getElementByteSize() {
 
   @Override
   public PCollection<KV<K, Iterable<InputT>>> expand(PCollection<KV<K, InputT>> input) {
-
+    Duration allowedLateness = input.getWindowingStrategy().getAllowedLateness();
     checkArgument(
         input.getCoder() instanceof KvCoder,
         "coder specified in the input PCollection is not a KvCoder");
@@ -344,6 +344,7 @@ public PCollection<KV<K, Iterable<InputT>>> expand(PCollection<KV<K, InputT>> in
                 params.getBatchSizeBytes(),
                 weigher,
                 params.getMaxBufferingDuration(),
+                allowedLateness,
                 valueCoder)));
   }
 
@@ -357,12 +358,20 @@ private static class GroupIntoBatchesDoFn<K, InputT>
     @Nullable private final SerializableFunction<InputT, Long> weigher;
     private final Duration maxBufferingDuration;
 
+    private final Duration allowedLateness;
+
     // The following timer is no longer set. We maintain the spec for update compatibility.
     private static final String END_OF_WINDOW_ID = "endOFWindow";
 
     @TimerId(END_OF_WINDOW_ID)
     private final TimerSpec windowTimer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
 
+    // This timer manages the watermark hold if there is no buffering timer.
+    private static final String TIMER_HOLD_ID = "watermarkHold";
+
+    @TimerId(TIMER_HOLD_ID)
+    private final TimerSpec holdTimerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
     // This timer expires when it's time to batch and output the buffered data.
     private static final String END_OF_BUFFERING_ID = "endOfBuffering";
 
@@ -410,11 +419,13 @@ private static class GroupIntoBatchesDoFn<K, InputT>
         long batchSizeBytes,
         @Nullable SerializableFunction<InputT, Long> weigher,
         Duration maxBufferingDuration,
+        Duration allowedLateness,
         Coder<InputT> inputValueCoder) {
       this.batchSize = batchSize;
       this.batchSizeBytes = batchSizeBytes;
       this.weigher = weigher;
       this.maxBufferingDuration = maxBufferingDuration;
+      this.allowedLateness = allowedLateness;
       this.batchSpec = StateSpecs.bag(inputValueCoder);
 
       Combine.BinaryCombineLongFn sumCombineFn =
@@ -452,9 +463,18 @@ public long apply(long left, long right) {
       this.prefetchFrequency = ((batchSize / 5) <= 1) ? Long.MAX_VALUE : (batchSize / 5);
     }
 
+    @Override
+    public Duration getAllowedTimestampSkew() {
+      // This is required since flush is sometimes called from processElement. This is safe because
+      // a watermark hold
+      // will always be set using timer.withOutputTimestamp.
+      return Duration.millis(Long.MAX_VALUE);
+    }
+
     @ProcessElement
     public void processElement(
         @TimerId(END_OF_BUFFERING_ID) Timer bufferingTimer,
+        @TimerId(TIMER_HOLD_ID) Timer holdTimer,
         @StateId(BATCH_ID) BagState<InputT> batch,
         @StateId(NUM_ELEMENTS_IN_BATCH_ID) CombiningState<Long, long[], Long> storedBatchSize,
         @StateId(NUM_BYTES_IN_BATCH_ID) CombiningState<Long, long[], Long> storedBatchSizeBytes,
@@ -473,9 +493,10 @@ public void processElement(
         storedBatchSizeBytes.readLater();
       }
       storedBatchSize.readLater();
-      if (shouldCareAboutMaxBufferingDuration) {
-        minBufferedTs.readLater();
-      }
+      minBufferedTs.readLater();
+
+      // Make sure we always include the current timestamp in the minBufferedTs.
+      minBufferedTs.add(elementTs.getMillis());
 
       LOG.debug("*** BATCH *** Add element for window {} ", window);
       if (shouldCareAboutWeight) {
@@ -505,23 +526,26 @@ public void processElement(
               timerTs,
               minBufferedTs);
           bufferingTimer.clear();
+          holdTimer.clear();
         }
         storedBatchSizeBytes.add(elementWeight);
       }
       batch.add(element.getValue());
       // Blind add is supported with combiningState
       storedBatchSize.add(1L);
+      // Add the timestamp back into minBufferedTs as it might be cleared by flushBatch above.
+      minBufferedTs.add(elementTs.getMillis());
 
       final long num = storedBatchSize.read();
-      if (shouldCareAboutMaxBufferingDuration) {
-        long oldOutputTs =
-            MoreObjects.firstNonNull(
-                minBufferedTs.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
-        minBufferedTs.add(elementTs.getMillis());
-        // If this is the first element in the batch or if the timer's output timestamp needs
-        // modifying, then set a
-        // timer.
-        if (num == 1 || minBufferedTs.read() != oldOutputTs) {
+
+      // If this is the first element in the batch or if the timer's output timestamp needs
+      // modifying, then set a timer.
+      long oldOutputTs =
+          MoreObjects.firstNonNull(
+              minBufferedTs.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
+      boolean needsNewTimer = num == 1 || minBufferedTs.read() != oldOutputTs;
+      if (needsNewTimer) {
+        if (shouldCareAboutMaxBufferingDuration) {
           long targetTs =
               MoreObjects.firstNonNull(
                   timerTs.read(),
@@ -530,6 +554,12 @@ public void processElement(
           bufferingTimer
               .withOutputTimestamp(Instant.ofEpochMilli(minBufferedTs.read()))
               .set(Instant.ofEpochMilli(targetTs));
+        } else {
+          // The only way to hold the watermark is to set a timer. Since there is no buffering
+          // timer, we set a dummy
+          // timer at the end of the window to manage the hold.
+          Instant windowEnd = window.maxTimestamp().plus(allowedLateness);
+          holdTimer.withOutputTimestamp(Instant.ofEpochMilli(minBufferedTs.read())).set(windowEnd);
         }
       }
 
@@ -585,6 +615,11 @@ public void onWindowExpiration(
           receiver, key, batch, storedBatchSize, storedBatchSizeBytes, timerTs, minBufferedTs);
     }
 
+    @OnTimer(TIMER_HOLD_ID)
+    public void onHoldTimer() {
+      // Do nothing. The associated watermark hold will be automatically removed.
+    }
+
     // We no longer set this timer, since OnWindowExpiration takes care of his. However we leave the
     // callback in place
     // for existing jobs that have already set these timers.
@@ -618,7 +653,8 @@ private void flushBatch(
       Iterable<InputT> values = batch.read();
       // When the timer fires, batch state might be empty
       if (!Iterables.isEmpty(values)) {
-        receiver.output(KV.of(key, values));
+        receiver.outputWithTimestamp(
+            KV.of(key, values), Instant.ofEpochMilli(minBufferedTs.read()));
       }
       clearState(batch, storedBatchSize, storedBatchSizeBytes, timerTs, minBufferedTs);
     }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java

@@ -865,6 +865,8 @@ private WriteResult writeResult(Pipeline p, PCollection<TableDestination> succes
         new TupleTag<>("successfulInserts"),
         successfulWrites,
         null,
+        null,
+        null,
         null);
   }
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java

@@ -2094,6 +2094,7 @@ public static <T> Write<T> write() {
         .setAutoSchemaUpdate(false)
         .setDeterministicRecordIdFn(null)
         .setMaxRetryJobs(1000)
+        .setPropagateSuccessfulStorageApiWrites(false)
         .build();
   }
 
@@ -2211,6 +2212,8 @@ public enum Method {
 
     abstract int getNumStorageWriteApiStreams();
 
+    abstract boolean getPropagateSuccessfulStorageApiWrites();
+
     abstract int getMaxFilesPerPartition();
 
     abstract long getMaxBytesPerPartition();
@@ -2306,6 +2309,9 @@ abstract Builder<T> setAvroSchemaFactory(
 
       abstract Builder<T> setNumStorageWriteApiStreams(int numStorageApiStreams);
 
+      abstract Builder<T> setPropagateSuccessfulStorageApiWrites(
+          boolean propagateSuccessfulStorageApiWrites);
+
       abstract Builder<T> setMaxFilesPerPartition(int maxFilesPerPartition);
 
       abstract Builder<T> setMaxBytesPerPartition(long maxBytesPerPartition);
@@ -2763,6 +2769,17 @@ public Write<T> withNumStorageWriteApiStreams(int numStorageWriteApiStreams) {
       return toBuilder().setNumStorageWriteApiStreams(numStorageWriteApiStreams).build();
     }
 
+    /**
+     * If set to true, then all successful writes will be propagated to {@link WriteResult} and
+     * accessible via the {@link WriteResult#getSuccessfulStorageApiInserts} method.
+     */
+    public Write<T> withPropagateSuccessfulStorageApiWrites(
+        boolean propagateSuccessfulStorageApiWrites) {
+      return toBuilder()
+          .setPropagateSuccessfulStorageApiWrites(propagateSuccessfulStorageApiWrites)
+          .build();
+    }
+
     /**
      * Provides a custom location on GCS for storing temporary files to be loaded via BigQuery batch
      * load jobs. See "Usage with templates" in {@link BigQueryIO} documentation for discussion.
@@ -3270,6 +3287,9 @@ private <DestinationT> WriteResult continueExpandTyped(
         checkArgument(
             getSchemaUpdateOptions() == null || getSchemaUpdateOptions().isEmpty(),
             "SchemaUpdateOptions are not supported when method == STREAMING_INSERTS");
+        checkArgument(
+            !getPropagateSuccessfulStorageApiWrites(),
+            "withPropagateSuccessfulStorageApiWrites only supported when using storage api writes.");
 
         RowWriterFactory.TableRowWriterFactory<T, DestinationT> tableRowWriterFactory =
             (RowWriterFactory.TableRowWriterFactory<T, DestinationT>) rowWriterFactory;
@@ -3301,6 +3321,9 @@ private <DestinationT> WriteResult continueExpandTyped(
               rowWriterFactory.getOutputType() == OutputType.AvroGenericRecord,
               "useAvroLogicalTypes can only be set with Avro output.");
         }
+        checkArgument(
+            !getPropagateSuccessfulStorageApiWrites(),
+            "withPropagateSuccessfulStorageApiWrites only supported when using storage api writes.");
 
         // Batch load jobs currently support JSON data insertion only with CSV files
         if (getJsonSchema() != null && getJsonSchema().isAccessible()) {
@@ -3406,7 +3429,8 @@ private <DestinationT> WriteResult continueExpandTyped(
                 method == Method.STORAGE_API_AT_LEAST_ONCE,
                 getAutoSharding(),
                 getAutoSchemaUpdate(),
-                getIgnoreUnknownValues());
+                getIgnoreUnknownValues(),
+                getPropagateSuccessfulStorageApiWrites());
         return input.apply("StorageApiLoads", storageApiLoads);
       } else {
         throw new RuntimeException("Unexpected write method " + method);

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java

@@ -18,20 +18,32 @@
 package org.apache.beam.sdk.io.gcp.bigquery;
 
 import com.google.api.services.bigquery.model.TableRow;
+import com.google.auto.value.AutoValue;
 import com.google.cloud.bigquery.storage.v1.ProtoRows;
 import com.google.protobuf.ByteString;
 import java.util.Iterator;
+import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
-import javax.annotation.Nullable;
+import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
 
 /**
  * Takes in an iterable and batches the results into multiple ProtoRows objects. The splitSize
  * parameter controls how many rows are batched into a single ProtoRows object before we move on to
  * the next one.
  */
-class SplittingIterable implements Iterable<ProtoRows> {
+class SplittingIterable implements Iterable<SplittingIterable.Value> {
+  @AutoValue
+  abstract static class Value {
+    abstract ProtoRows getProtoRows();
+
+    abstract List<Instant> getTimestamps();
+  }
+
   interface ConvertUnknownFields {
     ByteString convert(TableRow tableRow, boolean ignoreUnknownValues)
         throws TableRowToStorageApiProto.SchemaConversionException;
@@ -42,30 +54,34 @@ ByteString convert(TableRow tableRow, boolean ignoreUnknownValues)
 
   private final ConvertUnknownFields unknownFieldsToMessage;
   private final Function<ByteString, TableRow> protoToTableRow;
-  private final BiConsumer<TableRow, String> failedRowsConsumer;
+  private final BiConsumer<TimestampedValue<TableRow>, String> failedRowsConsumer;
   private final boolean autoUpdateSchema;
   private final boolean ignoreUnknownValues;
 
+  private final Instant elementsTimestamp;
+
   public SplittingIterable(
       Iterable<StorageApiWritePayload> underlying,
       long splitSize,
       ConvertUnknownFields unknownFieldsToMessage,
       Function<ByteString, TableRow> protoToTableRow,
-      BiConsumer<TableRow, String> failedRowsConsumer,
+      BiConsumer<TimestampedValue<TableRow>, String> failedRowsConsumer,
       boolean autoUpdateSchema,
-      boolean ignoreUnknownValues) {
+      boolean ignoreUnknownValues,
+      Instant elementsTimestamp) {
     this.underlying = underlying;
     this.splitSize = splitSize;
     this.unknownFieldsToMessage = unknownFieldsToMessage;
     this.protoToTableRow = protoToTableRow;
     this.failedRowsConsumer = failedRowsConsumer;
     this.autoUpdateSchema = autoUpdateSchema;
     this.ignoreUnknownValues = ignoreUnknownValues;
+    this.elementsTimestamp = elementsTimestamp;
   }
 
   @Override
-  public Iterator<ProtoRows> iterator() {
-    return new Iterator<ProtoRows>() {
+  public Iterator<Value> iterator() {
+    return new Iterator<Value>() {
       final Iterator<StorageApiWritePayload> underlyingIterator = underlying.iterator();
 
       @Override
@@ -74,11 +90,12 @@ public boolean hasNext() {
       }
 
       @Override
-      public ProtoRows next() {
+      public Value next() {
         if (!hasNext()) {
           throw new NoSuchElementException();
         }
 
+        List<Instant> timestamps = Lists.newArrayList();
         ProtoRows.Builder inserts = ProtoRows.newBuilder();
         long bytesSize = 0;
         while (underlyingIterator.hasNext()) {
@@ -107,7 +124,11 @@ public ProtoRows next() {
                   // 24926 is fixed, we need to merge the unknownFields back into the main row
                   // before outputting to the
                   // failed-rows consumer.
-                  failedRowsConsumer.accept(tableRow, e.toString());
+                  Instant timestamp = payload.getTimestamp();
+                  if (timestamp == null) {
+                    timestamp = elementsTimestamp;
+                  }
+                  failedRowsConsumer.accept(TimestampedValue.of(tableRow, timestamp), e.toString());
                   continue;
                 }
               }
@@ -116,12 +137,17 @@ public ProtoRows next() {
             }
           }
           inserts.addSerializedRows(byteString);
+          Instant timestamp = payload.getTimestamp();
+          if (timestamp == null) {
+            timestamp = elementsTimestamp;
+          }
+          timestamps.add(timestamp);
           bytesSize += byteString.size();
           if (bytesSize > splitSize) {
             break;
           }
         }
-        return inserts.build();
+        return new AutoValue_SplittingIterable_Value(inserts.build(), timestamps);
       }
     };
   }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java

@@ -33,6 +33,7 @@
 import org.apache.beam.sdk.values.TupleTagList;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
 
 /**
  * A transform that converts messages to protocol buffers in preparation for writing to BigQuery.
@@ -129,14 +130,16 @@ public void processElement(
         ProcessContext c,
         PipelineOptions pipelineOptions,
         @Element KV<DestinationT, ElementT> element,
+        @Timestamp Instant timestamp,
         MultiOutputReceiver o)
         throws Exception {
       dynamicDestinations.setSideInputAccessorFromProcessContext(c);
       MessageConverter<ElementT> messageConverter =
           messageConverters.get(
               element.getKey(), dynamicDestinations, getDatasetService(pipelineOptions));
       try {
-        StorageApiWritePayload payload = messageConverter.toMessage(element.getValue());
+        StorageApiWritePayload payload =
+            messageConverter.toMessage(element.getValue()).withTimestamp(timestamp);
         o.get(successfulWritesTag).output(KV.of(element.getKey(), payload));
       } catch (TableRowToStorageApiProto.SchemaConversionException e) {
         TableRow tableRow = messageConverter.toTableRow(element.getValue());

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java

@@ -62,9 +62,9 @@ public TableSchema getTableSchema() {
 
     @Override
     @SuppressWarnings("nullness")
-    public StorageApiWritePayload toMessage(T element) {
+    public StorageApiWritePayload toMessage(T element) throws Exception {
       Message msg = BeamRowToStorageApiProto.messageFromBeamRow(descriptor, toRow.apply(element));
-      return new AutoValue_StorageApiWritePayload(msg.toByteArray(), null);
+      return StorageApiWritePayload.of(msg.toByteArray(), null);
     }
 
     @Override