[BEAM-9602] Add timer family support for python SDK

Yichi Zhang · Yichi Zhang · commit a679babf88e6 · 2020-12-10T17:20:56.000-08:00
diff --git a/sdks/python/apache_beam/runners/common.pxd b/sdks/python/apache_beam/runners/common.pxd
@@ -44,6 +44,7 @@ cdef class MethodWrapper(object):
   cdef object restriction_provider_arg_name
   cdef object watermark_estimator_provider
   cdef object watermark_estimator_provider_arg_name
+  cdef object dynamic_timer_tag_arg_name
   cdef bint unbounded_per_element
 
 
diff --git a/sdks/python/apache_beam/runners/common.py b/sdks/python/apache_beam/runners/common.py
@@ -198,6 +198,7 @@ def __init__(self, obj_to_invoke, method_name):
     self.restriction_provider_arg_name = None
     self.watermark_estimator_provider = None
     self.watermark_estimator_provider_arg_name = None
+    self.dynamic_timer_tag_arg_name = None
 
     if hasattr(self.method_value, 'unbounded_per_element'):
       self.unbounded_per_element = True
@@ -223,14 +224,22 @@ def __init__(self, obj_to_invoke, method_name):
       elif isinstance(v, core.DoFn.WatermarkEstimatorParam):
         self.watermark_estimator_provider = v.watermark_estimator_provider
         self.watermark_estimator_provider_arg_name = kw
+      elif core.DoFn.DynamicTimerTagParam == v:
+        self.dynamic_timer_tag_arg_name = kw
 
     # Create NoOpWatermarkEstimatorProvider if there is no
     # WatermarkEstimatorParam provided.
     if self.watermark_estimator_provider is None:
       self.watermark_estimator_provider = NoOpWatermarkEstimatorProvider()
 
   def invoke_timer_callback(
-      self, user_state_context, key, window, timestamp, pane_info):
+      self,
+      user_state_context,
+      key,
+      window,
+      timestamp,
+      pane_info,
+      dynamic_timer_tag):
     # TODO(ccy): support side inputs.
     kwargs = {}
     if self.has_userstate_arguments:
