Add bigquery_to_bigtable yaml blueprint example and e2e test

sdks/python/apache_beam/yaml/examples/testing/examples_test.py

@@ -684,7 +684,8 @@ def _kafka_test_preprocessor(
     'test_wordCountInheritance_yaml',
     'test_iceberg_to_alloydb_yaml',
     'test_bigquery_write_yaml',
-    'test_mongodb_to_bigquery_yaml'
+    'test_mongodb_to_bigquery_yaml',
+    'test_bigquery_to_bigtable_yaml'
 ])
 def _io_write_test_preprocessor(
     test_spec: dict, expected: list[str], env: TestEnvironment):
@@ -828,7 +829,7 @@ def _spanner_io_read_test_preprocessor(
       if transform.get('type', '').startswith('ReadFromSpanner'):
         config = transform['config']
         instance, database = config['instance_id'], config['database_id']
-        if table := config.get('table', None) is None:
+        if (table := config.get('table')) is None:
           table = config.get('query', '').split('FROM')[-1].strip()
         transform['type'] = 'Create'
         transform['config'] = {
@@ -852,6 +853,44 @@ def _spanner_io_read_test_preprocessor(
   return test_spec
 
 
+@YamlExamplesTestSuite.register_test_preprocessor(
+    ['test_bigquery_to_bigtable_yaml'])
+def _bigquery_io_read_test_preprocessor(
+    test_spec: dict, expected: list[str], env: TestEnvironment):
+  """
+  Preprocessor for tests that involve reading from BigQuery.
+
+  This preprocessor replaces any ReadFromBigQuery transform with a Create
+  transform that reads from a predefined in-memory dictionary. This allows
+  the test to verify the pipeline's correctness without relying on external
+  BigQuery tables.
+
+  Args:
+    test_spec: The dictionary representation of the YAML pipeline specification.
+    expected: A list of strings representing the expected output of the
+      pipeline.
+    env: The TestEnvironment object providing utilities for creating temporary
+      files.
+
+  Returns:
+    The modified test_spec dictionary with ReadFromBigQuery transforms replaced.
+  """
+  if pipeline := test_spec.get('pipeline', None):
+    for transform in pipeline.get('transforms', []):
+      if transform.get('type', '') == 'ReadFromBigQuery':
+        config = transform['config']
+        dataset, table = config['table'].split('.')[-2:]
+        transform['type'] = 'Create'
+        transform['config'] = {
+            k: v
+            for k, v in config.items() if k.startswith('__')
+        }
+        transform['config']['elements'] = INPUT_TABLES[(
+            'BigQuery', str(dataset), str(table))]
+
+  return test_spec
+
+
 @YamlExamplesTestSuite.register_test_preprocessor(
     ['test_bigtable_enrichment_yaml', 'test_enrich_spanner_with_bigquery_yaml'])
 def _enrichment_test_preprocessor(
@@ -1369,6 +1408,7 @@ def _jinja_preprocessor(raw_spec_string: str, test_name: str):
     ('BigTable', 'beam-test', 'bigtable-enrichment-test'): input_data.
     bigtable_data(),
     ('BigQuery', 'ALL_TEST', 'customers'): input_data.bigquery_data(),
+    ('BigQuery', 'shipment', 'shipments'): input_data.shipments_data(),
     ('Jdbc', 'shipment', 'shipments'): input_data.shipments_data(),
     ('SqlServer', 'shipment', 'shipments'): input_data.shipments_data(),
     ('Postgres', 'shipment', 'shipments'): input_data.shipments_data(),

sdks/python/apache_beam/yaml/examples/transforms/blueprint/bigquery_to_bigtable.yaml

@@ -0,0 +1,64 @@
+# coding=utf-8
+#
+# 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.
+
+# This is an example of a Beam YAML pipeline that reads from BigQuery
+# and writes to Cloud Bigtable.
+
+pipeline:
+  type: chain
+  transforms:
+    # Step 1: Reading data from BigQuery
+    - type: ReadFromBigQuery
+      name: ReadFromBigQuery
+      config:
+        table: "apache-beam-testing.shipment.shipments"
+    # Step 2: Convert BigQuery rows to Bigtable SetCell mutations
+    - type: MapToFields
+      name: ConvertToBigTableMutations
+      config:
+        language: python
+        fields:
+          key:
+            callable: |
+              def to_key_bytes(row):
+                return bytes(row.shipment_id, 'utf-8')
+          type: "'SetCell'"
+          family_name: "'shipment'"
+          column_qualifier:
+            callable: |
+              def to_qualifier_bytes(row):
+                return b'customer_id'
+          value:
+            callable: |
+              def to_value_bytes(row):
+                return bytes(row.customer_id, 'utf-8')
+          timestamp_micros: -1
+    # Step 3: Write records out to Cloud Bigtable
+    - type: WriteToBigTable
+      name: WriteToBigTable
+      config:
+        project: "apache-beam-testing"
+        instance: "shipment-test"
+        table: "shipments"
+
+# Expected:
+#  Row(key=b'S1', type='SetCell', family_name='shipment', column_qualifier=b'customer_id', value=b'C1', timestamp_micros=-1)
+#  Row(key=b'S2', type='SetCell', family_name='shipment', column_qualifier=b'customer_id', value=b'C2', timestamp_micros=-1)
+#  Row(key=b'S3', type='SetCell', family_name='shipment', column_qualifier=b'customer_id', value=b'C1', timestamp_micros=-1)
+#  Row(key=b'S4', type='SetCell', family_name='shipment', column_qualifier=b'customer_id', value=b'C4', timestamp_micros=-1)
+#  Row(key=b'S5', type='SetCell', family_name='shipment', column_qualifier=b'customer_id', value=b'C5', timestamp_micros=-1)
+#  Row(key=b'S6', type='SetCell', family_name='shipment', column_qualifier=b'customer_id', value=b'C4', timestamp_micros=-1)

sdks/python/apache_beam/yaml/extended_tests/e2e/bigquery_to_bigtable.yaml

@@ -0,0 +1,140 @@
+#
+# 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.
+#
+
+fixtures:
+  - name: BQ_TABLE
+    type: "apache_beam.yaml.integration_tests.temp_bigquery_table"
+    config:
+      project: "apache-beam-testing"
+  - name: BT_INSTANCE
+    type: "apache_beam.yaml.integration_tests.temp_bigtable_table"
+    config:
+      project: "apache-beam-testing"
+  - name: TEMP_DIR
+    type: "apache_beam.yaml.integration_tests.gcs_temp_dir"
+    config:
+      bucket: "gs://temp-storage-for-end-to-end-tests/temp-it"
+
+pipelines:
+  # Pipeline 1: Write test records to BigQuery
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {key: "row1",
+                 type: "SetCell",
+                 family_name: "cf1",
+                 column_qualifier: "cq1",
+                 value: "value1",
+                 timestamp_micros: 5000}
+              - {key: "row1",
+                 type: "SetCell",
+                 family_name: "cf1",
+                 column_qualifier: "cq2",
+                 value: "value2",
+                 timestamp_micros: 1000}
+        - type: WriteToBigQuery
+          config:
+            table: "{BQ_TABLE}"
+    options:
+      project: "apache-beam-testing"
+      temp_location: "{TEMP_DIR}"
+
+  # Pipeline 2: Read from BigQuery and write to Cloud Bigtable
+  - pipeline:
+      type: chain
+      transforms:
+        - type: ReadFromBigQuery
+          config:
+            table: "{BQ_TABLE}"
+        - type: MapToFields
+          name: ConvertStringsToBytes
+          config:
+            language: python
+            fields:
+              key:
+                callable: |
+                  def convert_to_bytes(row):
+                    return bytes(row.key, 'utf-8') if "key" in row._fields else None
+              type:
+                type
+              family_name:
+                family_name
+              column_qualifier:
+                callable: |
+                  def convert_to_bytes(row):
+                    return bytes(row.column_qualifier, 'utf-8') if 'column_qualifier' in row._fields else None
+              value:
+                callable: |
+                  def convert_to_bytes(row):
+                    return bytes(row.value, 'utf-8') if 'value' in row._fields else None
+              timestamp_micros:
+                timestamp_micros
+        - type: WriteToBigTable
+          config:
+            project: "apache-beam-testing"
+            instance: "{BT_INSTANCE}"
+            table: "test-table"
+    options:
+      project: "apache-beam-testing"
+      temp_location: "{TEMP_DIR}"
+
+  # Pipeline 3: Read from Cloud Bigtable and verify records
+  - pipeline:
+      type: chain
+      transforms:
+        - type: ReadFromBigTable
+          config:
+            project: "apache-beam-testing"
+            instance: "{BT_INSTANCE}"
+            table: "test-table"
+            flatten: True
+        - type: MapToFields
+          config:
+            language: python
+            fields:
+              key:
+                callable: |
+                  def convert_to_string(row):
+                    return row.key.decode("utf-8") if "key" in row._fields else None
+              family_name:
+                family_name
+              column_qualifier:
+                callable: |
+                  def convert_to_string(row):
+                    return row.column_qualifier.decode("utf-8") if "column_qualifier" in row._fields else None
+              cells:
+                callable: |
+                  def convert_to_string(row):
+                    cell_bytes = []
+                    for (value, timestamp) in row.cells:
+                      value_bytes = value.decode("utf-8")
+                      cell_bytes.append(beam.Row(value=value_bytes, timestamp_micros=timestamp))
+                    return cell_bytes
+        - type: AssertEqual
+          config:
+            elements:
+              - {key: "row1",
+                 family_name: "cf1",
+                 column_qualifier: "cq1",
+                 cells: [{value: "value1", timestamp_micros: 5000}]}
+              - {key: "row1",
+                 family_name: "cf1",
+                 column_qualifier: "cq2",
+                 cells: [{value: "value2", timestamp_micros: 1000}]}