You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "Eliaaazzz (via GitHub)" <gi...@apache.org> on 2026/06/18 16:06:20 UTC
[PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Eliaaazzz opened a new pull request, #39023:
URL: https://github.com/apache/beam/pull/39023
This adds an experimental Watch transform to the Python SDK. It addresses #21521.
Watch watches a growing set of outputs for each input element. It calls a user poll function on an interval and emits an unbounded PCollection of input and output pairs. Polling stops per input when the poll reports completion or when a termination condition fires. The transform is a splittable DoFn, so each process call performs one poll round and then self checkpoints with defer_remainder, which keeps the polling loop durable across bundles. New outputs are deduplicated with a stable 128 bit blake2b hash of the encoded output, and a manual watermark estimator advances once per poll.
This is an initial minimal version intended for review and iteration. It implements growth_of, PollResult, PollFn, the never and after_total_of termination conditions, the growth state restriction with its coder, and the polling splittable DoFn. It is validated end to end on the DirectRunner. The output coder must be deterministic for dedup to hold across workers, and the transform logs a warning when the resolved coder is not deterministic.
The following are out of scope for this initial change and are not included here: side input poll functions, the remaining termination conditions, a separate bounded stage for exploding large poll results, and bounded growth state eviction.
Testing
- Unit tests for the termination conditions and for the growth state coder round trips, including ordered completed hashes and a fixed width 16 byte hash coder.
- Restriction tracker tests for claim, dedup, checkpoint, replay, boundedness, and a terminal split residual for every stop reason, so a checkpoint after a terminal poll cannot resume polling.
- DirectRunner end to end tests for single poll completion, multiple inputs, and multi round dedup.
- Local run reports 32 tests passing and pylint at 10.00 of 10.
------------------------
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
- [x] Mention the appropriate issue in your description (for example: `addresses #123`), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment `fixes #<ISSUE NUMBER>` instead.
- [ ] Update `CHANGES.md` with noteworthy changes.
- [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://github.com/apache/beam/blob/master/CONTRIBUTING.md#make-the-reviewers-job-easier).
To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
GitHub Actions Tests Status (on master branch)
------------------------------------------------------------------------------------------------
[](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI or the [workflows README](https://github.com/apache/beam/blob/master/.github/workflows/README.md) to see a list of phrases to trigger workflows.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3575511681
##########
sdks/python/apache_beam/io/watch_test.py:
##########
@@ -0,0 +1,307 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for the Watch transform."""
+
+import collections
+import unittest
+
+import apache_beam as beam
+from apache_beam.coders.coders import StrUtf8Coder
+from apache_beam.io.watch import PollResult
+from apache_beam.io.watch import Watch
+from apache_beam.io.watch import _GrowthRestrictionTracker
+from apache_beam.io.watch import _GrowthStateCoder
+from apache_beam.io.watch import _NonPollingGrowthState
+from apache_beam.io.watch import _plan_poll_round
+from apache_beam.io.watch import _PollingGrowthState
+from apache_beam.io.watch import _replay_plan
+from apache_beam.io.watch import after_total_of
+from apache_beam.io.watch import never
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.runners.sdf_utils import RestrictionTrackerView
+from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import TestWindowedValue
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+from apache_beam.transforms.window import FixedWindows
+from apache_beam.transforms.window import GlobalWindow
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.typehints import typehints
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+
+def _ts(value, timestamp):
+ return TimestampedValue(value, Timestamp(timestamp))
+
+
+def _plan(restriction, poll_fn, now=0.0, termination=None):
+ termination = termination or never()
+ return _plan_poll_round(
+ restriction,
+ poll_fn('input'),
+ termination,
+ StrUtf8Coder(),
+ Timestamp.of(now))
+
+
+def _initial_polling(termination=None, now=Timestamp(0)):
+ termination = termination or never()
+ return _PollingGrowthState(
+ collections.OrderedDict(), None, termination.for_new_input(now, 'input'))
+
+
+class GrowthStateCoderTest(unittest.TestCase):
+ def test_polling_round_trip_preserves_resume_state(self):
+ termination = after_total_of(Duration(30))
+ coder = _GrowthStateCoder(StrUtf8Coder(), termination)
+ completed = collections.OrderedDict([
+ (b'a' * 16, Timestamp(1)),
+ (b'b' * 16, Timestamp(2)),
+ (b'c' * 16, Timestamp(3)),
+ ])
+ termination_state = termination.for_new_input(Timestamp(7), 'input')
+ state = _PollingGrowthState(completed, Timestamp(5), termination_state)
+ decoded = coder.decode(coder.encode(state))
+ self.assertEqual(list(completed.items()), list(decoded.completed.items()))
+ self.assertEqual(Timestamp(5), decoded.poll_watermark)
+ self.assertEqual(termination_state, decoded.termination_state)
+
+ def test_non_polling_round_trip_preserves_pending_outputs(self):
+ coder = _GrowthStateCoder(StrUtf8Coder(), never())
+ pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP)
+ state = _NonPollingGrowthState(pending)
+ decoded = coder.decode(coder.encode(state))
+ self.assertEqual(MAX_TIMESTAMP, decoded.pending.watermark)
+ self.assertEqual([('a', Timestamp(1)), ('b', Timestamp(2))],
+ [(o.value, o.timestamp) for o in decoded.pending.outputs])
+
+
+class GrowthTrackerTest(unittest.TestCase):
+ def test_poll_claims_dedups_and_checkpoints(self):
+ def poll(unused_element):
+ return PollResult.incomplete([_ts('a', 1), _ts('a', 1), _ts('b', 2)])
+
+ plan = _plan(_initial_polling(), poll)
+ self.assertEqual(['a', 'b'], [o.value for o in plan.emit])
+ self.assertFalse(plan.stop)
+ self.assertEqual(Timestamp(1), plan.watermark)
+ self.assertIsInstance(plan.primary, _NonPollingGrowthState)
+ self.assertIsInstance(plan.residual, _PollingGrowthState)
+ self.assertEqual(2, len(plan.residual.completed))
+
+ tracker = _GrowthRestrictionTracker(_initial_polling())
+ self.assertFalse(tracker.is_bounded())
+ self.assertTrue(tracker.try_claim(plan))
+ primary, residual = tracker.try_split(0)
+ self.assertIs(plan.primary, primary)
+ self.assertIs(plan.residual, residual)
+ self.assertTrue(tracker.check_done())
+
+ def explicit_watermark_poll(unused_element):
+ return PollResult.incomplete([_ts('c', 3)]).with_watermark(5)
+
+ plan = _plan(_initial_polling(), explicit_watermark_poll)
+ self.assertEqual(Timestamp(5), plan.watermark)
+ self.assertEqual(Timestamp(5), plan.residual.poll_watermark)
+
+ def test_second_round_repolls_and_dedups_against_completed(self):
+ polls = []
+
+ def poll(unused_element):
+ polls.append(len(polls))
+ if len(polls) == 1:
+ return PollResult.incomplete([_ts('a', 1), _ts('b', 2)])
+ return PollResult.incomplete([_ts('a', 1), _ts('c', 3)])
+
+ first = _plan(_initial_polling(), poll)
+ resumed = _plan(first.residual, poll)
+ self.assertEqual(2, len(polls))
+ self.assertEqual(['c'], [o.value for o in resumed.emit])
+
+ def test_termination_condition_sets_stop(self):
+ def poll(unused_element):
+ return PollResult.incomplete([_ts('a', 1)])
+
+ termination = after_total_of(10)
+ for now, expected_stop in [(10.0, False), (11.0, True)]:
+ with self.subTest(now=now):
+ plan = _plan(
+ _initial_polling(termination, Timestamp(0)),
+ poll,
+ now=now,
+ termination=termination)
+ self.assertEqual(expected_stop, plan.stop)
+
+ def test_non_polling_replays(self):
+ pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP)
+ restriction = _NonPollingGrowthState(pending)
+ plan = _replay_plan(restriction)
+ self.assertEqual(['a', 'b'], [o.value for o in plan.emit])
+ self.assertTrue(plan.stop)
+
+ tracker = _GrowthRestrictionTracker(restriction)
+ self.assertTrue(tracker.is_bounded())
+ self.assertTrue(tracker.try_claim(plan))
+ # A checkpoint after replay resumes an empty state that emits nothing.
+ primary, residual = tracker.try_split(0)
+ self.assertIs(restriction, primary)
+ self.assertEqual((), residual.pending.outputs)
+ self.assertTrue(tracker.check_done())
+
+ def test_terminal_split_residual_is_empty_for_all_stop_causes(self):
+ termination = after_total_of(Duration(10))
+ cases = [
+ (
+ 'reached_max',
+ never(),
+ _initial_polling(), lambda element: PollResult(
+ (TimestampedValue('a', MAX_TIMESTAMP), ), None),
+ 0.0),
+ (
+ 'complete',
+ never(),
+ _initial_polling(),
+ lambda element: PollResult.complete([_ts('a', 1)]),
+ 0.0),
+ (
+ 'after_total_of',
+ termination,
+ _initial_polling(termination, Timestamp(0)),
+ lambda element: PollResult.incomplete([_ts('a', 1)]),
+ 100.0),
+ ]
+ for name, condition, restriction, poll_fn, now in cases:
+ with self.subTest(name=name):
+ plan = _plan(restriction, poll_fn, now=now, termination=condition)
+ self.assertTrue(plan.stop)
+ self.assertIsInstance(plan.residual, _NonPollingGrowthState)
+ self.assertEqual((), plan.residual.pending.outputs)
+
+ def test_wrapper_chain_defers_merged_residual(self):
+ def poll(unused_element):
+ return PollResult.incomplete([_ts('a', 1), _ts('b', 2)])
+
+ plan = _plan(_initial_polling(), poll)
+ threadsafe = ThreadsafeRestrictionTracker(
+ _GrowthRestrictionTracker(_initial_polling()))
+ view = RestrictionTrackerView(threadsafe)
+ self.assertTrue(view.try_claim(plan))
+ view.defer_remainder(Duration(5))
+ residual, _ = threadsafe.deferred_status()
+ self.assertIsInstance(residual, _PollingGrowthState)
+ self.assertEqual(2, len(residual.completed))
+
+
+# Module-level so the poll function pickles by reference; the call counter is
+# shared within the single in-memory DirectRunner process.
+_POLL_CALLS = collections.defaultdict(int)
+
+
+def _growing_poll(prefix):
+ _POLL_CALLS[prefix] += 1
+ count = _POLL_CALLS[prefix]
+ outputs = [_ts('%s%d' % (prefix, i), i + 1) for i in range(count)]
+ if count >= 3:
+ return PollResult.complete(outputs)
+ return PollResult.incomplete(outputs)
+
+
+def _complete_poll(prefix):
+ return PollResult.complete([_ts(prefix + 'a', 1), _ts(prefix + 'b', 2)])
+
+
+def _windowed_group(kv, window=beam.DoFn.WindowParam):
+ return ((window.start, window.end), sorted(kv[1]))
+
+
+class WatchEndToEndTest(unittest.TestCase):
+ def _in_memory_pipeline(self):
+ return TestPipeline(
+ options=PipelineOptions(direct_running_mode='in_memory'))
+
+ def test_complete_outputs_values_and_timestamps(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(
+ Duration(1)).with_output_coder(StrUtf8Coder()))
+ assert_that(
+ output,
+ equal_to([
+ TestWindowedValue(('k:', 'k:a'), Timestamp(1), [GlobalWindow()]),
+ TestWindowedValue(('k:', 'k:b'), Timestamp(2), [GlobalWindow()]),
+ ]),
+ reify_windows=True)
+
+ def test_complete_advances_watermark_for_windowed_pipeline(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(
+ Duration(1)).with_output_coder(StrUtf8Coder()))
+ grouped = (
+ output
+ | beam.WindowInto(FixedWindows(10))
+ | beam.Map(lambda kv: ('all', kv[1]))
+ | beam.GroupByKey()
+ | beam.Map(_windowed_group))
+ assert_that(
+ grouped,
+ equal_to([
+ ((Timestamp(0), Timestamp(10)), ['k:a', 'k:b']),
+ ]))
+
+ def test_multi_round_dedups_stops_and_is_per_input(self):
+ _POLL_CALLS.clear()
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['x:', 'y:'])
+ | Watch.growth_of(_growing_poll).with_poll_interval(
+ Duration(0.05)).with_output_coder(StrUtf8Coder()))
+ assert_that(
+ output,
+ equal_to([('x:', 'x:0'), ('x:', 'x:1'), ('x:', 'x:2'), ('y:', 'y:0'),
+ ('y:', 'y:1'), ('y:', 'y:2')]))
+ self.assertEqual(3, _POLL_CALLS['x:'])
+ self.assertEqual(3, _POLL_CALLS['y:'])
+
+ def test_rejects_non_deterministic_output_coder(self):
+ # No output coder resolves to PickleCoder, which is non-deterministic, so
+ # dedup could re-emit equal outputs. Expansion must reject it.
+ with self.assertRaises(ValueError):
+ with self._in_memory_pipeline() as p:
+ _ = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(Duration(1)))
+
+ def test_derives_output_type_from_input_and_coder(self):
+ # expand() types the (input, output) pairs from the input type and the
+ # resolved coder, so downstream stays typed without a manual hint.
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(
+ Duration(1)).with_output_coder(StrUtf8Coder()))
Review Comment:
I notice many `with_output_coder(StrUtf8Coder())` added in latest diff. This rather suggests coder inference is either broken or not working in the latest change.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ # Poll before claiming so a slow poll never holds the tracker lock, which
+ # would block runner progress checks and checkpoints.
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+ plan = _plan_poll_round(
+ restriction, result, self._termination, self._key_coder, now)
+ if not tracker.try_claim(plan):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in plan.emit:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if plan.stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ return
+ if plan.watermark is not None:
+ _set_watermark_if_greater(watermark_estimator, plan.watermark)
+ tracker.defer_remainder(self._poll_interval)
+
+
+def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None:
+ # set_watermark raises on regression, so only ever advance the watermark.
+ current = watermark_estimator.current_watermark()
+ if current is None or new_watermark > current:
+ watermark_estimator.set_watermark(new_watermark)
+
+
+# ------------------------------------------------------------------------------
+# Public PTransform.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
+ self._termination = termination or never()
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._now = now_fn
+
+ @classmethod
+ def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch':
+ return cls(poll_fn)
+
+ def _replace(self, **changes) -> 'Watch':
+ spec = dict(
+ poll_fn=self._poll_fn,
+ termination=self._termination,
+ poll_interval=self._poll_interval,
+ output_coder=self._output_coder,
+ now_fn=self._now)
+ spec.update(changes)
+ return Watch(**spec)
+
+ def with_poll_interval(self, poll_interval) -> 'Watch':
Review Comment:
Constructor already have "poll_interval", "output_coder", etc. No need providing another methods to set them (Java withXXX() way). Referring to example utility transforms (apache_beam/transform), we use Pythonic way to construct transform.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ # Poll before claiming so a slow poll never holds the tracker lock, which
+ # would block runner progress checks and checkpoints.
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+ plan = _plan_poll_round(
+ restriction, result, self._termination, self._key_coder, now)
+ if not tracker.try_claim(plan):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in plan.emit:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if plan.stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ return
+ if plan.watermark is not None:
+ _set_watermark_if_greater(watermark_estimator, plan.watermark)
+ tracker.defer_remainder(self._poll_interval)
+
+
+def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None:
+ # set_watermark raises on regression, so only ever advance the watermark.
+ current = watermark_estimator.current_watermark()
+ if current is None or new_watermark > current:
+ watermark_estimator.set_watermark(new_watermark)
+
+
+# ------------------------------------------------------------------------------
+# Public PTransform.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
+ self._termination = termination or never()
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._now = now_fn
+
+ @classmethod
+ def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch':
+ return cls(poll_fn)
+
+ def _replace(self, **changes) -> 'Watch':
+ spec = dict(
+ poll_fn=self._poll_fn,
+ termination=self._termination,
+ poll_interval=self._poll_interval,
+ output_coder=self._output_coder,
+ now_fn=self._now)
+ spec.update(changes)
+ return Watch(**spec)
+
+ def with_poll_interval(self, poll_interval) -> 'Watch':
+ return self._replace(poll_interval=_as_duration(poll_interval))
+
+ def with_termination_per_input(
+ self, termination: TerminationCondition) -> 'Watch':
+ return self._replace(termination=termination)
+
+ def with_output_coder(self, output_coder: Coder) -> 'Watch':
+ return self._replace(output_coder=output_coder)
+
+ def expand(self, pcoll):
+ if self._poll_interval is None:
+ raise ValueError('Watch requires with_poll_interval(...)')
+ output_coder = self._output_coder
+ if output_coder is None:
+ hint = self._poll_fn.default_output_coder() if isinstance(
+ self._poll_fn, PollFn) else None
+ output_coder = hint or coders.PickleCoder()
+ # Dedup hashes the encoded output as its key, so a non-deterministic coder
+ # would hash equal outputs differently and re-emit them, defeating the
+ # transform. Require determinism rather than silently emitting duplicates.
+ if not output_coder.is_deterministic():
Review Comment:
It's still not the right check. We only need to make sure key_coder isdeterministic here.
I think what happens to the code completion is that it doesn't understand the difference between key_coder and output_coder
This line in the
```
`self._key_coder = output_coder`
```
also suggests this.
Please check Java implication again. `outputKeyCoder` is involved when `OutputKeyFn` is configured. Otherwise dedup key defaults to be the same as value, and only in this case self._key_coder is equivalent to output_coder.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
Review Comment:
We need to check again this newly introduced PollPlan. From the pydoc it sounds an ad-hoc solution to put poll outside of tryClaim. Constructively we should just be able to use PollResult as its Restriction Tracker. It now needs other information now likely because the responsibility boundary of process and try_claim is still not clear
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Eliaaazzz (via GitHub)" <gi...@apache.org>.
Eliaaazzz commented on PR #39023:
URL: https://github.com/apache/beam/pull/39023#issuecomment-4969109100
Sorry about the force pushes earlier, I did not realize they break the incremental review, and I will keep fixes as follow up commits from now on. This round is commit c559963 on top of the reviewed one. It covers the four comments, and the tests and lint pass locally.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Eliaaazzz (via GitHub)" <gi...@apache.org>.
Eliaaazzz commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3504091699
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
Review Comment:
That is a cleaner design, thank you. I replaced the `Any` alias with an empty `_GrowthState` base class and had both restriction states inherit it, so it now reads as a real union type.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
Review Comment:
You have a point that the wrapper adds little in Python. I switched to `BytesCoder`, and since the digest is always 16 bytes from `blake2b`, I dropped the length guard as redundant.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
Review Comment:
Thank you for the explanation, that helps a lot. Since `TimestampedValue` gets unwrapped into a `WindowedValue`, I reworded the docstring to say the coder only exists to keep it inside the restriction state. Please let me know if you would like me to build on your public coder PR once it lands.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ The Python SDK ships no coder for ``TimestampedValue``, 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([_HashCode128Coder(), 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.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Drives one input's polling loop.
+
+ ``process()`` only sees a ``RestrictionTrackerView`` whose ``try_claim``
+ returns a bool, so the poll happens inside ``try_claim`` and its result is
+ returned through a two-slot holder list passed as the claim position:
+ ``holder[0]`` carries the input element in, ``holder[1]`` carries the work
+ out. At most one claim succeeds per ``process()``.
+
+ The poll runs while the tracker lock is held, so a ``PollFn`` must be bounded
+ or timeout-safe; a blocking poll delays runner-initiated checkpoints.
+ """
+ def __init__(
+ self,
+ restriction: _GrowthState,
+ poll_fn: Callable[[Any], PollResult],
+ key_coder: Coder,
+ termination: TerminationCondition,
+ now_fn: Callable[[], float]):
+ self._restriction = restriction
+ self._poll_fn = poll_fn
+ self._key_coder = key_coder
+ self._termination = termination
+ self._now = now_fn
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def _hash_output(self, value: Any) -> bytes:
+ return hashlib.blake2b(
+ self._key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+ def try_claim(self, holder: list) -> bool:
+ """Performs one poll round (or one replay) and reports it via ``holder``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation,
+ in which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ restriction = self._restriction
+ if isinstance(restriction, _NonPollingGrowthState):
+ holder[1] = ('replay', restriction.pending)
+ self._should_stop = True
+ return True
+
+ element = holder[0]
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = self._hash_output(output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = self._termination.on_seen_new_output(
+ now, termination_state)
+ termination_state = self._termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ self._termination.can_stop_polling(now, termination_state))
+
+ self._primary = _NonPollingGrowthState(
+ PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-
+ # initiated or via defer_remainder) resumes a state that emits nothing.
+ self._residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = self._max_watermark(
+ restriction.poll_watermark, watermark)
+ self._residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ holder[1] = ('poll', new_outputs, watermark, stop)
+ self._should_stop = True
+ return True
+
+ @staticmethod
+ 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 try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(
+ restriction,
+ self._poll_fn,
+ self._key_coder,
+ self._termination,
+ self._now)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ holder = [element, None]
+ if not tracker.try_claim(holder):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ work = holder[1]
+ if work[0] == 'replay':
+ for output in work[1].outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ new_outputs, watermark, stop = work[1], work[2], work[3]
+ for output in new_outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ 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.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
Review Comment:
Thank you, I took another pass at this. In `expand` I now derive the output element type from the input type and the resolved coder and set it with `with_output_types`, so the `(input, output)` pairs stay typed downstream. A `PollFn` can still declare its coder through `default_output_coder`, and I added a test that checks the derived type.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
Review Comment:
That is a helpful suggestion, thank you. I added a short example to the `PollFn` docstring that shows a callable alongside a `default_output_coder` implementation.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ The Python SDK ships no coder for ``TimestampedValue``, 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([_HashCode128Coder(), 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.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Drives one input's polling loop.
+
+ ``process()`` only sees a ``RestrictionTrackerView`` whose ``try_claim``
+ returns a bool, so the poll happens inside ``try_claim`` and its result is
+ returned through a two-slot holder list passed as the claim position:
+ ``holder[0]`` carries the input element in, ``holder[1]`` carries the work
+ out. At most one claim succeeds per ``process()``.
+
+ The poll runs while the tracker lock is held, so a ``PollFn`` must be bounded
+ or timeout-safe; a blocking poll delays runner-initiated checkpoints.
+ """
+ def __init__(
+ self,
+ restriction: _GrowthState,
+ poll_fn: Callable[[Any], PollResult],
+ key_coder: Coder,
+ termination: TerminationCondition,
+ now_fn: Callable[[], float]):
+ self._restriction = restriction
+ self._poll_fn = poll_fn
+ self._key_coder = key_coder
+ self._termination = termination
+ self._now = now_fn
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def _hash_output(self, value: Any) -> bytes:
+ return hashlib.blake2b(
+ self._key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+ def try_claim(self, holder: list) -> bool:
+ """Performs one poll round (or one replay) and reports it via ``holder``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation,
+ in which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ restriction = self._restriction
+ if isinstance(restriction, _NonPollingGrowthState):
+ holder[1] = ('replay', restriction.pending)
+ self._should_stop = True
+ return True
+
+ element = holder[0]
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = self._hash_output(output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = self._termination.on_seen_new_output(
+ now, termination_state)
+ termination_state = self._termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ self._termination.can_stop_polling(now, termination_state))
+
+ self._primary = _NonPollingGrowthState(
+ PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-
+ # initiated or via defer_remainder) resumes a state that emits nothing.
+ self._residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = self._max_watermark(
+ restriction.poll_watermark, watermark)
+ self._residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ holder[1] = ('poll', new_outputs, watermark, stop)
+ self._should_stop = True
+ return True
+
+ @staticmethod
+ 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 try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(
+ restriction,
+ self._poll_fn,
+ self._key_coder,
+ self._termination,
+ self._now)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ holder = [element, None]
+ if not tracker.try_claim(holder):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ work = holder[1]
+ if work[0] == 'replay':
+ for output in work[1].outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ new_outputs, watermark, stop = work[1], work[2], work[3]
+ for output in new_outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ 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.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
+ self._termination = termination or never()
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._now = now_fn
+
+ @classmethod
+ def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch':
+ return cls(poll_fn)
+
+ def _replace(self, **changes) -> 'Watch':
+ spec = dict(
+ poll_fn=self._poll_fn,
+ termination=self._termination,
+ poll_interval=self._poll_interval,
+ output_coder=self._output_coder,
+ now_fn=self._now)
+ spec.update(changes)
+ return Watch(**spec)
+
+ def with_poll_interval(self, poll_interval) -> 'Watch':
+ return self._replace(poll_interval=_as_duration(poll_interval))
+
+ def with_termination_per_input(
+ self, termination: TerminationCondition) -> 'Watch':
+ return self._replace(termination=termination)
+
+ def with_output_coder(self, output_coder: Coder) -> 'Watch':
+ return self._replace(output_coder=output_coder)
+
+ def expand(self, pcoll):
+ if self._poll_interval is None:
+ raise ValueError('Watch requires with_poll_interval(...)')
+ output_coder = self._output_coder
+ if output_coder is None:
+ hint = self._poll_fn.default_output_coder() if isinstance(
+ self._poll_fn, PollFn) else None
+ output_coder = hint or coders.PickleCoder()
+ if not output_coder.is_deterministic():
+ _LOGGER.warning(
+ 'Watch dedup uses a non-deterministic output coder (%s); equal '
Review Comment:
I agree that a warning is too soft here. I changed `expand` to raise when the resolved coder is non-deterministic. The dedup key is the encoded output, and the key coder is the output coder in this version, so this single check covers the key. I also updated the end-to-end tests to pass a deterministic coder and added a test for the new error.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ The Python SDK ships no coder for ``TimestampedValue``, 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([_HashCode128Coder(), 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.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Drives one input's polling loop.
+
+ ``process()`` only sees a ``RestrictionTrackerView`` whose ``try_claim``
+ returns a bool, so the poll happens inside ``try_claim`` and its result is
+ returned through a two-slot holder list passed as the claim position:
+ ``holder[0]`` carries the input element in, ``holder[1]`` carries the work
+ out. At most one claim succeeds per ``process()``.
+
+ The poll runs while the tracker lock is held, so a ``PollFn`` must be bounded
+ or timeout-safe; a blocking poll delays runner-initiated checkpoints.
+ """
+ def __init__(
+ self,
+ restriction: _GrowthState,
+ poll_fn: Callable[[Any], PollResult],
+ key_coder: Coder,
+ termination: TerminationCondition,
+ now_fn: Callable[[], float]):
+ self._restriction = restriction
+ self._poll_fn = poll_fn
+ self._key_coder = key_coder
+ self._termination = termination
+ self._now = now_fn
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def _hash_output(self, value: Any) -> bytes:
+ return hashlib.blake2b(
+ self._key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+ def try_claim(self, holder: list) -> bool:
+ """Performs one poll round (or one replay) and reports it via ``holder``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation,
+ in which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ restriction = self._restriction
+ if isinstance(restriction, _NonPollingGrowthState):
+ holder[1] = ('replay', restriction.pending)
+ self._should_stop = True
+ return True
+
+ element = holder[0]
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
Review Comment:
This was an important point, thank you for raising it. I moved the poll into `process` ahead of the claim, so it no longer runs while the tracker lock is held. The dedup and checkpoint logic now sits in a small pure function, and `try_claim` simply records the primary and residual. I reran the tests and lint afterward and both pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3524077546
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
Review Comment:
In the replay block for `_NonPollingGrowthState`, the watermark is unconditionally set to `timestamp` before yielding the outputs. If the pending watermark (`restriction.pending.watermark`) is less than `timestamp`, this will advance the watermark past the outputs' timestamps, causing them to be treated as late data. Additionally, since this is a terminal state that stops polling, the watermark hold must be released to `MAX_TIMESTAMP` after yielding the outputs to prevent the pipeline watermark from being permanently stuck.
```python
if isinstance(restriction, _NonPollingGrowthState):\n # Replay the outputs already emitted this round, then stop. No poll.\n if not tracker.try_claim(_replay_plan(restriction)):\n return\n watermark = restriction.pending.watermark\n if watermark is not None:\n _set_watermark_if_greater(watermark_estimator, min(timestamp, watermark))\n else:\n _set_watermark_if_greater(watermark_estimator, timestamp)\n for output in restriction.pending.outputs:\n yield TimestampedValue((element, output.value), output.timestamp)\n _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)\n return
```
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ # Poll before claiming so a slow poll never holds the tracker lock, which
+ # would block runner progress checks and checkpoints.
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+ plan = _plan_poll_round(
+ restriction, result, self._termination, self._key_coder, now)
+ if not tracker.try_claim(plan):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in plan.emit:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if plan.stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ return
+ if plan.watermark is not None:
+ _set_watermark_if_greater(watermark_estimator, plan.watermark)
+ tracker.defer_remainder(self._poll_interval)
Review Comment:
Unconditionally setting the watermark to `timestamp` before yielding the outputs can cause correctness issues. If `plan.watermark` is less than `timestamp`, setting the watermark to `timestamp` first will advance the watermark past the outputs' timestamps, causing them to be treated as late data. To prevent this, the watermark should be held at `min(timestamp, plan.watermark)` (if `plan.watermark` is not None) before yielding, and then advanced to `plan.watermark` after yielding.
```suggestion
# Seed the watermark hold from the input event time after the claim.\n if plan.watermark is not None:\n _set_watermark_if_greater(watermark_estimator, min(timestamp, plan.watermark))\n else:\n _set_watermark_if_greater(watermark_estimator, timestamp)\n for output in plan.emit:\n yield TimestampedValue((element, output.value), output.timestamp)\n if plan.stop:\n # The input is finished, so release the watermark hold to MAX.\n _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)\n return\n if plan.watermark is not None:\n _set_watermark_if_greater(watermark_estimator, plan.watermark)\n tracker.defer_remainder(self._poll_interval)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Eliaaazzz (via GitHub)" <gi...@apache.org>.
Eliaaazzz commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3622408317
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
Review Comment:
Thanks for pointing these out. The block quoted here was rewritten in c559963, so the code no longer looks like this.
The replay branch now leaves the watermark estimator untouched. It claims the pending state, replays the outputs and returns, at watch.py lines 566 to 572. Nothing advances the watermark ahead of the replayed outputs.
On releasing the hold to MAX_TIMESTAMP, I traced the runner before deciding. The estimator state is only read when an invocation produces a deferred residual, in runners/common.py around lines 1064 to 1082, and delayed applications and output watermarks are built only from residuals in bundle_processor.py. A terminal round produces no residual, so its last watermark is never reported as a hold and cannot pin the stage. The DirectRunner releases explicitly when there is no residual, in sdf_direct_runner.py around lines 367 to 374.
I added a test in aeafcb2, test_replay_round_leaves_the_watermark_alone, which pins that the replay round keeps the watermark at the seed and reports no residual.
Happy to revisit if you would rather have the explicit release.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ # Poll before claiming so a slow poll never holds the tracker lock, which
+ # would block runner progress checks and checkpoints.
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+ plan = _plan_poll_round(
+ restriction, result, self._termination, self._key_coder, now)
+ if not tracker.try_claim(plan):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in plan.emit:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if plan.stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ return
+ if plan.watermark is not None:
+ _set_watermark_if_greater(watermark_estimator, plan.watermark)
+ tracker.defer_remainder(self._poll_interval)
Review Comment:
Same here, this path was rewritten in c559963.
The per round watermark advance now happens after the outputs are emitted. Outputs are yielded at watch.py line 589 and the advance is at line 603, so the watermark never moves ahead of what was just emitted.
The only write before the yields is a single seed from the element timestamp, applied once while the estimator is still unset, at watch.py lines 563 to 564. Without it the stage reports a watermark of None, which a runner reads as the minimum timestamp and holds there until the first output.
I added a test in aeafcb2, test_terminal_round_after_deferring_leaves_no_residual. It defers twice, carries the estimator state across both resumes, then completes, and asserts the watermark stays where the last poll left it while the terminal round reports no residual.
I also checked the tests against deliberate regressions. Re-adding the MAX release, writing the watermark inside the replay block, or dropping the seed guard each make them fail.
Please let me know if you would like anything tightened.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3615705422
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
Review Comment:
Could you please check these two comments?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Eliaaazzz (via GitHub)" <gi...@apache.org>.
Eliaaazzz commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3578906397
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ # Poll before claiming so a slow poll never holds the tracker lock, which
+ # would block runner progress checks and checkpoints.
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+ plan = _plan_poll_round(
+ restriction, result, self._termination, self._key_coder, now)
+ if not tracker.try_claim(plan):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in plan.emit:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if plan.stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ return
+ if plan.watermark is not None:
+ _set_watermark_if_greater(watermark_estimator, plan.watermark)
+ tracker.defer_remainder(self._poll_interval)
+
+
+def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None:
+ # set_watermark raises on regression, so only ever advance the watermark.
+ current = watermark_estimator.current_watermark()
+ if current is None or new_watermark > current:
+ watermark_estimator.set_watermark(new_watermark)
+
+
+# ------------------------------------------------------------------------------
+# Public PTransform.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
+ self._termination = termination or never()
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._now = now_fn
+
+ @classmethod
+ def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch':
+ return cls(poll_fn)
+
+ def _replace(self, **changes) -> 'Watch':
+ spec = dict(
+ poll_fn=self._poll_fn,
+ termination=self._termination,
+ poll_interval=self._poll_interval,
+ output_coder=self._output_coder,
+ now_fn=self._now)
+ spec.update(changes)
+ return Watch(**spec)
+
+ def with_poll_interval(self, poll_interval) -> 'Watch':
+ return self._replace(poll_interval=_as_duration(poll_interval))
+
+ def with_termination_per_input(
+ self, termination: TerminationCondition) -> 'Watch':
+ return self._replace(termination=termination)
+
+ def with_output_coder(self, output_coder: Coder) -> 'Watch':
+ return self._replace(output_coder=output_coder)
+
+ def expand(self, pcoll):
+ if self._poll_interval is None:
+ raise ValueError('Watch requires with_poll_interval(...)')
+ output_coder = self._output_coder
+ if output_coder is None:
+ hint = self._poll_fn.default_output_coder() if isinstance(
+ self._poll_fn, PollFn) else None
+ output_coder = hint or coders.PickleCoder()
+ # Dedup hashes the encoded output as its key, so a non-deterministic coder
+ # would hash equal outputs differently and re-emit them, defeating the
+ # transform. Require determinism rather than silently emitting duplicates.
+ if not output_coder.is_deterministic():
Review Comment:
I went through the Java expansion again and you are right, I had taken the key coder name without the distinction behind it. The transform now takes `output_key_fn` and `output_key_coder`. When no key fn is given, the dedup key is the output itself and the key coder defaults to the output coder, and the determinism requirement applies only to the key coder. Expansion now converts it with `as_deterministic_coder`, the same conversion grouping keys go through, so the pickle based fallback hashes deterministically and only a coder with no deterministic form is rejected. The output coder itself no longer needs to be deterministic.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
Review Comment:
I agree, `_PollPlan` was patching around the boundary instead of fixing it, and it is gone now. The claimed position is the poll round itself, a pair of the deduped `PollResult` and the termination state. `try_claim` validates that pair against the restriction, and `try_split` derives the replay primary and the merged polling residual, so the checkpoint logic lives in the tracker again and `process` only polls, dedups, emits and decides whether to stop or resume. One behavior worth noting, a runner initiated split that lands right after a terminal round is claimed leaves a polling residual, so that input polls once more and stops on the next round. I also moved the watermark updates onto the resume path only, seeded from the element timestamp on the first round.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn merged PR #39023:
URL: https://github.com/apache/beam/pull/39023
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3499980025
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
Review Comment:
Python coders are not implemented identically as Java. In Java we have HashCode128Coder to write fix size 16 bytes. In Python encode() encodes to bytes, and will be length-prefixed in the raw stream.
Just use Beam's BytesCoder may have a better performance (reduced a layer of wrapper) while the length check can be done elsewhere
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
Review Comment:
What's the point of making an Any alias?
Would it be preferred to make it an empty dataclass and make _PollingGrowthState, _NonPollingGrowthState inherit it?
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
Review Comment:
> The Python SDK ships no coder for ``TimestampedValue``
A little bit surprising. Likely due to Python SDK has some special handling:
https://github.com/apache/beam/blob/a86e61152704e5b65bdedd2e968c8d64d0b2f944/sdks/python/apache_beam/runners/common.py#L1773
converting TimestampedValue to a WindowedValue, and WindowedValue has a standard coder
https://github.com/apache/beam/blob/a86e61152704e5b65bdedd2e968c8d64d0b2f944/sdks/python/apache_beam/coders/typecoders.py#L96
I'm considering make TimestampedValue a public interface like Java SDK we current have (could be a separate PR)
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
Review Comment:
As it's declared in `__all__`, consider add an example about creating a PollFn that is callable and has default_output_coder implemented
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ The Python SDK ships no coder for ``TimestampedValue``, 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([_HashCode128Coder(), 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.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Drives one input's polling loop.
+
+ ``process()`` only sees a ``RestrictionTrackerView`` whose ``try_claim``
+ returns a bool, so the poll happens inside ``try_claim`` and its result is
+ returned through a two-slot holder list passed as the claim position:
+ ``holder[0]`` carries the input element in, ``holder[1]`` carries the work
+ out. At most one claim succeeds per ``process()``.
+
+ The poll runs while the tracker lock is held, so a ``PollFn`` must be bounded
+ or timeout-safe; a blocking poll delays runner-initiated checkpoints.
+ """
+ def __init__(
+ self,
+ restriction: _GrowthState,
+ poll_fn: Callable[[Any], PollResult],
+ key_coder: Coder,
+ termination: TerminationCondition,
+ now_fn: Callable[[], float]):
+ self._restriction = restriction
+ self._poll_fn = poll_fn
+ self._key_coder = key_coder
+ self._termination = termination
+ self._now = now_fn
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def _hash_output(self, value: Any) -> bytes:
+ return hashlib.blake2b(
+ self._key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+ def try_claim(self, holder: list) -> bool:
+ """Performs one poll round (or one replay) and reports it via ``holder``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation,
+ in which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ restriction = self._restriction
+ if isinstance(restriction, _NonPollingGrowthState):
+ holder[1] = ('replay', restriction.pending)
+ self._should_stop = True
+ return True
+
+ element = holder[0]
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = self._hash_output(output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = self._termination.on_seen_new_output(
+ now, termination_state)
+ termination_state = self._termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ self._termination.can_stop_polling(now, termination_state))
+
+ self._primary = _NonPollingGrowthState(
+ PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-
+ # initiated or via defer_remainder) resumes a state that emits nothing.
+ self._residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = self._max_watermark(
+ restriction.poll_watermark, watermark)
+ self._residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ holder[1] = ('poll', new_outputs, watermark, stop)
+ self._should_stop = True
+ return True
+
+ @staticmethod
+ 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 try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(
+ restriction,
+ self._poll_fn,
+ self._key_coder,
+ self._termination,
+ self._now)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ holder = [element, None]
+ if not tracker.try_claim(holder):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ work = holder[1]
+ if work[0] == 'replay':
+ for output in work[1].outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ new_outputs, watermark, stop = work[1], work[2], work[3]
+ for output in new_outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ 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.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
Review Comment:
We should be able to wrap the poll_fn with a PollFn if it's typehint is compatible; or we should be able to derive typehints for the callables in expand beneath.
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ The Python SDK ships no coder for ``TimestampedValue``, 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([_HashCode128Coder(), 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.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Drives one input's polling loop.
+
+ ``process()`` only sees a ``RestrictionTrackerView`` whose ``try_claim``
+ returns a bool, so the poll happens inside ``try_claim`` and its result is
+ returned through a two-slot holder list passed as the claim position:
+ ``holder[0]`` carries the input element in, ``holder[1]`` carries the work
+ out. At most one claim succeeds per ``process()``.
+
+ The poll runs while the tracker lock is held, so a ``PollFn`` must be bounded
+ or timeout-safe; a blocking poll delays runner-initiated checkpoints.
+ """
+ def __init__(
+ self,
+ restriction: _GrowthState,
+ poll_fn: Callable[[Any], PollResult],
+ key_coder: Coder,
+ termination: TerminationCondition,
+ now_fn: Callable[[], float]):
+ self._restriction = restriction
+ self._poll_fn = poll_fn
+ self._key_coder = key_coder
+ self._termination = termination
+ self._now = now_fn
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def _hash_output(self, value: Any) -> bytes:
+ return hashlib.blake2b(
+ self._key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+ def try_claim(self, holder: list) -> bool:
+ """Performs one poll round (or one replay) and reports it via ``holder``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation,
+ in which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ restriction = self._restriction
+ if isinstance(restriction, _NonPollingGrowthState):
+ holder[1] = ('replay', restriction.pending)
+ self._should_stop = True
+ return True
+
+ element = holder[0]
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = self._hash_output(output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = self._termination.on_seen_new_output(
+ now, termination_state)
+ termination_state = self._termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ self._termination.can_stop_polling(now, termination_state))
+
+ self._primary = _NonPollingGrowthState(
+ PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-
+ # initiated or via defer_remainder) resumes a state that emits nothing.
+ self._residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = self._max_watermark(
+ restriction.poll_watermark, watermark)
+ self._residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ holder[1] = ('poll', new_outputs, watermark, stop)
+ self._should_stop = True
+ return True
+
+ @staticmethod
+ 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 try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(
+ restriction,
+ self._poll_fn,
+ self._key_coder,
+ self._termination,
+ self._now)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ holder = [element, None]
+ if not tracker.try_claim(holder):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ work = holder[1]
+ if work[0] == 'replay':
+ for output in work[1].outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ new_outputs, watermark, stop = work[1], work[2], work[3]
+ for output in new_outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ 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.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
+ self._termination = termination or never()
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._now = now_fn
+
+ @classmethod
+ def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch':
+ return cls(poll_fn)
+
+ def _replace(self, **changes) -> 'Watch':
+ spec = dict(
+ poll_fn=self._poll_fn,
+ termination=self._termination,
+ poll_interval=self._poll_interval,
+ output_coder=self._output_coder,
+ now_fn=self._now)
+ spec.update(changes)
+ return Watch(**spec)
+
+ def with_poll_interval(self, poll_interval) -> 'Watch':
+ return self._replace(poll_interval=_as_duration(poll_interval))
+
+ def with_termination_per_input(
+ self, termination: TerminationCondition) -> 'Watch':
+ return self._replace(termination=termination)
+
+ def with_output_coder(self, output_coder: Coder) -> 'Watch':
+ return self._replace(output_coder=output_coder)
+
+ def expand(self, pcoll):
+ if self._poll_interval is None:
+ raise ValueError('Watch requires with_poll_interval(...)')
+ output_coder = self._output_coder
+ if output_coder is None:
+ hint = self._poll_fn.default_output_coder() if isinstance(
+ self._poll_fn, PollFn) else None
+ output_coder = hint or coders.PickleCoder()
+ if not output_coder.is_deterministic():
+ _LOGGER.warning(
+ 'Watch dedup uses a non-deterministic output coder (%s); equal '
Review Comment:
This warning message needs more consideration --- if it's statement is true we should raise here, as the whole point of the Watch.growth is that it only emits newly seen elements.
I think as long as key is deterministic it should be fine.
In Java it made a distinction between "output" and "outputKey", as well as its own Fn and coders. KeyCoder needs to be deterministic.
I see Python implementation has a key_coder as well. Would check key coder be sufficient?
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,676 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import logging
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_LOGGER = logging.getLogger(__name__)
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`.
+ """
+ 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.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState:
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; the tracker builds a new
+ mapping for each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState:
+ """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
+
+
+_GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState]
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _HashCode128Coder(Coder):
+ """Fixed-width coder for a 16-byte output hash.
+
+ Encodes and decodes exactly 16 bytes and raises on any other length, so a
+ corrupt restriction surfaces at decode time.
+ """
+ def encode(self, value: bytes) -> bytes:
+ if len(value) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(value)))
+ return value
+
+ def decode(self, encoded: bytes) -> bytes:
+ if len(encoded) != _HASH_DIGEST_SIZE:
+ raise ValueError(
+ 'hash must be %d bytes, got %d' % (_HASH_DIGEST_SIZE, len(encoded)))
+ return encoded
+
+ def is_deterministic(self) -> bool:
+ return True
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ The Python SDK ships no coder for ``TimestampedValue``, 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([_HashCode128Coder(), 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.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Drives one input's polling loop.
+
+ ``process()`` only sees a ``RestrictionTrackerView`` whose ``try_claim``
+ returns a bool, so the poll happens inside ``try_claim`` and its result is
+ returned through a two-slot holder list passed as the claim position:
+ ``holder[0]`` carries the input element in, ``holder[1]`` carries the work
+ out. At most one claim succeeds per ``process()``.
+
+ The poll runs while the tracker lock is held, so a ``PollFn`` must be bounded
+ or timeout-safe; a blocking poll delays runner-initiated checkpoints.
+ """
+ def __init__(
+ self,
+ restriction: _GrowthState,
+ poll_fn: Callable[[Any], PollResult],
+ key_coder: Coder,
+ termination: TerminationCondition,
+ now_fn: Callable[[], float]):
+ self._restriction = restriction
+ self._poll_fn = poll_fn
+ self._key_coder = key_coder
+ self._termination = termination
+ self._now = now_fn
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def _hash_output(self, value: Any) -> bytes:
+ return hashlib.blake2b(
+ self._key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+ def try_claim(self, holder: list) -> bool:
+ """Performs one poll round (or one replay) and reports it via ``holder``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation,
+ in which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ restriction = self._restriction
+ if isinstance(restriction, _NonPollingGrowthState):
+ holder[1] = ('replay', restriction.pending)
+ self._should_stop = True
+ return True
+
+ element = holder[0]
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
Review Comment:
put poll inside try_claim is not recommended as long running try_claim could cause issues (see #36750, it's fixed in Java, not sure if it could also happens in Python, probably worth testing)
poll can take very long if it makes API calls (e.g. list all files in gcs). Would it be possible to put it inside DoFn's process element?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Eliaaazzz (via GitHub)" <gi...@apache.org>.
Eliaaazzz commented on PR #39023:
URL: https://github.com/apache/beam/pull/39023#issuecomment-4851620049
Thanks for the thorough review. I addressed all the comments and pushed. PTAL when you have time.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39023:
URL: https://github.com/apache/beam/pull/39023#issuecomment-4884326191
Assigning reviewers:
R: @shunping for label python.
Note: If you would like to opt out of this review, comment `assign to next reviewer`.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
The PR bot will only process comments in the main thread (not review comments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "codecov[bot] (via GitHub)" <gi...@apache.org>.
codecov[bot] commented on PR #39023:
URL: https://github.com/apache/beam/pull/39023#issuecomment-4747877340
## [Codecov](https://app.codecov.io/gh/apache/beam/pull/39023?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `88.21656%` with `37 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 54.15%. Comparing base ([`8801afd`](https://app.codecov.io/gh/apache/beam/commit/8801afdbd321d6de4a9dca23f0416b125a8ff14e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`5b039a0`](https://app.codecov.io/gh/apache/beam/commit/5b039a0d0c8dc61016b915abdcf726e9534ce18f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 1065 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/beam/pull/39023?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [sdks/python/apache\_beam/io/watch.py](https://app.codecov.io/gh/apache/beam/pull/39023?src=pr&el=tree&filepath=sdks%2Fpython%2Fapache_beam%2Fio%2Fwatch.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c2Rrcy9weXRob24vYXBhY2hlX2JlYW0vaW8vd2F0Y2gucHk=) | 88.21% | [37 Missing :warning: ](https://app.codecov.io/gh/apache/beam/pull/39023?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #39023 +/- ##
============================================
- Coverage 55.16% 54.15% -1.01%
Complexity 1676 1676
============================================
Files 1068 1071 +3
Lines 167265 166847 -418
Branches 1208 1208
============================================
- Hits 92265 90352 -1913
- Misses 72816 74311 +1495
Partials 2184 2184
```
| [Flag](https://app.codecov.io/gh/apache/beam/pull/39023/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [python](https://app.codecov.io/gh/apache/beam/pull/39023/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `78.76% <88.21%> (-2.02%)` | :arrow_down: |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/beam/pull/39023?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "Eliaaazzz (via GitHub)" <gi...@apache.org>.
Eliaaazzz commented on code in PR #39023:
URL: https://github.com/apache/beam/pull/39023#discussion_r3578904613
##########
sdks/python/apache_beam/io/watch.py:
##########
@@ -0,0 +1,691 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Experimental ``Watch`` transform for the Python SDK.
+
+``Watch`` continuously watches a growing set of outputs for each input element,
+calling a user poll function on an interval until a per-input termination
+condition fires. It is the engine behind periodic file-discovery and any
+periodic polling source.
+
+For every input element the transform runs an independent loop::
+
+ poll -> keep never-seen-before outputs -> emit them (timestamped) ->
+ update watermark -> check termination -> wait(poll_interval) -> poll -> ...
+
+The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each
+output carries the event time the poll function first reported it. Dedup uses a
+stable 128-bit hash of the encoded output, so the output coder must be
+deterministic for dedup to hold across workers and restarts.
+
+Example::
+
+ from apache_beam.io.watch import Watch, PollResult, after_total_of
+ from apache_beam.transforms.window import TimestampedValue
+ from apache_beam.utils.timestamp import Duration, Timestamp
+
+ def poll(prefix):
+ now = Timestamp.now()
+ outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)]
+ return PollResult.complete(outputs)
+
+ watched = (inputs
+ | Watch.growth_of(poll)
+ .with_poll_interval(Duration(seconds=5))
+ .with_termination_per_input(after_total_of(60)))
+
+This API is experimental and may change in backwards-incompatible ways.
+"""
+
+import collections
+import dataclasses
+import hashlib
+import time
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from apache_beam import coders
+from apache_beam.coders.coders import Coder
+from apache_beam.coders.coders import NullableCoder
+from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import TupleCoder
+from apache_beam.io import iobase
+from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
+from apache_beam.runners import sdf_utils
+from apache_beam.transforms import PTransform
+from apache_beam.transforms import core
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+__all__ = [
+ 'Watch',
+ 'PollResult',
+ 'PollFn',
+ 'TerminationCondition',
+ 'never',
+ 'after_total_of',
+]
+
+_HASH_DIGEST_SIZE = 16 # 128-bit digest width.
+
+# ------------------------------------------------------------------------------
+# Public API.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class PollResult:
+ """Outputs produced by one poll, plus an optional explicit watermark.
+
+ ``watermark`` of ``None`` lets the transform infer the watermark from the
+ earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by
+ :meth:`complete`) marks the input finished, so polling stops.
+ """
+ outputs: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp] = None
+
+ @property
+ def is_complete(self) -> bool:
+ return self.watermark == MAX_TIMESTAMP
+
+ @staticmethod
+ def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]:
+ if timestamp is None:
+ default_ts = Timestamp.now()
+ else:
+ default_ts = Timestamp.of(timestamp)
+ normalized = []
+ for output in outputs:
+ if isinstance(output, TimestampedValue):
+ normalized.append(output)
+ else:
+ normalized.append(TimestampedValue(output, default_ts))
+ return tuple(normalized)
+
+ @staticmethod
+ def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports outputs and expects more; the transform infers the watermark.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(PollResult._normalize(outputs, timestamp), watermark=None)
+
+ @staticmethod
+ def complete(outputs: Iterable, timestamp=None) -> 'PollResult':
+ """Reports the final outputs for an input, after which polling stops.
+
+ A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp``
+ when given, else with the current processing time.
+ """
+ return PollResult(
+ PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP)
+
+ def with_watermark(self, watermark) -> 'PollResult':
+ return dataclasses.replace(self, watermark=Timestamp.of(watermark))
+
+
+class PollFn(object):
+ """Optional base for a poll function ``input -> PollResult``.
+
+ Any callable with that signature works; subclass only to attach an output
+ coder hint via :meth:`default_output_coder`. The hint is used for dedup and
+ must be deterministic::
+
+ from apache_beam import coders
+
+ class ListFiles(PollFn):
+ def __call__(self, prefix):
+ return PollResult.incomplete(list_files(prefix))
+
+ def default_output_coder(self):
+ return coders.StrUtf8Coder()
+ """
+ def __call__(self, element: Any) -> PollResult:
+ raise NotImplementedError
+
+ def default_output_coder(self) -> Optional[Coder]:
+ return None
+
+
+class TerminationCondition(object):
+ """Per-input stop policy with immutable, encodable state.
+
+ Hooks follow the lifecycle of one input's polling loop. ``state`` flows from
+ :meth:`for_new_input` through the per-round hooks and is serialized with
+ :meth:`state_coder`.
+ """
+ def for_new_input(self, now: Timestamp, element: Any) -> Any:
+ raise NotImplementedError
+
+ def on_seen_new_output(self, now: Timestamp, state: Any) -> Any:
+ return state
+
+ def on_poll_complete(self, state: Any) -> Any:
+ return state
+
+ def can_stop_polling(self, now: Timestamp, state: Any) -> bool:
+ raise NotImplementedError
+
+ def state_coder(self) -> Coder:
+ raise NotImplementedError
+
+
+class _Never(TerminationCondition):
+ """Polls until the poll function returns :meth:`PollResult.complete`."""
+ def for_new_input(self, now, element):
+ return 0
+
+ def can_stop_polling(self, now, state):
+ return False
+
+ def state_coder(self):
+ return coders.VarIntCoder()
+
+
+class _AfterTotalOf(TerminationCondition):
+ """Stops once the wall-clock time since the input was first seen exceeds a
+ fixed duration."""
+ def __init__(self, duration: Duration):
+ self._duration_micros = duration.micros
+
+ def for_new_input(self, now, element):
+ return (now, self._duration_micros)
+
+ def can_stop_polling(self, now, state):
+ start, duration_micros = state
+ return (now - start).micros > duration_micros
+
+ def state_coder(self):
+ return TupleCoder([TimestampCoder(), coders.VarIntCoder()])
+
+
+def never() -> TerminationCondition:
+ """Polls until :meth:`PollResult.complete`."""
+ return _Never()
+
+
+def after_total_of(duration) -> TerminationCondition:
+ """Stops polling an input after ``duration`` (a :class:`Duration` or seconds)
+ has elapsed since it was first seen."""
+ return _AfterTotalOf(_as_duration(duration))
+
+
+# ------------------------------------------------------------------------------
+# Restriction state.
+# ------------------------------------------------------------------------------
+
+
+class _GrowthState:
+ """Base for the two restriction variants a Watch input can hold."""
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollingGrowthState(_GrowthState):
+ """Keep-polling state: emitted-output hashes, watermark, termination state.
+
+ ``completed`` maps a 16-byte output hash to the event time it was first seen.
+ It is insertion-ordered and treated as immutable; a new mapping is built for
+ each residual.
+ """
+ completed: 'collections.OrderedDict[bytes, Timestamp]'
+ poll_watermark: Optional[Timestamp]
+ termination_state: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class _NonPollingGrowthState(_GrowthState):
+ """Replay-then-stop state: the outputs already emitted this round.
+
+ Produced as the checkpoint primary so a bundle retry re-emits exactly those
+ outputs.
+ """
+ pending: PollResult
+
+
+# ------------------------------------------------------------------------------
+# Coders.
+# ------------------------------------------------------------------------------
+
+
+class _TimestampedValueCoder(Coder):
+ """Coder for :class:`TimestampedValue`.
+
+ ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the
+ wire, so the SDK ships no standalone coder for it. Watch keeps it inside the
+ restriction state, so this encodes the ``(value, timestamp)`` pair with a
+ :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode.
+ """
+ def __init__(self, value_coder: Coder):
+ self._tuple_coder = TupleCoder([value_coder, TimestampCoder()])
+
+ def encode(self, value: TimestampedValue) -> bytes:
+ return self._tuple_coder.encode((value.value, value.timestamp))
+
+ def decode(self, encoded: bytes) -> TimestampedValue:
+ value, timestamp = self._tuple_coder.decode(encoded)
+ return TimestampedValue(value, timestamp)
+
+ def is_deterministic(self) -> bool:
+ return self._tuple_coder.is_deterministic()
+
+
+class _GrowthStateCoder(Coder):
+ """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`.
+
+ A ``(tag, payload)`` envelope selects the variant; the payload is a
+ variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered
+ list of ``(hash, timestamp)`` pairs so insertion order survives a round trip.
+ This format is internal to the Python SDK.
+ """
+ def __init__(self, output_coder: Coder, termination: TerminationCondition):
+ nullable_ts = NullableCoder(TimestampCoder())
+ self._envelope_coder = TupleCoder(
+ [coders.VarIntCoder(), coders.BytesCoder()])
+ self._polling_coder = TupleCoder([
+ termination.state_coder(),
+ nullable_ts,
+ coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])),
+ ])
+ self._non_polling_coder = TupleCoder([
+ nullable_ts,
+ coders.ListCoder(_TimestampedValueCoder(output_coder)),
+ ])
+
+ def encode(self, state: _GrowthState) -> bytes:
+ if isinstance(state, _PollingGrowthState):
+ payload = self._polling_coder.encode((
+ state.termination_state,
+ state.poll_watermark,
+ list(state.completed.items())))
+ return self._envelope_coder.encode((0, payload))
+ payload = self._non_polling_coder.encode(
+ (state.pending.watermark, list(state.pending.outputs)))
+ return self._envelope_coder.encode((1, payload))
+
+ def decode(self, encoded: bytes) -> _GrowthState:
+ tag, payload = self._envelope_coder.decode(encoded)
+ if tag == 0:
+ termination_state, poll_watermark, items = self._polling_coder.decode(
+ payload)
+ return _PollingGrowthState(
+ collections.OrderedDict(items), poll_watermark, termination_state)
+ if tag == 1:
+ watermark, outputs = self._non_polling_coder.decode(payload)
+ return _NonPollingGrowthState(PollResult(tuple(outputs), watermark))
+ raise ValueError('unknown Watch growth state tag: %r' % (tag, ))
+
+ def is_deterministic(self) -> bool:
+ return False
+
+
+# ------------------------------------------------------------------------------
+# Restriction tracker.
+# ------------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class _PollPlan:
+ """One planned poll round: what to emit plus the self-checkpoint split.
+
+ ``process()`` builds this from a poll result before claiming, so the poll runs
+ outside the tracker lock. ``primary`` is the round just processed and
+ ``residual`` is the work a checkpoint resumes.
+ """
+ emit: Tuple[TimestampedValue, ...]
+ watermark: Optional[Timestamp]
+ stop: bool
+ primary: _GrowthState
+ residual: _GrowthState
+
+
+def _hash_output(key_coder: Coder, value: Any) -> bytes:
+ return hashlib.blake2b(
+ key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest()
+
+
+def _max_watermark(left: Optional[Timestamp],
+ right: Optional[Timestamp]) -> Optional[Timestamp]:
+ if left is None:
+ return right
+ if right is None:
+ return left
+ return max(left, right)
+
+
+def _plan_poll_round(
+ restriction: _PollingGrowthState,
+ result: PollResult,
+ termination: TerminationCondition,
+ key_coder: Coder,
+ now: Timestamp) -> _PollPlan:
+ """Dedups a poll result into new outputs and builds the checkpoint split.
+
+ Pure and side-effect free, so ``process()`` can call it after the poll and
+ before claiming, keeping the poll off the tracker lock.
+ """
+ new_outputs = [] # type: List[TimestampedValue]
+ claimed = [] # type: List[Tuple[bytes, Timestamp]]
+ seen_this_round = set() # type: set
+ for output in result.outputs:
+ key_hash = _hash_output(key_coder, output.value)
+ if key_hash in restriction.completed or key_hash in seen_this_round:
+ continue
+ seen_this_round.add(key_hash)
+ new_outputs.append(output)
+ claimed.append((key_hash, output.timestamp))
+ new_outputs.sort(key=lambda output: output.timestamp)
+
+ termination_state = restriction.termination_state
+ if new_outputs:
+ termination_state = termination.on_seen_new_output(now, termination_state)
+ termination_state = termination.on_poll_complete(termination_state)
+
+ if result.watermark is not None:
+ watermark = result.watermark
+ elif new_outputs:
+ watermark = new_outputs[0].timestamp
+ else:
+ watermark = None
+
+ # A watermark at MAX means no more output is possible, so polling stops.
+ reached_max = watermark is not None and watermark >= MAX_TIMESTAMP
+ stop = (
+ result.is_complete or reached_max or
+ termination.can_stop_polling(now, termination_state))
+
+ primary = _NonPollingGrowthState(PollResult(tuple(new_outputs), watermark))
+ if stop:
+ # Terminal round: no polling work remains, so a checkpoint (runner-initiated
+ # or via defer_remainder) resumes a state that emits nothing.
+ residual = _NonPollingGrowthState(PollResult((), watermark))
+ else:
+ merged = collections.OrderedDict(restriction.completed)
+ for key_hash, first_seen in claimed:
+ merged[key_hash] = first_seen
+ residual_watermark = _max_watermark(restriction.poll_watermark, watermark)
+ residual = _PollingGrowthState(
+ merged, residual_watermark, termination_state)
+ return _PollPlan(tuple(new_outputs), watermark, stop, primary, residual)
+
+
+def _replay_plan(restriction: _NonPollingGrowthState) -> _PollPlan:
+ """Plans a replay of the outputs already emitted this round, then stops."""
+ resumed = _NonPollingGrowthState(
+ PollResult((), restriction.pending.watermark))
+ return _PollPlan(
+ restriction.pending.outputs, None, True, restriction, resumed)
+
+
+class _GrowthRestrictionTracker(iobase.RestrictionTracker):
+ """Tracks one input's polling restriction and its self-checkpoints.
+
+ ``process()`` runs the poll and plans the round with
+ :func:`_plan_poll_round`, then claims the plan; the tracker only records the
+ plan's primary and residual so a checkpoint (via ``defer_remainder`` or the
+ runner) resumes correctly. Keeping the poll in ``process()`` means a slow
+ poll never holds the tracker lock, so it cannot delay runner progress checks
+ or checkpoints.
+ """
+ def __init__(self, restriction: _GrowthState):
+ self._restriction = restriction
+ self._should_stop = False
+ self._primary = None # type: Optional[_GrowthState]
+ self._residual = None # type: Optional[_GrowthState]
+
+ def current_restriction(self) -> _GrowthState:
+ return self._restriction
+
+ def try_claim(self, plan: _PollPlan) -> bool:
+ """Records a planned round's checkpoint split; one claim per ``process()``.
+
+ Returns ``False`` only when a checkpoint already stopped this invocation, in
+ which case ``process()`` must emit nothing.
+ """
+ if self._should_stop:
+ return False
+ self._primary = plan.primary
+ self._residual = plan.residual
+ self._should_stop = True
+ return True
+
+ def try_split(self, fraction_of_remainder):
+ # Only self-checkpoint (fraction 0) is supported; decline dynamic splits.
+ if fraction_of_remainder != 0:
+ return None
+ if self._primary is None:
+ # No claim happened this invocation: keep the whole state as the residual.
+ primary = _NonPollingGrowthState(PollResult((), None))
+ residual = self._restriction
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+ primary, residual = self._primary, self._residual
+ self._restriction = primary
+ self._should_stop = True
+ return primary, residual
+
+ def check_done(self) -> bool:
+ # Called after every process(); the single claim or a split sets the flag.
+ if self._should_stop:
+ return True
+ raise ValueError(
+ 'Watch restriction was neither claimed nor checkpointed: %r' %
+ (self._restriction, ))
+
+ def current_progress(self) -> 'iobase.RestrictionProgress':
+ if self._should_stop:
+ return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
+ return iobase.RestrictionProgress(completed=0.0, remaining=1.0)
+
+ def is_bounded(self) -> bool:
+ # A polling restriction is unbounded; a replay-then-stop one is bounded.
+ return isinstance(self._restriction, _NonPollingGrowthState)
+
+
+# ------------------------------------------------------------------------------
+# Splittable DoFn (its own restriction provider).
+# ------------------------------------------------------------------------------
+
+
+class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider):
+ """Polling SDF that emits ``(input, output)`` pairs.
+
+ The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no
+ argument resolves the provider to the DoFn instance, so the provider methods
+ read the transform-level spec (poll function, coders, termination) off
+ ``self``. Provider methods run on a separately deserialized copy and before
+ ``setup()``, so the spec is immutable state set in ``__init__``.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: TerminationCondition,
+ poll_interval: Duration,
+ output_coder: Coder,
+ now_fn: Optional[Callable[[], float]] = None):
+ self._poll_fn = poll_fn
+ self._termination = termination
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._key_coder = output_coder
+ self._now = now_fn or time.time
+ self._restriction_coder = _GrowthStateCoder(output_coder, termination)
+
+ def initial_restriction(self, element) -> _PollingGrowthState:
+ now = Timestamp.of(self._now())
+ return _PollingGrowthState(
+ collections.OrderedDict(),
+ None,
+ self._termination.for_new_input(now, element))
+
+ def create_tracker(self, restriction) -> _GrowthRestrictionTracker:
+ return _GrowthRestrictionTracker(restriction)
+
+ def split(self, element, restriction):
+ # Watch fans out by input element, so each restriction stays whole.
+ yield restriction
+
+ def restriction_coder(self) -> Coder:
+ return self._restriction_coder
+
+ def restriction_size(self, element, restriction) -> int:
+ return 1
+
+ def truncate(self, element, restriction):
+ # On drain, replay a pending NonPolling state and stop further polling.
+ if isinstance(restriction, _NonPollingGrowthState):
+ return restriction
+ return None
+
+ @core.DoFn.unbounded_per_element()
+ def process(
+ self,
+ element,
+ timestamp=core.DoFn.TimestampParam,
+ tracker=core.DoFn.RestrictionParam(),
+ watermark_estimator=core.DoFn.WatermarkEstimatorParam(
+ ManualWatermarkEstimator.default_provider())):
+ assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
+ restriction = tracker.current_restriction()
+ if isinstance(restriction, _NonPollingGrowthState):
+ # Replay the outputs already emitted this round, then stop. No poll.
+ if not tracker.try_claim(_replay_plan(restriction)):
+ return
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in restriction.pending.outputs:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ return
+ # Poll before claiming so a slow poll never holds the tracker lock, which
+ # would block runner progress checks and checkpoints.
+ now = Timestamp.of(self._now())
+ result = self._poll_fn(element)
+ plan = _plan_poll_round(
+ restriction, result, self._termination, self._key_coder, now)
+ if not tracker.try_claim(plan):
+ # A checkpoint already stopped this invocation; emit nothing.
+ return
+ # Seed the watermark hold from the input event time after the claim.
+ _set_watermark_if_greater(watermark_estimator, timestamp)
+ for output in plan.emit:
+ yield TimestampedValue((element, output.value), output.timestamp)
+ if plan.stop:
+ # The input is finished, so release the watermark hold to MAX.
+ _set_watermark_if_greater(watermark_estimator, MAX_TIMESTAMP)
+ return
+ if plan.watermark is not None:
+ _set_watermark_if_greater(watermark_estimator, plan.watermark)
+ tracker.defer_remainder(self._poll_interval)
+
+
+def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None:
+ # set_watermark raises on regression, so only ever advance the watermark.
+ current = watermark_estimator.current_watermark()
+ if current is None or new_watermark > current:
+ watermark_estimator.set_watermark(new_watermark)
+
+
+# ------------------------------------------------------------------------------
+# Public PTransform.
+# ------------------------------------------------------------------------------
+
+
+class Watch(PTransform):
+ """Watches a growing set of outputs per input via a periodic poll function.
+
+ Build with :meth:`growth_of` and the ``with_*`` methods. The output is an
+ unbounded ``PCollection`` of ``(input, output)`` pairs.
+ """
+ def __init__(
+ self,
+ poll_fn: Callable[[Any], PollResult],
+ termination: Optional[TerminationCondition] = None,
+ poll_interval: Optional[Duration] = None,
+ output_coder: Optional[Coder] = None,
+ now_fn: Optional[Callable[[], float]] = None):
+ super().__init__()
+ self._poll_fn = poll_fn
+ self._termination = termination or never()
+ self._poll_interval = poll_interval
+ self._output_coder = output_coder
+ self._now = now_fn
+
+ @classmethod
+ def growth_of(cls, poll_fn: Callable[[Any], PollResult]) -> 'Watch':
+ return cls(poll_fn)
+
+ def _replace(self, **changes) -> 'Watch':
+ spec = dict(
+ poll_fn=self._poll_fn,
+ termination=self._termination,
+ poll_interval=self._poll_interval,
+ output_coder=self._output_coder,
+ now_fn=self._now)
+ spec.update(changes)
+ return Watch(**spec)
+
+ def with_poll_interval(self, poll_interval) -> 'Watch':
Review Comment:
You are right, the setter methods added nothing on top of the constructor. I removed `growth_of` and all the `with_*` methods, so the transform is now built in one call like the other utility transforms, and `poll_interval` became a required constructor argument since polling cannot run without it.
##########
sdks/python/apache_beam/io/watch_test.py:
##########
@@ -0,0 +1,307 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for the Watch transform."""
+
+import collections
+import unittest
+
+import apache_beam as beam
+from apache_beam.coders.coders import StrUtf8Coder
+from apache_beam.io.watch import PollResult
+from apache_beam.io.watch import Watch
+from apache_beam.io.watch import _GrowthRestrictionTracker
+from apache_beam.io.watch import _GrowthStateCoder
+from apache_beam.io.watch import _NonPollingGrowthState
+from apache_beam.io.watch import _plan_poll_round
+from apache_beam.io.watch import _PollingGrowthState
+from apache_beam.io.watch import _replay_plan
+from apache_beam.io.watch import after_total_of
+from apache_beam.io.watch import never
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.runners.sdf_utils import RestrictionTrackerView
+from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import TestWindowedValue
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+from apache_beam.transforms.window import FixedWindows
+from apache_beam.transforms.window import GlobalWindow
+from apache_beam.transforms.window import TimestampedValue
+from apache_beam.typehints import typehints
+from apache_beam.utils.timestamp import MAX_TIMESTAMP
+from apache_beam.utils.timestamp import Duration
+from apache_beam.utils.timestamp import Timestamp
+
+
+def _ts(value, timestamp):
+ return TimestampedValue(value, Timestamp(timestamp))
+
+
+def _plan(restriction, poll_fn, now=0.0, termination=None):
+ termination = termination or never()
+ return _plan_poll_round(
+ restriction,
+ poll_fn('input'),
+ termination,
+ StrUtf8Coder(),
+ Timestamp.of(now))
+
+
+def _initial_polling(termination=None, now=Timestamp(0)):
+ termination = termination or never()
+ return _PollingGrowthState(
+ collections.OrderedDict(), None, termination.for_new_input(now, 'input'))
+
+
+class GrowthStateCoderTest(unittest.TestCase):
+ def test_polling_round_trip_preserves_resume_state(self):
+ termination = after_total_of(Duration(30))
+ coder = _GrowthStateCoder(StrUtf8Coder(), termination)
+ completed = collections.OrderedDict([
+ (b'a' * 16, Timestamp(1)),
+ (b'b' * 16, Timestamp(2)),
+ (b'c' * 16, Timestamp(3)),
+ ])
+ termination_state = termination.for_new_input(Timestamp(7), 'input')
+ state = _PollingGrowthState(completed, Timestamp(5), termination_state)
+ decoded = coder.decode(coder.encode(state))
+ self.assertEqual(list(completed.items()), list(decoded.completed.items()))
+ self.assertEqual(Timestamp(5), decoded.poll_watermark)
+ self.assertEqual(termination_state, decoded.termination_state)
+
+ def test_non_polling_round_trip_preserves_pending_outputs(self):
+ coder = _GrowthStateCoder(StrUtf8Coder(), never())
+ pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP)
+ state = _NonPollingGrowthState(pending)
+ decoded = coder.decode(coder.encode(state))
+ self.assertEqual(MAX_TIMESTAMP, decoded.pending.watermark)
+ self.assertEqual([('a', Timestamp(1)), ('b', Timestamp(2))],
+ [(o.value, o.timestamp) for o in decoded.pending.outputs])
+
+
+class GrowthTrackerTest(unittest.TestCase):
+ def test_poll_claims_dedups_and_checkpoints(self):
+ def poll(unused_element):
+ return PollResult.incomplete([_ts('a', 1), _ts('a', 1), _ts('b', 2)])
+
+ plan = _plan(_initial_polling(), poll)
+ self.assertEqual(['a', 'b'], [o.value for o in plan.emit])
+ self.assertFalse(plan.stop)
+ self.assertEqual(Timestamp(1), plan.watermark)
+ self.assertIsInstance(plan.primary, _NonPollingGrowthState)
+ self.assertIsInstance(plan.residual, _PollingGrowthState)
+ self.assertEqual(2, len(plan.residual.completed))
+
+ tracker = _GrowthRestrictionTracker(_initial_polling())
+ self.assertFalse(tracker.is_bounded())
+ self.assertTrue(tracker.try_claim(plan))
+ primary, residual = tracker.try_split(0)
+ self.assertIs(plan.primary, primary)
+ self.assertIs(plan.residual, residual)
+ self.assertTrue(tracker.check_done())
+
+ def explicit_watermark_poll(unused_element):
+ return PollResult.incomplete([_ts('c', 3)]).with_watermark(5)
+
+ plan = _plan(_initial_polling(), explicit_watermark_poll)
+ self.assertEqual(Timestamp(5), plan.watermark)
+ self.assertEqual(Timestamp(5), plan.residual.poll_watermark)
+
+ def test_second_round_repolls_and_dedups_against_completed(self):
+ polls = []
+
+ def poll(unused_element):
+ polls.append(len(polls))
+ if len(polls) == 1:
+ return PollResult.incomplete([_ts('a', 1), _ts('b', 2)])
+ return PollResult.incomplete([_ts('a', 1), _ts('c', 3)])
+
+ first = _plan(_initial_polling(), poll)
+ resumed = _plan(first.residual, poll)
+ self.assertEqual(2, len(polls))
+ self.assertEqual(['c'], [o.value for o in resumed.emit])
+
+ def test_termination_condition_sets_stop(self):
+ def poll(unused_element):
+ return PollResult.incomplete([_ts('a', 1)])
+
+ termination = after_total_of(10)
+ for now, expected_stop in [(10.0, False), (11.0, True)]:
+ with self.subTest(now=now):
+ plan = _plan(
+ _initial_polling(termination, Timestamp(0)),
+ poll,
+ now=now,
+ termination=termination)
+ self.assertEqual(expected_stop, plan.stop)
+
+ def test_non_polling_replays(self):
+ pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP)
+ restriction = _NonPollingGrowthState(pending)
+ plan = _replay_plan(restriction)
+ self.assertEqual(['a', 'b'], [o.value for o in plan.emit])
+ self.assertTrue(plan.stop)
+
+ tracker = _GrowthRestrictionTracker(restriction)
+ self.assertTrue(tracker.is_bounded())
+ self.assertTrue(tracker.try_claim(plan))
+ # A checkpoint after replay resumes an empty state that emits nothing.
+ primary, residual = tracker.try_split(0)
+ self.assertIs(restriction, primary)
+ self.assertEqual((), residual.pending.outputs)
+ self.assertTrue(tracker.check_done())
+
+ def test_terminal_split_residual_is_empty_for_all_stop_causes(self):
+ termination = after_total_of(Duration(10))
+ cases = [
+ (
+ 'reached_max',
+ never(),
+ _initial_polling(), lambda element: PollResult(
+ (TimestampedValue('a', MAX_TIMESTAMP), ), None),
+ 0.0),
+ (
+ 'complete',
+ never(),
+ _initial_polling(),
+ lambda element: PollResult.complete([_ts('a', 1)]),
+ 0.0),
+ (
+ 'after_total_of',
+ termination,
+ _initial_polling(termination, Timestamp(0)),
+ lambda element: PollResult.incomplete([_ts('a', 1)]),
+ 100.0),
+ ]
+ for name, condition, restriction, poll_fn, now in cases:
+ with self.subTest(name=name):
+ plan = _plan(restriction, poll_fn, now=now, termination=condition)
+ self.assertTrue(plan.stop)
+ self.assertIsInstance(plan.residual, _NonPollingGrowthState)
+ self.assertEqual((), plan.residual.pending.outputs)
+
+ def test_wrapper_chain_defers_merged_residual(self):
+ def poll(unused_element):
+ return PollResult.incomplete([_ts('a', 1), _ts('b', 2)])
+
+ plan = _plan(_initial_polling(), poll)
+ threadsafe = ThreadsafeRestrictionTracker(
+ _GrowthRestrictionTracker(_initial_polling()))
+ view = RestrictionTrackerView(threadsafe)
+ self.assertTrue(view.try_claim(plan))
+ view.defer_remainder(Duration(5))
+ residual, _ = threadsafe.deferred_status()
+ self.assertIsInstance(residual, _PollingGrowthState)
+ self.assertEqual(2, len(residual.completed))
+
+
+# Module-level so the poll function pickles by reference; the call counter is
+# shared within the single in-memory DirectRunner process.
+_POLL_CALLS = collections.defaultdict(int)
+
+
+def _growing_poll(prefix):
+ _POLL_CALLS[prefix] += 1
+ count = _POLL_CALLS[prefix]
+ outputs = [_ts('%s%d' % (prefix, i), i + 1) for i in range(count)]
+ if count >= 3:
+ return PollResult.complete(outputs)
+ return PollResult.incomplete(outputs)
+
+
+def _complete_poll(prefix):
+ return PollResult.complete([_ts(prefix + 'a', 1), _ts(prefix + 'b', 2)])
+
+
+def _windowed_group(kv, window=beam.DoFn.WindowParam):
+ return ((window.start, window.end), sorted(kv[1]))
+
+
+class WatchEndToEndTest(unittest.TestCase):
+ def _in_memory_pipeline(self):
+ return TestPipeline(
+ options=PipelineOptions(direct_running_mode='in_memory'))
+
+ def test_complete_outputs_values_and_timestamps(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(
+ Duration(1)).with_output_coder(StrUtf8Coder()))
+ assert_that(
+ output,
+ equal_to([
+ TestWindowedValue(('k:', 'k:a'), Timestamp(1), [GlobalWindow()]),
+ TestWindowedValue(('k:', 'k:b'), Timestamp(2), [GlobalWindow()]),
+ ]),
+ reify_windows=True)
+
+ def test_complete_advances_watermark_for_windowed_pipeline(self):
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(
+ Duration(1)).with_output_coder(StrUtf8Coder()))
+ grouped = (
+ output
+ | beam.WindowInto(FixedWindows(10))
+ | beam.Map(lambda kv: ('all', kv[1]))
+ | beam.GroupByKey()
+ | beam.Map(_windowed_group))
+ assert_that(
+ grouped,
+ equal_to([
+ ((Timestamp(0), Timestamp(10)), ['k:a', 'k:b']),
+ ]))
+
+ def test_multi_round_dedups_stops_and_is_per_input(self):
+ _POLL_CALLS.clear()
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['x:', 'y:'])
+ | Watch.growth_of(_growing_poll).with_poll_interval(
+ Duration(0.05)).with_output_coder(StrUtf8Coder()))
+ assert_that(
+ output,
+ equal_to([('x:', 'x:0'), ('x:', 'x:1'), ('x:', 'x:2'), ('y:', 'y:0'),
+ ('y:', 'y:1'), ('y:', 'y:2')]))
+ self.assertEqual(3, _POLL_CALLS['x:'])
+ self.assertEqual(3, _POLL_CALLS['y:'])
+
+ def test_rejects_non_deterministic_output_coder(self):
+ # No output coder resolves to PickleCoder, which is non-deterministic, so
+ # dedup could re-emit equal outputs. Expansion must reject it.
+ with self.assertRaises(ValueError):
+ with self._in_memory_pipeline() as p:
+ _ = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(Duration(1)))
+
+ def test_derives_output_type_from_input_and_coder(self):
+ # expand() types the (input, output) pairs from the input type and the
+ # resolved coder, so downstream stays typed without a manual hint.
+ with self._in_memory_pipeline() as p:
+ output = (
+ p | beam.Create(['k:'])
+ | Watch.growth_of(_complete_poll).with_poll_interval(
+ Duration(1)).with_output_coder(StrUtf8Coder()))
Review Comment:
Good catch, inference was missing. Expansion fell back to `PickleCoder`, and the determinism error from the last round then rejected it, which is what forced every test to pass a coder by hand. The output coder is now inferred. A `PollFn` can provide `default_output_coder`, and a plain function can annotate its return type as `PollResult[str]` so the element type resolves through the coder registry. Together with the key coder change, the fallback coder also works without any hint, so the explicit coders in the tests are gone.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Python] Add Watch transform with growth_of polling SDF [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #39023:
URL: https://github.com/apache/beam/pull/39023#issuecomment-4884314286
## Summary of Changes
Hello, I'm Gemini Code Assist[^1]! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces an experimental Watch transform to the Apache Beam Python SDK, designed to handle periodic polling of sources. It provides a robust mechanism for watching growing inputs, deduplicating results, and managing watermarks in an unbounded PCollection. The implementation leverages splittable DoFns to ensure polling remains durable across bundles and worker restarts, with comprehensive support for custom termination policies.
### Highlights
* **Experimental Watch Transform**: Introduced the Watch transform in the Python SDK, enabling periodic polling of growing sets of outputs for each input element.
* **Splittable DoFn Implementation**: Implemented the transform as a splittable DoFn, allowing for durable polling loops that self-checkpoint using defer_remainder.
* **Deduplication and Watermarking**: Added deduplication using 128-bit blake2b hashes of encoded outputs and a manual watermark estimator to manage progress.
* **Termination Conditions**: Included support for termination conditions such as 'never' and 'after_total_of' to control the polling lifecycle.
<details>
<summary><b>New Features</b></summary>
🧠 You can now enable **Memory** (public preview) to help **Gemini Code Assist** learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. **Click [here](https://codeassist.google/code-review/login) to enable Memory in your admin console.**
</details>
<details>
<summary><b>Using Gemini Code Assist</b></summary>
<br>
The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
<b>Invoking Gemini</b>
You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page.
Feature | Command | Description
--- | --- | ---
Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help | `/gemini help` | Displays a list of available commands.
<b>Customization</b>
To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
<b>Limitations & Feedback</b>
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with :thumbsup: and :thumbsdown: on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up [here](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
</details>
[^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org