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 dd4c6d12bfc2..f61c7bd836d0 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json +++ b/.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 +} \ No newline at end of file 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 8b8cd389b3d8..afdc7f7012a8 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json +++ b/.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 } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java index b866b89e9879..05aca9813e55 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java @@ -23,13 +23,23 @@ 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; @@ -37,15 +47,23 @@ 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 { +public class PortableBigQueryDestinations + extends DynamicDestinations> { public static final String DESTINATION = "destination"; public static final String RECORD = "record"; + public static final String SCHEMA = "schema"; + + private static final ConcurrentHashMap JSON_SCHEMA_CACHE = + new ConcurrentHashMap<>(); + private @MonotonicNonNull RowStringInterpolator interpolator = null; private final @Nullable List primaryKey; private final RowFilter rowFilter; @@ -73,30 +91,48 @@ public PortableBigQueryDestinations(Schema rowSchema, BigQueryWriteConfiguration } @Override - public String getDestination(@Nullable ValueInSingleWindow element) { + public KV getDestination(@Nullable ValueInSingleWindow 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> getDestinationCoder() { + return KvCoder.of(StringUtf8Coder.of(), NullableCoder.of(StringUtf8Coder.of())); + } + @Override + public TableDestination getTable(KV 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 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 destination) { if (primaryKey != null) { return new TableConstraints() .setPrimaryKey(new TableConstraints.PrimaryKey().setColumns(primaryKey)); @@ -104,13 +140,77 @@ public TableDestination getTable(String destination) { 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 tableRow, @Nullable List 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 entry : tableRow.entrySet()) { + if (entry.getKey().equalsIgnoreCase(fieldName)) { + value = entry.getValue(); + break; + } + } + } + if (value == null) { + continue; + } + List 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 mapItem = (Map) item; + filteredList.add(filterTableRowBySchema(mapItem, subfields)); + } else if (item != null) { + filteredList.add(item); + } + } + value = filteredList; + } else if (value instanceof Map) { + @SuppressWarnings("unchecked") + Map mapValue = (Map) value; + value = filterTableRowBySchema(mapValue, subfields); + } + } + filtered.set(fieldName, value); + } + return filtered; + } + public SerializableFunction 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, 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 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); }; } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProviderTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProviderTest.java index 81789f784255..4bab7c14a6cc 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProviderTest.java +++ b/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 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 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 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); + } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinationsTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinationsTest.java new file mode 100644 index 000000000000..a8c9a867a102 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinationsTest.java @@ -0,0 +1,345 @@ +/* + * 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.providers; + +import static org.apache.beam.sdk.io.gcp.bigquery.providers.PortableBigQueryDestinations.DESTINATION; +import static org.apache.beam.sdk.io.gcp.bigquery.providers.PortableBigQueryDestinations.RECORD; +import static org.apache.beam.sdk.io.gcp.bigquery.providers.PortableBigQueryDestinations.SCHEMA; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; + +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.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.CoderProperties; +import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.joda.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link PortableBigQueryDestinations}. */ +@RunWith(JUnit4.class) +public class PortableBigQueryDestinationsTest { + + private static final Schema UNION_RECORD_SCHEMA = + Schema.builder() + .addNullableStringField("name") + .addNullableInt64Field("number") + .addNullableDoubleField("score") + .build(); + + private static final String SCHEMA_JSON_1 = + BigQueryHelpers.toJsonString( + new TableSchema() + .setFields( + Arrays.asList( + new TableFieldSchema().setName("name").setType("STRING").setMode("NULLABLE"), + new TableFieldSchema() + .setName("number") + .setType("INTEGER") + .setMode("NULLABLE")))); + + private static final String SCHEMA_JSON_2 = + BigQueryHelpers.toJsonString( + new TableSchema() + .setFields( + Arrays.asList( + new TableFieldSchema().setName("name").setType("STRING").setMode("NULLABLE"), + new TableFieldSchema() + .setName("score") + .setType("FLOAT") + .setMode("NULLABLE")))); + + @Test + public void testGetDestinationAndSchema_staticSchema() { + Schema wrapperSchema = + Schema.builder() + .addStringField(DESTINATION) + .addRowField(RECORD, UNION_RECORD_SCHEMA) + .build(); + + BigQueryWriteConfiguration config = + BigQueryWriteConfiguration.builder() + .setTable(BigQueryWriteConfiguration.DYNAMIC_DESTINATIONS) + .build(); + + PortableBigQueryDestinations destinations = + new PortableBigQueryDestinations(UNION_RECORD_SCHEMA, config); + + Row recordRow = + Row.withSchema(UNION_RECORD_SCHEMA) + .withFieldValue("name", "alice") + .withFieldValue("number", 1L) + .withFieldValue("score", null) + .build(); + Row wrapperRow = + Row.withSchema(wrapperSchema) + .withFieldValue(DESTINATION, "proj:ds.table1") + .withFieldValue(RECORD, recordRow) + .build(); + + ValueInSingleWindow windowedRow = + ValueInSingleWindow.of( + wrapperRow, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING); + + KV dest = destinations.getDestination(windowedRow); + assertEquals("proj:ds.table1", dest.getKey()); + assertNull(dest.getValue()); + + TableSchema resolvedSchema = destinations.getSchema(dest); + assertNotNull(resolvedSchema); + assertEquals(3, resolvedSchema.getFields().size()); + } + + @Test + public void testGetDestinationAndSchema_dynamicSchemaWithJson() { + Schema wrapperSchema = + Schema.builder() + .addStringField(DESTINATION) + .addStringField(SCHEMA) + .addRowField(RECORD, UNION_RECORD_SCHEMA) + .build(); + + BigQueryWriteConfiguration config = + BigQueryWriteConfiguration.builder() + .setTable(BigQueryWriteConfiguration.DYNAMIC_DESTINATIONS) + .build(); + + PortableBigQueryDestinations destinations = + new PortableBigQueryDestinations(UNION_RECORD_SCHEMA, config); + + Row recordRow = + Row.withSchema(UNION_RECORD_SCHEMA) + .withFieldValue("name", "alice") + .withFieldValue("number", 1L) + .withFieldValue("score", null) + .build(); + Row wrapperRow = + Row.withSchema(wrapperSchema) + .withFieldValue(DESTINATION, "proj:ds.table1") + .withFieldValue(SCHEMA, SCHEMA_JSON_1) + .withFieldValue(RECORD, recordRow) + .build(); + + ValueInSingleWindow windowedRow = + ValueInSingleWindow.of( + wrapperRow, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING); + + KV dest = destinations.getDestination(windowedRow); + assertEquals("proj:ds.table1", dest.getKey()); + assertEquals(SCHEMA_JSON_1, dest.getValue()); + + TableSchema resolvedSchema = destinations.getSchema(dest); + assertNotNull(resolvedSchema); + assertEquals(2, resolvedSchema.getFields().size()); + assertEquals("name", resolvedSchema.getFields().get(0).getName()); + assertEquals("number", resolvedSchema.getFields().get(1).getName()); + + Row wrapperRow2 = + Row.withSchema(wrapperSchema) + .withFieldValue(DESTINATION, "proj:ds.table2") + .withFieldValue(SCHEMA, SCHEMA_JSON_2) + .withFieldValue(RECORD, recordRow) + .build(); + ValueInSingleWindow windowedRow2 = + ValueInSingleWindow.of( + wrapperRow2, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING); + KV dest2 = destinations.getDestination(windowedRow2); + assertEquals("proj:ds.table2", dest2.getKey()); + assertEquals(SCHEMA_JSON_2, dest2.getValue()); + + TableSchema resolvedSchema2 = destinations.getSchema(dest2); + assertNotNull(resolvedSchema2); + assertEquals(2, resolvedSchema2.getFields().size()); + assertEquals("name", resolvedSchema2.getFields().get(0).getName()); + assertEquals("score", resolvedSchema2.getFields().get(1).getName()); + } + + @Test + public void testGetDestinationAndSchema_dynamicSchemaEmptyFallsBackToNull() { + Schema wrapperSchema = + Schema.builder() + .addStringField(DESTINATION) + .addNullableField(SCHEMA, Schema.FieldType.STRING) + .addRowField(RECORD, UNION_RECORD_SCHEMA) + .build(); + + BigQueryWriteConfiguration config = + BigQueryWriteConfiguration.builder() + .setTable(BigQueryWriteConfiguration.DYNAMIC_DESTINATIONS) + .build(); + + PortableBigQueryDestinations destinations = + new PortableBigQueryDestinations(UNION_RECORD_SCHEMA, config); + + Row recordRow = + Row.withSchema(UNION_RECORD_SCHEMA) + .withFieldValue("name", "alice") + .withFieldValue("number", 1L) + .withFieldValue("score", null) + .build(); + Row wrapperRow = + Row.withSchema(wrapperSchema) + .withFieldValue(DESTINATION, "proj:ds.table1") + .withFieldValue(SCHEMA, "") + .withFieldValue(RECORD, recordRow) + .build(); + + ValueInSingleWindow windowedRow = + ValueInSingleWindow.of( + wrapperRow, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING); + + KV dest = destinations.getDestination(windowedRow); + assertEquals("proj:ds.table1", dest.getKey()); + assertEquals("", dest.getValue()); + + // When dynamic schema is empty/unspecified, getSchema should return null so that + // StorageApiDynamicDestinationsTableRow fetches the existing table schema from SCHEMA_CACHE. + assertNull(destinations.getSchema(dest)); + } + + @Test + public void testFilterFormatFunction_filtersExtraUnionFields() { + Schema wrapperSchema = + Schema.builder() + .addStringField(DESTINATION) + .addStringField(SCHEMA) + .addRowField(RECORD, UNION_RECORD_SCHEMA) + .build(); + + BigQueryWriteConfiguration config = + BigQueryWriteConfiguration.builder() + .setTable(BigQueryWriteConfiguration.DYNAMIC_DESTINATIONS) + .build(); + + PortableBigQueryDestinations destinations = + new PortableBigQueryDestinations(UNION_RECORD_SCHEMA, config); + + SerializableFunction formatFn = destinations.getFilterFormatFunction(true); + + // Even if 'score' has a non-null default/padding value in the union Row, + // SCHEMA_JSON_1 only includes ['name', 'number'], so 'score' must be filtered out. + Row recordRow = + Row.withSchema(UNION_RECORD_SCHEMA) + .withFieldValue("name", "alice") + .withFieldValue("number", 1L) + .withFieldValue("score", 99.5) + .build(); + Row wrapperRow = + Row.withSchema(wrapperSchema) + .withFieldValue(DESTINATION, "proj:ds.table1") + .withFieldValue(SCHEMA, SCHEMA_JSON_1) + .withFieldValue(RECORD, recordRow) + .build(); + + TableRow tableRow = formatFn.apply(wrapperRow); + assertEquals("alice", tableRow.get("name")); + assertEquals("1", tableRow.get("number").toString()); + assertFalse(tableRow.containsKey("score")); + } + + @Test + public void testFilterTableRowBySchema_nestedAndRepeatedRecords() { + TableSchema nestedSchema = + new TableSchema() + .setFields( + Arrays.asList( + new TableFieldSchema().setName("id").setType("INTEGER"), + new TableFieldSchema() + .setName("details") + .setType("RECORD") + .setFields( + Collections.singletonList( + new TableFieldSchema().setName("keep_field").setType("STRING"))), + new TableFieldSchema() + .setName("items") + .setType("RECORD") + .setMode("REPEATED") + .setFields( + Collections.singletonList( + new TableFieldSchema().setName("item_id").setType("INTEGER"))))); + + TableRow rawDetails = new TableRow().set("keep_field", "kept").set("drop_field", "dropped"); + TableRow item1 = new TableRow().set("item_id", 10).set("extra_item_field", "dropped"); + TableRow item2 = new TableRow().set("item_id", 20).set("extra_item_field", "dropped"); + + TableRow rawRow = + new TableRow() + .set("id", 1) + .set("extra_top_field", "dropped") + .set("details", rawDetails) + .set("items", Arrays.asList(item1, item2)); + + TableRow filtered = + PortableBigQueryDestinations.filterTableRowBySchema(rawRow, nestedSchema.getFields()); + + assertEquals(1, filtered.get("id")); + assertFalse(filtered.containsKey("extra_top_field")); + + TableRow filteredDetails = (TableRow) filtered.get("details"); + assertNotNull(filteredDetails); + assertEquals("kept", filteredDetails.get("keep_field")); + assertFalse(filteredDetails.containsKey("drop_field")); + + @SuppressWarnings("unchecked") + List filteredItems = (List) filtered.get("items"); + assertEquals(2, filteredItems.size()); + assertEquals(10, filteredItems.get(0).get("item_id")); + assertFalse(filteredItems.get(0).containsKey("extra_item_field")); + assertEquals(20, filteredItems.get(1).get("item_id")); + assertFalse(filteredItems.get(1).containsKey("extra_item_field")); + } + + @Test + public void testDestinationCoder() throws Exception { + BigQueryWriteConfiguration config = + BigQueryWriteConfiguration.builder() + .setTable(BigQueryWriteConfiguration.DYNAMIC_DESTINATIONS) + .build(); + PortableBigQueryDestinations destinations = + new PortableBigQueryDestinations(UNION_RECORD_SCHEMA, config); + + Coder> coder = destinations.getDestinationCoder(); + assertNotNull(coder); + coder.verifyDeterministic(); + + KV withSchema = KV.of("proj:ds.table1", SCHEMA_JSON_1); + KV withNullSchema = KV.of("proj:ds.table2", null); + KV withEmptySchema = KV.of("proj:ds.table3", ""); + + CoderProperties.coderDecodeEncodeEqual(coder, withSchema); + CoderProperties.coderDecodeEncodeEqual(coder, withNullSchema); + CoderProperties.coderDecodeEncodeEqual(coder, withEmptySchema); + CoderProperties.coderDeterministic(coder, withSchema, withSchema); + CoderProperties.coderDeterministic(coder, withNullSchema, withNullSchema); + } +} 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 38acd29da7d9..8782fec81532 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 @@ -366,6 +434,7 @@ def chain_after(result): import uuid import warnings from dataclasses import dataclass +from dataclasses import dataclass from enum import Enum from typing import Optional from typing import Union @@ -2013,6 +2082,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. @@ -2474,6 +2657,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, @@ -2485,7 +2669,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}") @@ -2703,7 +2888,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. """ @@ -2713,6 +2898,7 @@ class StorageWriteToBigQuery(PTransform): # fields for rows sent to Storage API with dynamic destinations DESTINATION = "destination" RECORD = "record" + SCHEMA = "schema" # field names for rows sent to Storage API for CDC functionality CDC_INFO = "row_mutation_info" CDC_MUTATION_TYPE = "mutation_type" @@ -2725,6 +2911,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, @@ -2736,10 +2923,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 @@ -2753,6 +2942,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: @@ -2764,9 +2954,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 @@ -2778,13 +2967,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. @@ -2816,7 +3015,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 @@ -2883,25 +3089,175 @@ 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) + + if record_schema is None or record_schema == SCHEMA_AUTODETECT: + schema_json = '' + else: + schema_dict = bigquery_tools.get_dict_table_schema(record_schema) + schema_json = json.dumps(schema_dict) + + 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.SCHEMA: schema_json, + 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 @@ -2912,17 +3268,21 @@ 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)) - ]) + if callable(self.schema): + type_hint = RowTypeConstraint.from_fields([ + (StorageWriteToBigQuery.DESTINATION, str), + (StorageWriteToBigQuery.SCHEMA, str), + (StorageWriteToBigQuery.RECORD, record_hint) + ]) + else: + type_hint = RowTypeConstraint.from_fields([ + (StorageWriteToBigQuery.DESTINATION, str), + (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..3f4eab8e5ba9 --- /dev/null +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -0,0 +1,436 @@ +# +# 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 json +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', + schema=json.dumps(bigquery_tools.get_dict_table_schema(schema1)), + record=beam.Row(id=1, name='foo')), + beam.Row( + destination='table2', + schema=json.dumps(bigquery_tools.get_dict_table_schema(schema2)), + 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', + schema=json.dumps(bigquery_tools.get_dict_table_schema(schema1)), + record=beam.Row(id=1, name='foo')), + beam.Row( + destination='table2', + schema=json.dumps(bigquery_tools.get_dict_table_schema(schema2)), + 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.SCHEMA, 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], (bigquery.StorageWriteToBigQuery.SCHEMA, str)) + self.assertEqual( + type_hint_dyn._fields[2][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[2][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.schema == json.dumps( + bigquery_tools.get_dict_table_schema('id:INTEGER,name:STRING')) + assert r1.record.id == 1 + assert r1.record.name == 'alice' + assert r1.record.score is None + assert r2.destination == 'dest_scores' + assert r2.schema == json.dumps( + bigquery_tools.get_dict_table_schema('id:INTEGER,score:INTEGER')) + 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)