[Dataflow Streaming] [Multi Key] Introduce KeyGroupWorkQueue and integrate with BoundedQueueExecutor by arunpandianp · Pull Request #38767 · apache/beam

arunpandianp · 2026-06-02T09:06:33Z

- Add Uint128Proto and key_group to WorkItem in windmill.proto.
- Update Work model to support key groups.
- Implement KeyGroupWorkQueue for grouping tasks by key group. KepGroupWorkQueue is a FIFO queue that allows polling elements based on global order and also by order within a key group.
- Integrate BoundedQueueExecutor with KepGroupWorkQueue and support targeted polling.

The changes are behind an experiment unstable_enable_multi_key_bundle and does not affect the default core logic.

…oundedQueueExecutor - Add Uint128Proto and key_group to WorkItem in windmill.proto. - Update Work model to support key groups. - Implement KeyGroupWorkQueue for grouping tasks by key group. KepGroupWorkQueue is a FIFO queue that allows polling elements based on global order and also by order within a key group. - Integrate BoundedQueueExecutor with KepGroupWorkQueue and support targeted polling.

arunpandianp · 2026-06-02T09:06:58Z

R: @scwhittle

github-actions · 2026-06-02T09:08:16Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

gemini-code-assist · 2026-06-02T09:10:25Z

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 a new KeyGroupWorkQueue implementation to support grouping and targeted polling of work items by key group. It updates the Work model and windmill.proto to include key group information and integrates this new queue mechanism into the BoundedQueueExecutor. These changes are gated behind an experimental flag, ensuring no impact on existing core logic.

Highlights

Protocol Buffers Update: Added Uint128Proto and key_group field to WorkItem in windmill.proto.
Work Model Enhancement: Updated the Work model to support and store key group information.
New Queue Implementation: Introduced KeyGroupWorkQueue, a thread-safe BlockingQueue that supports both global FIFO ordering and targeted polling by computation and key group.
Executor Integration: Integrated KeyGroupWorkQueue into BoundedQueueExecutor, gated behind the unstable_enable_multi_key_bundle experiment.

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 support for multi-key bundles in the Streaming Dataflow Worker, guarded by the unstable_enable_multi_key_bundle experiment. It adds a custom, thread-safe KeyGroupWorkQueue to allow polling work by computation and key group in FIFO order, along with corresponding updates to BoundedQueueExecutor, Work, and windmill.proto. Feedback on the new queue implementation suggests adding a defensive check in unlinkNode to prevent double-unlinking issues, and reassigning the result of checkStateNotNull in take() to satisfy static analysis nullness checks.

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.

arunpandianp · 2026-06-05T08:51:53Z

Run Java PreCommit

scwhittle

mostly test stuff remaining

scwhittle · 2026-06-05T12:32:11Z

remove? might be able to remove some null checks below then too

scwhittle · 2026-06-05T12:46:24Z

+                    }
+                    if (task != null) {
+                      consumedCount.incrementAndGet();
+                    }


if task is null, should we perform blocking take the next time?

Added an iteration that uses take.

arunpandianp · 2026-06-07T08:12:41Z

Run Java PreCommit

arunpandianp · 2026-06-07T08:13:24Z

Run Java Dataflow Non-portable Worker PreCommit

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

gemini-code-assist Bot reviewed Jun 2, 2026

View reviewed changes

Comment thread ...ava/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java

Comment thread ...ava/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java Outdated

scwhittle requested changes Jun 4, 2026

View reviewed changes

address comments

5fc07cc

arunpandianp requested a review from scwhittle June 5, 2026 08:20

scwhittle requested changes Jun 5, 2026

View reviewed changes

arunpandianp added 2 commits June 5, 2026 10:25

address comments

b0f8235

Add fairQueue test for KeyGroupWorkQueue

7f7d972

arunpandianp requested a review from scwhittle June 5, 2026 10:34

address comments

4749b08

scwhittle requested changes Jun 5, 2026

View reviewed changes

arunpandianp added 2 commits June 7, 2026 09:22

Address comments

6430094

Address comments

efbac1c

Conversation

arunpandianp commented Jun 2, 2026

Uh oh!

arunpandianp commented Jun 2, 2026

Uh oh!

github-actions Bot commented Jun 2, 2026

Uh oh!

gemini-code-assist Bot commented Jun 2, 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

arunpandianp commented Jun 5, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

scwhittle left a comment

Choose a reason for hiding this comment

Uh oh!

scwhittle Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

scwhittle Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

arunpandianp Jun 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

arunpandianp commented Jun 7, 2026

Uh oh!

arunpandianp commented Jun 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants