diff --git a/sdks/python/apache_beam/io/watch.py b/sdks/python/apache_beam/io/watch.py new file mode 100644 index 000000000000..2fcee7a8080f --- /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. +# ------------------------------------------------------------------------------ + + +@dataclasses.dataclass(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.""" + + +@dataclasses.dataclass(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 + + +@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 + + +# 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 000000000000..8c1f6571da66 --- /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()