[Python] Add Watch transform with growth_of polling SDF

[Eliaaazzz]

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:

• 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.

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.

[codecov]

Codecov Report

❌ Patch coverage is 88.21656% with 37 lines in your changes missing coverage. Please review.
✅ Project coverage is 54.15%. Comparing base (8801afd) to head (5b039a0).
⚠️ Report is 1065 commits behind head on master.

Files with missing lines	Patch %	Lines
sdks/python/apache_beam/io/watch.py	88.21%	37 Missing ⚠️

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              

Flag	Coverage Δ
python	78.76% <88.21%> (-2.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Abacn]

Thanks, went through watch.py changes

[Abacn]

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?

[Eliaaazzz]

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.

[Abacn]

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

[Eliaaazzz]

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.

[Abacn]

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

elif isinstance(result, TimestampedValue):

converting TimestampedValue to a WindowedValue, and WindowedValue has a standard coder

beam/sdks/python/apache_beam/coders/typecoders.py

Line 96 in a86e611

self._register_coder_internal(

I'm considering make TimestampedValue a public interface like Java SDK we current have (could be a separate PR)

[Eliaaazzz]

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.

[Abacn]

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.

[Eliaaazzz]

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.

[Abacn]

As it's declared in __all__, consider add an example about creating a PollFn that is callable and has default_output_coder implemented

[Eliaaazzz]

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.

[Abacn]

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?

[Eliaaazzz]

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.

[Abacn]

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?

[Eliaaazzz]

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.

[Eliaaazzz]

Thanks for the thorough review. I addressed all the comments and pushed. PTAL when you have time.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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.

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 Assist

The 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 /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.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

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.

[gemini-code-assist]

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

[Abacn]

Could you please check these two comments?

[Eliaaazzz]

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.

[gemini-code-assist]

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.

Suggested change

	# 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)

[Eliaaazzz]

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.

[github-actions]

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).

[Abacn]

Thanks, the changes seem to raising a few more questions, commented below

Please do not force push after PR received initial review, unless there is a merge conflict (even so preserve the number of commits during rebase)

[Abacn]

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.

[Eliaaazzz]

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.

[Abacn]

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.

[Eliaaazzz]

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.

[Abacn]

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.

[Eliaaazzz]

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.

[Abacn]

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

[Eliaaazzz]

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.

[Eliaaazzz]

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.