Adds Multimap support to JAVA FnApi by acrites · Pull Request #36218 · apache/beam

acrites · 2025-09-19T21:30:00Z

A MultimapUserState class already existed, but wasn't hooked up to the FnApiStateAccessor. In addition to hooking it up, we also added an API to the FnApi to get all entries in the multimap (which was part of the SDK's state API, but was missing in the FnApi).

Added a new validates runner test for the entries API as well as a MultimapUserStateTest test for the new behavior.

This PR is to help address #23616.

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.

…no guarantee on the order keys are returned. Also fixes a couple warnings from other FnApi tests.

… a missing FnApi state proto to get all of the entries of a multimap. This type of access is part of the state API (and supported by the non-portable harness), but was not present in the protos.

… turned my entire file in IntelliJ yellow.

gemini-code-assist · 2025-09-19T21:30:18Z

Summary of Changes

Hello @acrites, 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 significantly enhances the Java FnApi by introducing comprehensive support for Multimap state. It integrates the previously existing MultimapUserState with the FnApiStateAccessor and provides a new API to access all entries within a multimap. This change is crucial for enabling more advanced stateful operations in Beam pipelines and brings the FnApi's state capabilities closer to the full functionality offered by the Beam SDK.

Highlights

Multimap State Support: The Java FnApi now fully supports Multimap state, enabling more complex stateful processing patterns.
New API for Multimap Entries: An API has been added to the FnApi to retrieve all entries (key-value pairs) from a Multimap, aligning with the SDK's existing state API.
Integration of MultimapUserState: The existing MultimapUserState class has been successfully hooked up to the FnApiStateAccessor, making it usable within the FnApi framework.
Comprehensive Testing: New validates runner tests and unit tests (MultimapUserStateTest) have been introduced to ensure the correctness and robustness of the new Multimap entries API and overall behavior.
Issue Resolution: This pull request directly addresses and helps resolve issue [Task]: Add MultimapState support for Dataflow FnAPI #23616, which concerned the lack of Multimap support in the Fn API.

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

Customization

To customize 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 counter productive. 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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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

github-actions · 2025-09-19T23:06:32Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

acrites · 2025-09-21T17:23:54Z

The failing PreCommits don't seem related to this CL.

acrites · 2025-09-21T17:23:58Z

assign set of reviewers

github-actions · 2025-09-21T17:25:00Z

Assigning reviewers:

R: @Abacn for label java.

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

@nonnull

Removes non-needed @nonnull annotations.

scwhittle · 2025-09-25T08:19:11Z

+        return new PrefetchableIterator<Map.Entry<K, V>>() {
+          PrefetchableIterator<Map.Entry<K, V>> persistedEntriesIterator =
+              PrefetchableIterables.concat(
+                      Iterables.concat(


Would
PrefetchableIterator<Map.Entry<K, V>> persistedEntriesIterator =
PrefetchableIterables.maybePrefetchable(Iterables.transform(
persistedEntries,
entry ->
Iterables.transform(
entry.getValue(),
value -> Maps.immutableEntry(entry.getKey(), value)))

work? seems like single iterable so don't need the two concats. If I'm missing something it at least seems like you could just have PrefetchableIterables.concat and not the other one.

It seems to have something to do with PrefetchableIterator though. If I replace everything with just Iterators, it works fine.

I've tried to get rid of the inner concat, but then I get errors about trying to pass a PrefetchableIterator<Iterator<Map.Entry<K,V>>> to the outermost concat. I was able to get it to work by making maybePrefetchable public (it's currently private), but I still need a single concat of the combined transforms.

The WindmillMultimap implementation also needed a concat:

beam/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMultimap.java

Line 646 in bbf3613

Iterator<Map.Entry<K, V>> unnestWindmill =

Im a little worried we're ending up with something that is not actually plumbing prefetch down to the entries. It seems like we need PrefetchableIterables.transform as otherwise Iterables.transform is going to wrap the PrefetchableIterator and PrefetchableIterators.maybePrefetch wrapper is going to just make the prefetch a no-op.

I added a test similar to the other prefetch tests. It looks like a call to entries().prefetch() produces a data request across the FnApi.

scwhittle · 2025-09-25T08:25:14Z

+                                      value -> Maps.immutableEntry(entry.getKey(), value)))))
+                  .iterator();
+          Iterator<Map.Entry<K, V>> pendingAddsNowIterator;
+          boolean hasNext;


hasNextReady? it's confusing as is since hasNext=false doesn't mean it doesn't have next, just that we haven't calculated it.

or perhaps you can remove this variable and just mark nextEntry nullable and comment that if it is non-null it is calculated next element to return.

Sounds good. Got rid of hasNext and reset nextEntry to null when calling next().

scwhittle · 2025-09-25T08:29:47Z

