Auto-infer BigQuery schema from schema'd PCollections for FILE_LOADS

.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": "40241",
-  "modification": 63
+  "pr": "39662",
+  "modification": 12
 }

sdks/python/apache_beam/io/gcp/bigquery.py

@@ -2404,6 +2404,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(

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.

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()