From 85476e4eb3a178db27f3b9bf5195c5897e69b6f2 Mon Sep 17 00:00:00 2001 From: Eliaazzz Date: Wed, 24 Jun 2026 23:32:19 +1000 Subject: [PATCH] [Python] Bound Watch state with a timestamp cursor Follow-up to the experimental Watch transform that documents the event-time contract, makes late data observable, and adds an opt-in O(1) bound on the per-input state. Watermark/event-time: - Document the event-time/watermark contract on the module and PollResult. - Warn (throttled) when an output is emitted behind the current watermark, so out-of-order late data is observable instead of silently dropped, and point to PollResult.with_watermark for out-of-order sources. Scalability: - Add opt-in Watch.with_timestamp_cursor(): dedup by a high-water-mark timestamp instead of by value identity. Watch keeps only the greatest event time emitted per input and emits the polled outputs strictly past it, so the per-input state is a single timestamp that never grows and the poll result is not hashed, giving O(1) state and per-checkpoint encoding. For sources whose outputs carry strictly increasing event-time timestamps; an output at or below the cursor is treated as already seen, so out-of-order or equal-timestamp outputs are dropped (use the default exact dedup otherwise). - Reuse the parent dedup map on idle rounds to drop the O(N) per-round copy. Tests cover the watermark contract, the out-of-order late warning, cursor emit-after-high-water-mark, O(1) state, relist-once, complete, cursor coder round-trip, and end-to-end cursor dedup. 26 tests and 5 subtests pass. --- sdks/python/apache_beam/io/watch.py | 228 ++++++++++++++++++---- sdks/python/apache_beam/io/watch_test.py | 232 +++++++++++++++++++++++ 2 files changed, 424 insertions(+), 36 deletions(-) diff --git a/sdks/python/apache_beam/io/watch.py b/sdks/python/apache_beam/io/watch.py index 8ccbd4f8e6e5..c6e79e9746e9 100644 --- a/sdks/python/apache_beam/io/watch.py +++ b/sdks/python/apache_beam/io/watch.py @@ -28,9 +28,11 @@ update watermark -> check termination -> wait(poll_interval) -> poll -> ... The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each -output carries the event time the poll function first reported it. Dedup uses a -stable 128-bit hash of the encoded output, so the output coder must be -deterministic for dedup to hold across workers and restarts. +output carries the event time the poll function first reported it. By default +dedup is by value identity, using a stable 128-bit hash of the encoded output, +so the output coder must be deterministic for dedup to hold across workers and +restarts. :meth:`Watch.with_timestamp_cursor` switches to bounded-state dedup by +event time (see Scalability below). Example:: @@ -48,6 +50,45 @@ def poll(prefix): .with_poll_interval(Duration(seconds=5)) .with_termination_per_input(after_total_of(60))) +Watermark and event-time contract +--------------------------------- +Each emitted output carries the event time the poll function reported for it. +Raw (non-``TimestampedValue``) outputs default to processing time +(``Timestamp.now()`` at poll), so wrap outputs in ``TimestampedValue`` or pass +``timestamp=`` to :meth:`PollResult.incomplete`/:meth:`PollResult.complete` when +the data has a real event time. The watermark for an input is derived per round +as: (a) the explicit ``with_watermark`` if the poll supplies one; else (b) the +earliest event time of this round's new outputs; else (c) held unchanged when a +round yields nothing; and it is released to ``MAX_TIMESTAMP`` on +:meth:`PollResult.complete`. The watermark only ever advances. + +Policy (b) is safe only for poll functions that enumerate outputs in +non-decreasing event-time order. If a later round returns a brand-new output +whose event time is *below* the already-advanced watermark, that output is +emitted at its true (earlier) time and is therefore late: downstream event-time +windowing may drop it. Watch logs a throttled warning when this happens. For +out-of-order sources, have the poll supply an explicit watermark via +:meth:`PollResult.with_watermark` that bounds the earliest event time any future +output may have. Default dedup is by output value identity (a re-seen value +keeps its first-seen timestamp), and the output coder must be deterministic for +dedup to hold across workers and restarts. + +Scalability +----------- +Parallelism is per input element: each input is one restriction watched on one +worker, with no intra-key splitting, so throughput scales with the number of +input elements, not with a single input's growth. By default dedup is exact and +by value identity: the per-input state retains one hash per distinct output and +is not garbage-collected, so it grows with the number of distinct outputs an +input ever produces. For a long-lived, high-cardinality source whose outputs +carry strictly increasing event-time timestamps, call +:meth:`Watch.with_timestamp_cursor` to dedup by a high-water-mark timestamp +instead: the per-input state is a single timestamp that never grows and the poll +result is not hashed, so state and per-checkpoint encoding stay O(1) regardless +of how many outputs the input produces. See its docstring for the ordering +precondition. The poll function runs synchronously on the watch path, so keep it +bounded and timeout-safe. + This API is experimental and may change in backwards-incompatible ways. """ @@ -113,6 +154,11 @@ def is_complete(self) -> bool: @staticmethod def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]: + # The default timestamp is computed once per call (not per output) so all + # raw outputs in one poll share an event time and the inferred watermark is + # not jittered by wall-clock reads. ``timestamp=None`` means "use processing + # time"; pass ``timestamp=`` or wrap outputs in ``TimestampedValue`` to set + # a real event time. if timestamp is None: default_ts = Timestamp.now() else: @@ -130,7 +176,10 @@ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult': """Reports outputs and expects more; the transform infers the watermark. A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp`` - when given, else with the current processing time. + when given, else with the current processing time. With no explicit + watermark, the transform holds the watermark at the earliest event time of + this poll's new outputs, which is only safe for non-decreasing event-time + enumerations; out-of-order sources should call :meth:`with_watermark`. """ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None) @@ -139,12 +188,16 @@ def complete(outputs: Iterable, timestamp=None) -> 'PollResult': """Reports the final outputs for an input, after which polling stops. A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp`` - when given, else with the current processing time. + when given, else with the current processing time. The watermark is released + to ``MAX_TIMESTAMP`` so downstream event-time windows for this input close. """ return PollResult( PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP) def with_watermark(self, watermark) -> 'PollResult': + """Sets an explicit watermark: a promise that no future output for this + input will have an event time below ``watermark``. Use this for sources that + can surface outputs out of event-time order across poll rounds.""" return dataclasses.replace(self, watermark=Timestamp.of(watermark)) @@ -231,15 +284,17 @@ def after_total_of(duration) -> TerminationCondition: @dataclasses.dataclass(frozen=True) class _PollingGrowthState: - """Keep-polling state: emitted-output hashes, watermark, termination state. + """Keep-polling state: dedup state, watermark, termination state. - ``completed`` maps a 16-byte output hash to the event time it was first seen. - It is insertion-ordered and treated as immutable; the tracker builds a new - mapping for each residual. + In the default (hash) mode ``completed`` maps a 16-byte output hash to the + event time it was first seen; it is insertion-ordered and treated as + immutable. In timestamp-cursor mode ``completed`` is empty and ``cursor`` + holds the greatest event time emitted so far, so the state is O(1). """ completed: 'collections.OrderedDict[bytes, Timestamp]' poll_watermark: Optional[Timestamp] termination_state: Any + cursor: Optional[Timestamp] = None @dataclasses.dataclass(frozen=True) @@ -307,8 +362,9 @@ class _GrowthStateCoder(Coder): A ``(tag, payload)`` envelope selects the variant; the payload is a variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered - list of ``(hash, timestamp)`` pairs so insertion order survives a round trip. - This format is internal to the Python SDK. + list of ``(hash, timestamp)`` pairs so insertion order survives a round trip, + and the cursor is a nullable timestamp. This format is internal to the Python + SDK. """ def __init__(self, output_coder: Coder, termination: TerminationCondition): nullable_ts = NullableCoder(TimestampCoder()) @@ -318,6 +374,7 @@ def __init__(self, output_coder: Coder, termination: TerminationCondition): termination.state_coder(), nullable_ts, coders.ListCoder(TupleCoder([_HashCode128Coder(), TimestampCoder()])), + nullable_ts, ]) self._non_polling_coder = TupleCoder([ nullable_ts, @@ -329,7 +386,8 @@ def encode(self, state: _GrowthState) -> bytes: payload = self._polling_coder.encode(( state.termination_state, state.poll_watermark, - list(state.completed.items()))) + list(state.completed.items()), + state.cursor)) return self._envelope_coder.encode((0, payload)) payload = self._non_polling_coder.encode( (state.pending.watermark, list(state.pending.outputs))) @@ -338,10 +396,13 @@ def encode(self, state: _GrowthState) -> bytes: def decode(self, encoded: bytes) -> _GrowthState: tag, payload = self._envelope_coder.decode(encoded) if tag == 0: - termination_state, poll_watermark, items = self._polling_coder.decode( - payload) + termination_state, poll_watermark, items, cursor = ( + self._polling_coder.decode(payload)) return _PollingGrowthState( - collections.OrderedDict(items), poll_watermark, termination_state) + collections.OrderedDict(items), + poll_watermark, + termination_state, + cursor) if tag == 1: watermark, outputs = self._non_polling_coder.decode(payload) return _NonPollingGrowthState(PollResult(tuple(outputs), watermark)) @@ -374,12 +435,14 @@ def __init__( poll_fn: Callable[[Any], PollResult], key_coder: Coder, termination: TerminationCondition, - now_fn: Callable[[], float]): + now_fn: Callable[[], float], + cursor_mode: bool = False): self._restriction = restriction self._poll_fn = poll_fn self._key_coder = key_coder self._termination = termination self._now = now_fn + self._cursor_mode = cursor_mode self._should_stop = False self._primary = None # type: Optional[_GrowthState] self._residual = None # type: Optional[_GrowthState] @@ -409,16 +472,26 @@ def try_claim(self, holder: list) -> bool: now = Timestamp.of(self._now()) result = self._poll_fn(element) - new_outputs = [] # type: List[TimestampedValue] claimed = [] # type: List[Tuple[bytes, Timestamp]] - seen_this_round = set() # type: set - for output in result.outputs: - key_hash = self._hash_output(output.value) - if key_hash in restriction.completed or key_hash in seen_this_round: - continue - seen_this_round.add(key_hash) - new_outputs.append(output) - claimed.append((key_hash, output.timestamp)) + if self._cursor_mode: + # Dedup by a high-water-mark timestamp: keep only outputs strictly past + # the cursor, so the state is one timestamp and the poll is not hashed. + cursor = restriction.cursor + new_outputs = [ + output for output in result.outputs + if cursor is None or output.timestamp > cursor + ] + else: + # Dedup by value identity against the per-input hash set. + new_outputs = [] # type: List[TimestampedValue] + seen_this_round = set() # type: set + for output in result.outputs: + key_hash = self._hash_output(output.value) + if key_hash in restriction.completed or key_hash in seen_this_round: + continue + seen_this_round.add(key_hash) + new_outputs.append(output) + claimed.append((key_hash, output.timestamp)) new_outputs.sort(key=lambda output: output.timestamp) termination_state = restriction.termination_state @@ -446,14 +519,25 @@ def try_claim(self, holder: list) -> bool: # Terminal round: no polling work remains, so a checkpoint (runner- # initiated or via defer_remainder) resumes a state that emits nothing. self._residual = _NonPollingGrowthState(PollResult((), watermark)) + elif self._cursor_mode: + # new_outputs is sorted and all strictly above the old cursor, so the + # last one is the new high-water mark; an empty round leaves it unchanged. + new_cursor = restriction.cursor + if new_outputs: + new_cursor = new_outputs[-1].timestamp + residual_watermark = self._max_watermark( + restriction.poll_watermark, watermark) + self._residual = _PollingGrowthState( + collections.OrderedDict(), + residual_watermark, + termination_state, + new_cursor) else: - merged = collections.OrderedDict(restriction.completed) - for key_hash, first_seen in claimed: - merged[key_hash] = first_seen residual_watermark = self._max_watermark( restriction.poll_watermark, watermark) + completed = self._build_completed(restriction.completed, claimed) self._residual = _PollingGrowthState( - merged, residual_watermark, termination_state) + completed, residual_watermark, termination_state) holder[1] = ('poll', new_outputs, watermark, stop) self._should_stop = True return True @@ -467,6 +551,24 @@ def _max_watermark(left: Optional[Timestamp], return left return max(left, right) + @staticmethod + def _build_completed( + base: 'collections.OrderedDict[bytes, Timestamp]', + claimed: List[Tuple[bytes, Timestamp]], + ) -> 'collections.OrderedDict[bytes, Timestamp]': + """Builds the residual dedup map. + + Reuses the parent map untouched when a round adds nothing, so a steady-state + idle round avoids the O(N) copy; otherwise copies and appends the new + hashes. + """ + if not claimed: + return base + merged = collections.OrderedDict(base) + for key_hash, first_seen in claimed: + merged[key_hash] = first_seen + return merged + def try_split(self, fraction_of_remainder): # Only self-checkpoint (fraction 0) is supported; decline dynamic splits. if fraction_of_remainder != 0: @@ -521,14 +623,18 @@ def __init__( termination: TerminationCondition, poll_interval: Duration, output_coder: Coder, - now_fn: Optional[Callable[[], float]] = None): + now_fn: Optional[Callable[[], float]] = None, + cursor_mode: bool = False): self._poll_fn = poll_fn self._termination = termination self._poll_interval = poll_interval self._output_coder = output_coder self._key_coder = output_coder self._now = now_fn or time.time + self._cursor_mode = cursor_mode self._restriction_coder = _GrowthStateCoder(output_coder, termination) + # Count of late emissions seen on this worker, for throttled warnings. + self._late_count = 0 def initial_restriction(self, element) -> _PollingGrowthState: now = Timestamp.of(self._now()) @@ -543,7 +649,8 @@ def create_tracker(self, restriction) -> _GrowthRestrictionTracker: self._poll_fn, self._key_coder, self._termination, - self._now) + self._now, + self._cursor_mode) def split(self, element, restriction): # Watch fans out by input element, so each restriction stays whole. @@ -582,7 +689,14 @@ def process( yield TimestampedValue((element, output.value), output.timestamp) return new_outputs, watermark, stop = work[1], work[2], work[3] + # Emit outputs first, then advance the watermark (emit-then-advance), so a + # round that reports both outputs and a higher watermark cannot push the + # watermark past those outputs' event times. An output behind the current + # watermark is already late; warn rather than silently dropping it. + current_watermark = watermark_estimator.current_watermark() for output in new_outputs: + if current_watermark is not None and output.timestamp < current_watermark: + self._warn_late(element, output.timestamp, current_watermark) yield TimestampedValue((element, output.value), output.timestamp) if stop: # The input is finished, so release the watermark hold to MAX. @@ -592,6 +706,23 @@ def process( _set_watermark_if_greater(watermark_estimator, watermark) tracker.defer_remainder(self._poll_interval) + def _warn_late(self, element, output_timestamp, watermark) -> None: + # Throttle to powers of two so an ongoing problem stays visible without + # flooding worker logs. + self._late_count += 1 + if self._late_count & (self._late_count - 1) == 0: + _LOGGER.warning( + 'Watch emitted an output for input %r at event time %s, behind the ' + 'current watermark %s, so downstream event-time windowing may drop ' + 'it as late. This happens when the poll returns outputs out of ' + 'event-time order across rounds without an explicit watermark; call ' + 'PollResult.with_watermark(...) for out-of-order sources. ' + '(%d late emissions so far on this worker)', + element, + output_timestamp, + watermark, + self._late_count) + def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None: # set_watermark raises on regression, so only ever advance the watermark. @@ -617,13 +748,15 @@ def __init__( termination: Optional[TerminationCondition] = None, poll_interval: Optional[Duration] = None, output_coder: Optional[Coder] = None, - now_fn: Optional[Callable[[], float]] = None): + now_fn: Optional[Callable[[], float]] = None, + cursor_mode: bool = False): super().__init__() self._poll_fn = poll_fn self._termination = termination or never() self._poll_interval = poll_interval self._output_coder = output_coder self._now = now_fn + self._cursor_mode = cursor_mode @classmethod def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch': @@ -635,7 +768,8 @@ def _replace(self, **changes) -> 'Watch': termination=self._termination, poll_interval=self._poll_interval, output_coder=self._output_coder, - now_fn=self._now) + now_fn=self._now, + cursor_mode=self._cursor_mode) spec.update(changes) return Watch(**spec) @@ -649,6 +783,27 @@ def with_termination_per_input( def with_output_coder(self, output_coder: Coder) -> 'Watch': return self._replace(output_coder=output_coder) + def with_timestamp_cursor(self) -> 'Watch': + """Dedups by a high-water-mark timestamp instead of by value identity, so + the per-input state is a single timestamp that never grows. + + In this mode Watch keeps only the greatest event time it has emitted for an + input (the cursor) and, each round, emits exactly the polled outputs whose + event time is greater than the cursor, then advances the cursor to the new + maximum. No hash set is kept and the poll result is not hashed, so the + per-input state and the per-checkpoint encoding are O(1) regardless of how + many outputs the input produces. + + PRECONDITION: each distinct output must carry an event time strictly greater + than every output already emitted for that input (a monotonic cursor, e.g. + an ``updated_at`` column or an increasing file modification time). An output + whose event time is at or below the cursor is treated as already seen and + skipped, so out-of-order or equal-timestamp outputs are dropped. For sources + that re-list arbitrary collections or surface outputs out of event-time + order, use the default exact (hash) dedup instead. + """ + return self._replace(cursor_mode=True) + def expand(self, pcoll): if self._poll_interval is None: raise ValueError('Watch requires with_poll_interval(...)') @@ -657,11 +812,11 @@ def expand(self, pcoll): hint = self._poll_fn.default_output_coder() if isinstance( self._poll_fn, PollFn) else None output_coder = hint or coders.PickleCoder() - if not output_coder.is_deterministic(): + if not self._cursor_mode and not output_coder.is_deterministic(): _LOGGER.warning( 'Watch dedup uses a non-deterministic output coder (%s); equal ' 'outputs may be emitted more than once. Pass a deterministic coder ' - 'via with_output_coder() for reliable dedup.', + 'via with_output_coder() or switch to with_timestamp_cursor().', type(output_coder).__name__) return pcoll | core.ParDo( _WatchGrowthDoFn( @@ -669,7 +824,8 @@ def expand(self, pcoll): self._termination, self._poll_interval, output_coder, - self._now)) + self._now, + self._cursor_mode)) def _as_duration(value) -> Duration: diff --git a/sdks/python/apache_beam/io/watch_test.py b/sdks/python/apache_beam/io/watch_test.py index 07f2eadfbd91..def9d08b8856 100644 --- a/sdks/python/apache_beam/io/watch_test.py +++ b/sdks/python/apache_beam/io/watch_test.py @@ -60,6 +60,43 @@ def _initial_polling(termination=None, now=Timestamp(0)): collections.OrderedDict(), None, termination.for_new_input(now, 'input')) +def _cursor_tracker(restriction, poll_fn, now=0.0): + return _GrowthRestrictionTracker( + restriction, + poll_fn, + StrUtf8Coder(), + never(), lambda: now, + cursor_mode=True) + + +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)) @@ -75,6 +112,18 @@ def test_polling_round_trip_preserves_resume_state(self): 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(self): + coder = _GrowthStateCoder(StrUtf8Coder(), never()) + state = _PollingGrowthState( + collections.OrderedDict(), + 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(0, len(decoded.completed)) def test_non_polling_round_trip_preserves_pending_outputs(self): coder = _GrowthStateCoder(StrUtf8Coder(), never()) @@ -217,6 +266,143 @@ def poll(unused_element): self.assertIsInstance(residual, _PollingGrowthState) self.assertEqual(2, len(residual.completed)) + def test_out_of_order_new_output_infers_its_earlier_event_time(self): + # Round 1 surfaces a@10; round 2 surfaces a brand-new b@5 (earlier). The + # round-2 inferred watermark is b's own time; the process() estimator's + # monotonic guard is what leaves b late. This locks in the + # reference-consistent behavior the design documents and warns about. + polls = [] + + def poll(unused_element): + polls.append(len(polls)) + if len(polls) == 1: + return PollResult.incomplete([_ts('a', 10)]) + return PollResult.incomplete([_ts('b', 5)]) + + tracker = _new_tracker(_initial_polling(), poll) + holder = ['input', None] + tracker.try_claim(holder) + self.assertEqual(Timestamp(10), holder[1][2]) + _, residual = tracker.try_split(0) + + resumed = _new_tracker(residual, poll) + holder = ['input', None] + self.assertTrue(resumed.try_claim(holder)) + self.assertEqual(['b'], [o.value for o in holder[1][1]]) + self.assertEqual(Timestamp(5), holder[1][2]) + + 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) + + tracker = _new_tracker(_initial_polling(), poll) + holder = ['input', None] + tracker.try_claim(holder) + self.assertEqual(Timestamp(0), holder[1][2]) + _, residual = tracker.try_split(0) + self.assertEqual(Timestamp(0), residual.poll_watermark) + + 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). + polls = [] + + def poll(unused_element): + polls.append(len(polls)) + if len(polls) == 1: + return PollResult.incomplete([_ts('a', 1)]) + return PollResult.incomplete([]) + + tracker = _new_tracker(_initial_polling(), poll) + holder = ['input', None] + tracker.try_claim(holder) + _, residual1 = tracker.try_split(0) + + resumed = _new_tracker(residual1, poll) + holder = ['input', None] + resumed.try_claim(holder) + _, residual2 = resumed.try_split(0) + self.assertIs(residual1.completed, residual2.completed) + + def test_cursor_keeps_state_o1_and_tracks_high_water_mark(self): + # In cursor mode the dedup set stays empty and only the greatest emitted + # event time is retained, so the per-input state is O(1). + def poll(unused_element): + return PollResult.incomplete([_ts('a', 1), _ts('b', 2), _ts('c', 3)]) + + tracker = _cursor_tracker(_initial_polling(), poll) + holder = ['input', None] + tracker.try_claim(holder) + self.assertEqual(['a', 'b', 'c'], [o.value for o in holder[1][1]]) + _, residual = tracker.try_split(0) + self.assertIsInstance(residual, _PollingGrowthState) + self.assertEqual(0, len(residual.completed)) # no hash set + self.assertEqual(Timestamp(3), residual.cursor) # high-water mark + + def test_cursor_emits_only_outputs_after_the_cursor(self): + # A later round emits only outputs strictly past the cursor; a re-listed + # output (== cursor) and an earlier output (< cursor) are both dropped. + polls = [] + + def poll(unused_element): + polls.append(len(polls)) + if len(polls) == 1: + return PollResult.incomplete([_ts('a', 10)]) + return PollResult.incomplete( + [_ts('early', 5), _ts('a', 10), _ts('c', 20)]) + + tracker = _cursor_tracker(_initial_polling(), poll) + holder = ['input', None] + tracker.try_claim(holder) + _, residual = tracker.try_split(0) + self.assertEqual(Timestamp(10), residual.cursor) + + resumed = _cursor_tracker(residual, poll) + holder = ['input', None] + self.assertTrue(resumed.try_claim(holder)) + self.assertEqual(['c'], [o.value for o in holder[1][1]]) # only 20 > 10 + _, residual = resumed.try_split(0) + self.assertEqual(Timestamp(20), residual.cursor) + + def test_cursor_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 state never accumulates a hash set. + def poll_for(round_index): + def poll(unused_element): + return PollResult.incomplete( + [_ts('f%d' % i, i + 1) for i in range(round_index + 1)]) + + return poll + + state = _initial_polling() + emitted = collections.Counter() + for round_index in range(10): + tracker = _cursor_tracker(state, poll_for(round_index)) + holder = ['input', None] + tracker.try_claim(holder) + for output in holder[1][1]: + emitted[output.value] += 1 + _, state = tracker.try_split(0) + self.assertEqual(0, len(state.completed)) # O(1) throughout + self.assertEqual([1] * 10, [emitted['f%d' % i] for i in range(10)]) + self.assertEqual(Timestamp(10), state.cursor) + + def test_cursor_complete_stops_and_keeps_o1_state(self): + def poll(unused_element): + return PollResult.complete([_ts('a', 1), _ts('b', 2)]) + + tracker = _cursor_tracker(_initial_polling(), poll) + holder = ['input', None] + tracker.try_claim(holder) + self.assertEqual(['a', 'b'], [o.value for o in holder[1][1]]) + self.assertTrue(holder[1][3]) # stop + _, residual = tracker.try_split(0) + self.assertIsInstance(residual, _NonPollingGrowthState) + self.assertEqual((), residual.pending.outputs) + # Module-level so the poll function pickles by reference; the call counter is # shared within the single in-memory DirectRunner process. @@ -236,10 +422,27 @@ def _complete_poll(prefix): return PollResult.complete([_ts(prefix + 'a', 1), _ts(prefix + 'b', 2)]) +def _out_of_order_poll(prefix): + # Round 1 emits a 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 + count = _POLL_CALLS[prefix] + if count == 1: + return PollResult.incomplete([_ts(prefix + 'late_after', 10)]) + return PollResult.complete([_ts(prefix + 'early', 5)]) + + def _windowed_group(kv, window=beam.DoFn.WindowParam): return ((window.start, window.end), sorted(kv[1])) +class WatchBuilderTest(unittest.TestCase): + def test_with_timestamp_cursor_enables_cursor_mode(self): + watch = Watch.growth_of(_complete_poll) + self.assertFalse(watch._cursor_mode) + self.assertTrue(watch.with_timestamp_cursor()._cursor_mode) + + class WatchEndToEndTest(unittest.TestCase): def _in_memory_pipeline(self): return TestPipeline( @@ -288,6 +491,35 @@ def test_multi_round_dedups_stops_and_is_per_input(self): self.assertEqual(3, _POLL_CALLS['x:']) self.assertEqual(3, _POLL_CALLS['y:']) + def test_cursor_mode_dedups_growing_source(self): + _POLL_CALLS.clear() + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['x:', 'y:']) + | Watch.growth_of(_growing_poll).with_poll_interval( + Duration(0.05)).with_timestamp_cursor()) + # Each output emitted exactly once via the high-water-mark cursor, with no + # hash set kept, 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_out_of_order_across_rounds_warns_about_late_output(self): + _POLL_CALLS.clear() + with self.assertLogs(level='WARNING') as logs: + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['k:']) + | Watch.growth_of(_out_of_order_poll).with_poll_interval( + Duration(0.05))) + # Watch emits every new output; the lateness is a downstream concern. + assert_that( + output, equal_to([('k:', 'k:late_after'), ('k:', 'k:early')])) + self.assertTrue( + any('behind the current watermark' in line for line in logs.output), + 'expected a late-emission warning, got: %s' % logs.output) + if __name__ == '__main__': unittest.main()