only check expr

Signed-off-by: Jeffrey Kinard <[email protected]>
apache · May 28, 2024 · b8182e3 · b8182e3
1 parent cdee460
commit b8182e3
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/sdks/python/apache_beam/yaml/yaml_mapping.py b/sdks/python/apache_beam/yaml/yaml_mapping.py
@@ -281,15 +281,10 @@ def _as_callable_for_pcoll(
         list(input_schema.keys()), fn_spec, msg, language, input_schema)
 
 
-def _as_callable(
-    original_fields, expr, transform_name, language, input_schema=None):
+def _as_callable(original_fields, expr, transform_name, language, input_schema):
 
   # Extract original type from upstream pcoll when doing simple mappings
-  if input_schema is None:
-    input_schema = {}
-  original_type = input_schema.get(transform_name, None)
-  if expr in original_fields:
-    original_type = input_schema.get(expr, None)
+  original_type = input_schema.get(str(expr), None)
 
   # TODO(yaml): support an imports parameter
   # TODO(yaml): support a requirements parameter (possibly at a higher level)