feat(bigquery): add BigLake configuration support for Storage Write API

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json

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

.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json

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

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -3203,6 +3203,16 @@ class BeamModulePlugin implements Plugin<Project> {
             testJavaHome = project.findProperty("java${testJavaVersion}Home")
           }
 
+          // Detect macOS and append '-macos' to tox environment to avoid pip check issues
+          def actualToxEnv = tox_env
+          def osName = System.getProperty("os.name").toLowerCase()
+          if (osName.contains("mac")) {
+            // Only append -macos for standard python environments (py39, py310, etc.)
+            if (tox_env.matches("py\\d+")) {
+              actualToxEnv = "${tox_env}-macos"
+            }
+          }
+
           if (project.hasProperty('useWheelDistribution')) {
             def pythonVersionNumber  = project.ext.pythonVersion.replace('.', '')
             dependsOn ":sdks:python:bdistPy${pythonVersionNumber}linux"
@@ -3218,7 +3228,7 @@ class BeamModulePlugin implements Plugin<Project> {
                   environment "JAVA_HOME", testJavaHome
                 }
                 executable 'sh'
-                args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $tox_env ${packageFilename} '$posargs' "
+                args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $actualToxEnv ${packageFilename} '$posargs' "
               }
             }
           } else {
@@ -3231,12 +3241,12 @@ class BeamModulePlugin implements Plugin<Project> {
                   environment "JAVA_HOME", testJavaHome
                 }
                 executable 'sh'
-                args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $tox_env '$posargs'"
+                args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $actualToxEnv '$posargs'"
               }
             }
           }
           inputs.files project.pythonSdkDeps
-          outputs.files project.fileTree(dir: "${pythonRootDir}/target/.tox/${tox_env}/log/")
+          outputs.files project.fileTree(dir: "${pythonRootDir}/target/.tox/${actualToxEnv}/log/")
         }
       }
       // Run single or a set of integration tests with provided test options and pipeline options.

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

