diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json index 83346d34aee0..48379c82f94c 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 16 + "modification": 21 } diff --git a/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..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 @@ -483,6 +483,73 @@ 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_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 id, name FROM %s" % table_a, + data=self.parse_expected_data(elements_a)), + BigqueryFullResultMatcher( + project=self.project, + query="SELECT id, score, active 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(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]) + _ = ( + 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)) + 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..701fbc1fa39f 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -2388,6 +2388,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, @@ -2616,7 +2617,7 @@ def __getitem__(self, key): class StorageWriteToBigQuery(PTransform): """Writes data to BigQuery using Storage API. - Supports dynamic destinations. Dynamic schemas are not supported yet. + Supports dynamic destinations and dynamic schemas. Experimental; no backwards compatibility guarantees. """ @@ -2638,6 +2639,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 +2653,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 +2680,8 @@ def expand(self, input): "A schema is required in order to prepare rows " "for writing with STORAGE_WRITE_API.") from exn elif callable(self._schema): - raise NotImplementedError( - "Writing with dynamic schemas is not " - "supported for this write method.") + schema = self._schema + is_rows = False elif isinstance(self._schema, vp.ValueProvider): schema = self._schema.get() is_rows = False @@ -2691,6 +2693,10 @@ 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: @@ -2729,7 +2735,11 @@ 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,24 +2807,84 @@ 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): self.schema = schema self.dynamic_destinations = dynamic_destinations 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: - 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', ()) + ] + + 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( + convert_dynamic_row, *self.schema_side_inputs)) + else: + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.Map( + lambda row, schema=DoFn.SetupContextParam( + StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args= + [self.schema]): beam.Row( + **{ + StorageWriteToBigQuery.DESTINATION: row[0], + StorageWriteToBigQuery.RECORD: bigquery_tools. + beam_row_from_dict(row[1], schema) + }))) else: return ( input_dicts @@ -2825,17 +2895,14 @@ def expand(self, input_dicts): ]): bigquery_tools.beam_row_from_dict(row, schema))) def with_output_types(self): - row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( - self.schema, self.type_overrides) + record_hint = self._get_record_type_hint() if self.dynamic_destinations: type_hint = RowTypeConstraint.from_fields([ (StorageWriteToBigQuery.DESTINATION, str), - ( - StorageWriteToBigQuery.RECORD, - RowTypeConstraint.from_fields(row_type_hints)) + (StorageWriteToBigQuery.RECORD, record_hint) ]) else: - type_hint = RowTypeConstraint.from_fields(row_type_hints) + type_hint = record_hint return super().with_output_types(type_hint) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py new file mode 100644 index 000000000000..a84edd3c7ba1 --- /dev/null +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -0,0 +1,297 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Unit tests for BigQuery Storage Write API dynamic schemas.""" + +import unittest +from unittest import mock + +import apache_beam as beam +from apache_beam.io.gcp import bigquery +from apache_beam.io.gcp import bigquery_tools +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to +from apache_beam.typehints.row_type import RowTypeConstraint + +try: + from google.api_core.exceptions import GoogleAPICallError +except ImportError: + GoogleAPICallError = None + + +@unittest.skipIf( + GoogleAPICallError is None, 'GCP dependencies are not installed') +@mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') +class BigQueryStorageWriteDynamicSchemaTest(unittest.TestCase): + """Test dynamic schema support in BigQuery Storage Write API.""" + def test_storage_write_init_with_schema_side_inputs( + self, mock_expansion_service): + """Test that StorageWriteToBigQuery accepts schema_side_inputs.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', + schema=lambda dest: None, + schema_side_inputs=('side_input_1', )) + self.assertEqual(transform._schema_side_inputs, ('side_input_1', )) + self.assertEqual(transform._table_side_inputs, ()) + + def test_convert_to_beam_rows_dynamic_destinations_dynamic_schema( + self, mock_expansion_service): + """Test ConvertToBeamRows with dynamic destinations and dynamic schema.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, + ] + } + schema2 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'score', 'type': 'FLOAT' + }, + ] + } + schema_map = {'table1': schema1, 'table2': schema2} + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest: schema_map[dest], dynamic_destinations=True) + + with TestPipeline() as p: + input_data = [ + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), + ] + res = p | "CreateInput" >> beam.Create(input_data) | converter + + expected_rows = [ + beam.Row(destination='table1', record=beam.Row(id=1, name='foo')), + beam.Row(destination='table2', record=beam.Row(id=2, score=3.14)), + ] + assert_that(res, equal_to(expected_rows)) + + def test_convert_to_beam_rows_dynamic_destinations_with_side_inputs( + self, mock_expansion_service): + """Test ConvertToBeamRows with dynamic schema and side inputs.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, + ] + } + schema2 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'score', 'type': 'FLOAT' + }, + ] + } + + with TestPipeline() as p: + side_pcoll = ( + p + | "CreateSide" >> beam.Create([{ + 'table1': schema1, 'table2': schema2 + }])) + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest, side_map: side_map[dest], + dynamic_destinations=True, + schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll), )) + + input_data = [ + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), + ] + res = p | "CreateInput" >> beam.Create(input_data) | converter + + expected_rows = [ + beam.Row(destination='table1', record=beam.Row(id=1, name='foo')), + beam.Row(destination='table2', record=beam.Row(id=2, score=3.14)), + ] + assert_that(res, equal_to(expected_rows)) + + def test_storage_write_static_destination_dynamic_schema_raises_error( + self, mock_expansion_service): + """Test that static destination with dynamic schema raises ValueError.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', schema=lambda dest: None) + with self.assertRaisesRegex( + ValueError, + "Writing with a dynamic schema is only supported when writing to " + "dynamic destinations."): + with TestPipeline() as p: + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | transform + + def test_convert_to_beam_rows_with_output_types_dynamic_schema( + self, mock_expansion_service): + """Test with_output_types when schema is callable.""" + converter_dyn = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest: None, dynamic_destinations=True) + type_hint_dyn = converter_dyn.with_output_types().get_type_hints( + ).simple_output_type('') + self.assertIsInstance(type_hint_dyn, RowTypeConstraint) + self.assertEqual( + type_hint_dyn._fields, + ( + (bigquery.StorageWriteToBigQuery.DESTINATION, str), + ( + bigquery.StorageWriteToBigQuery.RECORD, + RowTypeConstraint.from_fields([])), + )) + + def test_convert_to_beam_rows_with_output_types_dynamic_schema_hint( + self, mock_expansion_service): + """Test with_output_types when schema is callable with _union_schema.""" + def dyn_schema(dest): + return None + + dyn_schema._union_schema = 'id:INTEGER,name:STRING' + converter_dyn = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, dynamic_destinations=True) + type_hint_dyn = converter_dyn.with_output_types().get_type_hints( + ).simple_output_type('') + self.assertIsInstance(type_hint_dyn, RowTypeConstraint) + self.assertEqual( + type_hint_dyn._fields[0], + (bigquery.StorageWriteToBigQuery.DESTINATION, str)) + self.assertEqual( + type_hint_dyn._fields[1][0], bigquery.StorageWriteToBigQuery.RECORD) + expected_record_hint = RowTypeConstraint.from_fields( + bigquery_tools.get_beam_typehints_from_tableschema( + 'id:INTEGER,name:STRING')) + self.assertEqual( + type_hint_dyn._fields[1][1]._fields, expected_record_hint._fields) + + def test_convert_to_beam_rows_union_schema_fills_missing_attributes( + self, mock_expansion_service): + """Test ConvertToBeamRows fills None for fields in union schema not in row.""" + def dyn_schema(dest): + if 'users' in dest: + return 'id:INTEGER,name:STRING' + return 'id:INTEGER,score:INTEGER' + + dyn_schema._union_schema = 'id:INTEGER,name:STRING,score:INTEGER' + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, dynamic_destinations=True) + + with TestPipeline() as p: + rows = ( + p + | beam.Create([ + ('dest_users', { + 'id': 1, 'name': 'alice' + }), + ('dest_scores', { + 'id': 2, 'score': 95 + }), + ]) + | converter) + + def check_rows(actual): + actual_list = list(actual) + assert len(actual_list) == 2 + r1, r2 = actual_list[0], actual_list[1] + if r1.destination == 'dest_scores': + r1, r2 = r2, r1 + assert r1.destination == 'dest_users' + assert r1.record.id == 1 + assert r1.record.name == 'alice' + assert r1.record.score is None + assert r2.destination == 'dest_scores' + assert r2.record.id == 2 + assert r2.record.name is None + assert r2.record.score == 95 + + assert_that(rows, check_rows) + + def test_storage_write_to_bigquery_expand_dynamic_schema( + self, mock_expansion_service): + """Test StorageWriteToBigQuery expand does not fail for callable schema.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + ] + } + + class _DummyExternalTransform(beam.PTransform): + def expand(self, pcoll): + return { + bigquery.StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS: ( + pcoll.pipeline | "CreateErrors" >> beam.Create([])) + } + + with mock.patch.object(bigquery, + 'SchemaAwareExternalTransform', + autospec=True) as mock_ext: + mock_ext.return_value = _DummyExternalTransform() + transform = bigquery.StorageWriteToBigQuery( + table=lambda record: 'table1', schema=lambda dest: schema1) + + with TestPipeline() as p: + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | transform + + mock_ext.assert_called_once() + _, kwargs = mock_ext.call_args + self.assertEqual( + kwargs['table'], bigquery.StorageWriteToBigQuery.DYNAMIC_DESTINATIONS) + + def test_write_to_bigquery_storage_api_passes_schema_side_inputs( + self, mock_expansion_service): + """Test WriteToBigQuery passes schema_side_inputs to StorageWriteToBigQuery.""" + with mock.patch.object(bigquery, 'StorageWriteToBigQuery', + autospec=True) as mock_storage_write: + mock_storage_write.return_value = beam.Map(lambda x: x) + with TestPipeline() as p: + side_pc = p | "CreateSide" >> beam.Create([1]) + write_transform = bigquery.WriteToBigQuery( + table='proj:ds.table', + method=bigquery.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=lambda dest: None, + schema_side_inputs=(beam.pvalue.AsSingleton(side_pc), )) + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | write_transform + + mock_storage_write.assert_called_once() + _, kwargs = mock_storage_write.call_args + self.assertEqual(len(kwargs['schema_side_inputs']), 1) + + +if __name__ == '__main__': + unittest.main()