Merge pull request #12782 from [BEAM-10950] Overriding Dataflow Native BQSource.

* Overriding Dataflow Native BQSource.

* Fixing tests

* Tryin to fix to_runner_proto transform

* fixup

* fixup

* fixup

* adding comment

Author: pabloem

sdks/python/apache_beam/examples/snippets/snippets_test.py

@@ -682,6 +682,7 @@ def test_model_bigqueryio(self):
         snippets.model_bigqueryio(p, project, dataset, table)
     else:
       p = TestPipeline()
+      p.options.view_as(GoogleCloudOptions).temp_location = 'gs://mylocation'
       snippets.model_bigqueryio(p)
 
   def _run_test_pipeline_for_options(self, fn):

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

@@ -302,6 +302,7 @@ def test_big_query_new_types_native(self):
         'output_schema': NEW_TYPES_OUTPUT_SCHEMA,
         'use_standard_sql': False,
         'native': True,
+        'use_json_exports': True,
         'wait_until_finish_duration': WAIT_UNTIL_FINISH_DURATION_MS,
         'on_success_matcher': all_of(*pipeline_verifiers),
         'experiments': 'use_legacy_bq_sink',

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

@@ -91,24 +91,16 @@ def run_bq_pipeline(argv=None):
         use_standard_sql=known_args.use_standard_sql,
         use_json_exports=known_args.use_json_exports,
         kms_key=kms_key)
-  if known_args.native:
-    _ = data | 'write' >> beam.io.Write(
-        beam.io.BigQuerySink(
-            known_args.output,
-            schema=table_schema,
-            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
-            write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY,
-            kms_key=kms_key))
-  else:
-    temp_file_format = (
-        'NEWLINE_DELIMITED_JSON' if known_args.use_json_exports else 'AVRO')
-    _ = data | 'write' >> beam.io.WriteToBigQuery(
-        known_args.output,
-        schema=table_schema,
-        create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
-        write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY,
-        temp_file_format=temp_file_format,
-        kms_key=kms_key)
+
+  temp_file_format = (
+      'NEWLINE_DELIMITED_JSON' if known_args.use_json_exports else 'AVRO')
+  _ = data | 'write' >> beam.io.WriteToBigQuery(
+      known_args.output,
+      schema=table_schema,
+      create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
+      write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY,
+      temp_file_format=temp_file_format,
+      kms_key=kms_key)
 
   result = p.run()
   result.wait_until_finish()

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

@@ -426,7 +426,45 @@ def validate_write(disposition):
 # BigQuerySource, BigQuerySink.
 
 
