def __init__(self, sdf, args_for_invoker, kwargs_for_invoker):
self.sdf = sdf
self._element_tag = _ReadModifyWriteStateTag('element')
self._restriction_tag = _ReadModifyWriteStateTag('restriction')
self.watermark_hold_tag = _ReadModifyWriteStateTag('watermark_hold')
self._process_element_invoker = None
self._output_processor = _OutputProcessor()
self.sdf_invoker = DoFnInvoker.create_invoker(
DoFnSignature(self.sdf),
context=DoFnContext('unused_context'),
output_processor=self._output_processor,
input_args=args_for_invoker,
input_kwargs=kwargs_for_invoker)
self._step_context = None
def __init__(self, step_context, dofn, key_coder):
self.step_context = step_context
self.dofn = dofn
self.key_coder = key_coder
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
if isinstance(state_spec, userstate.ReadModifyWriteStateSpec):
state_tag = _ReadModifyWriteStateTag(state_key)
elif isinstance(state_spec, userstate.BagStateSpec):
state_tag = _ListStateTag(state_key)
elif isinstance(state_spec, userstate.CombiningValueStateSpec):
state_tag = _ListStateTag(state_key)
elif isinstance(state_spec, userstate.SetStateSpec):
state_tag = _SetStateTag(state_key)
else:
raise ValueError('Invalid state spec: %s' % state_spec)
self.state_tags[state_spec] = state_tag
self.cached_states = {}
self.cached_timers = {}
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 = _ReadModifyWriteStateTag(
'_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().__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.
if isinstance(tv.value, WindowedValueHolder):
bundle.output(tv.value.windowed_value)
else:
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})