| # | |
| # Licensed to the Apache Software Foundation (ASF) under one or more | |
| # contributor license agreements. See the NOTICE file distributed with | |
| # this work for additional information regarding copyright ownership. | |
| # The ASF licenses this file to You under the Apache License, Version 2.0 | |
| # (the "License"); you may not use this file except in compliance with | |
| # the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| """Unit tests for the Beam State and Timer API interfaces.""" | |
| # pytype: skip-file | |
| import queue | |
| import threading | |
| import time | |
| import unittest | |
| from typing import Any | |
| import mock | |
| import pytest | |
| import apache_beam as beam | |
| from apache_beam.coders import BytesCoder | |
| from apache_beam.coders import ListCoder | |
| from apache_beam.coders import StrUtf8Coder | |
| from apache_beam.coders import VarIntCoder | |
| from apache_beam.options.pipeline_options import PipelineOptions | |
| from apache_beam.portability import common_urns | |
| from apache_beam.portability.api import beam_runner_api_pb2 | |
| from apache_beam.runners import pipeline_context | |
| from apache_beam.runners.common import DoFnSignature | |
| from apache_beam.runners.worker.sdk_worker import GrpcStateHandler | |
| from apache_beam.runners.worker.sdk_worker import _Future | |
| from apache_beam.testing.test_pipeline import TestPipeline | |
| from apache_beam.testing.test_stream import TestStream | |
| from apache_beam.testing.util import assert_that | |
| from apache_beam.testing.util import equal_to | |
| from apache_beam.transforms import trigger | |
| from apache_beam.transforms import userstate | |
| from apache_beam.transforms import window | |
| from apache_beam.transforms.combiners import ToListCombineFn | |
| from apache_beam.transforms.combiners import TopCombineFn | |
| from apache_beam.transforms.core import DoFn | |
| from apache_beam.transforms.timeutil import TimeDomain | |
| from apache_beam.transforms.userstate import BagStateSpec | |
| from apache_beam.transforms.userstate import CombiningValueStateSpec | |
| from apache_beam.transforms.userstate import ReadModifyWriteStateSpec | |
| from apache_beam.transforms.userstate import SetStateSpec | |
| from apache_beam.transforms.userstate import TimerSpec | |
| from apache_beam.transforms.userstate import get_dofn_specs | |
| from apache_beam.transforms.userstate import is_stateful_dofn | |
| from apache_beam.transforms.userstate import on_timer | |
| from apache_beam.transforms.userstate import validate_stateful_dofn | |
| class TestStatefulDoFn(DoFn): | |
| """An example stateful DoFn with state and timers.""" | |
| BUFFER_STATE_1 = BagStateSpec('buffer', BytesCoder()) | |
| BUFFER_STATE_2 = BagStateSpec('buffer2', VarIntCoder()) | |
| EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_3 = TimerSpec('expiry3', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_FAMILY = TimerSpec('expiry_family', TimeDomain.WATERMARK) | |
| 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), | |
| dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): | |
| yield element | |
| @on_timer(EXPIRY_TIMER_1) | |
| def on_expiry_1( | |
| self, | |
| window=DoFn.WindowParam, | |
| timestamp=DoFn.TimestampParam, | |
| key=DoFn.KeyParam, | |
| buffer=DoFn.StateParam(BUFFER_STATE_1), | |
| timer_1=DoFn.TimerParam(EXPIRY_TIMER_1), | |
| timer_2=DoFn.TimerParam(EXPIRY_TIMER_2), | |
| timer_3=DoFn.TimerParam(EXPIRY_TIMER_3)): | |
| yield 'expired1' | |
| @on_timer(EXPIRY_TIMER_2) | |
| def on_expiry_2( | |
| self, | |
| buffer=DoFn.StateParam(BUFFER_STATE_2), | |
| timer_2=DoFn.TimerParam(EXPIRY_TIMER_2), | |
| timer_3=DoFn.TimerParam(EXPIRY_TIMER_3)): | |
| yield 'expired2' | |
| @on_timer(EXPIRY_TIMER_3) | |
| 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' | |
| @on_timer(EXPIRY_TIMER_FAMILY) | |
| def on_expiry_family( | |
| self, | |
| dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY), | |
| dynamic_timer_tag=DoFn.DynamicTimerTagParam): | |
| yield (dynamic_timer_tag, 'expired_dynamic_timer') | |
| class InterfaceTest(unittest.TestCase): | |
| def _validate_dofn(self, dofn): | |
| # Construction of DoFnSignature performs validation of the given DoFn. | |
| # In particular, it ends up calling userstate._validate_stateful_dofn. | |
| # That behavior is explicitly tested below in test_validate_dofn() | |
| return DoFnSignature(dofn) | |
| @mock.patch('apache_beam.transforms.userstate.validate_stateful_dofn') | |
| def test_validate_dofn(self, unused_mock): | |
| dofn = TestStatefulDoFn() | |
| self._validate_dofn(dofn) | |
| userstate.validate_stateful_dofn.assert_called_with(dofn) | |
| def test_spec_construction(self): | |
| BagStateSpec('statename', VarIntCoder()) | |
| with self.assertRaises(TypeError): | |
| BagStateSpec(123, VarIntCoder()) | |
| CombiningValueStateSpec('statename', VarIntCoder(), TopCombineFn(10)) | |
| with self.assertRaises(TypeError): | |
| CombiningValueStateSpec(123, VarIntCoder(), TopCombineFn(10)) | |
| with self.assertRaises(TypeError): | |
| CombiningValueStateSpec('statename', VarIntCoder(), object()) | |
| SetStateSpec('setstatename', VarIntCoder()) | |
| with self.assertRaises(TypeError): | |
| SetStateSpec(123, VarIntCoder()) | |
| with self.assertRaises(TypeError): | |
| SetStateSpec('setstatename', object()) | |
| ReadModifyWriteStateSpec('valuestatename', VarIntCoder()) | |
| with self.assertRaises(TypeError): | |
| ReadModifyWriteStateSpec(123, VarIntCoder()) | |
| with self.assertRaises(TypeError): | |
| ReadModifyWriteStateSpec('valuestatename', object()) | |
| # TODO: add more spec tests | |
| with self.assertRaises(ValueError): | |
| DoFn.TimerParam(BagStateSpec('elements', BytesCoder())) | |
| TimerSpec('timer', TimeDomain.WATERMARK) | |
| TimerSpec('timer', TimeDomain.REAL_TIME) | |
| with self.assertRaises(ValueError): | |
| TimerSpec('timer', 'bogus_time_domain') | |
| with self.assertRaises(ValueError): | |
| DoFn.StateParam(TimerSpec('timer', TimeDomain.WATERMARK)) | |
| def test_state_spec_proto_conversion(self): | |
| context = pipeline_context.PipelineContext() | |
| state = BagStateSpec('statename', VarIntCoder()) | |
| state_proto = state.to_runner_api(context) | |
| self.assertEqual( | |
| beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), | |
| state_proto.protocol) | |
| context = pipeline_context.PipelineContext() | |
| state = CombiningValueStateSpec( | |
| 'statename', VarIntCoder(), TopCombineFn(10)) | |
| state_proto = state.to_runner_api(context) | |
| self.assertEqual( | |
| beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), | |
| state_proto.protocol) | |
| context = pipeline_context.PipelineContext() | |
| state = SetStateSpec('setstatename', VarIntCoder()) | |
| state_proto = state.to_runner_api(context) | |
| self.assertEqual( | |
| beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), | |
| state_proto.protocol) | |
| context = pipeline_context.PipelineContext() | |
| state = ReadModifyWriteStateSpec('valuestatename', VarIntCoder()) | |
| state_proto = state.to_runner_api(context) | |
| self.assertEqual( | |
| beam_runner_api_pb2.FunctionSpec(urn=common_urns.user_state.BAG.urn), | |
| state_proto.protocol) | |
| 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 test_stateful_dofn_detection(self): | |
| self.assertFalse(is_stateful_dofn(DoFn())) | |
| self.assertTrue(is_stateful_dofn(TestStatefulDoFn())) | |
| def test_good_signatures(self): | |
| class BasicStatefulDoFn(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| EXPIRY_TIMER = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_FAMILY = TimerSpec('expiry_family_1', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| buffer=DoFn.StateParam(BUFFER_STATE), | |
| timer1=DoFn.TimerParam(EXPIRY_TIMER), | |
| dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): | |
| yield element | |
| @on_timer(EXPIRY_TIMER) | |
| def expiry_callback(self, element, timer=DoFn.TimerParam(EXPIRY_TIMER)): | |
| yield element | |
| @on_timer(EXPIRY_TIMER_FAMILY) | |
| def expiry_family_callback( | |
| self, element, dynamic_timer=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): | |
| yield element | |
| # Validate get_dofn_specs() and timer callbacks in | |
| # DoFnSignature. | |
| stateful_dofn = BasicStatefulDoFn() | |
| signature = self._validate_dofn(stateful_dofn) | |
| expected_specs = ( | |
| set([BasicStatefulDoFn.BUFFER_STATE]), | |
| set([ | |
| BasicStatefulDoFn.EXPIRY_TIMER, | |
| BasicStatefulDoFn.EXPIRY_TIMER_FAMILY | |
| ]), | |
| ) | |
| self.assertEqual(expected_specs, get_dofn_specs(stateful_dofn)) | |
| self.assertEqual( | |
| stateful_dofn.expiry_callback, | |
| signature.timer_methods[BasicStatefulDoFn.EXPIRY_TIMER].method_value) | |
| self.assertEqual( | |
| stateful_dofn.expiry_family_callback, | |
| signature.timer_methods[ | |
| BasicStatefulDoFn.EXPIRY_TIMER_FAMILY].method_value) | |
| stateful_dofn = TestStatefulDoFn() | |
| signature = self._validate_dofn(stateful_dofn) | |
| expected_specs = ( | |
| set([TestStatefulDoFn.BUFFER_STATE_1, TestStatefulDoFn.BUFFER_STATE_2]), | |
| set([ | |
| TestStatefulDoFn.EXPIRY_TIMER_1, | |
| TestStatefulDoFn.EXPIRY_TIMER_2, | |
| TestStatefulDoFn.EXPIRY_TIMER_3, | |
| TestStatefulDoFn.EXPIRY_TIMER_FAMILY | |
| ])) | |
| self.assertEqual(expected_specs, get_dofn_specs(stateful_dofn)) | |
| self.assertEqual( | |
| stateful_dofn.on_expiry_1, | |
| signature.timer_methods[TestStatefulDoFn.EXPIRY_TIMER_1].method_value) | |
| self.assertEqual( | |
| stateful_dofn.on_expiry_2, | |
| signature.timer_methods[TestStatefulDoFn.EXPIRY_TIMER_2].method_value) | |
| self.assertEqual( | |
| stateful_dofn.on_expiry_3, | |
| signature.timer_methods[TestStatefulDoFn.EXPIRY_TIMER_3].method_value) | |
| self.assertEqual( | |
| stateful_dofn.on_expiry_family, | |
| signature.timer_methods[ | |
| TestStatefulDoFn.EXPIRY_TIMER_FAMILY].method_value) | |
| def test_bad_signatures(self): | |
| # (1) The same state parameter is duplicated on the process method. | |
| class BadStatefulDoFn1(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| def process( | |
| self, | |
| element, | |
| b1=DoFn.StateParam(BUFFER_STATE), | |
| b2=DoFn.StateParam(BUFFER_STATE)): | |
| yield element | |
| with self.assertRaises(ValueError): | |
| self._validate_dofn(BadStatefulDoFn1()) | |
| # (2) The same timer parameter is duplicated on the process method. | |
| class BadStatefulDoFn2(DoFn): | |
| TIMER = TimerSpec('timer', TimeDomain.WATERMARK) | |
| def process( | |
| self, element, t1=DoFn.TimerParam(TIMER), t2=DoFn.TimerParam(TIMER)): | |
| yield element | |
| with self.assertRaises(ValueError): | |
| self._validate_dofn(BadStatefulDoFn2()) | |
| # (3) The same state parameter is duplicated on the on_timer method. | |
| class BadStatefulDoFn3(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) | |
| @on_timer(EXPIRY_TIMER_1) | |
| def expiry_callback( | |
| self, | |
| element, | |
| b1=DoFn.StateParam(BUFFER_STATE), | |
| b2=DoFn.StateParam(BUFFER_STATE)): | |
| yield element | |
| with self.assertRaises(ValueError): | |
| self._validate_dofn(BadStatefulDoFn3()) | |
| # (4) The same timer parameter is duplicated on the on_timer method. | |
| class BadStatefulDoFn4(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) | |
| @on_timer(EXPIRY_TIMER_1) | |
| def expiry_callback( | |
| self, | |
| element, | |
| t1=DoFn.TimerParam(EXPIRY_TIMER_2), | |
| t2=DoFn.TimerParam(EXPIRY_TIMER_2)): | |
| yield element | |
| with self.assertRaises(ValueError): | |
| self._validate_dofn(BadStatefulDoFn4()) | |
| # (5) The same timer family parameter is duplicated on the process method. | |
| class BadStatefulDoFn5(DoFn): | |
| EXPIRY_TIMER_FAMILY = TimerSpec('dynamic_timer', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| dynamic_timer_1=DoFn.TimerParam(EXPIRY_TIMER_FAMILY), | |
| dynamic_timer_2=DoFn.TimerParam(EXPIRY_TIMER_FAMILY)): | |
| yield element | |
| with self.assertRaises(ValueError): | |
| self._validate_dofn(BadStatefulDoFn5()) | |
| def test_validation_typos(self): | |
| # (1) Here, the user mistakenly used the same timer spec twice for two | |
| # different timer callbacks. | |
| with self.assertRaisesRegex( | |
| ValueError, | |
| r'Multiple on_timer callbacks registered for TimerSpec\(.*expiry1\).'): | |
| class StatefulDoFnWithTimerWithTypo1(DoFn): # pylint: disable=unused-variable | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) | |
| def process(self, element): | |
| pass | |
| @on_timer(EXPIRY_TIMER_1) | |
| def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): | |
| yield 'expired1' | |
| # Note that we mistakenly associate this with the first timer. | |
| @on_timer(EXPIRY_TIMER_1) | |
| def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): | |
| yield 'expired2' | |
| # (2) Here, the user mistakenly used the same callback name and overwrote | |
| # the first on_expiry_1 callback. | |
| class StatefulDoFnWithTimerWithTypo2(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| timer1=DoFn.TimerParam(EXPIRY_TIMER_1), | |
| timer2=DoFn.TimerParam(EXPIRY_TIMER_2)): | |
| pass | |
| @on_timer(EXPIRY_TIMER_1) | |
| def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): | |
| yield 'expired1' | |
| # Note that we mistakenly reuse the "on_expiry_1" name; this is valid | |
| # syntactically in Python. | |
| @on_timer(EXPIRY_TIMER_2) | |
| def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): # pylint: disable=function-redefined | |
| yield 'expired2' | |
| # Use a stable string value for matching. | |
| def __repr__(self): | |
| return 'StatefulDoFnWithTimerWithTypo2' | |
| dofn = StatefulDoFnWithTimerWithTypo2() | |
| with self.assertRaisesRegex( | |
| ValueError, | |
| (r'The on_timer callback for TimerSpec\(.*expiry1\) is not the ' | |
| r'specified .on_expiry_1 method for DoFn ' | |
| r'StatefulDoFnWithTimerWithTypo2 \(perhaps it was overwritten\?\).')): | |
| validate_stateful_dofn(dofn) | |
| # (2) Here, the user forgot to add an on_timer decorator for 'expiry2' | |
| class StatefulDoFnWithTimerWithTypo3(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| EXPIRY_TIMER_1 = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| EXPIRY_TIMER_2 = TimerSpec('expiry2', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| timer1=DoFn.TimerParam(EXPIRY_TIMER_1), | |
| timer2=DoFn.TimerParam(EXPIRY_TIMER_2)): | |
| pass | |
| @on_timer(EXPIRY_TIMER_1) | |
| def on_expiry_1(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): | |
| yield 'expired1' | |
| def on_expiry_2(self, buffer_state=DoFn.StateParam(BUFFER_STATE)): | |
| yield 'expired2' | |
| # Use a stable string value for matching. | |
| def __repr__(self): | |
| return 'StatefulDoFnWithTimerWithTypo3' | |
| dofn = StatefulDoFnWithTimerWithTypo3() | |
| with self.assertRaisesRegex( | |
| ValueError, | |
| (r'DoFn StatefulDoFnWithTimerWithTypo3 has a TimerSpec without an ' | |
| r'associated on_timer callback: TimerSpec\(.*expiry2\).')): | |
| validate_stateful_dofn(dofn) | |
| class StatefulDoFnOnDirectRunnerTest(unittest.TestCase): | |
| # pylint: disable=expression-not-assigned | |
| all_records: list[Any] | |
| def setUp(self): | |
| # Use state on the TestCase class, since other references would be pickled | |
| # into a closure and not have the desired side effects. | |
| # | |
| # TODO(https://github.com/apache/beam/issues/18987): Use assert_that after | |
| # it works for the cases here in streaming mode. | |
| StatefulDoFnOnDirectRunnerTest.all_records = [] | |
| def record_dofn(self): | |
| class RecordDoFn(DoFn): | |
| def process(self, element): | |
| StatefulDoFnOnDirectRunnerTest.all_records.append(element) | |
| return RecordDoFn() | |
| def test_simple_stateful_dofn(self): | |
| class SimpleTestStatefulDoFn(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| EXPIRY_TIMER = TimerSpec('expiry', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| buffer=DoFn.StateParam(BUFFER_STATE), | |
| timer1=DoFn.TimerParam(EXPIRY_TIMER)): | |
| unused_key, value = element | |
| buffer.add(b'A' + str(value).encode('latin1')) | |
| timer1.set(20) | |
| @on_timer(EXPIRY_TIMER) | |
| def expiry_callback( | |
| self, | |
| buffer=DoFn.StateParam(BUFFER_STATE), | |
| timer=DoFn.TimerParam(EXPIRY_TIMER)): | |
| yield b''.join(sorted(buffer.read())) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(10).add_elements( | |
| [1, | |
| 2]).add_elements([3]).advance_watermark_to(25).add_elements([4])) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(SimpleTestStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| # Two firings should occur: once after element 3 since the timer should | |
| # fire after the watermark passes time 20, and another time after element | |
| # 4, since the timer issued at that point should fire immediately. | |
| self.assertEqual([b'A1A2A3', b'A1A2A3A4'], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_clearing_bag_state(self): | |
| class BagStateClearingStatefulDoFn(beam.DoFn): | |
| BAG_STATE = BagStateSpec('bag_state', StrUtf8Coder()) | |
| EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) | |
| CLEAR_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| bag_state=beam.DoFn.StateParam(BAG_STATE), | |
| emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), | |
| clear_timer=beam.DoFn.TimerParam(CLEAR_TIMER)): | |
| value = element[1] | |
| bag_state.add(value) | |
| clear_timer.set(100) | |
| emit_timer.set(1000) | |
| @on_timer(EMIT_TIMER) | |
| def emit_values(self, bag_state=beam.DoFn.StateParam(BAG_STATE)): | |
| for value in bag_state.read(): | |
| yield value | |
| yield 'extra' | |
| @on_timer(CLEAR_TIMER) | |
| def clear_values(self, bag_state=beam.DoFn.StateParam(BAG_STATE)): | |
| bag_state.clear() | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(0).add_elements([ | |
| ('key', 'value') | |
| ]).advance_watermark_to(100)) | |
| _ = ( | |
| p | |
| | test_stream | |
| | beam.ParDo(BagStateClearingStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual(['extra'], StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_two_timers_one_function(self): | |
| class BagStateClearingStatefulDoFn(beam.DoFn): | |
| BAG_STATE = BagStateSpec('bag_state', StrUtf8Coder()) | |
| EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) | |
| EMIT_TWICE_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| bag_state=beam.DoFn.StateParam(BAG_STATE), | |
| emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), | |
| emit_twice_timer=beam.DoFn.TimerParam(EMIT_TWICE_TIMER)): | |
| value = element[1] | |
| bag_state.add(value) | |
| emit_twice_timer.set(100) | |
| emit_timer.set(1000) | |
| @on_timer(EMIT_TWICE_TIMER) | |
| @on_timer(EMIT_TIMER) | |
| def emit_values(self, bag_state=beam.DoFn.StateParam(BAG_STATE)): | |
| for value in bag_state.read(): | |
| yield value | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(0).add_elements([ | |
| ('key', 'value') | |
| ]).advance_watermark_to(100)) | |
| _ = ( | |
| p | |
| | test_stream | |
| | beam.ParDo(BagStateClearingStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual(['value', 'value'], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_simple_read_modify_write_stateful_dofn(self): | |
| class SimpleTestReadModifyWriteStatefulDoFn(DoFn): | |
| VALUE_STATE = ReadModifyWriteStateSpec('value', StrUtf8Coder()) | |
| def process(self, element, last_element=DoFn.StateParam(VALUE_STATE)): | |
| last_element.write('%s:%s' % element) | |
| yield last_element.read() | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(0).add_elements([ | |
| ('a', 1) | |
| ]).advance_watermark_to(10).add_elements([ | |
| ('a', 3) | |
| ]).advance_watermark_to(20).add_elements([('a', 5)])) | |
| ( | |
| p | test_stream | |
| | beam.ParDo(SimpleTestReadModifyWriteStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual(['a:1', 'a:3', 'a:5'], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_clearing_read_modify_write_state(self): | |
| class SimpleClearingReadModifyWriteStatefulDoFn(DoFn): | |
| VALUE_STATE = ReadModifyWriteStateSpec('value', StrUtf8Coder()) | |
| def process(self, element, last_element=DoFn.StateParam(VALUE_STATE)): | |
| value = last_element.read() | |
| if value is not None: | |
| yield value | |
| last_element.clear() | |
| last_element.write("%s:%s" % (last_element.read(), element[1])) | |
| if element[1] == 5: | |
| yield last_element.read() | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(0).add_elements([ | |
| ('a', 1) | |
| ]).advance_watermark_to(10).add_elements([ | |
| ('a', 3) | |
| ]).advance_watermark_to(20).add_elements([('a', 5)])) | |
| ( | |
| p | test_stream | |
| | beam.ParDo(SimpleClearingReadModifyWriteStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual(['None:1', 'None:3', 'None:5'], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_simple_set_stateful_dofn(self): | |
| class SimpleTestSetStatefulDoFn(DoFn): | |
| BUFFER_STATE = SetStateSpec('buffer', VarIntCoder()) | |
| EXPIRY_TIMER = TimerSpec('expiry', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| buffer=DoFn.StateParam(BUFFER_STATE), | |
| timer1=DoFn.TimerParam(EXPIRY_TIMER)): | |
| unused_key, value = element | |
| buffer.add(value) | |
| timer1.set(20) | |
| @on_timer(EXPIRY_TIMER) | |
| def expiry_callback(self, buffer=DoFn.StateParam(BUFFER_STATE)): | |
| yield sorted(buffer.read()) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(10).add_elements( | |
| [1, 2, 3]).add_elements([2]).advance_watermark_to(24)) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(SimpleTestSetStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| # Two firings should occur: once after element 3 since the timer should | |
| # fire after the watermark passes time 20, and another time after element | |
| # 4, since the timer issued at that point should fire immediately. | |
| self.assertEqual([[1, 2, 3]], StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_clearing_set_state(self): | |
| class SetStateClearingStatefulDoFn(beam.DoFn): | |
| SET_STATE = SetStateSpec('buffer', StrUtf8Coder()) | |
| EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) | |
| CLEAR_TIMER = TimerSpec('clear_timer', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| set_state=beam.DoFn.StateParam(SET_STATE), | |
| emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), | |
| clear_timer=beam.DoFn.TimerParam(CLEAR_TIMER)): | |
| value = element[1] | |
| set_state.add(value) | |
| clear_timer.set(100) | |
| emit_timer.set(1000) | |
| @on_timer(EMIT_TIMER) | |
| def emit_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): | |
| for value in set_state.read(): | |
| yield value | |
| @on_timer(CLEAR_TIMER) | |
| def clear_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): | |
| set_state.clear() | |
| set_state.add('different-value') | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(0).add_elements([ | |
| ('key1', 'value1') | |
| ]).advance_watermark_to(100)) | |
| _ = ( | |
| p | |
| | test_stream | |
| | beam.ParDo(SetStateClearingStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual(['different-value'], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_stateful_set_state_portably(self): | |
| class SetStatefulDoFn(beam.DoFn): | |
| SET_STATE = SetStateSpec('buffer', VarIntCoder()) | |
| def process(self, element, set_state=beam.DoFn.StateParam(SET_STATE)): | |
| _, value = element | |
| aggregated_value = 0 | |
| set_state.add(value) | |
| for saved_value in set_state.read(): | |
| aggregated_value += saved_value | |
| yield aggregated_value | |
| with TestPipeline() as p: | |
| values = p | beam.Create([('key', 1), ('key', 2), ('key', 3), ('key', 4), | |
| ('key', 3)], | |
| reshuffle=False) | |
| actual_values = (values | beam.ParDo(SetStatefulDoFn())) | |
| assert_that(actual_values, equal_to([1, 3, 6, 10, 10])) | |
| # Mock random to always return 1.0 to force compaction on every add. | |
| @mock.patch( | |
| 'apache_beam.runners.worker.bundle_processor.random.random', lambda: 1.0) | |
| def test_stateful_set_state_compaction_race_portably(self): | |
| old_request = GrpcStateHandler._request | |
| request_queue = queue.Queue() | |
| def worker(): | |
| while True: | |
| handler, request, future, instruction_id = request_queue.get() | |
| time.sleep(0.1) # Simulate latency for each request sequentially. | |
| handler._context.process_instruction_id = instruction_id | |
| underlying_future = old_request(handler, request) | |
| underlying_future.wait() | |
| future.set(underlying_future.get()) | |
| request_queue.task_done() | |
| t = threading.Thread(target=worker, daemon=True) | |
| t.start() | |
| def delayed_request(self, request): | |
| if request.HasField('append') or request.HasField('clear'): | |
| future = _Future() | |
| instruction_id = getattr(self._context, 'process_instruction_id', None) | |
| request_queue.put((self, request, future, instruction_id)) | |
| return future | |
| else: | |
| return old_request(self, request) | |
| GrpcStateHandler._request = delayed_request | |
| class SetStatefulDoFn(beam.DoFn): | |
| SET_STATE = SetStateSpec('buffer', VarIntCoder()) | |
| def process(self, element, set_state=beam.DoFn.StateParam(SET_STATE)): | |
| _, value = element | |
| aggregated_value = 0 | |
| set_state.add(value) | |
| for saved_value in set_state.read(): | |
| aggregated_value += saved_value | |
| yield aggregated_value | |
| try: | |
| options = PipelineOptions([ | |
| '--max_cache_memory_usage_mb=100', | |
| '--environment_type=LOOPBACK', | |
| ]) | |
| with TestPipeline(options=options) as p: | |
| test_stream = ( | |
| TestStream( | |
| coder=beam.coders.TupleCoder(( | |
| beam.coders.StrUtf8Coder(), beam.coders.VarIntCoder() | |
| ))).advance_watermark_to(10).add_elements([ | |
| ('key', 1) | |
| ]).advance_watermark_to(20).add_elements([ | |
| ('key', 2) | |
| ]).advance_watermark_to(30)) | |
| actual_values = (p | test_stream | beam.ParDo(SetStatefulDoFn())) | |
| assert_that(actual_values, equal_to([1, 3])) | |
| finally: | |
| GrpcStateHandler._request = old_request | |
| def test_stateful_set_state_clean_portably(self): | |
| class SetStateClearingStatefulDoFn(beam.DoFn): | |
| SET_STATE = SetStateSpec('buffer', VarIntCoder()) | |
| EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| set_state=beam.DoFn.StateParam(SET_STATE), | |
| emit_timer=beam.DoFn.TimerParam(EMIT_TIMER)): | |
| _, value = element | |
| set_state.add(value) | |
| all_elements = [element for element in set_state.read()] | |
| if len(all_elements) == 5: | |
| set_state.clear() | |
| set_state.add(100) | |
| emit_timer.set(1) | |
| @on_timer(EMIT_TIMER) | |
| def emit_values(self, set_state=beam.DoFn.StateParam(SET_STATE)): | |
| yield sorted(set_state.read()) | |
| # Pin to FnApiRunner since this assumes a large bundle size to contain | |
| # all elements | |
| with TestPipeline('FnApiRunner') as p: | |
| values = p | beam.Create([('key', 1), ('key', 2), ('key', 3), ('key', 4), | |
| ('key', 5)]) | |
| actual_values = ( | |
| values | |
| | beam.Map(lambda t: window.TimestampedValue(t, 1)) | |
| | beam.WindowInto(window.FixedWindows(1)) | |
| | beam.ParDo(SetStateClearingStatefulDoFn())) | |
| assert_that(actual_values, equal_to([[100]])) | |
| def test_stateful_dofn_nonkeyed_input(self): | |
| p = TestPipeline() | |
| values = p | beam.Create([1, 2, 3]) | |
| with self.assertRaisesRegex( | |
| ValueError, | |
| ('Input elements to the transform .* with stateful DoFn must be ' | |
| 'key-value pairs.')): | |
| values | beam.ParDo(TestStatefulDoFn()) | |
| def test_generate_sequence_with_realtime_timer(self): | |
| from apache_beam.transforms.combiners import CountCombineFn | |
| class GenerateRecords(beam.DoFn): | |
| EMIT_TIMER = TimerSpec('emit_timer', TimeDomain.REAL_TIME) | |
| COUNT_STATE = CombiningValueStateSpec( | |
| 'count_state', VarIntCoder(), CountCombineFn()) | |
| def __init__(self, frequency, total_records): | |
| self.total_records = total_records | |
| self.frequency = frequency | |
| def process(self, element, emit_timer=beam.DoFn.TimerParam(EMIT_TIMER)): | |
| # Processing time timers should be set on ABSOLUTE TIME. | |
| emit_timer.set(self.frequency) | |
| yield element[1] | |
| @on_timer(EMIT_TIMER) | |
| def emit_values( | |
| self, | |
| emit_timer=beam.DoFn.TimerParam(EMIT_TIMER), | |
| count_state=beam.DoFn.StateParam(COUNT_STATE)): | |
| count = count_state.read() or 0 | |
| if self.total_records == count: | |
| return | |
| count_state.add(1) | |
| # Processing time timers should be set on ABSOLUTE TIME. | |
| emit_timer.set(count + 1 + self.frequency) | |
| yield 'value' | |
| TOTAL_RECORDS = 3 | |
| FREQUENCY = 1 | |
| test_stream = ( | |
| TestStream().advance_watermark_to(0).add_elements([ | |
| ('key', 0) | |
| ]).advance_processing_time(1) # Timestamp: 1 | |
| .add_elements([('key', 1)]).advance_processing_time(1) # Timestamp: 2 | |
| .add_elements([('key', 2)]).advance_processing_time(1) # Timestamp: 3 | |
| .add_elements([('key', 3)])) | |
| with beam.Pipeline(argv=['--streaming', '--runner=DirectRunner']) as p: | |
| _ = ( | |
| p | |
| | test_stream | |
| | beam.ParDo(GenerateRecords(FREQUENCY, TOTAL_RECORDS)) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual( | |
| # 4 RECORDS go through process | |
| # 3 values are emitted from timer | |
| # Timestamp moves gradually. | |
| [0, 'value', 1, 'value', 2, 'value', 3], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_simple_stateful_dofn_combining(self): | |
| class SimpleTestStatefulDoFn(DoFn): | |
| BUFFER_STATE = CombiningValueStateSpec( | |
| 'buffer', ListCoder(VarIntCoder()), ToListCombineFn()) | |
| EXPIRY_TIMER = TimerSpec('expiry1', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| buffer=DoFn.StateParam(BUFFER_STATE), | |
| timer1=DoFn.TimerParam(EXPIRY_TIMER)): | |
| unused_key, value = element | |
| buffer.add(value) | |
| timer1.set(20) | |
| @on_timer(EXPIRY_TIMER) | |
| def expiry_callback( | |
| self, | |
| buffer=DoFn.StateParam(BUFFER_STATE), | |
| timer=DoFn.TimerParam(EXPIRY_TIMER)): | |
| yield ''.join(str(x) for x in sorted(buffer.read())) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(10).add_elements( | |
| [1, | |
| 2]).add_elements([3]).advance_watermark_to(25).add_elements([4])) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(SimpleTestStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual(['123', '1234'], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_timer_output_timestamp(self): | |
| class TimerEmittingStatefulDoFn(DoFn): | |
| EMIT_TIMER_1 = TimerSpec('emit1', TimeDomain.WATERMARK) | |
| EMIT_TIMER_2 = TimerSpec('emit2', TimeDomain.WATERMARK) | |
| EMIT_TIMER_3 = TimerSpec('emit3', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| timer1=DoFn.TimerParam(EMIT_TIMER_1), | |
| timer2=DoFn.TimerParam(EMIT_TIMER_2), | |
| timer3=DoFn.TimerParam(EMIT_TIMER_3)): | |
| timer1.set(10) | |
| timer2.set(20) | |
| timer3.set(30) | |
| @on_timer(EMIT_TIMER_1) | |
| def emit_callback_1(self): | |
| yield 'timer1' | |
| @on_timer(EMIT_TIMER_2) | |
| def emit_callback_2(self): | |
| yield 'timer2' | |
| @on_timer(EMIT_TIMER_3) | |
| def emit_callback_3(self): | |
| yield 'timer3' | |
| class TimestampReifyingDoFn(DoFn): | |
| def process(self, element, ts=DoFn.TimestampParam): | |
| yield (element, int(ts)) | |
| with TestPipeline() as p: | |
| test_stream = (TestStream().advance_watermark_to(10).add_elements([1])) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(TimerEmittingStatefulDoFn()) | |
| | beam.ParDo(TimestampReifyingDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('timer1', 10), ('timer2', 20), ('timer3', 30)], | |
| sorted(StatefulDoFnOnDirectRunnerTest.all_records)) | |
| def test_timer_output_timestamp_and_window(self): | |
| class TimerEmittingStatefulDoFn(DoFn): | |
| EMIT_TIMER_1 = TimerSpec('emit1', TimeDomain.WATERMARK) | |
| def process(self, element, timer1=DoFn.TimerParam(EMIT_TIMER_1)): | |
| timer1.set(10) | |
| @on_timer(EMIT_TIMER_1) | |
| def emit_callback_1( | |
| self, | |
| window=DoFn.WindowParam, | |
| ts=DoFn.TimestampParam, | |
| key=DoFn.KeyParam): | |
| yield ( | |
| 'timer1-{key}'.format(key=key), | |
| int(ts), | |
| int(window.start), | |
| int(window.end)) | |
| pipeline_options = PipelineOptions() | |
| with TestPipeline(options=pipeline_options) as p: | |
| test_stream = (TestStream().advance_watermark_to(10).add_elements([1])) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | "window_into" >> beam.WindowInto( | |
| window.FixedWindows(5), | |
| accumulation_mode=trigger.AccumulationMode.DISCARDING) | |
| | beam.ParDo(TimerEmittingStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('timer1-mykey', 10, 10, 15)], | |
| sorted(StatefulDoFnOnDirectRunnerTest.all_records)) | |
| def test_timer_default_tag(self): | |
| class DynamicTimerDoFn(DoFn): | |
| EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) | |
| def process(self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)): | |
| emit.set(10) | |
| emit.set(20, dynamic_timer_tag='') | |
| @on_timer(EMIT_TIMER_FAMILY) | |
| def emit_callback( | |
| self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): | |
| yield (tag, ts) | |
| with TestPipeline() as p: | |
| test_stream = (TestStream().advance_watermark_to(10).add_elements( | |
| [1])).advance_watermark_to_infinity() | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(DynamicTimerDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('', 20)], | |
| sorted(StatefulDoFnOnDirectRunnerTest.all_records)) | |
| def test_dynamic_timer_simple_dofn(self): | |
| class DynamicTimerDoFn(DoFn): | |
| EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) | |
| 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') | |
| @on_timer(EMIT_TIMER_FAMILY) | |
| def emit_callback( | |
| self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): | |
| yield (tag, ts) | |
| with TestPipeline() as p: | |
| test_stream = (TestStream().advance_watermark_to(10).add_elements( | |
| [1])).advance_watermark_to_infinity() | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(DynamicTimerDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('emit1', 10), ('emit2', 20), ('emit3', 30)], | |
| sorted(StatefulDoFnOnDirectRunnerTest.all_records)) | |
| @pytest.mark.no_xdist | |
| @pytest.mark.timeout(60) | |
| def test_dynamic_timer_clear_then_set_timer(self): | |
| class EmitTwoEvents(DoFn): | |
| EMIT_CLEAR_SET_TIMER = TimerSpec('emitclear', TimeDomain.WATERMARK) | |
| def process(self, element, emit=DoFn.TimerParam(EMIT_CLEAR_SET_TIMER)): | |
| yield ('1', 'set') | |
| emit.set(1) | |
| @on_timer(EMIT_CLEAR_SET_TIMER) | |
| def emit_clear(self): | |
| yield ('1', 'clear') | |
| class DynamicTimerDoFn(DoFn): | |
| EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) | |
| def process(self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)): | |
| if element[1] == 'set': | |
| emit.set(10, dynamic_timer_tag='emit1') | |
| emit.set(20, dynamic_timer_tag='emit2') | |
| if element[1] == 'clear': | |
| emit.set(30, dynamic_timer_tag='emit3') | |
| emit.clear(dynamic_timer_tag='emit3') | |
| emit.set(40, dynamic_timer_tag='emit3') | |
| return [] | |
| @on_timer(EMIT_TIMER_FAMILY) | |
| def emit_callback( | |
| self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): | |
| yield (tag, ts) | |
| with TestPipeline() as p: | |
| res = ( | |
| p | |
| | beam.Create([('1', 'impulse')]) | |
| | beam.ParDo(EmitTwoEvents()) | |
| | beam.ParDo(DynamicTimerDoFn())) | |
| assert_that(res, equal_to([('emit1', 10), ('emit2', 20), ('emit3', 40)])) | |
| def test_dynamic_timer_clear_timer(self): | |
| class DynamicTimerDoFn(DoFn): | |
| EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) | |
| def process(self, element, emit=DoFn.TimerParam(EMIT_TIMER_FAMILY)): | |
| if element[1] == 'set': | |
| emit.set(10, dynamic_timer_tag='emit1') | |
| emit.set(20, dynamic_timer_tag='emit2') | |
| emit.set(30, dynamic_timer_tag='emit3') | |
| if element[1] == 'clear': | |
| emit.clear(dynamic_timer_tag='emit3') | |
| @on_timer(EMIT_TIMER_FAMILY) | |
| def emit_callback( | |
| self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): | |
| yield (tag, ts) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(5).add_elements( | |
| ['set']).advance_watermark_to(10).add_elements( | |
| ['clear']).advance_watermark_to_infinity()) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(DynamicTimerDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('emit1', 10), ('emit2', 20)], | |
| sorted(StatefulDoFnOnDirectRunnerTest.all_records)) | |
| def test_dynamic_timer_multiple(self): | |
| class DynamicTimerDoFn(DoFn): | |
| EMIT_TIMER_FAMILY1 = TimerSpec('emit_family_1', TimeDomain.WATERMARK) | |
| EMIT_TIMER_FAMILY2 = TimerSpec('emit_family_2', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| emit1=DoFn.TimerParam(EMIT_TIMER_FAMILY1), | |
| emit2=DoFn.TimerParam(EMIT_TIMER_FAMILY2)): | |
| emit1.set(10, dynamic_timer_tag='emit11') | |
| emit1.set(20, dynamic_timer_tag='emit12') | |
| emit1.set(30, dynamic_timer_tag='emit13') | |
| emit2.set(30, dynamic_timer_tag='emit21') | |
| emit2.set(20, dynamic_timer_tag='emit22') | |
| emit2.set(10, dynamic_timer_tag='emit23') | |
| @on_timer(EMIT_TIMER_FAMILY1) | |
| def emit_callback( | |
| self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): | |
| yield (tag, ts) | |
| @on_timer(EMIT_TIMER_FAMILY2) | |
| def emit_callback_2( | |
| self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): | |
| yield (tag, ts) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(5).add_elements( | |
| ['1']).advance_watermark_to_infinity()) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(DynamicTimerDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('emit11', 10), ('emit12', 20), ('emit13', 30), | |
| ('emit21', 30), ('emit22', 20), ('emit23', 10)], | |
| sorted(StatefulDoFnOnDirectRunnerTest.all_records)) | |
| def test_dynamic_timer_and_simple_timer(self): | |
| class DynamicTimerDoFn(DoFn): | |
| EMIT_TIMER_FAMILY = TimerSpec('emit', TimeDomain.WATERMARK) | |
| GC_TIMER = TimerSpec('gc', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| emit=DoFn.TimerParam(EMIT_TIMER_FAMILY), | |
| gc=DoFn.TimerParam(GC_TIMER)): | |
| emit.set(10, dynamic_timer_tag='emit1') | |
| emit.set(20, dynamic_timer_tag='emit2') | |
| emit.set(30, dynamic_timer_tag='emit3') | |
| gc.set(40) | |
| @on_timer(EMIT_TIMER_FAMILY) | |
| def emit_callback( | |
| self, ts=DoFn.TimestampParam, tag=DoFn.DynamicTimerTagParam): | |
| yield (tag, ts) | |
| @on_timer(GC_TIMER) | |
| def gc(self, ts=DoFn.TimestampParam): | |
| yield ('gc', ts) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(5).add_elements( | |
| ['1']).advance_watermark_to_infinity()) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(DynamicTimerDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('emit1', 10), ('emit2', 20), ('emit3', 30), ('gc', 40)], | |
| sorted(StatefulDoFnOnDirectRunnerTest.all_records)) | |
| def test_index_assignment(self): | |
| class IndexAssigningStatefulDoFn(DoFn): | |
| INDEX_STATE = CombiningValueStateSpec('index', sum) | |
| def process(self, element, state=DoFn.StateParam(INDEX_STATE)): | |
| unused_key, value = element | |
| current_index = state.read() | |
| yield (value, current_index) | |
| state.add(1) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(10).add_elements([ | |
| 'A', 'B' | |
| ]).add_elements(['C']).advance_watermark_to(25).add_elements(['D'])) | |
| ( | |
| p | |
| | test_stream | |
| | beam.Map(lambda x: ('mykey', x)) | |
| | beam.ParDo(IndexAssigningStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| self.assertEqual([('A', 0), ('B', 1), ('C', 2), ('D', 3)], | |
| StatefulDoFnOnDirectRunnerTest.all_records) | |
| def test_hash_join(self): | |
| class HashJoinStatefulDoFn(DoFn): | |
| BUFFER_STATE = BagStateSpec('buffer', BytesCoder()) | |
| UNMATCHED_TIMER = TimerSpec('unmatched', TimeDomain.WATERMARK) | |
| def process( | |
| self, | |
| element, | |
| state=DoFn.StateParam(BUFFER_STATE), | |
| timer=DoFn.TimerParam(UNMATCHED_TIMER)): | |
| key, value = element | |
| existing_values = list(state.read()) | |
| if not existing_values: | |
| state.add(value) | |
| timer.set(100) | |
| else: | |
| yield b'Record<%s,%s,%s>' % (key, existing_values[0], value) | |
| state.clear() | |
| timer.clear() | |
| @on_timer(UNMATCHED_TIMER) | |
| 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], ) | |
| with TestPipeline() as p: | |
| test_stream = ( | |
| TestStream().advance_watermark_to(10).add_elements([ | |
| (b'A', b'a'), (b'B', b'b') | |
| ]).add_elements([ | |
| (b'A', b'aa'), (b'C', b'c') | |
| ]).advance_watermark_to(25).add_elements([ | |
| (b'A', b'aaa'), (b'B', b'bb') | |
| ]).add_elements([ | |
| (b'D', b'd'), (b'D', b'dd'), (b'D', b'ddd'), (b'D', b'dddd') | |
| ]).advance_watermark_to(125).add_elements([(b'C', b'cc')])) | |
| ( | |
| p | |
| | test_stream | |
| | beam.ParDo(HashJoinStatefulDoFn()) | |
| | beam.ParDo(self.record_dofn())) | |
| equal_to(StatefulDoFnOnDirectRunnerTest.all_records)([ | |
| b'Record<A,a,aa>', | |
| b'Record<B,b,bb>', | |
| b'Record<D,d,dd>', | |
| b'Record<D,ddd,dddd>', | |
| b'Unmatched<aaa>', | |
| b'Unmatched<c>', | |
| b'Unmatched<cc>' | |
| ]) | |
| if __name__ == '__main__': | |
| unittest.main() |