From 877b911af6065572e9a712b2c56985241a3ee8cc Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Thu, 8 Oct 2026 20:25:08 +0000 Subject: [PATCH 1/3] Add gcs_avro_to_cloud_bigtable yaml blueprint example and e2e test --- .../yaml/examples/testing/examples_test.py | 43 +++++- .../blueprint/gcs_avro_to_cloud_bigtable.yaml | 44 ++++++ .../e2e/gcs_avro_to_cloud_bigtable.yaml | 130 ++++++++++++++++++ 3 files changed, 215 insertions(+), 2 deletions(-) create mode 100644 sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml create mode 100644 sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml diff --git a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py b/sdks/python/apache_beam/yaml/examples/testing/examples_test.py index daf72a5f8d59..cfa6929e9bdd 100644 --- a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py +++ b/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_gcs_avro_to_cloud_bigtable_yaml' ]) def _io_write_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): @@ -755,6 +756,43 @@ def _file_io_read_test_preprocessor( return test_spec +@YamlExamplesTestSuite.register_test_preprocessor( + ['test_gcs_avro_to_cloud_bigtable_yaml']) +def _avro_io_read_test_preprocessor( + test_spec: dict, expected: list[str], env: TestEnvironment): + """ + Preprocessor for tests that involve reading from Avro. + + This preprocessor replaces any ReadFromAvro 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 + Avro files. + + 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 ReadFromAvro transforms replaced. + """ + if pipeline := test_spec.get('pipeline', None): + for transform in pipeline.get('transforms', []): + if transform.get('type', '') == 'ReadFromAvro': + config = transform['config'] + file_name = config['path'].split('/')[-1] + transform['type'] = 'Create' + transform['config'] = { + k: v + for k, v in config.items() if k.startswith('__') + } + transform['config']['elements'] = INPUT_FILES[file_name] + + return test_spec + + @YamlExamplesTestSuite.register_test_preprocessor( ['test_iceberg_read_yaml', 'test_iceberg_to_alloydb_yaml']) def _iceberg_io_read_test_preprocessor( @@ -1351,7 +1389,8 @@ def _jinja_preprocessor(raw_spec_string: str, test_name: str): 'products.csv': input_data.products_csv(), 'kinglear.txt': input_data.text_data(), 'youtube-comments.csv': input_data.youtube_comments_csv(), - 'system-logs.csv': input_data.system_logs_csv() + 'system-logs.csv': input_data.system_logs_csv(), + 'shipments.avro': input_data.shipments_data() } KAFKA_TOPICS = {'test-topic': input_data.kafka_messages_data()} diff --git a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml new file mode 100644 index 000000000000..1c58d2869688 --- /dev/null +++ b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml @@ -0,0 +1,44 @@ +# 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 GCS Avro files +# and writes to Cloud Bigtable. This matches the Dataflow Template located +# here - https://cloud.google.com/dataflow/docs/guides/templates/provided/avro-to-cloud-bigtable + +pipeline: + type: chain + transforms: + # Step 1: Reading data from GCS Avro files + - type: ReadFromAvro + name: ReadFromGCSAvro + config: + path: "gs://apache-beam-testing/yaml/shipments.avro" + # Step 2: Write records out to Cloud Bigtable + - type: WriteToBigTable + name: WriteToBigTable + config: + project: "apache-beam-testing" + instance: "shipment-test" + table: "shipments" + +# Expected: +# Row(shipment_id='S1', customer_id='C1', shipment_date='2023-05-01', shipment_cost=150.0, customer_name='Alice', customer_email='alice@example.com') +# Row(shipment_id='S2', customer_id='C2', shipment_date='2023-06-12', shipment_cost=300.0, customer_name='Bob', customer_email='bob@example.com') +# Row(shipment_id='S3', customer_id='C1', shipment_date='2023-05-10', shipment_cost=20.0, customer_name='Alice', customer_email='alice@example.com') +# Row(shipment_id='S4', customer_id='C4', shipment_date='2024-07-01', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') +# Row(shipment_id='S5', customer_id='C5', shipment_date='2023-05-09', shipment_cost=300.0, customer_name='Erin', customer_email='erin@example.com') +# Row(shipment_id='S6', customer_id='C4', shipment_date='2024-07-02', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') diff --git a/sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml b/sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml new file mode 100644 index 000000000000..e492e697af3a --- /dev/null +++ b/sdks/python/apache_beam/yaml/extended_tests/e2e/gcs_avro_to_cloud_bigtable.yaml @@ -0,0 +1,130 @@ +# +# 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: 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 GCS Avro + - 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: WriteToAvro + config: + path: "{TEMP_DIR}/out.avro" + + # Pipeline 2: Read from GCS Avro and write to Cloud Bigtable + - pipeline: + type: chain + transforms: + - type: ReadFromAvro + config: + path: "{TEMP_DIR}/out.avro*" + - 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" + + # 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}]} From 9bc2827b4c4c3a6d40b5765ff5a7386b29e39d10 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Fri, 9 Oct 2026 12:36:24 +0000 Subject: [PATCH 2/3] update blueprint to follow similar pattern as e2e --- .../blueprint/gcs_avro_to_cloud_bigtable.yaml | 35 +++++++++++++++---- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml index 1c58d2869688..4268aeffbe41 100644 --- a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml +++ b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml @@ -27,7 +27,28 @@ pipeline: name: ReadFromGCSAvro config: path: "gs://apache-beam-testing/yaml/shipments.avro" - # Step 2: Write records out to Cloud Bigtable + # Step 2: Convert Avro 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: @@ -36,9 +57,9 @@ pipeline: table: "shipments" # Expected: -# Row(shipment_id='S1', customer_id='C1', shipment_date='2023-05-01', shipment_cost=150.0, customer_name='Alice', customer_email='alice@example.com') -# Row(shipment_id='S2', customer_id='C2', shipment_date='2023-06-12', shipment_cost=300.0, customer_name='Bob', customer_email='bob@example.com') -# Row(shipment_id='S3', customer_id='C1', shipment_date='2023-05-10', shipment_cost=20.0, customer_name='Alice', customer_email='alice@example.com') -# Row(shipment_id='S4', customer_id='C4', shipment_date='2024-07-01', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') -# Row(shipment_id='S5', customer_id='C5', shipment_date='2023-05-09', shipment_cost=300.0, customer_name='Erin', customer_email='erin@example.com') -# Row(shipment_id='S6', customer_id='C4', shipment_date='2024-07-02', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') +# 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) From 1c75f23d5a484c2fdb3e2c124a515209d5212517 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Tue, 13 Oct 2026 00:19:27 +0000 Subject: [PATCH 3/3] remove template link --- .../transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml index 4268aeffbe41..2ce2e5aaae5c 100644 --- a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml +++ b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/gcs_avro_to_cloud_bigtable.yaml @@ -16,8 +16,7 @@ # limitations under the License. # This is an example of a Beam YAML pipeline that reads from GCS Avro files -# and writes to Cloud Bigtable. This matches the Dataflow Template located -# here - https://cloud.google.com/dataflow/docs/guides/templates/provided/avro-to-cloud-bigtable +# and writes to Cloud Bigtable. pipeline: type: chain