diff --git a/sdks/python/apache_beam/io/fileio.py b/sdks/python/apache_beam/io/fileio.py index a333b7c89775..e57577d34769 100644 --- a/sdks/python/apache_beam/io/fileio.py +++ b/sdks/python/apache_beam/io/fileio.py @@ -93,7 +93,6 @@ import random import uuid from collections import namedtuple -from functools import partial from typing import Any from typing import BinaryIO # pylint: disable=unused-import from typing import Callable @@ -101,20 +100,28 @@ from typing import Union import apache_beam as beam +from apache_beam.coders.coders import StrUtf8Coder +from apache_beam.coders.coders import VarIntCoder from apache_beam.io import filesystem from apache_beam.io import filesystems from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressionTypes +from apache_beam.io.watch import PollFn +from apache_beam.io.watch import PollResult +from apache_beam.io.watch import TerminationCondition +from apache_beam.io.watch import Watch +from apache_beam.io.watch import never from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.value_provider import StaticValueProvider from apache_beam.options.value_provider import ValueProvider from apache_beam.transforms.periodicsequence import PeriodicImpulse -from apache_beam.transforms.userstate import CombiningValueStateSpec from apache_beam.transforms.window import BoundedWindow from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import GlobalWindow from apache_beam.transforms.window import IntervalWindow +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__ = [ @@ -251,6 +258,79 @@ def process( yield ReadableFile(metadata, self._compression) +class _WatchWindowTermination(TerminationCondition): + """Stops after the polls that fall in the ``[start, stop)`` window. + + ``max_polls`` is the ``PeriodicImpulse`` tick count + ``ceil((stop - start) / interval)``, so the number of polls is independent of + how fast the runner reschedules deferred work. Only polls at or after + ``start`` count toward the budget; earlier polls are deferred waits that must + not consume it, matching ``PeriodicImpulse``, which never advances a tick + while waiting for the start time. + """ + def __init__(self, start_micros: int, max_polls: int): + self._start_micros = start_micros + self._max_polls = max_polls + + def for_new_input(self, now, element): + return 0 + + def on_poll_complete(self, now, state): + if now.micros >= self._start_micros: + return state + 1 + return state + + def can_stop_polling(self, now, state): + return state >= self._max_polls + + def state_coder(self): + return VarIntCoder() + + +def _file_path_key(metadata: filesystem.FileMetadata) -> str: + return metadata.path + + +class _MatchContinuouslyPollFn(PollFn): + """Polls a file pattern for ``MatchContinuously``, honoring empty-match rules. + + Emits no outputs before ``start_timestamp`` so polling can start ahead of the + first intended match. Normally each match is stamped with the poll time as + its event time. When matching updated files, each match is stamped with the + file's last-modified time so ``Watch(use_timestamp=True)`` can emit the same + path again when that timestamp changes. The watermark advances to the poll + time so downstream event-time windows keep progressing even when a poll finds + no new files. + """ + def __init__( + self, + empty_match_treatment, + start_timestamp, + use_metadata_timestamp: bool = False): + self._empty_match_treatment = empty_match_treatment + self._start_micros = Timestamp.of(start_timestamp).micros + self._use_metadata_timestamp = use_metadata_timestamp + + def __call__(self, file_pattern: str) -> PollResult: + now = Timestamp.now() + if now.micros < self._start_micros: + return PollResult.incomplete(()) + match_result = filesystems.FileSystems.match([file_pattern])[0] + if (not match_result.metadata_list and + not EmptyMatchTreatment.allow_empty_match(file_pattern, + self._empty_match_treatment)): + raise BeamIOError( + 'Empty match for pattern %s. Disallowed.' % file_pattern) + if self._use_metadata_timestamp: + outputs = [ + TimestampedValue(metadata, metadata.last_updated_in_seconds) + for metadata in match_result.metadata_list + ] + return PollResult.incomplete(outputs).with_watermark(now) + return PollResult.incomplete( + match_result.metadata_list, timestamp=now).with_watermark(now) + + class MatchContinuously(beam.PTransform): """Checks for new files for a given pattern every interval. @@ -306,37 +386,72 @@ def __init__( 'if possible') def expand(self, pbegin) -> beam.PCollection[filesystem.FileMetadata]: - # invoke periodic impulse - impulse = pbegin | PeriodicImpulse( - start_timestamp=self.start_ts, - stop_timestamp=self.stop_ts, - fire_interval=self.interval) - - # match file pattern periodically - file_pattern = self.file_pattern - match_files = ( - impulse - | 'GetFilePattern' >> beam.Map(lambda x: file_pattern) - | MatchAll(self.empty_match_treatment)) - - # apply deduplication strategy if required if self.has_deduplication: - # Making a Key Value so each file has its own state. - match_files = match_files | 'ToKV' >> beam.Map(lambda x: (x.path, x)) - if self.match_upd: - match_files = match_files | 'RemoveOldAlreadyRead' >> beam.ParDo( - _RemoveOldDuplicates()) - else: - match_files = match_files | 'RemoveAlreadyRead' >> beam.ParDo( - _RemoveDuplicates()) - - # apply windowing if required. Apply at last because deduplication relies on - # the global window. + match_files = self._match_deduplicated(pbegin) + else: + match_files = self._match_all_each_poll(pbegin) + + # Apply windowing last because dedup relies on the global window. if self.apply_windowing: match_files = match_files | beam.WindowInto(FixedWindows(self.interval)) return match_files + def _match_deduplicated(self, + pbegin) -> beam.PCollection[filesystem.FileMetadata]: + # The Watch transform polls the pattern and emits each file once per path. + # When matching updated files, the poll function timestamps each output with + # the file's last-modified time and Watch includes that timestamp in dedup. + # stop_timestamp bounds the watch to the polls that fall in [start, stop). + if self.stop_ts == MAX_TIMESTAMP: + termination = never() + else: + start_ts = Timestamp.of(self.start_ts) + stop_ts = Timestamp.of(self.stop_ts) + if stop_ts < start_ts: + raise ValueError( + 'MatchContinuously stop_timestamp %s precedes start_timestamp %s' % + (stop_ts, start_ts)) + interval_micros = Duration.of(self.interval).micros + if interval_micros <= 0: + raise ValueError('MatchContinuously interval must be positive.') + span_micros = (stop_ts - start_ts).micros + # Ceiling division reproduces PeriodicImpulse's tick count; the window + # upper bound is exclusive. + max_polls = -(-span_micros // interval_micros) + termination = _WatchWindowTermination(start_ts.micros, max_polls) + file_pattern = self.file_pattern + poll_fn = _MatchContinuouslyPollFn( + self.empty_match_treatment, self.start_ts, self.match_upd) + watch = Watch( + poll_fn, + poll_interval=self.interval, + termination=termination, + output_key_fn=_file_path_key, + output_key_coder=StrUtf8Coder(), + use_timestamp=self.match_upd) + # Watch emits (pattern, file) pairs; keep the FileMetadata output type so + # downstream transforms stay typed instead of falling back to Any. + return ( + pbegin + | 'Impulse' >> beam.Create([file_pattern]) + | 'Watch' >> watch + | 'DropPattern' >> beam.Map(lambda kv: kv[1]).with_output_types( + filesystem.FileMetadata)) + + def _match_all_each_poll(self, + pbegin) -> beam.PCollection[filesystem.FileMetadata]: + # No deduplication: re-emit every match on each poll. + file_pattern = self.file_pattern + return ( + pbegin + | PeriodicImpulse( + start_timestamp=self.start_ts, + stop_timestamp=self.stop_ts, + fire_interval=self.interval) + | 'GetFilePattern' >> beam.Map(lambda x: file_pattern) + | MatchAll(self.empty_match_treatment)) + class ReadMatches(beam.PTransform): """Converts each result of MatchFiles() or MatchAll() to a ReadableFile. @@ -892,50 +1007,3 @@ def finish_bundle(self): timestamp=key[1].start, windows=[key[1]] # TODO(pabloem) HOW DO WE GET THE PANE )) - - -class _RemoveDuplicates(beam.DoFn): - """Internal DoFn that filters out filenames already seen (even though the file - has updated).""" - COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum) - - def process( - self, - element: tuple[str, filesystem.FileMetadata], - count_state=beam.DoFn.StateParam(COUNT_STATE) - ) -> Iterable[filesystem.FileMetadata]: - - path = element[0] - file_metadata = element[1] - counter = count_state.read() - - if counter == 0: - count_state.add(1) - _LOGGER.debug('Generated entry for file %s', path) - yield file_metadata - else: - _LOGGER.debug('File %s was already read, seen %d times', path, counter) - - -class _RemoveOldDuplicates(beam.DoFn): - """Internal DoFn that filters out filenames already seen and timestamp - unchanged.""" - TIME_STATE = CombiningValueStateSpec( - 'count', combine_fn=partial(max, default=0.0)) - - def process( - self, - element: tuple[str, filesystem.FileMetadata], - time_state=beam.DoFn.StateParam(TIME_STATE) - ) -> Iterable[filesystem.FileMetadata]: - path = element[0] - file_metadata = element[1] - new_ts = file_metadata.last_updated_in_seconds - old_ts = time_state.read() - - if old_ts < new_ts: - time_state.add(new_ts) - _LOGGER.debug('Generated entry for file %s', path) - yield file_metadata - else: - _LOGGER.debug('File %s was already read', path) diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py index ce535265ef2f..a8e4d71c0adc 100644 --- a/sdks/python/apache_beam/io/fileio_test.py +++ b/sdks/python/apache_beam/io/fileio_test.py @@ -34,6 +34,7 @@ import apache_beam as beam from apache_beam.io import fileio from apache_beam.io.filebasedsink_test import _TestCaseWithTempDirCleanUp +from apache_beam.io.filesystem import BeamIOError from apache_beam.io.filesystem import CompressionTypes from apache_beam.io.filesystems import FileSystems from apache_beam.options.pipeline_options import PipelineOptions @@ -420,6 +421,65 @@ def _create_extra_file(element): assert_that(match_continiously, equal_to(files)) + def test_poll_fn_gates_on_start_timestamp(self): + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + self._create_temp_file(dir=tempdir) + pattern = FileSystems.join(tempdir, '*') + + future_start = fileio._MatchContinuouslyPollFn( + fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() + 3600) + self.assertEqual((), future_start(pattern).outputs) + + past_start = fileio._MatchContinuouslyPollFn( + fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() - 3600) + self.assertEqual(1, len(past_start(pattern).outputs)) + + def test_poll_fn_disallows_empty_match(self): + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + poll_fn = fileio._MatchContinuouslyPollFn( + fileio.EmptyMatchTreatment.DISALLOW, Timestamp.now() - 3600) + with self.assertRaises(BeamIOError): + poll_fn(FileSystems.join(tempdir, 'no-such-file')) + + def test_poll_fn_can_timestamp_outputs_with_file_mtime(self): + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + self._create_temp_file(dir=tempdir) + poll_fn = fileio._MatchContinuouslyPollFn( + fileio.EmptyMatchTreatment.ALLOW, + Timestamp.now() - 3600, + use_metadata_timestamp=True) + result = poll_fn(FileSystems.join(tempdir, '*')) + self.assertEqual(1, len(result.outputs)) + output = result.outputs[0] + self.assertEqual( + Timestamp.of(output.value.last_updated_in_seconds), output.timestamp) + + def test_watch_window_termination_ignores_pre_start_polls(self): + # Polls before start_timestamp are deferred waits and must not consume the + # budget, otherwise a future start_timestamp silently drops all output. + start_micros = Timestamp.of(1000).micros + term = fileio._WatchWindowTermination(start_micros, max_polls=2) + before = Timestamp.of(999) + after = Timestamp.of(1000) + state = term.for_new_input(before, 'pattern') + state = term.on_poll_complete(before, state) + state = term.on_poll_complete(before, state) + self.assertFalse(term.can_stop_polling(before, state)) + state = term.on_poll_complete(after, state) + self.assertFalse(term.can_stop_polling(after, state)) + state = term.on_poll_complete(after, state) + self.assertTrue(term.can_stop_polling(after, state)) + + def test_poll_fn_advances_watermark_on_empty_match(self): + # An empty (but allowed) match still carries a watermark so downstream + # event-time windows keep progressing when no new files appear. + tempdir = '%s%s' % (self._new_tempdir(), os.sep) + poll_fn = fileio._MatchContinuouslyPollFn( + fileio.EmptyMatchTreatment.ALLOW, Timestamp.now() - 3600) + result = poll_fn(FileSystems.join(tempdir, '*')) + self.assertEqual((), result.outputs) + self.assertIsNotNone(result.watermark) + class WriteFilesTest(_TestCaseWithTempDirCleanUp): diff --git a/sdks/python/apache_beam/io/watch.py b/sdks/python/apache_beam/io/watch.py new file mode 100644 index 000000000000..5bbc83fe73ae --- /dev/null +++ b/sdks/python/apache_beam/io/watch.py @@ -0,0 +1,759 @@ +# +# 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. When ``use_timestamp`` is true, +the output event time is part of the dedup key. + +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, now: Timestamp, 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 + + +_HASH_KEY_AND_TIMESTAMP_CODER = TupleCoder( + [coders.BytesCoder(), TimestampCoder()]) + + +def _hash_output( + key_coder: Coder, + value: Any, + timestamp: Optional[Timestamp] = None, + use_timestamp: bool = False) -> bytes: + encoded_key = key_coder.encode(value) + if use_timestamp: + encoded_key = _HASH_KEY_AND_TIMESTAMP_CODER.encode( + (encoded_key, Timestamp.of(timestamp))) + return hashlib.blake2b(encoded_key, 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, + use_timestamp: bool = False) -> PollResult: + """Filters a poll result down to outputs whose key was never seen before. + + Dedup hashes ``key_fn(output.value)`` (or that key plus ``output.timestamp`` + when requested) 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), output.timestamp, use_timestamp) + 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, + use_timestamp: bool = False): + self._restriction = restriction + self._key_fn = key_fn + self._key_coder = key_coder + self._use_timestamp = use_timestamp + 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, output: TimestampedValue) -> bytes: + return _hash_output( + self._key_coder, + self._key_fn(output.value), + output.timestamp, + self._use_timestamp) + + 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)] = 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) 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, + use_timestamp: bool = False, + 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._use_timestamp = use_timestamp + 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, self._use_timestamp) + + 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, self._use_timestamp) + 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( + now, 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. + use_timestamp: whether to include each output timestamp in the dedup key. + Defaults to ``False``, so equal output keys are emitted only once. + 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, + use_timestamp: bool = False, + 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._use_timestamp = use_timestamp + 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._use_timestamp, + 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..f8d5ed7f9380 --- /dev/null +++ b/sdks/python/apache_beam/io/watch_test.py @@ -0,0 +1,416 @@ +# +# 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, use_timestamp=False): + return _never_seen_before( + restriction, result, key_fn or _identity, StrUtf8Coder(), use_timestamp) + + +def _tracker(restriction, use_timestamp=False): + return _GrowthRestrictionTracker( + restriction, _identity, StrUtf8Coder(), use_timestamp) + + +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_use_timestamp_dedups_by_key_and_timestamp(self): + result = PollResult.incomplete([_ts('a', 1), _ts('a', 2), _ts('a', 1)]) + new_results = _new_results(_initial_polling(), result, use_timestamp=True) + self.assertEqual([('a', Timestamp(1)), ('a', Timestamp(2))], + [(o.value, o.timestamp) 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): + 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( + 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()