| # |
| # 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. |
| # |
| |
| """Tests for the Watch transform.""" |
| |
| import collections |
| import typing |
| import unittest |
| |
| import apache_beam as beam |
| from apache_beam.coders.coders import BytesCoder |
| from apache_beam.coders.coders import Coder |
| from apache_beam.coders.coders import ListCoder |
| from apache_beam.coders.coders import NullableCoder |
| from apache_beam.coders.coders import StrUtf8Coder |
| from apache_beam.coders.coders import TimestampCoder |
| from apache_beam.coders.coders import TupleCoder |
| from apache_beam.coders.coders import VarIntCoder |
| from apache_beam.io.watch import PollFn |
| from apache_beam.io.watch import PollResult |
| from apache_beam.io.watch import Watch |
| from apache_beam.io.watch import _GrowthRestrictionTracker |
| from apache_beam.io.watch import _GrowthStateCoder |
| from apache_beam.io.watch import _never_seen_before |
| from apache_beam.io.watch import _NonPollingGrowthState |
| from apache_beam.io.watch import _PollingGrowthState |
| from apache_beam.io.watch import _retention_floor |
| from apache_beam.io.watch import _WatchGrowthDoFn |
| from apache_beam.io.watch import after_total_of |
| from apache_beam.io.watch import never |
| from apache_beam.io.watermark_estimators import ManualWatermarkEstimator |
| from apache_beam.options.pipeline_options import PipelineOptions |
| from apache_beam.runners.sdf_utils import RestrictionTrackerView |
| from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker |
| from apache_beam.runners.sdf_utils import ThreadsafeWatermarkEstimator |
| from apache_beam.testing.test_pipeline import TestPipeline |
| from apache_beam.testing.util import TestWindowedValue |
| from apache_beam.testing.util import assert_that |
| from apache_beam.testing.util import equal_to |
| from apache_beam.transforms.window import FixedWindows |
| from apache_beam.transforms.window import GlobalWindow |
| from apache_beam.transforms.window import TimestampedValue |
| from apache_beam.typehints import typehints |
| from apache_beam.utils.timestamp import MAX_TIMESTAMP |
| from apache_beam.utils.timestamp import MIN_TIMESTAMP |
| from apache_beam.utils.timestamp import Duration |
| from apache_beam.utils.timestamp import Timestamp |
| |
| |
| def _ts(value, timestamp): |
| return TimestampedValue(value, Timestamp(timestamp)) |
| |
| |
| def _identity(output): |
| return output |
| |
| |
| def _new_results(restriction, result, key_fn=None): |
| return _never_seen_before( |
| restriction, result, key_fn or _identity, StrUtf8Coder()) |
| |
| |
| def _tracker(restriction): |
| return _GrowthRestrictionTracker(restriction, _identity, StrUtf8Coder()) |
| |
| |
| def _cursor_tracker(restriction, allowed_lateness=Duration(0)): |
| return _GrowthRestrictionTracker( |
| restriction, |
| _identity, |
| StrUtf8Coder(), |
| timestamp_cursor=True, |
| allowed_lateness=allowed_lateness) |
| |
| |
| def _cursor_results(restriction, result, allowed_lateness=Duration(0)): |
| return _never_seen_before( |
| restriction, |
| result, |
| _identity, |
| StrUtf8Coder(), |
| _retention_floor(restriction, allowed_lateness)) |
| |
| |
| def _initial_polling(termination=None, now=Timestamp(0)): |
| termination = termination or never() |
| return _PollingGrowthState( |
| collections.OrderedDict(), None, termination.for_new_input(now, 'input')) |
| |
| |
| class PollResultTest(unittest.TestCase): |
| def test_normalize_stamps_one_processing_time_when_timestamp_none(self): |
| before = Timestamp.now() |
| result = PollResult.incomplete(['a', 'b']) |
| after = Timestamp.now() |
| # Raw outputs share a single processing-time stamp (no per-output jitter). |
| stamps = {o.timestamp for o in result.outputs} |
| self.assertEqual(1, len(stamps)) |
| ts = stamps.pop() |
| self.assertTrue(before <= ts <= after) |
| |
| def test_normalize_preserves_timestamped_and_applies_explicit_default(self): |
| result = PollResult.incomplete([_ts('a', 1), 'b'], timestamp=7) |
| by_value = {o.value: o.timestamp for o in result.outputs} |
| self.assertEqual(Timestamp(1), by_value['a']) # TimestampedValue preserved |
| self.assertEqual(Timestamp(7), by_value['b']) # raw stamped with default |
| |
| def test_complete_releases_watermark_to_max(self): |
| self.assertEqual( |
| MAX_TIMESTAMP, PollResult.complete([_ts('a', 1)]).watermark) |
| self.assertTrue(PollResult.complete([]).is_complete) |
| |
| def test_with_watermark_overrides(self): |
| self.assertEqual( |
| Timestamp(0), |
| PollResult.incomplete([_ts('a', 9)]).with_watermark(0).watermark) |
| |
| |
| class GrowthStateCoderTest(unittest.TestCase): |
| def test_polling_round_trip_preserves_resume_state(self): |
| termination = after_total_of(Duration(30)) |
| coder = _GrowthStateCoder(StrUtf8Coder(), termination) |
| completed = collections.OrderedDict([ |
| (b'a' * 16, Timestamp(1)), |
| (b'b' * 16, Timestamp(2)), |
| (b'c' * 16, Timestamp(3)), |
| ]) |
| termination_state = termination.for_new_input(Timestamp(7), 'input') |
| state = _PollingGrowthState(completed, Timestamp(5), termination_state) |
| decoded = coder.decode(coder.encode(state)) |
| self.assertEqual(list(completed.items()), list(decoded.completed.items())) |
| self.assertEqual(Timestamp(5), decoded.poll_watermark) |
| self.assertEqual(termination_state, decoded.termination_state) |
| self.assertIsNone(decoded.cursor) |
| |
| def test_polling_round_trip_preserves_cursor_and_retained_keys(self): |
| coder = _GrowthStateCoder(StrUtf8Coder(), never()) |
| completed = collections.OrderedDict([(b'a' * 16, Timestamp(42))]) |
| state = _PollingGrowthState( |
| completed, |
| Timestamp(5), |
| never().for_new_input(Timestamp(0), 'input'), |
| Timestamp(42)) |
| decoded = coder.decode(coder.encode(state)) |
| self.assertEqual(Timestamp(42), decoded.cursor) |
| self.assertEqual(list(completed.items()), list(decoded.completed.items())) |
| self.assertEqual(Timestamp(5), decoded.poll_watermark) |
| |
| def test_cursorless_state_keeps_the_pre_cursor_byte_format(self): |
| # A polling state without a cursor must encode exactly as before the |
| # cursor existed, so in-flight hash-mode restrictions decode across an |
| # upgrade in either direction. |
| termination = never() |
| coder = _GrowthStateCoder(StrUtf8Coder(), termination) |
| completed = collections.OrderedDict([(b'a' * 16, Timestamp(1))]) |
| termination_state = termination.for_new_input(Timestamp(0), 'input') |
| state = _PollingGrowthState(completed, Timestamp(5), termination_state) |
| legacy_polling_coder = TupleCoder([ |
| termination.state_coder(), |
| NullableCoder(TimestampCoder()), |
| ListCoder(TupleCoder([BytesCoder(), TimestampCoder()])), |
| ]) |
| legacy_payload = legacy_polling_coder.encode( |
| (termination_state, Timestamp(5), list(completed.items()))) |
| legacy_encoded = TupleCoder([VarIntCoder(), BytesCoder()]).encode( |
| (0, legacy_payload)) |
| self.assertEqual(legacy_encoded, coder.encode(state)) |
| decoded = coder.decode(legacy_encoded) |
| self.assertEqual(list(completed.items()), list(decoded.completed.items())) |
| self.assertIsNone(decoded.cursor) |
| |
| def test_non_polling_round_trip_preserves_pending_outputs(self): |
| coder = _GrowthStateCoder(StrUtf8Coder(), never()) |
| pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP) |
| state = _NonPollingGrowthState(pending) |
| decoded = coder.decode(coder.encode(state)) |
| self.assertEqual(MAX_TIMESTAMP, decoded.pending.watermark) |
| self.assertEqual([('a', Timestamp(1)), ('b', Timestamp(2))], |
| [(o.value, o.timestamp) for o in decoded.pending.outputs]) |
| |
| |
| class NeverSeenBeforeTest(unittest.TestCase): |
| def test_dedups_and_sorts_by_timestamp(self): |
| result = PollResult.incomplete([_ts('b', 2), _ts('a', 1), _ts('a', 1)]) |
| new_results = _new_results(_initial_polling(), result) |
| self.assertEqual(['a', 'b'], [o.value for o in new_results.outputs]) |
| |
| def test_dedups_against_completed_keys(self): |
| state = _initial_polling() |
| first = _new_results( |
| state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)])) |
| tracker = _tracker(state) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual = tracker.try_split(0) |
| second = _new_results( |
| residual, PollResult.incomplete([_ts('a', 1), _ts('c', 3)])) |
| self.assertEqual(['c'], [o.value for o in second.outputs]) |
| |
| def test_output_key_dedups_by_derived_key(self): |
| result = PollResult.incomplete([_ts('a1', 1), _ts('a2', 2), _ts('b1', 3)]) |
| # The key is the first character, so 'a1' and 'a2' collapse to one output. |
| new_results = _new_results( |
| _initial_polling(), result, key_fn=lambda output: output[0]) |
| self.assertEqual(['a1', 'b1'], [o.value for o in new_results.outputs]) |
| |
| def test_preserves_explicit_watermark(self): |
| result = PollResult.incomplete([_ts('c', 3)]).with_watermark(5) |
| new_results = _new_results(_initial_polling(), result) |
| self.assertEqual(Timestamp(5), new_results.watermark) |
| |
| |
| class GrowthTrackerTest(unittest.TestCase): |
| def test_claim_then_split_builds_replay_primary_and_merged_residual(self): |
| state = _initial_polling() |
| new_results = _new_results( |
| state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)])) |
| tracker = _tracker(state) |
| self.assertFalse(tracker.is_bounded()) |
| self.assertTrue(tracker.try_claim((new_results, 0))) |
| primary, residual = tracker.try_split(0) |
| self.assertIsInstance(primary, _NonPollingGrowthState) |
| self.assertEqual(new_results, primary.pending) |
| self.assertIsInstance(residual, _PollingGrowthState) |
| self.assertEqual(2, len(residual.completed)) |
| self.assertEqual(0, residual.termination_state) |
| self.assertTrue(tracker.check_done()) |
| |
| def test_split_merges_explicit_watermark_into_residual(self): |
| state = _initial_polling() |
| result = PollResult.incomplete([_ts('c', 3)]).with_watermark(5) |
| tracker = _tracker(state) |
| self.assertTrue(tracker.try_claim((_new_results(state, result), 0))) |
| _, residual = tracker.try_split(0) |
| self.assertEqual(Timestamp(5), residual.poll_watermark) |
| |
| def test_second_claim_is_rejected(self): |
| state = _initial_polling() |
| new_results = _new_results(state, PollResult.incomplete([_ts('a', 1)])) |
| tracker = _tracker(state) |
| self.assertTrue(tracker.try_claim((new_results, 0))) |
| self.assertFalse(tracker.try_claim((new_results, 0))) |
| |
| def test_claim_rejects_already_completed_keys(self): |
| # The tracker re-validates a claim, so a poll round that was not deduped |
| # against the restriction is rejected instead of emitting duplicates. |
| state = _initial_polling() |
| first = _new_results(state, PollResult.incomplete([_ts('a', 1)])) |
| tracker = _tracker(state) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual = tracker.try_split(0) |
| stale = PollResult.incomplete([_ts('a', 1)]) |
| self.assertFalse(_tracker(residual).try_claim((stale, 0))) |
| |
| def test_split_before_claim_moves_all_work_to_residual(self): |
| state = _initial_polling() |
| tracker = _tracker(state) |
| primary, residual = tracker.try_split(0) |
| self.assertIs(state, residual) |
| self.assertIsInstance(primary, _NonPollingGrowthState) |
| self.assertEqual((), primary.pending.outputs) |
| new_results = _new_results(state, PollResult.incomplete([_ts('a', 1)])) |
| self.assertFalse(tracker.try_claim((new_results, 0))) |
| self.assertTrue(tracker.check_done()) |
| |
| def test_non_polling_replays_exactly_the_pending_outputs(self): |
| pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP) |
| tracker = _tracker(_NonPollingGrowthState(pending)) |
| self.assertTrue(tracker.is_bounded()) |
| # A replay must claim the pending poll result exactly. |
| partial = PollResult((_ts('a', 1), ), None) |
| self.assertFalse(tracker.try_claim((partial, None))) |
| self.assertTrue(tracker.try_claim((pending, None))) |
| # A checkpoint after the replay leaves no residual work. |
| _, residual = tracker.try_split(0) |
| self.assertEqual((), residual.pending.outputs) |
| self.assertTrue(tracker.check_done()) |
| |
| def test_check_done_raises_without_claim_or_split(self): |
| tracker = _tracker(_initial_polling()) |
| with self.assertRaises(ValueError): |
| tracker.check_done() |
| |
| def test_wrapper_chain_defers_merged_residual(self): |
| state = _initial_polling() |
| new_results = _new_results( |
| state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)])) |
| threadsafe = ThreadsafeRestrictionTracker(_tracker(state)) |
| view = RestrictionTrackerView(threadsafe) |
| self.assertTrue(view.try_claim((new_results, 0))) |
| view.defer_remainder(Duration(5)) |
| residual, _ = threadsafe.deferred_status() |
| self.assertIsInstance(residual, _PollingGrowthState) |
| self.assertEqual(2, len(residual.completed)) |
| |
| def test_idle_round_reuses_completed_map_object(self): |
| # A round that discovers nothing must reuse the parent dedup map rather |
| # than copying it O(N), so a steady-state empty poll stays cheap. |
| state = _initial_polling() |
| first = _new_results(state, PollResult.incomplete([_ts('a', 1)])) |
| tracker = _tracker(state) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual1 = tracker.try_split(0) |
| resumed = _tracker(residual1) |
| empty = _new_results(residual1, PollResult.incomplete([])) |
| self.assertTrue(resumed.try_claim((empty, 0))) |
| _, residual2 = resumed.try_split(0) |
| self.assertIs(residual1.completed, residual2.completed) |
| |
| |
| class TimestampCursorTest(unittest.TestCase): |
| """Cursor-mode dedup: hash dedup whose keys the cursor retires.""" |
| def test_bounds_the_key_set_to_the_newest_event_time(self): |
| state = _initial_polling() |
| result = PollResult.incomplete([_ts('a', 1), _ts('b', 2), _ts('c', 3)]) |
| new_results = _cursor_results(state, result) |
| self.assertEqual(['a', 'b', 'c'], [o.value for o in new_results.outputs]) |
| tracker = _cursor_tracker(state) |
| self.assertTrue(tracker.try_claim((new_results, 0))) |
| _, residual = tracker.try_split(0) |
| self.assertIsInstance(residual, _PollingGrowthState) |
| self.assertEqual(Timestamp(3), residual.cursor) |
| # The two older keys are retired; only the one at the cursor is kept. |
| self.assertEqual(1, len(residual.completed)) |
| |
| def test_outputs_sharing_an_event_time_are_each_emitted_once(self): |
| # A bare cursor cannot tell two outputs at one event time apart, so it |
| # either drops the second or repeats both on the next re-list. The keys |
| # the cursor still retains are what distinguishes them. |
| state = _initial_polling() |
| first = _cursor_results( |
| state, PollResult.incomplete([_ts('a', 10), _ts('b', 10)])) |
| self.assertEqual(['a', 'b'], sorted(o.value for o in first.outputs)) |
| tracker = _cursor_tracker(state) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual = tracker.try_split(0) |
| self.assertEqual(Timestamp(10), residual.cursor) |
| relist = _cursor_results( |
| residual, |
| PollResult.incomplete([_ts('a', 10), _ts('b', 10), _ts('c', 10)])) |
| self.assertEqual(['c'], [o.value for o in relist.outputs]) |
| |
| def test_drops_outputs_the_cursor_retired(self): |
| state = _initial_polling() |
| tracker = _cursor_tracker(state) |
| first = _cursor_results(state, PollResult.incomplete([_ts('a', 10)])) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual = tracker.try_split(0) |
| self.assertEqual(Timestamp(10), residual.cursor) |
| second = _cursor_results( |
| residual, |
| PollResult.incomplete([_ts('early', 5), _ts('a', 10), _ts('c', 20)])) |
| # 'early' is below the floor, 'a' is a retained key, only 'c' is new. |
| self.assertEqual(['c'], [o.value for o in second.outputs]) |
| resumed = _cursor_tracker(residual) |
| self.assertTrue(resumed.try_claim((second, 0))) |
| _, residual = resumed.try_split(0) |
| self.assertEqual(Timestamp(20), residual.cursor) |
| |
| def test_allowed_lateness_retains_keys_below_the_cursor(self): |
| # A wider window keeps deduping outputs that arrive behind the cursor |
| # instead of taking them as already seen. |
| lateness = Duration(10) |
| state = _initial_polling() |
| tracker = _cursor_tracker(state, lateness) |
| first = _cursor_results( |
| state, PollResult.incomplete([_ts('a', 20)]), lateness) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual = tracker.try_split(0) |
| late = _cursor_results( |
| residual, |
| PollResult.incomplete([_ts('a', 20), _ts('late', 12), _ts('old', 5)]), |
| lateness) |
| self.assertEqual(['late'], [o.value for o in late.outputs]) |
| |
| def test_a_retired_key_returning_later_is_emitted_again(self): |
| # What bounding the state costs. A key is retired by the event time it was |
| # recorded with, so a key that comes back at a later event time, after the |
| # cursor has moved past the one it was recorded with, has nothing left to |
| # prove it was seen. This is the case for a file modified after the cursor |
| # passed it: it is emitted a second time, whatever the key function says |
| # about updates. Keep the default hash dedup where that matters. |
| state = _initial_polling() |
| tracker = _cursor_tracker(state) |
| first = _cursor_results(state, PollResult.incomplete([_ts('a', 10)])) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual = tracker.try_split(0) |
| # 'b' moves the cursor past the event time 'a' was recorded with, which |
| # retires 'a'. |
| second = _cursor_results( |
| residual, PollResult.incomplete([_ts('a', 10), _ts('b', 20)])) |
| self.assertEqual(['b'], [o.value for o in second.outputs]) |
| resumed = _cursor_tracker(residual) |
| self.assertTrue(resumed.try_claim((second, 0))) |
| _, residual = resumed.try_split(0) |
| self.assertEqual([Timestamp(20)], list(residual.completed.values())) |
| # 'a' now returns above the floor, so it reads as new. |
| third = _cursor_results( |
| residual, PollResult.incomplete([_ts('a', 30), _ts('b', 20)])) |
| self.assertEqual(['a'], [o.value for o in third.outputs]) |
| |
| def test_relist_emits_each_output_exactly_once(self): |
| # A full re-list of a growing collection at strictly increasing event |
| # times emits each output once; the key set stays bounded throughout. |
| state = _initial_polling() |
| emitted = collections.Counter() |
| for round_index in range(10): |
| result = PollResult.incomplete( |
| [_ts('f%d' % i, i + 1) for i in range(round_index + 1)]) |
| new_results = _cursor_results(state, result) |
| tracker = _cursor_tracker(state) |
| self.assertTrue(tracker.try_claim((new_results, 0))) |
| for output in new_results.outputs: |
| emitted[output.value] += 1 |
| _, state = tracker.try_split(0) |
| self.assertEqual(1, len(state.completed)) |
| self.assertEqual([1] * 10, [emitted['f%d' % i] for i in range(10)]) |
| self.assertEqual(Timestamp(10), state.cursor) |
| |
| def test_round_below_the_cursor_leaves_it_unchanged(self): |
| state = _initial_polling() |
| tracker = _cursor_tracker(state) |
| first = _cursor_results(state, PollResult.incomplete([_ts('a', 10)])) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual1 = tracker.try_split(0) |
| stale = _cursor_results( |
| residual1, PollResult.incomplete([_ts('a', 10), _ts('old', 4)])) |
| self.assertEqual((), stale.outputs) |
| resumed = _cursor_tracker(residual1) |
| self.assertTrue(resumed.try_claim((stale, 0))) |
| _, residual2 = resumed.try_split(0) |
| self.assertEqual(Timestamp(10), residual2.cursor) # unchanged |
| self.assertIs(residual1.completed, residual2.completed) |
| |
| def test_claim_rejects_retained_keys_and_retired_outputs(self): |
| # The tracker re-validates a claim, so a round that was not filtered is |
| # rejected instead of emitting already-seen outputs. |
| state = _initial_polling() |
| tracker = _cursor_tracker(state) |
| first = _cursor_results(state, PollResult.incomplete([_ts('a', 10)])) |
| self.assertTrue(tracker.try_claim((first, 0))) |
| _, residual = tracker.try_split(0) |
| self.assertFalse( |
| _cursor_tracker(residual).try_claim( |
| (PollResult.incomplete([_ts('a', 10)]), 0))) |
| self.assertFalse( |
| _cursor_tracker(residual).try_claim( |
| (PollResult.incomplete([_ts('old', 4)]), 0))) |
| |
| def test_switching_hash_state_to_cursor_keeps_the_keys(self): |
| # A restriction resumed in cursor mode still holds the hashes from its |
| # hash rounds, so nothing re-emits; the cursor retires them from there on. |
| state = _initial_polling() |
| hash_tracker = _tracker(state) |
| first = _new_results(state, PollResult.incomplete([_ts('a', 5)])) |
| self.assertTrue(hash_tracker.try_claim((first, 0))) |
| _, legacy = hash_tracker.try_split(0) |
| self.assertIsNone(legacy.cursor) |
| relist = _cursor_results( |
| legacy, PollResult.incomplete([_ts('a', 5), _ts('c', 20)])) |
| self.assertEqual(['c'], [o.value for o in relist.outputs]) |
| tracker = _cursor_tracker(legacy) |
| self.assertTrue(tracker.try_claim((relist, 0))) |
| _, residual = tracker.try_split(0) |
| self.assertEqual(Timestamp(20), residual.cursor) |
| self.assertEqual(1, len(residual.completed)) |
| |
| def test_hash_round_drops_a_stale_cursor(self): |
| # The reverse switch: a hash round retains every key, so the cursor that |
| # would retire them is dropped. |
| state = _PollingGrowthState( |
| collections.OrderedDict(), None, 0, cursor=Timestamp(10)) |
| tracker = _tracker(state) |
| result = _new_results(state, PollResult.incomplete([_ts('a', 20)])) |
| self.assertTrue(tracker.try_claim((result, 0))) |
| _, residual = tracker.try_split(0) |
| self.assertIsNone(residual.cursor) |
| self.assertEqual(1, len(residual.completed)) |
| |
| def test_cursor_state_encoding_size_is_independent_of_outputs(self): |
| coder = _GrowthStateCoder(StrUtf8Coder(), never()) |
| |
| def encoded_residual_after_claiming(count): |
| state = _initial_polling() |
| result = PollResult.incomplete( |
| [_ts('output%d' % i, i + 1) for i in range(count)]) |
| tracker = _cursor_tracker(state) |
| self.assertTrue(tracker.try_claim((_cursor_results(state, result), 0))) |
| _, residual = tracker.try_split(0) |
| return coder.encode(residual) |
| |
| self.assertEqual( |
| len(encoded_residual_after_claiming(1)), |
| len(encoded_residual_after_claiming(100))) |
| |
| |
| class TerminationConditionTest(unittest.TestCase): |
| def test_never_does_not_stop(self): |
| termination = never() |
| state = termination.for_new_input(Timestamp(0), 'input') |
| self.assertFalse(termination.can_stop_polling(MAX_TIMESTAMP, state)) |
| |
| def test_after_total_of_stops_once_duration_elapsed(self): |
| termination = after_total_of(10) |
| state = termination.for_new_input(Timestamp(0), 'input') |
| self.assertFalse(termination.can_stop_polling(Timestamp(10), state)) |
| self.assertTrue(termination.can_stop_polling(Timestamp(11), state)) |
| |
| |
| # Module-level so the poll function pickles by reference; the call counter is |
| # shared within the single in-memory DirectRunner process. |
| _POLL_CALLS = collections.defaultdict(int) |
| |
| |
| def _growing_poll(prefix): |
| # Unannotated on purpose: dedup must hold on the inferred fallback coder. |
| _POLL_CALLS[prefix] += 1 |
| count = _POLL_CALLS[prefix] |
| outputs = [_ts('%s%d' % (prefix, i), i + 1) for i in range(count)] |
| if count >= 3: |
| return PollResult.complete(outputs) |
| return PollResult.incomplete(outputs) |
| |
| |
| def _complete_poll(prefix) -> PollResult[str]: |
| return PollResult.complete([_ts(prefix + 'a', 1), _ts(prefix + 'b', 2)]) |
| |
| |
| def _first_char(output): |
| return output[0] |
| |
| |
| def _empty_poll(unused_element): |
| return PollResult.incomplete([]) |
| |
| |
| def _out_of_order_poll(prefix): |
| # Round 1 emits late_after@10 (advances the watermark to 10); round 2 emits |
| # early@5, which is behind the watermark and therefore late. |
| _POLL_CALLS[prefix] += 1 |
| if _POLL_CALLS[prefix] == 1: |
| return PollResult.incomplete([_ts(prefix + 'late_after', 10)]) |
| return PollResult.complete([_ts(prefix + 'early', 5)]) |
| |
| |
| def _max_timestamp_poll(unused_element): |
| return PollResult.incomplete( |
| [_ts('a', 10), TimestampedValue('b', MAX_TIMESTAMP)]) |
| |
| |
| def _keyed_poll(prefix): |
| # 'a1' and 'a2' share the dedup key 'a', so only 'a1' is emitted. |
| return PollResult.complete([_ts('a1', 1), _ts('a2', 2), _ts('b1', 3)]) |
| |
| |
| class _StrCoderPollFn(PollFn): |
| def __call__(self, element): |
| return PollResult.complete([_ts(element + 'a', 1)]) |
| |
| def default_output_coder(self): |
| return StrUtf8Coder() |
| |
| |
| class _NoDeterministicFormCoder(Coder): |
| def encode(self, value): |
| return b'' |
| |
| def decode(self, encoded): |
| return None |
| |
| def is_deterministic(self): |
| return False |
| |
| |
| def _windowed_group(kv, window=beam.DoFn.WindowParam): |
| return ((window.start, window.end), sorted(kv[1])) |
| |
| |
| class WatchDoFnProcessTest(unittest.TestCase): |
| def _process( |
| self, |
| poll_fn, |
| element, |
| timestamp, |
| restriction=None, |
| watermark=None, |
| timestamp_cursor=False): |
| dofn = _WatchGrowthDoFn( |
| poll_fn, |
| never(), |
| Duration(1), |
| StrUtf8Coder(), |
| _identity, |
| StrUtf8Coder(), |
| timestamp_cursor) |
| if restriction is None: |
| restriction = dofn.initial_restriction(element) |
| threadsafe = ThreadsafeRestrictionTracker(dofn.create_tracker(restriction)) |
| estimator = ThreadsafeWatermarkEstimator( |
| ManualWatermarkEstimator(watermark)) |
| outputs = list( |
| dofn.process( |
| element, |
| timestamp=timestamp, |
| tracker=RestrictionTrackerView(threadsafe), |
| watermark_estimator=estimator)) |
| return outputs, threadsafe, estimator |
| |
| def test_empty_round_holds_watermark_at_input_timestamp(self): |
| outputs, threadsafe, estimator = self._process( |
| _empty_poll, 'in', Timestamp(7)) |
| self.assertEqual([], outputs) |
| # The estimator is seeded from the input timestamp, so the deferred |
| # residual holds the watermark there instead of at MIN_TIMESTAMP. |
| self.assertEqual(Timestamp(7), estimator.current_watermark()) |
| residual, _ = threadsafe.deferred_status() |
| self.assertIsInstance(residual, _PollingGrowthState) |
| |
| def test_complete_round_stops_without_residual(self): |
| outputs, threadsafe, _ = self._process(_complete_poll, 'k:', Timestamp(0)) |
| self.assertEqual([('k:', 'k:a'), ('k:', 'k:b')], |
| [value.value for value in outputs]) |
| self.assertIsNone(threadsafe.deferred_status()) |
| self.assertTrue(threadsafe.check_done()) |
| |
| def test_replay_round_leaves_the_watermark_alone(self): |
| pending = PollResult((_ts('k:a', 1), _ts('k:b', 2)), MAX_TIMESTAMP) |
| outputs, threadsafe, estimator = self._process( |
| _empty_poll, |
| 'k:', |
| Timestamp(7), |
| restriction=_NonPollingGrowthState(pending)) |
| self.assertEqual([('k:', 'k:a'), ('k:', 'k:b')], |
| [value.value for value in outputs]) |
| # The replay branch holds the watermark at the seed, so it never runs ahead |
| # of the replayed outputs and never releases to MAX_TIMESTAMP itself. |
| self.assertEqual(Timestamp(7), estimator.current_watermark()) |
| self.assertIsNone(threadsafe.deferred_status()) |
| self.assertTrue(threadsafe.check_done()) |
| |
| def test_terminal_round_after_deferring_leaves_no_residual(self): |
| _POLL_CALLS.clear() |
| # Round one defers and parks the watermark on the new output's time. |
| _, threadsafe, estimator = self._process(_growing_poll, 'd:', Timestamp(0)) |
| residual, _ = threadsafe.deferred_status() |
| self.assertIsInstance(residual, _PollingGrowthState) |
| self.assertEqual(Timestamp(1), estimator.current_watermark()) |
| # Round two resumes from that residual, carrying the watermark forward. |
| _, threadsafe, estimator = self._process( |
| _growing_poll, |
| 'd:', |
| Timestamp(0), |
| restriction=residual, |
| watermark=estimator.current_watermark()) |
| residual, _ = threadsafe.deferred_status() |
| self.assertEqual(Timestamp(2), estimator.current_watermark()) |
| # Round three completes. The watermark stays where round two left it and |
| # the round reports no residual, so nothing carries that hold forward. |
| outputs, threadsafe, estimator = self._process( |
| _growing_poll, |
| 'd:', |
| Timestamp(0), |
| restriction=residual, |
| watermark=estimator.current_watermark()) |
| self.assertEqual([('d:', 'd:2')], [value.value for value in outputs]) |
| self.assertEqual(Timestamp(2), estimator.current_watermark()) |
| self.assertIsNone(threadsafe.deferred_status()) |
| self.assertTrue(threadsafe.check_done()) |
| |
| def test_cursor_at_max_timestamp_stops_polling(self): |
| # A cursor reaching MAX is terminal: nothing can be strictly past it, so |
| # the round stops instead of polling forever and dropping every output. |
| outputs, threadsafe, _ = self._process( |
| _max_timestamp_poll, 'k:', Timestamp(0), timestamp_cursor=True) |
| self.assertEqual([('k:', 'a'), ('k:', 'b')], |
| [value.value for value in outputs]) |
| self.assertIsNone(threadsafe.deferred_status()) |
| self.assertTrue(threadsafe.check_done()) |
| |
| def test_resumed_cursor_at_max_stops_without_polling(self): |
| # A restriction resumed with the cursor already at MAX (persisted by a |
| # checkpoint after a MAX-timestamped round) must stop without invoking the |
| # poll function at all. |
| polls = [] |
| |
| def poll(unused_element): |
| polls.append(1) |
| return PollResult.incomplete([]) |
| |
| resumed = _PollingGrowthState( |
| collections.OrderedDict(), |
| None, |
| never().for_new_input(Timestamp(0), 'input'), |
| MAX_TIMESTAMP) |
| outputs, threadsafe, _ = self._process( |
| poll, 'k:', Timestamp(0), restriction=resumed, timestamp_cursor=True) |
| self.assertEqual([], outputs) |
| self.assertEqual([], polls) # the poll function never ran |
| self.assertIsNone(threadsafe.deferred_status()) |
| self.assertTrue(threadsafe.check_done()) |
| |
| def test_out_of_order_new_output_emits_late_and_warns(self): |
| # Round 1 surfaces late_after@10 and parks the watermark there; round 2 |
| # surfaces a brand-new early@5. The output is emitted at its true (earlier) |
| # time, so it is late for downstream windowing, and Watch warns about it. |
| _POLL_CALLS.clear() |
| _, threadsafe, estimator = self._process( |
| _out_of_order_poll, 'k:', Timestamp(0)) |
| self.assertEqual(Timestamp(10), estimator.current_watermark()) |
| residual, _ = threadsafe.deferred_status() |
| with self.assertLogs('apache_beam.io.watch', level='WARNING') as logs: |
| outputs, _, _ = self._process( |
| _out_of_order_poll, |
| 'k:', |
| Timestamp(0), |
| restriction=residual, |
| watermark=estimator.current_watermark()) |
| self.assertEqual([('k:', 'k:early')], [value.value for value in outputs]) |
| self.assertEqual([Timestamp(5)], [value.timestamp for value in outputs]) |
| self.assertTrue( |
| any('behind the watermark' in line for line in logs.output), |
| 'expected a late-emission warning, got: %s' % logs.output) |
| |
| def test_first_round_early_output_does_not_warn(self): |
| # While the estimator holds the input element's timestamp seed, an output |
| # behind it must not trigger the out-of-order warning: the seed is not a |
| # poll-order signal. |
| def poll(unused_element): |
| return PollResult.incomplete([_ts('a', 5)]) |
| |
| with self.assertNoLogs('apache_beam.io.watch', level='WARNING'): |
| outputs, _, _ = self._process(poll, 'k:', Timestamp(10)) |
| self.assertEqual([Timestamp(5)], [value.timestamp for value in outputs]) |
| |
| def test_early_output_after_empty_poll_does_not_warn(self): |
| # An empty first poll defers with the watermark still at the element seed; |
| # the next round's first real output must not be treated as out-of-order |
| # either; the watermark has not advanced past the seed. |
| polls = [] |
| |
| def poll(unused_element): |
| polls.append(len(polls)) |
| if len(polls) == 1: |
| return PollResult.incomplete([]) |
| return PollResult.incomplete([_ts('a', 5)]) |
| |
| _, threadsafe, estimator = self._process(poll, 'k:', Timestamp(10)) |
| self.assertEqual(Timestamp(10), estimator.current_watermark()) |
| residual, _ = threadsafe.deferred_status() |
| with self.assertNoLogs('apache_beam.io.watch', level='WARNING'): |
| outputs, _, _ = self._process( |
| poll, |
| 'k:', |
| Timestamp(10), |
| restriction=residual, |
| watermark=estimator.current_watermark()) |
| self.assertEqual([Timestamp(5)], [value.timestamp for value in outputs]) |
| |
| def test_explicit_watermark_holds_below_output_time(self): |
| # An explicit watermark below the output's own event time is honored, so |
| # a later, earlier-timestamped output stays on time (the out-of-order-safe |
| # path). |
| def poll(unused_element): |
| return PollResult.incomplete([_ts('a', 10)]).with_watermark(0) |
| |
| _, threadsafe, estimator = self._process(poll, 'k:', Timestamp(0)) |
| self.assertEqual(Timestamp(0), estimator.current_watermark()) |
| residual, _ = threadsafe.deferred_status() |
| self.assertEqual(Timestamp(0), residual.poll_watermark) |
| |
| |
| class WatchEndToEndTest(unittest.TestCase): |
| def _in_memory_pipeline(self): |
| return TestPipeline( |
| options=PipelineOptions(direct_running_mode='in_memory')) |
| |
| def test_complete_outputs_values_and_timestamps(self): |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['k:']) |
| | Watch(_complete_poll, poll_interval=Duration(1))) |
| assert_that( |
| output, |
| equal_to([ |
| TestWindowedValue(('k:', 'k:a'), Timestamp(1), [GlobalWindow()]), |
| TestWindowedValue(('k:', 'k:b'), Timestamp(2), [GlobalWindow()]), |
| ]), |
| reify_windows=True) |
| |
| def test_complete_advances_watermark_for_windowed_pipeline(self): |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['k:']) |
| | Watch(_complete_poll, poll_interval=Duration(1))) |
| grouped = ( |
| output |
| | beam.WindowInto(FixedWindows(10)) |
| | beam.Map(lambda kv: ('all', kv[1])) |
| | beam.GroupByKey() |
| | beam.Map(_windowed_group)) |
| assert_that( |
| grouped, |
| equal_to([ |
| ((Timestamp(0), Timestamp(10)), ['k:a', 'k:b']), |
| ])) |
| |
| def test_multi_round_dedups_stops_and_is_per_input(self): |
| _POLL_CALLS.clear() |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['x:', 'y:']) |
| | Watch(_growing_poll, poll_interval=Duration(0.05))) |
| assert_that( |
| output, |
| equal_to([('x:', 'x:0'), ('x:', 'x:1'), ('x:', 'x:2'), ('y:', 'y:0'), |
| ('y:', 'y:1'), ('y:', 'y:2')])) |
| self.assertEqual(3, _POLL_CALLS['x:']) |
| self.assertEqual(3, _POLL_CALLS['y:']) |
| |
| def test_timestamp_cursor_dedups_growing_source(self): |
| _POLL_CALLS.clear() |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['x:', 'y:']) |
| | Watch( |
| _growing_poll, |
| poll_interval=Duration(0.05), |
| timestamp_cursor=True)) |
| # Each output is emitted exactly once, with the cursor retiring keys as |
| # it advances, across poll rounds and checkpoints. |
| assert_that( |
| output, |
| equal_to([('x:', 'x:0'), ('x:', 'x:1'), ('x:', 'x:2'), ('y:', 'y:0'), |
| ('y:', 'y:1'), ('y:', 'y:2')])) |
| |
| def test_timestamp_cursor_composes_with_an_output_key(self): |
| # The cursor bounds the state; the key still decides what counts as seen. |
| _POLL_CALLS.clear() |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['x:']) |
| | Watch( |
| _growing_poll, |
| poll_interval=Duration(0.05), |
| output_key_fn=_first_char, |
| timestamp_cursor=True)) |
| # Every output shares a key, so only the first one is ever emitted. |
| assert_that(output, equal_to([('x:', 'x:0')])) |
| |
| def test_output_key_dedups_across_pipeline(self): |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['k']) |
| | Watch( |
| _keyed_poll, poll_interval=Duration(1), |
| output_key_fn=_first_char)) |
| assert_that(output, equal_to([('k', 'a1'), ('k', 'b1')])) |
| |
| def test_rejects_key_coder_without_deterministic_form(self): |
| with self.assertRaises(ValueError): |
| with self._in_memory_pipeline() as p: |
| _ = ( |
| p | beam.Create(['k:']) |
| | Watch( |
| _complete_poll, |
| poll_interval=Duration(1), |
| output_key_coder=_NoDeterministicFormCoder())) |
| |
| def test_infers_output_coder_from_return_annotation(self): |
| # _complete_poll is annotated ``-> PollResult[str]``, so the output coder |
| # and with it the (input, output) element type are inferred without hints. |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['k:']) |
| | Watch(_complete_poll, poll_interval=Duration(1))) |
| self.assertEqual(typehints.Tuple[str, str], output.element_type) |
| |
| def test_infers_coder_from_generic_annotations(self): |
| # tuple[str, float] and typing.Tuple[str, float] resolve to a tuple coder, |
| # not the pickling fallback. |
| def native_poll(element) -> PollResult[tuple[str, float]]: |
| return PollResult.complete([(element, 1.0)]) |
| |
| def typing_poll( |
| element) -> PollResult[typing.Tuple[str, float]]: # noqa: UP006 |
| return PollResult.complete([(element, 1.0)]) |
| |
| for poll in (native_poll, typing_poll): |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['k:']) | Watch(poll, poll_interval=Duration(1))) |
| self.assertEqual( |
| typehints.Tuple[str, typehints.Tuple[str, float]], |
| output.element_type) |
| |
| def test_uses_poll_fn_default_output_coder(self): |
| with self._in_memory_pipeline() as p: |
| output = ( |
| p | beam.Create(['k:']) |
| | Watch(_StrCoderPollFn(), poll_interval=Duration(1))) |
| self.assertEqual(typehints.Tuple[str, str], output.element_type) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
