def process(self, element, t=DoFn.TimestampParam,
buffer_1=DoFn.StateParam(BUFFER_STATE_1),
buffer_2=DoFn.StateParam(BUFFER_STATE_2),
timer_1=DoFn.TimerParam(EXPIRY_TIMER_1),
timer_2=DoFn.TimerParam(EXPIRY_TIMER_2)):
yield element
def process(self,
element,
emit=DoFn.TimerParam(EMIT_CLEAR_SET_TIMER)):
yield ('1', 'set')
emit.set(1)
def process(self,
element,
timer1=DoFn.TimerParam(EXPIRY_TIMER_1),
timer2=DoFn.TimerParam(EXPIRY_TIMER_2)):
pass
def expiry_callback(self,
buffer=DoFn.StateParam(BUFFER_STATE),
timer=DoFn.TimerParam(EXPIRY_TIMER)):
yield b''.join(sorted(buffer.read()))
def process(self,
element,
b1=DoFn.StateParam(BUFFER_STATE),
b2=DoFn.StateParam(BUFFER_STATE)):
yield element
def expiry_callback(self,
element,
b1=DoFn.StateParam(BUFFER_STATE),
b2=DoFn.StateParam(BUFFER_STATE)):
yield element
def expiry_family_callback(
self,
element,
dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)):
yield element
def process(self,
element,
buffer=DoFn.StateParam(BUFFER_STATE),
timer1=DoFn.TimerParam(EXPIRY_TIMER)):
yield element
def process(self, element, state=DoFn.StateParam(INDEX_STATE)):
unused_key, value = element
current_index = state.read()
yield (value, current_index)
state.add(1)
def expiry_callback(self, state=DoFn.StateParam(BUFFER_STATE)):
buffered = list(state.read())
assert len(buffered) == 1, buffered
state.clear()
yield b'Unmatched<%s>' % (buffered[0], )
def on_expiry_family(self,
dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY),
dynamic_timer_tag=DoFn.DynamicTimerTagParam):
yield (dynamic_timer_tag, 'expired_dynamic_timer')
def on_window_timer(self,
element_state=DoFn.StateParam(ELEMENT_STATE),
count_state=DoFn.StateParam(COUNT_STATE),
buffering_timer=DoFn.TimerParam(BUFFERING_TIMER)):
return self.flush_batch(element_state, count_state,
buffering_timer)
def process(self, element, buffer=DoFn.StateParam(BUFFER_STATE),
timer1=DoFn.TimerParam(EXPIRY_TIMER)):
unused_key, value = element
buffer.add(value)
timer1.set(20)
def test_param_construction(self):
with self.assertRaises(ValueError):
DoFn.StateParam(TimerSpec('timer', TimeDomain.WATERMARK))
with self.assertRaises(ValueError):
DoFn.TimerParam(BagStateSpec('elements', BytesCoder()))
def process(self,
element,
dynamic_timer_1=DoFn.TimerParam(EXPIRY_TIMER_FAMILY),
dynamic_timer_2=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)):
yield element
def test_stateful_dofn_detection(self):
self.assertFalse(is_stateful_dofn(DoFn()))
self.assertTrue(is_stateful_dofn(TestStatefulDoFn()))
def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): # pylint: disable=function-redefined
yield 'expired2'
def expiry_callback(self,
element,
timer=DoFn.TimerParam(EXPIRY_TIMER)):
yield element
def process(self,
element,
last_element=DoFn.StateParam(VALUE_STATE)):
last_element.write('%s:%s' % element)
yield last_element.read()
def process(self,
element,
t1=DoFn.TimerParam(TIMER),
t2=DoFn.TimerParam(TIMER)):
yield element
def process(self, element, timer1=DoFn.TimerParam(EMIT_TIMER_1)):
timer1.set(10)
def expiry_callback(self,
element,
t1=DoFn.TimerParam(EXPIRY_TIMER_2),
t2=DoFn.TimerParam(EXPIRY_TIMER_2)):
yield element
def process(self, element,
emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)):
emit.set(10, dynamic_timer_tag='emit1')
emit.set(20, dynamic_timer_tag='emit2')
emit.set(30, dynamic_timer_tag='emit3')
def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)):
yield 'expired2'
def on_expiry_3(self,
buffer_1=DoFn.StateParam(BUFFER_STATE_1),
buffer_2=DoFn.StateParam(BUFFER_STATE_2),
timer_3=DoFn.TimerParam(EXPIRY_TIMER_3)):
yield 'expired3'
def expiry_callback(self, buffer=DoFn.StateParam(BUFFER_STATE)):
yield sorted(buffer.read())
def process(self, element, state=DoFn.StateParam(INDEX_STATE)):
unused_key, value = element
next_index, = list(state.read()) or [0]
yield (value, next_index)
state.clear()
state.add(next_index + 1)