-class BigQuerySource(dataflow_io.NativeSource):
+@deprecated(since='2.25.0', current="ReadFromBigQuery")
+def BigQuerySource(
+    table=None,
+    dataset=None,
+    project=None,
+    query=None,
+    validate=False,
+    coder=None,
+    use_standard_sql=False,
+    flatten_results=True,
+    kms_key=None,
+    use_dataflow_native_source=False):
+  if use_dataflow_native_source:
+    return _BigQuerySource(
+        table,
+        dataset,
+        project,
+        query,
+        validate,
+        coder,
+        use_standard_sql,
+        flatten_results,
+        kms_key)
+  else:
+    return ReadFromBigQuery(
+        table=table,
+        dataset=dataset,
+        project=project,
+        query=query,
+        validate=validate,
+        coder=coder,
+        use_standard_sql=use_standard_sql,
+        flatten_results=flatten_results,
+        use_json_exports=True,
+        kms_key=kms_key)
+
+
+@deprecated(since='2.25.0', current="ReadFromBigQuery")
+class _BigQuerySource(dataflow_io.NativeSource):
   """A source based on a BigQuery table."""
   def __init__(
       self,

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

@@ -168,7 +168,8 @@ def test_invalid_json_neg_inf(self):
 @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
 class TestBigQuerySource(unittest.TestCase):
   def test_display_data_item_on_validate_true(self):
-    source = beam.io.BigQuerySource('dataset.table', validate=True)
+    source = beam.io.BigQuerySource(
+        'dataset.table', validate=True, use_dataflow_native_source=True)
 
     dd = DisplayData.create_from(source)
     expected_items = [
@@ -178,23 +179,26 @@ def test_display_data_item_on_validate_true(self):
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
 
   def test_table_reference_display_data(self):
-    source = beam.io.BigQuerySource('dataset.table')
+    source = beam.io.BigQuerySource(
+        'dataset.table', use_dataflow_native_source=True)
     dd = DisplayData.create_from(source)
     expected_items = [
         DisplayDataItemMatcher('validation', False),
         DisplayDataItemMatcher('table', 'dataset.table')
     ]
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
 
-    source = beam.io.BigQuerySource('project:dataset.table')
+    source = beam.io.BigQuerySource(
+        'project:dataset.table', use_dataflow_native_source=True)
     dd = DisplayData.create_from(source)
     expected_items = [
         DisplayDataItemMatcher('validation', False),
         DisplayDataItemMatcher('table', 'project:dataset.table')
     ]
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
 
-    source = beam.io.BigQuerySource('xyz.com:project:dataset.table')
+    source = beam.io.BigQuerySource(
+        'xyz.com:project:dataset.table', use_dataflow_native_source=True)
     dd = DisplayData.create_from(source)
     expected_items = [
         DisplayDataItemMatcher('validation', False),
@@ -203,27 +207,32 @@ def test_table_reference_display_data(self):
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
 
   def test_parse_table_reference(self):
-    source = beam.io.BigQuerySource('dataset.table')
+    source = beam.io.BigQuerySource(
+        'dataset.table', use_dataflow_native_source=True)
     self.assertEqual(source.table_reference.datasetId, 'dataset')
     self.assertEqual(source.table_reference.tableId, 'table')
 
-    source = beam.io.BigQuerySource('project:dataset.table')
+    source = beam.io.BigQuerySource(
+        'project:dataset.table', use_dataflow_native_source=True)
     self.assertEqual(source.table_reference.projectId, 'project')
     self.assertEqual(source.table_reference.datasetId, 'dataset')
     self.assertEqual(source.table_reference.tableId, 'table')
 
-    source = beam.io.BigQuerySource('xyz.com:project:dataset.table')
+    source = beam.io.BigQuerySource(
+        'xyz.com:project:dataset.table', use_dataflow_native_source=True)
     self.assertEqual(source.table_reference.projectId, 'xyz.com:project')
     self.assertEqual(source.table_reference.datasetId, 'dataset')
     self.assertEqual(source.table_reference.tableId, 'table')
 
-    source = beam.io.BigQuerySource(query='my_query')
+    source = beam.io.BigQuerySource(
+        query='my_query', use_dataflow_native_source=True)
     self.assertEqual(source.query, 'my_query')
     self.assertIsNone(source.table_reference)
     self.assertTrue(source.use_legacy_sql)
 
   def test_query_only_display_data(self):
-    source = beam.io.BigQuerySource(query='my_query')
+    source = beam.io.BigQuerySource(
+        query='my_query', use_dataflow_native_source=True)
     dd = DisplayData.create_from(source)
     expected_items = [
         DisplayDataItemMatcher('validation', False),
@@ -232,25 +241,36 @@ def test_query_only_display_data(self):
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
 
   def test_specify_query_sql_format(self):
-    source = beam.io.BigQuerySource(query='my_query', use_standard_sql=True)
+    source = beam.io.BigQuerySource(
+        query='my_query',
+        use_standard_sql=True,
+        use_dataflow_native_source=True)
     self.assertEqual(source.query, 'my_query')
     self.assertFalse(source.use_legacy_sql)
 
   def test_specify_query_flattened_records(self):
-    source = beam.io.BigQuerySource(query='my_query', flatten_results=False)
+    source = beam.io.BigQuerySource(
+        query='my_query',
+        flatten_results=False,
+        use_dataflow_native_source=True)
     self.assertFalse(source.flatten_results)
 
   def test_specify_query_unflattened_records(self):
-    source = beam.io.BigQuerySource(query='my_query', flatten_results=True)
+    source = beam.io.BigQuerySource(
+        query='my_query', flatten_results=True, use_dataflow_native_source=True)
     self.assertTrue(source.flatten_results)
 
   def test_specify_query_without_table(self):
-    source = beam.io.BigQuerySource(query='my_query')
+    source = beam.io.BigQuerySource(
+        query='my_query', use_dataflow_native_source=True)
     self.assertEqual(source.query, 'my_query')
     self.assertIsNone(source.table_reference)
 
   def test_date_partitioned_table_name(self):
-    source = beam.io.BigQuerySource('dataset.table$20030102', validate=True)
+    source = beam.io.BigQuerySource(
+        'dataset.table$20030102',
+        validate=True,
+        use_dataflow_native_source=True)
     dd = DisplayData.create_from(source)
     expected_items = [
         DisplayDataItemMatcher('validation', True),

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

@@ -471,7 +471,9 @@ def test_read_from_table(self):
     client.jobs.GetQueryResults.return_value = bigquery.GetQueryResultsResponse(
         jobComplete=True, rows=table_rows, schema=schema)
     actual_rows = []
-    with beam.io.BigQuerySource('dataset.table').reader(client) as reader:
+    with beam.io.BigQuerySource(
+        'dataset.table',
+        use_dataflow_native_source=True).reader(client) as reader:
       for row in reader:
         actual_rows.append(row)
     self.assertEqual(actual_rows, expected_rows)
@@ -485,7 +487,9 @@ def test_read_from_query(self):
     client.jobs.GetQueryResults.return_value = bigquery.GetQueryResultsResponse(
         jobComplete=True, rows=table_rows, schema=schema)
     actual_rows = []
-    with beam.io.BigQuerySource(query='query').reader(client) as reader:
+    with beam.io.BigQuerySource(
+        query='query',
+        use_dataflow_native_source=True).reader(client) as reader:
       for row in reader:
         actual_rows.append(row)
     self.assertEqual(actual_rows, expected_rows)
@@ -501,8 +505,9 @@ def test_read_from_query_sql_format(self):
     client.jobs.GetQueryResults.return_value = bigquery.GetQueryResultsResponse(
         jobComplete=True, rows=table_rows, schema=schema)
     actual_rows = []
-    with beam.io.BigQuerySource(query='query',
-                                use_standard_sql=True).reader(client) as reader:
+    with beam.io.BigQuerySource(
+        query='query', use_standard_sql=True,
+        use_dataflow_native_source=True).reader(client) as reader:
       for row in reader:
         actual_rows.append(row)
     self.assertEqual(actual_rows, expected_rows)
@@ -518,8 +523,9 @@ def test_read_from_query_unflatten_records(self):
     client.jobs.GetQueryResults.return_value = bigquery.GetQueryResultsResponse(
         jobComplete=True, rows=table_rows, schema=schema)
     actual_rows = []
-    with beam.io.BigQuerySource(query='query',
-                                flatten_results=False).reader(client) as reader:
+    with beam.io.BigQuerySource(
+        query='query', flatten_results=False,
+        use_dataflow_native_source=True).reader(client) as reader:
       for row in reader:
         actual_rows.append(row)
     self.assertEqual(actual_rows, expected_rows)
@@ -532,12 +538,13 @@ def test_using_both_query_and_table_fails(self):
         ValueError,
         r'Both a BigQuery table and a query were specified\. Please specify '
         r'only one of these'):
-      beam.io.BigQuerySource(table='dataset.table', query='query')
+      beam.io.BigQuerySource(
+          table='dataset.table', query='query', use_dataflow_native_source=True)
 
   def test_using_neither_query_nor_table_fails(self):
     with self.assertRaisesRegex(
         ValueError, r'A BigQuery table or a query must be specified'):
-      beam.io.BigQuerySource()
+      beam.io.BigQuerySource(use_dataflow_native_source=True)
 
   def test_read_from_table_as_tablerows(self):
     client = mock.Mock()
@@ -550,7 +557,9 @@ def test_read_from_table_as_tablerows(self):
     # We set the coder to TableRowJsonCoder, which is a signal that
     # the caller wants to see the rows as TableRows.
     with beam.io.BigQuerySource(
-        'dataset.table', coder=TableRowJsonCoder).reader(client) as reader:
+        'dataset.table',
+        coder=TableRowJsonCoder,
+        use_dataflow_native_source=True).reader(client) as reader:
       for row in reader:
         actual_rows.append(row)
     self.assertEqual(actual_rows, table_rows)
@@ -570,7 +579,9 @@ def test_read_from_table_and_job_complete_retry(self, patched_time_sleep):
             jobComplete=True, rows=table_rows, schema=schema)
     ]
     actual_rows = []
-    with beam.io.BigQuerySource('dataset.table').reader(client) as reader:
+    with beam.io.BigQuerySource(
+        'dataset.table',
+        use_dataflow_native_source=True).reader(client) as reader:
       for row in reader:
         actual_rows.append(row)
     self.assertEqual(actual_rows, expected_rows)
@@ -590,7 +601,9 @@ def test_read_from_table_and_multiple_pages(self):
             jobComplete=True, rows=table_rows, schema=schema)
     ]
     actual_rows = []
