From 2b4586554f6b82433ef84d29ce24340728103034 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 7 Jul 2026 11:09:35 -0400 Subject: [PATCH 01/11] [WIP] Add Dynamic Schema support to StorageWriteToBigQuery --- .../io/external/xlang_bigqueryio_it_test.py | 50 ++++ sdks/python/apache_beam/io/gcp/bigquery.py | 100 ++++++-- .../io/gcp/bigquery_storage_write_test.py | 227 ++++++++++++++++++ 3 files changed, 352 insertions(+), 25 deletions(-) create mode 100644 sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py 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..5c771de23c7a 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 @@ -483,6 +483,56 @@ 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) + tables = [base_table_spec + str(record['int']) for record in self.ELEMENTS] + + bq_matchers = [ + BigqueryFullResultMatcher( + project=self.project, + query="SELECT * FROM %s" % tables[i], + data=self.parse_expected_data(self.ELEMENTS[i])) + for i in range(len(tables)) + ] + + with beam.Pipeline(argv=self.args) as p: + schema_pc = p | "CreateSchema" >> beam.Create([self.ALL_TYPES_SCHEMA]) + _ = ( + p + | "CreateElements" >> beam.Create(self.ELEMENTS) + | beam.io.WriteToBigQuery( + table=lambda record: spec_with_project + str(record['int']), + method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=lambda dest, schema: schema, + schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc),), + use_at_least_once=False)) + hamcrest_assert(p, all_of(*bq_matchers)) + + def test_write_with_dynamic_schema(self): + table = 'write_with_dynamic_schema' + table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) + + bq_matcher = BigqueryFullResultMatcher( + project=self.project, + query="SELECT * FROM {}.{}".format(self.dataset_id, table), + data=self.parse_expected_data(self.ELEMENTS)) + + with beam.Pipeline(argv=self.args) as p: + schema_pc = p | "CreateSchema" >> beam.Create([self.ALL_TYPES_SCHEMA]) + _ = ( + p + | "Create test data" >> beam.Create(self.ELEMENTS) + | beam.io.WriteToBigQuery( + table=table_id, + method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=lambda dest, schema: schema, + schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc),), + create_disposition='CREATE_IF_NEEDED', + write_disposition='WRITE_TRUNCATE')) + + hamcrest_assert(p, bq_matcher) + 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 a2d17f12569e..26f8e9999fae 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -366,6 +366,7 @@ def chain_after(result): import uuid import warnings from dataclasses import dataclass +from apache_beam.typehints.typehints import Any from typing import Optional from typing import Union @@ -2388,6 +2389,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, @@ -2638,6 +2640,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, @@ -2651,8 +2654,9 @@ def __init__( expansion_service=None, type_overrides=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 @@ -2677,9 +2681,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 @@ -2697,7 +2700,11 @@ def expand(self, input): input_beam_rows = ( input | "Convert dict to Beam Row" >> self.ConvertToBeamRows( - schema, False, self._type_overrides).with_output_types()) + schema, + False, + self._type_overrides, + schema_side_inputs=self._schema_side_inputs, + destination=table).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 +2736,10 @@ 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).with_output_types()) # communicate to Java that this write should use dynamic destinations table = StorageWriteToBigQuery.DYNAMIC_DESTINATIONS @@ -2797,34 +2807,74 @@ def __exit__(self, *args): pass 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, + destination=None): self.schema = schema self.dynamic_destinations = dynamic_destinations self.type_overrides = type_overrides + self.schema_side_inputs = schema_side_inputs or () + self.destination = destination 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): + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.Map( + lambda row, *schema_side_inputs: beam.Row( + **{ + StorageWriteToBigQuery.DESTINATION: row[0], + StorageWriteToBigQuery.RECORD: bigquery_tools. + beam_row_from_dict( + row[1], self.schema(row[0], *schema_side_inputs)) + }), + *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 - | "Convert dict to Beam Row" >> beam.Map( - lambda row, schema=DoFn.SetupContextParam( - StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args=[ - self.schema - ]): bigquery_tools.beam_row_from_dict(row, schema))) + if callable(self.schema): + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.Map( + lambda row, *schema_side_inputs: bigquery_tools. + beam_row_from_dict( + row, self.schema(self.destination, *schema_side_inputs)), + *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]): bigquery_tools.beam_row_from_dict( + row, schema))) def with_output_types(self): + if callable(self.schema): + if self.dynamic_destinations: + type_hint = RowTypeConstraint.from_fields([ + (StorageWriteToBigQuery.DESTINATION, str), + (StorageWriteToBigQuery.RECORD, Any) + ]) + else: + type_hint = Any + return super().with_output_types(type_hint) + row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( self.schema, self.type_overrides) if self.dynamic_destinations: 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..ae7845725bb2 --- /dev/null +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -0,0 +1,227 @@ +# +# 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 + +from apache_beam.typehints.typehints import Any + +import apache_beam as beam +from apache_beam.io.gcp import bigquery +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 + + +@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_convert_to_beam_rows_static_destination_with_side_inputs( + self, mock_expansion_service): + """Test ConvertToBeamRows with static destination and dynamic schema.""" + schema1 = { + 'fields': [ + {'name': 'id', 'type': 'INTEGER'}, + {'name': 'name', 'type': 'STRING'}, + ] + } + + with TestPipeline() as p: + side_pcoll = ( + p + | "CreateSide" >> beam.Create([{'proj:ds.table': schema1}]) + ) + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest, side_map: side_map[dest], + dynamic_destinations=False, + schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll),), + destination='proj:ds.table') + + input_data = [ + {'id': 1, 'name': 'foo'}, + {'id': 2, 'name': 'bar'}, + ] + res = p | "CreateInput" >> beam.Create(input_data) | converter + + expected_rows = [ + beam.Row(id=1, name='foo'), + beam.Row(id=2, name='bar'), + ] + assert_that(res, equal_to(expected_rows)) + + 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, Any), + )) + + converter_static = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest: None, + dynamic_destinations=False) + type_hint_static = converter_static.with_output_types().get_type_hints( + ).simple_output_type('') + self.assertEqual(type_hint_static, Any) + + 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) + + +if __name__ == '__main__': + unittest.main() From 7808954678eed0661ed94faf552fa7964e1245ec Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 7 Jul 2026 11:14:38 -0400 Subject: [PATCH 02/11] yapf --- .../io/external/xlang_bigqueryio_it_test.py | 4 +- sdks/python/apache_beam/io/gcp/bigquery.py | 9 +- .../io/gcp/bigquery_storage_write_test.py | 120 +++++++++++------- 3 files changed, 84 insertions(+), 49 deletions(-) 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 5c771de23c7a..fe5ed3e3443e 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 @@ -505,7 +505,7 @@ def test_write_to_dynamic_destinations_with_dynamic_schema(self): table=lambda record: spec_with_project + str(record['int']), method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=lambda dest, schema: schema, - schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc),), + schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) @@ -527,7 +527,7 @@ def test_write_with_dynamic_schema(self): table=table_id, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=lambda dest, schema: schema, - schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc),), + schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), create_disposition='CREATE_IF_NEEDED', write_disposition='WRITE_TRUNCATE')) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 26f8e9999fae..6c7efd70711d 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -366,7 +366,6 @@ def chain_after(result): import uuid import warnings from dataclasses import dataclass -from apache_beam.typehints.typehints import Any from typing import Optional from typing import Union @@ -418,6 +417,7 @@ def chain_after(result): from apache_beam.transforms.util import ReshufflePerKey from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.schemas import schema_from_element_type +from apache_beam.typehints.typehints import Any from apache_beam.utils import retry from apache_beam.utils.annotations import deprecated @@ -2739,7 +2739,8 @@ def expand(self, input): schema, True, self._type_overrides, - schema_side_inputs=self._schema_side_inputs).with_output_types()) + schema_side_inputs=self._schema_side_inputs).with_output_types( + )) # communicate to Java that this write should use dynamic destinations table = StorageWriteToBigQuery.DYNAMIC_DESTINATIONS @@ -2828,8 +2829,8 @@ def expand(self, input_dicts): | "Convert dict to Beam Row" >> beam.Map( lambda row, *schema_side_inputs: beam.Row( **{ - StorageWriteToBigQuery.DESTINATION: row[0], - StorageWriteToBigQuery.RECORD: bigquery_tools. + StorageWriteToBigQuery.DESTINATION: row[ + 0], StorageWriteToBigQuery.RECORD: bigquery_tools. beam_row_from_dict( row[1], self.schema(row[0], *schema_side_inputs)) }), 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 index ae7845725bb2..b971989e573e 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -33,15 +33,14 @@ @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',)) + 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( @@ -49,26 +48,37 @@ def test_convert_to_beam_rows_dynamic_destinations_dynamic_schema( """Test ConvertToBeamRows with dynamic destinations and dynamic schema.""" schema1 = { 'fields': [ - {'name': 'id', 'type': 'INTEGER'}, - {'name': 'name', 'type': 'STRING'}, + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, ] } schema2 = { 'fields': [ - {'name': 'id', 'type': 'INTEGER'}, - {'name': 'score', 'type': 'FLOAT'}, + { + '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) + 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}), + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), ] res = p | "CreateInput" >> beam.Create(input_data) | converter @@ -83,31 +93,44 @@ def test_convert_to_beam_rows_dynamic_destinations_with_side_inputs( """Test ConvertToBeamRows with dynamic schema and side inputs.""" schema1 = { 'fields': [ - {'name': 'id', 'type': 'INTEGER'}, - {'name': 'name', 'type': 'STRING'}, + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, ] } schema2 = { 'fields': [ - {'name': 'id', 'type': 'INTEGER'}, - {'name': 'score', 'type': 'FLOAT'}, + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'score', 'type': 'FLOAT' + }, ] } with TestPipeline() as p: side_pcoll = ( p - | "CreateSide" >> beam.Create([{'table1': schema1, 'table2': schema2}]) - ) + | "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),)) + schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll), )) input_data = [ - ('table1', {'id': 1, 'name': 'foo'}), - ('table2', {'id': 2, 'score': 3.14}), + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), ] res = p | "CreateInput" >> beam.Create(input_data) | converter @@ -122,26 +145,35 @@ def test_convert_to_beam_rows_static_destination_with_side_inputs( """Test ConvertToBeamRows with static destination and dynamic schema.""" schema1 = { 'fields': [ - {'name': 'id', 'type': 'INTEGER'}, - {'name': 'name', 'type': 'STRING'}, + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, ] } with TestPipeline() as p: side_pcoll = ( p - | "CreateSide" >> beam.Create([{'proj:ds.table': schema1}]) - ) + | "CreateSide" >> beam.Create([{ + 'proj:ds.table': schema1 + }])) converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( schema=lambda dest, side_map: side_map[dest], dynamic_destinations=False, - schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll),), + schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll), ), destination='proj:ds.table') input_data = [ - {'id': 1, 'name': 'foo'}, - {'id': 2, 'name': 'bar'}, + { + 'id': 1, 'name': 'foo' + }, + { + 'id': 2, 'name': 'bar' + }, ] res = p | "CreateInput" >> beam.Create(input_data) | converter @@ -155,19 +187,19 @@ 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) + 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, Any), - )) + self.assertEqual( + type_hint_dyn._fields, + ( + (bigquery.StorageWriteToBigQuery.DESTINATION, str), + (bigquery.StorageWriteToBigQuery.RECORD, Any), + )) converter_static = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( - schema=lambda dest: None, - dynamic_destinations=False) + schema=lambda dest: None, dynamic_destinations=False) type_hint_static = converter_static.with_output_types().get_type_hints( ).simple_output_type('') self.assertEqual(type_hint_static, Any) @@ -177,7 +209,9 @@ def test_storage_write_to_bigquery_expand_dynamic_schema( """Test StorageWriteToBigQuery expand does not fail for callable schema.""" schema1 = { 'fields': [ - {'name': 'id', 'type': 'INTEGER'}, + { + 'name': 'id', 'type': 'INTEGER' + }, ] } @@ -188,12 +222,12 @@ def expand(self, pcoll): pcoll.pipeline | "CreateErrors" >> beam.Create([])) } - with mock.patch.object( - bigquery, 'SchemaAwareExternalTransform', autospec=True) as mock_ext: + 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) + table=lambda record: 'table1', schema=lambda dest: schema1) with TestPipeline() as p: _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | transform @@ -206,8 +240,8 @@ def expand(self, pcoll): 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: + 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]) @@ -215,7 +249,7 @@ def test_write_to_bigquery_storage_api_passes_schema_side_inputs( table='proj:ds.table', method=bigquery.WriteToBigQuery.Method.STORAGE_WRITE_API, schema=lambda dest: None, - schema_side_inputs=(beam.pvalue.AsSingleton(side_pc),)) + schema_side_inputs=(beam.pvalue.AsSingleton(side_pc), )) _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | write_transform mock_storage_write.assert_called_once() From 1d2ceca5a2e62c3b92c323b37b29d66e537b9245 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 7 Jul 2026 11:20:31 -0400 Subject: [PATCH 03/11] Trigger xlang integration tests --- .../beam_PostCommit_Python_Xlang_Gcp_Dataflow.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 } From 85b85464a56ddd3dc2a9130de3b905e911ee9eae Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 7 Jul 2026 11:43:17 -0400 Subject: [PATCH 04/11] remove dynamic schema + static table case --- .../io/external/xlang_bigqueryio_it_test.py | 24 -------- sdks/python/apache_beam/io/gcp/bigquery.py | 48 ++++++--------- .../io/gcp/bigquery_storage_write_test.py | 60 ++++--------------- 3 files changed, 28 insertions(+), 104 deletions(-) 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 fe5ed3e3443e..f0f78ceceb5a 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 @@ -509,30 +509,6 @@ def test_write_to_dynamic_destinations_with_dynamic_schema(self): use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) - def test_write_with_dynamic_schema(self): - table = 'write_with_dynamic_schema' - table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) - - bq_matcher = BigqueryFullResultMatcher( - project=self.project, - query="SELECT * FROM {}.{}".format(self.dataset_id, table), - data=self.parse_expected_data(self.ELEMENTS)) - - with beam.Pipeline(argv=self.args) as p: - schema_pc = p | "CreateSchema" >> beam.Create([self.ALL_TYPES_SCHEMA]) - _ = ( - p - | "Create test data" >> beam.Create(self.ELEMENTS) - | beam.io.WriteToBigQuery( - table=table_id, - method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, - schema=lambda dest, schema: schema, - schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), - create_disposition='CREATE_IF_NEEDED', - write_disposition='WRITE_TRUNCATE')) - - hamcrest_assert(p, bq_matcher) - 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 6c7efd70711d..b08cc75b7c02 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -2694,17 +2694,17 @@ 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, - schema_side_inputs=self._schema_side_inputs, - destination=table).with_output_types()) + schema, False, self._type_overrides).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. @@ -2813,13 +2813,11 @@ def __init__( schema, dynamic_destinations, type_overrides=None, - schema_side_inputs=None, - destination=None): + schema_side_inputs=None): self.schema = schema self.dynamic_destinations = dynamic_destinations self.type_overrides = type_overrides self.schema_side_inputs = schema_side_inputs or () - self.destination = destination def expand(self, input_dicts): if self.dynamic_destinations: @@ -2848,32 +2846,20 @@ def expand(self, input_dicts): beam_row_from_dict(row[1], schema) }))) else: - if callable(self.schema): - return ( - input_dicts - | "Convert dict to Beam Row" >> beam.Map( - lambda row, *schema_side_inputs: bigquery_tools. - beam_row_from_dict( - row, self.schema(self.destination, *schema_side_inputs)), - *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]): bigquery_tools.beam_row_from_dict( - row, schema))) + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.Map( + lambda row, schema=DoFn.SetupContextParam( + StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args=[ + self.schema + ]): bigquery_tools.beam_row_from_dict(row, schema))) def with_output_types(self): if callable(self.schema): - if self.dynamic_destinations: - type_hint = RowTypeConstraint.from_fields([ - (StorageWriteToBigQuery.DESTINATION, str), - (StorageWriteToBigQuery.RECORD, Any) - ]) - else: - type_hint = Any + type_hint = RowTypeConstraint.from_fields([ + (StorageWriteToBigQuery.DESTINATION, str), + (StorageWriteToBigQuery.RECORD, Any) + ]) return super().with_output_types(type_hint) row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( 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 index b971989e573e..a00ffe2f8e8f 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -20,14 +20,13 @@ import unittest from unittest import mock -from apache_beam.typehints.typehints import Any - import apache_beam as beam from apache_beam.io.gcp import bigquery 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 +from apache_beam.typehints.typehints import Any @mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') @@ -140,48 +139,17 @@ def test_convert_to_beam_rows_dynamic_destinations_with_side_inputs( ] assert_that(res, equal_to(expected_rows)) - def test_convert_to_beam_rows_static_destination_with_side_inputs( + def test_storage_write_static_destination_dynamic_schema_raises_error( self, mock_expansion_service): - """Test ConvertToBeamRows with static destination and dynamic schema.""" - schema1 = { - 'fields': [ - { - 'name': 'id', 'type': 'INTEGER' - }, - { - 'name': 'name', 'type': 'STRING' - }, - ] - } - - with TestPipeline() as p: - side_pcoll = ( - p - | "CreateSide" >> beam.Create([{ - 'proj:ds.table': schema1 - }])) - - converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( - schema=lambda dest, side_map: side_map[dest], - dynamic_destinations=False, - schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll), ), - destination='proj:ds.table') - - input_data = [ - { - 'id': 1, 'name': 'foo' - }, - { - 'id': 2, 'name': 'bar' - }, - ] - res = p | "CreateInput" >> beam.Create(input_data) | converter - - expected_rows = [ - beam.Row(id=1, name='foo'), - beam.Row(id=2, name='bar'), - ] - assert_that(res, equal_to(expected_rows)) + """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): @@ -198,12 +166,6 @@ def test_convert_to_beam_rows_with_output_types_dynamic_schema( (bigquery.StorageWriteToBigQuery.RECORD, Any), )) - converter_static = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( - schema=lambda dest: None, dynamic_destinations=False) - type_hint_static = converter_static.with_output_types().get_type_hints( - ).simple_output_type('') - self.assertEqual(type_hint_static, Any) - def test_storage_write_to_bigquery_expand_dynamic_schema( self, mock_expansion_service): """Test StorageWriteToBigQuery expand does not fail for callable schema.""" From 8eb3754e7103411a12a6e1b1630066f1868a0fca Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 7 Jul 2026 11:46:06 -0400 Subject: [PATCH 05/11] Fix StorageWriteToBigQuery docstring --- sdks/python/apache_beam/io/gcp/bigquery.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index b08cc75b7c02..f531df19fcc5 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -2618,7 +2618,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. """ From d0c9aa673e02822252e0617bdb45e0aee2451099 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 11:20:54 -0400 Subject: [PATCH 06/11] Make IT test more robust --- .../io/external/xlang_bigqueryio_it_test.py | 51 ++++++++++++++++--- 1 file changed, 43 insertions(+), 8 deletions(-) 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 f0f78ceceb5a..7abe43bff7ea 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 @@ -486,25 +486,60 @@ def test_write_to_dynamic_destinations(self): 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) - tables = [base_table_spec + str(record['int']) for record in self.ELEMENTS] + 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" + + 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" % tables[i], - data=self.parse_expected_data(self.ELEMENTS[i])) - for i in range(len(tables)) + 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' + with beam.Pipeline(argv=self.args) as p: - schema_pc = p | "CreateSchema" >> beam.Create([self.ALL_TYPES_SCHEMA]) + schema_pc = p | "CreateSchema" >> beam.Create([schema_map]) _ = ( p - | "CreateElements" >> beam.Create(self.ELEMENTS) + | "CreateElements" >> beam.Create(elements) | beam.io.WriteToBigQuery( - table=lambda record: spec_with_project + str(record['int']), + table=get_destination, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, - schema=lambda dest, schema: schema, + schema=lambda dest, side_map: side_map[dest], schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) From 9f7f24bbd1546fac3da248930312685293a0f27c Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 11:43:36 -0400 Subject: [PATCH 07/11] Skip unit tests if GCP dependencies are not installed --- .../apache_beam/io/gcp/bigquery_storage_write_test.py | 6 ++++++ 1 file changed, 6 insertions(+) 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 index a00ffe2f8e8f..fc100e0644ea 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -28,7 +28,13 @@ from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.typehints import Any +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.""" From 19c296a61db5c5cc159e2bdcd9face74904dfbe7 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 8 Jul 2026 12:07:09 -0400 Subject: [PATCH 08/11] yapf --- sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 index fc100e0644ea..cbb7a7135d22 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -34,7 +34,8 @@ GoogleAPICallError = None -@unittest.skipIf(GoogleAPICallError is None, 'GCP dependencies are not installed') +@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.""" From 9aface19854dac5c7748f912899ebc70a283000a Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 9 Jul 2026 09:27:19 -0400 Subject: [PATCH 09/11] hint update for coders --- sdks/python/apache_beam/io/gcp/bigquery.py | 3 +-- .../python/apache_beam/io/gcp/bigquery_storage_write_test.py | 5 +++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index f531df19fcc5..f75a2071144d 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -417,7 +417,6 @@ def chain_after(result): from apache_beam.transforms.util import ReshufflePerKey from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.typehints.schemas import schema_from_element_type -from apache_beam.typehints.typehints import Any from apache_beam.utils import retry from apache_beam.utils.annotations import deprecated @@ -2858,7 +2857,7 @@ def with_output_types(self): if callable(self.schema): type_hint = RowTypeConstraint.from_fields([ (StorageWriteToBigQuery.DESTINATION, str), - (StorageWriteToBigQuery.RECORD, Any) + (StorageWriteToBigQuery.RECORD, RowTypeConstraint.from_fields([])) ]) 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 index cbb7a7135d22..779fe73978dc 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -26,7 +26,6 @@ from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.typehints.row_type import RowTypeConstraint -from apache_beam.typehints.typehints import Any try: from google.api_core.exceptions import GoogleAPICallError @@ -170,7 +169,9 @@ def test_convert_to_beam_rows_with_output_types_dynamic_schema( type_hint_dyn._fields, ( (bigquery.StorageWriteToBigQuery.DESTINATION, str), - (bigquery.StorageWriteToBigQuery.RECORD, Any), + ( + bigquery.StorageWriteToBigQuery.RECORD, + RowTypeConstraint.from_fields([])), )) def test_storage_write_to_bigquery_expand_dynamic_schema( From 8645501c0d15ddf0cbc44cbac6b26a43138a51fb Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 9 Jul 2026 13:23:54 -0400 Subject: [PATCH 10/11] union schema fix --- .../io/external/xlang_bigqueryio_it_test.py | 8 ++++++- sdks/python/apache_beam/io/gcp/bigquery.py | 24 ++++++++++++++++++- .../io/gcp/bigquery_storage_write_test.py | 24 +++++++++++++++++++ 3 files changed, 54 insertions(+), 2 deletions(-) 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 7abe43bff7ea..7e7798cd0080 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 @@ -531,6 +531,12 @@ def get_destination(record): return spec_with_project + 'users' return spec_with_project + 'scores' + def get_schema(dest, side_map): + return side_map[dest] + + get_schema._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]) _ = ( @@ -539,7 +545,7 @@ def get_destination(record): | beam.io.WriteToBigQuery( table=get_destination, method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, - schema=lambda dest, side_map: side_map[dest], + schema=get_schema, schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index f75a2071144d..989667cdfb44 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -2855,9 +2855,31 @@ def expand(self, input_dicts): def with_output_types(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) + record_hint = RowTypeConstraint.from_fields(row_type_hints) + elif isinstance(schema_hint, RowTypeConstraint): + record_hint = schema_hint + else: + record_hint = RowTypeConstraint.from_fields([]) + else: + record_hint = RowTypeConstraint.from_fields([]) + type_hint = RowTypeConstraint.from_fields([ (StorageWriteToBigQuery.DESTINATION, str), - (StorageWriteToBigQuery.RECORD, RowTypeConstraint.from_fields([])) + (StorageWriteToBigQuery.RECORD, 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 index 779fe73978dc..3900182ba102 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -22,6 +22,7 @@ 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 @@ -174,6 +175,29 @@ def test_convert_to_beam_rows_with_output_types_dynamic_schema( 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_storage_write_to_bigquery_expand_dynamic_schema( self, mock_expansion_service): """Test StorageWriteToBigQuery expand does not fail for callable schema.""" From ab39151f6ba847e577c66567eca26af66f23aa45 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Fri, 10 Jul 2026 11:38:51 -0400 Subject: [PATCH 11/11] Another coder fix --- .../io/external/xlang_bigqueryio_it_test.py | 4 +- sdks/python/apache_beam/io/gcp/bigquery.py | 95 ++++++++++--------- .../io/gcp/bigquery_storage_write_test.py | 42 ++++++++ 3 files changed, 96 insertions(+), 45 deletions(-) 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 7e7798cd0080..2609468c2d23 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 @@ -518,11 +518,11 @@ def test_write_to_dynamic_destinations_with_dynamic_schema(self): bq_matchers = [ BigqueryFullResultMatcher( project=self.project, - query="SELECT * FROM %s" % table_a, + query="SELECT id, name FROM %s" % table_a, data=self.parse_expected_data(elements_a)), BigqueryFullResultMatcher( project=self.project, - query="SELECT * FROM %s" % table_b, + query="SELECT id, score, active FROM %s" % table_b, data=self.parse_expected_data(elements_b)), ] diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 989667cdfb44..701fbc1fa39f 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -2818,20 +2818,61 @@ def __init__( self.type_overrides = type_overrides self.schema_side_inputs = schema_side_inputs or () + 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: if callable(self.schema): + record_hint = self._get_record_type_hint() + union_field_names = [ + name for name, _ in getattr(record_hint, '_fields', ()) + ] + + def convert_dynamic_row(row, *schema_side_inputs): + dest, dict_row = row[0], row[1] + record_schema = self.schema(dest, *schema_side_inputs) + record_row = bigquery_tools.beam_row_from_dict( + dict_row, record_schema) + if union_field_names: + record_dict = record_row._asdict() + record_row = beam.Row( + **{ + name: record_dict.get(name, None) + for name in union_field_names + }) + return beam.Row( + **{ + StorageWriteToBigQuery.DESTINATION: dest, + StorageWriteToBigQuery.RECORD: record_row + }) + return ( input_dicts | "Convert dict to Beam Row" >> beam.Map( - lambda row, *schema_side_inputs: beam.Row( - **{ - StorageWriteToBigQuery.DESTINATION: row[ - 0], StorageWriteToBigQuery.RECORD: bigquery_tools. - beam_row_from_dict( - row[1], self.schema(row[0], *schema_side_inputs)) - }), - *self.schema_side_inputs)) + convert_dynamic_row, *self.schema_side_inputs)) else: return ( input_dicts @@ -2854,46 +2895,14 @@ def expand(self, input_dicts): ]): bigquery_tools.beam_row_from_dict(row, schema))) def with_output_types(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) - record_hint = RowTypeConstraint.from_fields(row_type_hints) - elif isinstance(schema_hint, RowTypeConstraint): - record_hint = schema_hint - else: - record_hint = RowTypeConstraint.from_fields([]) - else: - record_hint = RowTypeConstraint.from_fields([]) - - type_hint = RowTypeConstraint.from_fields([ - (StorageWriteToBigQuery.DESTINATION, str), - (StorageWriteToBigQuery.RECORD, record_hint) - ]) - return super().with_output_types(type_hint) - - 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 index 3900182ba102..a84edd3c7ba1 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -198,6 +198,48 @@ def dyn_schema(dest): 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."""