Handle multi-key bundles in side-input runner - #38988
Conversation
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 implements support for multi-key bundles within the side-input runner. By introducing an 'activeKey' tracking mechanism, the changes allow the runner to correctly manage side-input processing across multiple keys within a single bundle, moving away from the previous assumption that bundles were strictly single-key. 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 activeKey boolean flag across SimpleParDoFn, StreamingKeyedWorkItemSideInputParDoFn, and StreamingSideInputDoFnRunner to track key lifecycle state and manage bundle initialization correctly. Feedback on these changes highlights a critical bug in StreamingKeyedWorkItemSideInputParDoFn where onStartKey() is invoked before the current key is written to state, as well as a suggestion to make the activeKey field private in StreamingSideInputDoFnRunner to improve encapsulation.
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.
| public void processElement(Object untypedElem) throws Exception { | ||
| if (helpers.fnRunner == null) { | ||
| boolean initialize = !activeKey || helpers.fnRunner == null; | ||
| if (initialize) { | ||
| // If we need to run reallyStartBundle in here, we need to make sure to switch the state | ||
| // sampler into the start state. | ||
| try (Closeable start = helpers.operationContext.enterStart()) { | ||
| helpers.reallyStartBundle(); | ||
| if (helpers.fnRunner == null) { | ||
| helpers.reallyStartBundle(); | ||
| } | ||
| onStartKey(); | ||
| } | ||
| } |
There was a problem hiding this comment.
In StreamingKeyedWorkItemSideInputParDoFn, onStartKey() reads the current key from state using keyValue().read(). However, when onStartKey() is called during processElement initialization, the new key has not yet been written to keyValue() (which normally happens later in onProcessWindowedValue). This means keyValue().read() will either return null (if the key is being processed for the first time) or the key from a previous bundle, leading to incorrect behavior or missed unblocking of elements/timers for the new key.
To fix this, we should extract the key from untypedElem and write it to keyValue() before calling onStartKey() in processElement.
| public void processElement(Object untypedElem) throws Exception { | |
| if (helpers.fnRunner == null) { | |
| boolean initialize = !activeKey || helpers.fnRunner == null; | |
| if (initialize) { | |
| // If we need to run reallyStartBundle in here, we need to make sure to switch the state | |
| // sampler into the start state. | |
| try (Closeable start = helpers.operationContext.enterStart()) { | |
| helpers.reallyStartBundle(); | |
| if (helpers.fnRunner == null) { | |
| helpers.reallyStartBundle(); | |
| } | |
| onStartKey(); | |
| } | |
| } | |
| public void processElement(Object untypedElem) throws Exception { | |
| boolean initialize = !activeKey || helpers.fnRunner == null; | |
| if (initialize) { | |
| // If we need to run reallyStartBundle in here, we need to make sure to switch the state | |
| // sampler into the start state. | |
| try (Closeable start = helpers.operationContext.enterStart()) { | |
| if (helpers.fnRunner == null) { | |
| helpers.reallyStartBundle(); | |
| } | |
| WindowedValue<KeyedWorkItem<K, InputT>> elem = (WindowedValue<KeyedWorkItem<K, InputT>>) untypedElem; | |
| keyValue().write(elem.getValue().key()); | |
| onStartKey(); | |
| } | |
| } |
| implements DoFnRunner<InputT, OutputT> { | ||
| private final DoFnRunner<InputT, OutputT> simpleDoFnRunner; | ||
| private final StreamingSideInputProcessor<InputT, W> sideInputProcessor; | ||
| boolean activeKey = false; |
There was a problem hiding this comment.
The activeKey field has default (package-private) visibility. It should be declared as private to maintain proper encapsulation and consistency with other classes like SimpleParDoFn and StreamingKeyedWorkItemSideInputParDoFn.
| boolean activeKey = false; | |
| private boolean activeKey = false; |
arunpandianp
left a comment
arunpandianp
left a comment
There was a problem hiding this comment.
Moving onStartKey out of startBundle could simplify things. Not sure if there is a reason that prevents us from doing so.
There was a problem hiding this comment.
There could be keys with only timers and no elements, do we need to call onStartKey in processTimers?
There was a problem hiding this comment.
I think helpers.processTimers already does this
| @Override | ||
| public void finishKey(Object key) throws Exception {} | ||
| public void finishKey(Object key) throws Exception { | ||
| this.activeKey = false; |
There was a problem hiding this comment.
IIUC there could be keys with no elements or timers. Should we move the onStartKey call here to handle such keys?
Instead of calling onStartKey in startBundle and processElements, we could call it in processElements, processTimers and finishKey.
| boolean hasState = helpers.hasState(); | ||
|
|
||
| // TODO(relax): We should be able to get this without writing it to state! | ||
| // TODO(relax): We should be able to get this without writing it to state! To make this work, |
There was a problem hiding this comment.
finishKey has the key. If we move onStartKey to finishKey, we should be able to make this optimization.
| @Override | ||
| public void finishKey(Object key) throws Exception {} | ||
| public void finishKey(Object key) throws Exception { | ||
| this.activeKey = false; |
There was a problem hiding this comment.
should we call sideInputFetcher.persist(); in finishKey
|
Assigning reviewers: R: @chamikaramj added as fallback since no labels match configuration 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). |
| if (sideInputFetcher != null) { | ||
| sideInputFetcher.persist(); | ||
| } | ||
| this.activeKey = false; |
There was a problem hiding this comment.
nit: redundant with the logic in finishKey and can be removed.
| public void processTimers() {} | ||
| public void processTimers() throws Exception { | ||
| if (!activeKey) { | ||
| onStartKey(null); |
There was a problem hiding this comment.
can we pass key from stepContext instead of null?
There was a problem hiding this comment.
Removed parameter instead, as we never use it.
| this.onStartKey = | ||
| k -> { | ||
| onStartKey.accept(k); | ||
| this.activeKey = false; |
There was a problem hiding this comment.
should this be this.activeKey = true;
| void finishBundle(StreamingSideInputProcessor<?, ?> sideInputProcessor) throws Exception { | ||
| if (fnRunner != null) { | ||
| fnRunner.finishBundle(); | ||
| if (sideInputProcessor != null) { |
There was a problem hiding this comment.
I think the handleFinishKeyOrBundle calls can be removed from finishBundle
There was a problem hiding this comment.
do we need thesideInputProcessor.handleFinishKeyOrBundle call here?
| try (Closeable start = operationContext.enterStart()) { | ||
| reallyStartBundle(); | ||
| startKey.run(); | ||
| this.onStartKey.accept((K) context.stateInternals().getKey()); |
There was a problem hiding this comment.
this should to be under a if (!activeKey) { check.
In processElements onStartKey is called outside operationContext.enterStart(), here we call it inside operationContext.enterStart(). Move it outside the enterStart block?
| boolean activeKey = false; | ||
| Consumer<K> onStartKey; |
There was a problem hiding this comment.
make these private? onStartKey can also be final
| protected void onProcessWindowedValue(WindowedValue<KeyedWorkItem<K, InputT>> elem) { | ||
| // TODO: Get rid of this! | ||
| // TODO: Get rid of this once we know the current key. | ||
| final K key = elem.getValue().key(); |
There was a problem hiding this comment.
can we remove keyValue().write(key);?
If it is kept for update compatibility, need to update comment and maybe add a TODO to remove on newer jobs with updateCompatibility flag.
There was a problem hiding this comment.
Yes - it's update compatible, it just means that we won't support downgrading the Beam version.
I also considered adding code to delete the key state if it's there, but I think it's best not to
| implements DoFnRunner<InputT, OutputT> { | ||
| private final DoFnRunner<InputT, OutputT> simpleDoFnRunner; | ||
| private final StreamingSideInputProcessor<InputT, W> sideInputProcessor; | ||
| boolean activeKey = false; |
| } | ||
|
|
||
| private void tryUnblockElements() { | ||
| sideInputProcessor.tryUnblockElements( |
There was a problem hiding this comment.
do we need to recreate sideInputProcessor every key like in other classses?
5050955 to
0359ed9
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors streaming side input handling in the Dataflow worker to support multi-key bundles by introducing lazy initialization of side input fetchers and per-key lifecycle hooks. The review feedback highlights a logic bug in SimpleParDoFnHelpers.processTimers where key initialization could be skipped for keys with only timers. Additionally, the reviewer identified several potential memory leaks across multiple classes (StreamingSideInputDoFnRunner, SimpleParDoFn, StreamingKeyedWorkItemSideInputParDoFn, and PartialGroupByKeyParDoFns) where stale side input processors or fetchers are not cleared at the end of finishKey or finishBundle execution.
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 (timer != null && fnRunner == null) { | ||
| // If we need to run reallyStartBundle in here, we need to make sure to switch the state | ||
| // sampler into the start state. | ||
| try (Closeable start = operationContext.enterStart()) { | ||
| reallyStartBundle(); | ||
| startKey.run(); | ||
| } | ||
|
|
||
| if (!activeKey) { | ||
| this.onStartKey.accept((K) context.stateInternals().getKey()); | ||
| } | ||
| } |
There was a problem hiding this comment.
The if (!activeKey) check and onStartKey.accept(...) call are nested inside the if (timer != null && fnRunner == null) block. However, if fnRunner is already initialized (which is the case when reallyStartBundle is called in startBundle), fnRunner == null is false, so onStartKey is never called even if activeKey is false. This means the side input processor won't be initialized for keys that only have timers, leading to potential NPEs or incorrect behavior. We should move the if (!activeKey) block outside of the fnRunner == null check, but still within the timer != null check.
if (timer != null) {
if (fnRunner == null) {
// If we need to run reallyStartBundle in here, we need to make sure to switch the state
// sampler into the start state.
try (Closeable start = operationContext.enterStart()) {
reallyStartBundle();
}
}
if (!activeKey) {
this.onStartKey.accept((K) context.stateInternals().getKey());
}
}| @Override | ||
| public <KeyT extends @Nullable Object> void finishKey(KeyT key) { | ||
| if (!activeKey) { | ||
| // This means that there were no elements for this key. Try to unblock any queued elements. | ||
| onNewKey(); | ||
| } | ||
| Preconditions.checkStateNotNull(sideInputProcessor).handleFinishKeyOrBundle(); | ||
| simpleDoFnRunner.finishKey(key); | ||
| this.activeKey = false; | ||
| } | ||
|
|
||
| @Override | ||
| public void finishBundle() { | ||
| simpleDoFnRunner.finishBundle(); | ||
| sideInputProcessor.handleFinishBundle(); | ||
| this.activeKey = false; | ||
| } |
There was a problem hiding this comment.
To prevent memory leaks, sideInputProcessor should be set to null at the end of finishKey and finishBundle. Otherwise, the runner will hold onto the StreamingSideInputProcessor (and its associated fetcher and state) of the last processed key indefinitely. Additionally, the onTimer method (not shown in the diff but present in the class) also needs to call onNewKey() if !activeKey to ensure the key context is initialized and any blocked elements are processed before the timer runs.
| @Override | |
| public <KeyT extends @Nullable Object> void finishKey(KeyT key) { | |
| if (!activeKey) { | |
| // This means that there were no elements for this key. Try to unblock any queued elements. | |
| onNewKey(); | |
| } | |
| Preconditions.checkStateNotNull(sideInputProcessor).handleFinishKeyOrBundle(); | |
| simpleDoFnRunner.finishKey(key); | |
| this.activeKey = false; | |
| } | |
| @Override | |
| public void finishBundle() { | |
| simpleDoFnRunner.finishBundle(); | |
| sideInputProcessor.handleFinishBundle(); | |
| this.activeKey = false; | |
| } | |
| @Override | |
| public <KeyT extends @Nullable Object> void finishKey(KeyT key) { | |
| if (!activeKey) { | |
| // This means that there were no elements for this key. Try to unblock any queued elements. | |
| onNewKey(); | |
| } | |
| Preconditions.checkStateNotNull(sideInputProcessor).handleFinishKeyOrBundle(); | |
| simpleDoFnRunner.finishKey(key); | |
| this.activeKey = false; | |
| this.sideInputProcessor = null; | |
| } | |
| @Override | |
| public void finishBundle() { | |
| simpleDoFnRunner.finishBundle(); | |
| this.activeKey = false; | |
| this.sideInputProcessor = null; | |
| } |
| @Override | ||
| public void finishKey(Object key) throws Exception {} | ||
| public void finishKey(Object key) throws Exception { | ||
| helpers.finishKey(key, sideInputProcessor); | ||
| } |
There was a problem hiding this comment.
To prevent memory leaks, sideInputProcessor should be set to null after helpers.finishKey is called, so that the SimpleParDoFn instance does not hold onto the stale processor of the last processed key.
| @Override | |
| public void finishKey(Object key) throws Exception {} | |
| public void finishKey(Object key) throws Exception { | |
| helpers.finishKey(key, sideInputProcessor); | |
| } | |
| @Override | |
| public void finishKey(Object key) throws Exception { | |
| helpers.finishKey(key, sideInputProcessor); | |
| this.sideInputProcessor = null; | |
| } |
| @Override | ||
| public void finishKey(Object key) throws Exception {} | ||
| public void finishKey(Object key) throws Exception { | ||
| helpers.finishKey((K) key, sideInputProcessor); | ||
| } |
There was a problem hiding this comment.
To prevent memory leaks, sideInputProcessor should be set to null after helpers.finishKey is called, so that the StreamingKeyedWorkItemSideInputParDoFn instance does not hold onto the stale processor of the last processed key.
| @Override | |
| public void finishKey(Object key) throws Exception {} | |
| public void finishKey(Object key) throws Exception { | |
| helpers.finishKey((K) key, sideInputProcessor); | |
| } | |
| @Override | |
| public void finishKey(Object key) throws Exception { | |
| helpers.finishKey((K) key, sideInputProcessor); | |
| this.sideInputProcessor = null; | |
| } |
| @Override | ||
| public void finishBundle() throws Exception { | ||
| groupingTable.flush(receiver); | ||
| sideInputFetcher.persist(); | ||
| if (sideInputFetcher != null) { | ||
| sideInputFetcher.persist(); | ||
| } | ||
| } |
There was a problem hiding this comment.
To prevent memory leaks, sideInputFetcher should be set to null in finishBundle if it is not null.
| @Override | |
| public void finishBundle() throws Exception { | |
| groupingTable.flush(receiver); | |
| sideInputFetcher.persist(); | |
| if (sideInputFetcher != null) { | |
| sideInputFetcher.persist(); | |
| } | |
| } | |
| @Override | |
| public void finishBundle() throws Exception { | |
| groupingTable.flush(receiver); | |
| if (sideInputFetcher != null) { | |
| sideInputFetcher.persist(); | |
| this.sideInputFetcher = null; | |
| } | |
| } |
|
@arunpandianp friendly ping |
|
Can we add a CHANGES.md entry about StreamingKeyedWorkItemSideInputParDoFn.java breaking compatibility? Few nits other than that the changes look good to me. |
| public void finishBundle() throws Exception { | ||
| groupingTable.flush(receiver); | ||
| sideInputFetcher.persist(); | ||
| if (sideInputFetcher != null) { |
| void finishBundle(StreamingSideInputProcessor<?, ?> sideInputProcessor) throws Exception { | ||
| if (fnRunner != null) { | ||
| fnRunner.finishBundle(); | ||
| if (sideInputProcessor != null) { |
There was a problem hiding this comment.
do we need thesideInputProcessor.handleFinishKeyOrBundle call here?
| implements DoFnRunner<InputT, OutputT> { | ||
| private final DoFnRunner<InputT, OutputT> simpleDoFnRunner; | ||
| private final StreamingSideInputProcessor<InputT, W> sideInputProcessor; | ||
| boolean activeKey = false; |
|
Reminder, please take a look at this pr: @chamikaramj |
|
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: @chamikaramj added as fallback since no labels match configuration Available commands:
|
|
Reminder, please take a look at this pr: @chamikaramj |
|
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: @chamikaramj added as fallback since no labels match configuration Available commands:
|
|
Reminder, please take a look at this pr: @chamikaramj |
2 participants
R: @arunpandianp