Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
61 changes: 35 additions & 26 deletions sdks/python/apache_beam/yaml/tests/bigtable.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ pipelines:
project: 'apache-beam-testing'
instance: "{BT_INSTANCE}"
table: 'test-table'
flatten: True
- type: MapToFields
config:
language: python
Expand Down Expand Up @@ -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 }
]
}
}


20 changes: 15 additions & 5 deletions sdks/python/apache_beam/yaml/yaml_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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)]))
Expand Down
44 changes: 44 additions & 0 deletions sdks/python/apache_beam/yaml/yaml_provider_unit_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
}])
10 changes: 3 additions & 7 deletions sdks/python/apache_beam/yaml/yaml_testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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.
Loading