Use registered type for Row (#38108)

* Use registered type for Row

* Introduce register_row to register with both coder and schema registry
  Save schema registry id->type mapping

* Allow decorator usage; document register_row is preferred

Author: Abacn

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 1
+  "modification": 16
 }

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 1
+  "modification": 17
 }

.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 1
+  "modification": 3
 }

sdks/python/apache_beam/coders/row_coder_test.py

@@ -60,13 +60,13 @@
      ("one_more_field", typing.Optional[str])])
 
 
+@coders_registry.register_row
 class People(typing.NamedTuple):
   primary: Person
   partner: typing.Optional[Person]
 
 
-coders_registry.register_coder(Person, RowCoder)
-coders_registry.register_coder(People, RowCoder)
+coders_registry.register_row(Person)
 
 
 class RowCoderTest(unittest.TestCase):

sdks/python/apache_beam/coders/typecoders.py

@@ -124,6 +124,19 @@ def _normalize_typehint_type(typehint_type):
   def register_coder(
       self, typehint_type: Any,
       typehint_coder_class: Type[coders.Coder]) -> None:
+    """
+    Register a user type with a coder.
+
+    Typical usage::
+
+      class MyCustomType:
+        pass
+
+      coders.registry.register_coder(MyCustomType, MyCustomCoder)
+
+    To register a supported user type (data class or named tuple) with Beam Row,
+    use :meth:`register_row` instead, as it registers both coder and schema.
+    """
     if not isinstance(typehint_coder_class, type):
       raise TypeError(
           'Coder registration requires a coder class object. '
@@ -133,6 +146,34 @@ def register_coder(
     self._register_coder_internal(
         self._normalize_typehint_type(typehint_type), typehint_coder_class)
 
+  def register_row(self, typehint_type: type[Any]) -> type[Any]:
+    """
+    Register a user type with a Beam Row.
+
+    This registers the type with a RowCoder and register its schema.
+
+    Register a dataclass::
+
+      @coders.registry.register_row
+      @dataclass
+      class MyDataClass:
+        id: int
+        name: str
+
+    Register a named tuple::
+
+      coders.registry.register_row(MyNamedTuple)
+    """
+    from apache_beam.coders import RowCoder
+    from apache_beam.typehints.schemas import typing_to_runner_api
+
+    # Register with row coder
+    self.register_coder(typehint_type, RowCoder)
+    # This call generated a schema id for the type and register it with
+    # schema registry
+    typing_to_runner_api(typehint_type)
+    return typehint_type
+
   def get_coder(self, typehint: Any) -> coders.Coder:
     if typehint and typehint.__module__ == '__main__':
       # See https://github.com/apache/beam/issues/21541

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -256,20 +256,27 @@ def dump_session(file_path):
   # dump supported Beam Registries (currently only logical type registry)
   from apache_beam.coders import typecoders
   from apache_beam.typehints import schemas
+  from apache_beam.typehints.schema_registry import SCHEMA_REGISTRY
 
   with _pickle_lock, open(file_path, 'wb') as file:
     coder_reg = typecoders.registry.get_custom_type_coder_tuples()
     logical_type_reg = schemas.LogicalType._known_logical_types.copy_custom()
+    schema_reg = SCHEMA_REGISTRY.get_registered_typings()
 
     pickler = cloudpickle.CloudPickler(file)
     # TODO(https://github.com/apache/beam/issues/18500) add file system registry
     # once implemented
-    pickler.dump({"coder": coder_reg, "logical_type": logical_type_reg})
+    pickler.dump({
+        "coder": coder_reg,
+        "logical_type": logical_type_reg,
+        "schema": schema_reg
+    })
 
 
 def load_session(file_path):
   from apache_beam.coders import typecoders
   from apache_beam.typehints import schemas
+  from apache_beam.typehints.schema_registry import SCHEMA_REGISTRY
 
   with _pickle_lock, open(file_path, 'rb') as file:
     registries = cloudpickle.load(file)
@@ -284,3 +291,7 @@ def load_session(file_path):
       schemas.LogicalType._known_logical_types.load(registries["logical_type"])
     else:
       _LOGGER.warning('No logical type registry found in saved session')
+    if "schema" in registries:
+      SCHEMA_REGISTRY.load_registered_typings(registries["schema"])
+    else:
+      _LOGGER.warning('No schema registry found in saved session')

sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py

@@ -30,6 +30,7 @@
 
 import apache_beam as beam
 from apache_beam import coders
+from apache_beam.io import jdbc
 from apache_beam.io.jdbc import ReadFromJdbc
 from apache_beam.io.jdbc import WriteToJdbc
 from apache_beam.options.pipeline_options import StandardOptions
@@ -64,7 +65,7 @@
      ("f_timestamp", Timestamp), ("f_decimal", Decimal),
      ("f_date", datetime.date), ("f_time", datetime.time)],
 )
