[WIP] Add Dynamic Schema support to StorageWriteToBigQuery

.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
 }

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)

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)