satybald commented on a change in pull request #15185:
URL: https://github.com/apache/beam/pull/15185#discussion_r683832113
##########
File path: sdks/python/apache_beam/io/gcp/bigquery.py
##########
@@ -883,6 +892,272 @@ def _export_files(self, bq):
return table.schema, metadata_list
+class _CustomBigQueryStorageSourceBase(BoundedSource):
+ """A base class for BoundedSource implementations which read from BigQuery
+ using the BigQuery Storage API.
+
+ Args:
+ table (str, TableReference): 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:
+ ``'PROJECT:DATASET.TABLE'`` or must specify a TableReference.
+ dataset (str): The ID of the dataset containing this table or
+ :data:`None` if the table argument specifies a TableReference.
+ project (str): The ID of the project containing this table or
+ :data:`None` if the table argument specifies a TableReference.
+ selected_fields (List[str]): Names of the fields in the table that should
be
+ read. If empty, all fields will be read. If the specified field is a
+ nested field, all the sub-fields in the field will be selected. The
output
+ field order is unrelated to the order of fields in selected_fields.
+ row_restriction (str): SQL text filtering statement, similar to a WHERE
+ clause in a query. Aggregates are not supported.Restricted to a maximum
+ length for 1 MB.
+ """
+
+ # The maximum number of streams which will be requested when creating a read
+ # session, regardless of the desired bundle size.
+ MAX_SPLIT_COUNT = 10000
+ # The minimum number of streams which will be requested when creating a read
+ # session, regardless of the desired bundle size. Note that the server may
+ # still choose to return fewer than ten streams based on the layout of the
+ # table.
+ MIN_SPLIT_COUNT = 10
+
+ def __init__(
+ self,
+ table: Union[str, TableReference],
+ dataset: str = None,
+ project: str = None,
+ selected_fields: List[str] = None,
+ row_restriction: str = None,
+ use_fastavro: bool = None,
+ pipeline_options: GoogleCloudOptions = None):
+
+ self.table_reference = bigquery_tools.parse_table_reference(
+ table, dataset, project)
+ self.table = self.table_reference.tableId
+ self.dataset = self.table_reference.datasetId
+ self.project = self.table_reference.projectId
+ self.selected_fields = selected_fields
+ self.row_restriction = row_restriction
+ self.use_fastavro = True if use_fastavro is not None else use_fastavro
+ self.pipeline_options = pipeline_options
+ self.split_result = None
+
+ def _get_parent_project(self):
+ """Returns the project that will be billed."""
+ project = self.pipeline_options.view_as(GoogleCloudOptions).project
+ if isinstance(project, vp.ValueProvider):
+ project = project.get()
+ if not project:
+ project = self.project
+ return project
+
+ def _get_table_size(self, table, dataset, project):
+ if project is None:
+ project = self._get_parent_project()
+
+ bq = bigquery_tools.BigQueryWrapper()
+ table = bq.get_table(project, dataset, table)
+ return table.numBytes
+
+ def display_data(self):
+ return {
+ 'project': str(self.project),
+ 'dataset': str(self.dataset),
+ 'table': str(self.table),
+ 'selected_fields': str(self.selected_fields),
+ 'row_restriction': str(self.row_restriction)
Review comment:
Should we also log `use_fastavro` flag option for troubleshooting
purpose?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
