Add Dynamic Schema support to StorageWriteToBigQuery

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 16
+  "modification": 21
 }

.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 2
+  "modification": 7
 }

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

@@ -168,9 +168,11 @@ public static DynamicMessage messageFromBeamRow(
     for (int i = 0; i < row.getFieldCount(); ++i) {
       Field beamField = beamSchema.getField(i);
       FieldDescriptor fieldDescriptor =
-          Preconditions.checkNotNull(
-              descriptor.findFieldByName(beamField.getName().toLowerCase()),
-              beamField.getName().toLowerCase());
+          descriptor.findFieldByName(beamField.getName().toLowerCase());
+      if (fieldDescriptor == null) {
+        // Field in the union row is not present in the destination table's descriptor; skip it.
+        continue;
+      }
       @Nullable Object value = messageValueFromRowValue(fieldDescriptor, beamField, i, row);
       if (value != null) {
         builder.setField(fieldDescriptor, value);
@@ -330,7 +332,8 @@ private static Object toProtoValue(
       FieldDescriptor fieldDescriptor, FieldType beamFieldType, Object value) {
     switch (beamFieldType.getTypeName()) {
       case ROW:
-        return messageFromBeamRow(fieldDescriptor.getMessageType(), (Row) value, null, -1);
+        return messageFromBeamRow(
+            fieldDescriptor.getMessageType(), (Row) value, null, (String) null);
       case ARRAY:
       case ITERABLE:
         Iterable<Object> iterable = (Iterable<Object>) value;
@@ -419,14 +422,21 @@ static Object mapEntryToProtoValue(
     DynamicMessage.Builder builder = DynamicMessage.newBuilder(descriptor);
     FieldDescriptor keyFieldDescriptor =
         Preconditions.checkNotNull(descriptor.findFieldByName("key"));
-    @Nullable Object key = toProtoValue(keyFieldDescriptor, keyFieldType, entryValue.getKey());
+    @Nullable
+    Object key =
+        entryValue.getKey() != null
+            ? toProtoValue(keyFieldDescriptor, keyFieldType, entryValue.getKey())
+            : null;
     if (key != null) {
       builder.setField(keyFieldDescriptor, key);
     }
     FieldDescriptor valueFieldDescriptor =
         Preconditions.checkNotNull(descriptor.findFieldByName("value"));
     @Nullable
-    Object value = toProtoValue(valueFieldDescriptor, valueFieldType, entryValue.getValue());
+    Object value =
+        entryValue.getValue() != null
+            ? toProtoValue(valueFieldDescriptor, valueFieldType, entryValue.getValue())
+            : null;
     if (value != null) {
       builder.setField(valueFieldDescriptor, value);
     }

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

@@ -3990,11 +3990,13 @@ private <DestinationT> WriteResult expandTyped(
         // TODO: If the user provided a schema, we should use that. There are things that can be
         // specified in a
         // BQ schema that don't have exact matches in a Beam schema (e.g. GEOGRAPHY types).
-        TableSchema tableSchema = BigQueryUtils.toTableSchema(input.getSchema());
-        dynamicDestinations =
-            new ConstantSchemaDestinations<>(
-                dynamicDestinations,
-                StaticValueProvider.of(BigQueryHelpers.toJsonString(tableSchema)));
+        if (!hasSchema) {
+          TableSchema tableSchema = BigQueryUtils.toTableSchema(input.getSchema());
+          dynamicDestinations =
+              new ConstantSchemaDestinations<>(
+                  dynamicDestinations,
+                  StaticValueProvider.of(BigQueryHelpers.toJsonString(tableSchema)));
+        }
       } else if (writeProtoClass != null) {
         if (!hasSchema) {
           try {
@@ -4467,6 +4469,7 @@ static void clearStaticCaches() throws ExecutionException, InterruptedException
     CreateTables.clearCreatedTables();
     TwoLevelMessageConverterCache.clear();
     StorageApiDynamicDestinationsTableRow.clearSchemaCache();
+    StorageApiDynamicDestinationsBeamRow.clearSchemaCache();
     StorageApiWriteUnshardedRecords.clearCache();
     StorageApiWritesShardedRecords.clearCache();
   }

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

@@ -17,13 +17,16 @@
  */
 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 com.google.cloud.bigquery.storage.v1.TableSchema;
 import com.google.protobuf.DescriptorProtos;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Message;
 import java.io.IOException;
+import java.util.concurrent.ExecutionException;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.WriteStreamService;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.transforms.SerializableBiFunction;
@@ -32,10 +35,22 @@
 import org.apache.beam.sdk.values.Row;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** Storage API DynamicDestinations used when the input is a Beam Row. */
 class StorageApiDynamicDestinationsBeamRow<T, DestinationT extends @NonNull Object>
     extends StorageApiDynamicDestinations<T, DestinationT> {
+  private static final Logger LOG =
+      LoggerFactory.getLogger(StorageApiDynamicDestinationsBeamRow.class);
+  private static final TableSchemaCache SCHEMA_CACHE =
+      new TableSchemaCache(Duration.standardSeconds(1));
+
+  static {
+    SCHEMA_CACHE.start();
+  }
+
   private final TableSchema tableSchema;
   private final SerializableFunction<T, Row> toRow;
   private final @Nullable SerializableBiFunction<
@@ -59,21 +74,56 @@ class StorageApiDynamicDestinationsBeamRow<T, DestinationT extends @NonNull Obje
     this.usesCdc = usesCdc;
   }
 
+  static void clearSchemaCache() throws ExecutionException, InterruptedException {
+    SCHEMA_CACHE.clear();
+  }
+
   @Override
   public MessageConverter<T> getMessageConverter(
       DestinationT destination,
       PipelineOptions pipelineOptions,
-      DatasetService datasetService,
-      BigQueryServices.WriteStreamService writeStreamService)
+      @Nullable DatasetService datasetService,
+      @Nullable WriteStreamService writeStreamService)
       throws Exception {
-    return new BeamRowConverter();
+    TableSchema destinationProtoSchema = null;
+    TableDestination tableDestination = getTable(destination);
+    TableReference tableReference =
+        tableDestination != null ? tableDestination.getTableReference() : null;
+
+    if (tableReference != null && datasetService != null) {
+      try {
+        com.google.api.services.bigquery.model.TableSchema fetchedSchema =
+            SCHEMA_CACHE.getSchema(tableReference, datasetService);
+        if (fetchedSchema != null) {
+          destinationProtoSchema =
+              TableRowToStorageApiProto.schemaToProtoTableSchema(fetchedSchema);
+        }
+      } catch (Exception e) {
+        LOG.warn("Could not fetch schema from BigQuery for table {}", tableReference, e);
+      }
+    }
+
+    if (destinationProtoSchema == null) {
+      com.google.api.services.bigquery.model.TableSchema destSchema = getSchema(destination);
+      if (destSchema != null) {
+        destinationProtoSchema = TableRowToStorageApiProto.schemaToProtoTableSchema(destSchema);
+      }
+    }
+
+    if (destinationProtoSchema == null) {
+      destinationProtoSchema = this.tableSchema;
+    }
+
+    return new BeamRowConverter(destinationProtoSchema);
   }
 
   class BeamRowConverter implements MessageConverter<T> {
+    final TableSchema tableSchema;
     final Descriptor descriptor;
     final @Nullable Descriptor cdcDescriptor;
 
-    BeamRowConverter() throws Exception {
+    BeamRowConverter(TableSchema tableSchema) throws Exception {
+      this.tableSchema = tableSchema;
       this.descriptor =
           TableRowToStorageApiProto.getDescriptorFromTableSchema(tableSchema, true, false);
       if (usesCdc) {
@@ -130,5 +180,5 @@ public TableRow toFailsafeTableRow(T element) {
         return BigQueryUtils.toTableRow(toRow.apply(element));
       }
     }
-  };
+  }
 }

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

@@ -38,9 +38,13 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class StorageApiDynamicDestinationsTableRow<T, DestinationT extends @NonNull Object>
     extends StorageApiDynamicDestinations<T, DestinationT> {
+  private static final Logger LOG =
+      LoggerFactory.getLogger(StorageApiDynamicDestinationsTableRow.class);
   private final BigQueryIO.TableRowFormatFunction<T> formatFunction;
   private final BigQueryIO.@Nullable TableRowFormatFunction<T> formatRecordOnFailureFunction;
 
@@ -94,8 +98,23 @@ public MessageConverter<T> getMessageConverter(
           }
         };
 
+    TableSchema schemaToUse = null;
+    TableDestination tableDestination = getTable(destination);
+    TableReference tableReference =
+        tableDestination != null ? tableDestination.getTableReference() : null;
+    if (tableReference != null && datasetService != null) {
+      try {
+        schemaToUse = SCHEMA_CACHE.getSchema(tableReference, datasetService);
+      } catch (Exception e) {
+        LOG.debug("Could not fetch schema from BigQuery for destination {}", destination, e);
+      }
+    }
+    if (schemaToUse == null) {
+      schemaToUse = getSchema(destination);
+    }
+
     return schemaUpdateOptions.isEmpty()
-        ? getConverter.apply(getSchema(destination))
+        ? getConverter.apply(schemaToUse)
         : new SchemaUpgradingTableRowConverter(
             getConverter, options, datasetService, writeStreamService);
   }

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

@@ -888,7 +888,7 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto)
         if (unknownFields != null) {
           unknownFields.set(key, entry.getValue());
         }
-        if (ignoreUnknownValues) {
+        if (ignoreUnknownValues || entry.getValue() == null) {
           continue;
         } else {
           String prefix = schemaInformation.getFullName();

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

@@ -53,6 +53,7 @@
 import org.apache.beam.sdk.values.TypeDescriptors;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
 
 /**
@@ -241,20 +242,36 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
             Schema.of(
                 Field.of("failed_row", FieldType.row(inputSchema)),
                 Field.of("error_message", FieldType.STRING));
+        boolean isDynamicDestinations = configuration.getTable().equals(DYNAMIC_DESTINATIONS);
+        @Nullable
+        Schema recordSchema =
+            isDynamicDestinations ? inputSchema.getField(RECORD).getType().getRowSchema() : null;
         PCollection<Row> failedRowsWithErrors =
             result
                 .getFailedStorageApiInserts()
                 .apply(
                     "Construct failed rows and errors",
                     MapElements.into(TypeDescriptors.rows())
                         .via(
-                            (storageError) ->
-                                Row.withSchema(errorSchema)
-                                    .withFieldValue("error_message", storageError.getErrorMessage())
-                                    .withFieldValue(
-                                        "failed_row",
-                                        BigQueryUtils.toBeamRow(inputSchema, storageError.getRow()))
-                                    .build()))
+                            (storageError) -> {
+                              Row failedRow;
+                              if (isDynamicDestinations && recordSchema != null) {
+                                Row recordRow =
+                                    BigQueryUtils.toBeamRow(recordSchema, storageError.getRow());
+                                failedRow =
+                                    Row.withSchema(inputSchema)
+                                        .withFieldValue(DESTINATION, "")
+                                        .withFieldValue(RECORD, recordRow)
+                                        .build();
+                              } else {
+                                failedRow =
+                                    BigQueryUtils.toBeamRow(inputSchema, storageError.getRow());
+                              }
+                              return Row.withSchema(errorSchema)
+                                  .withFieldValue("error_message", storageError.getErrorMessage())
+                                  .withFieldValue("failed_row", failedRow)
+                                  .build();
+                            }))
                 .setRowSchema(errorSchema);
         return PCollectionRowTuple.of("post_write", postWrite)
             .and(configuration.getErrorHandling().getOutput(), failedRowsWithErrors);