@@ -246,6 +255,8 @@ def invoke_timer_callback(
       kwargs[self.window_arg_name] = window
     if self.key_arg_name:
       kwargs[self.key_arg_name] = key
+    if self.dynamic_timer_tag_arg_name:
+      kwargs[self.dynamic_timer_tag_arg_name] = dynamic_timer_tag
 
     if kwargs:
       return self.method_value(**kwargs)
@@ -526,12 +537,18 @@ def invoke_teardown(self):
     """
     self.signature.teardown_lifecycle_method.method_value()
 
-  def invoke_user_timer(self, timer_spec, key, window, timestamp, pane_info):
+  def invoke_user_timer(
+      self, timer_spec, key, window, timestamp, pane_info, dynamic_timer_tag):
     # self.output_processor is Optional, but in practice it won't be None here
     self.output_processor.process_outputs(
         WindowedValue(None, timestamp, (window, )),
         self.signature.timer_methods[timer_spec].invoke_timer_callback(
-            self.user_state_context, key, window, timestamp, pane_info))
+            self.user_state_context,
+            key,
+            window,
+            timestamp,
+            pane_info,
+            dynamic_timer_tag))
 
   def invoke_create_watermark_estimator(self, estimator_state):
     return self.signature.create_watermark_estimator_method.method_value(
@@ -1233,10 +1250,11 @@ def current_element_progress(self):
     assert isinstance(self.do_fn_invoker, PerWindowInvoker)
     return self.do_fn_invoker.current_element_progress()
 
-  def process_user_timer(self, timer_spec, key, window, timestamp, pane_info):
+  def process_user_timer(
+      self, timer_spec, key, window, timestamp, pane_info, dynamic_timer_tag):
     try:
       self.do_fn_invoker.invoke_user_timer(
-          timer_spec, key, window, timestamp, pane_info)
+          timer_spec, key, window, timestamp, pane_info, dynamic_timer_tag)
     except BaseException as exn:
       self._reraise_augmented(exn)
 
diff --git a/sdks/python/apache_beam/runners/direct/direct_userstate.py b/sdks/python/apache_beam/runners/direct/direct_userstate.py
@@ -206,8 +206,7 @@ def __init__(self, step_context, dofn, key_coder):
     self.dofn = dofn
     self.key_coder = key_coder
 
-    self.all_state_specs, self.all_timer_specs = (
-        userstate.get_dofn_specs(dofn))
+    self.all_state_specs, self.all_timer_specs = userstate.get_dofn_specs(dofn)
     self.state_tags = {}
     for state_spec in self.all_state_specs:
       state_key = 'user/%s' % state_spec.name
@@ -226,12 +225,14 @@ def __init__(self, step_context, dofn, key_coder):
     self.cached_states = {}
     self.cached_timers = {}
 
-  def get_timer(self, timer_spec, key, window, timestamp, pane):
+  def get_timer(
+      self, timer_spec: userstate.TimerSpec, key, window, timestamp,
+      pane) -> userstate.RuntimeTimer:
     assert timer_spec in self.all_timer_specs
     encoded_key = self.key_coder.encode(key)
     cache_key = (encoded_key, window, timer_spec)
     if cache_key not in self.cached_timers:
-      self.cached_timers[cache_key] = userstate.RuntimeTimer(timer_spec)
+      self.cached_timers[cache_key] = userstate.RuntimeTimer()
     return self.cached_timers[cache_key]
 
   def get_state(self, state_spec, key, window):
@@ -291,16 +292,22 @@ def commit(self):
       encoded_key, window, timer_spec = cache_key
       state = self.step_context.get_keyed_state(encoded_key)
       timer_name = 'user/%s' % timer_spec.name
-      if runtime_timer._cleared:
-        state.clear_timer(window, timer_name, timer_spec.time_domain)
-      if runtime_timer._new_timestamp is not None:
-        # TODO(ccy): add corresponding watermark holds after the DirectRunner
-        # allows for keyed watermark holds.
-        state.set_timer(
-            window,
-            timer_name,
-            timer_spec.time_domain,
-            runtime_timer._new_timestamp)
+      for dynamic_timer_tag, timer in runtime_timer._timer_recordings.items():
+        if timer.cleared:
+          state.clear_timer(
+              window,
+              timer_name,
+              timer_spec.time_domain,
+              dynamic_timer_tag=dynamic_timer_tag)
+        if timer.timestamp:
+          # TODO(ccy): add corresponding watermark holds after the DirectRunner
+          # allows for keyed watermark holds.
+          state.set_timer(
+              window,
+              timer_name,
+              timer_spec.time_domain,
+              timer.timestamp,
+              dynamic_timer_tag=dynamic_timer_tag)
 
   def reset(self):
     for state in self.cached_states.values():
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
@@ -311,7 +311,10 @@ def process_timer_wrapper(self, timer_firing):
     """
     state = self._step_context.get_keyed_state(timer_firing.encoded_key)
     state.clear_timer(
-        timer_firing.window, timer_firing.name, timer_firing.time_domain)
+        timer_firing.window,
+        timer_firing.name,
+        timer_firing.time_domain,
+        dynamic_timer_tag=timer_firing.dynamic_timer_tag)
     self.process_timer(timer_firing)
 
   def process_timer(self, timer_firing):
@@ -875,7 +878,8 @@ def process_timer(self, timer_firing):
         timer_firing.window,
         timer_firing.timestamp,
         # TODO Add paneinfo to timer_firing in DirectRunner
-        None)
+        None,
+        timer_firing.dynamic_timer_tag)
 
   def process_element(self, element):
     self.runner.process(element)
diff --git a/sdks/python/apache_beam/runners/direct/util.py b/sdks/python/apache_beam/runners/direct/util.py
@@ -61,16 +61,28 @@ def __init__(
 
 class TimerFiring(object):
   """A single instance of a fired timer."""
-  def __init__(self, encoded_key, window, name, time_domain, timestamp):
+  def __init__(
+      self,
+      encoded_key,
+      window,
+      name,
+      time_domain,
+      timestamp,
+      dynamic_timer_tag=''):
     self.encoded_key = encoded_key
     self.window = window
     self.name = name
     self.time_domain = time_domain
     self.timestamp = timestamp
+    self.dynamic_timer_tag = dynamic_timer_tag
 
   def __repr__(self):
-    return 'TimerFiring({!r}, {!r}, {}, {})'.format(
-        self.encoded_key, self.name, self.time_domain, self.timestamp)
+    return 'TimerFiring({!r}, {!r}, {}, {}, {})'.format(
+        self.encoded_key,
+        self.name,
+        self.time_domain,
+        self.timestamp,
+        self.dynamic_timer_tag)
 
 
 class KeyedWorkItem(object):
diff --git a/sdks/python/apache_beam/runners/direct/watermark_manager.py b/sdks/python/apache_beam/runners/direct/watermark_manager.py
@@ -325,8 +325,14 @@ def extract_transform_timers(self):
         if had_realtime_timer:
           has_realtime_timer = True
         for expired in timers:
-          window, (name, time_domain, timestamp) = expired
+          window, (name, time_domain, timestamp, dynamic_timer_tag) = expired
           fired_timers.append(
-              TimerFiring(encoded_key, window, name, time_domain, timestamp))
+              TimerFiring(
+                  encoded_key,
+                  window,
+                  name,
+                  time_domain,
+                  timestamp,
+                  dynamic_timer_tag=dynamic_timer_tag))
       self._fired_timers.update(fired_timers)
       return fired_timers, has_realtime_timer
