[Python] Add Watch transform with growth_of polling SDF#39023
Conversation
Eliaaazzz
force-pushed
the
python-watch-transform
branch
from
0cd6937 to
5b039a0
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ 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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Eliaaazzz
force-pushed
the
python-watch-transform
branch
from
5b039a0 to
9acab87
Compare
Abacn
left a comment
Abacn
left a comment
There was a problem hiding this comment.
Thanks, went through watch.py changes
| pending: PollResult | ||
|
|
||
|
|
||
| _GrowthState = Any # Union[_PollingGrowthState, _NonPollingGrowthState] |
There was a problem hiding this 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?
There was a problem hiding this 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.
| # ------------------------------------------------------------------------------ | ||
|
|
||
|
|
||
| class _HashCode128Coder(Coder): |
There was a problem hiding this 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
There was a problem hiding this 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.
| return True | ||
|
|
||
|
|
||
| class _TimestampedValueCoder(Coder): |
There was a problem hiding this comment.
The Python SDK ships no coder for
TimestampedValue
A little bit surprising. Likely due to Python SDK has some special handling:
beam/sdks/python/apache_beam/runners/common.py
Line 1773 in a86e611
converting TimestampedValue to a WindowedValue, and WindowedValue has a standard coder
I'm considering make TimestampedValue a public interface like Java SDK we current have (could be a separate PR)
There was a problem hiding this 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.
| output_coder: Optional[Coder] = None, | ||
| now_fn: Optional[Callable[[], float]] = None): | ||
| super().__init__() | ||
| self._poll_fn = poll_fn |
There was a problem hiding this 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.
There was a problem hiding this 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.
| """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`. |
There was a problem hiding this comment.
As it's declared in __all__, consider add an example about creating a PollFn that is callable and has default_output_coder implemented
There was a problem hiding this 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.
| output_coder = hint or coders.PickleCoder() | ||
| if not output_coder.is_deterministic(): | ||
| _LOGGER.warning( | ||
| 'Watch dedup uses a non-deterministic output coder (%s); equal ' |
There was a problem hiding this 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?
There was a problem hiding this 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.
|
|
||
| element = holder[0] | ||
| now = Timestamp.of(self._now()) | ||
| result = self._poll_fn(element) |
There was a problem hiding this 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?
There was a problem hiding this 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.
Add an experimental Watch transform that watches a growing set of outputs per input element. Watch.growth_of(poll_fn) runs a periodic poll loop as a splittable DoFn and emits an unbounded PCollection of (input, output) pairs. Each process() performs one poll round and self-checkpoints via defer_remainder. New outputs are deduplicated with a stable 128-bit blake2b hash of the encoded output, and a manual watermark estimator advances per poll. Per-input termination supports never() and after_total_of(); polling also stops when a poll returns PollResult.complete(). The DoFn is its own RestrictionProvider, and restriction state serializes through a tagged GrowthState coder. Tests cover termination conditions, coder round-trips, the restriction tracker claim/checkpoint/dedup logic, and DirectRunner end-to-end runs.
Eliaaazzz
force-pushed
the
python-watch-transform
branch
from
9acab87 to
5a54aa6
Compare
|
Thanks for the thorough review. I addressed all the comments and pushed. PTAL when you have time. |
Summary of ChangesHello, I'm Gemini Code Assist1! 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
New Features🧠 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 to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization 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 Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 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. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an experimental Watch transform for the Apache Beam Python SDK, which continuously polls for new outputs per input element using a Splittable DoFn. The review feedback highlights critical issues with watermark estimation during both the replay and active polling phases. Specifically, unconditionally setting the watermark to the input element's timestamp before yielding outputs can prematurely advance the watermark and cause outputs to be treated as late data. The reviewer recommends holding the watermark at the minimum of the input timestamp and the planned watermark before yielding, and releasing the watermark hold to MAX_TIMESTAMP upon termination to prevent the pipeline watermark from getting stuck.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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 |
There was a problem hiding this 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.
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| # 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) |
There was a problem hiding this 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.
| # 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) | |
| # 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) |
|
Assigning reviewers: R: @shunping for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
2 participants
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
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
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, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.