[Python] Support Lakehouse runtime catalog tables in ReadFromBigQuery…

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
 

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.

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

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,24 +298,84 @@ 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>{_PROJECT_PATTERN})[:\\.])?'
-        f'(?P<dataset>{_DATASET_PATTERN})\\.(?P<table>{_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
     table_reference.tableId = table
   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>{_PROJECT_PATTERN})[:\\.])?'
+      f'(?P<dataset>{_DATASET_PATTERN})\\.(?P<table>{_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.
 

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'