This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 9c3f51e56e0 [Python] Add Watch transform with growth_of polling SDF
(#39023)
9c3f51e56e0 is described below
commit 9c3f51e56e09f9b5dd7e1b30c5f78dc55d51142b
Author: Elia Liu <[email protected]>
AuthorDate: Fri Jul 24 04:17:55 2026 +1000
[Python] Add Watch transform with growth_of polling SDF (#39023)
* [Python] Add Watch transform with growth_of polling SDF
---
sdks/python/apache_beam/io/watch.py | 730 +++++++++++++++++++++++++++++++
sdks/python/apache_beam/io/watch_test.py | 456 +++++++++++++++++++
2 files changed, 1186 insertions(+)
diff --git a/sdks/python/apache_beam/io/watch.py
b/sdks/python/apache_beam/io/watch.py
new file mode 100644
index 00000000000..2fcee7a8080
--- /dev/null
+++ b/sdks/python/apache_beam/io/watch.py
@@ -0,0 +1,730 @@
+#
+# 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
+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::
+
+ 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) -> PollResult[str]:
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ 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.
+"""
+
+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 Optional
+from typing import Tuple
+from typing import TypeVar
+
+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.
+
+OutputT = TypeVar('OutputT')
+
+#
------------------------------------------------------------------------------
+# Public API.
+#
------------------------------------------------------------------------------
+
+
[email protected](frozen=True)
+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
+
+ @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`::
+
+ 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()
+
+ 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
+
+ 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."""
+
+
[email protected](frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``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]
+ termination_state: Any
+
+
[email protected](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
+
+
+# Primary used when a checkpoint arrives before any claim; replays nothing.
+_EMPTY_STATE = _NonPollingGrowthState(PollResult((), None))
+
+#
------------------------------------------------------------------------------
+# 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.
+#
------------------------------------------------------------------------------
+
+
+def _identity(value: Any) -> Any:
+ return value
+
+
+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 _never_seen_before(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ key_fn: Callable[[Any], Any],
+ key_coder: Coder) -> PollResult:
+ """Filters a poll result down to outputs whose key was never seen before.
+
+ 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 = []
+ seen_this_round = set()
+ for output in result.outputs:
+ 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)
+ new_outputs.sort(key=lambda output: output.timestamp)
+ return dataclasses.replace(result, outputs=tuple(new_outputs))
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """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,
+ 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
+
+ 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, position: Tuple[PollResult, Any]) -> bool:
+ """Claims one poll round; at most one claim succeeds per ``process()``.
+
+ 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
+ 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):
+ # 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 = _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 self._restriction, 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,
+ 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_fn = key_fn
+ self._key_coder = key_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, self._key_fn,
self._key_coder)
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ @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)
+ # 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((restriction.pending, None)):
+ return
+ 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.
+ result = self._poll_fn(element)
+ # 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
+ for output in new_results.outputs:
+ yield TimestampedValue((element, output.value), output.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 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)
+
+
+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.
+#
------------------------------------------------------------------------------
+
+
+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.
+
+ 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,
+ 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._output_coder = output_coder
+ self._output_key_fn = output_key_fn
+ self._output_key_coder = output_key_coder
+ self._now = now_fn
+
+ def expand(self, pcoll):
+ 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:
+ 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.
+ 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,
+ key_fn,
+ key_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 00000000000..8c1f6571da6
--- /dev/null
+++ b/sdks/python/apache_beam/io/watch_test.py
@@ -0,0 +1,456 @@
+#
+# 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 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 _PollingGrowthState
+from apache_beam.io.watch import _WatchGrowthDoFn
+from apache_beam.io.watch import after_total_of
+from apache_beam.io.watch import never
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.runners.sdf_utils import RestrictionTrackerView
+from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker
+from apache_beam.runners.sdf_utils import ThreadsafeWatermarkEstimator
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import TestWindowedValue
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+from apache_beam.transforms.window import FixedWindows
+from apache_beam.transforms.window import GlobalWindow
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.typehints import typehints
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+
+def _ts(value, timestamp):
+ return TimestampedValue(value, Timestamp(timestamp))
+
+
+def _identity(output):
+ return output
+
+
+def _new_results(restriction, result, key_fn=None):
+ return _never_seen_before(
+ restriction, result, key_fn or _identity, StrUtf8Coder())
+
+
+def _tracker(restriction):
+ return _GrowthRestrictionTracker(restriction, _identity, StrUtf8Coder())
+
+
+def _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 NeverSeenBeforeTest(unittest.TestCase):
+ def test_dedups_and_sorts_by_timestamp(self):
+ result = PollResult.incomplete([_ts('b', 2), _ts('a', 1), _ts('a', 1)])
+ new_results = _new_results(_initial_polling(), result)
+ self.assertEqual(['a', 'b'], [o.value for o in new_results.outputs])
+
+ def test_dedups_against_completed_keys(self):
+ state = _initial_polling()
+ first = _new_results(
+ state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)]))
+ tracker = _tracker(state)
+ self.assertTrue(tracker.try_claim((first, 0)))
+ _, residual = tracker.try_split(0)
+ second = _new_results(
+ residual, PollResult.incomplete([_ts('a', 1), _ts('c', 3)]))
+ self.assertEqual(['c'], [o.value for o in second.outputs])
+
+ def test_output_key_dedups_by_derived_key(self):
+ result = PollResult.incomplete([_ts('a1', 1), _ts('a2', 2), _ts('b1', 3)])
+ # The key is the first character, so 'a1' and 'a2' collapse to one output.
+ new_results = _new_results(
+ _initial_polling(), result, key_fn=lambda output: output[0])
+ self.assertEqual(['a1', 'b1'], [o.value for o in new_results.outputs])
+
+ def test_preserves_explicit_watermark(self):
+ result = PollResult.incomplete([_ts('c', 3)]).with_watermark(5)
+ new_results = _new_results(_initial_polling(), result)
+ self.assertEqual(Timestamp(5), new_results.watermark)
+
+
+class GrowthTrackerTest(unittest.TestCase):
+ def test_claim_then_split_builds_replay_primary_and_merged_residual(self):
+ state = _initial_polling()
+ new_results = _new_results(
+ state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)]))
+ tracker = _tracker(state)
+ self.assertFalse(tracker.is_bounded())
+ self.assertTrue(tracker.try_claim((new_results, 0)))
+ primary, residual = tracker.try_split(0)
+ self.assertIsInstance(primary, _NonPollingGrowthState)
+ self.assertEqual(new_results, primary.pending)
+ self.assertIsInstance(residual, _PollingGrowthState)
+ self.assertEqual(2, len(residual.completed))
+ self.assertEqual(0, residual.termination_state)
+ self.assertTrue(tracker.check_done())
+
+ def test_split_merges_explicit_watermark_into_residual(self):
+ state = _initial_polling()
+ result = PollResult.incomplete([_ts('c', 3)]).with_watermark(5)
+ tracker = _tracker(state)
+ self.assertTrue(tracker.try_claim((_new_results(state, result), 0)))
+ _, residual = tracker.try_split(0)
+ self.assertEqual(Timestamp(5), residual.poll_watermark)
+
+ def test_second_claim_is_rejected(self):
+ state = _initial_polling()
+ new_results = _new_results(state, PollResult.incomplete([_ts('a', 1)]))
+ tracker = _tracker(state)
+ self.assertTrue(tracker.try_claim((new_results, 0)))
+ self.assertFalse(tracker.try_claim((new_results, 0)))
+
+ def test_claim_rejects_already_completed_keys(self):
+ # The tracker re-validates a claim, so a poll round that was not deduped
+ # against the restriction is rejected instead of emitting duplicates.
+ state = _initial_polling()
+ first = _new_results(state, PollResult.incomplete([_ts('a', 1)]))
+ tracker = _tracker(state)
+ self.assertTrue(tracker.try_claim((first, 0)))
+ _, residual = tracker.try_split(0)
+ stale = PollResult.incomplete([_ts('a', 1)])
+ self.assertFalse(_tracker(residual).try_claim((stale, 0)))
+
+ def test_split_before_claim_moves_all_work_to_residual(self):
+ state = _initial_polling()
+ tracker = _tracker(state)
+ primary, residual = tracker.try_split(0)
+ self.assertIs(state, residual)
+ self.assertIsInstance(primary, _NonPollingGrowthState)
+ self.assertEqual((), primary.pending.outputs)
+ new_results = _new_results(state, PollResult.incomplete([_ts('a', 1)]))
+ self.assertFalse(tracker.try_claim((new_results, 0)))
+ self.assertTrue(tracker.check_done())
+
+ def test_non_polling_replays_exactly_the_pending_outputs(self):
+ pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP)
+ tracker = _tracker(_NonPollingGrowthState(pending))
+ self.assertTrue(tracker.is_bounded())
+ # A replay must claim the pending poll result exactly.
+ partial = PollResult((_ts('a', 1), ), None)
+ self.assertFalse(tracker.try_claim((partial, None)))
+ self.assertTrue(tracker.try_claim((pending, None)))
+ # A checkpoint after the replay leaves no residual work.
+ _, residual = tracker.try_split(0)
+ self.assertEqual((), residual.pending.outputs)
+ self.assertTrue(tracker.check_done())
+
+ def test_check_done_raises_without_claim_or_split(self):
+ tracker = _tracker(_initial_polling())
+ with self.assertRaises(ValueError):
+ tracker.check_done()
+
+ def test_wrapper_chain_defers_merged_residual(self):
+ state = _initial_polling()
+ new_results = _new_results(
+ state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)]))
+ threadsafe = ThreadsafeRestrictionTracker(_tracker(state))
+ view = RestrictionTrackerView(threadsafe)
+ self.assertTrue(view.try_claim((new_results, 0)))
+ view.defer_remainder(Duration(5))
+ residual, _ = threadsafe.deferred_status()
+ self.assertIsInstance(residual, _PollingGrowthState)
+ self.assertEqual(2, len(residual.completed))
+
+
+class TerminationConditionTest(unittest.TestCase):
+ def test_never_does_not_stop(self):
+ termination = never()
+ state = termination.for_new_input(Timestamp(0), 'input')
+ self.assertFalse(termination.can_stop_polling(MAX_TIMESTAMP, state))
+
+ def test_after_total_of_stops_once_duration_elapsed(self):
+ termination = after_total_of(10)
+ state = termination.for_new_input(Timestamp(0), 'input')
+ self.assertFalse(termination.can_stop_polling(Timestamp(10), state))
+ self.assertTrue(termination.can_stop_polling(Timestamp(11), state))
+
+
+# Module-level so the poll function pickles by reference; the call counter is
+# shared within the single in-memory DirectRunner process.
+_POLL_CALLS = collections.defaultdict(int)
+
+
+def _growing_poll(prefix):
+ # Unannotated on purpose: dedup must hold on the inferred fallback coder.
+ _POLL_CALLS[prefix] += 1
+ count = _POLL_CALLS[prefix]
+ outputs = [_ts('%s%d' % (prefix, i), i + 1) for i in range(count)]
+ if count >= 3:
+ return PollResult.complete(outputs)
+ return PollResult.incomplete(outputs)
+
+
+def _complete_poll(prefix) -> PollResult[str]:
+ return PollResult.complete([_ts(prefix + 'a', 1), _ts(prefix + 'b', 2)])
+
+
+def _first_char(output):
+ return output[0]
+
+
+def _empty_poll(unused_element):
+ return PollResult.incomplete([])
+
+
+def _keyed_poll(prefix):
+ # 'a1' and 'a2' share the dedup key 'a', so only 'a1' is emitted.
+ return PollResult.complete([_ts('a1', 1), _ts('a2', 2), _ts('b1', 3)])
+
+
+class _StrCoderPollFn(PollFn):
+ def __call__(self, element):
+ return PollResult.complete([_ts(element + 'a', 1)])
+
+ def default_output_coder(self):
+ return StrUtf8Coder()
+
+
+class _NoDeterministicFormCoder(Coder):
+ def encode(self, value):
+ return b''
+
+ def decode(self, encoded):
+ return None
+
+ def is_deterministic(self):
+ return False
+
+
+def _windowed_group(kv, window=beam.DoFn.WindowParam):
+ return ((window.start, window.end), sorted(kv[1]))
+
+
+class WatchDoFnProcessTest(unittest.TestCase):
+ def _process(
+ self, poll_fn, element, timestamp, restriction=None, watermark=None):
+ dofn = _WatchGrowthDoFn(
+ poll_fn,
+ never(),
+ Duration(1),
+ StrUtf8Coder(),
+ _identity,
+ StrUtf8Coder())
+ if restriction is None:
+ restriction = dofn.initial_restriction(element)
+ threadsafe = ThreadsafeRestrictionTracker(dofn.create_tracker(restriction))
+ estimator = ThreadsafeWatermarkEstimator(
+ ManualWatermarkEstimator(watermark))
+ outputs = list(
+ dofn.process(
+ element,
+ timestamp=timestamp,
+ tracker=RestrictionTrackerView(threadsafe),
+ watermark_estimator=estimator))
+ return outputs, threadsafe, estimator
+
+ def test_empty_round_holds_watermark_at_input_timestamp(self):
+ outputs, threadsafe, estimator = self._process(
+ _empty_poll, 'in', Timestamp(7))
+ self.assertEqual([], outputs)
+ # The estimator is seeded from the input timestamp, so the deferred
+ # residual holds the watermark there instead of at MIN_TIMESTAMP.
+ self.assertEqual(Timestamp(7), estimator.current_watermark())
+ residual, _ = threadsafe.deferred_status()
+ self.assertIsInstance(residual, _PollingGrowthState)
+
+ def test_complete_round_stops_without_residual(self):
+ outputs, threadsafe, _ = self._process(_complete_poll, 'k:', Timestamp(0))
+ self.assertEqual([('k:', 'k:a'), ('k:', 'k:b')],
+ [value.value for value in outputs])
+ self.assertIsNone(threadsafe.deferred_status())
+ self.assertTrue(threadsafe.check_done())
+
+ def test_replay_round_leaves_the_watermark_alone(self):
+ pending = PollResult((_ts('k:a', 1), _ts('k:b', 2)), MAX_TIMESTAMP)
+ outputs, threadsafe, estimator = self._process(
+ _empty_poll,
+ 'k:',
+ Timestamp(7),
+ restriction=_NonPollingGrowthState(pending))
+ self.assertEqual([('k:', 'k:a'), ('k:', 'k:b')],
+ [value.value for value in outputs])
+ # The replay branch holds the watermark at the seed, so it never runs ahead
+ # of the replayed outputs and never releases to MAX_TIMESTAMP itself.
+ self.assertEqual(Timestamp(7), estimator.current_watermark())
+ self.assertIsNone(threadsafe.deferred_status())
+ self.assertTrue(threadsafe.check_done())
+
+ def test_terminal_round_after_deferring_leaves_no_residual(self):
+ _POLL_CALLS.clear()
+ # Round one defers and parks the watermark on the new output's time.
+ _, threadsafe, estimator = self._process(_growing_poll, 'd:', Timestamp(0))
+ residual, _ = threadsafe.deferred_status()
+ self.assertIsInstance(residual, _PollingGrowthState)
+ self.assertEqual(Timestamp(1), estimator.current_watermark())
+ # Round two resumes from that residual, carrying the watermark forward.
+ _, threadsafe, estimator = self._process(
+ _growing_poll,
+ 'd:',
+ Timestamp(0),
+ restriction=residual,
+ watermark=estimator.current_watermark())
+ residual, _ = threadsafe.deferred_status()
+ self.assertEqual(Timestamp(2), estimator.current_watermark())
+ # Round three completes. The watermark stays where round two left it and
+ # the round reports no residual, so nothing carries that hold forward.
+ outputs, threadsafe, estimator = self._process(
+ _growing_poll,
+ 'd:',
+ Timestamp(0),
+ restriction=residual,
+ watermark=estimator.current_watermark())
+ self.assertEqual([('d:', 'd:2')], [value.value for value in outputs])
+ self.assertEqual(Timestamp(2), estimator.current_watermark())
+ self.assertIsNone(threadsafe.deferred_status())
+ self.assertTrue(threadsafe.check_done())
+
+
+class WatchEndToEndTest(unittest.TestCase):
+ def _in_memory_pipeline(self):
+ return TestPipeline(
+ options=PipelineOptions(direct_running_mode='in_memory'))
+
+ def test_complete_outputs_values_and_timestamps(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch(_complete_poll, poll_interval=Duration(1)))
+ assert_that(
+ output,
+ equal_to([
+ TestWindowedValue(('k:', 'k:a'), Timestamp(1), [GlobalWindow()]),
+ TestWindowedValue(('k:', 'k:b'), Timestamp(2), [GlobalWindow()]),
+ ]),
+ reify_windows=True)
+
+ def test_complete_advances_watermark_for_windowed_pipeline(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch(_complete_poll, poll_interval=Duration(1)))
+ grouped = (
+ output
+ | beam.WindowInto(FixedWindows(10))
+ | beam.Map(lambda kv: ('all', kv[1]))
+ | beam.GroupByKey()
+ | beam.Map(_windowed_group))
+ assert_that(
+ grouped,
+ equal_to([
+ ((Timestamp(0), Timestamp(10)), ['k:a', 'k:b']),
+ ]))
+
+ def test_multi_round_dedups_stops_and_is_per_input(self):
+ _POLL_CALLS.clear()
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['x:', 'y:'])
+ | Watch(_growing_poll, poll_interval=Duration(0.05)))
+ assert_that(
+ output,
+ equal_to([('x:', 'x:0'), ('x:', 'x:1'), ('x:', 'x:2'), ('y:', 'y:0'),
+ ('y:', 'y:1'), ('y:', 'y:2')]))
+ self.assertEqual(3, _POLL_CALLS['x:'])
+ self.assertEqual(3, _POLL_CALLS['y:'])
+
+ def test_output_key_dedups_across_pipeline(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k'])
+ | Watch(
+ _keyed_poll, poll_interval=Duration(1),
+ output_key_fn=_first_char))
+ assert_that(output, equal_to([('k', 'a1'), ('k', 'b1')]))
+
+ def test_rejects_key_coder_without_deterministic_form(self):
+ with self.assertRaises(ValueError):
+ with self._in_memory_pipeline() as p:
+ _ = (
+ p | beam.Create(['k:'])
+ | Watch(
+ _complete_poll,
+ poll_interval=Duration(1),
+ output_key_coder=_NoDeterministicFormCoder()))
+
+ def test_infers_output_coder_from_return_annotation(self):
+ # _complete_poll is annotated ``-> PollResult[str]``, so the output coder
+ # and with it the (input, output) element type are inferred without hints.
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch(_complete_poll, poll_interval=Duration(1)))
+ self.assertEqual(typehints.Tuple[str, str], output.element_type)
+
+ def test_uses_poll_fn_default_output_coder(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch(_StrCoderPollFn(), poll_interval=Duration(1)))
+ self.assertEqual(typehints.Tuple[str, str], output.element_type)
+
+
+if __name__ == '__main__':
+ unittest.main()