-    with beam.io.BigQuerySource('dataset.table').reader(client) as reader:
+    with beam.io.BigQuerySource(
+        'dataset.table',
+        use_dataflow_native_source=True).reader(client) as reader:
       for row in reader:
         actual_rows.append(row)
     # We return expected rows for each of the two pages of results so we
@@ -599,7 +612,8 @@ def test_read_from_table_and_multiple_pages(self):
 
   def test_table_schema_without_project(self):
     # Reader should pick executing project by default.
-    source = beam.io.BigQuerySource(table='mydataset.mytable')
+    source = beam.io.BigQuerySource(
+        table='mydataset.mytable', use_dataflow_native_source=True)
     options = PipelineOptions(flags=['--project', 'myproject'])
     source.pipeline_options = options
     reader = source.reader()

sdks/python/apache_beam/io/iobase.py

@@ -890,6 +890,11 @@ def get_desired_chunk_size(total_size):
   def expand(self, pbegin):
     if isinstance(self.source, BoundedSource):
       return pbegin | _SDFBoundedSourceWrapper(self.source)
+    elif isinstance(self.source, ptransform.PTransform):
+      # The Read transform can also admit a full PTransform as an input
+      # rather than an anctual source. If the input is a PTransform, then
+      # just apply it directly.
+      return pbegin.pipeline | self.source
     else:
       # Treat Read itself as a primitive.
       return pvalue.PCollection(
@@ -917,13 +922,17 @@ def display_data(self):
 
   def to_runner_api_parameter(self, context):
     # type: (PipelineContext) -> Tuple[str, beam_runner_api_pb2.ReadPayload]
-    return (
-        common_urns.deprecated_primitives.READ.urn,
-        beam_runner_api_pb2.ReadPayload(
-            source=self.source.to_runner_api(context),
-            is_bounded=beam_runner_api_pb2.IsBounded.BOUNDED
-            if self.source.is_bounded() else
-            beam_runner_api_pb2.IsBounded.UNBOUNDED))
+    from apache_beam.runners.dataflow.native_io import iobase as dataflow_io
+    if isinstance(self.source, (BoundedSource, dataflow_io.NativeSource)):
+      return (
+          common_urns.deprecated_primitives.READ.urn,
+          beam_runner_api_pb2.ReadPayload(
+              source=self.source.to_runner_api(context),
+              is_bounded=beam_runner_api_pb2.IsBounded.BOUNDED
+              if self.source.is_bounded() else
+              beam_runner_api_pb2.IsBounded.UNBOUNDED))
+    elif isinstance(self.source, ptransform.PTransform):
+      return self.source.to_runner_api_parameter(context)
 
   @staticmethod
   def from_runner_api_parameter(unused_ptransform, parameter, context):