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,67 @@ 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 * 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([schema_map])
+      _ = (
+          p
+          | "CreateElements" >> beam.Create(elements)
+          | beam.io.WriteToBigQuery(
+              table=get_destination,
+              method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API,
+              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))
+
   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

@@ -417,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
 
@@ -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,
@@ -2616,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.
   """
@@ -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
@@ -2691,6 +2694,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 +2736,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 +2808,43 @@ 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 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
@@ -2825,6 +2855,13 @@ def expand(self, input_dicts):
                     ]): bigquery_tools.beam_row_from_dict(row, schema)))
 
     def with_output_types(self):
+      if callable(self.schema):
+        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(
           self.schema, self.type_overrides)
       if self.dynamic_destinations: