diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json index 83346d34aee0..48379c82f94c 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json +++ b/.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 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json index b26833333238..d6a91b7e2e86 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json +++ b/.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 } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProto.java index 674329fc6847..3e7292fec89c 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProto.java +++ b/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 iterable = (Iterable) 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); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java index c2ac4efb5b5d..cbf5c4681f81 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java @@ -3990,11 +3990,13 @@ private 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(); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java index 401395030542..fec2e234a99b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRow.java +++ b/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 extends StorageApiDynamicDestinations { + 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 toRow; private final @Nullable SerializableBiFunction< @@ -59,21 +74,56 @@ class StorageApiDynamicDestinationsBeamRow 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 { + 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)); } } - }; + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java index 1710d32689c9..5207fe2e9f45 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java +++ b/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 extends StorageApiDynamicDestinations { + private static final Logger LOG = + LoggerFactory.getLogger(StorageApiDynamicDestinationsTableRow.class); private final BigQueryIO.TableRowFormatFunction formatFunction; private final BigQueryIO.@Nullable TableRowFormatFunction formatRecordOnFailureFunction; @@ -94,8 +98,23 @@ public MessageConverter 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); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java index ba72bb8682fd..32f48481d745 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java +++ b/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(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java index 1d618ba685ed..35bb48848af9 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java +++ b/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,6 +242,10 @@ 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 failedRowsWithErrors = result .getFailedStorageApiInserts() @@ -248,13 +253,25 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { "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); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProtoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProtoTest.java index 50fb2c073621..53e9b2c2f776 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProtoTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProtoTest.java @@ -19,6 +19,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import com.google.cloud.bigquery.storage.v1.TableFieldSchema; @@ -675,4 +676,255 @@ public void testTimestampNanosMessage() throws Exception { TEST_INSTANT_NANOS.getNano() * 1000L, picos.getField(picosDesc.findFieldByName("picoseconds"))); } + + @Test + public void testMessageFromBeamRow_withFieldSubset() throws Exception { + Schema unionSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("name", FieldType.STRING.withNullable(true)) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .addField("active", FieldType.BOOLEAN.withNullable(true)) + .build(); + + Schema subsetSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .addField("active", FieldType.BOOLEAN.withNullable(true)) + .build(); + + Row row = + Row.withSchema(unionSchema) + .withFieldValue("id", 42L) + .withFieldValue("name", "Alice") + .withFieldValue("score", 99.5) + .withFieldValue("active", true) + .build(); + + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + BeamRowToStorageApiProto.protoTableSchemaFromBeamSchema(subsetSchema), true, false); + + DynamicMessage msg = BeamRowToStorageApiProto.messageFromBeamRow(descriptor, row, null, -1); + + assertEquals(3, msg.getAllFields().size()); + FieldDescriptor idField = descriptor.findFieldByName("id"); + FieldDescriptor scoreField = descriptor.findFieldByName("score"); + FieldDescriptor activeField = descriptor.findFieldByName("active"); + + assertNotNull(idField); + assertNotNull(scoreField); + assertNotNull(activeField); + assertNull(descriptor.findFieldByName("name")); + + assertEquals(42L, msg.getField(idField)); + assertEquals(99.5, msg.getField(scoreField)); + assertEquals(true, msg.getField(activeField)); + } + + @Test + public void testMessageFromBeamRow_withNestedSubset() throws Exception { + Schema innerUnionSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("name", FieldType.STRING.withNullable(true)) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .build(); + + Schema innerSubsetSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .build(); + + Schema topUnionSchema = + Schema.builder() + .addField("nested", FieldType.row(innerUnionSchema).withNullable(true)) + .addField("topName", FieldType.STRING.withNullable(true)) + .build(); + + Schema topSubsetSchema = + Schema.builder() + .addField("nested", FieldType.row(innerSubsetSchema).withNullable(true)) + .build(); + + Row innerRow = + Row.withSchema(innerUnionSchema) + .withFieldValue("id", 100L) + .withFieldValue("name", "Bob") + .withFieldValue("score", 85.0) + .build(); + + Row topRow = + Row.withSchema(topUnionSchema) + .withFieldValue("nested", innerRow) + .withFieldValue("topName", "TopLevelName") + .build(); + + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + BeamRowToStorageApiProto.protoTableSchemaFromBeamSchema(topSubsetSchema), true, false); + + DynamicMessage msg = BeamRowToStorageApiProto.messageFromBeamRow(descriptor, topRow, null, -1); + + assertEquals(1, msg.getAllFields().size()); + FieldDescriptor nestedField = descriptor.findFieldByName("nested"); + assertNotNull(nestedField); + assertNull(descriptor.findFieldByName("topname")); + + DynamicMessage nestedMsg = (DynamicMessage) msg.getField(nestedField); + assertEquals(2, nestedMsg.getAllFields().size()); + + FieldDescriptor innerIdField = nestedField.getMessageType().findFieldByName("id"); + FieldDescriptor innerScoreField = nestedField.getMessageType().findFieldByName("score"); + assertNotNull(innerIdField); + assertNotNull(innerScoreField); + assertNull(nestedField.getMessageType().findFieldByName("name")); + + assertEquals(100L, nestedMsg.getField(innerIdField)); + assertEquals(85.0, nestedMsg.getField(innerScoreField)); + } + + @Test + public void testMessageFromBeamRow_withArrayOfNestedRowsSubset() throws Exception { + Schema innerUnionSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("name", FieldType.STRING.withNullable(true)) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .build(); + + Schema innerSubsetSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .build(); + + Schema topUnionSchema = + Schema.builder() + .addField("nestedArray", FieldType.array(FieldType.row(innerUnionSchema))) + .addField("nestedIterable", FieldType.iterable(FieldType.row(innerUnionSchema))) + .build(); + + Schema topSubsetSchema = + Schema.builder() + .addField("nestedArray", FieldType.array(FieldType.row(innerSubsetSchema))) + .addField("nestedIterable", FieldType.iterable(FieldType.row(innerSubsetSchema))) + .build(); + + Row innerRow1 = + Row.withSchema(innerUnionSchema) + .withFieldValue("id", 1L) + .withFieldValue("name", "Alice") + .withFieldValue("score", 90.0) + .build(); + + Row innerRow2 = + Row.withSchema(innerUnionSchema) + .withFieldValue("id", 2L) + .withFieldValue("name", "Bob") + .withFieldValue("score", 80.0) + .build(); + + Row topRow = + Row.withSchema(topUnionSchema) + .withFieldValue("nestedArray", ImmutableList.of(innerRow1, innerRow2)) + .withFieldValue("nestedIterable", ImmutableList.of(innerRow1, innerRow2)) + .build(); + + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + BeamRowToStorageApiProto.protoTableSchemaFromBeamSchema(topSubsetSchema), true, false); + + DynamicMessage msg = BeamRowToStorageApiProto.messageFromBeamRow(descriptor, topRow, null, -1); + + FieldDescriptor arrayField = descriptor.findFieldByName("nestedarray"); + assertNotNull(arrayField); + assertEquals(2, msg.getRepeatedFieldCount(arrayField)); + + DynamicMessage elem0 = (DynamicMessage) msg.getRepeatedField(arrayField, 0); + assertEquals(2, elem0.getAllFields().size()); + assertEquals(1L, elem0.getField(elem0.getDescriptorForType().findFieldByName("id"))); + assertEquals(90.0, elem0.getField(elem0.getDescriptorForType().findFieldByName("score"))); + assertNull(elem0.getDescriptorForType().findFieldByName("name")); + + FieldDescriptor iterField = descriptor.findFieldByName("nestediterable"); + assertNotNull(iterField); + assertEquals(2, msg.getRepeatedFieldCount(iterField)); + } + + @Test + public void testMessageFromBeamRow_withMapOfNestedRowsSubsetAndNullValues() throws Exception { + Schema innerUnionSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("name", FieldType.STRING.withNullable(true)) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .build(); + + Schema innerSubsetSchema = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .build(); + + Schema topUnionSchema = + Schema.builder() + .addField( + "nestedMap", + FieldType.map(FieldType.STRING, FieldType.row(innerUnionSchema).withNullable(true))) + .addField( + "primitiveMapWithNulls", + FieldType.map(FieldType.STRING, FieldType.INT32.withNullable(true)) + .withNullable(true)) + .build(); + + Schema topSubsetSchema = + Schema.builder() + .addField( + "nestedMap", + FieldType.map( + FieldType.STRING, FieldType.row(innerSubsetSchema).withNullable(true))) + .addField( + "primitiveMapWithNulls", + FieldType.map(FieldType.STRING, FieldType.INT32.withNullable(true)) + .withNullable(true)) + .build(); + + Row innerRow = + Row.withSchema(innerUnionSchema) + .withFieldValue("id", 10L) + .withFieldValue("name", "Carol") + .withFieldValue("score", 75.0) + .build(); + + Map rowMap = new HashMap<>(); + rowMap.put("k1", innerRow); + rowMap.put("k2_null", null); + + Map primMap = new HashMap<>(); + primMap.put("a", 100); + primMap.put("b_null", null); + + Row topRow = + Row.withSchema(topUnionSchema) + .withFieldValue("nestedMap", rowMap) + .withFieldValue("primitiveMapWithNulls", primMap) + .build(); + + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + BeamRowToStorageApiProto.protoTableSchemaFromBeamSchema(topSubsetSchema), true, false); + + DynamicMessage msg = BeamRowToStorageApiProto.messageFromBeamRow(descriptor, topRow, null, -1); + + FieldDescriptor mapField = descriptor.findFieldByName("nestedmap"); + assertNotNull(mapField); + assertEquals(2, msg.getRepeatedFieldCount(mapField)); + + FieldDescriptor primMapField = descriptor.findFieldByName("primitivemapwithnulls"); + assertNotNull(primMapField); + assertEquals(2, msg.getRepeatedFieldCount(primMapField)); + } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRowTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRowTest.java new file mode 100644 index 000000000000..8ecf0e35ea4f --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsBeamRowTest.java @@ -0,0 +1,331 @@ +/* + * 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableFieldSchema; +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.api.services.bigquery.model.TableSchema; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.DynamicMessage; +import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService; +import org.apache.beam.sdk.io.gcp.bigquery.StorageApiDynamicDestinations.MessageConverter; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.FieldType; +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.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link StorageApiDynamicDestinationsBeamRow}. */ +@RunWith(JUnit4.class) +public class StorageApiDynamicDestinationsBeamRowTest { + + private static final Schema UNION_SCHEMA = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("name", FieldType.STRING.withNullable(true)) + .addField("score", FieldType.DOUBLE.withNullable(true)) + .addField("active", FieldType.BOOLEAN.withNullable(true)) + .build(); + + private static final Schema SCHEMA_USERS = + Schema.builder() + .addField("id", FieldType.INT64) + .addField("name", FieldType.STRING.withNullable(true)) + .build(); + + private static final TableSchema BQ_SCHEMA_USERS = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("id").setType("INTEGER"), + new TableFieldSchema().setName("name").setType("STRING"))); + + private static final TableSchema BQ_SCHEMA_SCORES = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("id").setType("INTEGER"), + new TableFieldSchema().setName("score").setType("FLOAT"), + new TableFieldSchema().setName("active").setType("BOOLEAN"))); + + private PipelineOptions pipelineOptions; + + @Before + public void setUp() { + pipelineOptions = PipelineOptionsFactory.create(); + } + + @After + public void tearDown() throws Exception { + StorageApiDynamicDestinationsBeamRow.clearSchemaCache(); + } + + static class FakeDynamicDestinations extends DynamicDestinations { + private final Map schemas; + + FakeDynamicDestinations(Map schemas) { + this.schemas = schemas; + } + + @Override + public String getDestination(@Nullable ValueInSingleWindow element) { + return ""; + } + + @Override + public TableDestination getTable(String destination) { + return new TableDestination(destination, null); + } + + @Override + public @Nullable TableSchema getSchema(String destination) { + return schemas.get(destination); + } + } + + @Test + public void testPerDestinationMessageConverterWithInnerSchemas() throws Exception { + Map schemas = new HashMap<>(); + schemas.put("project:dataset.users", BQ_SCHEMA_USERS); + schemas.put("project:dataset.scores", BQ_SCHEMA_SCORES); + + FakeDynamicDestinations inner = new FakeDynamicDestinations(schemas); + StorageApiDynamicDestinationsBeamRow destinations = + new StorageApiDynamicDestinationsBeamRow<>(inner, UNION_SCHEMA, row -> row, null, false); + + MessageConverter converterUsers = + destinations.getMessageConverter("project:dataset.users", pipelineOptions, null, null); + MessageConverter converterScores = + destinations.getMessageConverter("project:dataset.scores", pipelineOptions, null, null); + + assertEquals(2, converterUsers.getTableSchema().getFieldsCount()); + assertEquals(3, converterScores.getTableSchema().getFieldsCount()); + + Descriptor descriptorUsers = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + converterUsers.getTableSchema(), true, false); + Descriptor descriptorScores = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + converterScores.getTableSchema(), true, false); + + assertNotNull(descriptorUsers.findFieldByName("id")); + assertNotNull(descriptorUsers.findFieldByName("name")); + assertNull(descriptorUsers.findFieldByName("score")); + assertNull(descriptorUsers.findFieldByName("active")); + + assertNotNull(descriptorScores.findFieldByName("id")); + assertNull(descriptorScores.findFieldByName("name")); + assertNotNull(descriptorScores.findFieldByName("score")); + assertNotNull(descriptorScores.findFieldByName("active")); + + Row testRow = + Row.withSchema(UNION_SCHEMA) + .withFieldValue("id", 1L) + .withFieldValue("name", "Alice") + .withFieldValue("score", 95.0) + .withFieldValue("active", true) + .build(); + + StorageApiWritePayload payloadUsers = + converterUsers.toMessage( + testRow, null, TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT); + DynamicMessage msgUsers = DynamicMessage.parseFrom(descriptorUsers, payloadUsers.getPayload()); + assertEquals(2, msgUsers.getAllFields().size()); + assertEquals(1L, msgUsers.getField(descriptorUsers.findFieldByName("id"))); + assertEquals("Alice", msgUsers.getField(descriptorUsers.findFieldByName("name"))); + + StorageApiWritePayload payloadScores = + converterScores.toMessage( + testRow, null, TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT); + DynamicMessage msgScores = + DynamicMessage.parseFrom(descriptorScores, payloadScores.getPayload()); + assertEquals(3, msgScores.getAllFields().size()); + assertEquals(1L, msgScores.getField(descriptorScores.findFieldByName("id"))); + assertEquals(95.0, msgScores.getField(descriptorScores.findFieldByName("score"))); + assertEquals(true, msgScores.getField(descriptorScores.findFieldByName("active"))); + } + + @Test + public void testSchemaResolutionFromDatasetService() throws Exception { + FakeDynamicDestinations inner = new FakeDynamicDestinations(Collections.emptyMap()); + StorageApiDynamicDestinationsBeamRow destinations = + new StorageApiDynamicDestinationsBeamRow<>(inner, UNION_SCHEMA, row -> row, null, false); + + DatasetService mockDatasetService = mock(DatasetService.class); + TableReference usersRef = BigQueryHelpers.parseTableSpec("project:dataset.users"); + TableReference scoresRef = BigQueryHelpers.parseTableSpec("project:dataset.scores"); + + when(mockDatasetService.getTable(eq(usersRef), any(), any())) + .thenReturn(new Table().setSchema(BQ_SCHEMA_USERS)); + when(mockDatasetService.getTable(eq(scoresRef), any(), any())) + .thenReturn(new Table().setSchema(BQ_SCHEMA_SCORES)); + + MessageConverter converterUsers = + destinations.getMessageConverter( + "project:dataset.users", pipelineOptions, mockDatasetService, null); + MessageConverter converterScores = + destinations.getMessageConverter( + "project:dataset.scores", pipelineOptions, mockDatasetService, null); + + assertEquals(2, converterUsers.getTableSchema().getFieldsCount()); + assertEquals(3, converterScores.getTableSchema().getFieldsCount()); + } + + @Test + public void testFallbackToStaticSchemaWhenResolutionFails() throws Exception { + FakeDynamicDestinations inner = new FakeDynamicDestinations(Collections.emptyMap()); + StorageApiDynamicDestinationsBeamRow destinations = + new StorageApiDynamicDestinationsBeamRow<>(inner, SCHEMA_USERS, row -> row, null, false); + + MessageConverter converter = + destinations.getMessageConverter("project:dataset.unknown", pipelineOptions, null, null); + + assertEquals(2, converter.getTableSchema().getFieldsCount()); + } + + @Test + public void testCdcWritesWithDynamicPerDestinationSchemas() throws Exception { + Map schemas = new HashMap<>(); + schemas.put("project:dataset.users", BQ_SCHEMA_USERS); + + FakeDynamicDestinations inner = new FakeDynamicDestinations(schemas); + StorageApiDynamicDestinationsBeamRow destinations = + new StorageApiDynamicDestinationsBeamRow<>(inner, UNION_SCHEMA, row -> row, null, true); + + MessageConverter converter = + destinations.getMessageConverter("project:dataset.users", pipelineOptions, null, null); + + DescriptorProtos.DescriptorProto proto = converter.getDescriptor(true); + assertNotNull(proto); + + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + converter.getTableSchema(), true, true); + assertNotNull(descriptor.findFieldByName(StorageApiCDC.CHANGE_TYPE_COLUMN)); + assertNotNull(descriptor.findFieldByName(StorageApiCDC.CHANGE_SQN_COLUMN)); + + Row testRow = + Row.withSchema(UNION_SCHEMA) + .withFieldValue("id", 1L) + .withFieldValue("name", "Alice") + .withFieldValue("score", 95.0) + .withFieldValue("active", true) + .build(); + + RowMutationInformation mutationInfo = + RowMutationInformation.of(RowMutationInformation.MutationType.UPSERT, 42L); + + StorageApiWritePayload payload = + converter.toMessage( + testRow, mutationInfo, TableRowToStorageApiProto.ErrorCollector.DONT_COLLECT); + DynamicMessage msg = DynamicMessage.parseFrom(descriptor, payload.getPayload()); + + assertEquals( + "UPSERT", msg.getField(descriptor.findFieldByName(StorageApiCDC.CHANGE_TYPE_COLUMN))); + assertEquals( + Long.toHexString(42L), + msg.getField(descriptor.findFieldByName(StorageApiCDC.CHANGE_SQN_COLUMN))); + assertEquals(1L, msg.getField(descriptor.findFieldByName("id"))); + assertEquals("Alice", msg.getField(descriptor.findFieldByName("name"))); + } + + @Test + public void testDatasetServiceFailureFallsBackGracefully() throws Exception { + Map schemas = new HashMap<>(); + schemas.put("project:dataset.users", BQ_SCHEMA_USERS); + FakeDynamicDestinations inner = new FakeDynamicDestinations(schemas); + + StorageApiDynamicDestinationsBeamRow destinations = + new StorageApiDynamicDestinationsBeamRow<>(inner, UNION_SCHEMA, row -> row, null, false); + + DatasetService mockDatasetService = mock(DatasetService.class); + TableReference usersRef = BigQueryHelpers.parseTableSpec("project:dataset.users"); + when(mockDatasetService.getTable(eq(usersRef), any(), any())) + .thenThrow(new IOException("BigQuery quota exceeded")); + + MessageConverter converter = + destinations.getMessageConverter( + "project:dataset.users", pipelineOptions, mockDatasetService, null); + + assertEquals(2, converter.getTableSchema().getFieldsCount()); + Descriptor descriptor = + TableRowToStorageApiProto.getDescriptorFromTableSchema( + converter.getTableSchema(), true, false); + assertNotNull(descriptor.findFieldByName("id")); + assertNotNull(descriptor.findFieldByName("name")); + } + + @Test + public void testToFailsafeTableRow() throws Exception { + Row testRow = + Row.withSchema(UNION_SCHEMA) + .withFieldValue("id", 1L) + .withFieldValue("name", "Alice") + .withFieldValue("score", 95.0) + .withFieldValue("active", true) + .build(); + + FakeDynamicDestinations inner = new FakeDynamicDestinations(Collections.emptyMap()); + StorageApiDynamicDestinationsBeamRow destinationsDefault = + new StorageApiDynamicDestinationsBeamRow<>(inner, UNION_SCHEMA, row -> row, null, false); + + MessageConverter converterDefault = + destinationsDefault.getMessageConverter( + "project:dataset.users", pipelineOptions, null, null); + TableRow defaultTableRow = converterDefault.toFailsafeTableRow(testRow); + assertEquals("1", defaultTableRow.get("id").toString()); + assertEquals("Alice", defaultTableRow.get("name")); + + StorageApiDynamicDestinationsBeamRow destinationsCustom = + new StorageApiDynamicDestinationsBeamRow<>( + inner, + UNION_SCHEMA, + row -> row, + (schemaInfo, element) -> new TableRow().set("custom_id", element.getInt64("id")), + false); + + MessageConverter converterCustom = + destinationsCustom.getMessageConverter( + "project:dataset.users", pipelineOptions, null, null); + TableRow customTableRow = converterCustom.toFailsafeTableRow(testRow); + assertEquals(1L, customTableRow.get("custom_id")); + } +} diff --git a/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py b/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py index 49725d54e990..2bd8bc192452 100644 --- a/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py +++ b/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py @@ -31,6 +31,8 @@ from hamcrest.core.core.allof import all_of import apache_beam as beam +from apache_beam.io.gcp import bigquery +from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery import StorageWriteToBigQuery from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher @@ -483,6 +485,104 @@ def test_write_to_dynamic_destinations(self): use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) + def test_write_to_dynamic_destinations_with_dynamic_schema(self): + base_table_spec = '{}.dynamic_dest_dyn_schema_'.format(self.dataset_id) + spec_with_project = '{}:{}'.format(self.project, base_table_spec) + table_id_a = 'dynamic_dest_dyn_schema_users' + table_id_b = 'dynamic_dest_dyn_schema_scores' + table_a = base_table_spec + 'users' + table_b = base_table_spec + 'scores' + + schema_a = "id:INTEGER,name:STRING" + schema_b = "id:INTEGER,score:INTEGER,active:BOOLEAN" + + # Pre-create destination tables with their distinct specific schemas prior + # to pipeline execution to ensure tables only contain their specific fields. + self.bigquery_client.get_or_create_table( + project_id=self.project, + dataset_id=self.dataset_id, + table_id=table_id_a, + schema=bigquery_tools.get_table_schema_from_string(schema_a), + create_disposition='CREATE_IF_NEEDED', + write_disposition='WRITE_APPEND') + self.bigquery_client.get_or_create_table( + project_id=self.project, + dataset_id=self.dataset_id, + table_id=table_id_b, + schema=bigquery_tools.get_table_schema_from_string(schema_b), + create_disposition='CREATE_IF_NEEDED', + write_disposition='WRITE_APPEND') + + elements_a = [ + { + 'id': 1, 'name': 'alice' + }, + { + 'id': 2, 'name': 'bob' + }, + ] + elements_b = [ + { + 'id': 101, 'score': 95, 'active': True + }, + { + 'id': 102, 'score': 80, 'active': False + }, + ] + elements = elements_a + elements_b + + schema_map = { + spec_with_project + 'users': schema_a, + spec_with_project + 'scores': schema_b, + } + + bq_matchers = [ + BigqueryFullResultMatcher( + project=self.project, + query="SELECT * FROM %s" % table_a, + data=self.parse_expected_data(elements_a)), + BigqueryFullResultMatcher( + project=self.project, + query="SELECT * FROM %s" % table_b, + data=self.parse_expected_data(elements_b)), + ] + + def get_destination(record): + if 'name' in record: + return spec_with_project + 'users' + return spec_with_project + 'scores' + + def get_schema_raw(dest, side_map): + return side_map[dest] + + get_schema = bigquery.dynamic_schema( + get_schema_raw, + union_schema="id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN") + + with beam.Pipeline(argv=self.args) as p: + schema_pc = p | "CreateSchema" >> beam.Create([schema_map]) + _ = ( + p + | "CreateElements" >> beam.Create(elements) + | beam.io.WriteToBigQuery( + table=get_destination, + method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=get_schema, + schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), + use_at_least_once=False)) + hamcrest_assert(p, all_of(*bq_matchers)) + + # Verify destination tables retained their specific schemas and were not + # created or altered to the union schema + fetched_table_a = self.bigquery_client.get_table( + self.project, self.dataset_id, table_id_a) + fetched_table_b = self.bigquery_client.get_table( + self.project, self.dataset_id, table_id_b) + self.assertEqual([f.name for f in fetched_table_a.schema.fields], + ['id', 'name']) + self.assertEqual([f.name for f in fetched_table_b.schema.fields], + ['id', 'score', 'active']) + def test_write_to_dynamic_destinations_with_beam_rows(self): base_table_spec = '{}.dynamic_dest_'.format(self.dataset_id) spec_with_project = '{}:{}'.format(self.project, base_table_spec) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 314effad5520..5e2e108122c9 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -197,6 +197,73 @@ def compute_table_name(row): a tuple of PCollectionViews to be passed to the schema callable (much like the `table_side_inputs` parameter). +Dynamic Schemas with Storage Write API +-------------------------------------- +When writing to dynamic destinations with `method=STORAGE_WRITE_API`, a union schema +containing all fields across destination tables is required at the PCollection level +for cross-language type inference and runtime row serialization. + +The recommended best-practice is to use the `dynamic_schema` helper: + +* **Using a dictionary map**: If schemas are known at pipeline construction time, pass + a dictionary mapping destinations to schemas. The helper automatically infers and merges + all fields into the required union schema:: + + schema_map = { + 'my_project:dataset.users': 'id:INTEGER,name:STRING', + 'my_project:dataset.scores': 'id:INTEGER,score:INTEGER,active:BOOLEAN' + } + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=dynamic_schema(schema_map)) + +* **Using a callable function or side inputs**: If schemas are determined dynamically + via a callable function, wrap the callable with `dynamic_schema` and explicitly pass + `union_schema`:: + + def get_schema(destination, schema_dict): + return schema_dict[destination] + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=dynamic_schema( + get_schema, + union_schema='id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN'), + schema_side_inputs=(schema_dict_side_input,)) + +Differences and Limitations Compared to Native Java BigQueryIO +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Because Python uses a cross-language expansion (`SchemaAwareExternalTransform`) +to invoke the Java Storage Write API implementation, certain behaviors differ +from native Java `DynamicDestinations`: + +1. **Table Creation on Write (`CREATE_IF_NEEDED`)**: + When writing to dynamic destinations with callable schemas, `WriteToBigQuery` + automatically creates destination tables in BigQuery using each table's exact + specific schema on the Python side before delegating writes to the Storage + Write API, ensuring that auto-created tables only contain their specific + columns (matching the behavior of native Java). + +2. **Streaming Schema Evolution**: + In native Java, pipelines can write dynamic `TableRow` objects and leverage BigQuery's + automatic schema update capabilities (`autoSchemaUpdates`) to append new fields to + existing tables at runtime. In Python cross-language execution, elements must pass + through a static `RowCoder` compiled at pipeline submission time. Introducing new fields + or destination tables not represented in the initial `union_schema` requires draining + and updating/restarting the pipeline. + +3. **Side Inputs for Schemas**: + Side inputs (`schema_side_inputs`) can dynamically dictate destination-to-schema + mappings and select subsets of fields per destination at runtime. However, side inputs + cannot introduce new fields that were omitted from the statically declared `union_schema`. + +4. **Field Type Compatibility Across Destinations**: + Overlapping column names across different destination tables must share compatible + BigQuery types (e.g. `status` cannot be `INTEGER` in one table and `STRING` in another). + Additional Parameters for BigQuery Tables ----------------------------------------- @@ -356,6 +423,7 @@ def chain_after(result): # pytype: skip-file import collections +import copy import io import itertools import json @@ -1956,6 +2024,120 @@ def _restore_table_ref(sharded_table_ref_elems_kv): SCHEMA_AUTODETECT = 'SCHEMA_AUTODETECT' +def dynamic_schema(schema_fn_or_map, union_schema=None): + """Helper to construct a dynamic schema callable with a union schema hint. + + When using the BigQuery Storage Write API (`method=STORAGE_WRITE_API`) with + dynamic destinations, the cross-language transform requires a PCollection-level + union schema containing all fields across all target tables for protobuf + serialization and type inference. + + This helper provides the recommended best practice for constructing dynamic + schemas: + + 1. **Dictionary Map**: If destination table schemas are provided as a dictionary + mapping table names/specs to schemas (str, dict, or TableSchema), this helper + automatically merges all fields into a single union schema. + 2. **Callable**: If a callable function is used, this helper attaches the provided + `union_schema` to the callable as a schema hint (`_union_schema`). + + Example using a dictionary map (union schema is auto-inferred):: + + schema_map = { + 'project:dataset.users': 'id:INTEGER,name:STRING', + 'project:dataset.scores': 'id:INTEGER,score:INTEGER,active:BOOLEAN' + } + + def get_destination(record): + if 'name' in record: + return 'project:dataset.users' + return 'project:dataset.scores' + + schema_callable = dynamic_schema(schema_map) + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=schema_callable) + + Example using a callable with explicit union schema:: + + def get_schema(destination, schema_side_input): + return schema_side_input[destination] + + schema_callable = dynamic_schema( + get_schema, + union_schema='id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN') + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=schema_callable, + schema_side_inputs=(schema_side_input,)) + + Args: + schema_fn_or_map: A callable `(destination, *side_inputs) -> schema` + or a dictionary mapping destination strings to schemas (str, dict, or + TableSchema). + union_schema: (Optional) The union schema containing all fields across + target tables. Can be a string, dict, or TableSchema object. Required if + `schema_fn_or_map` is a callable. + + Returns: + A callable with the attached `_union_schema` attribute for Storage Write API. + """ + if isinstance(schema_fn_or_map, dict): + if union_schema is None: + bq_schemas = [ + copy.deepcopy(bigquery_tools.get_bq_tableschema(s)) + for s in schema_fn_or_map.values() + ] + + def _merge_fields(field_a, field_b): + if field_a.type != field_b.type: + raise ValueError( + f"Conflicting types for field '{field_a.name}': " + f"{field_a.type} vs {field_b.type}") + if field_a.type in ('RECORD', 'STRUCT'): + merged_subfields = {} + for f in (field_a.fields or []): + merged_subfields[f.name] = f + for f in (field_b.fields or []): + if f.name in merged_subfields: + merged_subfields[f.name] = _merge_fields( + merged_subfields[f.name], f) + else: + merged_subfields[f.name] = f + field_a.fields = list(merged_subfields.values()) + return field_a + + merged_fields = {} + for schema in bq_schemas: + for field in schema.fields: + name = field.name + if name in merged_fields: + merged_fields[name] = _merge_fields(merged_fields[name], field) + else: + merged_fields[name] = field + union_schema = bigquery.TableSchema(fields=list(merged_fields.values())) + + def lookup_schema(destination, *args): + return schema_fn_or_map[destination] + + schema_callable = lookup_schema + elif callable(schema_fn_or_map): + if union_schema is None: + raise ValueError( + "union_schema must be explicitly provided when schema_fn_or_map " + "is a callable.") + schema_callable = schema_fn_or_map + else: + raise TypeError("schema_fn_or_map must be a callable or a dictionary.") + + schema_callable._union_schema = union_schema + return schema_callable + + class WriteToBigQuery(PTransform): """Write data to BigQuery. @@ -2388,6 +2570,7 @@ def find_in_nested_dict(schema): table=self.table_reference, schema=self.schema, table_side_inputs=self.table_side_inputs, + schema_side_inputs=self.schema_side_inputs, create_disposition=self.create_disposition, write_disposition=self.write_disposition, additional_bq_parameters=self.additional_bq_parameters, @@ -2399,7 +2582,8 @@ def find_in_nested_dict(schema): primary_key=self._primary_key, big_lake_configuration=self._big_lake_configuration, expansion_service=self.expansion_service, - type_overrides=self._type_overrides) + type_overrides=self._type_overrides, + test_client=self.test_client) else: raise ValueError(f"Unsupported method {method_to_use}") @@ -2616,7 +2800,7 @@ def __getitem__(self, key): class StorageWriteToBigQuery(PTransform): """Writes data to BigQuery using Storage API. - Supports dynamic destinations. Dynamic schemas are not supported yet. + Supports dynamic destinations and dynamic schemas. Experimental; no backwards compatibility guarantees. """ @@ -2638,6 +2822,7 @@ def __init__( table, table_side_inputs=None, schema=None, + schema_side_inputs=None, create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=BigQueryDisposition.WRITE_APPEND, additional_bq_parameters=None, @@ -2649,10 +2834,12 @@ def __init__( primary_key: list[str] = None, big_lake_configuration=None, expansion_service=None, - type_overrides=None): + type_overrides=None, + test_client=None): self._table = table - self._table_side_inputs = table_side_inputs + self._table_side_inputs = table_side_inputs or () self._schema = schema + self._schema_side_inputs = schema_side_inputs or () self._create_disposition = create_disposition self._write_disposition = write_disposition self.additional_bq_parameters = additional_bq_parameters @@ -2666,6 +2853,7 @@ def __init__( self._type_overrides = type_overrides self._expansion_service = expansion_service or BeamJarExpansionService( 'sdks:java:io:google-cloud-platform:expansion-service:build') + self._test_client = test_client def expand(self, input): if self._schema is None: @@ -2677,9 +2865,8 @@ def expand(self, input): "A schema is required in order to prepare rows " "for writing with STORAGE_WRITE_API.") from exn elif callable(self._schema): - raise NotImplementedError( - "Writing with dynamic schemas is not " - "supported for this write method.") + schema = self._schema + is_rows = False elif isinstance(self._schema, vp.ValueProvider): schema = self._schema.get() is_rows = False @@ -2691,13 +2878,23 @@ def expand(self, input): # if writing to one destination, just convert to Beam rows and send over if not callable(table): + if callable(schema): + raise ValueError( + "Writing with a dynamic schema is only supported when writing to " + "dynamic destinations.") if is_rows: input_beam_rows = input else: input_beam_rows = ( input | "Convert dict to Beam Row" >> self.ConvertToBeamRows( - schema, False, self._type_overrides).with_output_types()) + schema, + False, + self._type_overrides, + create_disposition=self._create_disposition, + write_disposition=self._write_disposition, + additional_bq_parameters=self.additional_bq_parameters, + test_client=self._test_client).with_output_types()) # For dynamic destinations, we first figure out where each row is going. # Then we send (destination, record) rows over to Java SchemaTransform. @@ -2729,7 +2926,14 @@ def expand(self, input): input_beam_rows = ( input_rows | "Convert dict to Beam Row" >> self.ConvertToBeamRows( - schema, True, self._type_overrides).with_output_types()) + schema, + True, + self._type_overrides, + schema_side_inputs=self._schema_side_inputs, + create_disposition=self._create_disposition, + write_disposition=self._write_disposition, + additional_bq_parameters=self.additional_bq_parameters, + test_client=self._test_client).with_output_types()) # communicate to Java that this write should use dynamic destinations table = StorageWriteToBigQuery.DYNAMIC_DESTINATIONS @@ -2796,25 +3000,168 @@ def __enter__(self): def __exit__(self, *args): pass + class _ConvertDynamicRowDoFn(DoFn): + def __init__( + self, + schema, + union_field_names, + create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, + write_disposition=BigQueryDisposition.WRITE_APPEND, + additional_bq_parameters=None, + test_client=None): + self.schema = schema + self.union_field_names = union_field_names + self.create_disposition = create_disposition + self.write_disposition = write_disposition + self.additional_bq_parameters = additional_bq_parameters + self.test_client = test_client + self.bigquery_wrapper = None + + def start_bundle(self): + if self.create_disposition == BigQueryDisposition.CREATE_IF_NEEDED: + if self.test_client: + self.bigquery_wrapper = bigquery_tools.BigQueryWrapper( + client=self.test_client) + else: + try: + self.bigquery_wrapper = bigquery_tools.BigQueryWrapper() + except Exception: + self.bigquery_wrapper = None + + def _create_table_if_needed(self, dest, record_schema): + if self.create_disposition != BigQueryDisposition.CREATE_IF_NEEDED: + return + try: + table_ref = bigquery_tools.parse_table_reference(dest) + except ValueError: + return + if not table_ref.datasetId or not table_ref.tableId: + return + + if self.bigquery_wrapper is None: + if self.test_client: + self.bigquery_wrapper = bigquery_tools.BigQueryWrapper( + client=self.test_client) + else: + try: + self.bigquery_wrapper = bigquery_tools.BigQueryWrapper() + except Exception: + return + + project_id = ( + table_ref.projectId or self.bigquery_wrapper._get_project_id()) + str_table_ref = '%s:%s.%s' % ( + project_id, table_ref.datasetId, table_ref.tableId) + if str_table_ref in _KNOWN_TABLES or dest in _KNOWN_TABLES: + return + + table_schema = bigquery_tools.get_bq_tableschema(record_schema) + self.bigquery_wrapper.get_or_create_table( + project_id=project_id, + dataset_id=table_ref.datasetId, + table_id=table_ref.tableId, + schema=table_schema, + create_disposition=self.create_disposition, + write_disposition=self.write_disposition, + additional_create_parameters=self.additional_bq_parameters) + _KNOWN_TABLES.add(str_table_ref) + _KNOWN_TABLES.add(dest) + + def process(self, row, *schema_side_inputs): + dest, dict_row = row[0], row[1] + record_schema = self.schema(dest, *schema_side_inputs) + self._create_table_if_needed(dest, record_schema) + + record_row = bigquery_tools.beam_row_from_dict(dict_row, record_schema) + if self.union_field_names: + record_dict = record_row._asdict() + record_row = beam.Row( + **{ + name: record_dict.get(name, None) + for name in self.union_field_names + }) + yield beam.Row( + **{ + StorageWriteToBigQuery.DESTINATION: dest, + StorageWriteToBigQuery.RECORD: record_row + }) + class ConvertToBeamRows(PTransform): - def __init__(self, schema, dynamic_destinations, type_overrides=None): + def __init__( + self, + schema, + dynamic_destinations, + type_overrides=None, + schema_side_inputs=None, + create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, + write_disposition=BigQueryDisposition.WRITE_APPEND, + additional_bq_parameters=None, + test_client=None): self.schema = schema self.dynamic_destinations = dynamic_destinations self.type_overrides = type_overrides + self.schema_side_inputs = schema_side_inputs or () + self.create_disposition = create_disposition + self.write_disposition = write_disposition + self.additional_bq_parameters = additional_bq_parameters + self.test_client = test_client + + def _get_record_type_hint(self): + if callable(self.schema): + schema_hint = ( + getattr(self.schema, '_union_schema', None) or + getattr(self.schema, '_table_schema', None) or + getattr(self.schema, '_beam_schema', None) or + getattr(self.schema, '_schema_hint', None) or + getattr(self.schema, '_output_types', None) or + getattr(self.schema, 'table_schema', None) or + getattr(self.schema, 'schema', None)) + if schema_hint is not None: + if isinstance( + schema_hint, + (bigquery.TableSchema, bigquery.TableFieldSchema, str, dict)): + row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( + schema_hint, self.type_overrides) + return RowTypeConstraint.from_fields(row_type_hints) + elif isinstance(schema_hint, RowTypeConstraint): + return schema_hint + return RowTypeConstraint.from_fields([]) + else: + row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( + self.schema, self.type_overrides) + return RowTypeConstraint.from_fields(row_type_hints) def expand(self, input_dicts): if self.dynamic_destinations: - return ( - input_dicts - | "Convert dict to Beam Row" >> beam.Map( - lambda row, schema=DoFn.SetupContextParam( - StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args= - [self.schema]): beam.Row( - **{ - StorageWriteToBigQuery.DESTINATION: row[0], - StorageWriteToBigQuery.RECORD: bigquery_tools. - beam_row_from_dict(row[1], schema) - }))) + if callable(self.schema): + record_hint = self._get_record_type_hint() + union_field_names = [ + name for name, _ in getattr(record_hint, '_fields', ()) + ] + + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.ParDo( + StorageWriteToBigQuery._ConvertDynamicRowDoFn( + self.schema, + union_field_names, + create_disposition=self.create_disposition, + write_disposition=self.write_disposition, + additional_bq_parameters=self.additional_bq_parameters, + test_client=self.test_client), + *self.schema_side_inputs)) + else: + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.Map( + lambda row, schema=DoFn.SetupContextParam( + StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args= + [self.schema]): beam.Row( + **{ + StorageWriteToBigQuery.DESTINATION: row[0], + StorageWriteToBigQuery.RECORD: bigquery_tools. + beam_row_from_dict(row[1], schema) + }))) else: return ( input_dicts @@ -2825,17 +3172,14 @@ def expand(self, input_dicts): ]): bigquery_tools.beam_row_from_dict(row, schema))) def with_output_types(self): - row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( - self.schema, self.type_overrides) + record_hint = self._get_record_type_hint() if self.dynamic_destinations: type_hint = RowTypeConstraint.from_fields([ (StorageWriteToBigQuery.DESTINATION, str), - ( - StorageWriteToBigQuery.RECORD, - RowTypeConstraint.from_fields(row_type_hints)) + (StorageWriteToBigQuery.RECORD, record_hint) ]) else: - type_hint = RowTypeConstraint.from_fields(row_type_hints) + type_hint = record_hint return super().with_output_types(type_hint) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py new file mode 100644 index 000000000000..33d4a98bfec9 --- /dev/null +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -0,0 +1,416 @@ +# +# 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. +# + +"""Unit tests for BigQuery Storage Write API dynamic schemas.""" + +import unittest +from unittest import mock + +import apache_beam as beam +from apache_beam.io.gcp import bigquery +from apache_beam.io.gcp import bigquery_tools +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to +from apache_beam.typehints.row_type import RowTypeConstraint + +try: + from google.api_core.exceptions import GoogleAPICallError +except ImportError: + GoogleAPICallError = None + + +@unittest.skipIf( + GoogleAPICallError is None, 'GCP dependencies are not installed') +@mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') +class BigQueryStorageWriteDynamicSchemaTest(unittest.TestCase): + """Test dynamic schema support in BigQuery Storage Write API.""" + def test_storage_write_init_with_schema_side_inputs( + self, mock_expansion_service): + """Test that StorageWriteToBigQuery accepts schema_side_inputs.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', + schema=lambda dest: None, + schema_side_inputs=('side_input_1', )) + self.assertEqual(transform._schema_side_inputs, ('side_input_1', )) + self.assertEqual(transform._table_side_inputs, ()) + + def test_convert_to_beam_rows_dynamic_destinations_dynamic_schema( + self, mock_expansion_service): + """Test ConvertToBeamRows with dynamic destinations and dynamic schema.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, + ] + } + schema2 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'score', 'type': 'FLOAT' + }, + ] + } + schema_map = {'table1': schema1, 'table2': schema2} + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest: schema_map[dest], dynamic_destinations=True) + + with TestPipeline() as p: + input_data = [ + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), + ] + res = p | "CreateInput" >> beam.Create(input_data) | converter + + expected_rows = [ + beam.Row(destination='table1', record=beam.Row(id=1, name='foo')), + beam.Row(destination='table2', record=beam.Row(id=2, score=3.14)), + ] + assert_that(res, equal_to(expected_rows)) + + def test_convert_to_beam_rows_dynamic_destinations_with_side_inputs( + self, mock_expansion_service): + """Test ConvertToBeamRows with dynamic schema and side inputs.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, + ] + } + schema2 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'score', 'type': 'FLOAT' + }, + ] + } + + with TestPipeline() as p: + side_pcoll = ( + p + | "CreateSide" >> beam.Create([{ + 'table1': schema1, 'table2': schema2 + }])) + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest, side_map: side_map[dest], + dynamic_destinations=True, + schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll), )) + + input_data = [ + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), + ] + res = p | "CreateInput" >> beam.Create(input_data) | converter + + expected_rows = [ + beam.Row(destination='table1', record=beam.Row(id=1, name='foo')), + beam.Row(destination='table2', record=beam.Row(id=2, score=3.14)), + ] + assert_that(res, equal_to(expected_rows)) + + def test_storage_write_static_destination_dynamic_schema_raises_error( + self, mock_expansion_service): + """Test that static destination with dynamic schema raises ValueError.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', schema=lambda dest: None) + with self.assertRaisesRegex( + ValueError, + "Writing with a dynamic schema is only supported when writing to " + "dynamic destinations."): + with TestPipeline() as p: + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | transform + + def test_convert_to_beam_rows_with_output_types_dynamic_schema( + self, mock_expansion_service): + """Test with_output_types when schema is callable.""" + converter_dyn = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest: None, dynamic_destinations=True) + type_hint_dyn = converter_dyn.with_output_types().get_type_hints( + ).simple_output_type('') + self.assertIsInstance(type_hint_dyn, RowTypeConstraint) + self.assertEqual( + type_hint_dyn._fields, + ( + (bigquery.StorageWriteToBigQuery.DESTINATION, str), + ( + bigquery.StorageWriteToBigQuery.RECORD, + RowTypeConstraint.from_fields([])), + )) + + def test_convert_to_beam_rows_with_output_types_dynamic_schema_hint( + self, mock_expansion_service): + """Test with_output_types when schema is callable with _union_schema.""" + def dyn_schema(dest): + return None + + dyn_schema._union_schema = 'id:INTEGER,name:STRING' + converter_dyn = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, dynamic_destinations=True) + type_hint_dyn = converter_dyn.with_output_types().get_type_hints( + ).simple_output_type('') + self.assertIsInstance(type_hint_dyn, RowTypeConstraint) + self.assertEqual( + type_hint_dyn._fields[0], + (bigquery.StorageWriteToBigQuery.DESTINATION, str)) + self.assertEqual( + type_hint_dyn._fields[1][0], bigquery.StorageWriteToBigQuery.RECORD) + expected_record_hint = RowTypeConstraint.from_fields( + bigquery_tools.get_beam_typehints_from_tableschema( + 'id:INTEGER,name:STRING')) + self.assertEqual( + type_hint_dyn._fields[1][1]._fields, expected_record_hint._fields) + + def test_convert_to_beam_rows_union_schema_fills_missing_attributes( + self, mock_expansion_service): + """Test ConvertToBeamRows fills None for fields in union schema not in row.""" + def dyn_schema(dest): + if 'users' in dest: + return 'id:INTEGER,name:STRING' + return 'id:INTEGER,score:INTEGER' + + dyn_schema._union_schema = 'id:INTEGER,name:STRING,score:INTEGER' + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, dynamic_destinations=True) + + with TestPipeline() as p: + rows = ( + p + | beam.Create([ + ('dest_users', { + 'id': 1, 'name': 'alice' + }), + ('dest_scores', { + 'id': 2, 'score': 95 + }), + ]) + | converter) + + def check_rows(actual): + actual_list = list(actual) + assert len(actual_list) == 2 + r1, r2 = actual_list[0], actual_list[1] + if r1.destination == 'dest_scores': + r1, r2 = r2, r1 + assert r1.destination == 'dest_users' + assert r1.record.id == 1 + assert r1.record.name == 'alice' + assert r1.record.score is None + assert r2.destination == 'dest_scores' + assert r2.record.id == 2 + assert r2.record.name is None + assert r2.record.score == 95 + + assert_that(rows, check_rows) + + def test_storage_write_to_bigquery_expand_dynamic_schema( + self, mock_expansion_service): + """Test StorageWriteToBigQuery expand does not fail for callable schema.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + ] + } + + class _DummyExternalTransform(beam.PTransform): + def expand(self, pcoll): + return { + bigquery.StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS: ( + pcoll.pipeline | "CreateErrors" >> beam.Create([])) + } + + with mock.patch.object(bigquery, + 'SchemaAwareExternalTransform', + autospec=True) as mock_ext: + mock_ext.return_value = _DummyExternalTransform() + transform = bigquery.StorageWriteToBigQuery( + table=lambda record: 'table1', schema=lambda dest: schema1) + + with TestPipeline() as p: + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | transform + + mock_ext.assert_called_once() + _, kwargs = mock_ext.call_args + self.assertEqual( + kwargs['table'], bigquery.StorageWriteToBigQuery.DYNAMIC_DESTINATIONS) + + def test_write_to_bigquery_storage_api_passes_schema_side_inputs( + self, mock_expansion_service): + """Test WriteToBigQuery passes schema_side_inputs to StorageWriteToBigQuery.""" + with mock.patch.object(bigquery, 'StorageWriteToBigQuery', + autospec=True) as mock_storage_write: + mock_storage_write.return_value = beam.Map(lambda x: x) + with TestPipeline() as p: + side_pc = p | "CreateSide" >> beam.Create([1]) + write_transform = bigquery.WriteToBigQuery( + table='proj:ds.table', + method=bigquery.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=lambda dest: None, + schema_side_inputs=(beam.pvalue.AsSingleton(side_pc), )) + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | write_transform + + mock_storage_write.assert_called_once() + _, kwargs = mock_storage_write.call_args + self.assertEqual(len(kwargs['schema_side_inputs']), 1) + + def test_dynamic_schema_helper_with_dictionary(self, mock_expansion_service): + """Test dynamic_schema auto-merges fields from a dictionary map.""" + schema_map = { + 'table_a': 'id:INTEGER,name:STRING', + 'table_b': 'id:INTEGER,score:INTEGER,active:BOOLEAN' + } + schema_callable = bigquery.dynamic_schema(schema_map) + + # 1. Verify it returns the correct schema per destination + self.assertEqual(schema_callable('table_a'), 'id:INTEGER,name:STRING') + self.assertEqual( + schema_callable('table_b'), 'id:INTEGER,score:INTEGER,active:BOOLEAN') + + # 2. Verify it auto-merged all unique fields into _union_schema + expected_union = bigquery_tools.get_bq_tableschema( + 'id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN') + self.assertEqual(schema_callable._union_schema, expected_union) + + def test_dynamic_schema_helper_with_callable_and_explicit_union( + self, mock_expansion_service): + """Test dynamic_schema attaches union_schema explicitly to a callable.""" + def get_schema(dest): + return 'id:INTEGER,name:STRING' + + union_schema = 'id:INTEGER,name:STRING,score:INTEGER' + schema_callable = bigquery.dynamic_schema( + get_schema, union_schema=union_schema) + + self.assertEqual(schema_callable('table_a'), 'id:INTEGER,name:STRING') + self.assertEqual(schema_callable._union_schema, union_schema) + + def test_dynamic_schema_helper_missing_union_on_callable_raises_error( + self, mock_expansion_service): + """Test dynamic_schema raises ValueError if union_schema is missing on callable.""" + def get_schema(dest): + return 'id:INTEGER' + + with self.assertRaises(ValueError): + bigquery.dynamic_schema(get_schema) + + def test_convert_to_beam_rows_creates_tables_with_specific_schemas( + self, mock_expansion_service): + """Test ConvertToBeamRows creates destination tables with specific schemas.""" + def dyn_schema(dest): + if 'users' in dest: + return 'id:INTEGER,name:STRING' + return 'id:INTEGER,score:INTEGER' + + dyn_schema._union_schema = 'id:INTEGER,name:STRING,score:INTEGER' + + created_tables = [] + + def mock_get_or_create_table( + project_id, dataset_id, table_id, schema, *args, **kwargs): + created_tables.append((table_id, [f.name for f in schema.fields])) + return mock.Mock() + + with mock.patch.object(bigquery_tools.BigQueryWrapper, + 'get_or_create_table', + side_effect=mock_get_or_create_table): + # Clear known tables to ensure fresh creation check + bigquery._KNOWN_TABLES.clear() + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, + dynamic_destinations=True, + create_disposition=bigquery.BigQueryDisposition.CREATE_IF_NEEDED) + + with TestPipeline() as p: + _ = ( + p + | beam.Create([ + ('project:ds.users', { + 'id': 1, 'name': 'alice' + }), + ('project:ds.scores', { + 'id': 2, 'score': 95 + }), + ('project:ds.users', { + 'id': 3, 'name': 'bob' + }), + ]) + | converter) + + # Verify get_or_create_table was called exactly twice (once per distinct destination) + self.assertEqual(len(created_tables), 2) + table_map = dict(created_tables) + # Verify 'users' was created with ['id', 'name'] (NOT union schema!) + self.assertEqual(table_map['users'], ['id', 'name']) + # Verify 'scores' was created with ['id', 'score'] (NOT union schema!) + self.assertEqual(table_map['scores'], ['id', 'score']) + + def test_convert_to_beam_rows_create_never_does_not_create_tables( + self, mock_expansion_service): + """Test ConvertToBeamRows does not call get_or_create_table when CREATE_NEVER.""" + def dyn_schema(dest): + return 'id:INTEGER,name:STRING' + + dyn_schema._union_schema = 'id:INTEGER,name:STRING' + + with mock.patch.object(bigquery_tools.BigQueryWrapper, + 'get_or_create_table') as mock_create: + bigquery._KNOWN_TABLES.clear() + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, + dynamic_destinations=True, + create_disposition=bigquery.BigQueryDisposition.CREATE_NEVER) + + with TestPipeline() as p: + _ = ( + p + | beam.Create([('project:ds.table', { + 'id': 1, 'name': 'alice' + })]) + | converter) + + mock_create.assert_not_called() + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/rag/ingestion/cloudsql_it_test.py b/sdks/python/apache_beam/ml/rag/ingestion/cloudsql_it_test.py index 1d4b988a5db0..f5a2a452d3d7 100644 --- a/sdks/python/apache_beam/ml/rag/ingestion/cloudsql_it_test.py +++ b/sdks/python/apache_beam/ml/rag/ingestion/cloudsql_it_test.py @@ -397,6 +397,9 @@ def verify_standard_operations( os.environ.get('EXPANSION_JARS'), "EXPANSION_JARS environment var is not provided, " "indicating that jars have not been built") +@unittest.skipUnless( + os.environ.get('ALLOYDB_PASSWORD'), + "ALLOYDB_PASSWORD environment var is not provided") class CloudSQLVectorWriterConfigTest(unittest.TestCase): def setUp(self): self.write_test_pipeline = TestPipeline(is_integration_test=True)