Save schema registry id->type mapping

Author: Abacn

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/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