def __init__( self, sdf, args_for_invoker, kwargs_for_invoker): self.sdf = sdf self._element_tag = _ValueStateTag('element') self._restriction_tag = _ValueStateTag('restriction') self.watermark_hold_tag = _ValueStateTag('watermark_hold') self._process_element_invoker = None self.sdf_invoker = DoFnInvoker.create_invoker( DoFnSignature(self.sdf), context=DoFnContext('unused_context'), input_args=args_for_invoker, input_kwargs=kwargs_for_invoker) self._step_context = None
def __init__( self, sdf, args_for_invoker, kwargs_for_invoker): self.sdf = sdf self._element_tag = _ValueStateTag('element') self._restriction_tag = _ValueStateTag('restriction') self.watermark_hold_tag = _ValueStateTag('watermark_hold') self._process_element_invoker = None self.sdf_invoker = DoFnInvoker.create_invoker( DoFnSignature(self.sdf), context=DoFnContext('unused_context'), input_args=args_for_invoker, input_kwargs=kwargs_for_invoker) self._step_context = None
class _WatermarkControllerEvaluator(_TransformEvaluator): """TransformEvaluator for the _WatermarkController transform. This is used to enable multiple output watermarks for the TestStream. """ # The state tag used to store the watermark. WATERMARK_TAG = _ValueStateTag( '_WatermarkControllerEvaluator_Watermark_Tag') def __init__(self, evaluation_context, applied_ptransform, input_committed_bundle, side_inputs): assert not side_inputs self.transform = applied_ptransform.transform super(_WatermarkControllerEvaluator, self).__init__(evaluation_context, applied_ptransform, input_committed_bundle, side_inputs) self._init_state() def _init_state(self): """Gets and sets the initial state. This is used to keep track of the watermark hold between calls. """ transform_states = self._evaluation_context._transform_keyed_states state = transform_states[self._applied_ptransform] if self.WATERMARK_TAG not in state: watermark_state = InMemoryUnmergedState() watermark_state.set_global_state(self.WATERMARK_TAG, MIN_TIMESTAMP) state[self.WATERMARK_TAG] = watermark_state self._state = state[self.WATERMARK_TAG] @property def _watermark(self): return self._state.get_global_state(self.WATERMARK_TAG) @_watermark.setter def _watermark(self, watermark): self._state.set_global_state(self.WATERMARK_TAG, watermark) def start_bundle(self): self.bundles = [] def process_element(self, element): # In order to keep the order of the elements between the script and what # flows through the pipeline the same, emit the elements here. event = element.value if isinstance(event, WatermarkEvent): self._watermark = event.new_watermark elif isinstance(event, ElementEvent): main_output = list(self._outputs)[0] bundle = self._evaluation_context.create_bundle(main_output) for tv in event.timestamped_values: # Unreify the value into the correct window. try: bundle.output(WindowedValue(**tv.value)) except TypeError: bundle.output( GlobalWindows.windowed_value(tv.value, timestamp=tv.timestamp)) self.bundles.append(bundle) def finish_bundle(self): # The watermark hold we set here is the way we allow the TestStream events # to control the output watermark. return TransformResult(self, self.bundles, [], None, {None: self._watermark})