-coders.registry.register_coder(JdbcTestRow, coders.RowCoder)
+coders.registry.register_row(JdbcTestRow)
 
 CustomSchemaRow = typing.NamedTuple(
     "CustomSchemaRow",
@@ -82,11 +83,17 @@
         ("renamed_time", datetime.time),
     ],
 )
-coders.registry.register_coder(CustomSchemaRow, coders.RowCoder)
+
+# Need to put inside enforce_millis_instant_for_timestamp context to align
+# with the same setup in ReadFromJdbc.__init__. Remove once Beam moved to
+# micros instant for timestamp
+# Alternatively, use coders.registry.register_coder(CustomSchemaRow, RowCoder)
+with jdbc.enforce_millis_instant_for_timestamp():
+  coders.registry.register_row(CustomSchemaRow)
 
 SimpleRow = typing.NamedTuple(
     "SimpleRow", [("id", int), ("name", str), ("value", float)])
-coders.registry.register_coder(SimpleRow, coders.RowCoder)
+coders.registry.register_row(SimpleRow)
 
 
 @pytest.mark.uses_gcp_java_expansion_service

sdks/python/apache_beam/io/jdbc.py

@@ -264,6 +264,7 @@ def enforce_millis_instant_for_timestamp():
   LogicalType._known_logical_types = old_registry.copy()
   try:
     LogicalType.register_logical_type(MillisInstant)
+    LogicalType.register_logical_type(JdbcDateType)
     yield
   finally:
     LogicalType._known_logical_types = old_registry

sdks/python/apache_beam/typehints/row_type_test.py

@@ -33,17 +33,18 @@ class RowTypeTest(unittest.TestCase):
   @staticmethod
   def _check_key_type_and_count(x) -> int:
     key_type = type(x[0])
-    if not row_type._user_type_is_generated(key_type):
-      raise RuntimeError("Expect type after GBK to be generated user type")
+    if row_type._user_type_is_generated(key_type):
+      raise RuntimeError("Type after GBK not preserved, get generated type")
+    if not hasattr(key_type, row_type._BEAM_SCHEMA_ID):
+      raise RuntimeError("Type after GBK missing Beam schema ID")
 
     return len(x[1])
 
   def test_group_by_key_namedtuple(self):
     MyNamedTuple = typing.NamedTuple(
         "MyNamedTuple", [("id", int), ("name", str)])
 
-    beam.coders.typecoders.registry.register_coder(
-        MyNamedTuple, beam.coders.RowCoder)
+    beam.coders.typecoders.registry.register_row(MyNamedTuple)
 
     def generate(num: int):
       for i in range(100):
@@ -67,8 +68,7 @@ class MyDataClass:
       id: int
       name: str
 
-    beam.coders.typecoders.registry.register_coder(
-        MyDataClass, beam.coders.RowCoder)
+    beam.coders.typecoders.registry.register_row(MyDataClass)
 
     def generate(num: int):
       for i in range(100):
@@ -120,10 +120,8 @@ class DataClassInt:
     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)
+    beam.coders.typecoders.registry.register_row(DataClassInt)
+    beam.coders.typecoders.registry.register_row(DataClassStr)
 
     def generate(num: int):
       for i in range(10):

sdks/python/apache_beam/typehints/schema_registry.py

@@ -26,7 +26,7 @@
 class SchemaTypeRegistry(object):
   def __init__(self):
     self.by_id = {}
-    self.by_typing = {}
+    self.by_typing = {}  # currently not used
 
   def generate_new_id(self):
     for _ in range(100):
@@ -43,6 +43,15 @@ def add(self, typing, schema):
     if schema.id:
       self.by_id[schema.id] = (typing, schema)
 
+  def load_registered_typings(self, by_id):
+    for id, typing in by_id.items():
+      if id not in self.by_id:
+        self.by_id[id] = (typing, None)
+
+  def get_registered_typings(self):
+    # Used by save_main_session, as pb2.schema isn't picklable
+    return {k: v[0] for k, v in self.by_id.items()}
+
   def get_typing_by_id(self, unique_id):
     if not unique_id:
       return None