Handle multi-key bundles in side-input runner by reuvenlax · Pull Request #38988 · apache/beam

reuvenlax · 2026-06-16T18:02:52Z

gemini-code-assist · 2026-06-16T18:03:07Z

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

Key Tracking State: Introduced an 'activeKey' boolean flag across multiple runner classes to track the state of the current key being processed.
Bundle Initialization Logic: Refactored 'processElement' and 'startBundle' methods to conditionally initialize and trigger side-input processing based on the 'activeKey' state, allowing for multi-key bundle support.
Lifecycle Management: Updated 'finishKey' and 'finishBundle' methods to reset the 'activeKey' flag, ensuring proper cleanup between keys and bundles.

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.

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

gemini-code-assist · 2026-06-16T18:05:12Z

  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();
      }
    }


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.

Suggested change

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();

}

}

gemini-code-assist · 2026-06-16T18:05:12Z

    implements DoFnRunner<InputT, OutputT> {
  private final DoFnRunner<InputT, OutputT> simpleDoFnRunner;
  private final StreamingSideInputProcessor<InputT, W> sideInputProcessor;
+  boolean activeKey = false;


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.

Suggested change

boolean activeKey = false;

private boolean activeKey = false;

arunpandianp

Moving onStartKey out of startBundle could simplify things. Not sure if there is a reason that prevents us from doing so.

arunpandianp · 2026-06-16T18:13:05Z


  @Override
  public void processTimers() throws Exception {



There could be keys with only timers and no elements, do we need to call onStartKey in processTimers?

arunpandianp · 2026-06-16T18:15:02Z

  @Override
-  public void finishKey(Object key) throws Exception {}
+  public void finishKey(Object key) throws Exception {
+    this.activeKey = false;


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.

arunpandianp · 2026-06-16T18:16:49Z

      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,


finishKey has the key. If we move onStartKey to finishKey, we should be able to make this optimization.

arunpandianp · 2026-06-16T18:21:04Z

  @Override
-  public void finishKey(Object key) throws Exception {}
+  public void finishKey(Object key) throws Exception {
+    this.activeKey = false;


should we call sideInputFetcher.persist(); in finishKey

github-actions · 2026-06-16T21:07:10Z

Assigning reviewers:

R: @chamikaramj added as fallback since no labels match configuration

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

reuvenlax added 2 commits June 16, 2026 10:59

handle multi-key bundles in side-input ronnre

1575cad

one more place

68f4db2

github-actions Bot added runners dataflow labels Jun 16, 2026

gemini-code-assist Bot reviewed Jun 16, 2026

View reviewed changes

arunpandianp reviewed Jun 16, 2026

View reviewed changes

spotless

e2fdc7f

github-actions Bot added the Next Action: Reviewers label Jun 16, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Handle multi-key bundles in side-input runner#38988

Handle multi-key bundles in side-input runner#38988
reuvenlax wants to merge 3 commits into
apache:masterfrom
reuvenlax:side_input_finish_key

reuvenlax commented Jun 16, 2026

Uh oh!

gemini-code-assist Bot commented Jun 16, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jun 16, 2026

Uh oh!

gemini-code-assist Bot Jun 16, 2026

Uh oh!

arunpandianp left a comment

Uh oh!

arunpandianp Jun 16, 2026

Uh oh!

arunpandianp Jun 16, 2026

Uh oh!

arunpandianp Jun 16, 2026

Uh oh!

arunpandianp Jun 16, 2026

Uh oh!

github-actions Bot commented Jun 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	boolean activeKey = false;
	private boolean activeKey = false;

Conversation

reuvenlax commented Jun 16, 2026

Uh oh!

gemini-code-assist Bot commented Jun 16, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

arunpandianp left a comment

Choose a reason for hiding this comment

Uh oh!

arunpandianp Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

arunpandianp Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

arunpandianp Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

arunpandianp Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants