From 1eb92846a9fadee630afa06391389ddd11442cec Mon Sep 17 00:00:00 2001 From: Eliaazzz Date: Fri, 19 Jun 2026 00:55:06 +1000 Subject: [PATCH 1/3] [Python] Add Watch transform with growth_of polling SDF Add an experimental Watch transform that watches a growing set of outputs per input element. Watch.growth_of(poll_fn) runs a periodic poll loop as a splittable DoFn and emits an unbounded PCollection of (input, output) pairs. Each process() performs one poll round and self-checkpoints via defer_remainder. New outputs are deduplicated with a stable 128-bit blake2b hash of the encoded output, and a manual watermark estimator advances per poll. Per-input termination supports never() and after_total_of(); polling also stops when a poll returns PollResult.complete(). The DoFn is its own RestrictionProvider, and restriction state serializes through a tagged GrowthState coder. Tests cover termination conditions, coder round-trips, the restriction tracker claim/checkpoint/dedup logic, and DirectRunner end-to-end runs. --- sdks/python/apache_beam/io/watch.py | 691 +++++++++++++++++++++++ sdks/python/apache_beam/io/watch_test.py | 307 ++++++++++ 2 files changed, 998 insertions(+) create mode 100644 sdks/python/apache_beam/io/watch.py create mode 100644 sdks/python/apache_beam/io/watch_test.py diff --git a/sdks/python/apache_beam/io/watch.py b/sdks/python/apache_beam/io/watch.py new file mode 100644 index 000000000000..68eef1596290 --- /dev/null +++ b/sdks/python/apache_beam/io/watch.py @@ -0,0 +1,691 @@ +# +# 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. +# + +"""Experimental ``Watch`` transform for the Python SDK. + +``Watch`` continuously watches a growing set of outputs for each input element, +calling a user poll function on an interval until a per-input termination +condition fires. It is the engine behind periodic file-discovery and any +periodic polling source. + +For every input element the transform runs an independent loop:: + + poll -> keep never-seen-before outputs -> emit them (timestamped) -> + 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. + +Example:: + + from apache_beam.io.watch import Watch, PollResult, after_total_of + from apache_beam.transforms.window import TimestampedValue + from apache_beam.utils.timestamp import Duration, Timestamp + + def poll(prefix): + now = Timestamp.now() + outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)] + return PollResult.complete(outputs) + + watched = (inputs + | Watch.growth_of(poll) + .with_poll_interval(Duration(seconds=5)) + .with_termination_per_input(after_total_of(60))) + +This API is experimental and may change in backwards-incompatible ways. +""" + +import collections +import dataclasses +import hashlib +import time +from typing import Any +from typing import Callable +from typing import Iterable +from typing import List +from typing import Optional +from typing import Tuple + +from apache_beam import coders +from apache_beam.coders.coders import Coder +from apache_beam.coders.coders import NullableCoder +from apache_beam.coders.coders import TimestampCoder +from apache_beam.coders.coders import TupleCoder +from apache_beam.io import iobase +from apache_beam.io.watermark_estimators import ManualWatermarkEstimator +from apache_beam.runners import sdf_utils +from apache_beam.transforms import PTransform +from apache_beam.transforms import core +from apache_beam.transforms.window import TimestampedValue +from apache_beam.utils.timestamp import MAX_TIMESTAMP +from apache_beam.utils.timestamp import Duration +from apache_beam.utils.timestamp import Timestamp + +__all__ = [ + 'Watch', + 'PollResult', + 'PollFn', + 'TerminationCondition', + 'never', + 'after_total_of', +] + +_HASH_DIGEST_SIZE = 16 # 128-bit digest width. + +# ------------------------------------------------------------------------------ +# Public API. +# ------------------------------------------------------------------------------ + + +@dataclasses.dataclass(frozen=True) +class PollResult: + """Outputs produced by one poll, plus an optional explicit watermark. + + ``watermark`` of ``None`` lets the transform infer the watermark from the + earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by + :meth:`complete`) marks the input finished, so polling stops. + """ + outputs: Tuple[TimestampedValue, ...] + watermark: Optional[Timestamp] = None + + @property + def is_complete(self) -> bool: + return self.watermark == MAX_TIMESTAMP + + @staticmethod + def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]: + if timestamp is None: + default_ts = Timestamp.now() + else: + default_ts = Timestamp.of(timestamp) + normalized = [] + for output in outputs: + if isinstance(output, TimestampedValue): + normalized.append(output) + else: + normalized.append(TimestampedValue(output, default_ts)) + return tuple(normalized) + + @staticmethod + 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. + """ + return PollResult(PollResult._normalize(outputs, timestamp), watermark=None) + + @staticmethod + 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. + """ + return PollResult( + PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP) + + def with_watermark(self, watermark) -> 'PollResult': + return dataclasses.replace(self, watermark=Timestamp.of(watermark)) + + +class PollFn(object): + """Optional base for a poll function ``input -> PollResult``. + + Any callable with that signature works; subclass only to attach an output + coder hint via :meth:`default_output_coder`. The hint is used for dedup and + must be deterministic:: + + from apache_beam import coders + + class ListFiles(PollFn): + def __call__(self, prefix): + return PollResult.incomplete(list_files(prefix)) + + def default_output_coder(self): + return coders.StrUtf8Coder() + """ + def __call__(self, element: Any) -> PollResult: + raise NotImplementedError + + def default_output_coder(self) -> Optional[Coder]: + return None + + +class TerminationCondition(object): + """Per-input stop policy with immutable, encodable state. + + Hooks follow the lifecycle of one input's polling loop. ``state`` flows from + :meth:`for_new_input` through the per-round hooks and is serialized with + :meth:`state_coder`. + """ + def for_new_input(self, now: Timestamp, element: Any) -> Any: + raise NotImplementedError + + def on_seen_new_output(self, now: Timestamp, state: Any) -> Any: + return state + + def on_poll_complete(self, state: Any) -> Any: + return state + + def can_stop_polling(self, now: Timestamp, state: Any) -> bool: + raise NotImplementedError + + def state_coder(self) -> Coder: + raise NotImplementedError + + +class _Never(TerminationCondition): + """Polls until the poll function returns :meth:`PollResult.complete`.""" + def for_new_input(self, now, element): + return 0 + + def can_stop_polling(self, now, state): + return False + + def state_coder(self): + return coders.VarIntCoder() + + +class _AfterTotalOf(TerminationCondition): + """Stops once the wall-clock time since the input was first seen exceeds a + fixed duration.""" + def __init__(self, duration: Duration): + self._duration_micros = duration.micros + + def for_new_input(self, now, element): + return (now, self._duration_micros) + + def can_stop_polling(self, now, state): + start, duration_micros = state + return (now - start).micros > duration_micros + + def state_coder(self): + return TupleCoder([TimestampCoder(), coders.VarIntCoder()]) + + +def never() -> TerminationCondition: + """Polls until :meth:`PollResult.complete`.""" + return _Never() + + +def after_total_of(duration) -> TerminationCondition: + """Stops polling an input after ``duration`` (a :class:`Duration` or seconds) + has elapsed since it was first seen.""" + return _AfterTotalOf(_as_duration(duration)) + + +# ------------------------------------------------------------------------------ +# Restriction state. +# ------------------------------------------------------------------------------ + + +class _GrowthState: + """Base for the two restriction variants a Watch input can hold.""" + + +@dataclasses.dataclass(frozen=True) +class _PollingGrowthState(_GrowthState): + """Keep-polling state: emitted-output hashes, 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; a new mapping is built for + each residual. + """ + completed: 'collections.OrderedDict[bytes, Timestamp]' + poll_watermark: Optional[Timestamp] + termination_state: Any + + +@dataclasses.dataclass(frozen=True) +class _NonPollingGrowthState(_GrowthState): + """Replay-then-stop state: the outputs already emitted this round. + + Produced as the checkpoint primary so a bundle retry re-emits exactly those + outputs. + """ + pending: PollResult + + +# ------------------------------------------------------------------------------ +# Coders. +# ------------------------------------------------------------------------------ + + +class _TimestampedValueCoder(Coder): + """Coder for :class:`TimestampedValue`. + + ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the + wire, so the SDK ships no standalone coder for it. Watch keeps it inside the + restriction state, so this encodes the ``(value, timestamp)`` pair with a + :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode. + """ + def __init__(self, value_coder: Coder): + self._tuple_coder = TupleCoder([value_coder, TimestampCoder()]) + + def encode(self, value: TimestampedValue) -> bytes: + return self._tuple_coder.encode((value.value, value.timestamp)) + + def decode(self, encoded: bytes) -> TimestampedValue: + value, timestamp = self._tuple_coder.decode(encoded) + return TimestampedValue(value, timestamp) + + def is_deterministic(self) -> bool: + return self._tuple_coder.is_deterministic() + + +class _GrowthStateCoder(Coder): + """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`. + + 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. + """ + def __init__(self, output_coder: Coder, termination: TerminationCondition): + nullable_ts = NullableCoder(TimestampCoder()) + self._envelope_coder = TupleCoder( + [coders.VarIntCoder(), coders.BytesCoder()]) + self._polling_coder = TupleCoder([ + termination.state_coder(), + nullable_ts, + coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])), + ]) + self._non_polling_coder = TupleCoder([ + nullable_ts, + coders.ListCoder(_TimestampedValueCoder(output_coder)), + ]) + + def encode(self, state: _GrowthState) -> bytes: + if isinstance(state, _PollingGrowthState): + payload = self._polling_coder.encode(( + state.termination_state, + state.poll_watermark, + list(state.completed.items()))) + return self._envelope_coder.encode((0, payload)) + payload = self._non_polling_coder.encode( + (state.pending.watermark, list(state.pending.outputs))) + return self._envelope_coder.encode((1, payload)) + + 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) + return _PollingGrowthState( + collections.OrderedDict(items), poll_watermark, termination_state) + if tag == 1: + watermark, outputs = self._non_polling_coder.decode(payload) + return _NonPollingGrowthState(PollResult(tuple(outputs), watermark)) + raise ValueError('unknown Watch growth state tag: %r' % (tag, )) + + def is_deterministic(self) -> bool: + return False + + +# ------------------------------------------------------------------------------ +# Restriction tracker. +# ------------------------------------------------------------------------------ + + +@dataclasses.dataclass(frozen=True) +class _PollPlan: + """One planned poll round: what to emit plus the self-checkpoint split. + + ``process()`` builds this from a poll result before claiming, so the poll runs + outside the tracker lock. ``primary`` is the round just processed and + ``residual`` is the work a checkpoint resumes. + """ + emit: Tuple[TimestampedValue, ...] + watermark: Optional[Timestamp] + stop: bool + primary: _GrowthState + residual: _GrowthState + + +def _hash_output(key_coder: Coder, value: Any) -> bytes: + return hashlib.blake2b( + key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest() + + +def _max_watermark(left: Optional[Timestamp], + right: Optional[Timestamp]) -> Optional[Timestamp]: + if left is None: + return right + if right is None: + return left + return max(left, right) + + +def _plan_poll_round( + restriction: _PollingGrowthState, + result: PollResult, + termination: TerminationCondition, + key_coder: Coder, + now: Timestamp) -> _PollPlan: + """Dedups a poll result into new outputs and builds the checkpoint split. + + Pure and side-effect free, so ``process()`` can call it after the poll and + before claiming, keeping the poll off the tracker lock. + """ + new_outputs = [] # type: List[TimestampedValue] + claimed = [] # type: List[Tuple[bytes, Timestamp]] + seen_this_round = set() # type: set + for output in result.outputs: + key_hash = _hash_output(key_coder, 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 + if new_outputs: + termination_state = termination.on_seen_new_output(now, termination_state) + termination_state = termination.on_poll_complete(termination_state) + + if result.watermark is not None: + watermark = result.watermark + elif new_outputs: + watermark = new_outputs[0].timestamp + else: + watermark = None + + # A watermark at MAX means no more output is possible, so polling stops. + reached_max = watermark is not None and watermark >= MAX_TIMESTAMP + stop = ( + result.is_complete or reached_max or + termination.can_stop_polling(now, termination_state)) + + primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark)) + if stop: + # Terminal round: no polling work remains, so a checkpoint (runner-initiated + # or via defer_remainder) resumes a state that emits nothing. + residual = _NonPollingGrowthState(PollResult((), watermark)) + else: + merged = collections.OrderedDict(restriction.completed) + for key_hash, first_seen in claimed: + merged[key_hash] = first_seen + residual_watermark = _max_watermark(restriction.poll_watermark, watermark) + residual = _PollingGrowthState( + merged, residual_watermark, termination_state) + return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual) + + +def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan: + """Plans a replay of the outputs already emitted this round, then stops.""" + resumed = _NonPollingGrowthState( + PollResult((), restriction.pending.watermark)) + return _PollPlan( + restriction.pending.outputs, None, True, restriction, resumed) + + +class _GrowthRestrictionTracker(iobase.RestrictionTracker): + """Tracks one input's polling restriction and its self-checkpoints. + + ``process()`` runs the poll and plans the round with + :func:`_plan_poll_round`, then claims the plan; the tracker only records the + plan's primary and residual so a checkpoint (via ``defer_remainder`` or the + runner) resumes correctly. Keeping the poll in ``process()`` means a slow + poll never holds the tracker lock, so it cannot delay runner progress checks + or checkpoints. + """ + def __init__(self, restriction: _GrowthState): + self._restriction = restriction + self._should_stop = False + self._primary = None # type: Optional[_GrowthState] + self._residual = None # type: Optional[_GrowthState] + + def current_restriction(self) -> _GrowthState: + return self._restriction + + def try_claim(self, plan: _PollPlan) -> bool: + """Records a planned round's checkpoint split; one claim per ``process()``. + + Returns ``False`` only when a checkpoint already stopped this invocation, in + which case ``process()`` must emit nothing. + """ + if self._should_stop: + return False + self._primary = plan.primary + self._residual = plan.residual + self._should_stop = True + return True + + def try_split(self, fraction_of_remainder): + # Only self-checkpoint (fraction 0) is supported; decline dynamic splits. + if fraction_of_remainder != 0: + return None + if self._primary is None: + # No claim happened this invocation: keep the whole state as the residual. + primary = _NonPollingGrowthState(PollResult((), None)) + residual = self._restriction + self._restriction = primary + self._should_stop = True + return primary, residual + primary, residual = self._primary, self._residual + self._restriction = primary + self._should_stop = True + return primary, residual + + def check_done(self) -> bool: + # Called after every process(); the single claim or a split sets the flag. + if self._should_stop: + return True + raise ValueError( + 'Watch restriction was neither claimed nor checkpointed: %r' % + (self._restriction, )) + + def current_progress(self) -> 'iobase.RestrictionProgress': + if self._should_stop: + return iobase.RestrictionProgress(completed=1.0, remaining=0.0) + return iobase.RestrictionProgress(completed=0.0, remaining=1.0) + + def is_bounded(self) -> bool: + # A polling restriction is unbounded; a replay-then-stop one is bounded. + return isinstance(self._restriction, _NonPollingGrowthState) + + +# ------------------------------------------------------------------------------ +# Splittable DoFn (its own restriction provider). +# ------------------------------------------------------------------------------ + + +class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider): + """Polling SDF that emits ``(input, output)`` pairs. + + The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no + argument resolves the provider to the DoFn instance, so the provider methods + read the transform-level spec (poll function, coders, termination) off + ``self``. Provider methods run on a separately deserialized copy and before + ``setup()``, so the spec is immutable state set in ``__init__``. + """ + def __init__( + self, + poll_fn: Callable[[Any], PollResult], + termination: TerminationCondition, + poll_interval: Duration, + output_coder: Coder, + now_fn: Optional[Callable[[], float]] = None): + 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._restriction_coder = _GrowthStateCoder(output_coder, termination) + + def initial_restriction(self, element) -> _PollingGrowthState: + now = Timestamp.of(self._now()) + return _PollingGrowthState( + collections.OrderedDict(), + None, + self._termination.for_new_input(now, element)) + + def create_tracker(self, restriction) -> _GrowthRestrictionTracker: + return _GrowthRestrictionTracker(restriction) + + def split(self, element, restriction): + # Watch fans out by input element, so each restriction stays whole. + yield restriction + + def restriction_coder(self) -> Coder: + return self._restriction_coder + + def restriction_size(self, element, restriction) -> int: + return 1 + + def truncate(self, element, restriction): + # On drain, replay a pending NonPolling state and stop further polling. + if isinstance(restriction, _NonPollingGrowthState): + return restriction + return None + + @core.DoFn.unbounded_per_element() + def process( + self, + element, + timestamp=core.DoFn.TimestampParam, + tracker=core.DoFn.RestrictionParam(), + watermark_estimator=core.DoFn.WatermarkEstimatorParam( + ManualWatermarkEstimator.default_provider())): + assert isinstance(tracker, sdf_utils.RestrictionTrackerView) + restriction = tracker.current_restriction() + if isinstance(restriction, _NonPollingGrowthState): + # Replay the outputs already emitted this round, then stop. No poll. + if not tracker.try_claim(_replay_plan(restriction)): + return + _set_watermark_if_greater(watermark_estimator, timestamp) + for output in restriction.pending.outputs: + yield TimestampedValue((element, output.value), output.timestamp) + return + # Poll before claiming so a slow poll never holds the tracker lock, which + # would block runner progress checks and checkpoints. + now = Timestamp.of(self._now()) + result = self._poll_fn(element) + plan = _plan_poll_round( + restriction, result, self._termination, self._key_coder, now) + if not tracker.try_claim(plan): + # A checkpoint already stopped this invocation; emit nothing. + return + # Seed the watermark hold from the input event time after the claim. + _set_watermark_if_greater(watermark_estimator, timestamp) + for output in plan.emit: + yield TimestampedValue((element, output.value), output.timestamp) + if plan.stop: + # The input is finished, so release the watermark hold to MAX. + _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP) + return + if plan.watermark is not None: + _set_watermark_if_greater(watermark_estimator, plan.watermark) + tracker.defer_remainder(self._poll_interval) + + +def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None: + # set_watermark raises on regression, so only ever advance the watermark. + current = watermark_estimator.current_watermark() + if current is None or new_watermark > current: + watermark_estimator.set_watermark(new_watermark) + + +# ------------------------------------------------------------------------------ +# Public PTransform. +# ------------------------------------------------------------------------------ + + +class Watch(PTransform): + """Watches a growing set of outputs per input via a periodic poll function. + + Build with :meth:`growth_of` and the ``with_*`` methods. The output is an + unbounded ``PCollection`` of ``(input, output)`` pairs. + """ + def __init__( + self, + poll_fn: Callable[[Any], PollResult], + termination: Optional[TerminationCondition] = None, + poll_interval: Optional[Duration] = None, + output_coder: Optional[Coder] = None, + now_fn: Optional[Callable[[], float]] = None): + 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 + + @classmethod + def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch': + return cls(poll_fn) + + def _replace(self, **changes) -> 'Watch': + spec = dict( + poll_fn=self._poll_fn, + termination=self._termination, + poll_interval=self._poll_interval, + output_coder=self._output_coder, + now_fn=self._now) + spec.update(changes) + return Watch(**spec) + + def with_poll_interval(self, poll_interval) -> 'Watch': + return self._replace(poll_interval=_as_duration(poll_interval)) + + def with_termination_per_input( + self, termination: TerminationCondition) -> 'Watch': + return self._replace(termination=termination) + + def with_output_coder(self, output_coder: Coder) -> 'Watch': + return self._replace(output_coder=output_coder) + + def expand(self, pcoll): + if self._poll_interval is None: + raise ValueError('Watch requires with_poll_interval(...)') + output_coder = self._output_coder + if output_coder is None: + hint = self._poll_fn.default_output_coder() if isinstance( + self._poll_fn, PollFn) else None + output_coder = hint or coders.PickleCoder() + # Dedup hashes the encoded output as its key, so a non-deterministic coder + # would hash equal outputs differently and re-emit them, defeating the + # transform. Require determinism rather than silently emitting duplicates. + if not output_coder.is_deterministic(): + raise ValueError( + 'Watch dedup requires a deterministic output coder, but %s is not ' + 'deterministic. Pass one via with_output_coder().' % + type(output_coder).__name__) + # Type the (input, output) pairs from the input type and the resolved + # coder's type, so downstream transforms are typed and coder inference does + # not fall back to pickling. + input_type = pcoll.element_type or Any + try: + value_type = output_coder.to_type_hint() + except NotImplementedError: + value_type = Any + return pcoll | core.ParDo( + _WatchGrowthDoFn( + self._poll_fn, + self._termination, + self._poll_interval, + output_coder, + self._now)).with_output_types(Tuple[input_type, value_type]) + + +def _as_duration(value) -> Duration: + return value if isinstance(value, Duration) else Duration(value) diff --git a/sdks/python/apache_beam/io/watch_test.py b/sdks/python/apache_beam/io/watch_test.py new file mode 100644 index 000000000000..e530a42df7e2 --- /dev/null +++ b/sdks/python/apache_beam/io/watch_test.py @@ -0,0 +1,307 @@ +# +# 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 unittest + +import apache_beam as beam +from apache_beam.coders.coders import StrUtf8Coder +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 _NonPollingGrowthState +from apache_beam.io.watch import _plan_poll_round +from apache_beam.io.watch import _PollingGrowthState +from apache_beam.io.watch import _replay_plan +from apache_beam.io.watch import after_total_of +from apache_beam.io.watch import never +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.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 Duration +from apache_beam.utils.timestamp import Timestamp + + +def _ts(value, timestamp): + return TimestampedValue(value, Timestamp(timestamp)) + + +def _plan(restriction, poll_fn, now=0.0, termination=None): + termination = termination or never() + return _plan_poll_round( + restriction, + poll_fn('input'), + termination, + StrUtf8Coder(), + Timestamp.of(now)) + + +def _initial_polling(termination=None, now=Timestamp(0)): + termination = termination or never() + return _PollingGrowthState( + collections.OrderedDict(), None, termination.for_new_input(now, 'input')) + + +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) + + 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 GrowthTrackerTest(unittest.TestCase): + def test_poll_claims_dedups_and_checkpoints(self): + def poll(unused_element): + return PollResult.incomplete([_ts('a', 1), _ts('a', 1), _ts('b', 2)]) + + plan = _plan(_initial_polling(), poll) + self.assertEqual(['a', 'b'], [o.value for o in plan.emit]) + self.assertFalse(plan.stop) + self.assertEqual(Timestamp(1), plan.watermark) + self.assertIsInstance(plan.primary, _NonPollingGrowthState) + self.assertIsInstance(plan.residual, _PollingGrowthState) + self.assertEqual(2, len(plan.residual.completed)) + + tracker = _GrowthRestrictionTracker(_initial_polling()) + self.assertFalse(tracker.is_bounded()) + self.assertTrue(tracker.try_claim(plan)) + primary, residual = tracker.try_split(0) + self.assertIs(plan.primary, primary) + self.assertIs(plan.residual, residual) + self.assertTrue(tracker.check_done()) + + def explicit_watermark_poll(unused_element): + return PollResult.incomplete([_ts('c', 3)]).with_watermark(5) + + plan = _plan(_initial_polling(), explicit_watermark_poll) + self.assertEqual(Timestamp(5), plan.watermark) + self.assertEqual(Timestamp(5), plan.residual.poll_watermark) + + def test_second_round_repolls_and_dedups_against_completed(self): + polls = [] + + def poll(unused_element): + polls.append(len(polls)) + if len(polls) == 1: + return PollResult.incomplete([_ts('a', 1), _ts('b', 2)]) + return PollResult.incomplete([_ts('a', 1), _ts('c', 3)]) + + first = _plan(_initial_polling(), poll) + resumed = _plan(first.residual, poll) + self.assertEqual(2, len(polls)) + self.assertEqual(['c'], [o.value for o in resumed.emit]) + + def test_termination_condition_sets_stop(self): + def poll(unused_element): + return PollResult.incomplete([_ts('a', 1)]) + + termination = after_total_of(10) + for now, expected_stop in [(10.0, False), (11.0, True)]: + with self.subTest(now=now): + plan = _plan( + _initial_polling(termination, Timestamp(0)), + poll, + now=now, + termination=termination) + self.assertEqual(expected_stop, plan.stop) + + def test_non_polling_replays(self): + pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP) + restriction = _NonPollingGrowthState(pending) + plan = _replay_plan(restriction) + self.assertEqual(['a', 'b'], [o.value for o in plan.emit]) + self.assertTrue(plan.stop) + + tracker = _GrowthRestrictionTracker(restriction) + self.assertTrue(tracker.is_bounded()) + self.assertTrue(tracker.try_claim(plan)) + # A checkpoint after replay resumes an empty state that emits nothing. + primary, residual = tracker.try_split(0) + self.assertIs(restriction, primary) + self.assertEqual((), residual.pending.outputs) + self.assertTrue(tracker.check_done()) + + def test_terminal_split_residual_is_empty_for_all_stop_causes(self): + termination = after_total_of(Duration(10)) + cases = [ + ( + 'reached_max', + never(), + _initial_polling(), lambda element: PollResult( + (TimestampedValue('a', MAX_TIMESTAMP), ), None), + 0.0), + ( + 'complete', + never(), + _initial_polling(), + lambda element: PollResult.complete([_ts('a', 1)]), + 0.0), + ( + 'after_total_of', + termination, + _initial_polling(termination, Timestamp(0)), + lambda element: PollResult.incomplete([_ts('a', 1)]), + 100.0), + ] + for name, condition, restriction, poll_fn, now in cases: + with self.subTest(name=name): + plan = _plan(restriction, poll_fn, now=now, termination=condition) + self.assertTrue(plan.stop) + self.assertIsInstance(plan.residual, _NonPollingGrowthState) + self.assertEqual((), plan.residual.pending.outputs) + + def test_wrapper_chain_defers_merged_residual(self): + def poll(unused_element): + return PollResult.incomplete([_ts('a', 1), _ts('b', 2)]) + + plan = _plan(_initial_polling(), poll) + threadsafe = ThreadsafeRestrictionTracker( + _GrowthRestrictionTracker(_initial_polling())) + view = RestrictionTrackerView(threadsafe) + self.assertTrue(view.try_claim(plan)) + view.defer_remainder(Duration(5)) + residual, _ = threadsafe.deferred_status() + self.assertIsInstance(residual, _PollingGrowthState) + self.assertEqual(2, len(residual.completed)) + + +# 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): + _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): + return PollResult.complete([_ts(prefix + 'a', 1), _ts(prefix + 'b', 2)]) + + +def _windowed_group(kv, window=beam.DoFn.WindowParam): + return ((window.start, window.end), sorted(kv[1])) + + +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.growth_of(_complete_poll).with_poll_interval( + Duration(1)).with_output_coder(StrUtf8Coder())) + 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.growth_of(_complete_poll).with_poll_interval( + Duration(1)).with_output_coder(StrUtf8Coder())) + 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.growth_of(_growing_poll).with_poll_interval( + Duration(0.05)).with_output_coder(StrUtf8Coder())) + 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_rejects_non_deterministic_output_coder(self): + # No output coder resolves to PickleCoder, which is non-deterministic, so + # dedup could re-emit equal outputs. Expansion must reject it. + with self.assertRaises(ValueError): + with self._in_memory_pipeline() as p: + _ = ( + p | beam.Create(['k:']) + | Watch.growth_of(_complete_poll).with_poll_interval(Duration(1))) + + def test_derives_output_type_from_input_and_coder(self): + # expand() types the (input, output) pairs from the input type and the + # resolved coder, so downstream stays typed without a manual hint. + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['k:']) + | Watch.growth_of(_complete_poll).with_poll_interval( + Duration(1)).with_output_coder(StrUtf8Coder())) + self.assertEqual(typehints.Tuple[str, str], output.element_type) + + +if __name__ == '__main__': + unittest.main() From c559963d11aeeaa8857d04b4fb912c58211eaf20 Mon Sep 17 00:00:00 2001 From: Eliaazzz Date: Tue, 14 Jul 2026 21:24:31 +1000 Subject: [PATCH 2/3] [Python] Rework Watch API, coder inference, and tracker per review - Replace growth_of and the with_* builders with constructor arguments. - Infer the output coder from PollFn.default_output_coder() or a PollResult[V] return annotation via the coder registry; require determinism only of the dedup key coder, converting it GroupByKey-style with as_deterministic_coder, and add output_key_fn/output_key_coder mirroring Java's outputKeyFn/outputKeyCoder. - Rework the restriction tracker to Java's design: try_claim takes the (PollResult, termination_state) round and validates it against the restriction; try_split derives the replay primary and merged residual. - Match Java's watermark handling: seed the estimator from the input timestamp and advance it to the poll watermark or earliest new output when resuming. --- sdks/python/apache_beam/io/watch.py | 391 +++++++++++++---------- sdks/python/apache_beam/io/watch_test.py | 346 +++++++++++++------- 2 files changed, 439 insertions(+), 298 deletions(-) diff --git a/sdks/python/apache_beam/io/watch.py b/sdks/python/apache_beam/io/watch.py index 68eef1596290..2fcee7a8080f 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. Dedup +hashes each output's key: the output itself by default, or +``output_key_fn(output)`` when one is given. The key coder is inferred when +not passed explicitly and converted to its deterministic form, so equal keys +hash equally across workers and restarts. Example:: @@ -38,15 +40,15 @@ from apache_beam.transforms.window import TimestampedValue from apache_beam.utils.timestamp import Duration, Timestamp - def poll(prefix): + def poll(prefix) -> PollResult[str]: now = Timestamp.now() outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)] return PollResult.complete(outputs) - watched = (inputs - | Watch.growth_of(poll) - .with_poll_interval(Duration(seconds=5)) - .with_termination_per_input(after_total_of(60))) + watched = inputs | Watch( + poll, + poll_interval=Duration(seconds=5), + termination=after_total_of(60)) This API is experimental and may change in backwards-incompatible ways. """ @@ -54,13 +56,16 @@ def poll(prefix): import collections import dataclasses import hashlib +import inspect import time +import typing from typing import Any from typing import Callable +from typing import Generic from typing import Iterable -from typing import List from typing import Optional from typing import Tuple +from typing import TypeVar from apache_beam import coders from apache_beam.coders.coders import Coder @@ -88,18 +93,23 @@ def poll(prefix): _HASH_DIGEST_SIZE = 16 # 128-bit digest width. +OutputT = TypeVar('OutputT') + # ------------------------------------------------------------------------------ # Public API. # ------------------------------------------------------------------------------ @dataclasses.dataclass(frozen=True) -class PollResult: +class PollResult(Generic[OutputT]): """Outputs produced by one poll, plus an optional explicit watermark. ``watermark`` of ``None`` lets the transform infer the watermark from the earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by :meth:`complete`) marks the input finished, so polling stops. + + The ``OutputT`` type parameter can annotate a poll function's return type, + as in ``-> PollResult[str]``; the transform infers the output coder from it. """ outputs: Tuple[TimestampedValue, ...] watermark: Optional[Timestamp] = None @@ -149,8 +159,7 @@ class PollFn(object): """Optional base for a poll function ``input -> PollResult``. Any callable with that signature works; subclass only to attach an output - coder hint via :meth:`default_output_coder`. The hint is used for dedup and - must be deterministic:: + coder hint via :meth:`default_output_coder`:: from apache_beam import coders @@ -160,6 +169,9 @@ def __call__(self, prefix): def default_output_coder(self): return coders.StrUtf8Coder() + + A plain function can instead annotate its return type as ``PollResult[V]`` + and have the output coder inferred from ``V``. """ def __call__(self, element: Any) -> PollResult: raise NotImplementedError @@ -244,9 +256,9 @@ class _GrowthState: class _PollingGrowthState(_GrowthState): """Keep-polling state: emitted-output hashes, 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; a new mapping is built for - each residual. + ``completed`` maps a 16-byte output-key hash to the event time it was first + seen. It is insertion-ordered and treated as immutable; a new mapping is + built for each residual. """ completed: 'collections.OrderedDict[bytes, Timestamp]' poll_watermark: Optional[Timestamp] @@ -263,6 +275,9 @@ class _NonPollingGrowthState(_GrowthState): pending: PollResult +# Primary used when a checkpoint arrives before any claim; replays nothing. +_EMPTY_STATE = _NonPollingGrowthState(PollResult((), None)) + # ------------------------------------------------------------------------------ # Coders. # ------------------------------------------------------------------------------ @@ -344,19 +359,8 @@ def is_deterministic(self) -> bool: # ------------------------------------------------------------------------------ -@dataclasses.dataclass(frozen=True) -class _PollPlan: - """One planned poll round: what to emit plus the self-checkpoint split. - - ``process()`` builds this from a poll result before claiming, so the poll runs - outside the tracker lock. ``primary`` is the round just processed and - ``residual`` is the work a checkpoint resumes. - """ - emit: Tuple[TimestampedValue, ...] - watermark: Optional[Timestamp] - stop: bool - primary: _GrowthState - residual: _GrowthState +def _identity(value: Any) -> Any: + return value def _hash_output(key_coder: Coder, value: Any) -> bytes: @@ -373,117 +377,110 @@ def _max_watermark(left: Optional[Timestamp], return max(left, right) -def _plan_poll_round( +def _never_seen_before( restriction: _PollingGrowthState, result: PollResult, - termination: TerminationCondition, - key_coder: Coder, - now: Timestamp) -> _PollPlan: - """Dedups a poll result into new outputs and builds the checkpoint split. + key_fn: Callable[[Any], Any], + key_coder: Coder) -> PollResult: + """Filters a poll result down to outputs whose key was never seen before. - Pure and side-effect free, so ``process()`` can call it after the poll and - before claiming, keeping the poll off the tracker lock. + Dedup hashes ``key_fn(output.value)`` against the restriction's completed + set, also dropping in-round duplicates. Outputs are sorted by timestamp so + the earliest one can serve as the inferred watermark. """ - new_outputs = [] # type: List[TimestampedValue] - claimed = [] # type: List[Tuple[bytes, Timestamp]] - seen_this_round = set() # type: set + new_outputs = [] + seen_this_round = set() for output in result.outputs: - key_hash = _hash_output(key_coder, output.value) + key_hash = _hash_output(key_coder, key_fn(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 - if new_outputs: - termination_state = termination.on_seen_new_output(now, termination_state) - termination_state = termination.on_poll_complete(termination_state) - - if result.watermark is not None: - watermark = result.watermark - elif new_outputs: - watermark = new_outputs[0].timestamp - else: - watermark = None - - # A watermark at MAX means no more output is possible, so polling stops. - reached_max = watermark is not None and watermark >= MAX_TIMESTAMP - stop = ( - result.is_complete or reached_max or - termination.can_stop_polling(now, termination_state)) - - primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark)) - if stop: - # Terminal round: no polling work remains, so a checkpoint (runner-initiated - # or via defer_remainder) resumes a state that emits nothing. - residual = _NonPollingGrowthState(PollResult((), watermark)) - else: - merged = collections.OrderedDict(restriction.completed) - for key_hash, first_seen in claimed: - merged[key_hash] = first_seen - residual_watermark = _max_watermark(restriction.poll_watermark, watermark) - residual = _PollingGrowthState( - merged, residual_watermark, termination_state) - return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual) - - -def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan: - """Plans a replay of the outputs already emitted this round, then stops.""" - resumed = _NonPollingGrowthState( - PollResult((), restriction.pending.watermark)) - return _PollPlan( - restriction.pending.outputs, None, True, restriction, resumed) + return dataclasses.replace(result, outputs=tuple(new_outputs)) class _GrowthRestrictionTracker(iobase.RestrictionTracker): - """Tracks one input's polling restriction and its self-checkpoints. - - ``process()`` runs the poll and plans the round with - :func:`_plan_poll_round`, then claims the plan; the tracker only records the - plan's primary and residual so a checkpoint (via ``defer_remainder`` or the - runner) resumes correctly. Keeping the poll in ``process()`` means a slow - poll never holds the tracker lock, so it cannot delay runner progress checks - or checkpoints. + """Tracks one input's polling restriction over claimed poll rounds. + + The claimed position is one poll round: a ``(PollResult, termination_state)`` + pair whose ``PollResult`` holds only never-seen-before outputs. ``process()`` + polls and dedups before claiming, so a slow poll never holds the tracker + lock; the tracker validates each claim against the restriction and derives + the checkpoint split from the claimed round in :meth:`try_split`. """ - def __init__(self, restriction: _GrowthState): + def __init__( + self, + restriction: _GrowthState, + key_fn: Callable[[Any], Any], + key_coder: Coder): self._restriction = restriction + self._key_fn = key_fn + self._key_coder = key_coder + self._claimed_result = None # type: Optional[PollResult] + self._claimed_termination_state = None # type: Any + self._claimed_hashes = None # type: Optional[collections.OrderedDict] self._should_stop = False - self._primary = None # type: Optional[_GrowthState] - self._residual = None # type: Optional[_GrowthState] + + def _hash(self, value: Any) -> bytes: + return _hash_output(self._key_coder, self._key_fn(value)) def current_restriction(self) -> _GrowthState: return self._restriction - def try_claim(self, plan: _PollPlan) -> bool: - """Records a planned round's checkpoint split; one claim per ``process()``. + def try_claim(self, position: Tuple[PollResult, Any]) -> bool: + """Claims one poll round; at most one claim succeeds per ``process()``. - Returns ``False`` only when a checkpoint already stopped this invocation, in - which case ``process()`` must emit nothing. + The claim is rejected after a checkpoint already stopped this invocation, + when a claimed output key was already completed, or when a replay does not + match the pending outputs exactly. """ if self._should_stop: return False - self._primary = plan.primary - self._residual = plan.residual + result, termination_state = position + claimed_hashes = collections.OrderedDict() + for output in result.outputs: + claimed_hashes[self._hash(output.value)] = output.timestamp + if isinstance(self._restriction, _PollingGrowthState): + if any(key_hash in self._restriction.completed + for key_hash in claimed_hashes): + return False + else: + expected = set( + self._hash(output.value) + for output in self._restriction.pending.outputs) + if expected != set(claimed_hashes): + return False self._should_stop = True + self._claimed_result = result + self._claimed_termination_state = termination_state + self._claimed_hashes = claimed_hashes return True def try_split(self, fraction_of_remainder): - # Only self-checkpoint (fraction 0) is supported; decline dynamic splits. - if fraction_of_remainder != 0: - return None - if self._primary is None: - # No claim happened this invocation: keep the whole state as the residual. - primary = _NonPollingGrowthState(PollResult((), None)) + # Every split checkpoints at the claimed poll round; splitting a round + # further is not supported. + if self._claimed_result is None: + # No claim happened this invocation: the residual is all the work and + # the primary replays nothing. residual = self._restriction - self._restriction = primary - self._should_stop = True - return primary, residual - primary, residual = self._primary, self._residual - self._restriction = primary + self._restriction = _EMPTY_STATE + elif isinstance(self._restriction, _NonPollingGrowthState): + # The claimed replay was the entire restriction, so nothing remains. + residual = _EMPTY_STATE + else: + # The primary becomes a replay of the claimed round; the residual + # resumes polling with the claimed keys marked completed. + merged = collections.OrderedDict(self._restriction.completed) + merged.update(self._claimed_hashes) + residual = _PollingGrowthState( + merged, + _max_watermark( + self._restriction.poll_watermark, self._claimed_result.watermark), + self._claimed_termination_state) + self._restriction = _NonPollingGrowthState(self._claimed_result) self._should_stop = True - return primary, residual + return self._restriction, residual def check_done(self) -> bool: # Called after every process(); the single claim or a split sets the flag. @@ -523,12 +520,15 @@ def __init__( termination: TerminationCondition, poll_interval: Duration, output_coder: Coder, + key_fn: Callable[[Any], Any], + key_coder: Coder, now_fn: Optional[Callable[[], float]] = None): self._poll_fn = poll_fn self._termination = termination self._poll_interval = poll_interval self._output_coder = output_coder - self._key_coder = output_coder + self._key_fn = key_fn + self._key_coder = key_coder self._now = now_fn or time.time self._restriction_coder = _GrowthStateCoder(output_coder, termination) @@ -540,11 +540,7 @@ def initial_restriction(self, element) -> _PollingGrowthState: self._termination.for_new_input(now, element)) def create_tracker(self, restriction) -> _GrowthRestrictionTracker: - return _GrowthRestrictionTracker(restriction) - - def split(self, element, restriction): - # Watch fans out by input element, so each restriction stays whole. - yield restriction + return _GrowthRestrictionTracker(restriction, self._key_fn, self._key_coder) def restriction_coder(self) -> Coder: return self._restriction_coder @@ -552,12 +548,6 @@ def restriction_coder(self) -> Coder: def restriction_size(self, element, restriction) -> int: return 1 - def truncate(self, element, restriction): - # On drain, replay a pending NonPolling state and stop further polling. - if isinstance(restriction, _NonPollingGrowthState): - return restriction - return None - @core.DoFn.unbounded_per_element() def process( self, @@ -567,34 +557,50 @@ def process( watermark_estimator=core.DoFn.WatermarkEstimatorParam( ManualWatermarkEstimator.default_provider())): assert isinstance(tracker, sdf_utils.RestrictionTrackerView) + # Java seeds the manual estimator with the element timestamp; the Python + # default provider starts at None, which a runner reads as MIN_TIMESTAMP + # and would pin the stage's output watermark until the first output. + if watermark_estimator.current_watermark() is None: + watermark_estimator.set_watermark(timestamp) restriction = tracker.current_restriction() if isinstance(restriction, _NonPollingGrowthState): # Replay the outputs already emitted this round, then stop. No poll. - if not tracker.try_claim(_replay_plan(restriction)): + if not tracker.try_claim((restriction.pending, None)): return - _set_watermark_if_greater(watermark_estimator, timestamp) for output in restriction.pending.outputs: yield TimestampedValue((element, output.value), output.timestamp) return # Poll before claiming so a slow poll never holds the tracker lock, which # would block runner progress checks and checkpoints. - now = Timestamp.of(self._now()) result = self._poll_fn(element) - plan = _plan_poll_round( - restriction, result, self._termination, self._key_coder, now) - if not tracker.try_claim(plan): + # Read the clock after the poll so a slow poll counts against termination. + now = Timestamp.of(self._now()) + new_results = _never_seen_before( + restriction, result, self._key_fn, self._key_coder) + termination_state = restriction.termination_state + if new_results.outputs: + termination_state = self._termination.on_seen_new_output( + now, termination_state) + termination_state = self._termination.on_poll_complete(termination_state) + if not tracker.try_claim((new_results, termination_state)): # A checkpoint already stopped this invocation; emit nothing. return - # Seed the watermark hold from the input event time after the claim. - _set_watermark_if_greater(watermark_estimator, timestamp) - for output in plan.emit: + for output in new_results.outputs: yield TimestampedValue((element, output.value), output.timestamp) - if plan.stop: - # The input is finished, so release the watermark hold to MAX. - _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP) + if new_results.watermark is not None: + watermark = new_results.watermark + elif new_results.outputs: + # Outputs are timestamp-sorted, so the first one is the earliest. + watermark = new_results.outputs[0].timestamp + else: + watermark = None + if self._termination.can_stop_polling(now, termination_state): return - if plan.watermark is not None: - _set_watermark_if_greater(watermark_estimator, plan.watermark) + if watermark is not None and watermark >= MAX_TIMESTAMP: + # No more output is possible (PollResult.complete), so polling stops. + return + if watermark is not None: + _set_watermark_if_greater(watermark_estimator, watermark) tracker.defer_remainder(self._poll_interval) @@ -610,66 +616,97 @@ def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None: # ------------------------------------------------------------------------------ +def _return_type(fn) -> Any: + """The return type annotation of ``fn`` or its ``__call__``, else ``Any``.""" + target = fn if inspect.isroutine(fn) else getattr(type(fn), '__call__', None) + if target is None: + return Any + try: + hints = typing.get_type_hints(target) + except (NameError, TypeError): + return Any + return hints.get('return', Any) + + +def _poll_output_type(poll_fn) -> Any: + """The ``V`` of a ``PollResult[V]`` return annotation on ``poll_fn``. + + This mirrors the Java SDK, which infers the output coder from the + ``PollFn``'s ``OutputT`` type parameter. Returns ``Any`` when ``poll_fn`` + carries no such annotation. + """ + hint = _return_type(poll_fn) + if typing.get_origin(hint) is PollResult: + args = typing.get_args(hint) + if len(args) == 1: + return args[0] + return Any + + class Watch(PTransform): """Watches a growing set of outputs per input via a periodic poll function. - Build with :meth:`growth_of` and the ``with_*`` methods. The output is an - unbounded ``PCollection`` of ``(input, output)`` pairs. + The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. + + Args: + poll_fn: callable ``input -> PollResult``, invoked once per poll round. + poll_interval: delay between two poll rounds for one input, as a + :class:`Duration` or in seconds. + termination: per-input stop policy; defaults to :func:`never`. + output_coder: coder for the poll outputs, used to keep them in the + restriction state. Inferred when omitted: from a :class:`PollFn`'s + :meth:`~PollFn.default_output_coder`, else from the registered coder for + the ``V`` of a ``PollResult[V]`` return annotation on ``poll_fn``. + output_key_fn: derives the dedup key from an output; an output is emitted + only when its key was never seen before. Defaults to the output itself. + output_key_coder: coder whose encoding of the key is hashed for dedup; + inferred like ``output_coder`` when omitted. It is converted with + ``as_deterministic_coder`` so equal keys always hash equally; a coder + with no deterministic form is rejected. + now_fn: clock used for termination decisions; tests can inject one. """ def __init__( self, poll_fn: Callable[[Any], PollResult], + poll_interval, termination: Optional[TerminationCondition] = None, - poll_interval: Optional[Duration] = None, output_coder: Optional[Coder] = None, + output_key_fn: Optional[Callable[[Any], Any]] = None, + output_key_coder: Optional[Coder] = None, now_fn: Optional[Callable[[], float]] = None): super().__init__() + if poll_interval is None: + raise ValueError('Watch requires a poll_interval') self._poll_fn = poll_fn + self._poll_interval = _as_duration(poll_interval) self._termination = termination or never() - self._poll_interval = poll_interval self._output_coder = output_coder + self._output_key_fn = output_key_fn + self._output_key_coder = output_key_coder self._now = now_fn - @classmethod - def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch': - return cls(poll_fn) - - def _replace(self, **changes) -> 'Watch': - spec = dict( - poll_fn=self._poll_fn, - termination=self._termination, - poll_interval=self._poll_interval, - output_coder=self._output_coder, - now_fn=self._now) - spec.update(changes) - return Watch(**spec) - - def with_poll_interval(self, poll_interval) -> 'Watch': - return self._replace(poll_interval=_as_duration(poll_interval)) - - def with_termination_per_input( - self, termination: TerminationCondition) -> 'Watch': - return self._replace(termination=termination) - - def with_output_coder(self, output_coder: Coder) -> 'Watch': - return self._replace(output_coder=output_coder) - def expand(self, pcoll): - if self._poll_interval is None: - raise ValueError('Watch requires with_poll_interval(...)') output_coder = self._output_coder + if output_coder is None and isinstance(self._poll_fn, PollFn): + output_coder = self._poll_fn.default_output_coder() if output_coder is None: - hint = self._poll_fn.default_output_coder() if isinstance( - self._poll_fn, PollFn) else None - output_coder = hint or coders.PickleCoder() - # Dedup hashes the encoded output as its key, so a non-deterministic coder - # would hash equal outputs differently and re-emit them, defeating the - # transform. Require determinism rather than silently emitting duplicates. - if not output_coder.is_deterministic(): - raise ValueError( - 'Watch dedup requires a deterministic output coder, but %s is not ' - 'deterministic. Pass one via with_output_coder().' % - type(output_coder).__name__) + output_coder = coders.registry.get_coder(_poll_output_type(self._poll_fn)) + if self._output_key_fn is None: + # The output is its own dedup key, so the key coder is the output coder. + key_fn = _identity + key_coder = self._output_key_coder or output_coder + else: + key_fn = self._output_key_fn + key_coder = self._output_key_coder or coders.registry.get_coder( + _return_type(self._output_key_fn)) + # Dedup hashes the encoded key, so equal keys must encode equally; use the + # coder's deterministic form and reject coders that have none. + key_coder = key_coder.as_deterministic_coder( + self.label, + 'Watch dedups by hashing the encoded output key, so the key coder ' + 'must be deterministic. %s has no deterministic form; pass a ' + 'deterministic output_key_coder (or output_coder).' % + type(key_coder).__name__) # Type the (input, output) pairs from the input type and the resolved # coder's type, so downstream transforms are typed and coder inference does # not fall back to pickling. @@ -684,6 +721,8 @@ def expand(self, pcoll): self._termination, self._poll_interval, output_coder, + key_fn, + key_coder, self._now)).with_output_types(Tuple[input_type, value_type]) diff --git a/sdks/python/apache_beam/io/watch_test.py b/sdks/python/apache_beam/io/watch_test.py index e530a42df7e2..2b41ff1cf319 100644 --- a/sdks/python/apache_beam/io/watch_test.py +++ b/sdks/python/apache_beam/io/watch_test.py @@ -21,20 +21,24 @@ import unittest import apache_beam as beam +from apache_beam.coders.coders import Coder from apache_beam.coders.coders import StrUtf8Coder +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 _plan_poll_round from apache_beam.io.watch import _PollingGrowthState -from apache_beam.io.watch import _replay_plan +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 @@ -52,14 +56,17 @@ def _ts(value, timestamp): return TimestampedValue(value, Timestamp(timestamp)) -def _plan(restriction, poll_fn, now=0.0, termination=None): - termination = termination or never() - return _plan_poll_round( - restriction, - poll_fn('input'), - termination, - StrUtf8Coder(), - Timestamp.of(now)) +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 _initial_polling(termination=None, now=Timestamp(0)): @@ -94,128 +101,140 @@ def test_non_polling_round_trip_preserves_pending_outputs(self): [(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_poll_claims_dedups_and_checkpoints(self): - def poll(unused_element): - return PollResult.incomplete([_ts('a', 1), _ts('a', 1), _ts('b', 2)]) - - plan = _plan(_initial_polling(), poll) - self.assertEqual(['a', 'b'], [o.value for o in plan.emit]) - self.assertFalse(plan.stop) - self.assertEqual(Timestamp(1), plan.watermark) - self.assertIsInstance(plan.primary, _NonPollingGrowthState) - self.assertIsInstance(plan.residual, _PollingGrowthState) - self.assertEqual(2, len(plan.residual.completed)) - - tracker = _GrowthRestrictionTracker(_initial_polling()) + 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(plan)) + self.assertTrue(tracker.try_claim((new_results, 0))) primary, residual = tracker.try_split(0) - self.assertIs(plan.primary, primary) - self.assertIs(plan.residual, residual) + 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 explicit_watermark_poll(unused_element): - return PollResult.incomplete([_ts('c', 3)]).with_watermark(5) - - plan = _plan(_initial_polling(), explicit_watermark_poll) - self.assertEqual(Timestamp(5), plan.watermark) - self.assertEqual(Timestamp(5), plan.residual.poll_watermark) - - def test_second_round_repolls_and_dedups_against_completed(self): - polls = [] - - def poll(unused_element): - polls.append(len(polls)) - if len(polls) == 1: - return PollResult.incomplete([_ts('a', 1), _ts('b', 2)]) - return PollResult.incomplete([_ts('a', 1), _ts('c', 3)]) - - first = _plan(_initial_polling(), poll) - resumed = _plan(first.residual, poll) - self.assertEqual(2, len(polls)) - self.assertEqual(['c'], [o.value for o in resumed.emit]) - - def test_termination_condition_sets_stop(self): - def poll(unused_element): - return PollResult.incomplete([_ts('a', 1)]) + 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()) - termination = after_total_of(10) - for now, expected_stop in [(10.0, False), (11.0, True)]: - with self.subTest(now=now): - plan = _plan( - _initial_polling(termination, Timestamp(0)), - poll, - now=now, - termination=termination) - self.assertEqual(expected_stop, plan.stop) - - def test_non_polling_replays(self): + def test_non_polling_replays_exactly_the_pending_outputs(self): pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP) - restriction = _NonPollingGrowthState(pending) - plan = _replay_plan(restriction) - self.assertEqual(['a', 'b'], [o.value for o in plan.emit]) - self.assertTrue(plan.stop) - - tracker = _GrowthRestrictionTracker(restriction) + tracker = _tracker(_NonPollingGrowthState(pending)) self.assertTrue(tracker.is_bounded()) - self.assertTrue(tracker.try_claim(plan)) - # A checkpoint after replay resumes an empty state that emits nothing. - primary, residual = tracker.try_split(0) - self.assertIs(restriction, primary) + # 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_terminal_split_residual_is_empty_for_all_stop_causes(self): - termination = after_total_of(Duration(10)) - cases = [ - ( - 'reached_max', - never(), - _initial_polling(), lambda element: PollResult( - (TimestampedValue('a', MAX_TIMESTAMP), ), None), - 0.0), - ( - 'complete', - never(), - _initial_polling(), - lambda element: PollResult.complete([_ts('a', 1)]), - 0.0), - ( - 'after_total_of', - termination, - _initial_polling(termination, Timestamp(0)), - lambda element: PollResult.incomplete([_ts('a', 1)]), - 100.0), - ] - for name, condition, restriction, poll_fn, now in cases: - with self.subTest(name=name): - plan = _plan(restriction, poll_fn, now=now, termination=condition) - self.assertTrue(plan.stop) - self.assertIsInstance(plan.residual, _NonPollingGrowthState) - self.assertEqual((), plan.residual.pending.outputs) + 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): - def poll(unused_element): - return PollResult.incomplete([_ts('a', 1), _ts('b', 2)]) - - plan = _plan(_initial_polling(), poll) - threadsafe = ThreadsafeRestrictionTracker( - _GrowthRestrictionTracker(_initial_polling())) + 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(plan)) + 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)) +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)] @@ -224,14 +243,84 @@ def _growing_poll(prefix): return PollResult.incomplete(outputs) -def _complete_poll(prefix): +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 _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): + dofn = _WatchGrowthDoFn( + poll_fn, + never(), + Duration(1), + StrUtf8Coder(), + _identity, + StrUtf8Coder()) + threadsafe = ThreadsafeRestrictionTracker( + dofn.create_tracker(dofn.initial_restriction(element))) + estimator = ThreadsafeWatermarkEstimator(ManualWatermarkEstimator(None)) + 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()) + + class WatchEndToEndTest(unittest.TestCase): def _in_memory_pipeline(self): return TestPipeline( @@ -241,8 +330,7 @@ def test_complete_outputs_values_and_timestamps(self): with self._in_memory_pipeline() as p: output = ( p | beam.Create(['k:']) - | Watch.growth_of(_complete_poll).with_poll_interval( - Duration(1)).with_output_coder(StrUtf8Coder())) + | Watch(_complete_poll, poll_interval=Duration(1))) assert_that( output, equal_to([ @@ -255,8 +343,7 @@ def test_complete_advances_watermark_for_windowed_pipeline(self): with self._in_memory_pipeline() as p: output = ( p | beam.Create(['k:']) - | Watch.growth_of(_complete_poll).with_poll_interval( - Duration(1)).with_output_coder(StrUtf8Coder())) + | Watch(_complete_poll, poll_interval=Duration(1))) grouped = ( output | beam.WindowInto(FixedWindows(10)) @@ -274,8 +361,7 @@ def test_multi_round_dedups_stops_and_is_per_input(self): 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_output_coder(StrUtf8Coder())) + | 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'), @@ -283,23 +369,39 @@ 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_rejects_non_deterministic_output_coder(self): - # No output coder resolves to PickleCoder, which is non-deterministic, so - # dedup could re-emit equal outputs. Expansion must reject it. + 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.growth_of(_complete_poll).with_poll_interval(Duration(1))) + | 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_derives_output_type_from_input_and_coder(self): - # expand() types the (input, output) pairs from the input type and the - # resolved coder, so downstream stays typed without a manual hint. + def test_uses_poll_fn_default_output_coder(self): with self._in_memory_pipeline() as p: output = ( p | beam.Create(['k:']) - | Watch.growth_of(_complete_poll).with_poll_interval( - Duration(1)).with_output_coder(StrUtf8Coder())) + | Watch(_StrCoderPollFn(), poll_interval=Duration(1))) self.assertEqual(typehints.Tuple[str, str], output.element_type) From aeafcb2f01f6f5011f16e0176e248cd1bb1db791 Mon Sep 17 00:00:00 2001 From: Eliaazzz Date: Tue, 21 Jul 2026 22:36:17 +1000 Subject: [PATCH 3/3] [Python] Add Watch tests for terminal and replay round watermarks The replay round leaves the watermark at the seed and reports no residual. A run that defers twice before completing keeps the watermark where the last poll left it and ends without a residual. --- sdks/python/apache_beam/io/watch_test.py | 55 ++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/io/watch_test.py b/sdks/python/apache_beam/io/watch_test.py index 2b41ff1cf319..8c1f6571da66 100644 --- a/sdks/python/apache_beam/io/watch_test.py +++ b/sdks/python/apache_beam/io/watch_test.py @@ -284,7 +284,8 @@ def _windowed_group(kv, window=beam.DoFn.WindowParam): class WatchDoFnProcessTest(unittest.TestCase): - def _process(self, poll_fn, element, timestamp): + def _process( + self, poll_fn, element, timestamp, restriction=None, watermark=None): dofn = _WatchGrowthDoFn( poll_fn, never(), @@ -292,9 +293,11 @@ def _process(self, poll_fn, element, timestamp): StrUtf8Coder(), _identity, StrUtf8Coder()) - threadsafe = ThreadsafeRestrictionTracker( - dofn.create_tracker(dofn.initial_restriction(element))) - estimator = ThreadsafeWatermarkEstimator(ManualWatermarkEstimator(None)) + if restriction is None: + restriction = dofn.initial_restriction(element) + threadsafe = ThreadsafeRestrictionTracker(dofn.create_tracker(restriction)) + estimator = ThreadsafeWatermarkEstimator( + ManualWatermarkEstimator(watermark)) outputs = list( dofn.process( element, @@ -320,6 +323,50 @@ def test_complete_round_stops_without_residual(self): 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()) + class WatchEndToEndTest(unittest.TestCase): def _in_memory_pipeline(self):