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.