Fix #37862: fixed named tuple and effectively fails dataclass inside union typehint

Author: Abacn

sdks/python/apache_beam/typehints/row_type_test.py

@@ -86,6 +86,62 @@ def generate(num: int):
           | 'Count Elements' >> beam.Map(self._check_key_type_and_count))
       assert_that(result, equal_to([10] * 100))
 
+  def test_group_by_key_namedtuple_union(self):
+    Tuple1 = typing.NamedTuple("Tuple1", [("id", int)])
+
+    Tuple2 = typing.NamedTuple("Tuple2", [("id", int), ("name", str)])
+
+    def generate(num: int):
+      for i in range(2):
+        yield (Tuple1(i), num)
+        yield (Tuple2(i, 'a'), num)
+
+    pipeline = TestPipeline(is_integration_test=False)
+
+    with pipeline as p:
+      result = (
+          p
+          | 'Create' >> beam.Create([i for i in range(2)])
+          | 'Generate' >> beam.ParDo(generate).with_output_types(
+              tuple[(Tuple1 | Tuple2), int])
+          | 'GBK' >> beam.GroupByKey()
+          | 'Count' >> beam.Map(lambda x: len(x[1])))
+      assert_that(result, equal_to([2] * 4))
+
+  # Union of dataclasses as type hint currently result in FastPrimitiveCoder
+  # fails at GBK
+  @unittest.skip("https://github.com/apache/beam/issues/22085")
+  def test_group_by_key_inherited_dataclass(self):
+    @dataclass
+    class DataClassInt:
+      id: int
+
+    @dataclass
+    class DataClassStr(DataClassInt):
+      name: str
+
+    beam.coders.typecoders.registry.register_coder(
+        DataClassInt, beam.coders.RowCoder)
+    beam.coders.typecoders.registry.register_coder(
+        DataClassStr, beam.coders.RowCoder)
+
+    def generate(num: int):
+      for i in range(10):
+        yield (DataClassInt(i), num)
+        yield (DataClassStr(i, 'a'), num)
+
+    pipeline = TestPipeline(is_integration_test=False)
+
+    with pipeline as p:
+      result = (
+          p
+          | 'Create' >> beam.Create([i for i in range(2)])
+          | 'Generate' >> beam.ParDo(generate).with_output_types(
+              tuple[(DataClassInt | DataClassStr), int])
+          | 'GBK' >> beam.GroupByKey()
+          | 'Count Elements' >> beam.Map(self._check_key_type_and_count))
+      assert_that(result, equal_to([2] * 4))
+
   def test_derived_dataclass_schema_id(self):
     @dataclass
     class BaseDataClass:

sdks/python/apache_beam/typehints/schemas.py

@@ -663,8 +663,15 @@ def union_schema_type(element_types):
   element_types must be a set of schema-aware types whose fields have the
   same naming and ordering.
   """
+  named_fields_and_types = []
+  for t in element_types:
+    n = named_fields_from_element_type(t)
+    if named_fields_and_types and len(named_fields_and_types[-1]) != len(n):
+      raise TypeError("element types has different number of fields")
+    named_fields_and_types.append(n)
+
   union_fields_and_types = []
-  for field in zip(*[named_fields_from_element_type(t) for t in element_types]):
+  for field in zip(*named_fields_and_types):
     names, types = zip(*field)
     name_set = set(names)
     if len(name_set) != 1: