[bvolpato]

AsyncWrapper tracks local futures by value (or id_fn result), while runner state and timers are scoped to a key and window. For inputs ('key1', 7) and ('key2', 7), the second key can consume the first key's result and clear its own state without running its input. Equal identifiers across windows have the same problem. Timer cleanup can also cancel work belonging to another window of the same key.

Track futures by (encoded_key, window, identifier) and pass the window through scheduling, timer callbacks, and retries. The key uses the existing bag-state key coder: its encoded identity survives state decoding, including NaN keys, and distinguishes numeric keys that compare equal in Python, such as -0.0 and 0.0. Cleanup only removes orphaned work from the current encoded key and window. The persisted bag-state format stays unchanged.

Testing

From sdks/python, with the SDK and test dependencies installed:

python -m pytest apache_beam/transforms/async_dofn_test.py -q -n 2

All 48 tests and 10 subtests pass on Python 3.12.

The cases exercise both the thread-pool and asyncio implementations: equal values and custom IDs across keys/windows, cleanup preserving another window's pending and completed work, and timer-driven rescheduling. The four cross-key regression cases fail against the upstream implementation. Additional state-coder roundtrip cases cover NaN completion, signed-zero and mixed numeric-key separation, and cancellation preserving another encoded key's work. Those cases fail with raw Python keys in the tracking map.

YAPF 0.43.0, Ruff 0.15.22, and git diff --check pass for the changed files. Validation uses the actual wrapper with state/timer doubles, including the real bag-state coder for the numeric-key cases; no end-to-end runner test was added.

Downsides

Each local tracking entry now includes encoded key bytes and its window. Inputs with matching identifiers in different keys or windows execute independently and can occupy separate buffer slots. The configured buffer limit is unchanged.

---

• Describe the bug and include reproducible regression tests.
• Update CHANGES.md with the behavior change.
• Apache Individual Contributor License Agreement, if required for this contribution.

[bvolpato]

Run Python_ML PreCommit 3.12

[bvolpato]

Run Python_ML PreCommit 3.10

[bvolpato]

I reproduced the key collision with the two regression tests added here. Each test runs in both asyncio and thread-pool modes. The test file was copied from this PR unchanged and run against each revision's async_dofn.py with Python 3.12.13:

• Original pre-fix parent ffcbd7e464f8244bd623d23e7fa9fcc750e81bbb: 4 failed.
• Current Apache Beam master b4db8d2a11d6b2796a98a149ea0c6709d97d6b36: 4 failed.
• This PR head ec16d879834ca4b7da40980cbac7dbfbbb51abf7: 4 passed.

On both failing revisions, committing key 2's state returned key 1's result for equal values and for matching custom IDs.

To rerun, from the repository root:

git show ec16d879834ca4b7da40980cbac7dbfbbb51abf7:sdks/python/apache_beam/transforms/async_dofn_test.py > /tmp/async_dofn_pr_test.py
cd sdks/python
PYTHONPATH=. python -m pytest /tmp/async_dofn_pr_test.py -k 'equal_values_in_different_keys or custom_ids_are_scoped_to_key' -vv

Run those commands with the source checkout at each revision above; the PYTHONPATH makes the tests exercise that checkout's SDK code.

[github-actions]

Assigning reviewers:

R: @claudevdm 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).

[github-actions]

Reminder, please take a look at this pr: @claudevdm

[github-actions]

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @jrmccluskey for label python.

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)

[claudevdm]

Thanks!