Allow non-frozen dataclass register with other coders as a backup for backward compatibility; add tests

Author: Abacn

CHANGES.md

@@ -69,6 +69,7 @@
 ## New Features / Improvements
 
 * Added support for large pipeline options via a file (Python) ([#37370](https://github.com/apache/beam/issues/37370)).
+* Supported infer schema from dataclass (Python) ([#22085](https://github.com/apache/beam/issues/37370)). Default coder for typehint-ed (or set with_output_type) for non-frozen dataclasses changed to RowCoder. To preserve the old behavior (fast primitive coder), explicitly register the type with FastPrimitiveCoder.
 
 ## Breaking Changes
 

sdks/python/apache_beam/typehints/native_type_compatibility.py

@@ -179,27 +179,38 @@ def match_is_named_tuple(user_type):
 def match_dataclass_for_row(user_type):
   """Match whether the type is a dataclass handled by row coder.
 
-  for frozen dataclasses, only true when explicitly registered with row coder:
+  For frozen dataclasses, only true when explicitly registered with row coder:
 
     beam.coders.typecoders.registry.register_coder(
         MyDataClass, beam.coders.RowCoder)
+
+  (for backward-compatibility reason).
+
+  For non-frozen dataclasses, default to true otherwise explicitly registered
+  with a coder other than the row coder.
   """
+
   if not dataclasses.is_dataclass(user_type):
     return False
 
-  if not user_type.__dataclass_params__.frozen:
-    return True
-
+  is_frozen = user_type.__dataclass_params__.frozen
   # avoid circular import
-  # pylint: disable=wrong-import-position
-  from apache_beam.coders.typecoders import registry as coders_registry
-  from apache_beam.coders import RowCoder
+  try:
+    # pylint: disable=wrong-import-position
+    from apache_beam.coders.typecoders import registry as coders_registry
+    from apache_beam.coders import RowCoder
+  except AttributeError:
+    # coder registery not yet initialized so it must be absent
+    return not is_frozen
 
-  # check _coders (not get_coder) to get the registered coder directly without
-  # fallback
-  return (
-      user_type in coders_registry._coders and
-      coders_registry._coders[user_type] == RowCoder)
+  if is_frozen:
+    return (
+        user_type in coders_registry._coders and
+        coders_registry._coders[user_type] == RowCoder)
+  else:
+    return (
+        user_type not in coders_registry._coders or
+        coders_registry._coders[user_type] == RowCoder)
 
 
 def _match_is_optional(user_type):

sdks/python/apache_beam/typehints/native_type_compatibility_test.py

@@ -20,6 +20,7 @@
 # pytype: skip-file
 
 import collections.abc
+import dataclasses
 import enum
 import re
 import typing
@@ -33,6 +34,7 @@
 from apache_beam.typehints.native_type_compatibility import convert_to_python_types
 from apache_beam.typehints.native_type_compatibility import convert_typing_to_builtin
 from apache_beam.typehints.native_type_compatibility import is_any
+from apache_beam.typehints.native_type_compatibility import match_dataclass_for_row
 
 _TestNamedTuple = typing.NamedTuple(
     '_TestNamedTuple', [('age', int), ('name', bytes)])
@@ -509,6 +511,39 @@ def test_type_alias_type_unwrapped(self):
     self.assertEqual(
         typehints.Tuple[int, ...], convert_to_beam_type(AliasTuple))
 
+  def test_dataclass_default(self):
+    @dataclasses.dataclass(frozen=True)
+    class FrozenDC:
+      foo: int
+
+    @dataclasses.dataclass
+    class NonFrozenDC:
+      foo: int
+
+    self.assertFalse(match_dataclass_for_row(FrozenDC))
+    self.assertTrue(match_dataclass_for_row(NonFrozenDC))
+
+  def test_dataclass_registered(self):
+    @dataclasses.dataclass(frozen=True)
+    class FrozenRegisteredDC:
+      foo: int
+
+    @dataclasses.dataclass
+    class NonFrozenRegisteredDC:
+      foo: int
+
+    # pylint: disable=wrong-import-position
+    from apache_beam.coders import typecoders
+    from apache_beam.coders.coders import FastPrimitivesCoder
+    from apache_beam.coders import RowCoder
+
+    typecoders.registry.register_coder(FrozenRegisteredDC, RowCoder)
+    typecoders.registry.register_coder(
+        NonFrozenRegisteredDC, FastPrimitivesCoder)
+
+    self.assertTrue(match_dataclass_for_row(FrozenRegisteredDC))
+    self.assertFalse(match_dataclass_for_row(NonFrozenRegisteredDC))
+
 
 if __name__ == '__main__':
   unittest.main()