diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index 89cec619b020..ea130ded4ea6 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.github/trigger_files/beam_PostCommit_Python.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "pr": "38701", - "modification": 56 + "pr": "39662", + "modification": 57 } diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 314effad5520..83175d9864af 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -2319,6 +2319,21 @@ def expand(self, pcoll): BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS]) elif method_to_use == WriteToBigQuery.Method.FILE_LOADS: + if self.schema is None: + # If the input PCollection carries a Beam schema (e.g. it was + # produced by ReadFromBigQuery(..., output_type='BEAM_ROW'), or is + # otherwise a PCollection of NamedTuples, dataclasses, or Beam Rows), + # auto-infer the destination table's schema from it. This mirrors + # the schema auto-inference already performed for the + # STORAGE_WRITE_API method. + try: + beam_schema = schema_from_element_type(pcoll.element_type) + except TypeError: + beam_schema = None + if beam_schema is not None: + self.schema = bigquery_schema_tools.beam_schema_to_bq_table_schema( + beam_schema) + if self._temp_file_format == bigquery_tools.FileFormat.AVRO: if self.schema == SCHEMA_AUTODETECT: raise ValueError( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py index d3d608b1fc6f..8f148d749628 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py @@ -116,6 +116,122 @@ def bq_field_to_type(field, mode, type_overrides=None): raise ValueError(f"Encountered an unsupported mode: {mode!r}") +_ATOMIC_TYPE_TO_BQ_TYPE = { + schema_pb2.BOOLEAN: 'BOOL', + schema_pb2.BYTES: 'BYTES', + schema_pb2.STRING: 'STRING', + schema_pb2.BYTE: 'INT64', + schema_pb2.INT16: 'INT64', + schema_pb2.INT32: 'INT64', + schema_pb2.INT64: 'INT64', + schema_pb2.FLOAT: 'FLOAT64', + schema_pb2.DOUBLE: 'FLOAT64', +} + + +def _logical_type_to_bq_type_map(): + # Built lazily to avoid a module-level import cycle between + # bigquery_schema_tools and typehints.schemas. + schemas = apache_beam.typehints.schemas + return { + schemas.MillisInstant.urn(): 'TIMESTAMP', + schemas.MicrosInstant.urn(): 'TIMESTAMP', + schemas.Date.urn(): 'DATE', + schemas.DecimalLogicalType.urn(): 'NUMERIC', + schemas.FixedPrecisionDecimalLogicalType.urn(): 'NUMERIC', + } + + +def _bq_type_mode_and_fields(field_type: schema_pb2.FieldType): + """Maps a Beam schema_pb2.FieldType to a + ``(bq_type, bq_mode, nested_bq_fields)`` tuple, where ``nested_bq_fields`` + is a list of BigQuery field dicts for STRUCT types, and ``None`` otherwise. + """ + type_info = field_type.WhichOneof('type_info') + mode = 'NULLABLE' if field_type.nullable else 'REQUIRED' + + if type_info in ('array_type', 'iterable_type'): + element_type = ( + field_type.array_type.element_type + if type_info == 'array_type' else field_type.iterable_type.element_type) + if element_type.WhichOneof('type_info') in ('array_type', 'iterable_type'): + raise ValueError( + 'BigQuery does not support nested (repeated-of-repeated) fields; ' + 'please provide an explicit schema instead of relying on ' + 'auto-inference.') + bq_type, _, nested_fields = _bq_type_mode_and_fields(element_type) + return bq_type, 'REPEATED', nested_fields + elif type_info == 'row_type': + nested_fields = [ + beam_field_to_bq_field(f) for f in field_type.row_type.schema.fields + ] + return 'STRUCT', mode, nested_fields + elif type_info == 'logical_type': + urn = field_type.logical_type.urn + logical_type_to_bq_type = _logical_type_to_bq_type_map() + if urn in logical_type_to_bq_type: + return logical_type_to_bq_type[urn], mode, None + # Fall back to the logical type's representation type, which covers + # e.g. fixed/variable-length strings and bytes. + representation_type = field_type.logical_type.representation + if representation_type.WhichOneof('type_info') == 'atomic_type': + atomic = representation_type.atomic_type + if atomic in _ATOMIC_TYPE_TO_BQ_TYPE: + return _ATOMIC_TYPE_TO_BQ_TYPE[atomic], mode, None + raise ValueError( + f'Cannot automatically infer a BigQuery type for the logical type ' + f'with urn {urn!r}. Please provide an explicit schema.') + elif type_info == 'atomic_type': + atomic = field_type.atomic_type + if atomic not in _ATOMIC_TYPE_TO_BQ_TYPE: + raise ValueError( + f'Cannot automatically infer a BigQuery type for the atomic type ' + f'{atomic!r}. Please provide an explicit schema.') + return _ATOMIC_TYPE_TO_BQ_TYPE[atomic], mode, None + elif type_info == 'map_type': + raise ValueError( + 'BigQuery schema auto-inference does not support MapType fields; ' + 'please provide an explicit schema.') + else: + raise ValueError( + f'Cannot automatically infer a BigQuery type for field type ' + f'{field_type!r}. Please provide an explicit schema.') + + +def beam_field_to_bq_field(field: schema_pb2.Field) -> dict: + """Convert a single Beam schema field (schema_pb2.Field) into a BigQuery + TableFieldSchema, in dictionary form.""" + bq_type, mode, nested_fields = _bq_type_mode_and_fields(field.type) + bq_field = {'name': field.name, 'type': bq_type, 'mode': mode} + if nested_fields is not None: + bq_field['fields'] = nested_fields + return bq_field + + +def beam_schema_to_bq_table_schema(schema: schema_pb2.Schema) -> dict: + """Convert a Beam schema (schema_pb2.Schema) into a BigQuery TableSchema, + in dictionary form. + + This is the reverse of `generate_user_type_from_bq_schema`, and is used to + auto-infer the destination table's schema from a schema'd PCollection -- + for example a PCollection of NamedTuples, dataclasses, or Beam Rows, such + as the ones produced by + ``ReadFromBigQuery(..., output_type='BEAM_ROW')``. + + Args: + schema: a `schema_pb2.Schema` instance, as returned by + `apache_beam.typehints.schemas.schema_from_element_type`. + + Returns: + Dict[str, Any]: A BigQuery TableSchema in dictionary form, e.g. + ``{'fields': [{'name': 'a', 'type': 'INT64', 'mode': 'NULLABLE'}, ...]}`` + + Raises: + ValueError: if a field's type has no known BigQuery equivalent. + """ + return {'fields': [beam_field_to_bq_field(f) for f in schema.fields]} + + def convert_to_usertype( table_schema, selected_fields=None, type_overrides=None): """Convert a BigQuery table schema to a user type. diff --git a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py index 73cedb3a6aee..34f355ac7ef6 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py @@ -512,6 +512,151 @@ def test_type_overrides_can_override_default_types(self): bq_field_to_type("GEOGRAPHY", "REQUIRED", overrides), bytes) +class TestBeamSchemaToBqTableSchema(unittest.TestCase): + """Tests for beam_schema_to_bq_table_schema, the reverse of + generate_user_type_from_bq_schema, used to auto-infer a destination + table's schema from a schema'd PCollection (e.g. for FILE_LOADS).""" + def test_atomic_types(self): + import decimal + + class MyRow(typing.NamedTuple): + name: str + age: np.int64 + score: float + active: bool + raw: bytes + amount: decimal.Decimal + + schema = beam.typehints.schemas.schema_from_element_type(MyRow) + bq_schema = bigquery_schema_tools.beam_schema_to_bq_table_schema(schema) + + self.assertEqual( + bq_schema, + { + 'fields': [ + { + 'name': 'name', 'type': 'STRING', 'mode': 'REQUIRED' + }, + { + 'name': 'age', 'type': 'INT64', 'mode': 'REQUIRED' + }, + { + 'name': 'score', 'type': 'FLOAT64', 'mode': 'REQUIRED' + }, + { + 'name': 'active', 'type': 'BOOL', 'mode': 'REQUIRED' + }, + { + 'name': 'raw', 'type': 'BYTES', 'mode': 'REQUIRED' + }, + { + 'name': 'amount', 'type': 'NUMERIC', 'mode': 'REQUIRED' + }, + ] + }) + + def test_timestamp_type(self): + from apache_beam.utils.timestamp import Timestamp + + class MyRow(typing.NamedTuple): + when: Timestamp + + schema = beam.typehints.schemas.schema_from_element_type(MyRow) + bq_schema = bigquery_schema_tools.beam_schema_to_bq_table_schema(schema) + + self.assertEqual( + bq_schema, + {'fields': [{ + 'name': 'when', 'type': 'TIMESTAMP', 'mode': 'REQUIRED' + }]}) + + def test_nullable_field(self): + class MyRow(typing.NamedTuple): + name: typing.Optional[str] + + schema = beam.typehints.schemas.schema_from_element_type(MyRow) + bq_schema = bigquery_schema_tools.beam_schema_to_bq_table_schema(schema) + + self.assertEqual( + bq_schema, + {'fields': [{ + 'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE' + }]}) + + def test_repeated_field(self): + class MyRow(typing.NamedTuple): + tags: typing.Sequence[str] + + schema = beam.typehints.schemas.schema_from_element_type(MyRow) + bq_schema = bigquery_schema_tools.beam_schema_to_bq_table_schema(schema) + + self.assertEqual( + bq_schema, + {'fields': [{ + 'name': 'tags', 'type': 'STRING', 'mode': 'REPEATED' + }]}) + + def test_nested_record(self): + class Nested(typing.NamedTuple): + x: np.int64 + y: typing.Optional[str] + + class MyRow(typing.NamedTuple): + nested: Nested + + schema = beam.typehints.schemas.schema_from_element_type(MyRow) + bq_schema = bigquery_schema_tools.beam_schema_to_bq_table_schema(schema) + + self.assertEqual( + bq_schema, + { + 'fields': [{ + 'name': 'nested', + 'type': 'STRUCT', + 'mode': 'REQUIRED', + 'fields': [ + { + 'name': 'x', 'type': 'INT64', 'mode': 'REQUIRED' + }, + { + 'name': 'y', 'type': 'STRING', 'mode': 'NULLABLE' + }, + ] + }] + }) + + def test_repeated_record(self): + class Nested(typing.NamedTuple): + x: np.int64 + + class MyRow(typing.NamedTuple): + nested_list: typing.Sequence[Nested] + + schema = beam.typehints.schemas.schema_from_element_type(MyRow) + bq_schema = bigquery_schema_tools.beam_schema_to_bq_table_schema(schema) + + self.assertEqual( + bq_schema, + { + 'fields': [{ + 'name': 'nested_list', + 'type': 'STRUCT', + 'mode': 'REPEATED', + 'fields': [{ + 'name': 'x', 'type': 'INT64', 'mode': 'REQUIRED' + }] + }] + }) + + def test_map_type_raises(self): + class MyRow(typing.NamedTuple): + attrs: typing.Mapping[str, str] + + schema = beam.typehints.schemas.schema_from_element_type(MyRow) + with self.assertRaises(ValueError): + bigquery_schema_tools.beam_schema_to_bq_table_schema(schema) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index 51d13d96b73a..255badf8fad4 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -28,11 +28,13 @@ import re import secrets import time +import typing import unittest import uuid import hamcrest as hc import mock +import numpy as np import pytest import pytz import requests @@ -101,6 +103,12 @@ _LOGGER = logging.getLogger(__name__) + +class _MyRowForFileLoadsSchemaTest(typing.NamedTuple): + name: str + age: np.int64 + + _ELEMENTS = [ { 'name': 'beam', 'language': 'py' @@ -1007,6 +1015,92 @@ def test_schema_autodetect_not_allowed_with_avro_file_loads(self): schema=beam.io.gcp.bigquery.SCHEMA_AUTODETECT, temp_file_format=bigquery_tools.FileFormat.AVRO)) + def test_schema_autoinferred_for_file_loads_from_schema_pcoll(self): + p = beam.Pipeline() + pc = p | beam.Create([_MyRowForFileLoadsSchemaTest('a', 1) + ]).with_output_types(_MyRowForFileLoadsSchemaTest) + + transform = beam.io.gcp.bigquery.WriteToBigQuery( + "dataset.table", + schema=None, + method=beam.io.gcp.bigquery.WriteToBigQuery.Method.FILE_LOADS, + temp_file_format=bigquery_tools.FileFormat.JSON) + + with mock.patch( + 'apache_beam.io.gcp.bigquery_file_loads.BigQueryBatchFileLoads' + ) as mock_batch_file_loads: + mock_batch_file_loads.side_effect = RuntimeError('stop-here') + with self.assertRaisesRegex(RuntimeError, 'stop-here'): + transform.expand(pc) + + self.assertEqual( + transform.schema, + { + 'fields': [ + { + 'name': 'name', 'type': 'STRING', 'mode': 'REQUIRED' + }, + { + 'name': 'age', 'type': 'INT64', 'mode': 'REQUIRED' + }, + ] + }) + + def test_schema_autoinferred_for_avro_file_loads(self): + # Prior to auto-inference, FILE_LOADS with AVRO always required an + # explicit schema, even when the input PCollection carried a Beam + # schema. This confirms that requirement is now satisfied by inference. + p = beam.Pipeline() + pc = p | beam.Create([_MyRowForFileLoadsSchemaTest('a', 1) + ]).with_output_types(_MyRowForFileLoadsSchemaTest) + + transform = beam.io.gcp.bigquery.WriteToBigQuery( + "dataset.table", + schema=None, + method=beam.io.gcp.bigquery.WriteToBigQuery.Method.FILE_LOADS, + temp_file_format=bigquery_tools.FileFormat.AVRO) + + with mock.patch( + 'apache_beam.io.gcp.bigquery_file_loads.BigQueryBatchFileLoads' + ) as mock_batch_file_loads: + mock_batch_file_loads.side_effect = RuntimeError('stop-here') + # Should reach BigQueryBatchFileLoads (and hit our stub error) rather + # than raising "A schema must be provided" from the AVRO check. + with self.assertRaisesRegex(RuntimeError, 'stop-here'): + transform.expand(pc) + + self.assertEqual( + transform.schema, + { + 'fields': [ + { + 'name': 'name', 'type': 'STRING', 'mode': 'REQUIRED' + }, + { + 'name': 'age', 'type': 'INT64', 'mode': 'REQUIRED' + }, + ] + }) + + def test_schema_not_autoinferred_for_file_loads_without_schema_pcoll(self): + p = beam.Pipeline() + pc = p | beam.Create([{'name': 'a', 'age': 1}]) + + transform = beam.io.gcp.bigquery.WriteToBigQuery( + "dataset.table", + schema=None, + method=beam.io.gcp.bigquery.WriteToBigQuery.Method.FILE_LOADS, + temp_file_format=bigquery_tools.FileFormat.JSON) + + with mock.patch( + 'apache_beam.io.gcp.bigquery_file_loads.BigQueryBatchFileLoads' + ) as mock_batch_file_loads: + mock_batch_file_loads.side_effect = RuntimeError('stop-here') + with self.assertRaisesRegex(RuntimeError, 'stop-here'): + transform.expand(pc) + + self.assertIsNone(transform.schema) + def test_to_from_runner_api(self): """Tests that serialization of WriteToBigQuery is correct.