@@ -315,6 +315,9 @@ BigQueryIO.Write<Row> createStorageWriteApiTransform(Schema schema) {
       if (!Strings.isNullOrEmpty(configuration.getKmsKey())) {
         write = write.withKmsKey(configuration.getKmsKey());
       }
+      if (configuration.getBigLakeConfiguration() != null) {
+        write = write.withBigLakeConfiguration(configuration.getBigLakeConfiguration());
+      }
       if (this.testBigQueryServices != null) {
         write = write.withTestServices(testBigQueryServices);
       }

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

@@ -197,6 +197,14 @@ public static Builder builder() {
   @SchemaFieldDescription("A list of columns to cluster the BigQuery table by.")
   public abstract @Nullable List<String> getClusteringFields();
 
+  @SchemaFieldDescription(
+      "Configuration for creating BigLake tables. The following options are available:"
+          + "\n - connectionId (REQUIRED): the name of your cloud resource connection,"
+          + "\n - storageUri (REQUIRED): the path to your GCS folder where data will be written to,"
+          + "\n - fileFormat (OPTIONAL): defaults to 'parquet',"
+          + "\n - tableFormat (OPTIONAL): defaults to 'iceberg'.")
+  public abstract java.util.@Nullable Map<String, String> getBigLakeConfiguration();
+
   /** Builder for {@link BigQueryWriteConfiguration}. */
   @AutoValue.Builder
   public abstract static class Builder {
@@ -231,6 +239,9 @@ public abstract static class Builder {
 
     public abstract Builder setClusteringFields(List<String> clusteringFields);
 
+    public abstract Builder setBigLakeConfiguration(
+        java.util.Map<String, String> bigLakeConfiguration);
+
     /** Builds a {@link BigQueryWriteConfiguration} instance. */
     public abstract BigQueryWriteConfiguration build();
   }

sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py

@@ -43,6 +43,11 @@
 # Protect against environments where bigquery library is not available.
 # pylint: disable=wrong-import-order, wrong-import-position
 
+try:
+  from apache_beam.io.gcp.gcsio import GcsIO
+except ImportError:
+  GcsIO = None
+
 try:
   from apitools.base.py.exceptions import HttpError
 except ImportError:
@@ -145,6 +150,62 @@ def parse_expected_data(self, expected_elements):
 
     return data
 
+  def assert_iceberg_tables_created(
+      self, table_prefix, storage_uri, expected_count=1):
+    """Verify that Iceberg table directories are created in
+    the warehouse location.
+
+    Args:
+      table_prefix: The table name prefix to look for
+      storage_uri: The GCS storage URI (e.g., 'gs://bucket/path')
+      expected_count: Expected number of table directories
+    """
+    if GcsIO is None:
+      _LOGGER.warning(
+          "GcsIO not available, skipping warehouse location verification")
+      return
+
+    gcs_io = GcsIO()
+
+    # Parse the storage URI to get bucket and prefix
+    if not storage_uri.startswith('gs://'):
+      raise ValueError(f'Storage URI must start with gs://, got: {storage_uri}')
+
+    # Remove 'gs://' prefix and split bucket from path
+    path_parts = storage_uri[5:].split('/', 1)
+    bucket_name = path_parts[0]
+    base_prefix = path_parts[1] if len(path_parts) > 1 else ''
+
+    # Construct the full prefix to search for table directories
+    # Following the pattern:
+    # {base_prefix}/{project}/{dataset}/{table_prefix}
+    search_prefix = (
+        f"{base_prefix}/"
+        f"{self.project}/{self.dataset_id}/{table_prefix}")
+
+    # List objects in the bucket with the constructed prefix
+    try:
+      objects = gcs_io.list_prefix(f"gs://{bucket_name}/{search_prefix}")
+      object_count = len(list(objects))
+
+      if object_count < expected_count:
+        raise AssertionError(
+            f"Expected at least {expected_count} objects in warehouse "
+            f"location gs://{bucket_name}/{search_prefix}, but found "
+            f"{object_count}")
+
+      _LOGGER.info(
+          "Successfully verified %s objects created in "
+          "warehouse location gs://%s/%s",
+          object_count,
+          bucket_name,
+          search_prefix)
+
+    except Exception as e:
+      raise AssertionError(
+          f"Failed to verify table creation in warehouse location "
+          f"gs://{bucket_name}/{search_prefix}: {str(e)}")
+
   def run_storage_write_test(
       self, table_name, items, schema, use_at_least_once=False):
     table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table_name)
@@ -511,6 +572,41 @@ def test_streaming_with_at_least_once(self):
     table = 'streaming_with_at_least_once'
     self.run_streaming(table_name=table, use_at_least_once=True)
 
+  def test_write_with_big_lake_configuration(self):
+    """Test BigQuery Storage Write API with BigLake configuration."""
+    table = 'write_with_big_lake_config'
+    table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table)
+
+    # BigLake configuration with required parameters (matching Java test)
+    big_lake_config = {
+        'connectionId': 'projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete',  # pylint: disable=line-too-long
+        'storageUri': 'gs://apache-beam-testing-bq-biglake/BigQueryXlangStorageWriteIT',  # pylint: disable=line-too-long
+        'fileFormat': 'parquet',
+        'tableFormat': 'iceberg'
+    }
+
+    bq_matcher = BigqueryFullResultMatcher(
+        project=self.project,
+        query="SELECT * FROM {}.{}".format(self.dataset_id, table),
+        data=self.parse_expected_data(self.ELEMENTS))
+
+    with beam.Pipeline(argv=self.args) as p:
+      _ = (
+          p
+          | "Create test data" >> beam.Create(self.ELEMENTS)
+          | beam.io.WriteToBigQuery(
+              table=table_id,
+              method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API,
+              schema=self.ALL_TYPES_SCHEMA,
+              create_disposition='CREATE_IF_NEEDED',
+              write_disposition='WRITE_TRUNCATE',
+              big_lake_configuration=big_lake_config))
+
+    hamcrest_assert(p, bq_matcher)
+
+    # Verify that the table directory was created in the warehouse location
+    self.assert_iceberg_tables_created(table, big_lake_config['storageUri'])
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)

sdks/python/apache_beam/io/gcp/bigquery.py

