Beam Python Dataflow Runner Uses deprecated BigQuerySink instead of WriteToBigQuery in apply_WriteToBigQuery

gabriel1836_20178 0
Viewed 2005

In terms of implementation details within the DataflowRunner, a lot of people may not care about whether BigQuerySink or WriteToBigQuery is used.

However, in my case I am trying to deploy my code to dataflow templates with RunTimeValueProvider's for the arguments. This behavior is supported in WriteToBigQuery:

class WriteToBigQuery(PTransform):
....

 table (str, callable, ValueProvider): The ID of the table, or a callable
         that returns it. The ID must contain only letters ``a-z``, ``A-Z``,
         numbers ``0-9``, or underscores ``_``. If dataset argument is
         :data:`None` then the table argument must contain the entire table
         reference specified as: ``'DATASET.TABLE'``
         or ``'PROJECT:DATASET.TABLE'``. If it's a callable, it must receive one
         argument representing an element to be written to BigQuery, and return
         a TableReference, or a string table name as specified above.
         Multiple destinations are only supported on Batch pipelines at the
         moment.

It is not supported in BigQuerySink:

class BigQuerySink(dataflow_io.NativeSink):
      table (str): The ID of the table. The ID must contain only letters
        ``a-z``, ``A-Z``, numbers ``0-9``, or underscores ``_``. If
        **dataset** argument is :data:`None` then the table argument must
        contain the entire table reference specified as: ``'DATASET.TABLE'`` or
        ``'PROJECT:DATASET.TABLE'``.

What's even more fun is that BigQuerySink in the code has been deprecated since 2.11.0.

@deprecated(since='2.11.0', current="WriteToBigQuery")

However in DataFlowRunner, the current code and comments seem completely out of line with the expectation that WriteToBigQuery is the default class to use over BigQuerySink:

  def apply_WriteToBigQuery(self, transform, pcoll, options):
    # Make sure this is the WriteToBigQuery class that we expected, and that
    # users did not specifically request the new BQ sink by passing experiment
    # flag.

    # TODO(BEAM-6928): Remove this function for release 2.14.0.
    experiments = options.view_as(DebugOptions).experiments or []
    if (not isinstance(transform, beam.io.WriteToBigQuery)
        or 'use_beam_bq_sink' in experiments):
      return self.apply_PTransform(transform, pcoll, options)
    if transform.schema == beam.io.gcp.bigquery.SCHEMA_AUTODETECT:
      raise RuntimeError(
          'Schema auto-detection is not supported on the native sink')
    standard_options = options.view_as(StandardOptions)
    if standard_options.streaming:
      if (transform.write_disposition ==
          beam.io.BigQueryDisposition.WRITE_TRUNCATE):
        raise RuntimeError('Can not use write truncation mode in streaming')
      return self.apply_PTransform(transform, pcoll, options)
    else:
      from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json
      schema = None
      if transform.schema:
        schema = parse_table_schema_from_json(json.dumps(transform.schema))
      return pcoll  | 'WriteToBigQuery' >> beam.io.Write(
          beam.io.BigQuerySink(
              transform.table_reference.tableId,
              transform.table_reference.datasetId,
              transform.table_reference.projectId,
              schema,
              transform.create_disposition,
              transform.write_disposition,
              kms_key=transform.kms_key))

My questions are twofold:

  1. Why the difference in contract/expectations between the DataflowRunner and the io.BigQuery class?
  2. Without waiting for a bug fix, does anyone have suggestions on how to force the DataflowRunner to use WriteToBigQuery over BigQuerySink?
apache-beam google-cloud-dataflow google-bigquery python
1 Answers

The WriteToBigQuery transform has two different strategies to write to BigQuery:

  • Streaming inserts into a BigQuery endpoint
  • File Load jobs triggered periodically (or once for batch pipelines)

For the Python SDK, we initially only had support for Streaming inserts, and we had a runner-native implementation for File Loads that only worked on Dataflow (this is the BigQuerySink).

For Batch pipelines running on Dataflow, BigQuerySink is replaced in - as you properly found. For all other cases, streaming inserts were used.

In a recent version of Beam, we added support for File Loads natively in the SDK - the implementation for this is in BigQueryBatchFileLoads.

Because we didn't want to break users relying on the old behavior, we masked BigQueryBatchFileLoads behind an experiment flag. (the flag is use_beam_bq_sink).

So:

  • In a future version, we will automatically use BigQueryBatchFileLoads, but for now, you have two options to have access to it:

    1. Use it directly in your pipeline (e.g. input | BigQueryBatchFileLoads(...)).
    2. Pass the option --experiments use_beam_bq_sink, while using WriteToBigQuery.

I hope that helps!

user_12553560
Related