From c75730f7f70a8ee652a8aab39e5a27830b92e86d Mon Sep 17 00:00:00 2001 From: vishal Date: Thu, 12 Nov 2026 17:02:25 +0530 Subject: [PATCH 1/2] [YAML] Fix Flatten schema unification ordering and type degradation (#41238) --- CHANGES.md | 1 + sdks/python/apache_beam/yaml/yaml_provider.py | 40 +++++++++--- .../yaml/yaml_provider_unit_test.py | 62 +++++++++++++++++++ 3 files changed, 96 insertions(+), 7 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 19acae3a9313..81c7bb439a98 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -135,6 +135,7 @@ ## Bugfixes +* (Python/YAML) Fixed YAML Flatten schema unification causing non-deterministic field ordering, type degradation from Optional[T] to Optional[Any], and unnecessary per-element overhead on identical schemas ([#41238](https://github.com/apache/beam/issues/41238)). * (Python) Fixed cross-language `ReadFromKafka(redistribute=True)` creating one key per record when `redistribute_num_keys` is not set. It now uses the KafkaIO default of 32768 keys, as Java `withRedistribute()` does. On Flink the per-record keys left timer state that grew until the job ran out of memory ([#41103](https://github.com/apache/beam/issues/41103)). * (Python/YAML) Beam YAML `generic` filters now accept `>=`, `<=`, `==` and `!=`; previously only `<` and `>` were recognized ([#41043](https://github.com/apache/beam/issues/41043)). * (Python) Fixed batched DoFns whose input must be rebatched into a `pyarrow.Table` with more than one column, which failed with `ArrowInvalid: Column 1 ... expected length N but got length 0`. This affected, for example, Beam YAML `ReadFromCsv` or `ReadFromJson` followed directly by `WriteToParquet` ([#41042](https://github.com/apache/beam/issues/41042)). diff --git a/sdks/python/apache_beam/yaml/yaml_provider.py b/sdks/python/apache_beam/yaml/yaml_provider.py index 6409c295f182..a4525d7af7e8 100755 --- a/sdks/python/apache_beam/yaml/yaml_provider.py +++ b/sdks/python/apache_beam/yaml/yaml_provider.py @@ -1030,7 +1030,11 @@ def _merge_schemas(self, pcolls): input PCollections. Fields are made optional to handle missing values. If fields have different types, they are unified to Optional[Any]. """ + from apache_beam.typehints.native_type_compatibility import _match_is_optional + from apache_beam.typehints.native_type_compatibility import extract_optional_type from apache_beam.typehints.schemas import named_fields_from_element_type + from apache_beam.typehints.schemas import named_fields_to_schema + from apache_beam.typehints.schemas import named_tuple_from_schema # Collect all schemas schemas = [] @@ -1046,20 +1050,42 @@ def _merge_schemas(self, pcolls): if not schemas: return None - # Merge all field names and types. - all_field_names = set().union(*(s.keys() for s in schemas)) + # If all PCollections have identical schemas, no unification is needed. + if len(schemas) == len(pcolls) and all( + list(s.items()) == list(schemas[0].items()) for s in schemas): + return None + + def unwrap_type(t): + if _match_is_optional(t): + return extract_optional_type(t), True + return t, False + + # Maintain deterministic field ordering across all input schemas. + all_field_names = list( + dict.fromkeys(name for s in schemas for name in s.keys())) unified_fields = {} + has_unschematized_pcoll = len(schemas) < len(pcolls) + for name in all_field_names: - present_types = {s[name] for s in schemas if name in s} - if len(present_types) > 1: + is_optional = has_unschematized_pcoll or ( + len([s for s in schemas if name in s]) < len(schemas)) + unwrapped_types = set() + for s in schemas: + if name in s: + inner_type, was_optional = unwrap_type(s[name]) + unwrapped_types.add(inner_type) + if was_optional: + is_optional = True + + if len(unwrapped_types) > 1: unified_fields[name] = Optional[Any] else: - unified_fields[name] = Optional[present_types.pop()] + concrete_type = unwrapped_types.pop() + unified_fields[name] = ( + Optional[concrete_type] if is_optional else concrete_type) # Create unified schema if unified_fields: - from apache_beam.typehints.schemas import named_fields_to_schema - from apache_beam.typehints.schemas import named_tuple_from_schema unified_schema = named_fields_to_schema(list(unified_fields.items())) return named_tuple_from_schema(unified_schema) 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 37676d1a6930..ec26ea28bc68 100644 --- a/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py +++ b/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py @@ -20,6 +20,9 @@ import sys import tempfile import unittest +from typing import Any +from typing import NamedTuple +from typing import Optional import mock import yaml @@ -421,3 +424,62 @@ def test_assert_equal_nested_rows(self): } } }]) + + +class FlattenSchemaUnificationTest(unittest.TestCase): + def test_identical_schemas_return_none(self): + RowType = NamedTuple('RowType', [('a', int), ('b', str)]) + pcoll1 = mock.MagicMock(element_type=RowType) + pcoll2 = mock.MagicMock(element_type=RowType) + flatten = YamlProviders.Flatten() + self.assertIsNone(flatten._merge_schemas((pcoll1, pcoll2))) + + def test_deterministic_field_ordering(self): + RowType1 = NamedTuple('RowType1', [('alpha', int), ('beta', str)]) + RowType2 = NamedTuple('RowType2', [('gamma', float), ('alpha', int)]) + pcoll1 = mock.MagicMock(element_type=RowType1) + pcoll2 = mock.MagicMock(element_type=RowType2) + flatten = YamlProviders.Flatten() + unified = flatten._merge_schemas((pcoll1, pcoll2)) + self.assertIsNotNone(unified) + self.assertEqual(list(unified._fields), ['alpha', 'beta', 'gamma']) + + def test_unwrap_optional_preserves_concrete_type(self): + RowType1 = NamedTuple('RowType1', [('id', int)]) + RowType2 = NamedTuple('RowType2', [('id', Optional[int])]) + pcoll1 = mock.MagicMock(element_type=RowType1) + pcoll2 = mock.MagicMock(element_type=RowType2) + flatten = YamlProviders.Flatten() + unified = flatten._merge_schemas((pcoll1, pcoll2)) + self.assertIsNotNone(unified) + from apache_beam.typehints.schemas import named_fields_from_element_type + fields_dict = dict(named_fields_from_element_type(unified)) + import numpy as np + self.assertIn(fields_dict['id'], (Optional[int], Optional[np.int64])) + + def test_required_field_preservation(self): + RowType1 = NamedTuple('RowType1', [('id', int), ('a', str)]) + RowType2 = NamedTuple('RowType2', [('id', int), ('b', float)]) + pcoll1 = mock.MagicMock(element_type=RowType1) + pcoll2 = mock.MagicMock(element_type=RowType2) + flatten = YamlProviders.Flatten() + unified = flatten._merge_schemas((pcoll1, pcoll2)) + self.assertIsNotNone(unified) + from apache_beam.typehints.schemas import named_fields_from_element_type + fields_dict = dict(named_fields_from_element_type(unified)) + import numpy as np + self.assertIn(fields_dict['id'], (int, np.int64)) + self.assertEqual(fields_dict['a'], Optional[str]) + self.assertIn(fields_dict['b'], (Optional[float], Optional[np.float64])) + + def test_conflicting_types_fallback_to_optional_any(self): + RowType1 = NamedTuple('RowType1', [('id', int)]) + RowType2 = NamedTuple('RowType2', [('id', str)]) + pcoll1 = mock.MagicMock(element_type=RowType1) + pcoll2 = mock.MagicMock(element_type=RowType2) + flatten = YamlProviders.Flatten() + unified = flatten._merge_schemas((pcoll1, pcoll2)) + self.assertIsNotNone(unified) + from apache_beam.typehints.schemas import named_fields_from_element_type + fields_dict = dict(named_fields_from_element_type(unified)) + self.assertIn(fields_dict['id'], (Any, Optional[Any])) From 82fcde38d4ab4ab74c0bc5a7cf15be2aa2fcbe7e Mon Sep 17 00:00:00 2001 From: vishal Date: Thu, 12 Nov 2026 23:24:15 +0530 Subject: [PATCH 2/2] Move unittest.main() to end of file in yaml_provider_unit_test.py --- .../python/apache_beam/yaml/yaml_provider_unit_test.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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 ec26ea28bc68..7601394a1b6e 100644 --- a/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py +++ b/sdks/python/apache_beam/yaml/yaml_provider_unit_test.py @@ -364,11 +364,6 @@ def test_empty_base(self): yaml_provider._join_url_or_filepath(None, 'a/b.yaml'), 'a/b.yaml') -if __name__ == '__main__': - logging.getLogger().setLevel(logging.INFO) - unittest.main() - - class YamlProvidersCreateTest(unittest.TestCase): def test_create_mixed_types(self): with beam.Pipeline() as p: @@ -483,3 +478,8 @@ def test_conflicting_types_fallback_to_optional_any(self): from apache_beam.typehints.schemas import named_fields_from_element_type fields_dict = dict(named_fields_from_element_type(unified)) self.assertIn(fields_dict['id'], (Any, Optional[Any])) + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + unittest.main()