Support lakehouse PCNT format in BigQueryIO storage read.

.github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run.",
-    "modification": 2
+    "modification": 3
 }

CHANGES.md

@@ -65,6 +65,7 @@
 ## I/Os
 
 * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* BigQueryIO now supports reading BigQuery Lakehouse runtime catalog (BigLake metastore) Iceberg tables with the Storage Read API, using 4-part `project.catalog.namespace.table` identifiers (or a `TableReference` with a composite `catalog.namespace` dataset id). Previously such references were silently mis-parsed (Java) ([#39597](https://github.com/apache/beam/issues/39597)) .
 
 ## New Features / Improvements
 

runners/google-cloud-dataflow-java/build.gradle

@@ -683,6 +683,7 @@ task googleCloudPlatformLegacyWorkerIntegrationTest(type: Test, dependsOn: copyG
   exclude '**/BigQueryIODynamicReadIT.class'
   exclude '**/BigQueryIODynamicReadTableRowIT.class'
   exclude '**/BigQueryTimestampPicosIT.java'
+  exclude '**/BigQueryIOIcebergManagedTableIT.class'
   exclude '**/PubsubReadIT.class'
   exclude '**/FhirIOReadIT.class'
   exclude '**/DicomIOReadIT.class'
@@ -740,6 +741,7 @@ task googleCloudPlatformRunnerV2IntegrationTest(type: Test) {
   exclude '**/BigQueryIODynamicReadIT.class'
   exclude '**/BigQueryIODynamicReadTableRowIT.class'
   exclude '**/BigQueryTimestampPicosIT.java'
+  exclude '**/BigQueryIOIcebergManagedTableIT.class'
   exclude '**/SpannerWriteIT.class'
   exclude '**/*KmsKeyIT.class'
   exclude '**/FhirIOReadIT.class'

sdks/java/io/google-cloud-platform/build.gradle

@@ -203,6 +203,7 @@ task integrationTest(type: Test, dependsOn: processTestResources) {
   outputs.upToDateWhen { false }
 
   include '**/*IT.class'
+  exclude '**/BigQueryIOIcebergManagedTableIT.class'
   exclude '**/BigQueryIOReadIT.class'
   exclude '**/BigQueryIOStorageQueryIT.class'
   exclude '**/BigQueryIOStorageReadIT.class'
@@ -252,6 +253,37 @@ task integrationTestKms(type: Test) {
   }
 }
 
+/*
+  Integration tests for BigQueryIO against managed Apache Iceberg tables. Kept out of the
+  generic integrationTest task because they need a provisioned CLOUD_RESOURCE connection;
+  override the apache-beam-testing defaults with -PgcpProject, -PbqImtConnection and
+  -PbqImtStorageUri.
+ */
+task bigQueryManagedIcebergIntegrationTest(type: Test, dependsOn: processTestResources) {
+  group = "Verification"
+  def gcpProject = project.findProperty('gcpProject') ?: 'apache-beam-testing'
+  def gcpTempRoot = project.findProperty('gcpTempRoot') ?: 'gs://temp-storage-for-end-to-end-tests'
+  def bqImtConnection = project.findProperty('bqImtConnection') ?:
+          'projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete'
+  def bqImtStorageUri = project.findProperty('bqImtStorageUri') ?: 'gs://apache-beam-testing-bq-biglake'
+  systemProperty "beamTestPipelineOptions", JsonOutput.toJson([
+          "--runner=DirectRunner",
+          "--project=${gcpProject}",
+          "--tempRoot=${gcpTempRoot}",
+          "--tempLocation=${gcpTempRoot}",
+  ])
+  systemProperty "beam.bq.imt.connection", bqImtConnection
+  systemProperty "beam.bq.imt.storageUri", bqImtStorageUri
+
+  outputs.upToDateWhen { false }
+
+  include '**/BigQueryIOIcebergManagedTableIT.class'
+
+  maxParallelForks 4
+  classpath = sourceSets.test.runtimeClasspath
+  testClassesDirs = sourceSets.test.output.classesDirs
+}
+
 /*
   Integration tests for BigQueryIO that run on BigQuery's early rollout region (us-east7)
   with the intended purpose of catching breaking changes from new BigQuery releases.
@@ -358,4 +390,5 @@ task postCommit {
   description = "Integration tests of GCP connectors using the DirectRunner."
   dependsOn integrationTest
   dependsOn integrationTestKms
+  dependsOn bigQueryManagedIcebergIntegrationTest
 }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java

@@ -45,6 +45,7 @@
 import java.util.Map;
 import java.util.UUID;
 import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.extensions.gcp.util.BackOffAdapter;
 import org.apache.beam.sdk.io.FileSystems;
@@ -79,6 +80,9 @@ public class BigQueryHelpers {
 
   private static final Logger LOG = LoggerFactory.getLogger(BigQueryHelpers.class);
 
+  /** Matches valid Project ID patterns. */
+  private static final Pattern PROJECT_NAME_SEGMENT_PATTERN = Pattern.compile("[-a-z0-9]*[a-z0-9]");
+
   // Given a potential failure and a current job-id, return the next job-id to be used on retry.
   // Algorithm is as follows (given input of job_id_prefix-N)
   //   If BigQuery has no status for job_id_prefix-n, we should retry with the same id.
@@ -458,6 +462,13 @@ static <K, V> List<V> getOrCreateMapListValue(Map<K, List<V>> map, K key) {
    * Parse a table specification in the form {@code "[project_id]:[dataset_id].[table_id]"} or
    * {@code "[project_id].[dataset_id].[table_id]"} or {@code "[dataset_id].[table_id]"}.
    *
+   * <p>Lakehouse runtime catalog (BigLake metastore) tables are referenced with four parts, {@code
+   * "[project_id].[catalog_id].[namespace_id].[table_id]"} (or {@code
+   * "[project_id]:[catalog_id].[namespace_id].[table_id]"}); these parse to a composite {@code
+   * "[catalog_id].[namespace_id]"} dataset id, which is the form the BigQuery APIs accept for such
+   * tables. More generally, when a specification contains more than three segments, everything
+   * between the project id and the final (table) segment becomes the dataset id.
+   *
    * <p>If the project id is omitted, the default project id is used.
    */
   @SuppressWarnings({
@@ -466,19 +477,88 @@ static <K, V> List<V> getOrCreateMapListValue(Map<K, List<V>> map, K key) {
   public static TableReference parseTableSpec(String tableSpec) {
     Matcher match = BigQueryIO.TABLE_SPEC.matcher(tableSpec);
     if (!match.matches()) {
-      throw new IllegalArgumentException(
-          String.format(
-              "Table specification [%s] is not in one of the expected formats ("
-                  + " [project_id]:[dataset_id].[table_id],"
-                  + " [project_id].[dataset_id].[table_id],"
-                  + " [dataset_id].[table_id])",
-              tableSpec));
+      throw invalidTableSpec(tableSpec);
+    }
+
+    // Table ids cannot contain '.', so the table is always the segment after
+    // the last dot.
+    int lastDot = tableSpec.lastIndexOf('.');
+    String table = tableSpec.substring(lastDot + 1);
+    String prefix = tableSpec.substring(0, lastDot);
+
+    String project = null;
+    String dataset;
+    long colonCount = prefix.chars().filter(c -> c == ':').count();
+    if (colonCount == 0) {
+      // No colon means the purely dotted form ("p.d.t", "d.t", "p.catalog.ns.t"): the
+      // leading segment is the project id when it is a plausible project id.
+      // (Dataset ids may contain characters such as '_' that project ids may
+      // not, in which case the whole prefix is the dataset id.)
+      // The firstDot < length-1 guard keeps degenerate trailing-dot specs
+      // ("pp..t", accepted by the character-set gate with dataset "pp.")
+      // instead of producing an empty dataset id.
+      int firstDot = prefix.indexOf('.');
+      if (firstDot >= 0
+          && firstDot < prefix.length() - 1
+          && BigQueryIO.PROJECT_ID_PATTERN.matcher(prefix.substring(0, firstDot)).matches()) {
+        project = prefix.substring(0, firstDot);
+        dataset = prefix.substring(firstDot + 1);
+      } else {
+        dataset = prefix;
+      }
+    } else if (colonCount == 1) {
+      // One colon ("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 ("example.com:proj") spelled with '.',
+      // the first segment after the colon completes the project id, and any
+      // remaining middle segments bind as a (possibly composite) dataset.
+      // (Only the domain half of a domain-scoped id may contain dots, so
+      // exactly one segment can complete it.)
+      int colon = prefix.indexOf(':');
+      project = prefix.substring(0, colon);
+      dataset = prefix.substring(colon + 1);
+      int firstDot = dataset.indexOf('.');
+      // Absorb the first dataset segment into a dotted (domain-scoped) project
+      // only when the split leaves a non-empty dataset.
+      if (firstDot >= 0
+          && firstDot < dataset.length() - 1
+          && project.indexOf('.') >= 0
+          && PROJECT_NAME_SEGMENT_PATTERN.matcher(dataset.substring(0, firstDot)).matches()) {
+        project = project + ":" + dataset.substring(0, firstDot);
+        dataset = dataset.substring(firstDot + 1);
+      }
+    } else if (colonCount == 2) {
+      // Two colons - the last colon is an explicit project terminator. This is
+      // the canonical spelling for a domain-scoped project, whose id itself
+      // contains a colon ("example.com:proj:ds.t"), including with a composite
+      // Lakehouse catalog dataset ("example.com:proj:catalog.ns.t"). Both
+      // domain-scoped spellings keep toTableSpec/parseTableSpec a round trip
+      // for composite dataset ids.
+      int lastColon = prefix.lastIndexOf(':');
+      project = prefix.substring(0, lastColon);
+      dataset = prefix.substring(lastColon + 1);
+    } else {
+      // Project ids contain at most one colon (domain-scoped), so no valid
+      // reference has more than two. Fail early rather than binding an
+      // invalid project id.
+      throw invalidTableSpec(tableSpec);
     }
 
     TableReference ref = new TableReference();
-    ref.setProjectId(match.group("PROJECT"));
-
-    return ref.setDatasetId(match.group("DATASET")).setTableId(match.group("TABLE"));
+    ref.setProjectId(project);
+    return ref.setDatasetId(dataset).setTableId(table);
+  }
+
+  private static IllegalArgumentException invalidTableSpec(String tableSpec) {
+    return new IllegalArgumentException(
+        String.format(
+            "Table specification [%s] is not in one of the expected formats ("
+                + " [project_id]:[dataset_id].[table_id],"
+                + " [project_id].[dataset_id].[table_id],"
+                + " [dataset_id].[table_id],"
+                + " [project_id]:[catalog_id].[namespace_id].[table_id],"
+                + " [project_id].[catalog_id].[namespace_id].[table_id])",
+            tableSpec));
   }
 
   @SuppressWarnings({

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java

@@ -589,6 +589,13 @@ public class BigQueryIO {
    */
   private static final String PROJECT_ID_REGEXP = "[a-z][-a-z0-9:.]{0,61}[a-z0-9]";
 
+  /**
+   * Matches a whole string against {@link #PROJECT_ID_REGEXP}. Used by {@link
+   * BigQueryHelpers#parseTableSpec} to decide whether the leading segment of a dotted table
+   * specification is a project id.
+   */
+  static final Pattern PROJECT_ID_PATTERN = Pattern.compile(PROJECT_ID_REGEXP);
+
   /** Regular expression that matches Dataset IDs. */
   private static final String DATASET_REGEXP = "[-\\w.]{1,1024}";
 
@@ -604,6 +611,12 @@ public class BigQueryIO {
   /**
    * Matches table specifications in the form {@code "[project_id]:[dataset_id].[table_id]"}, {@code
    * "[project_id].[dataset_id].[table_id]"}, or {@code "[dataset_id].[table_id]"}.
+   *
+   * <p>This pattern is used for syntactic validation only; the assignment of the matched string's
+   * segments to the project/dataset/table fields is done by {@link BigQueryHelpers#parseTableSpec},
+   * which additionally understands 4-part Lakehouse runtime catalog (BigLake metastore) references
+   * {@code "[project_id].[catalog_id].[namespace_id].[table_id]"}, mapping them to a composite
+   * {@code "[catalog_id].[namespace_id]"} dataset id.
    */
   private static final String DATASET_TABLE_REGEXP =
       String.format(

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageSourceBase.java

@@ -150,12 +150,15 @@ public List<BigQueryStorageStreamSource<T>> split(
     // an appropriate number of streams for the Session to produce reasonable throughput.
     // This is required when using the Read API Source V2.
     int streamCount = 0;
-    if (!bqOptions.getEnableStorageReadApiV2()) {
+    // numBytes is null for tables that don't report storage statistics, e.g. Lakehouse
+    // runtime catalog (BigLake metastore) tables. Let BQ read API decide number of streams
+    // by specifying streamCount == 0.
+    boolean isNumBytesMissing = targetTable != null && targetTable.getNumBytes() == null;
+    if (!bqOptions.getEnableStorageReadApiV2() && !isNumBytesMissing) {
       if (desiredBundleSizeBytes > 0) {
         long tableSizeBytes = (targetTable != null) ? targetTable.getNumBytes() : 0;
         streamCount = (int) Math.min(tableSizeBytes / desiredBundleSizeBytes, MAX_SPLIT_COUNT);
       }
-
       streamCount = Math.max(streamCount, MIN_SPLIT_COUNT);
     }
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageTableSource.java

@@ -163,12 +163,13 @@ public void populateDisplayData(DisplayData.Builder builder) {
   @Override
   public long getEstimatedSizeBytes(PipelineOptions options) throws Exception {
     Table table = getTargetTable(options.as(BigQueryOptions.class));
-    if (table != null) {
+    if (table != null && table.getNumBytes() != null) {
       return table.getNumBytes();
     }
     // If the table does not exist, then it will be null.
     // Avoid the NullPointerException here, allow a more meaningful table "not_found"
     // error to be shown to the user, upon table read.
+    // Lakehouse runtime catalog (BigLake metastore) tables exist but report no numBytes.
     return 0;
   }
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryTableSource.java

@@ -82,6 +82,11 @@ public synchronized long getEstimatedSizeBytes(PipelineOptions options) throws E
         }
 
         Long numBytes = table.getNumBytes();
+        if (numBytes == null) {
+          // Tables that don't report storage statistics, e.g. Lakehouse runtime catalog
+          // (BigLake metastore) tables.
+          numBytes = 0L;
+        }
         if (table.getStreamingBuffer() != null
             && table.getStreamingBuffer().getEstimatedBytes() != null) {
           numBytes += table.getStreamingBuffer().getEstimatedBytes().longValue();