From 576282aba7378fa208ed5cd8c8aefee1daa7d0f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 22 Sep 2026 12:02:22 -0400 Subject: [PATCH 1/3] [Python] Support Lakehouse runtime catalog tables in ReadFromBigQuery DIRECT_READ Port of the Java change in #39597 to the Python SDK. - parse_table_reference now assigns segments by explicit splitting (mirroring BigQueryHelpers.parseTableSpec) so 4-part `project.catalog.namespace.table` and `project:catalog.namespace.table` specs map to a composite `catalog.namespace` dataset id. Previously the dotted form raised and the colon form silently bound the project to `project:catalog`. - Storage Read split tolerates tables that report no numBytes by requesting zero streams so the Read API picks the stream count; the export source's size estimate returns 0 instead of raising. - Unit tests for the new spec forms, domain-scoped projects, invalid specs, and the numBytes-missing paths. --- CHANGES.md | 1 + sdks/python/apache_beam/io/gcp/bigquery.py | 39 +++++--- .../apache_beam/io/gcp/bigquery_test.py | 50 ++++++++++ .../apache_beam/io/gcp/bigquery_tools.py | 94 ++++++++++++++++--- .../apache_beam/io/gcp/bigquery_tools_test.py | 56 +++++++++++ 5 files changed, 214 insertions(+), 26 deletions(-) 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..802270e68f4d 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -29,7 +29,10 @@ 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 or + DATASET.TABLE string. Lakehouse runtime catalog (BigLake metastore) tables + use PROJECT.CATALOG.NAMESPACE.TABLE, which maps to a composite + CATALOG.NAMESPACE dataset id. - TableSchema can be a NAME:TYPE{,NAME:TYPE}* string (e.g. 'month:STRING,event_count:INTEGER'). @@ -789,6 +792,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: + # Tables that don't report storage statistics, e.g. Lakehouse runtime + # catalog (BigLake metastore) tables. + return 0 return int(table.numBytes) elif self.query is not None and self.query.is_accessible(): project = self._get_project() @@ -1111,6 +1118,8 @@ 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 (BigLake metastore) tables. return table.numBytes def _get_bq_metadata(self): @@ -1256,14 +1265,18 @@ 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() + # A stream_count of 0 lets the Storage Read API choose the number of + # streams; used when the table reports no size (e.g. Lakehouse runtime + # catalog tables). 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) + table_size = self._get_table_size(bq, self.table_reference) + if table_size is not None: + if desired_bundle_size > 0: + stream_count = min( + int(table_size / desired_bundle_size), + _CustomBigQueryStorageSource.MAX_SPLIT_COUNT) + stream_count = max( + stream_count, _CustomBigQueryStorageSource.MIN_SPLIT_COUNT) parent = 'projects/{}'.format(self.table_reference.projectId) read_session = storage_client.create_read_session( @@ -2953,10 +2966,12 @@ 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'`` or + ``'PROJECT.DATASET.TABLE'``. Lakehouse runtime catalog (BigLake + metastore) tables are specified as ``'PROJECT.CATALOG.NAMESPACE.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 (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..b024c9cf08d1 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -754,6 +754,56 @@ 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([ + param(num_bytes=None, expected_size=0), + 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..2842e3c421fb 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 (BigLake metastore) 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,8 +256,11 @@ 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 + 'DATASET.TABLE', 'PROJECT:DATASET.TABLE' or 'PROJECT.DATASET.TABLE'. + Lakehouse runtime catalog (BigLake metastore) tables use four parts, + 'PROJECT.CATALOG.NAMESPACE.TABLE' or 'PROJECT:CATALOG.NAMESPACE.TABLE', + which parse to a composite 'CATALOG.NAMESPACE' dataset id. 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'. @@ -288,17 +295,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 +306,72 @@ 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: 'p.d.t', 'd.t', 'p.catalog.ns.t'. The leading + # segment is the project when it looks like one; dataset ids may contain + # characters such as '_' that project ids may not, in which case the + # whole prefix is the dataset. The first_dot < len - 1 guard keeps + # degenerate trailing-dot specs from producing an empty dataset id. + first_dot = prefix.find('.') + leading = prefix[:first_dot] if first_dot >= 0 else None + if (leading is not None and first_dot < len(prefix) - 1 and + regex.fullmatch(_PROJECT_PATTERN, leading)): + return leading, prefix[first_dot + 1:], table + return None, prefix, 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..3c1f60cae941 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,26 @@ def test_calling_with_partially_qualified_table_ref(self): self.assertEqual(parsed_ref.datasetId, datasetId) self.assertEqual(parsed_ref.tableId, tableId) + def test_calling_with_partially_qualified_composite_dataset(self): + # A leading segment that cannot be a project id (underscore) binds the + # whole prefix as a composite dataset id. + parsed_ref = parse_table_reference('my_catalog.namespace.test_table') + self.assertIsNone(parsed_ref.projectId) + self.assertEqual(parsed_ref.datasetId, 'my_catalog.namespace') + self.assertEqual(parsed_ref.tableId, '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', ), + ]) + def test_calling_with_invalid_table_ref(self, table): + self.assertRaises(ValueError, parse_table_reference, table) + def test_calling_with_all_arguments(self): projectId = 'test_project' datasetId = 'test_dataset' From bcd884236d48f0e68ed31f4bb4b728ad9e75e17b Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 2 Oct 2026 17:55:43 -0400 Subject: [PATCH 2/3] comments --- sdks/python/apache_beam/io/gcp/bigquery.py | 50 +++++++++++-------- .../apache_beam/io/gcp/bigquery_test.py | 4 +- .../apache_beam/io/gcp/bigquery_tools.py | 11 ++-- 3 files changed, 39 insertions(+), 26 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 802270e68f4d..adae12335bf5 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -29,10 +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, PROJECT.DATASET.TABLE or - DATASET.TABLE string. Lakehouse runtime catalog (BigLake metastore) tables - use PROJECT.CATALOG.NAMESPACE.TABLE, which maps to a composite - CATALOG.NAMESPACE dataset id. + - TableReference can be a PROJECT:DATASET.TABLE, PROJECT.DATASET.TABLE, + DATASET.TABLE or, for Lakehouse runtime catalog (BigLake metastore) + tables, PROJECT.CATALOG.NAMESPACE.TABLE, which maps to a composite + CATALOG.NAMESPACE dataset id. A Lakehouse reference must include the + project id: a three-part string is read as PROJECT.DATASET.TABLE. - TableSchema can be a NAME:TYPE{,NAME:TYPE}* string (e.g. 'month:STRING,event_count:INTEGER'). @@ -793,9 +794,9 @@ def estimate_size(self): table = bq.get_table( table_ref.projectId, table_ref.datasetId, table_ref.tableId) if table.numBytes is None: - # Tables that don't report storage statistics, e.g. Lakehouse runtime + # Some tables don't report storage statistics, e.g. Lakehouse runtime # catalog (BigLake metastore) tables. - return 0 + return None return int(table.numBytes) elif self.query is not None and self.query.is_accessible(): project = self._get_project() @@ -1122,6 +1123,22 @@ def _get_table_size(self, bq, table_reference): # runtime catalog (BigLake metastore) tables. return table.numBytes + def _get_stream_count(self, bq, desired_bundle_size): + """Number of streams to request, or 0 to let the Storage Read API decide. + + A table that reports no size, e.g. a Lakehouse runtime catalog (BigLake + metastore) table, cannot be split by size. + """ + table_size = self._get_table_size(bq, self.table_reference) + if table_size is None: + 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) @@ -1265,18 +1282,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() - # A stream_count of 0 lets the Storage Read API choose the number of - # streams; used when the table reports no size (e.g. Lakehouse runtime - # catalog tables). - stream_count = 0 - table_size = self._get_table_size(bq, self.table_reference) - if table_size is not None: - if desired_bundle_size > 0: - 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( @@ -2966,9 +2972,11 @@ 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'``, ``'PROJECT:DATASET.TABLE'`` or - ``'PROJECT.DATASET.TABLE'``. Lakehouse runtime catalog (BigLake - metastore) tables are specified as ``'PROJECT.CATALOG.NAMESPACE.TABLE'``. + ``'DATASET.TABLE'``, ``'PROJECT:DATASET.TABLE'``, + ``'PROJECT.DATASET.TABLE'`` or, for Lakehouse runtime catalog (BigLake + metastore) tables, ``'PROJECT.CATALOG.NAMESPACE.TABLE'``. A Lakehouse + reference must include the project id, because a three-part string is + read as ``'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. diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index b024c9cf08d1..d7d4cab09d9e 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -788,7 +788,9 @@ class DummyTable: 'projects/project/datasets/catalog.namespace/tables/table') @parameterized.expand([ - param(num_bytes=None, expected_size=0), + # 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): diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 2842e3c421fb..d07b09045a7a 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -258,10 +258,13 @@ def parse_table_reference(table, dataset=None, project=None): then the table argument must contain the entire table reference: 'DATASET.TABLE', 'PROJECT:DATASET.TABLE' or 'PROJECT.DATASET.TABLE'. Lakehouse runtime catalog (BigLake metastore) tables use four parts, - 'PROJECT.CATALOG.NAMESPACE.TABLE' or 'PROJECT:CATALOG.NAMESPACE.TABLE', - which parse to a composite 'CATALOG.NAMESPACE' dataset id. 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 + '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 not optional for + these: a three-part string binds as 'PROJECT.DATASET.TABLE' whenever its + first segment is a valid project id. 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 From 91aac7e86ffa8f8f943e47387d6e029f594d4b49 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 4 Oct 2026 14:07:01 -0400 Subject: [PATCH 3/3] comments --- sdks/python/apache_beam/io/gcp/bigquery.py | 25 ++++++------- .../apache_beam/io/gcp/bigquery_tools.py | 37 ++++++++++--------- .../apache_beam/io/gcp/bigquery_tools_test.py | 21 +++++++---- 3 files changed, 44 insertions(+), 39 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index adae12335bf5..0caac19fb792 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -30,10 +30,10 @@ have a string representation that can be used for the corresponding arguments: - TableReference can be a PROJECT:DATASET.TABLE, PROJECT.DATASET.TABLE, - DATASET.TABLE or, for Lakehouse runtime catalog (BigLake metastore) - tables, PROJECT.CATALOG.NAMESPACE.TABLE, which maps to a composite + 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: a three-part string is read as PROJECT.DATASET.TABLE. + project id. - TableSchema can be a NAME:TYPE{,NAME:TYPE}* string (e.g. 'month:STRING,event_count:INTEGER'). @@ -795,7 +795,7 @@ def estimate_size(self): 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 (BigLake metastore) tables. + # catalog tables. return None return int(table.numBytes) elif self.query is not None and self.query.is_accessible(): @@ -1120,17 +1120,15 @@ def _get_table_size(self, bq, table_reference): 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 (BigLake metastore) tables. + # runtime catalog tables. return table.numBytes def _get_stream_count(self, bq, desired_bundle_size): - """Number of streams to request, or 0 to let the Storage Read API decide. - - A table that reports no size, e.g. a Lakehouse runtime catalog (BigLake - metastore) table, cannot be split by 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: @@ -2973,10 +2971,9 @@ class ReadFromBigQuery(PTransform): that returns it. If dataset argument is :data:`None` then the table argument must contain the entire table reference specified as: ``'DATASET.TABLE'``, ``'PROJECT:DATASET.TABLE'``, - ``'PROJECT.DATASET.TABLE'`` or, for Lakehouse runtime catalog (BigLake - metastore) tables, ``'PROJECT.CATALOG.NAMESPACE.TABLE'``. A Lakehouse - reference must include the project id, because a three-part string is - read as ``'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. diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index d07b09045a7a..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,8 @@ BQ_STREAMING_INSERT_TIMEOUT_SEC = 120 _PROJECT_PATTERN = r'([a-z0-9.-]+:)?[a-z][a-z0-9-]*[a-z0-9]' -# Dots are allowed because Lakehouse runtime catalog (BigLake metastore) tables -# are addressed with a composite 'catalog.namespace' dataset id. +# 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). @@ -257,14 +257,14 @@ def parse_table_reference(table, dataset=None, project=None): (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', 'PROJECT:DATASET.TABLE' or 'PROJECT.DATASET.TABLE'. - Lakehouse runtime catalog (BigLake metastore) tables use four parts, + 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 not optional for - these: a three-part string binds as 'PROJECT.DATASET.TABLE' whenever its - first segment is a valid project id. 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 + '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 @@ -338,17 +338,18 @@ def _split_table_spec(table_spec): colon_count = prefix.count(':') if colon_count == 0: - # Purely dotted form: 'p.d.t', 'd.t', 'p.catalog.ns.t'. The leading - # segment is the project when it looks like one; dataset ids may contain - # characters such as '_' that project ids may not, in which case the - # whole prefix is the dataset. The first_dot < len - 1 guard keeps - # degenerate trailing-dot specs from producing an empty dataset id. + # 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('.') - leading = prefix[:first_dot] if first_dot >= 0 else None - if (leading is not None and first_dot < len(prefix) - 1 and - regex.fullmatch(_PROJECT_PATTERN, leading)): - return leading, prefix[first_dot + 1:], table - return None, prefix, table + 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 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 3c1f60cae941..9f8b78ee5c3b 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -199,13 +199,12 @@ def test_calling_with_partially_qualified_table_ref(self): self.assertEqual(parsed_ref.datasetId, datasetId) self.assertEqual(parsed_ref.tableId, tableId) - def test_calling_with_partially_qualified_composite_dataset(self): - # A leading segment that cannot be a project id (underscore) binds the - # whole prefix as a composite dataset id. - parsed_ref = parse_table_reference('my_catalog.namespace.test_table') - self.assertIsNone(parsed_ref.projectId) - self.assertEqual(parsed_ref.datasetId, 'my_catalog.namespace') - self.assertEqual(parsed_ref.tableId, 'test_table') + 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' @@ -215,8 +214,16 @@ def test_calling_with_insufficient_table_ref(self): ('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):