[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
| 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
| 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
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 drops its dead hash map, 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.