Fix Flatten schema unification ordering and type degradation

CHANGES.md

@@ -135,7 +135,7 @@
 
 ## Bugfixes
 
-* (Python) Fixed Python and Beam YAML pipelines submitted to Spark with `--spark_submit_uber_jar` failing on the driver unless `--storage_level` was set. The Spark runner now defaults to `MEMORY_ONLY` when no Java runner is available to choose the storage level from ([#41080](https://github.com/apache/beam/issues/41080)).
+* (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)).
 * (YAML) Fixed SQL-backed Join generating invalid SQL syntax (missing ON clause or forward-referenced tables) when joining 3+ inputs ([#41052](https://github.com/apache/beam/issues/41052)).
 * (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)).

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)
 

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
@@ -361,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:
@@ -421,3 +419,67 @@ 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]))
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  unittest.main()