[Python] Bound Watch state with a timestamp cursor - #39090
Conversation
85476e4 to
71d6282
Compare
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 PR introduces a more scalable deduping strategy for the Watch transform in the Python SDK. By allowing users to opt into a timestamp-based cursor, the transform can now handle high-cardinality sources with O(1) state growth, avoiding the performance degradation associated with the default hash-based deduping set. 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 Python SDK, which continuously polls and watches a growing set of outputs for each input element. It supports exact deduplication using 128-bit hashes or O(1) state deduplication using a high-water-mark timestamp cursor. The feedback suggests two improvements: first, ensuring that timestamps are normalized to Timestamp objects before comparison in cursor mode to prevent TypeError exceptions in Python 3; second, converting the input outputs iterable to a list during normalization to safely handle single-use generators or iterators.
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.
| new_outputs = [ | ||
| output for output in result.outputs | ||
| if cursor is None or output.timestamp > cursor | ||
| ] |
There was a problem hiding this comment.
In _GrowthRestrictionTracker.try_claim, when self._cursor_mode is enabled, the code compares output.timestamp > cursor. However, output.timestamp can be a raw float or integer if the user directly instantiated TimestampedValue without using Timestamp objects (as TimestampedValue accepts int or float for its timestamp). Since cursor is a Timestamp object (or None), comparing a Timestamp with a float or integer will raise a TypeError in Python 3. To prevent comparison errors, normalize output.timestamp to a Timestamp object using Timestamp.of(output.timestamp) before comparison.
| new_outputs = [ | |
| output for output in result.outputs | |
| if cursor is None or output.timestamp > cursor | |
| ] | |
| cursor = restriction.cursor | |
| new_outputs = [ | |
| output for output in result.outputs | |
| if cursor is None or Timestamp.of(output.timestamp) > cursor | |
| ] |
There was a problem hiding this comment.
Thanks — this can't happen in practice: TimestampedValue.__init__ coerces its argument with Timestamp.of(timestamp), so by the time an output reaches the tracker its timestamp is always a Timestamp. (The rebased branch reworked this code path, but the invariant is the same.)
| for output in outputs: | ||
| if isinstance(output, TimestampedValue): | ||
| normalized.append(output) | ||
| else: | ||
| normalized.append(TimestampedValue(output, default_ts)) |
There was a problem hiding this comment.
When iterating over outputs, if the user passes a single-use generator or iterator, iterating over it to normalize elements will consume it. Since outputs is iterated again when constructing the PollResult or during normalization, this can lead to empty outputs or unexpected behavior. Consider converting outputs to a list or tuple first before normalization.
| for output in outputs: | |
| if isinstance(output, TimestampedValue): | |
| normalized.append(output) | |
| else: | |
| normalized.append(TimestampedValue(output, default_ts)) | |
| outputs = list(outputs) | |
| normalized = [] | |
| for output in outputs: | |
| if isinstance(output, TimestampedValue): | |
| normalized.append(output) | |
| else: | |
| normalized.append(TimestampedValue(output, default_ts)) |
There was a problem hiding this comment.
After the rebase onto the merged Watch (#39023), PollResult._normalize materializes outputs into a tuple in one pass, so a single-use iterator is only consumed once and the double-iteration this pointed at no longer exists.
|
Assigning reviewers: R: @jrmccluskey 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). |
71d6282 to
ac7d761
Compare
|
Reminder, please take a look at this pr: @jrmccluskey |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @claudevdm for label python. Available commands:
|
Opt-in timestamp_cursor=True dedups by a high-water-mark timestamp instead of by key identity: Watch keeps only the greatest event time it has emitted for an input and emits the polled outputs strictly past it, so the per-input state and per-checkpoint encoding are O(1) regardless of how many outputs the input produces. For sources whose outputs carry strictly increasing event-time timestamps; the default exact hash dedup remains for arbitrary-relisting or out-of-order sources. Also documents the event-time/watermark contract, adds a throttled late-emission warning, and reuses the parent dedup map on idle rounds so steady-state empty polls stay O(1) in hash mode too.
ac7d761 to
b30751a
Compare
| poll_interval=Duration(seconds=5), | ||
| termination=after_total_of(60)) | ||
|
|
||
| Watermark and event-time contract |
There was a problem hiding this comment.
Please keep doc and comments concise thoughout this PR.
Agent has a tendancy of stacking chat history into doc. Sometimes need a fresh session to rewrite things
There was a problem hiding this comment.
You are right, that is exactly what happened here. I use AI tooling for drafting and I let the docs absorb the review discussion instead of keeping them clean. I rewrote the docstrings and comments this PR adds from a clean slate, each fact is stated once now, and I have gone through every line myself.
| if tag == 1: | ||
| watermark, outputs = self._non_polling_coder.decode(payload) | ||
| return _NonPollingGrowthState(PollResult(tuple(outputs), watermark)) | ||
| if tag == 2: |
There was a problem hiding this comment.
Prefer enum over raw numbers
There was a problem hiding this comment.
Done, replaced the raw tag numbers with an IntEnum.
| state.poll_watermark, | ||
| list(state.completed.items()))) | ||
| return self._envelope_coder.encode((0, payload)) | ||
| list(state.completed.items()), |
There was a problem hiding this comment.
We should not encode the items when use cursors. The whole point of cursor (with_timestamp) mode is that the saved state doesn't grow with the number of items in poll therefore scalable
There was a problem hiding this comment.
Agreed, the cursor state should not carry the items at all. The completed map was already empty at runtime in cursor mode, but the byte format still reserved a slot for it. Cursor states now encode only the termination state, the poll watermark and the cursor, and a new test asserts the encoded size stays constant no matter how many outputs a round claims.
There was a problem hiding this comment.
Went one step further and re-derived the payload from what resume actually reads: poll_watermark was never read either, the watermark is restored from the estimator state the runner persists. The cursor payload is now just the termination state and the cursor.
There was a problem hiding this comment.
This round is a single follow up commit on top of the reviewed commit, bc56cb0, covering the three comments. I force pushed by mistake earlier today and restored the reviewed commit so the delta since your review is visible.
There was a problem hiding this comment.
One more follow up commit, 098623b: the hash to cursor switch now seeds the cursor from the hash map's greatest recorded event time instead of starting empty, so the switch cannot re-emit already seen outputs.
dde3905 to
3d93c6f
Compare
Encode cursor states as (termination_state, cursor) behind an IntEnum envelope tag, drop a stale cursor on hash rounds so the two dedup states never mix, pin the O(1) encoding with a size test, and rewrite the docstrings and comments this PR adds to be concise.
3d93c6f to
bc56cb0
Compare
A restriction carried over from hash dedup starts with no cursor, so its first cursor round re-emitted already seen outputs. The cursor now seeds from the greatest event time the hash map recorded, keeping the switch exactly once. Also clarifies the relisting guidance in the timestamp_cursor docs.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #39090 +/- ##
============================================
+ Coverage 60.52% 60.56% +0.04%
Complexity 20659 20659
============================================
Files 3324 3325 +1
Lines 318309 318704 +395
Branches 17249 17249
============================================
+ Hits 192644 193018 +374
- Misses 117036 117057 +21
Partials 8629 8629
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:
|
| ``output_key_fn(output)`` when one is given. The key coder is inferred when | ||
| not passed explicitly and converted to its deterministic form, so equal keys | ||
| hash equally across workers and restarts. | ||
| hash equally across workers and restarts. ``timestamp_cursor=True`` replaces |
There was a problem hiding this comment.
The last sentence sounds cryptic to users. How about this:
(a new paragraph): By default, watch transform internally stores the hash of all items seen. If the incremented items returned by the poll function guarantees monotonic timestamp growth (new items on next poll have timestamp larger than the largest of the previous poll), consider set timestamp_cursor=True for better performance, as it replaces the hash dedup with an O(1) event-time cursor; see :class:Watch.
There was a problem hiding this comment.
Thanks, that reads much better. Applied your wording as a new paragraph in f80a95b, with two small grammar touch ups.
Resolves the watch.py conflict with fc0d989, keeping the PEP 585 spellings from apache#39527.
2 participants
Update (rebased): #39023 merged, so this is again a single commit on master. The merged Watch reworked the internals this PR was written against (polling moved from the tracker into the SDF's
process()), so the cursor was re-implemented on the new architecture rather than mechanically rebased. The API is nowWatch(..., timestamp_cursor=True)to match the merged constructor-kwargs style. Semantics and the load-test rationale are unchanged. Hash-mode restriction states keep the exact pre-cursor byte format (cursor states use a new envelope tag), a hash-mode restriction switched to cursor mode retires its hash map and seeds the cursor from the map's greatest recorded event time, so the switch stays exactly once, and the out-of-order warning no longer fires while the watermark still holds the element-timestamp seed. Relates to #18459 (the unboundedcompleteddedup set).By default Watch dedups by value identity and keeps one hash per distinct output, so the per-input state grows without bound. This adds an opt-in
Watch(..., timestamp_cursor=True)that dedups by a high-water-mark timestamp instead: Watch keeps only the greatest event time it has emitted for an input and emits the polled outputs strictly past it, then advances the cursor. No hash set is kept and the poll result is not hashed, so the per-input state and per-checkpoint encoding are O(1) regardless of how many outputs the input produces. It is for sources whose outputs carry strictly increasing event-time timestamps; an output at or below the cursor is treated as already seen, so the default exact dedup remains for arbitrary-relisting or out-of-order sources. This also documents the event-time/watermark contract and adds a throttled late-emission warning.Testing: 28 unit tests and 5 subtests on the in-memory DirectRunner; yapf 0.43.0, isort 7.0.0, and pylint clean. Load tested on Dataflow Runner v2 (us-east1) against the default hash dedup, 4 inputs emitting 15000 outputs per 1s round over a 240s budget: the cursor processed 6,825,000 outputs versus the default's 1,650,000 in the same budget (4.1x), because the default re-serializes its growing dedup set on every checkpoint (it reached about 300-495K entries per input) while the cursor writes one timestamp; the gap widens with runtime. Both stayed exactly-once and reached JOB_STATE_DONE. Jobs on apache-beam-testing: cursor
2026-06-24_08_45_05-17501676526076065162, default2026-06-24_08_45_23-5060488636914380336.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
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.