def _ReadFromBigQueryImpl( # pylint: disable=invalid-name pipeline: beam.Pipeline, query: Text, use_bigquery_source: bool = False) -> beam.pvalue.PCollection: """Read from BigQuery. Args: pipeline: beam pipeline. query: a BigQuery sql string. use_bigquery_source: Whether to use BigQuerySource instead of experimental `ReadFromBigQuery` PTransform. Returns: PCollection of dict. """ # TODO(b/155441037): Consolidate to ReadFromBigQuery once its performance # on dataflow runner is on par with BigQuerySource. if use_bigquery_source: return ( pipeline | 'ReadFromBigQuerySource' >> beam.io.Read( beam.io.BigQuerySource(query=query, use_standard_sql=True))) return (pipeline | 'ReadFromBigQuery' >> beam_bigquery.ReadFromBigQuery( query=query, use_standard_sql=True, bigquery_job_labels=telemetry_utils.get_labels_dict()))
def ReadFromBigQuery( pipeline: beam.Pipeline, query: Text) -> beam.pvalue.PCollection: """Read data from BigQuery. Args: pipeline: Beam pipeline. query: A BigQuery sql string. Returns: PCollection of dict. """ return (pipeline | 'ReadFromBigQuery' >> bigquery.ReadFromBigQuery( query=query, use_standard_sql=True, bigquery_job_labels=telemetry_utils.get_labels_dict()))
def _ReadFromBigQueryImpl( # pylint: disable=invalid-name pipeline: beam.Pipeline, query: Text) -> beam.pvalue.PCollection: """Read from BigQuery. Args: pipeline: beam pipeline. query: a BigQuery sql string. Returns: PCollection of dict. """ return (pipeline | 'ReadFromBigQuery' >> beam_bigquery.ReadFromBigQuery( query=query, use_standard_sql=True, bigquery_job_labels=telemetry_utils.get_labels_dict()))
def ReadFromBigQuery(pipeline: beam.Pipeline, query_project: Text, query_dataset: Text, query_table: Text, gcs_location: Text, dest_project: Text, query_limit: int = None) -> beam.pvalue.PCollection: """ The Beam PTransform used to read data from a specific BQ table. Args: pipeline: Input beam.Pipeline object coming from a TFX Executor. query_project: Google Cloud project where the target table is located. query_dataset: Google Cloud project where the target dataset is located. query_table: Name of the target BigQuery table. gcs_location: Name of the Google Cloud Storage bucket where the extracted table should be written as a string. dest_project: Additional Google Cloud Project identifier. query_limit: Optional, maximum limit of how many datapoints to read from the specified BQ table. Returns: A beam.PCollection of data points. Each row in the BigQuery table represents a single data point. """ query = f'SELECT * FROM `{query_project}.{query_dataset}.{query_table}`' if query_limit is not None: query += f'\nLIMIT {query_limit}' return (pipeline | 'ReadFromBigQuery' >> beam_bigquery.ReadFromBigQuery( project=dest_project, gcs_location=gcs_location, query=query, use_standard_sql=True))
def __init__(self, *args, klio_message_columns=None, **kwargs): self.__reader = beam_bq.ReadFromBigQuery(*args, **kwargs) self.__mapper = _KlioReadFromBigQueryMapper(klio_message_columns)