[acrites]

Changes Beam's MapState.computeIfAbsent contract to be the same as Java's Map.computeIfAbsent, which returns the inserted or computed value. Updates WindmillMapViaMultimap.computeIfAbsent() to call the mapping function when the key exists, but is mapped to null.

Added new unit tests to demonstrate these bugs and added some basic unit tests since none previously existed.

[github-actions]

📋 Review Summary

This PR fixes WindmillMapViaMultimap.computeIfAbsent to conform to the MapState contract by returning null upon inserting a computed value and by treating keys mapped to null as absent. It also adds unit test coverage in WindmillMapViaMultimapTest.

🔍 General Feedback

• The logic in computeIfAbsent is sound and properly replaces existing null entries via put(key, inserted).
• There is an unused import in the new test file that will likely trigger Checkstyle/Spotless failures in CI, as well as a couple of minor cleanups in WindmillMapViaMultimap and WindmillMapViaMultimapTest.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 75.97%. Comparing base (76f6e6d) to head (557716c).
⚠️ Report is 341 commits behind head on master.

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #40750       +/-   ##
=============================================
+ Coverage     56.16%   75.97%   +19.80%     
+ Complexity     2288     2273       -15     
=============================================
  Files          1124      259      -865     
  Lines        178336    14377   -163959     
  Branches       1489     1479       -10     
=============================================
- Hits         100170    10923    -89247     
+ Misses        75645     2763    -72882     
+ Partials       2521      691     -1830     

Flag	Coverage Δ
go	?
java	75.97% <ø> (+0.18%)	⬆️
python	?

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

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

[acrites]

/gemini review

[github-actions]

📋 Review Summary

This PR fixes WindmillMapViaMultimap.computeIfAbsent to align with the MapState contract by returning ReadableStates.immediate(null) upon inserting a computed value and handling keys mapped to null. It also introduces a comprehensive unit test suite in WindmillMapViaMultimapTest covering core map operations and edge cases, with all previous review feedback addressed.

🔍 General Feedback

• The implementation in WindmillMapViaMultimap is clean, avoids redundant iterator creation, and properly overwrites null-mapped keys via put(key, inserted).
• The unit test coverage in WindmillMapViaMultimapTest thoroughly validates the contract fixes and standard MapState operations.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

Assigning reviewers:

R: @damccorm added as fallback since no labels match configuration
Backup reviewer: tvalentyn (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

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

[damccorm]

R: @scwhittle

[github-actions]

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

[scwhittle]

I find it pretty confusing that MapState.computeIfAbsent doesn't match Java Map.computeIfAbsent behavior where it returns the existing or newly computed value. https://docs.oracle.com/javase/8/docs/api/java/util/Map.html#computeIfAbsent-K-java.util.function.Function-

I think it would be better to update the comment and match Map. Especially if the existing implementations are already behaving the other way.

[acrites]

I find it pretty confusing that MapState.computeIfAbsent doesn't match Java Map.computeIfAbsent behavior where it returns the existing or newly computed value. https://docs.oracle.com/javase/8/docs/api/java/util/Map.html#computeIfAbsent-K-java.util.function.Function-

I think it would be better to update the comment and match Map. Especially if the existing implementations are already behaving the other way.

The only other implementation I've found so far is the InMemoryState: https://github.com/apache/beam/blob/master/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryStateInternals.java#L843. It looks like it returns the newly computed value.

So if we're comfortable changing the interface contract, I can do that instead. I think we'll still need the changes in this PR to handle the case where the current key is mapped to null though.

[acrites]

I wonder if the confusion was because Java's Map.putIfAbsent, does return null if the element wasn't already in the map.

default V putIfAbsent(K key, V value)

If the specified key is not already associated with a value (or is mapped to null) associates it with the given value and returns null, else returns the current value.

[scwhittle]

Hmm, there is some difficulty here because null is a valid value for some coders within Beam (https://beam.apache.org/releases/javadoc/2.0.0/org/apache/beam/sdk/coders/NullableCoder.html for example)

So if the iterable exists it seems we should just return the value as it was.

However when evaluating mappingFunction above we need to decide how we interpret null return values:

• are they real values that should be cached (as written here)
• or do they mean to not insert? in which case we shouldn't add the key and value to the map

I think either can be valid interpretation, maybe we should see what the existing implementations do and update the comment to match that behavior.

[acrites]

Some other implementations I found were InMemoryMap:

    public ReadableState<V> computeIfAbsent(
        K key, Function<? super K, ? extends V> mappingFunction) {
      V v = contents.get(key);
      if (v == null) {
        v = contents.put(key, mappingFunction.apply(key));
      }

      return ReadableStates.immediate(v);
    }

Here is the Spark runner's implementation (note it's return value matches the original version of MapState.computeIfAbsent()):

    public ReadableState<MapValueT> computeIfAbsent(
        MapKeyT key, Function<? super MapKeyT, ? extends MapValueT> mappingFunction) {
      Map<MapKeyT, MapValueT> sparkMapState = readAsMap();
      MapValueT current = sparkMapState.get(key);
      if (current == null) {
        put(key, mappingFunction.apply(key));
      }
      return ReadableStates.immediate(current);
    }

It looks like the Flink runner has the same implementation as the Spark runner:

    public ReadableState<ValueT> computeIfAbsent(
        final KeyT key, Function<? super KeyT, ? extends ValueT> mappingFunction) {
      try {
        ValueT current =
            flinkStateBackend
                .getPartitionedState(namespace, namespaceSerializer, flinkStateDescriptor)
                .get(key);

        if (current == null) {
          flinkStateBackend
              .getPartitionedState(namespace, namespaceSerializer, flinkStateDescriptor)
              .put(key, mappingFunction.apply(key));
        }
        return ReadableStates.immediate(current);
      } catch (Exception e) {
        throw new RuntimeException("Error put kv to state.", e);
      }
    }

[acrites]

In all of these cases the behavior seems to be "if the value is currently mapped to null, apply the mapping function and overwrite it with whatever the mapping function returns". I think they all return the previous value that was stored in the map. So it seems like WindmillMapViaMultimap is the only class not implementing computeIfAbsent based on the original comments.