[Python] Merge runtime type-check wrappers

CHANGES.md

@@ -106,6 +106,7 @@
 
 ## New Features / Improvements
 
+* (Python) Reduced runtime type-check overhead by combining input and output checks into a single wrapper ([#20226](https://github.com/apache/beam/issues/20226)).
 * (Java) Flink runner: sources now forward Beam metrics to Flink at most once per second instead of after every record, which was a significant share of the CPU of reading. Metrics in the pipeline result are not affected ([#40607](https://github.com/apache/beam/issues/40607)).
 * (Go) Added `wait.On`, which delays each input window until the corresponding windows in its signal PCollections have closed ([#39909](https://github.com/apache/beam/issues/39909)).
 * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)).

sdks/python/apache_beam/typehints/typecheck.py

@@ -84,25 +84,68 @@ def teardown(self):
     return self.dofn.teardown()
 
 
-class OutputCheckWrapperDoFn(AbstractDoFnWrapper):
-  """A DoFn that verifies against common errors in the output type."""
-  def __init__(self, dofn, full_label):
+class RuntimeTypeCheckWrapperDoFn(AbstractDoFnWrapper):
+  """A wrapper around a DoFn which performs type-checking of input and output."""
+  def __init__(self, dofn, type_hints, full_label):
     super().__init__(dofn)
     self.full_label = full_label
 
+    process_fn = self.dofn._process_argspec_fn()
+    if hasattr(process_fn, '__func__'):
+      self._process_fn = process_fn.__func__
+      self._process_fn_owner = process_fn.__self__
+    else:
+      self._process_fn = process_fn
+      self._process_fn_owner = None
+
+    if type_hints.input_types:
+      input_args, input_kwargs = type_hints.input_types
+      self._input_hints = getcallargs_forhints(
+          process_fn, *input_args, **input_kwargs)
+    else:
+      self._input_hints = None
+
+    # TODO(robertwb): Multi-output.
+    self._output_type_hint = type_hints.simple_output_type(full_label)
+
   def wrapper(self, method, args, kwargs):
     try:
-      result = method(*args, **kwargs)
+      result = self._type_check_result(method(*args, **kwargs))
     except TypeCheckError as e:
-      # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name
-      error_msg = (
-          'Runtime type violation detected within ParDo(%s): '
-          '%s' % (self.full_label, e))
-      _, _, tb = sys.exc_info()
-      raise TypeCheckError(error_msg).with_traceback(tb)
+      raise self._add_label(e)
     else:
       return self._check_type(result)
 
+  def process(self, *args, **kwargs):
+    try:
+      if self._input_hints:
+        if self._process_fn_owner is not None:
+          actual_inputs = inspect.getcallargs(  # pylint: disable=deprecated-method
+              self._process_fn, self._process_fn_owner, *args, **kwargs)
+        else:
+          actual_inputs = inspect.getcallargs(  # pylint: disable=deprecated-method
+              self._process_fn, *args, **kwargs)
+
+        for var, hint in self._input_hints.items():
+          if hint is actual_inputs[var]:
+            # self parameter
+            continue
+          _check_instance_type(hint, actual_inputs[var], var, True)
+
+      result = self._type_check_result(self.dofn.process(*args, **kwargs))
+    except TypeCheckError as e:
+      raise self._add_label(e)
+    else:
+      return self._check_type(result)
+
+  def _add_label(self, e):
+    # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name.
+    error_msg = (
+        'Runtime type violation detected within ParDo(%s): '
+        '%s' % (self.full_label, e))
+    _, _, tb = sys.exc_info()
+    return TypeCheckError(error_msg).with_traceback(tb)
+
   @staticmethod
   def _check_type(output):
     if output is None:
@@ -120,36 +163,6 @@ def _check_type(output):
           'iterable. %s was returned instead.' % type(output))
     return output
 
-
-class TypeCheckWrapperDoFn(AbstractDoFnWrapper):
-  """A wrapper around a DoFn which performs type-checking of input and output.
-  """
-  def __init__(self, dofn, type_hints, label=None):
-    super().__init__(dofn)
-    self._process_fn = self.dofn._process_argspec_fn()
-    if type_hints.input_types:
-      input_args, input_kwargs = type_hints.input_types
-      self._input_hints = getcallargs_forhints(
-          self._process_fn, *input_args, **input_kwargs)
-    else:
-      self._input_hints = None
-    # TODO(robertwb): Multi-output.
-    self._output_type_hint = type_hints.simple_output_type(label)
-
-  def wrapper(self, method, args, kwargs):
-    result = method(*args, **kwargs)
-    return self._type_check_result(result)
-
-  def process(self, *args, **kwargs):
-    if self._input_hints:
-      actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs)  # pylint: disable=deprecated-method
-      for var, hint in self._input_hints.items():
-        if hint is actual_inputs[var]:
-          # self parameter
-          continue
-        _check_instance_type(hint, actual_inputs[var], var, True)
-    return self._type_check_result(self.dofn.process(*args, **kwargs))
-
   def _type_check_result(self, transform_results):
     if self._output_type_hint is None or transform_results is None:
       return transform_results
@@ -289,11 +302,9 @@ def visit_transform(self, applied_transform):
         if isinstance(transform.fn, core.CombineValuesDoFn):
           transform.fn.combinefn = self._wrapped_fn
       else:
-        transform.fn = transform.dofn = OutputCheckWrapperDoFn(
-            TypeCheckWrapperDoFn(
-                transform.fn,
-                transform.get_type_hints(),
-                applied_transform.full_label),
+        transform.fn = transform.dofn = RuntimeTypeCheckWrapperDoFn(
+            transform.fn,
+            transform.get_type_hints(),
             applied_transform.full_label)
 
 

sdks/python/apache_beam/typehints/typecheck_test.py

@@ -34,10 +34,15 @@
 from apache_beam import Pipeline
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import TypeOptions
+from apache_beam.pipeline import PipelineVisitor
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
+from apache_beam.transforms import combiners
+from apache_beam.transforms import userstate
+from apache_beam.transforms.core import CallableWrapperDoFn
 from apache_beam.typehints import decorators
+from apache_beam.typehints import typecheck
 from apache_beam.typehints import with_input_types
 from apache_beam.typehints import with_output_types
 
@@ -300,5 +305,128 @@ def process(self, element, *args, **kwargs):
       self.p.run().wait_until_finish()
 
 
+@with_input_types(int)
+@with_output_types(int)
+class _AddOneDoFn(beam.DoFn):
+  def process(self, element):
+    yield element + 1
+
+
+class RuntimeTypeCheckWrapperDoFnTest(unittest.TestCase):
+  def test_visitor_applies_single_wrapper_layer(self):
+    p = beam.Pipeline(options=PipelineOptions(runtime_type_check=True))
+    _ = (p | beam.Create([1, 2]) | 'TypedStep' >> beam.ParDo(_AddOneDoFn()))
+
+    p.visit(typecheck.TypeCheckVisitor())
+
+    wrapped_dofns = []
+
+    class _Collector(PipelineVisitor):
+      def visit_transform(self, applied_transform):
+        transform = applied_transform.transform
+        if isinstance(transform, beam.ParDo) and isinstance(
+            getattr(transform, 'fn', None), typecheck.AbstractDoFnWrapper):
+          wrapped_dofns.append(transform.fn)
+
+    p.visit(_Collector())
+
+    self.assertTrue(wrapped_dofns)
+    for wrapper in wrapped_dofns:
+      self.assertIsInstance(wrapper, typecheck.RuntimeTypeCheckWrapperDoFn)
+      self.assertNotIsInstance(wrapper.dofn, typecheck.AbstractDoFnWrapper)
+
+  def test_wrapper_labels_type_check_errors(self):
+    dofn = _AddOneDoFn()
+    wrapper = typecheck.RuntimeTypeCheckWrapperDoFn(
+        dofn, dofn.get_type_hints(), 'MyStep')
+
+    with self.assertRaisesRegex(
+        typecheck.TypeCheckError,
+        r'Runtime type violation detected within ParDo\(MyStep\)'):
+      wrapper.process('not-an-int')
+
+  def test_wrapper_labels_output_type_check_error_once(self):
+    @with_output_types(int)
+    class BadOutputDoFn(beam.DoFn):
+      def process(self, element):
+        return ['not-an-int']
+
+    dofn = BadOutputDoFn()
+    wrapper = typecheck.RuntimeTypeCheckWrapperDoFn(
+        dofn, dofn.get_type_hints(), 'MyStep')
+
+    with self.assertRaises(typecheck.TypeCheckError) as context:
+      wrapper.process(1)
+
+    prefix = 'Runtime type violation detected within ParDo(MyStep):'
+    error_message = str(context.exception)
+
+    self.assertEqual(error_message.count(prefix), 1)
+    self.assertIn('not-an-int', error_message)
+
+  def test_wrapper_preserves_results(self):
+    dofn = _AddOneDoFn()
+    wrapper = typecheck.RuntimeTypeCheckWrapperDoFn(
+        dofn, dofn.get_type_hints(), 'MyStep')
+
+    self.assertEqual(list(wrapper.process(1)), [2])
+
+  def test_wrapper_preserves_bound_method_owner_with_varargs(self):
+    class MethodOwner:
+      pass
+
+    class BoundMethod(MethodOwner):
+      @with_input_types(MethodOwner, int)
+      @with_output_types(Iterable[int])
+      def process(*args):
+        return [args[-1]]
+
+    owner = BoundMethod()
+    dofn = CallableWrapperDoFn(owner.process)
+
+    wrapper = typecheck.RuntimeTypeCheckWrapperDoFn(
+        dofn, dofn.get_type_hints(), 'MyStep')
+
+    self.assertIs(wrapper._process_fn_owner, owner)
+    self.assertEqual(list(wrapper.process(1)), [1])
+
+
+def _make_stateful_dofn():
+  count_state = userstate.CombiningValueStateSpec(
+      'count', combiners.CountCombineFn())
+
+  class _CountingStatefulDoFn(beam.DoFn):
+    def process(self, element, count=beam.DoFn.StateParam(count_state)):
+      count.add(1)
+      yield element[1]
+
+  return _CountingStatefulDoFn()
+
+
+class RuntimeTypeCheckStatefulDoFnTest(unittest.TestCase):
+  def test_wrapper_does_not_cache_bound_process(self):
+    dofn = _make_stateful_dofn()
+    wrapper = typecheck.RuntimeTypeCheckWrapperDoFn(
+        dofn, dofn.get_type_hints(), 'Step')
+
+    for attr, value in wrapper.__dict__.items():
+      self.assertFalse(
+          hasattr(value, '__self__'), 'wrapper caches bound method %r' % attr)
+
+    userstate.validate_stateful_dofn(wrapper)
+
+  def test_stateful_dofn_with_runtime_type_check(self):
+    options = PipelineOptions()
+    options.view_as(TypeOptions).runtime_type_check = True
+
+    with TestPipeline(options=options) as p:
+      result = (
+          p
+          | beam.Create([('k', 1), ('k', 2), ('k', 3)])
+          | beam.ParDo(_make_stateful_dofn()))
+
+      assert_that(result, equal_to([1, 2, 3]))
+
+
 if __name__ == '__main__':
   unittest.main()