@@ -1995,7 +1995,8 @@ def __init__(
       num_streaming_keys=DEFAULT_SHARDS_PER_DESTINATION,
       use_cdc_writes: bool = False,
       primary_key: List[str] = None,
-      expansion_service=None):
+      expansion_service=None,
+      big_lake_configuration=None):
     """Initialize a WriteToBigQuery transform.
 
     Args:
@@ -2216,6 +2217,7 @@ def __init__(
     self._num_streaming_keys = num_streaming_keys
     self._use_cdc_writes = use_cdc_writes
     self._primary_key = primary_key
+    self._big_lake_configuration = big_lake_configuration
 
   # Dict/schema methods were moved to bigquery_tools, but keep references
   # here for backward compatibility.
@@ -2378,6 +2380,7 @@ def find_in_nested_dict(schema):
           num_storage_api_streams=self._num_storage_api_streams,
           use_cdc_writes=self._use_cdc_writes,
           primary_key=self._primary_key,
+          big_lake_configuration=self._big_lake_configuration,
           expansion_service=self.expansion_service)
     else:
       raise ValueError(f"Unsupported method {method_to_use}")
@@ -2626,6 +2629,7 @@ def __init__(
       num_storage_api_streams=0,
       use_cdc_writes: bool = False,
       primary_key: List[str] = None,
+      big_lake_configuration=None,
       expansion_service=None):
     self._table = table
     self._table_side_inputs = table_side_inputs
@@ -2639,6 +2643,7 @@ def __init__(
     self._num_storage_api_streams = num_storage_api_streams
     self._use_cdc_writes = use_cdc_writes
     self._primary_key = primary_key
+    self._big_lake_configuration = big_lake_configuration
     self._expansion_service = expansion_service or BeamJarExpansionService(
         'sdks:java:io:google-cloud-platform:expansion-service:build')
 
@@ -2733,6 +2738,7 @@ def expand(self, input):
             use_cdc_writes=self._use_cdc_writes,
             primary_key=self._primary_key,
             clustering_fields=clustering_fields,
+            big_lake_configuration=self._big_lake_configuration,
             error_handling={
                 'output': StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS
             }))

sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py

@@ -0,0 +1,117 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Unit tests for BigQuery BigLake configuration."""
+
+import unittest
+from unittest import mock
+
+from apache_beam.io.gcp import bigquery
+
+
+@mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService')
+class BigQueryBigLakeTest(unittest.TestCase):
+  """Test BigLake configuration support in BigQuery Storage Write API."""
+  def test_storage_write_to_bigquery_with_biglake_config(
+      self, mock_expansion_service):
+    """Test that StorageWriteToBigQuery accepts bigLakeConfiguration."""
+    big_lake_config = {
+        'connectionId': (
+            'projects/test-project/locations/us/connections/test-connection'),
+        'storageUri': 'gs://test-bucket/test-path',
+        'fileFormat': 'parquet',
+        'tableFormat': 'iceberg'
+    }
+
+    # Test that the constructor accepts the bigLakeConfiguration parameter
+    transform = bigquery.StorageWriteToBigQuery(
+        table='test-project:test_dataset.test_table',
+        big_lake_configuration=big_lake_config)
+
+    # Verify the configuration is stored
+    self.assertEqual(transform._big_lake_configuration, big_lake_config)
+
+  def test_storage_write_to_bigquery_without_biglake_config(
+      self, mock_expansion_service):
+    """Test that StorageWriteToBigQuery works without bigLakeConfiguration."""
+    transform = bigquery.StorageWriteToBigQuery(
+        table='test-project:test_dataset.test_table')
+
+    # Verify the configuration is None by default
+    self.assertIsNone(transform._big_lake_configuration)
+
+  def test_biglake_config_passed_to_external_transform(
+      self, mock_expansion_service):
+    """Test that StorageWriteToBigQuery accepts bigLakeConfiguration."""
+    big_lake_config = {
+        'connection_id': 'projects/my-project/locations/us/connections/my-conn',
+        'table_format': 'ICEBERG'
+    }
+
+    # Mock the expansion service to avoid JAR dependency
+    mock_expansion_service.return_value = mock.MagicMock()
+
+    # Create the transform
+    transform = bigquery.StorageWriteToBigQuery(
+        table='my-project:my_dataset.my_table',
+        big_lake_configuration=big_lake_config)
+
+    # Verify the big_lake_configuration is stored correctly
+    self.assertEqual(transform._big_lake_configuration, big_lake_config)
+
+    # Verify that the transform has the expected identifier
+    self.assertEqual(
+        transform.IDENTIFIER,
+        "beam:schematransform:org.apache.beam:bigquery_storage_write:v2")
+
+    # Verify that the expansion service was created (mocked)
+    mock_expansion_service.assert_called_once_with(
+        'sdks:java:io:google-cloud-platform:expansion-service:build')
+
+  def test_biglake_config_validation(self, mock_expansion_service):
+    """Test validation of bigLakeConfiguration parameters."""
+    # Test with minimal required configuration
+    minimal_config = {
+        'connectionId': (
+            'projects/test-project/locations/us/connections/test-connection'),
+        'storageUri': 'gs://test-bucket/test-path'
+    }
+
+    transform = bigquery.StorageWriteToBigQuery(
+        table='test-project:test_dataset.test_table',
+        big_lake_configuration=minimal_config)
+
+    self.assertEqual(transform._big_lake_configuration, minimal_config)
+
+    # Test with full configuration
+    full_config = {
+        'connectionId': (
+            'projects/test-project/locations/us/connections/test-connection'),
+        'storageUri': 'gs://test-bucket/test-path',
+        'fileFormat': 'parquet',
+        'tableFormat': 'iceberg'
+    }
+
+    transform = bigquery.StorageWriteToBigQuery(
+        table='test-project:test_dataset.test_table',
+        big_lake_configuration=full_config)
+
+    self.assertEqual(transform._big_lake_configuration, full_config)
+
+
+if __name__ == '__main__':
+  unittest.main()