From 3abd877772af5959a436d094fcac6d0dc82b5d74 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 10:06:48 -0700 Subject: [PATCH 1/4] Fix WindmillMapViaMultimap.computeIfAbsent return value and add unit tests --- .../state/WindmillMapViaMultimap.java | 17 +- .../state/WindmillMapViaMultimapTest.java | 311 ++++++++++++++++++ 2 files changed, 323 insertions(+), 5 deletions(-) create mode 100644 runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java index e24e449d2598..0c2da6fe17da 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java @@ -62,19 +62,26 @@ public void put(KeyT key, ValueT value) { multimap.put(key, value); } + @SuppressWarnings("nullness") @Override public ReadableState computeIfAbsent( KeyT key, Function mappingFunction) { // Note that computeIfAbsent comments indicate that the read is lazy but this matches the - // existing eager - // behavior of WindmillMap. + // existing eager behavior of WindmillMap. Iterable existingValues = multimap.get(key).read(); - if (Iterables.isEmpty(existingValues)) { + Iterator iterator = existingValues.iterator(); + if (!iterator.hasNext()) { ValueT inserted = mappingFunction.apply(key); multimap.put(key, inserted); - return ReadableStates.immediate(inserted); + return ReadableStates.immediate(null); } else { - return ReadableStates.immediate(Iterables.getOnlyElement(existingValues)); + ValueT existing = Iterables.getOnlyElement(existingValues); + if (existing == null) { + ValueT inserted = mappingFunction.apply(key); + put(key, inserted); + return ReadableStates.immediate(null); + } + return ReadableStates.immediate(existing); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java new file mode 100644 index 000000000000..3ea123c846fe --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java @@ -0,0 +1,311 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.dataflow.worker.windmill.state; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.when; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.AbstractMap; +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaceForTest; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.runners.core.StateTags; +import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; +import org.apache.beam.runners.dataflow.worker.WindmillComputationKey; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.Coder.Context; +import org.apache.beam.sdk.coders.NullableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Futures; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.SettableFuture; +import org.hamcrest.Matchers; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.Timeout; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +/** + * Unit tests demonstrating correctness issues and contract violations in {@link + * WindmillMapViaMultimap}. + */ +@RunWith(JUnit4.class) +public class WindmillMapViaMultimapTest { + @Rule public transient Timeout globalTimeout = Timeout.seconds(600); + + private static final String STATE_FAMILY = "family"; + private static final StateNamespace NAMESPACE = new StateNamespaceForTest("ns"); + + @Mock private WindmillStateReader mockReader; + @Mock private Supplier readStateSupplier; + + private WindmillStateCache cacheViaMultimap; + private WindmillStateInternals underTestMapViaMultimap; + + private static ByteString key(StateNamespace namespace, String addrId) { + return ByteString.copyFromUtf8(namespace.stringKey() + "+u" + addrId); + } + + private static ByteString encodeWithCoder(T value, Coder coder) { + ByteStringOutputStream out = new ByteStringOutputStream(); + try { + coder.encode(value, out, Context.OUTER); + } catch (IOException e) { + throw new RuntimeException(e); + } + return out.toByteString(); + } + + @Before + public void setUp() { + MockitoAnnotations.initMocks(this); + DataflowWorkerHarnessOptions options = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + cacheViaMultimap = + WindmillStateCache.builder() + .setSizeMb(options.getWorkerCacheMb()) + .setSupportMapViaMultimap(true) + .build(); + + WindmillTagEncoding tagEncoding = WindmillTagEncodingV1.instance(); + underTestMapViaMultimap = + new WindmillStateInternals<>( + "key", + STATE_FAMILY, + mockReader, + false, + cacheViaMultimap + .forComputation("comp") + .forKey( + WindmillComputationKey.create( + "comp", ByteString.copyFrom("key", StandardCharsets.UTF_8), 123), + 17L, + 1L) + .forFamily(STATE_FAMILY), + tagEncoding, + readStateSupplier); + + when(mockReader.multimapFetchAllFuture(anyBoolean(), any(), any(), any())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + when(mockReader.multimapFetchSingleEntryFuture(any(), any(), any(), any())) + .thenReturn(Futures.immediateFuture(Collections.emptyList())); + } + + /** + * Tests that {@link WindmillMapViaMultimap#computeIfAbsent} returns {@code null} on insertion, + * conforming to the {@link MapState#computeIfAbsent} specification. + */ + @Test + public void testComputeIfAbsentReturnsNullOnInsert() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + final String key = "tag1"; + SettableFuture> future = SettableFuture.create(); + when(mockReader.multimapFetchSingleEntryFuture( + encodeWithCoder(key, StringUtf8Coder.of()), + key(NAMESPACE, "map"), + STATE_FAMILY, + VarIntCoder.of())) + .thenReturn(future); + future.set(Collections.emptyList()); + + ReadableState result = mapState.computeIfAbsent(key, k -> 42); + Integer readValue = result.read(); + + assertNull( + "computeIfAbsent returns null when key was absent and computed value is inserted", + readValue); + assertEquals(Integer.valueOf(42), mapState.get(key).read()); + } + + /** + * Tests that {@link WindmillMapViaMultimap} throws {@link IllegalArgumentException} if multiple + * values exist for a key, enforcing at-most-one-value map invariants. + */ + @Test(expected = IllegalArgumentException.class) + public void testGetThrowsOnMultiValuedKey() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + final String key = "tag1"; + SettableFuture> future = SettableFuture.create(); + when(mockReader.multimapFetchSingleEntryFuture( + encodeWithCoder(key, StringUtf8Coder.of()), + key(NAMESPACE, "map"), + STATE_FAMILY, + VarIntCoder.of())) + .thenReturn(future); + // Storage returns multiple values for key + future.set(Arrays.asList(10, 20)); + + mapState.get(key).read(); + } + + /** + * Demonstrates Issue 3.3: Every {@link WindmillMapViaMultimap#put} on an unread key executes + * {@code multimap.remove(key)} followed by {@code multimap.put(key, value)}, which marks the key + * as locally removed and forces a redundant {@code DeleteAll = true} update request to Windmill + * on commit. + */ + @Test + public void testPutEmitsDeleteAllAndAddValues() throws Exception { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + // Call put without prior read + mapState.put("key1", 42); + + Windmill.WorkItemCommitRequest.Builder commitBuilder = + Windmill.WorkItemCommitRequest.newBuilder(); + underTestMapViaMultimap.persist(commitBuilder); + + assertEquals(1, commitBuilder.getMultimapUpdatesCount()); + Windmill.TagMultimapUpdateRequest update = commitBuilder.getMultimapUpdates(0); + assertEquals(1, update.getUpdatesCount()); + Windmill.TagMultimapEntry entry = update.getUpdates(0); + + // DEMONSTRATION OF OVERHEAD: entry has DeleteAll = true AND AddValues, because put() called + // remove() first! + assertTrue( + "Demonstrates Issue 3.3: put() without prior read forced deleteAll=true update request", + entry.getDeleteAll()); + assertEquals(1, entry.getValuesCount()); + } + + /** Tests basic put, get, overwrite, and remove on {@link WindmillMapViaMultimap}. */ + @Test + public void testBasicPutGetOverwriteAndRemove() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("a", 1); + assertEquals(Integer.valueOf(1), mapState.get("a").read()); + + mapState.put("a", 2); + assertEquals(Integer.valueOf(2), mapState.get("a").read()); + + mapState.remove("a"); + assertNull(mapState.get("a").read()); + } + + /** Tests keys, values, and entries on {@link WindmillMapViaMultimap}. */ + @Test + public void testKeysValuesAndEntries() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("k1", 10); + mapState.put("k2", 20); + + assertThat(mapState.keys().read(), Matchers.containsInAnyOrder("k1", "k2")); + assertThat(mapState.values().read(), Matchers.containsInAnyOrder(10, 20)); + assertThat( + mapState.entries().read(), + Matchers.containsInAnyOrder( + new AbstractMap.SimpleEntry<>("k1", 10), new AbstractMap.SimpleEntry<>("k2", 20))); + } + + /** Tests clear() and isEmpty() on {@link WindmillMapViaMultimap}. */ + @Test + public void testClearAndIsEmpty() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("key1", 10); + assertFalse(mapState.isEmpty().read()); + + mapState.clear(); + assertTrue(mapState.isEmpty().read()); + } + + /** Tests that computeIfAbsent on an existing key returns the existing value. */ + @Test + public void testComputeIfAbsentExistingKeyReturnsExistingValue() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("key1", 10); + AtomicBoolean called = new AtomicBoolean(false); + ReadableState result = + mapState.computeIfAbsent( + "key1", + k -> { + called.set(true); + return 99; + }); + + assertEquals(Integer.valueOf(10), result.read()); + assertFalse("Mapping function should not be called for existing key", called.get()); + } + + /** + * Tests that computeIfAbsent computes and inserts a new value when the key is currently mapped to + * null. + */ + @Test + public void testComputeIfAbsentWhenMappedToNull() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), NullableCoder.of(StringUtf8Coder.of())); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + mapState.put("key1", null); + AtomicBoolean called = new AtomicBoolean(false); + ReadableState result = + mapState.computeIfAbsent( + "key1", + k -> { + called.set(true); + return "computed"; + }); + + assertNull(result.read()); + assertTrue("Mapping function should be called when key is mapped to null", called.get()); + assertEquals("computed", mapState.get("key1").read()); + } +} From 612960c7a48c3823e30bfc83ed3c3d1949f19ea5 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 10:13:39 -0700 Subject: [PATCH 2/4] Clean up tests. --- .../state/WindmillMapViaMultimapTest.java | 38 ++----------------- 1 file changed, 3 insertions(+), 35 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java index 3ea123c846fe..ca53afa4c2c4 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java @@ -150,11 +150,9 @@ public void testComputeIfAbsentReturnsNullOnInsert() { future.set(Collections.emptyList()); ReadableState result = mapState.computeIfAbsent(key, k -> 42); - Integer readValue = result.read(); - assertNull( "computeIfAbsent returns null when key was absent and computed value is inserted", - readValue); + result.read()); assertEquals(Integer.valueOf(42), mapState.get(key).read()); } @@ -182,38 +180,6 @@ public void testGetThrowsOnMultiValuedKey() { mapState.get(key).read(); } - /** - * Demonstrates Issue 3.3: Every {@link WindmillMapViaMultimap#put} on an unread key executes - * {@code multimap.remove(key)} followed by {@code multimap.put(key, value)}, which marks the key - * as locally removed and forces a redundant {@code DeleteAll = true} update request to Windmill - * on commit. - */ - @Test - public void testPutEmitsDeleteAllAndAddValues() throws Exception { - StateTag> addr = - StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); - MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); - - // Call put without prior read - mapState.put("key1", 42); - - Windmill.WorkItemCommitRequest.Builder commitBuilder = - Windmill.WorkItemCommitRequest.newBuilder(); - underTestMapViaMultimap.persist(commitBuilder); - - assertEquals(1, commitBuilder.getMultimapUpdatesCount()); - Windmill.TagMultimapUpdateRequest update = commitBuilder.getMultimapUpdates(0); - assertEquals(1, update.getUpdatesCount()); - Windmill.TagMultimapEntry entry = update.getUpdates(0); - - // DEMONSTRATION OF OVERHEAD: entry has DeleteAll = true AND AddValues, because put() called - // remove() first! - assertTrue( - "Demonstrates Issue 3.3: put() without prior read forced deleteAll=true update request", - entry.getDeleteAll()); - assertEquals(1, entry.getValuesCount()); - } - /** Tests basic put, get, overwrite, and remove on {@link WindmillMapViaMultimap}. */ @Test public void testBasicPutGetOverwriteAndRemove() { @@ -256,6 +222,8 @@ public void testClearAndIsEmpty() { StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + assertTrue(mapState.isEmpty().read()); + mapState.put("key1", 10); assertFalse(mapState.isEmpty().read()); From 1669ffda20cbf7943be96b2db0bb69b117a10945 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 23 Oct 2026 10:23:33 -0700 Subject: [PATCH 3/4] Fixes extra copy of Iterator --- .../worker/windmill/state/WindmillMapViaMultimap.java | 2 +- .../worker/windmill/state/WindmillMapViaMultimapTest.java | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java index 0c2da6fe17da..795d114bdf56 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java @@ -75,7 +75,7 @@ public ReadableState computeIfAbsent( multimap.put(key, inserted); return ReadableStates.immediate(null); } else { - ValueT existing = Iterables.getOnlyElement(existingValues); + ValueT existing = Iterators.getOnlyElement(iterator); if (existing == null) { ValueT inserted = mappingFunction.apply(key); put(key, inserted); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java index ca53afa4c2c4..d0210ce579da 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java @@ -39,7 +39,6 @@ import org.apache.beam.runners.core.StateTags; import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; import org.apache.beam.runners.dataflow.worker.WindmillComputationKey; -import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.Coder.Context; import org.apache.beam.sdk.coders.NullableCoder; @@ -63,10 +62,7 @@ import org.mockito.Mock; import org.mockito.MockitoAnnotations; -/** - * Unit tests demonstrating correctness issues and contract violations in {@link - * WindmillMapViaMultimap}. - */ +/** Unit tests for {@link WindmillMapViaMultimap}. */ @RunWith(JUnit4.class) public class WindmillMapViaMultimapTest { @Rule public transient Timeout globalTimeout = Timeout.seconds(600); From 557716c7c91e934516603bd2b0859d46d97c913a Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 26 Oct 2026 11:20:14 -0700 Subject: [PATCH 4/4] Update MapState.computeIfAbsent javadoc and WindmillMapViaMultimap to return current/computed value --- .../state/WindmillMapViaMultimap.java | 4 +-- .../state/WindmillMapViaMultimapTest.java | 28 ++++++++++++++----- .../org/apache/beam/sdk/state/MapState.java | 3 +- 3 files changed, 25 insertions(+), 10 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java index 795d114bdf56..66a3f7736eca 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimap.java @@ -73,13 +73,13 @@ public ReadableState computeIfAbsent( if (!iterator.hasNext()) { ValueT inserted = mappingFunction.apply(key); multimap.put(key, inserted); - return ReadableStates.immediate(null); + return ReadableStates.immediate(inserted); } else { ValueT existing = Iterators.getOnlyElement(iterator); if (existing == null) { ValueT inserted = mappingFunction.apply(key); put(key, inserted); - return ReadableStates.immediate(null); + return ReadableStates.immediate(inserted); } return ReadableStates.immediate(existing); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java index d0210ce579da..c629bb102a65 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillMapViaMultimapTest.java @@ -126,11 +126,11 @@ public void setUp() { } /** - * Tests that {@link WindmillMapViaMultimap#computeIfAbsent} returns {@code null} on insertion, - * conforming to the {@link MapState#computeIfAbsent} specification. + * Tests that {@link WindmillMapViaMultimap#computeIfAbsent} returns the computed value on + * insertion, conforming to the {@link MapState#computeIfAbsent} specification. */ @Test - public void testComputeIfAbsentReturnsNullOnInsert() { + public void testComputeIfAbsentReturnsComputedValueOnInsert() { StateTag> addr = StateTags.map("map", StringUtf8Coder.of(), VarIntCoder.of()); MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); @@ -146,8 +146,9 @@ public void testComputeIfAbsentReturnsNullOnInsert() { future.set(Collections.emptyList()); ReadableState result = mapState.computeIfAbsent(key, k -> 42); - assertNull( - "computeIfAbsent returns null when key was absent and computed value is inserted", + assertEquals( + "computeIfAbsent returns computed value when key was absent and computed value is inserted", + Integer.valueOf(42), result.read()); assertEquals(Integer.valueOf(42), mapState.get(key).read()); } @@ -250,7 +251,7 @@ public void testComputeIfAbsentExistingKeyReturnsExistingValue() { /** * Tests that computeIfAbsent computes and inserts a new value when the key is currently mapped to - * null. + * null, returning the computed value. */ @Test public void testComputeIfAbsentWhenMappedToNull() { @@ -268,8 +269,21 @@ public void testComputeIfAbsentWhenMappedToNull() { return "computed"; }); - assertNull(result.read()); + assertEquals("computed", result.read()); assertTrue("Mapping function should be called when key is mapped to null", called.get()); assertEquals("computed", mapState.get("key1").read()); } + + /** Tests that computeIfAbsent returns null if the computed value is null. */ + @Test + public void testComputeIfAbsentReturnsNullWhenComputedValueIsNull() { + StateTag> addr = + StateTags.map("map", StringUtf8Coder.of(), NullableCoder.of(StringUtf8Coder.of())); + MapState mapState = underTestMapViaMultimap.state(NAMESPACE, addr); + + ReadableState result = mapState.computeIfAbsent("key1", k -> null); + + assertNull(result.read()); + assertNull(mapState.get("key1").read()); + } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/MapState.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/MapState.java index 0b4ed36a6c66..4bdf64442801 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/MapState.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/MapState.java @@ -65,7 +65,8 @@ default ReadableState putIfAbsent(K key, V value) { * map and reconciliation with any pending modifications. * *

If the specified key is not already associated with a value (or is mapped to {@code null}) - * associates it with the computed and returns {@code null}, else returns the current value. + * associates it with the computed value. Returns the current (existing or computed) value + * associated with the specified key, or null if the computed value is null. * *

Changes will not be reflected in the results returned by previous calls to {@link * ReadableState#read} on the results any of the reading methods ({@link #get}, {@link #keys},