Add Dynamic Schema support to StorageWriteToBigQuery

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json

@@ -1,5 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 16,
-  "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
-}
+  "modification": 21
+}

.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json

@@ -1,5 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 4,
-  "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
+  "modification": 11
 }

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

@@ -23,29 +23,47 @@
 
 import com.google.api.services.bigquery.model.Clustering;
 import com.google.api.services.bigquery.model.TableConstraints;
+import com.google.api.services.bigquery.model.TableFieldSchema;
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.api.services.bigquery.model.TableSchema;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.NullableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.extensions.avro.schemas.utils.AvroUtils;
 import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils;
 import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinations;
 import org.apache.beam.sdk.io.gcp.bigquery.TableDestination;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.util.RowFilter;
 import org.apache.beam.sdk.util.RowStringInterpolator;
+import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 @Internal
-public class PortableBigQueryDestinations extends DynamicDestinations<Row, String> {
+public class PortableBigQueryDestinations
+    extends DynamicDestinations<Row, KV<String, @Nullable String>> {
   public static final String DESTINATION = "destination";
   public static final String RECORD = "record";
+  public static final String SCHEMA = "schema";
+
+  private static final ConcurrentHashMap<String, TableSchema> JSON_SCHEMA_CACHE =
+      new ConcurrentHashMap<>();
+
   private @MonotonicNonNull RowStringInterpolator interpolator = null;
   private final @Nullable List<String> primaryKey;
   private final RowFilter rowFilter;
@@ -73,44 +91,126 @@ public PortableBigQueryDestinations(Schema rowSchema, BigQueryWriteConfiguration
   }
 
   @Override
-  public String getDestination(@Nullable ValueInSingleWindow<Row> element) {
+  public KV<String, @Nullable String> getDestination(@Nullable ValueInSingleWindow<Row> element) {
     if (interpolator != null) {
-      return interpolator.interpolate(checkArgumentNotNull(element));
+      return KV.of(interpolator.interpolate(checkArgumentNotNull(element)), null);
+    }
+    Row row = checkStateNotNull(checkStateNotNull(element).getValue());
+    String destination = checkStateNotNull(row.getString(DESTINATION));
+    if (row.getSchema().hasField(SCHEMA)) {
+      @Nullable String schemaJson = row.getString(SCHEMA);
+      return KV.of(destination, schemaJson == null ? "" : schemaJson);
     }
-    return checkStateNotNull(checkStateNotNull(element).getValue().getString(DESTINATION));
+    return KV.of(destination, null);
   }
 
   @Override
-  public TableDestination getTable(String destination) {
+  public Coder<KV<String, @Nullable String>> getDestinationCoder() {
+    return KvCoder.of(StringUtf8Coder.of(), NullableCoder.of(StringUtf8Coder.of()));
+  }
 
+  @Override
+  public TableDestination getTable(KV<String, @Nullable String> destination) {
+    String tableSpec = destination.getKey();
     if (clusteringFields != null && !clusteringFields.isEmpty()) {
       Clustering clustering = new Clustering().setFields(clusteringFields);
-      return new TableDestination(destination, null, null, clustering);
+      return new TableDestination(tableSpec, null, null, clustering);
     }
-    return new TableDestination(destination, null);
+    return new TableDestination(tableSpec, null);
   }
 
   @Override
-  public @Nullable TableSchema getSchema(String destination) {
+  public @Nullable TableSchema getSchema(KV<String, @Nullable String> destination) {
+    @Nullable String schemaJson = destination.getValue();
+    if (schemaJson != null) {
+      if (schemaJson.isEmpty()) {
+        return null;
+      }
+      return parseTableSchema(schemaJson);
+    }
     return BigQueryUtils.toTableSchema(rowFilter.outputSchema());
   }
 
   @Override
-  public @Nullable TableConstraints getTableConstraints(String destination) {
+  public @Nullable TableConstraints getTableConstraints(KV<String, @Nullable String> destination) {
     if (primaryKey != null) {
       return new TableConstraints()
           .setPrimaryKey(new TableConstraints.PrimaryKey().setColumns(primaryKey));
     }
     return null;
   }
 
+  private static TableSchema parseTableSchema(String schemaJson) {
+    return JSON_SCHEMA_CACHE.computeIfAbsent(
+        schemaJson, json -> BigQueryHelpers.fromJsonString(json, TableSchema.class));
+  }
+
+  @VisibleForTesting
+  static TableRow filterTableRowBySchema(
+      Map<String, Object> tableRow, @Nullable List<TableFieldSchema> fields) {
+    if (fields == null || fields.isEmpty()) {
+      TableRow copy = new TableRow();
+      copy.putAll(tableRow);
+      return copy;
+    }
+    TableRow filtered = new TableRow();
+    for (TableFieldSchema field : fields) {
+      String fieldName = field.getName();
+      @Nullable Object value = null;
+      if (tableRow.containsKey(fieldName)) {
+        value = tableRow.get(fieldName);
+      } else {
+        for (Map.Entry<String, Object> entry : tableRow.entrySet()) {
+          if (entry.getKey().equalsIgnoreCase(fieldName)) {
+            value = entry.getValue();
+            break;
+          }
+        }
+      }
+      if (value == null) {
+        continue;
+      }
+      List<TableFieldSchema> subfields = field.getFields();
+      if (subfields != null && !subfields.isEmpty()) {
+        if ("REPEATED".equalsIgnoreCase(field.getMode()) && value instanceof Iterable) {
+          List<@Nullable Object> filteredList = new ArrayList<>();
+          for (Object item : (Iterable<?>) value) {
+            if (item instanceof Map) {
+              @SuppressWarnings("unchecked")
+              Map<String, Object> mapItem = (Map<String, Object>) item;
+              filteredList.add(filterTableRowBySchema(mapItem, subfields));
+            } else if (item != null) {
+              filteredList.add(item);
+            }
+          }
+          value = filteredList;
+        } else if (value instanceof Map) {
+          @SuppressWarnings("unchecked")
+          Map<String, Object> mapValue = (Map<String, Object>) value;
+          value = filterTableRowBySchema(mapValue, subfields);
+        }
+      }
+      filtered.set(fieldName, value);
+    }
+    return filtered;
+  }
+
   public SerializableFunction<Row, TableRow> getFilterFormatFunction(boolean fetchNestedRecord) {
     return row -> {
+      @Nullable String schemaJson = null;
       if (fetchNestedRecord) {
+        if (row.getSchema().hasField(SCHEMA)) {
+          schemaJson = row.getString(SCHEMA);
+        }
         row = checkStateNotNull(row.getRow(RECORD));
       }
       Row filtered = rowFilter.filter(row);
-      return BigQueryUtils.toTableRow(filtered);
+      TableRow tableRow = BigQueryUtils.toTableRow(filtered);
+      if (schemaJson != null && !schemaJson.isEmpty()) {
+        TableSchema tableSchema = parseTableSchema(schemaJson);
+        tableRow = filterTableRowBySchema(tableRow, tableSchema.getFields());
+      }
+      return tableRow;
     };
   }
 
@@ -122,7 +222,16 @@ public SerializableFunction<AvroWriteRequest<Row>, GenericRecord> getAvroFilterF
         row = checkStateNotNull(row.getRow(RECORD));
       }
       Row filtered = rowFilter.filter(row);
-      return AvroUtils.toGenericRecord(filtered, request.getSchema());
+      org.apache.avro.Schema avroSchema = request.getSchema();
+      if (avroSchema != null
+          && avroSchema.getFields().size() != filtered.getSchema().getFieldCount()) {
+        List<String> fieldNames =
+            avroSchema.getFields().stream()
+                .map(org.apache.avro.Schema.Field::name)
+                .collect(Collectors.toList());
+        filtered = new RowFilter(filtered.getSchema()).keep(fieldNames).filter(filtered);
+      }
+      return AvroUtils.toGenericRecord(filtered, avroSchema);
     };
   }
 }

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

@@ -26,7 +26,9 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
+import com.google.api.services.bigquery.model.TableFieldSchema;
 import com.google.api.services.bigquery.model.TableRow;
+import com.google.api.services.bigquery.model.TableSchema;
 import java.io.Serializable;
 import java.time.LocalDateTime;
 import java.util.ArrayList;
@@ -505,4 +507,122 @@ public void testManagedChoosesStorageApiForUnboundedWrites() {
     assertThat(writeTransformProto.size(), greaterThan(0));
     p.enableAbandonedNodeEnforcement(false);
   }
+
+  @Test
+  public void testDynamicDestinationsWithDynamicSchemas() throws Exception {
+    Schema unionRecordSchema =
+        Schema.builder()
+            .addNullableStringField("name")
+            .addNullableInt64Field("number")
+            .addNullableDoubleField("score")
+            .build();
+
+    Schema wrapperSchema =
+        Schema.builder()
+            .addStringField(DESTINATION)
+            .addStringField(PortableBigQueryDestinations.SCHEMA)
+            .addRowField(RECORD, unionRecordSchema)
+            .build();
+
+    String schemaJson1 =
+        BigQueryHelpers.toJsonString(
+            new TableSchema()
+                .setFields(
+                    Arrays.asList(
+                        new TableFieldSchema()
+                            .setName("name")
+                            .setType("STRING")
+                            .setMode("NULLABLE"),
+                        new TableFieldSchema()
+                            .setName("number")
+                            .setType("INTEGER")
+                            .setMode("NULLABLE"))));
+
+    String schemaJson2 =
+        BigQueryHelpers.toJsonString(
+            new TableSchema()
+                .setFields(
+                    Arrays.asList(
+                        new TableFieldSchema()
+                            .setName("name")
+                            .setType("STRING")
+                            .setMode("NULLABLE"),
+                        new TableFieldSchema()
+                            .setName("score")
+                            .setType("FLOAT")
+                            .setMode("NULLABLE"))));
+
+    Row row1 =
+        Row.withSchema(wrapperSchema)
+            .withFieldValue(DESTINATION, "project:dataset.dyn_schema_table_1")
+            .withFieldValue(PortableBigQueryDestinations.SCHEMA, schemaJson1)
+            .withFieldValue(
+                RECORD,
+                Row.withSchema(unionRecordSchema)
+                    .withFieldValue("name", "alice")
+                    .withFieldValue("number", 10L)
+                    .withFieldValue("score", null)
+                    .build())
+            .build();
+
+    Row row2 =
+        Row.withSchema(wrapperSchema)
+            .withFieldValue(DESTINATION, "project:dataset.dyn_schema_table_2")
+            .withFieldValue(PortableBigQueryDestinations.SCHEMA, schemaJson2)
+            .withFieldValue(
+                RECORD,
+                Row.withSchema(unionRecordSchema)
+                    .withFieldValue("name", "bob")
+                    .withFieldValue("number", null)
+                    .withFieldValue("score", 95.5)
+                    .build())
+            .build();
+
+    BigQueryWriteConfiguration config =
+        BigQueryWriteConfiguration.builder()
+            .setTable(BigQueryWriteConfiguration.DYNAMIC_DESTINATIONS)
+            .build();
+
+    BigQueryStorageWriteApiSchemaTransformProvider provider =
+        new BigQueryStorageWriteApiSchemaTransformProvider();
+    BigQueryStorageWriteApiSchemaTransform writeTransform =
+        (BigQueryStorageWriteApiSchemaTransform) provider.from(config);
+    writeTransform.setBigQueryServices(fakeBigQueryServices);
+
+    PCollection<Row> inputRows =
+        p.apply(Create.of(Arrays.asList(row1, row2)).withRowSchema(wrapperSchema));
+    PCollectionRowTuple.of("input", inputRows).apply(writeTransform);
+
+    p.run().waitUntilFinish();
+
+    // Verify table 1 was created with ONLY ['name', 'number'] schema and has the expected row
+    com.google.api.services.bigquery.model.Table table1 =
+        fakeDatasetService.getTable(
+            BigQueryHelpers.parseTableSpec("project:dataset.dyn_schema_table_1"));
+    assertNotNull(table1);
+    assertEquals(2, table1.getSchema().getFields().size());
+    assertEquals("name", table1.getSchema().getFields().get(0).getName());
+    assertEquals("number", table1.getSchema().getFields().get(1).getName());
+
+    List<TableRow> table1Rows =
+        fakeDatasetService.getAllRows("project", "dataset", "dyn_schema_table_1");
+    assertEquals(1, table1Rows.size());
+    assertEquals("alice", table1Rows.get(0).get("name"));
+    assertEquals("10", table1Rows.get(0).get("number").toString());
+
+    // Verify table 2 was created with ONLY ['name', 'score'] schema and has the expected row
+    com.google.api.services.bigquery.model.Table table2 =
+        fakeDatasetService.getTable(
+            BigQueryHelpers.parseTableSpec("project:dataset.dyn_schema_table_2"));
+    assertNotNull(table2);
+    assertEquals(2, table2.getSchema().getFields().size());
+    assertEquals("name", table2.getSchema().getFields().get(0).getName());
+    assertEquals("score", table2.getSchema().getFields().get(1).getName());
+
+    List<TableRow> table2Rows =
+        fakeDatasetService.getAllRows("project", "dataset", "dyn_schema_table_2");
+    assertEquals(1, table2Rows.size());
+    assertEquals("bob", table2Rows.get(0).get("name"));
+    assertEquals(95.5, Double.parseDouble(table2Rows.get(0).get("score").toString()), 0.001);
+  }
 }