diff --git a/CHANGES.md b/CHANGES.md index 67db578fc273..ae400b7e0519 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -70,6 +70,7 @@ * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)). * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)). +* ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)). ## Breaking Changes diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 38acd29da7d9..0caac19fb792 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -29,7 +29,11 @@ Also, for programming convenience, instances of TableReference and TableSchema have a string representation that can be used for the corresponding arguments: - - TableReference can be a PROJECT:DATASET.TABLE or DATASET.TABLE string. + - TableReference can be a PROJECT:DATASET.TABLE, PROJECT.DATASET.TABLE, + DATASET.TABLE or, for Lakehouse runtime catalog tables, + PROJECT.CATALOG.NAMESPACE.TABLE, which maps to a composite + CATALOG.NAMESPACE dataset id. A Lakehouse reference must include the + project id. - TableSchema can be a NAME:TYPE{,NAME:TYPE}* string (e.g. 'month:STRING,event_count:INTEGER'). @@ -789,6 +793,10 @@ def estimate_size(self): table_ref.projectId = self._get_project() table = bq.get_table( table_ref.projectId, table_ref.datasetId, table_ref.tableId) + if table.numBytes is None: + # Some tables don't report storage statistics, e.g. Lakehouse runtime + # catalog tables. + return None return int(table.numBytes) elif self.query is not None and self.query.is_accessible(): project = self._get_project() @@ -1111,8 +1119,24 @@ def _get_table_size(self, bq, table_reference): if table_reference.projectId else self._get_parent_project()) table = bq.get_table( project, table_reference.datasetId, table_reference.tableId) + # None for tables that don't report storage statistics, e.g. Lakehouse + # runtime catalog tables. return table.numBytes + def _get_stream_count(self, bq, desired_bundle_size): + """Number of streams to request; 0 lets the Storage Read API decide.""" + table_size = self._get_table_size(bq, self.table_reference) + if table_size is None: + # A table that reports no size, e.g. a Lakehouse runtime catalog table, + # cannot be split by size. + return 0 + stream_count = 0 + if desired_bundle_size > 0: + stream_count = min( + int(table_size / desired_bundle_size), + _CustomBigQueryStorageSource.MAX_SPLIT_COUNT) + return max(stream_count, _CustomBigQueryStorageSource.MIN_SPLIT_COUNT) + def _get_bq_metadata(self): if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) @@ -1256,14 +1280,7 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None): requested_session.read_options.row_restriction = self.row_restriction storage_client = bq_storage.BigQueryReadClient() - stream_count = 0 - if desired_bundle_size > 0: - table_size = self._get_table_size(bq, self.table_reference) - stream_count = min( - int(table_size / desired_bundle_size), - _CustomBigQueryStorageSource.MAX_SPLIT_COUNT) - stream_count = max( - stream_count, _CustomBigQueryStorageSource.MIN_SPLIT_COUNT) + stream_count = self._get_stream_count(bq, desired_bundle_size) parent = 'projects/{}'.format(self.table_reference.projectId) read_session = storage_client.create_read_session( @@ -2953,10 +2970,13 @@ class ReadFromBigQuery(PTransform): table (str, callable, ValueProvider): The ID of the table, or a callable that returns it. 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. + ``'DATASET.TABLE'``, ``'PROJECT:DATASET.TABLE'``, + ``'PROJECT.DATASET.TABLE'`` or, for Lakehouse runtime catalog tables, + ``'PROJECT.CATALOG.NAMESPACE.TABLE'``. A Lakehouse reference must + include the project id. + 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. dataset (str): The ID of the dataset containing this table or :data:`None` if the table reference is specified entirely by the table argument. diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index 874ae68362e8..d7d4cab09d9e 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -754,6 +754,58 @@ def test_read_direct_lineage(self): Lineage.query(p.result.metrics(), Lineage.SOURCE), set(["bigquery:project.dataset.table"])) + @parameterized.expand([ + # Tables without storage statistics (e.g. Lakehouse runtime catalog + # tables) let the Storage Read API pick the stream count. + param(num_bytes=None, expected_max_stream_count=0), + param( + num_bytes=5, + expected_max_stream_count=beam_bq._CustomBigQueryStorageSource. + MIN_SPLIT_COUNT), + ]) + def test_direct_read_split_stream_count( + self, num_bytes, expected_max_stream_count): + class DummyTable: + numBytes = num_bytes + + with mock.patch.object(BigQueryWrapper, '_bigquery_client'), \ + mock.patch.object(BigQueryWrapper, 'get_table', + return_value=DummyTable()), \ + mock.patch.object(bq_storage.BigQueryReadClient, + 'create_read_session') as mock_create_session: + mock_create_session.return_value = mock.Mock(streams=[]) + source = beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project.catalog.namespace.table', + pipeline_options=PipelineOptions(['--project=project'])) + self.assertEqual(source.estimate_size(), num_bytes) + self.assertEqual(list(source.split(desired_bundle_size=1 << 20)), []) + + _, kwargs = mock_create_session.call_args + self.assertEqual(kwargs['max_stream_count'], expected_max_stream_count) + self.assertEqual( + kwargs['read_session'].table, + 'projects/project/datasets/catalog.namespace/tables/table') + + @parameterized.expand([ + # A table without storage statistics has an unknown size, not a size + # of zero. + param(num_bytes=None, expected_size=None), + param(num_bytes=5, expected_size=5), + ]) + def test_export_estimate_size(self, num_bytes, expected_size): + class DummyTable: + numBytes = num_bytes + + with mock.patch.object(BigQueryWrapper, '_bigquery_client'), \ + mock.patch.object(BigQueryWrapper, 'get_table', + return_value=DummyTable()): + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project.catalog.namespace.table', + pipeline_options=PipelineOptions(['--project=project'])) + self.assertEqual(source.estimate_size(), expected_size) + def test_read_all_lineage(self): # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index abe204c237a4..cc1a16a659fe 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -106,8 +106,12 @@ BQ_STREAMING_INSERT_TIMEOUT_SEC = 120 _PROJECT_PATTERN = r'([a-z0-9.-]+:)?[a-z][a-z0-9-]*[a-z0-9]' -_DATASET_PATTERN = r'\w{1,1024}' +# Dots are allowed because Lakehouse runtime catalog tables are addressed with +# a composite 'catalog.namespace' dataset id. +_DATASET_PATTERN = r'[-\w.]{1,1024}' _TABLE_PATTERN = r'[\p{L}\p{M}\p{N}\p{Pc}\p{Pd}\p{Zs}$]{1,1024}' +# A single segment of a project id (no dots or colons). +_PROJECT_NAME_SEGMENT_PATTERN = r'[-a-z0-9]*[a-z0-9]' # For CI: temp dataset of this name are automatically deleted _TEMP_DATASET_PREFIX = 'beam_temp_dataset_' @@ -252,9 +256,15 @@ def parse_table_reference(table, dataset=None, project=None): table: The ID of the table. The ID must contain only letters (a-z, A-Z), numbers (0-9), connectors (-_). If dataset argument is None then the table argument must contain the entire table reference: - 'DATASET.TABLE' or 'PROJECT:DATASET.TABLE'. This argument can be a - TableReference instance in which case dataset and project are - ignored and the reference is returned as a result. Additionally, for date + 'DATASET.TABLE', 'PROJECT:DATASET.TABLE' or 'PROJECT.DATASET.TABLE'. + Lakehouse runtime catalog tables use four parts, + 'PROJECT.CATALOG.NAMESPACE.TABLE' (the canonical spelling) or + 'PROJECT:CATALOG.NAMESPACE.TABLE', which parse to a composite + 'CATALOG.NAMESPACE' dataset id. The project id is required for these, + because a composite dataset id without one would be ambiguous with + 'PROJECT.DATASET.TABLE'. This argument can be a TableReference instance + in which case dataset and project are ignored and the reference is + returned as a result. Additionally, for date partitioned tables, appending '$YYYYmmdd' to the table name is supported, e.g. 'DATASET.TABLE$YYYYmmdd'. dataset: The ID of the dataset containing this table or null if the table @@ -288,17 +298,10 @@ def parse_table_reference(table, dataset=None, project=None): # table argument will contain a full table reference instead of just a # table name. if dataset is None: - pattern = ( - f'((?P{_PROJECT_PATTERN})[:\\.])?' - f'(?P{_DATASET_PATTERN})\\.(?P{_TABLE_PATTERN})') - match = regex.fullmatch(pattern, table) - if not match: - raise ValueError( - 'Expected a table reference (PROJECT:DATASET.TABLE or ' - 'DATASET.TABLE) instead of %s.' % table) - table_reference.projectId = match.group('project') - table_reference.datasetId = match.group('dataset') - table_reference.tableId = match.group('table') + project_id, dataset_id, table_id = _split_table_spec(table) + table_reference.projectId = project_id + table_reference.datasetId = dataset_id + table_reference.tableId = table_id else: table_reference.projectId = project table_reference.datasetId = dataset @@ -306,6 +309,73 @@ def parse_table_reference(table, dataset=None, project=None): return table_reference +def _invalid_table_spec(table_spec): + return ValueError( + 'Expected a table reference (PROJECT:DATASET.TABLE, ' + 'PROJECT.DATASET.TABLE, DATASET.TABLE, ' + 'PROJECT:CATALOG.NAMESPACE.TABLE or PROJECT.CATALOG.NAMESPACE.TABLE) ' + 'instead of %s.' % table_spec) + + +def _split_table_spec(table_spec): + """Splits a table spec string into (project, dataset, table). + + The regex only validates the character set; segment assignment is done by + explicit splitting so that composite Lakehouse dataset ids ('catalog.ns') + and domain-scoped project ids ('example.com:proj') are both handled. This + mirrors BigQueryHelpers.parseTableSpec in the Java SDK. + """ + pattern = ( + f'((?P{_PROJECT_PATTERN})[:\\.])?' + f'(?P{_DATASET_PATTERN})\\.(?P
{_TABLE_PATTERN})') + if not regex.fullmatch(pattern, table_spec): + raise _invalid_table_spec(table_spec) + + # Table ids cannot contain '.', so the table is always the last segment. + last_dot = table_spec.rfind('.') + table = table_spec[last_dot + 1:] + prefix = table_spec[:last_dot] + + colon_count = prefix.count(':') + if colon_count == 0: + # Purely dotted form: 'd.t', 'p.d.t', 'p.catalog.ns.t'. A dataset id may + # itself contain dots (a composite 'catalog.namespace'), so omitting the + # project id would be ambiguous with 'PROJECT.DATASET.TABLE'. Only the + # plain two-segment 'DATASET.TABLE' may leave the project id out. + first_dot = prefix.find('.') + if first_dot < 0: + return None, prefix, table + project = prefix[:first_dot] + dataset = prefix[first_dot + 1:] + if not regex.fullmatch(_PROJECT_PATTERN, project) or not dataset: + raise _invalid_table_spec(table_spec) + return project, dataset, table + + if colon_count == 1: + # 'p:d.t', 'p:catalog.ns.t', 'example.com:proj.ds.t'. If the part before + # the colon is dotted it is the domain half of a legacy domain-scoped + # project id; the first segment after the colon completes the project + # and any remaining segments are the (possibly composite) dataset. + colon = prefix.find(':') + project = prefix[:colon] + dataset = prefix[colon + 1:] + first_dot = dataset.find('.') + if (first_dot >= 0 and first_dot < len(dataset) - 1 and '.' in project and + regex.fullmatch(_PROJECT_NAME_SEGMENT_PATTERN, dataset[:first_dot])): + project = project + ':' + dataset[:first_dot] + dataset = dataset[first_dot + 1:] + return project, dataset, table + + if colon_count == 2: + # The last colon terminates a domain-scoped project id spelled with its + # own colon: 'example.com:proj:ds.t' or 'example.com:proj:catalog.ns.t'. + last_colon = prefix.rfind(':') + return prefix[:last_colon], prefix[last_colon + 1:], table + + # Project ids contain at most one colon, so more than two is never valid. + raise _invalid_table_spec(table_spec) + + # ----------------------------------------------------------------------------- # BigQueryWrapper. diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index c2a5c73f6755..9f8b78ee5c3b 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -136,6 +136,46 @@ def test_calling_with_value_provider(self): ('project:dataset.test- table', 'project', 'dataset', 'test- table'), ('project.dataset. test_table', 'project', 'dataset', ' test_table'), ('project.dataset.test$table', 'project', 'dataset', 'test$table'), + ('project.dataset.table', 'project', 'dataset', 'table'), + ('my-project.my_dataset.table', 'my-project', 'my_dataset', 'table'), + # Lakehouse runtime catalog tables: composite 'catalog.namespace' + # dataset id. + ( + 'project.catalog.namespace.table', + 'project', + 'catalog.namespace', + 'table'), + ( + 'project:catalog.namespace.table', + 'project', + 'catalog.namespace', + 'table'), + ( + 'my-project.my-catalog.ns.events$20240101', + 'my-project', + 'my-catalog.ns', + 'events$20240101'), + # Domain-scoped project ids, in both spellings. + ( + 'example.com:proj.dataset.table', + 'example.com:proj', + 'dataset', + 'table'), + ( + 'example.com:proj:dataset.table', + 'example.com:proj', + 'dataset', + 'table'), + ( + 'example.com:proj.catalog.namespace.table', + 'example.com:proj', + 'catalog.namespace', + 'table'), + ( + 'example.com:proj:catalog.namespace.table', + 'example.com:proj', + 'catalog.namespace', + 'table'), ]) def test_calling_with_fully_qualified_table_ref( self, @@ -159,10 +199,33 @@ def test_calling_with_partially_qualified_table_ref(self): self.assertEqual(parsed_ref.datasetId, datasetId) self.assertEqual(parsed_ref.tableId, tableId) + def test_composite_dataset_requires_project(self): + # A composite 'catalog.namespace' dataset id is only recognised with an + # explicit project id. Without one it would be ambiguous with + # 'PROJECT.DATASET.TABLE', so it is rejected rather than guessed at. + self.assertRaises( + ValueError, parse_table_reference, 'my_catalog.namespace.test_table') + def test_calling_with_insufficient_table_ref(self): table = 'test_table' self.assertRaises(ValueError, parse_table_reference, table) + @parameterized.expand([ + ('a:b:c:d.table', ), + ('project:.table', ), + ('project:dataset', ), + ('MyProject.MyDataset.table', ), + ('1project.dataset.table', ), + ('c.namespace.table', ), + ('.dataset.table', ), + ('project..table', ), + ('a.b.c.d.e.f.g', ), + ]) + def test_calling_with_invalid_table_ref(self, table): + # The dotted specs here have a leading segment that cannot be a project + # id; they are rejected rather than bound as a composite dataset id. + self.assertRaises(ValueError, parse_table_reference, table) + def test_calling_with_all_arguments(self): projectId = 'test_project' datasetId = 'test_dataset'