[Python] Cache opcode dispatch tables in trivial_inference and atomic types in SchemaTranslation

sdks/python/apache_beam/typehints/schemas.py

@@ -307,6 +307,12 @@ def __init__(self, schema_registry: SchemaTypeRegistry = SCHEMA_REGISTRY):
     self.schema_registry = schema_registry
 
   def typing_to_runner_api(self, type_: type) -> schema_pb2.FieldType:
+    try:
+      if type_ in PRIMITIVE_TO_ATOMIC_TYPE:
+        return schema_pb2.FieldType(atomic_type=PRIMITIVE_TO_ATOMIC_TYPE[type_])
+    except TypeError:
+      pass
+
     if isinstance(type_, schema_pb2.Schema):
       return schema_pb2.FieldType(row_type=schema_pb2.RowType(schema=type_))
 

sdks/python/apache_beam/typehints/trivial_inference.py

@@ -37,6 +37,9 @@
 from apache_beam.typehints import typehints
 from apache_beam.utils import python_callable
 
+_SIMPLE_OPS = None
+_INT_ONE_OPS = None
+
 
 class TypeInferenceError(ValueError):
   """Error to raise when type inference failed."""
@@ -360,9 +363,13 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
       dis.dis(f, show_caches=True)
     else:
       dis.dis(f)
-  from . import opcodes
-  simple_ops = dict((k.upper(), v) for k, v in opcodes.__dict__.items())
-  from . import intrinsic_one_ops
+  global _SIMPLE_OPS, _INT_ONE_OPS
+  if _SIMPLE_OPS is None:
+    from . import intrinsic_one_ops
+    from . import opcodes
+    _SIMPLE_OPS = {k.upper(): v for k, v in opcodes.__dict__.items()}
+    _INT_ONE_OPS = intrinsic_one_ops.INT_ONE_OPS
+  simple_ops = _SIMPLE_OPS
 
   co = f.__code__
   code = co.co_code
@@ -753,7 +760,7 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
       # Introduced in 3.12. The arg is an index into a table of
       # operations reproduced in INT_ONE_OPS. Not all ops are
       # relevant for our type checking infrastructure.
-      int_op = intrinsic_one_ops.INT_ONE_OPS[arg]
+      int_op = _INT_ONE_OPS[arg]
       if debug:
         print("Executing intrinsic one op", int_op.__name__.upper())
       int_op(state, arg)