github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/gcp/bigquery_write_perf_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """
    19  A pipeline that writes data from Synthetic Source to a BigQuery table.
    20  Besides of the standard options, there are options with special meaning:
    21  * output_dataset - BQ dataset name.
    22  * output_table - BQ table name. The table will be removed after test completion,
    23  * input_options - options for Synthetic Source:
    24  num_records - number of rows to be inserted,
    25  value_size - the length of a single row,
    26  key_size - required option, but its value has no meaning.
    27  
    28  Example test run on DataflowRunner:
    29  
    30  python -m apache_beam.io.gcp.bigquery_write_perf_test \
    31      --test-pipeline-options="
    32      --runner=TestDataflowRunner
    33      --project=...
    34      --region=...
    35      --staging_location=gs://...
    36      --temp_location=gs://...
    37      --sdk_location=.../dist/apache-beam-x.x.x.dev0.tar.gz
    38      --publish_to_big_query=true
    39      --metrics_dataset=gs://...
    40      --metrics_table=...
    41      --output_dataset=...
    42      --output_table=...
    43      --input_options='{
    44      \"num_records\": 1024,
    45      \"key_size\": 1,
    46      \"value_size\": 1024,
    47      }'"
    48  
    49  This setup will result in a table of 1MB size.
    50  """
    51  
    52  # pytype: skip-file
    53  
    54  import logging
    55  
    56  from apache_beam import Map
    57  from apache_beam import ParDo
    58  from apache_beam.io import BigQueryDisposition
    59  from apache_beam.io import Read
    60  from apache_beam.io import WriteToBigQuery
    61  from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json
    62  from apache_beam.io.gcp.tests import utils
    63  from apache_beam.testing.load_tests.load_test import LoadTest
    64  from apache_beam.testing.load_tests.load_test_metrics_utils import CountMessages
    65  from apache_beam.testing.load_tests.load_test_metrics_utils import MeasureTime
    66  from apache_beam.testing.synthetic_pipeline import SyntheticSource
    67  
    68  
    69  class BigQueryWritePerfTest(LoadTest):
    70    def __init__(self):
    71      super().__init__()
    72      self.output_dataset = self.pipeline.get_option('output_dataset')
    73      self.output_table = self.pipeline.get_option('output_table')
    74  
    75    def test(self):
    76      SCHEMA = parse_table_schema_from_json(
    77          '{"fields": [{"name": "data", "type": "BYTES"}]}')
    78  
    79      def format_record(record):
    80        # Since Synthetic Source returns data as a dictionary, we should skip one
    81        # of the part
    82        import base64
    83        return {'data': base64.b64encode(record[1])}
    84  
    85      (  # pylint: disable=expression-not-assigned
    86          self.pipeline
    87          | 'Produce rows' >> Read(
    88              SyntheticSource(self.parse_synthetic_source_options()))
    89          | 'Count messages' >> ParDo(CountMessages(self.metrics_namespace))
    90          | 'Format' >> Map(format_record)
    91          | 'Measure time' >> ParDo(MeasureTime(self.metrics_namespace))
    92          | 'Write to BigQuery' >> WriteToBigQuery(
    93              dataset=self.output_dataset,
    94              table=self.output_table,
    95              schema=SCHEMA,
    96              create_disposition=BigQueryDisposition.CREATE_IF_NEEDED,
    97              write_disposition=BigQueryDisposition.WRITE_TRUNCATE))
    98  
    99    def cleanup(self):
   100      """Removes an output BQ table."""
   101      utils.delete_bq_table(
   102          self.project_id, self.output_dataset, self.output_table)
   103  
   104  
   105  if __name__ == '__main__':
   106    logging.basicConfig(level=logging.INFO)
   107    BigQueryWritePerfTest().run()