+              }
+            }
+
+            if (pendingAddsNowIterator == null) {


do we need to merge the persisted and local adds? Does the iterator have any guarantees on ordering of elements for a key?

I changed this up so it merges in the local adds with backend values. Still needed two concats for some reason though. Otherwise, it complains that it's getting PrefetchableIterator<Iterable<Entry<K, V>>> instead of a PrefetchableIterator<Entry<K,V>>.

I'm not really sure if the prefetching propagates through all these Iterables.transform() calls though.

…om the backend and local adds together. Also needed to make maybePrefetchable public.

…into multimap-streaming

scwhittle · 2025-10-10T08:38:31Z

+        return new PrefetchableIterator<Map.Entry<K, V>>() {
+          PrefetchableIterator<Map.Entry<K, V>> persistedEntriesIterator =
+              PrefetchableIterables.concat(
+                      Iterables.concat(


Im a little worried we're ending up with something that is not actually plumbing prefetch down to the entries. It seems like we need PrefetchableIterables.transform as otherwise Iterables.transform is going to wrap the PrefetchableIterator and PrefetchableIterators.maybePrefetch wrapper is going to just make the prefetch a no-op.

github-actions · 2025-10-17T12:15:31Z

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

…st sent across FnApi.

github-actions · 2025-10-23T12:15:52Z

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: @ahmedabu98 for label java.

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)

scwhittle

Implementation looks good, just a question about capability protecting it

scwhittle · 2025-10-24T08:46:50Z

+  // The response data stream will be a concatenation of pairs, where the first
+  // component is the map key and the second component is a concatenation of
+  // values associated with that map key.
+  message MultimapEntriesUserState {


maybe we should add a capability for Multimap support similar to ordered list state?

Good idea! I added a capability to beam_runner_api.proto and specified in for the java sdk. I don't think python or go support it yet.

… java sdk.

codecov · 2025-10-24T21:22:23Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 36.12%. Comparing base (d26dbac) to head (555638a).
⚠️ Report is 324 commits behind head on master.

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #36218       +/-   ##
=============================================
- Coverage     54.93%   36.12%   -18.81%     
  Complexity     1617     1617               
=============================================
  Files          1057     1058        +1     
  Lines        164425   165037      +612     
  Branches       1165     1165               
=============================================
- Hits          90320    59613    -30707     
- Misses        71955   103274    +31319     
  Partials       2150     2150

Flag	Coverage Δ
python	`40.51% <ø> (-40.40%)`	⬇️

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

☔ View full report in Codecov by Sentry.
📢 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.

scwhittle · 2025-10-25T08:37:58Z

Run Java_Amazon-Web-Services2_IO_Direct PreCommit

scwhittle · 2025-10-25T08:38:29Z

Run Java PreCommit

Abacn · 2025-10-27T18:29:43Z

I understand it's ready to be merged?

acrites added 5 commits September 16, 2025 21:51

Changes multimap state key() tests to not care about order. There is …

c692d60

…no guarantee on the order keys are returned. Also fixes a couple warnings from other FnApi tests.

Adds Multimap user state support to the Java FnApi harness. Also adds…

70e2106

… a missing FnApi state proto to get all of the entries of a multimap. This type of access is part of the state API (and supported by the non-portable harness), but was not present in the protos.

Adds FnApi binding for entries() method.

9708b12

Removes extra unchecked warning supression. I had added it because it…

19468cb

… turned my entire file in IntelliJ yellow.

Removes extra comma.

e0c4b67

github-actions Bot added java model labels Sep 19, 2025

Merge branch 'master' into multimap-streaming

b0480fa

github-actions Bot added the Next Action: Reviewers label Sep 21, 2025

Abacn reviewed Sep 22, 2025

View reviewed changes

Comment thread sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java Outdated

acrites and others added 3 commits September 22, 2025 20:43

Removes not-needed @nonnull annotations.

442efba

Merge of upstream

5083d61

Update FnApiStateAccessor.java

88c34f1

Removes non-needed @nonnull annotations.

scwhittle requested changes Sep 25, 2025

View reviewed changes

acrites added 2 commits October 1, 2025 23:22

Changes multimap entries() iterable to put values for the same key fr…

1ba5282

…om the backend and local adds together. Also needed to make maybePrefetchable public.

Merge branch 'multimap-streaming' of https://github.com/acrites/beam …

45d3601

…into multimap-streaming

acrites requested review from Abacn and scwhittle October 6, 2025 21:42

scwhittle requested changes Oct 10, 2025

View reviewed changes

github-actions Bot added the slow-review label Oct 17, 2025

Adds a test that prefetching multimap entries results in a StateReque…

cf3c8f4

…st sent across FnApi.

github-actions Bot added reassigned-reviewers and removed slow-review labels Oct 23, 2025

scwhittle requested changes Oct 24, 2025

View reviewed changes

Adds an environment capability for multimap state and sets in for the…

555638a

… java sdk.

scwhittle approved these changes Oct 25, 2025

View reviewed changes

scwhittle merged commit 21dbf59 into apache:master Oct 28, 2025
143 of 146 checks passed

Conversation

acrites commented Sep 19, 2025

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot commented Sep 19, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions Bot commented Sep 19, 2025

Uh oh!

acrites commented Sep 21, 2025

Uh oh!

acrites commented Sep 21, 2025

Uh oh!

github-actions Bot commented Sep 21, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 17, 2025

Uh oh!

github-actions Bot commented Oct 23, 2025

Uh oh!

scwhittle left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Oct 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

scwhittle commented Oct 25, 2025

Uh oh!

scwhittle commented Oct 25, 2025

Uh oh!

Abacn commented Oct 27, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Oct 24, 2025 •

edited

Loading