Use schema registry directly.

Author: shunping

sdks/python/apache_beam/typehints/schema_registry.py

@@ -40,7 +40,7 @@ def generate_new_id(self):
         "schemas.")
 
   def add(self, typing, schema):
-    if not schema.id:
+    if schema.id:
       self.by_id[schema.id] = (typing, schema)
 
   def get_typing_by_id(self, unique_id):

sdks/python/apache_beam/typehints/schemas.py

@@ -135,8 +135,6 @@
 # Serialized schema_pb2.Schema w/o id to id.
 _SCHEMA_ID_CACHE = {}
 
-# Schema id to Schema
-_SCHEMA_CACHE = {}
 
 def named_fields_to_schema(
     names_and_types: Union[Dict[str, type], Sequence[Tuple[str, type]]],
@@ -171,7 +169,6 @@ def named_fields_to_schema(
     if key not in _SCHEMA_ID_CACHE:
       _SCHEMA_ID_CACHE[key] = schema_registry.generate_new_id()
     schema_id = _SCHEMA_ID_CACHE[key]
-    _SCHEMA_CACHE[schema_id] = schema
 
   schema.id = schema_id
   return schema
@@ -581,8 +578,6 @@ def named_tuple_from_schema(self, schema: schema_pb2.Schema) -> type:
         _named_tuple_reduce_method(schema.SerializeToString()))
     setattr(user_type, "_field_descriptions", descriptions)
     setattr(user_type, row_type._BEAM_SCHEMA_ID, schema.id)
-    if schema.id not in _SCHEMA_CACHE:
-      _SCHEMA_CACHE[schema.id] = schema
 
     self.schema_registry.add(user_type, schema)
     coders.registry.register_coder(user_type, coders.RowCoder)
@@ -622,11 +617,12 @@ def schema_from_element_type(element_type: type) -> schema_pb2.Schema:
     return named_fields_to_schema(element_type._fields)
   elif match_is_named_tuple(element_type):
     if hasattr(element_type, row_type._BEAM_SCHEMA_ID):
-      # if the named tuple's schema is in cache, we just use it instead of
+      # if the named tuple's schema is in registry, we just use it instead of
       # regenerating one.
       schema_id = getattr(element_type, row_type._BEAM_SCHEMA_ID)
-      if schema_id in _SCHEMA_CACHE:
-        return _SCHEMA_CACHE[schema_id]
+      schema = SCHEMA_REGISTRY.get_schema_by_id(schema_id)
+      if schema is not None:
+        return schema
     return named_tuple_to_schema(element_type)
   else:
     raise TypeError(