foo

Author: reuvenlax

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

@@ -66,6 +66,7 @@ class BufferMismatchedRows<DestinationT extends @NonNull Object, ElementT>
   // StorageApiWriteRecordsInconsistent.
   private final TupleTag<KV<String, String>> finalizeTag = new TupleTag<>("finalizeTag");
   private static final int NUM_DEFAULT_SHARDS = 20;
+  private static final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1);
 
   public BufferMismatchedRows(
       Coder<BigQueryStorageApiInsertError> failedRowsCoder,
@@ -155,8 +156,6 @@ class BufferingDoFn
     private final Counter rowsSentToFailedRowsCollection =
         Metrics.counter(BufferMismatchedRows.BufferingDoFn.class, "rowsSentToFailedRowsCollection");
 
-    private final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1);
-
     public BufferingDoFn(
         StorageApiWriteUnshardedRecords.WriteRecordsDoFn<DestinationT, ElementT> writeDoFn) {
       this.writeDoFn = writeDoFn;

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

@@ -46,10 +46,12 @@ static MismatchedRow of(StorageApiWritePayload payload, org.joda.time.Instant de
   // inside of state
   // variables (mostly involving the Dataflow runner and update). Therefore we use a custom coder.
   static class Coder extends CustomCoder<MismatchedRow> {
-    static final ByteArrayCoder byteArrayCoder = ByteArrayCoder.of();
-    static final InstantCoder instantCoder = InstantCoder.of();
-    static final NullableCoder<byte[]> nullableByteArrayCoder = NullableCoder.of(byteArrayCoder);
-    static final NullableCoder<Instant> nullableInstantCoder = NullableCoder.of(InstantCoder.of());
+    static final ByteArrayCoder BYTE_ARRAY_CODER = ByteArrayCoder.of();
+    static final InstantCoder INSTANT_CODER = InstantCoder.of();
+    static final NullableCoder<byte[]> NULLABLE_BYTE_ARRAY_CODER =
+        NullableCoder.of(BYTE_ARRAY_CODER);
+    static final NullableCoder<Instant> NULLABLE_INSTANT_CODER =
+        NullableCoder.of(InstantCoder.of());
 
     static MismatchedRow.Coder of() {
       return new Coder();
@@ -58,22 +60,22 @@ static MismatchedRow.Coder of() {
     @Override
     public void encode(MismatchedRow value, OutputStream outStream)
         throws CoderException, IOException {
-      byteArrayCoder.encode(value.getPayload().getPayload(), outStream);
-      nullableByteArrayCoder.encode(value.getPayload().getUnknownFieldsPayload(), outStream);
-      nullableInstantCoder.encode(value.getPayload().getTimestamp(), outStream);
-      nullableByteArrayCoder.encode(value.getPayload().getFailsafeTableRowPayload(), outStream);
-      nullableByteArrayCoder.encode(value.getPayload().getSchemaHash(), outStream);
-      instantCoder.encode(value.getDeadline(), outStream);
+      BYTE_ARRAY_CODER.encode(value.getPayload().getPayload(), outStream);
+      NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getUnknownFieldsPayload(), outStream);
+      NULLABLE_INSTANT_CODER.encode(value.getPayload().getTimestamp(), outStream);
+      NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getFailsafeTableRowPayload(), outStream);
+      NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getSchemaHash(), outStream);
+      INSTANT_CODER.encode(value.getDeadline(), outStream);
     }
 
     @Override
     public MismatchedRow decode(InputStream inStream) throws CoderException, IOException {
-      byte[] innerPayload = byteArrayCoder.decode(inStream);
-      byte @Nullable [] unknownFieldsPayload = nullableByteArrayCoder.decode(inStream);
-      @Nullable Instant timestamp = nullableInstantCoder.decode(inStream);
-      byte @Nullable [] failsafeTableRowPayload = nullableByteArrayCoder.decode(inStream);
-      byte @Nullable [] schemaHash = nullableByteArrayCoder.decode(inStream);
-      Instant deadline = instantCoder.decode(inStream);
+      byte[] innerPayload = BYTE_ARRAY_CODER.decode(inStream);
+      byte @Nullable [] unknownFieldsPayload = NULLABLE_BYTE_ARRAY_CODER.decode(inStream);
+      @Nullable Instant timestamp = NULLABLE_INSTANT_CODER.decode(inStream);
+      byte @Nullable [] failsafeTableRowPayload = NULLABLE_BYTE_ARRAY_CODER.decode(inStream);
+      byte @Nullable [] schemaHash = NULLABLE_BYTE_ARRAY_CODER.decode(inStream);
+      Instant deadline = INSTANT_CODER.decode(inStream);
 
       StorageApiWritePayload payload =
           StorageApiWritePayload.of(

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

@@ -581,6 +581,11 @@ long flush(
 
         pendingMessages.clear();
 
+        if (rowsToProcess.getProtoRows().getSerializedRowsCount() == 0) {
+          // All rows either failed or were buffered with schema mismatches.
+          return 0;
+        }
+
         // Handle the case where the request is too large.
         if (rowsToProcess.getProtoRows().getSerializedSize() >= maxRequestSize) {
           if (rowsToProcess.getProtoRows().getSerializedRowsCount() > 1) {

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

@@ -126,6 +126,7 @@ public class StorageApiWritesShardedRecords<DestinationT extends @NonNull Object
         PCollectionTuple> {
   private static final Logger LOG = LoggerFactory.getLogger(StorageApiWritesShardedRecords.class);
   private static final Duration DEFAULT_STREAM_IDLE_TIME = Duration.standardHours(1);
+  private static final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1);
 
   private final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
   private final CreateDisposition createDisposition;
@@ -373,8 +374,6 @@ class WriteRecordsDoFn
     @TimerId("idleTimer")
     private final TimerSpec idleTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
 
-    private final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1);
-
     private final Duration streamIdleTime;
     private final long splitSize;
     private final long maxRequestSize;
@@ -800,9 +799,6 @@ public void process(
         @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp,
         final MultiOutputReceiver o)
         throws Exception {
-      Instant now = Instant.now();
-      System.err.println("PROCESS " + now);
-
       dynamicDestinations.setSideInputAccessorFromProcessContext(c);
       TableDestination tableDestination =
           destinations.computeIfAbsent(
@@ -861,7 +857,6 @@ public void process(
           idleTimer,
           o,
           processMismatchedRows);
-      System.err.println("PROCESS DONE AFTER " + java.time.Duration.between(now, Instant.now()));
     }
 
     private void processPayloads(
@@ -1259,7 +1254,6 @@ public void onMismatchedRowsTimer(
         @StateId("currentMismatchedRowTimerValue") ValueState<Long> currentTimerValue,
         @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp)
         throws Exception {
-      System.err.println("RETRY TIMER " + Instant.now());
       mismatchedRowsBag.readLater();
       currentTimerValue.readLater();
       minPendingTimestamp.readLater();

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

@@ -18,7 +18,6 @@
 package org.apache.beam.sdk.io.gcp.bigquery;
 
 import com.google.api.services.bigquery.model.TableCell;
-import com.google.api.services.bigquery.model.TableRow;
 import com.google.cloud.bigquery.storage.v1.BigQuerySchemaUtil;
 import com.google.cloud.bigquery.storage.v1.TableFieldSchema;
 import com.google.cloud.bigquery.storage.v1.TableSchema;
@@ -317,7 +316,8 @@ private static boolean hasUnknownFields(Message message) {
   }
 
   public static boolean missingUnknownField(
-      TableRow unknownFields, ThrowingSupplier<Descriptors.Descriptor> schemaDescriptor)
+      AbstractMap<String, Object> unknownFields,
+      ThrowingSupplier<Descriptors.Descriptor> schemaDescriptor)
       throws Exception {
     @Nullable Object fValue = unknownFields.get("f");
     if (fValue instanceof List) {
@@ -392,8 +392,9 @@ private static boolean missingUnknownFieldObject(
         }
       }
       return false;
-    } else if (value instanceof TableRow) {
-      return missingUnknownField((TableRow) value, fieldDescriptor.getMessageType());
+    } else if (value instanceof AbstractMap) {
+      return missingUnknownField(
+          (AbstractMap<String, Object>) value, fieldDescriptor.getMessageType());
     }
     return false;
   }