diff --git a/sdks/python/apache_beam/runners/portability/flink_runner_test.py b/sdks/python/apache_beam/runners/portability/flink_runner_test.py
@@ -395,6 +395,9 @@ def test_callbacks_with_exception(self):
   def test_register_finalizations(self):
     raise unittest.SkipTest("BEAM-11021")
 
+  def test_pardo_dynamic_timer(self):
+    raise unittest.SkipTest("BEAM-10120")
+
   # Inherits all other tests.
 
 
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
@@ -493,6 +493,30 @@ def is_buffered_correctly(actual):
 
       assert_that(actual, is_buffered_correctly)
 
+  def test_pardo_dynamic_timer(self):
+    class DynamicTimerDoFn(beam.DoFn):
+      dynamic_timer_spec = userstate.TimerSpec(
+          'dynamic_timer', userstate.TimeDomain.WATERMARK)
+
+      def process(
+          self, element,
+          dynamic_timer=beam.DoFn.TimerParam(dynamic_timer_spec)):
+        dynamic_timer.set(element[1], dynamic_timer_tag=element[0])
+
+      @userstate.on_timer(dynamic_timer_spec)
+      def dynamic_timer_callback(
+          self,
+          tag=beam.DoFn.DynamicTimerTagParam,
+          timestamp=beam.DoFn.TimestampParam):
+        yield (tag, timestamp)
+
+    with self.create_pipeline() as p:
+      actual = (
+          p
+          | beam.Create([('key1', 10), ('key2', 20), ('key3', 30)])
+          | beam.ParDo(DynamicTimerDoFn()))
+      assert_that(actual, equal_to([('key1', 10), ('key2', 20), ('key3', 30)]))
+
   def test_sdf(self):
     class ExpandingStringsDoFn(beam.DoFn):
       def process(
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -88,8 +88,8 @@
   from apache_beam.runners.sdf_utils import SplitResultResidual
   from apache_beam.runners.worker import data_plane
   from apache_beam.runners.worker import sdk_worker
-  from apache_beam.transforms import window
   from apache_beam.transforms.core import Windowing
+  from apache_beam.transforms.window import BoundedWindow
   from apache_beam.utils import windowed_value
 
 # This module is experimental. No backwards-compatibility guarantees.
@@ -381,7 +381,7 @@ def __init__(self,
     self._element_coder = coder.wrapped_value_coder
     self._target_window_coder = coder.window_coder
     # TODO(robertwb): Limit the cache size.
-    self._cache = {}  # type: Dict[window.BoundedWindow, Any]
+    self._cache = {}  # type: Dict[BoundedWindow, Any]
 
   def __getitem__(self, window):
     target_window = self._side_input_data.window_mapping_fn(window)
@@ -644,14 +644,14 @@ def commit(self):
 class OutputTimer(userstate.BaseTimer):
   def __init__(self,
                key,
-               window,  # type: windowed_value.BoundedWindow
+               window,  # type: BoundedWindow
                timestamp,  # type: timestamp.Timestamp
                paneinfo,  # type: windowed_value.PaneInfo
                time_domain, # type: str
                timer_family_id,  # type: str
                timer_coder_impl,  # type: coder_impl.TimerCoderImpl
                output_stream  # type: data_plane.ClosableOutputStream
-              ):
+               ):
     self._key = key
     self._window = window
     self._input_timestamp = timestamp
@@ -661,12 +661,11 @@ def __init__(self,
     self._output_stream = output_stream
     self._timer_coder_impl = timer_coder_impl
 
-  def set(self, ts):
-    # type: (timestamp.TimestampTypes) -> None
+  def set(self, ts: timestamp.TimestampTypes, dynamic_timer_tag='') -> None:
     ts = timestamp.Timestamp.of(ts)
     timer = userstate.Timer(
         user_key=self._key,
-        dynamic_timer_tag='',
+        dynamic_timer_tag=dynamic_timer_tag,
         windows=(self._window, ),
         clear_bit=False,
         fire_timestamp=ts,
@@ -676,11 +675,10 @@ def set(self, ts):
     self._timer_coder_impl.encode_to_stream(timer, self._output_stream, True)
     self._output_stream.maybe_flush()
 
-  def clear(self):
-    # type: () -> None
+  def clear(self, dynamic_timer_tag='') -> None:
     timer = userstate.Timer(
         user_key=self._key,
-        dynamic_timer_tag='',
+        dynamic_timer_tag=dynamic_timer_tag,
         windows=(self._window, ),
         clear_bit=True,
         fire_timestamp=None,
@@ -713,10 +711,8 @@ def __init__(self,
     Args:
       state_handler: A StateServicer object.
       transform_id: The name of the PTransform that this context is associated.
-      key_coder:
-      window_coder:
-      timer_family_specs: A list of ``userstate.TimerSpec`` objects specifying
-        the timers associated with this operation.
+      key_coder: Coder for the key type.
+      window_coder: Coder for the window type.
     """
     self._state_handler = state_handler
     self._transform_id = transform_id
@@ -731,23 +727,20 @@ def add_timer_info(self, timer_family_id, timer_info):
     self._timers_info[timer_family_id] = timer_info
 
   def get_timer(
-      self,
-      timer_spec,
-      key,
-      window,  # type: windowed_value.BoundedWindow
-      timestamp,
-      pane):
-    # type: (...) -> OutputTimer
+      self, timer_spec: userstate.TimerSpec, key, window, timestamp,
+      pane) -> OutputTimer:
     assert self._timers_info[timer_spec.name].output_stream is not None
+    timer_coder_impl = self._timers_info[timer_spec.name].timer_coder_impl
+    output_stream = self._timers_info[timer_spec.name].output_stream
     return OutputTimer(
         key,
         window,
         timestamp,
         pane,
         timer_spec.time_domain,
         timer_spec.name,
-        self._timers_info[timer_spec.name].timer_coder_impl,
-        self._timers_info[timer_spec.name].output_stream)
+        timer_coder_impl,
+        output_stream)
 
   def get_state(self, *args):
     # type: (*Any) -> FnApiUserRuntimeStateTypes
@@ -759,7 +752,7 @@ def get_state(self, *args):
   def _create_state(self,
                     state_spec,  # type: userstate.StateSpec
                     key,
-                    window  # type: windowed_value.BoundedWindow
+                    window  # type: BoundedWindow
                    ):
     # type: (...) -> FnApiUserRuntimeStateTypes
     if isinstance(state_spec,
diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py
@@ -73,6 +73,7 @@
   from apache_beam.runners.sdf_utils import SplitResultResidual
   from apache_beam.runners.worker.bundle_processor import ExecutionContext
   from apache_beam.runners.worker.statesampler import StateSampler
+  from apache_beam.transforms.userstate import TimerSpec
 
 # Allow some "pure mode" declarations.
 try:
@@ -685,7 +686,7 @@ def setup(self):
         self.timer_specs = {
             spec.name: spec
             for spec in userstate.get_dofn_specs(fn)[1]
-        }
+        }  # type: Dict[str, TimerSpec]
 
       if self.side_input_maps is None:
         if tags_and_types:
@@ -741,7 +742,8 @@ def process_timer(self, tag, timer_data):
         timer_data.user_key,
         timer_data.windows[0],
         timer_data.fire_timestamp,
-        timer_data.paneinfo)
+        timer_data.paneinfo,
+        timer_data.dynamic_timer_tag)
 
   def finish(self):
     # type: () -> None
diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py
diff --git a/sdks/python/apache_beam/transforms/trigger.py b/sdks/python/apache_beam/transforms/trigger.py
diff --git a/sdks/python/apache_beam/transforms/trigger_test.py b/sdks/python/apache_beam/transforms/trigger_test.py
diff --git a/sdks/python/apache_beam/transforms/userstate.py b/sdks/python/apache_beam/transforms/userstate.py
diff --git a/sdks/python/apache_beam/transforms/userstate_test.py b/sdks/python/apache_beam/transforms/userstate_test.py
