From 3467cabda7a3cb87d0853874d1fe5fcf159a6fb2 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Fri, 6 Nov 2026 20:30:39 +0000 Subject: [PATCH] Cache opcode dispatch tables in trivial_inference and atomic types in SchemaTranslation --- sdks/python/apache_beam/typehints/schemas.py | 6 ++++++ .../apache_beam/typehints/trivial_inference.py | 15 +++++++++++---- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/typehints/schemas.py b/sdks/python/apache_beam/typehints/schemas.py index a023d4f4cd66..e07d8c36cb49 100644 --- a/sdks/python/apache_beam/typehints/schemas.py +++ b/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_)) diff --git a/sdks/python/apache_beam/typehints/trivial_inference.py b/sdks/python/apache_beam/typehints/trivial_inference.py index e34dbb98fbbc..e51535de07fc 100644 --- a/sdks/python/apache_beam/typehints/trivial_inference.py +++ b/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)