From 722e613f82d06d100fe0e1fb3aa6f378835bc774 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Fri, 4 Sep 2026 19:48:59 +0000 Subject: [PATCH 1/3] fix bigtable schema issue --- .../BigtableReadSchemaTransformProvider.java | 2 +- .../apache_beam/yaml/tests/bigtable.yaml | 61 +++++++++++-------- sdks/python/apache_beam/yaml/yaml_provider.py | 20 ++++-- .../yaml/yaml_provider_unit_test.py | 44 +++++++++++++ sdks/python/apache_beam/yaml/yaml_testing.py | 10 +-- 5 files changed, 98 insertions(+), 39 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadSchemaTransformProvider.java index ca4caee2e469..292ede316d51 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadSchemaTransformProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadSchemaTransformProvider.java @@ -149,7 +149,7 @@ public abstract static class Builder { public abstract Builder setProjectId(String projectId); - public abstract Builder setFlatten(Boolean flatten); + public abstract Builder setFlatten(@Nullable Boolean flatten); /** Builds a {@link BigtableReadSchemaTransformConfiguration} instance. */ public abstract BigtableReadSchemaTransformConfiguration build(); diff --git a/sdks/python/apache_beam/yaml/tests/bigtable.yaml b/sdks/python/apache_beam/yaml/tests/bigtable.yaml index 2f97b83c6e92..6850bd45a098 100644 --- a/sdks/python/apache_beam/yaml/tests/bigtable.yaml +++ b/sdks/python/apache_beam/yaml/tests/bigtable.yaml @@ -91,6 +91,7 @@ pipelines: project: 'apache-beam-testing' instance: "{BT_INSTANCE}" table: 'test-table' + flatten: True - type: MapToFields config: language: python @@ -145,33 +146,41 @@ pipelines: fields: key: callable: | - def convert_to_bytes(row): - return row.key.decode("utf-8") if "key" in row._fields else None + def convert_to_string(row): + k = getattr(row, 'key', None) + return k.decode("utf-8") if hasattr(k, 'decode') else k column_families: - column_families -# TODO: issue #35790, once fixed we can uncomment this assert -# - type: AssertEqual -# config: -# elements: -# - {key: 'row1', -# # Use explicit map syntax to match the actual output -# column_families: { -# cf1: { -# cq1: [ -# { value: "value1", timestamp_micros: 5000 } -# ], -# cq2: [ -# { value: "value2", timestamp_micros: 1000 } -# ] -# } -# } -# } - # - {'key': 'row1', - # column_families: {cf1: {cq2: - # [BeamSchema_3281a0ae_fe85_474b_9030_86fbed58833a(value=b'value2', timestamp_micros=1000)], 'cq1': [BeamSchema_3281a0ae_fe85_474b_9030_86fbed58833a(value=b'value1', timestamp_micros=5000)]}}} - - -# - type: LogForTesting + callable: | + def convert_cells_to_string(row): + cf = getattr(row, 'column_families', None) + if not cf: + return None + return { + fam: { + col: [ + beam.Row(value=c.value.decode("utf-8") if hasattr(c.value, 'decode') else c.value, + timestamp_micros=c.timestamp_micros) + for c in cells + ] + for col, cells in cols.items() + } + for fam, cols in cf.items() + } + - type: AssertEqual + config: + elements: + - key: 'row1' + # Use explicit map syntax to match the actual output + column_families: { + cf1: { + cq1: [ + { value: "value1", timestamp_micros: 5000 } + ], + cq2: [ + { value: "value2", timestamp_micros: 1000 } + ] + } + } diff --git a/sdks/python/apache_beam/yaml/yaml_provider.py b/sdks/python/apache_beam/yaml/yaml_provider.py index 6dd7a502eba6..8af23d365897 100755 --- a/sdks/python/apache_beam/yaml/yaml_provider.py +++ b/sdks/python/apache_beam/yaml/yaml_provider.py @@ -759,6 +759,21 @@ def dicts_to_rows(o): return o +def to_dict(value): + if value is None: + return None + if hasattr(value, '_asdict'): + return {k: to_dict(v) for k, v in value._asdict().items() if v is not None} + elif hasattr(value, 'as_dict'): + return {k: to_dict(v) for k, v in value.as_dict().items() if v is not None} + elif isinstance(value, (list, tuple)): + return [to_dict(v) for v in value] + elif isinstance(value, Mapping): + return {k: to_dict(v) for k, v in value.items() if v is not None} + else: + return value + + def _unify_element_with_schema(element, target_schema): """Convert an element to match the target schema, preserving existing fields only.""" @@ -828,11 +843,6 @@ def __init__(self, elements: Iterable[Any]): self._elements = elements def expand(self, pcoll): - def to_dict(row): - # filter None when comparing - temp_dict = {k: v for k, v in row._asdict().items() if v is not None} - return dict(temp_dict.items()) - return assert_that( pcoll | beam.Map(to_dict), equal_to([to_dict(e) for e in dicts_to_rows(self._elements)])) diff --git a/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py b/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py index e1e3ee847d96..37676d1a6930 100644 --- a/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py +++ b/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py @@ -377,3 +377,47 @@ def test_create_mixed_types(self): [('a', None), ('element', 1)], [('a', 2), ('element', None)], ])) + + +class YamlProvidersAssertEqualTest(unittest.TestCase): + def test_assert_equal_nested_mapping(self): + # Issue #35790: elements with nested dictionaries / MapFields + with beam.Pipeline() as p: + input_data = [ + beam.Row( + key='row1', + column_families={ + 'cf1': { + 'cq1': [beam.Row(value='value1', timestamp_micros=5000)], + 'cq2': [beam.Row(value='value2', timestamp_micros=1000)] + } + }) + ] + pcoll = p | beam.Create(input_data) + _ = pcoll | YamlProviders.AssertEqual( + elements=[{ + 'key': 'row1', + 'column_families': { + 'cf1': { + 'cq1': [{ + 'value': 'value1', 'timestamp_micros': 5000 + }], + 'cq2': [{ + 'value': 'value2', 'timestamp_micros': 1000 + }] + } + } + }]) + + def test_assert_equal_nested_rows(self): + with beam.Pipeline() as p: + input_data = [beam.Row(key='row1', nested=beam.Row(sub=beam.Row(val=42)))] + pcoll = p | beam.Create(input_data) + _ = pcoll | YamlProviders.AssertEqual( + elements=[{ + 'key': 'row1', 'nested': { + 'sub': { + 'val': 42 + } + } + }]) diff --git a/sdks/python/apache_beam/yaml/yaml_testing.py b/sdks/python/apache_beam/yaml/yaml_testing.py index c7f5f5f4e93c..364cceeacf3e 100644 --- a/sdks/python/apache_beam/yaml/yaml_testing.py +++ b/sdks/python/apache_beam/yaml/yaml_testing.py @@ -370,7 +370,7 @@ def __init__(self, elements, recording_id): def expand(self, pcoll): # Convert elements to rows outside the matcher to avoid capturing # any grpc channels that might be created during the conversion - expected_rows = yaml_provider.dicts_to_rows(self._elements) + expected_rows = [yaml_provider.to_dict(e) for e in self._elements] recording_id = self._recording_id # Create a serializable matcher function that doesn't capture @@ -392,8 +392,7 @@ def __call__(self, actual): raise matcher = SerializableMatcher(expected_rows, recording_id) - return assert_that( - pcoll | beam.Map(lambda row: beam.Row(**row._asdict())), matcher) + return assert_that(pcoll | beam.Map(yaml_provider.to_dict), matcher) def create_test( @@ -548,10 +547,7 @@ def _composite_key_to_nested( def _try_row_as_dict(row): - try: - return row._asdict() - except AttributeError: - return row + return yaml_provider.to_dict(row) # Linter: No need for unittest.main here. From c6846b3c9266a7015e71e9737a0d9c4fc7e51c89 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Fri, 25 Sep 2026 15:26:31 +0000 Subject: [PATCH 2/3] add comments about pipelines and add docstring --- sdks/python/apache_beam/yaml/tests/bigtable.yaml | 11 +++++++++++ sdks/python/apache_beam/yaml/yaml_provider.py | 2 ++ 2 files changed, 13 insertions(+) diff --git a/sdks/python/apache_beam/yaml/tests/bigtable.yaml b/sdks/python/apache_beam/yaml/tests/bigtable.yaml index 6850bd45a098..5b0355ac1a16 100644 --- a/sdks/python/apache_beam/yaml/tests/bigtable.yaml +++ b/sdks/python/apache_beam/yaml/tests/bigtable.yaml @@ -30,6 +30,9 @@ fixtures: # Tests for BigTable YAML IO pipelines: + # Pipeline 1: Write test data (SetCell mutations) to Bigtable, converting + # YAML string fields (key, column_qualifier, value) to UTF-8 bytes expected + # by WriteToBigTable. - pipeline: type: chain transforms: @@ -83,6 +86,10 @@ pipelines: project: 'apache-beam-testing' instance: "{BT_INSTANCE}" table: 'test-table' + + # Pipeline 2: Read from Bigtable with flatten=True (one output row per column + # qualifier), decode byte fields back to UTF-8 strings, and verify the + # flattened rows with AssertEqual. - pipeline: type: chain transforms: @@ -131,6 +138,10 @@ pipelines: timestamp_micros: 1000 } ] } - type: LogForTesting + # Pipeline 3: Read from Bigtable with flatten=False (one output row per + # Bigtable row key with nested column_families map), decode the key and nested + # cell values from bytes to UTF-8 strings, and verify the nested structure + # with AssertEqual (Issue #35790). - pipeline: type: chain transforms: diff --git a/sdks/python/apache_beam/yaml/yaml_provider.py b/sdks/python/apache_beam/yaml/yaml_provider.py index 8af23d365897..093cba52e29f 100755 --- a/sdks/python/apache_beam/yaml/yaml_provider.py +++ b/sdks/python/apache_beam/yaml/yaml_provider.py @@ -760,6 +760,8 @@ def dicts_to_rows(o): def to_dict(value): + """Recursively converts Row, NamedTuple, or Mapping objects to dicts, omitting + fields with None values.""" if value is None: return None if hasattr(value, '_asdict'): From 281531d71d37b76e2d2728825abdb8e9990e019a Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Fri, 25 Sep 2026 17:54:16 +0000 Subject: [PATCH 3/3] add random uuid to minimize chance of table collisions --- .../BigtableSimpleWriteSchemaTransformProviderIT.java | 5 ++++- .../org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java | 5 ++++- .../gcp/bigtable/BigtableWriteSchemaTransformProviderIT.java | 5 ++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java index de6a4d54f374..bdca791d6d2b 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java @@ -36,6 +36,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.UUID; import java.util.stream.Collectors; import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; import org.apache.beam.sdk.io.gcp.bigtable.BigtableWriteSchemaTransformProvider.BigtableWriteSchemaTransformConfiguration; @@ -63,7 +64,9 @@ public class BigtableSimpleWriteSchemaTransformProviderIT { private BigtableTableAdminClient tableAdminClient; private BigtableDataClient dataClient; private String tableId = - String.format("BigtableWriteIT-%tF-% writeTransform; diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java index 32b747f01a72..0c64b0c36823 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java @@ -39,6 +39,7 @@ import java.util.Iterator; import java.util.List; import java.util.Objects; +import java.util.UUID; import java.util.stream.Collectors; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; @@ -79,7 +80,9 @@ public class BigtableWriteIT implements Serializable { private static BigtableDataClient client; private static BigtableTableAdminClient tableAdminClient; private final String tableId = - String.format("BigtableWriteIT-%tF-% writeTransform;