feat: support biglake tables in pandas_gbq.sample

.gitignore

@@ -51,6 +51,7 @@ docs.metadata
 # Virtual environment
 env/
 venv/
+.venv/
 
 # Test logs
 coverage.xml

pandas_gbq/core/biglake.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2026 pandas-gbq Authors All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.
+
+"""
+Utilities for working with BigLake tables.
+"""
+
+# TODO(tswast): Synchronize with bigframes/session/iceberg.py, which uses
+# pyiceberg and the BigLake APIs, rather than relying on dry run.
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Sequence
+
+import google.auth.transport.requests
+import google.cloud.bigquery
+import google.oauth2.credentials
+
+import pandas_gbq.core.resource_references
+
+
+_DRY_RUN_TEMPLATE = """
+SELECT *
+FROM `{project}.{catalog}.{namespace}.{table}`
+"""
+
+
+_COUNT_TEMPLATE = """
+SELECT COUNT(*) as total_rows
+FROM `{project}.{catalog}.{namespace}.{table}`
+"""
+
+@dataclasses.dataclass(frozen=True)
+class BigLakeTableMetadata:
+    schema: Sequence[google.cloud.bigquery.SchemaField]
+    num_rows: int
+
+
+def get_table_metadata(
+    *,
+    reference: pandas_gbq.core.resource_references.BigLakeTableId,
+    bqclient: google.cloud.bigquery.Client,
+) -> BigLakeTableMetadata:
+    """
+    Get the schema for a BigLake table.
+
+    Currently, this does some BigQuery queries. In the future, we'll want to get
+    other metadata like the number of rows and storage bytes so that we can do a
+    more accurate estimate of how many rows to sample.
+    """
+    dry_run_config = google.cloud.bigquery.QueryJobConfig(dry_run=True)
+    query = _DRY_RUN_TEMPLATE.format(
+        project=reference.project,
+        catalog=reference.catalog,
+        namespace=".".join(reference.namespace),
+        table=reference.table,
+    )
+    job = bqclient.query(query, job_config=dry_run_config)
+    job.result()
+    schema = job.schema
+
+    count_rows = list(bqclient.query_and_wait(_COUNT_TEMPLATE.format(
+        project=reference.project,
+        catalog=reference.catalog,
+        namespace=".".join(reference.namespace),
+        table=reference.table,
+    )))
+    assert len(count_rows) == 1, "got unexpected query response when determining number of rows"
+    total_rows = count_rows[0].total_rows
+
+    return BigLakeTableMetadata(
+        schema=schema if schema is not None else [],
+        num_rows=total_rows,
+    )
+

pandas_gbq/core/resource_references.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2026 pandas-gbq Authors All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.
+
+import dataclasses
+import re
+
+
+_TABLE_REFEREENCE_PATTERN = re.compile(
+    # In the past, organizations could prefix their project IDs with a domain
+    # name. Such projects still exist, especially at Google.
+    r"^(?P<legacy_project_domain>[^:]+:)?"
+    r"(?P<project>[^.]+)\."
+    # Dataset for native BigQuery tables, catalog + namespace(s) for  BigLake.
+    r"(?P<inner_parts>([^.\s]+\.?)+)\."
+    # Table names can't contain ".", as that's used as the separator.
+    r"(?P<table>[^.]+)$"
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class BigLakeTableId:
+    project: str
+    catalog: str
+    namespace: tuple[str, ...]
+    table: str
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryTableId:
+    project_id: str
+    dataset_id: str
+    table_id: str
+
+
+def parse_table_id(table_id: str) -> BigLakeTableId | BigQueryTableId:
+    """Turn a string into a BigLakeTableId or BigQueryTableId.
+
+    Raises:
+        ValueError: If the table ID is invalid.
+    """
+    regex_match = _TABLE_REFEREENCE_PATTERN.match(table_id)
+    if not regex_match:
+        raise ValueError(f"Invalid table ID: {table_id}")
+
+    inner_parts = regex_match.group("inner_parts").split(".")
+    if len(inner_parts) == 1:
+        return BigQueryTableId(
+            project_id=regex_match.group("project"),
+            dataset_id=inner_parts[0],
+            table_id=regex_match.group("table"),
+        )
+
+    return BigLakeTableId(
+        project=regex_match.group("project"),
+        catalog=inner_parts[0],
+        namespace=tuple(inner_parts[1:]),
+        table=regex_match.group("table"),
+    )