From c692d60691b3e685e03439609f3a4d2a136b362d Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Tue, 16 Sep 2025 21:51:29 +0000 Subject: [PATCH 01/10] Changes multimap state key() tests to not care about order. There is no guarantee on the order keys are returned. Also fixes a couple warnings from other FnApi tests. --- ...nessMonitoringInfosInstructionHandlerTest.java | 3 +++ .../harness/control/ProcessBundleHandlerTest.java | 1 + .../fn/harness/state/MultimapUserStateTest.java | 15 ++++++++++----- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/HarnessMonitoringInfosInstructionHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/HarnessMonitoringInfosInstructionHandlerTest.java index ac69ed29a565..9e69cb2ec700 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/HarnessMonitoringInfosInstructionHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/HarnessMonitoringInfosInstructionHandlerTest.java @@ -30,7 +30,10 @@ import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.MetricsEnvironment; import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +@RunWith(JUnit4.class) public class HarnessMonitoringInfosInstructionHandlerTest { @Test diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index 8a35351fdb25..a7a62571e38e 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -236,6 +236,7 @@ public void finishBundle(FinishBundleContext context) { } } + @SuppressWarnings("ExtendsAutoValue") private static class TestBundleProcessor extends BundleProcessor { static int resetCnt = 0; diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java index 17550793a8b2..48c9ce43bdf0 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java @@ -21,6 +21,7 @@ import static java.util.Collections.singletonList; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.emptyIterable; +import static org.hamcrest.collection.ArrayMatching.arrayContainingInAnyOrder; import static org.hamcrest.core.Is.is; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -167,7 +168,9 @@ public void testKeys() throws Exception { userState.put(A3, "V1"); userState.put(A1, "V3"); assertArrayEquals(new byte[][] {A1, A2}, Iterables.toArray(initKeys, byte[].class)); - assertArrayEquals(new byte[][] {A1, A2, A3}, Iterables.toArray(userState.keys(), byte[].class)); + assertThat( + Iterables.toArray(userState.keys(), byte[].class), + is(arrayContainingInAnyOrder(A1, A2, A3))); userState.clear(); assertArrayEquals(new byte[][] {A1, A2}, Iterables.toArray(initKeys, byte[].class)); @@ -822,8 +825,9 @@ public void testKeysCached() throws Exception { userState.put(A2, "V1"); userState.put(A3, "V1"); - assertArrayEquals( - new byte[][] {A1, A2, A3}, Iterables.toArray(userState.keys(), byte[].class)); + assertThat( + Iterables.toArray(userState.keys(), byte[].class), + is(arrayContainingInAnyOrder(A1, A2, A3))); userState.asyncClose(); } @@ -841,8 +845,9 @@ public void testKeysCached() throws Exception { ByteArrayCoder.of(), StringUtf8Coder.of()); - assertArrayEquals( - new byte[][] {A1, A2, A3}, Iterables.toArray(userState.keys(), byte[].class)); + assertThat( + Iterables.toArray(userState.keys(), byte[].class), + is(arrayContainingInAnyOrder(A1, A2, A3))); userState.asyncClose(); } } From 70e21062ef437c5a26e1cdf37a864272e09e7fd7 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 19 Sep 2025 21:07:31 +0000 Subject: [PATCH 02/10] Adds Multimap user state support to the Java FnApi harness. Also adds 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. --- .../model/fn_execution/v1/beam_fn_api.proto | 24 +++ .../apache/beam/sdk/transforms/ParDoTest.java | 66 ++++++++ .../fn/harness/state/FnApiStateAccessor.java | 159 +++++++++++++++--- .../fn/harness/state/MultimapUserState.java | 141 +++++++++++++++- .../harness/state/MultimapUserStateTest.java | 76 +++++++++ 5 files changed, 443 insertions(+), 23 deletions(-) diff --git a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto index 9360522ab409..fde3657b667d 100644 --- a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto +++ b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto @@ -1009,6 +1009,29 @@ message StateKey { bytes key = 4; } + // Represents a request for all of the entries of a multimap associated with a + // specified user key and window for a PTransform. See + // https://s.apache.org/beam-fn-state-api-and-bundle-processing for further + // details. + // + // Can only be used to perform StateGetRequests and StateClearRequests on the + // user state. + // + // 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 { + // (Required) The id of the PTransform containing user state. + string transform_id = 1; + // (Required) The id of the user state. + string user_state_id = 2; + // (Required) The window encoded in a nested context. + bytes window = 3; + // (Required) The key of the currently executing element encoded in a + // nested context. + bytes key = 4; + } + // Represents a request for the values of the map key associated with a // specified user key and window for a PTransform. See // https://s.apache.org/beam-fn-state-api-and-bundle-processing for further @@ -1064,6 +1087,7 @@ message StateKey { MultimapKeysSideInput multimap_keys_side_input = 5; MultimapKeysValuesSideInput multimap_keys_values_side_input = 8; MultimapKeysUserState multimap_keys_user_state = 6; + MultimapEntriesUserState multimap_entries_user_state = 10; MultimapUserState multimap_user_state = 7; OrderedListUserState ordered_list_user_state = 9; } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java index 8409133772eb..e3464d98526c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java @@ -2917,6 +2917,72 @@ public void processElement( pipeline.run(); } + @Test + @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMultimapState.class}) + public void testMultimapStateEntries() { + final String stateId = "foo:"; + final String countStateId = "count"; + DoFn>, KV> fn = + new DoFn>, KV>() { + + @StateId(stateId) + private final StateSpec> multimapState = + StateSpecs.multimap(StringUtf8Coder.of(), VarIntCoder.of()); + + @StateId(countStateId) + private final StateSpec> countState = + StateSpecs.combiningFromInputInternal(VarIntCoder.of(), Sum.ofIntegers()); + + @ProcessElement + public void processElement( + ProcessContext c, + @Element KV> element, + @StateId(stateId) MultimapState state, + @StateId(countStateId) CombiningState count, + OutputReceiver> r) { + // Empty before we process any elements. + if (count.read() == 0) { + assertThat(state.entries().read(), emptyIterable()); + } + assertEquals(count.read().intValue(), Iterables.size(state.entries().read())); + + KV value = element.getValue(); + state.put(value.getKey(), value.getValue()); + count.add(1); + + if (count.read() >= 4) { + // Those should be evaluated only when ReadableState.read is called. + ReadableState>> entriesView = state.entries(); + ReadableState containsBadView = state.containsKey("BadKey"); + + // Those are evaluated immediately. + Iterable> entries = state.entries().read(); + + state.remove("b"); + assertEquals(4, Iterables.size(entries)); + state.put("a", 2); + + // Note we output the view of state before the modifications in this if statement. + for (Entry entry : entries) { + r.output(KV.of(entry.getKey(), entry.getValue())); + } + } + } + }; + PCollection> output = + pipeline + .apply( + Create.of( + KV.of("hello", KV.of("a", 97)), KV.of("hello", KV.of("a", 97)), + KV.of("hello", KV.of("a", 98)), KV.of("hello", KV.of("b", 33)))) + .apply(ParDo.of(fn)); + PAssert.that(output) + .containsInAnyOrder( + KV.of("a", 97), KV.of("a", 97), + KV.of("a", 98), KV.of("b", 33)); + pipeline.run(); + } + @Test @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMultimapState.class}) public void testMultimapStateRemove() { diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java index e06a82c8e25f..c379153c0a00 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java @@ -81,13 +81,15 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; +import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** Provides access to side inputs and state via a {@link BeamFnStateClient}. */ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness" // TODO(https://github.com/apache/beam/issues/20497) + "nullness", // TODO(https://github.com/apache/beam/issues/20497) + "unchecked" }) public class FnApiStateAccessor implements SideInputReader, StateBinder { @@ -117,7 +119,7 @@ public static class Factory { public Factory( PipelineOptions pipelineOptions, - Set runnerCapabilites, + Set runnerCapabilities, String ptransformId, Supplier processBundleInstructionId, Supplier> cacheTokens, @@ -128,7 +130,7 @@ public Factory( Coder keyCoder, Coder windowCoder) { this.pipelineOptions = pipelineOptions; - this.runnerCapabilities = runnerCapabilites; + this.runnerCapabilities = runnerCapabilities; this.ptransformId = ptransformId; this.processBundleInstructionId = processBundleInstructionId; this.cacheTokens = cacheTokens; @@ -240,7 +242,7 @@ public FnApiStateAccessor create() { } private final PipelineOptions pipelineOptions; - private final Set runnerCapabilites; + private final Set runnerCapabilities; private final Map stateKeyObjectCache; private final Map, SideInputSpec> sideInputSpecMap; private final BeamFnStateClient beamFnStateClient; @@ -259,7 +261,7 @@ public FnApiStateAccessor create() { public FnApiStateAccessor( PipelineOptions pipelineOptions, - Set runnerCapabilites, + Set runnerCapabilities, String ptransformId, Supplier processBundleInstructionId, Supplier> cacheTokens, @@ -270,7 +272,7 @@ public FnApiStateAccessor( Coder keyCoder, Coder windowCoder) { this.pipelineOptions = pipelineOptions; - this.runnerCapabilites = runnerCapabilites; + this.runnerCapabilities = runnerCapabilities; this.stateKeyObjectCache = Maps.newHashMap(); this.sideInputSpecMap = sideInputSpecMap; this.beamFnStateClient = beamFnStateClient; @@ -414,7 +416,7 @@ public T get(PCollectionView view, BoundedWindow window) { key, ((KvCoder) sideInputSpec.getCoder()).getKeyCoder(), ((KvCoder) sideInputSpec.getCoder()).getValueCoder(), - runnerCapabilites.contains( + runnerCapabilities.contains( BeamUrns.getUrn( RunnerApi.StandardRunnerProtocols.Enum .MULTIMAP_KEYS_VALUES_SIDE_INPUT)))); @@ -620,11 +622,12 @@ public SetState readLater() { } @Override + @NonNull public MapState bindMap( - String id, - StateSpec> spec, - Coder mapKeyCoder, - Coder mapValueCoder) { + @NonNull String id, + @NonNull StateSpec<@NonNull MapState> spec, + @NonNull Coder mapKeyCoder, + @NonNull Coder mapValueCoder) { return (MapState) stateKeyObjectCache.computeIfAbsent( createMultimapKeysUserStateKey(id), @@ -647,8 +650,8 @@ public void put(KeyT key, ValueT value) { } @Override - public ReadableState computeIfAbsent( - KeyT key, Function mappingFunction) { + public @NonNull ReadableState computeIfAbsent( + KeyT key, @NonNull Function mappingFunction) { Iterable values = impl.get(key); if (Iterables.isEmpty(values)) { impl.put(key, mappingFunction.apply(key)); @@ -662,11 +665,13 @@ public void remove(KeyT key) { } @Override + @NonNull public ReadableState get(KeyT key) { return getOrDefault(key, null); } @Override + @NonNull public ReadableState getOrDefault( KeyT key, @Nullable ValueT defaultValue) { return new ReadableState() { @@ -677,6 +682,7 @@ public ReadableState getOrDefault( } @Override + @NonNull public ReadableState readLater() { impl.get(key).prefetch(); return this; @@ -685,6 +691,7 @@ public ReadableState readLater() { } @Override + @NonNull public ReadableState> keys() { return new ReadableState>() { @Override @@ -693,6 +700,7 @@ public Iterable read() { } @Override + @NonNull public ReadableState> readLater() { impl.keys().prefetch(); return this; @@ -701,6 +709,7 @@ public ReadableState> readLater() { } @Override + @NonNull public ReadableState> values() { return new ReadableState>() { @Override @@ -709,6 +718,7 @@ public Iterable read() { } @Override + @NonNull public ReadableState> readLater() { entries().readLater(); return this; @@ -717,6 +727,7 @@ public ReadableState> readLater() { } @Override + @NonNull public ReadableState>> entries() { return new ReadableState>>() { @Override @@ -727,6 +738,7 @@ public Iterable> read() { } @Override + @NonNull public ReadableState>> readLater() { // Start prefetching the keys. We would need to block to start prefetching // the values. @@ -737,6 +749,7 @@ public ReadableState>> readLater() { } @Override + @NonNull public ReadableState isEmpty() { return new ReadableState() { @Override @@ -745,6 +758,7 @@ public Boolean read() { } @Override + @NonNull public ReadableState readLater() { keys().readLater(); return this; @@ -757,13 +771,118 @@ public ReadableState readLater() { } @Override - public MultimapState bindMultimap( - String id, - StateSpec> spec, - Coder keyCoder, - Coder valueCoder) { - // TODO(https://github.com/apache/beam/issues/23616) - throw new UnsupportedOperationException("Multimap is not currently supported with Fn API."); + public @NonNull MultimapState bindMultimap( + @NonNull String id, + @NonNull StateSpec<@NonNull MultimapState> spec, + @NonNull Coder keyCoder, + @NonNull Coder valueCoder) { + return (MultimapState) + stateKeyObjectCache.computeIfAbsent( + createMultimapKeysUserStateKey(id), + new Function() { + @Override + public Object apply(StateKey stateKey) { + return new MultimapState() { + private final MultimapUserState impl = + createMultimapUserState(stateKey, keyCoder, valueCoder); + + @Override + public void put(KeyT key, ValueT value) { + impl.put(key, value); + } + + @Override + @NonNull + public ReadableState> get(KeyT key) { + return new ReadableState>() { + @Override + public Iterable read() { + return impl.get(key); + } + + @Override + @NonNull + public ReadableState> readLater() { + impl.get(key).prefetch(); + return this; + } + }; + } + + @Override + public void remove(KeyT key) { + impl.remove(key); + } + + @Override + @NonNull + public ReadableState> keys() { + return new ReadableState>() { + @Override + public Iterable read() { + return impl.keys(); + } + + @Override + @NonNull + public ReadableState> readLater() { + impl.keys().prefetch(); + return this; + } + }; + } + + @Override + @NonNull + public ReadableState>> entries() { + // TODO(https://github.com/apache/beam/issues/23616) + throw new UnsupportedOperationException( + "Multimap entries() is not currently supported with Fn API."); + } + + @Override + @NonNull + public ReadableState containsKey(KeyT key) { + return new ReadableState() { + @Override + public Boolean read() { + return !Iterables.isEmpty(impl.get(key)); + } + + @Override + @NonNull + public ReadableState readLater() { + impl.get(key).prefetch(); + return this; + } + }; + } + + @Override + @NonNull + public ReadableState isEmpty() { + return new ReadableState() { + @Override + public Boolean read() { + return Iterables.isEmpty(impl.keys()); + } + + @Override + @NonNull + public ReadableState readLater() { + impl.keys().prefetch(); + return this; + } + }; + } + + @Override + public void clear() { + impl.clear(); + } + }; + } + }); } @Override diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java index 617faba87cc0..967cb76af1f3 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java @@ -38,12 +38,15 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.fn.stream.PrefetchableIterable; import org.apache.beam.sdk.fn.stream.PrefetchableIterables; import org.apache.beam.sdk.fn.stream.PrefetchableIterator; import org.apache.beam.sdk.util.ByteStringOutputStream; import org.apache.beam.sdk.values.KV; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; /** @@ -52,9 +55,6 @@ * *

Calling {@link #asyncClose()} schedules any required persistence changes. This object should * no longer be used after it is closed. - * - *

TODO: Move to an async persist model where persistence is signalled based upon cache memory - * pressure and its need to flush. */ public class MultimapUserState { @@ -63,8 +63,10 @@ public class MultimapUserState { private final Coder mapKeyCoder; private final Coder valueCoder; private final StateRequest keysStateRequest; + private final StateRequest entriesStateRequest; private final StateRequest userStateRequest; private final CachingStateIterable persistedKeys; + private final CachingStateIterable>> persistedEntries; private boolean isClosed; private boolean isCleared; @@ -106,6 +108,23 @@ public MultimapUserState( .setWindow(stateKey.getMultimapKeysUserState().getWindow()) .setKey(stateKey.getMultimapKeysUserState().getKey()); this.userStateRequest = userStateRequestBuilder.build(); + + StateRequest.Builder entriesStateRequestBuilder = StateRequest.newBuilder(); + entriesStateRequestBuilder + .setInstructionId(instructionId) + .getStateKeyBuilder() + .getMultimapEntriesUserStateBuilder() + .setTransformId(stateKey.getMultimapKeysUserState().getTransformId()) + .setUserStateId(stateKey.getMultimapKeysUserState().getUserStateId()) + .setWindow(stateKey.getMultimapKeysUserState().getWindow()) + .setKey(stateKey.getMultimapKeysUserState().getKey()); + this.entriesStateRequest = entriesStateRequestBuilder.build(); + this.persistedEntries = + StateFetchingIterators.readAllAndDecodeStartingFrom( + Caches.subCache(this.cache, "AllEntries"), + beamFnStateClient, + entriesStateRequest, + KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder))); } public void clear() { @@ -235,6 +254,122 @@ public K next() { }; } + @SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/21068) + }) + /* + * Returns an Iterable containing all entries in this multimap. + */ + public PrefetchableIterable> entries() { + checkState( + !isClosed, + "Multimap user state is no longer usable because it is closed for %s", + keysStateRequest.getStateKey()); + if (isCleared) { + Map>> pendingAddsNow = new HashMap<>(pendingAdds); + return PrefetchableIterables.concat( + Iterables.concat( + Iterables.transform( + pendingAddsNow.entrySet(), + entry -> + Iterables.transform( + entry.getValue().getValue(), + value -> Maps.immutableEntry(entry.getValue().getKey(), value))))); + } + + Set pendingRemovesNow = new HashSet<>(pendingRemoves.keySet()); + // Make a deep copy of pendingAdds so this iterator represents a snapshot of state at the time + // it was created. + Map> pendingAddsNow = new HashMap<>(); + for (Map.Entry>> entry : pendingAdds.entrySet()) { + pendingAddsNow.put( + entry.getValue().getKey(), new ArrayList<>()); // entry.getValue().getValue()); + for (V value : entry.getValue().getValue()) { + List values = pendingAddsNow.get(entry.getValue().getKey()); + values.add(value); + } + } + return new PrefetchableIterables.Default>() { + @Override + public PrefetchableIterator> createIterator() { + return new PrefetchableIterator>() { + PrefetchableIterator> persistedEntriesIterator = + PrefetchableIterables.concat( + Iterables.concat( + Iterables.transform( + persistedEntries, + entry -> + Iterables.transform( + entry.getValue(), + value -> Maps.immutableEntry(entry.getKey(), value))))) + .iterator(); + Iterator> pendingAddsNowIterator; + boolean hasNext; + Map.Entry nextEntry; + + @Override + public boolean isReady() { + return persistedEntriesIterator.isReady(); + } + + @Override + public void prefetch() { + if (!isReady()) { + persistedEntriesIterator.prefetch(); + } + } + + @Override + public boolean hasNext() { + if (hasNext) { + return true; + } + + while (persistedEntriesIterator.hasNext()) { + Map.Entry nextPersistedEntry = persistedEntriesIterator.next(); + Object nextKeyStructuralValue = + mapKeyCoder.structuralValue(nextPersistedEntry.getKey()); + if (!pendingRemovesNow.contains(nextKeyStructuralValue)) { + nextEntry = + Maps.immutableEntry(nextPersistedEntry.getKey(), nextPersistedEntry.getValue()); + hasNext = true; + return true; + } + } + + if (pendingAddsNowIterator == null) { + pendingAddsNowIterator = + Iterables.concat( + Iterables.transform( + pendingAddsNow.entrySet(), + entry -> + Iterables.transform( + entry.getValue(), + value -> Maps.immutableEntry(entry.getKey(), value)))) + .iterator(); + } + while (pendingAddsNowIterator.hasNext()) { + nextEntry = pendingAddsNowIterator.next(); + hasNext = true; + return true; + } + + return false; + } + + @Override + public Map.Entry next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + hasNext = false; + return nextEntry; + } + }; + } + }; + } + /* * Store a key-value pair in the multimap. * Allows duplicate key-value pairs. diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java index 48c9ce43bdf0..457362815cf6 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java @@ -22,6 +22,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.emptyIterable; import static org.hamcrest.collection.ArrayMatching.arrayContainingInAnyOrder; +import static org.hamcrest.collection.IsIterableContainingInAnyOrder.containsInAnyOrder; import static org.hamcrest.core.Is.is; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -34,11 +35,15 @@ import java.util.Collections; import java.util.Iterator; import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.NullableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.fn.stream.PrefetchableIterable; @@ -179,6 +184,66 @@ public void testKeys() throws Exception { assertThrows(IllegalStateException.class, () -> userState.keys()); } + @Test + public void testEntries() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapEntriesStateKey(), + KV.of( + KvCoder.of(ByteArrayCoder.of(), IterableCoder.of(StringUtf8Coder.of())), + asList(KV.of(A1, asList("V1", "V2")), KV.of(A2, asList("V3")))))); + MultimapUserState userState = + new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createMultimapKeyStateKey(), + ByteArrayCoder.of(), + StringUtf8Coder.of()); + + assertArrayEquals(A1, userState.entries().iterator().next().getKey()); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A1), "V1"), + KV.of(ByteString.copyFrom(A1), "V2"), + KV.of(ByteString.copyFrom(A2), "V3"))); + + userState.put(A1, "V4"); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A1), "V1"), + KV.of(ByteString.copyFrom(A1), "V2"), + KV.of(ByteString.copyFrom(A2), "V3"), + KV.of(ByteString.copyFrom(A1), "V4"))); + + userState.remove(A1); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder(KV.of(ByteString.copyFrom(A2), "V3"))); + + userState.put(A1, "V5"); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A2), "V3"), KV.of(ByteString.copyFrom(A1), "V5"))); + + userState.clear(); + assertThat(userState.entries(), emptyIterable()); + userState.asyncClose(); + assertThrows(IllegalStateException.class, () -> userState.entries()); + } + @Test public void testPut() throws Exception { FakeBeamFnStateClient fakeClient = @@ -1053,6 +1118,17 @@ private StateKey createMultimapKeyStateKey() throws IOException { .build(); } + private StateKey createMultimapEntriesStateKey() throws IOException { + return StateKey.newBuilder() + .setMultimapEntriesUserState( + StateKey.MultimapEntriesUserState.newBuilder() + .setWindow(encode(encodedWindow)) + .setKey(encode(encodedKey)) + .setTransformId(pTransformId) + .setUserStateId(stateId)) + .build(); + } + private StateKey createMultimapValueStateKey(byte[] key) throws IOException { return StateKey.newBuilder() .setMultimapUserState( From 9708b1248fdd494bcd9257d8b85184db89e3cbd6 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 19 Sep 2025 21:23:08 +0000 Subject: [PATCH 03/10] Adds FnApi binding for entries() method. --- .../fn/harness/state/FnApiStateAccessor.java | 78 ++++++++----------- 1 file changed, 33 insertions(+), 45 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java index c379153c0a00..563da1c9ebf9 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java @@ -622,8 +622,7 @@ public SetState readLater() { } @Override - @NonNull - public MapState bindMap( + public @NonNull MapState bindMap( @NonNull String id, @NonNull StateSpec<@NonNull MapState> spec, @NonNull Coder mapKeyCoder, @@ -665,14 +664,12 @@ public void remove(KeyT key) { } @Override - @NonNull - public ReadableState get(KeyT key) { + public @NonNull ReadableState get(KeyT key) { return getOrDefault(key, null); } @Override - @NonNull - public ReadableState getOrDefault( + public @NonNull ReadableState getOrDefault( KeyT key, @Nullable ValueT defaultValue) { return new ReadableState() { @Override @@ -682,8 +679,7 @@ public ReadableState getOrDefault( } @Override - @NonNull - public ReadableState readLater() { + public @NonNull ReadableState readLater() { impl.get(key).prefetch(); return this; } @@ -691,8 +687,7 @@ public ReadableState readLater() { } @Override - @NonNull - public ReadableState> keys() { + public @NonNull ReadableState> keys() { return new ReadableState>() { @Override public Iterable read() { @@ -700,8 +695,7 @@ public Iterable read() { } @Override - @NonNull - public ReadableState> readLater() { + public @NonNull ReadableState> readLater() { impl.keys().prefetch(); return this; } @@ -709,8 +703,7 @@ public ReadableState> readLater() { } @Override - @NonNull - public ReadableState> values() { + public @NonNull ReadableState> values() { return new ReadableState>() { @Override public Iterable read() { @@ -718,8 +711,7 @@ public Iterable read() { } @Override - @NonNull - public ReadableState> readLater() { + public @NonNull ReadableState> readLater() { entries().readLater(); return this; } @@ -727,8 +719,7 @@ public ReadableState> readLater() { } @Override - @NonNull - public ReadableState>> entries() { + public @NonNull ReadableState>> entries() { return new ReadableState>>() { @Override public Iterable> read() { @@ -738,8 +729,7 @@ public Iterable> read() { } @Override - @NonNull - public ReadableState>> readLater() { + public @NonNull ReadableState>> readLater() { // Start prefetching the keys. We would need to block to start prefetching // the values. keys().readLater(); @@ -749,8 +739,7 @@ public ReadableState>> readLater() { } @Override - @NonNull - public ReadableState isEmpty() { + public @NonNull ReadableState isEmpty() { return new ReadableState() { @Override public Boolean read() { @@ -758,8 +747,7 @@ public Boolean read() { } @Override - @NonNull - public ReadableState readLater() { + public @NonNull ReadableState readLater() { keys().readLater(); return this; } @@ -792,8 +780,7 @@ public void put(KeyT key, ValueT value) { } @Override - @NonNull - public ReadableState> get(KeyT key) { + public @NonNull ReadableState> get(KeyT key) { return new ReadableState>() { @Override public Iterable read() { @@ -801,8 +788,7 @@ public Iterable read() { } @Override - @NonNull - public ReadableState> readLater() { + public @NonNull ReadableState> readLater() { impl.get(key).prefetch(); return this; } @@ -815,8 +801,7 @@ public void remove(KeyT key) { } @Override - @NonNull - public ReadableState> keys() { + public @NonNull ReadableState> keys() { return new ReadableState>() { @Override public Iterable read() { @@ -824,8 +809,7 @@ public Iterable read() { } @Override - @NonNull - public ReadableState> readLater() { + public @NonNull ReadableState> readLater() { impl.keys().prefetch(); return this; } @@ -833,16 +817,23 @@ public ReadableState> readLater() { } @Override - @NonNull - public ReadableState>> entries() { - // TODO(https://github.com/apache/beam/issues/23616) - throw new UnsupportedOperationException( - "Multimap entries() is not currently supported with Fn API."); + public @NonNull ReadableState>> entries() { + return new ReadableState>>() { + @Override + public Iterable> read() { + return impl.entries(); + } + + @Override + public @NonNull ReadableState>> readLater() { + impl.entries().prefetch(); + return this; + } + }; } @Override - @NonNull - public ReadableState containsKey(KeyT key) { + public @NonNull ReadableState containsKey(KeyT key) { return new ReadableState() { @Override public Boolean read() { @@ -850,8 +841,7 @@ public Boolean read() { } @Override - @NonNull - public ReadableState readLater() { + public @NonNull ReadableState readLater() { impl.get(key).prefetch(); return this; } @@ -859,8 +849,7 @@ public ReadableState readLater() { } @Override - @NonNull - public ReadableState isEmpty() { + public @NonNull ReadableState isEmpty() { return new ReadableState() { @Override public Boolean read() { @@ -868,8 +857,7 @@ public Boolean read() { } @Override - @NonNull - public ReadableState readLater() { + public @NonNull ReadableState readLater() { impl.keys().prefetch(); return this; } From 19468cb239dcc6122ce7971420716b6cc42a417f Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 19 Sep 2025 21:24:57 +0000 Subject: [PATCH 04/10] Removes extra unchecked warning supression. I had added it because it turned my entire file in IntelliJ yellow. --- .../org/apache/beam/fn/harness/state/FnApiStateAccessor.java | 1 - 1 file changed, 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java index 563da1c9ebf9..b088c25404e3 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java @@ -89,7 +89,6 @@ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) "nullness", // TODO(https://github.com/apache/beam/issues/20497) - "unchecked" }) public class FnApiStateAccessor implements SideInputReader, StateBinder { From e0c4b678c462e516febc7f1fae3830543a0d7289 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 19 Sep 2025 21:28:47 +0000 Subject: [PATCH 05/10] Removes extra comma. --- .../org/apache/beam/fn/harness/state/FnApiStateAccessor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java index b088c25404e3..a341a675b294 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java @@ -88,7 +88,7 @@ /** Provides access to side inputs and state via a {@link BeamFnStateClient}. */ @SuppressWarnings({ "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) - "nullness", // TODO(https://github.com/apache/beam/issues/20497) + "nullness" // TODO(https://github.com/apache/beam/issues/20497) }) public class FnApiStateAccessor implements SideInputReader, StateBinder { From 442efba7cd3556ad7554f2304355d9d0abd2e2a5 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 22 Sep 2025 20:43:41 +0000 Subject: [PATCH 06/10] Removes not-needed @NonNull annotations. --- .../fn/harness/state/FnApiStateAccessor.java | 67 +++++++++---------- 1 file changed, 33 insertions(+), 34 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java index a341a675b294..6913c75a5f2d 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java @@ -81,7 +81,6 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; -import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; @@ -621,11 +620,11 @@ public SetState readLater() { } @Override - public @NonNull MapState bindMap( - @NonNull String id, - @NonNull StateSpec<@NonNull MapState> spec, - @NonNull Coder mapKeyCoder, - @NonNull Coder mapValueCoder) { + public MapState bindMap( + String id, + StateSpec> spec, + Coder mapKeyCoder, + Coder mapValueCoder) { return (MapState) stateKeyObjectCache.computeIfAbsent( createMultimapKeysUserStateKey(id), @@ -648,8 +647,8 @@ public void put(KeyT key, ValueT value) { } @Override - public @NonNull ReadableState computeIfAbsent( - KeyT key, @NonNull Function mappingFunction) { + public ReadableState computeIfAbsent( + KeyT key, Function mappingFunction) { Iterable values = impl.get(key); if (Iterables.isEmpty(values)) { impl.put(key, mappingFunction.apply(key)); @@ -663,12 +662,12 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState get(KeyT key) { + public ReadableState get(KeyT key) { return getOrDefault(key, null); } @Override - public @NonNull ReadableState getOrDefault( + public ReadableState getOrDefault( KeyT key, @Nullable ValueT defaultValue) { return new ReadableState() { @Override @@ -678,7 +677,7 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { impl.get(key).prefetch(); return this; } @@ -686,7 +685,7 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState> keys() { + public ReadableState> keys() { return new ReadableState>() { @Override public Iterable read() { @@ -694,7 +693,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { impl.keys().prefetch(); return this; } @@ -702,7 +701,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> values() { + public ReadableState> values() { return new ReadableState>() { @Override public Iterable read() { @@ -710,7 +709,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { entries().readLater(); return this; } @@ -718,7 +717,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState>> entries() { + public ReadableState>> entries() { return new ReadableState>>() { @Override public Iterable> read() { @@ -728,7 +727,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState>> readLater() { + public ReadableState>> readLater() { // Start prefetching the keys. We would need to block to start prefetching // the values. keys().readLater(); @@ -738,7 +737,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState isEmpty() { + public ReadableState isEmpty() { return new ReadableState() { @Override public Boolean read() { @@ -746,7 +745,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { keys().readLater(); return this; } @@ -758,11 +757,11 @@ public Boolean read() { } @Override - public @NonNull MultimapState bindMultimap( - @NonNull String id, - @NonNull StateSpec<@NonNull MultimapState> spec, - @NonNull Coder keyCoder, - @NonNull Coder valueCoder) { + public MultimapState bindMultimap( + String id, + StateSpec> spec, + Coder keyCoder, + Coder valueCoder) { return (MultimapState) stateKeyObjectCache.computeIfAbsent( createMultimapKeysUserStateKey(id), @@ -779,7 +778,7 @@ public void put(KeyT key, ValueT value) { } @Override - public @NonNull ReadableState> get(KeyT key) { + public ReadableState> get(KeyT key) { return new ReadableState>() { @Override public Iterable read() { @@ -787,7 +786,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { impl.get(key).prefetch(); return this; } @@ -800,7 +799,7 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState> keys() { + public ReadableState> keys() { return new ReadableState>() { @Override public Iterable read() { @@ -808,7 +807,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { impl.keys().prefetch(); return this; } @@ -816,7 +815,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState>> entries() { + public ReadableState>> entries() { return new ReadableState>>() { @Override public Iterable> read() { @@ -824,7 +823,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState>> readLater() { + public ReadableState>> readLater() { impl.entries().prefetch(); return this; } @@ -832,7 +831,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState containsKey(KeyT key) { + public ReadableState containsKey(KeyT key) { return new ReadableState() { @Override public Boolean read() { @@ -840,7 +839,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { impl.get(key).prefetch(); return this; } @@ -848,7 +847,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState isEmpty() { + public ReadableState isEmpty() { return new ReadableState() { @Override public Boolean read() { @@ -856,7 +855,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { impl.keys().prefetch(); return this; } From 88c34f183b4e1175c931df3d1acdac3241d1d9aa Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 22 Sep 2025 13:52:26 -0700 Subject: [PATCH 07/10] Update FnApiStateAccessor.java Removes non-needed @NonNull annotations. --- .../fn/harness/state/FnApiStateAccessor.java | 67 +++++++++---------- 1 file changed, 33 insertions(+), 34 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java index a341a675b294..6913c75a5f2d 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java @@ -81,7 +81,6 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; -import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; @@ -621,11 +620,11 @@ public SetState readLater() { } @Override - public @NonNull MapState bindMap( - @NonNull String id, - @NonNull StateSpec<@NonNull MapState> spec, - @NonNull Coder mapKeyCoder, - @NonNull Coder mapValueCoder) { + public MapState bindMap( + String id, + StateSpec> spec, + Coder mapKeyCoder, + Coder mapValueCoder) { return (MapState) stateKeyObjectCache.computeIfAbsent( createMultimapKeysUserStateKey(id), @@ -648,8 +647,8 @@ public void put(KeyT key, ValueT value) { } @Override - public @NonNull ReadableState computeIfAbsent( - KeyT key, @NonNull Function mappingFunction) { + public ReadableState computeIfAbsent( + KeyT key, Function mappingFunction) { Iterable values = impl.get(key); if (Iterables.isEmpty(values)) { impl.put(key, mappingFunction.apply(key)); @@ -663,12 +662,12 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState get(KeyT key) { + public ReadableState get(KeyT key) { return getOrDefault(key, null); } @Override - public @NonNull ReadableState getOrDefault( + public ReadableState getOrDefault( KeyT key, @Nullable ValueT defaultValue) { return new ReadableState() { @Override @@ -678,7 +677,7 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { impl.get(key).prefetch(); return this; } @@ -686,7 +685,7 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState> keys() { + public ReadableState> keys() { return new ReadableState>() { @Override public Iterable read() { @@ -694,7 +693,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { impl.keys().prefetch(); return this; } @@ -702,7 +701,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> values() { + public ReadableState> values() { return new ReadableState>() { @Override public Iterable read() { @@ -710,7 +709,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { entries().readLater(); return this; } @@ -718,7 +717,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState>> entries() { + public ReadableState>> entries() { return new ReadableState>>() { @Override public Iterable> read() { @@ -728,7 +727,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState>> readLater() { + public ReadableState>> readLater() { // Start prefetching the keys. We would need to block to start prefetching // the values. keys().readLater(); @@ -738,7 +737,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState isEmpty() { + public ReadableState isEmpty() { return new ReadableState() { @Override public Boolean read() { @@ -746,7 +745,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { keys().readLater(); return this; } @@ -758,11 +757,11 @@ public Boolean read() { } @Override - public @NonNull MultimapState bindMultimap( - @NonNull String id, - @NonNull StateSpec<@NonNull MultimapState> spec, - @NonNull Coder keyCoder, - @NonNull Coder valueCoder) { + public MultimapState bindMultimap( + String id, + StateSpec> spec, + Coder keyCoder, + Coder valueCoder) { return (MultimapState) stateKeyObjectCache.computeIfAbsent( createMultimapKeysUserStateKey(id), @@ -779,7 +778,7 @@ public void put(KeyT key, ValueT value) { } @Override - public @NonNull ReadableState> get(KeyT key) { + public ReadableState> get(KeyT key) { return new ReadableState>() { @Override public Iterable read() { @@ -787,7 +786,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { impl.get(key).prefetch(); return this; } @@ -800,7 +799,7 @@ public void remove(KeyT key) { } @Override - public @NonNull ReadableState> keys() { + public ReadableState> keys() { return new ReadableState>() { @Override public Iterable read() { @@ -808,7 +807,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState> readLater() { + public ReadableState> readLater() { impl.keys().prefetch(); return this; } @@ -816,7 +815,7 @@ public Iterable read() { } @Override - public @NonNull ReadableState>> entries() { + public ReadableState>> entries() { return new ReadableState>>() { @Override public Iterable> read() { @@ -824,7 +823,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState>> readLater() { + public ReadableState>> readLater() { impl.entries().prefetch(); return this; } @@ -832,7 +831,7 @@ public Iterable> read() { } @Override - public @NonNull ReadableState containsKey(KeyT key) { + public ReadableState containsKey(KeyT key) { return new ReadableState() { @Override public Boolean read() { @@ -840,7 +839,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { impl.get(key).prefetch(); return this; } @@ -848,7 +847,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState isEmpty() { + public ReadableState isEmpty() { return new ReadableState() { @Override public Boolean read() { @@ -856,7 +855,7 @@ public Boolean read() { } @Override - public @NonNull ReadableState readLater() { + public ReadableState readLater() { impl.keys().prefetch(); return this; } From 1ba5282550712c6e1802dbc9b093844789b9c90e Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Wed, 1 Oct 2025 23:22:47 +0000 Subject: [PATCH 08/10] Changes multimap entries() iterable to put values for the same key from the backend and local adds together. Also needed to make maybePrefetchable public. --- .../sdk/fn/stream/PrefetchableIterables.java | 2 +- .../apache/beam/sdk/transforms/ParDoTest.java | 7 +- .../fn/harness/state/MultimapUserState.java | 128 +++++++++--------- .../harness/state/MultimapUserStateTest.java | 15 ++ 4 files changed, 82 insertions(+), 70 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/PrefetchableIterables.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/PrefetchableIterables.java index dd7ec6b0f65a..1f7451e72a21 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/PrefetchableIterables.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/PrefetchableIterables.java @@ -94,7 +94,7 @@ public PrefetchableIterator createIterator() { * constructed that ensures that {@link PrefetchableIterator#prefetch()} is a no-op and {@link * PrefetchableIterator#isReady()} always returns true. */ - private static PrefetchableIterable maybePrefetchable(Iterable iterable) { + public static PrefetchableIterable maybePrefetchable(Iterable iterable) { if (iterable instanceof PrefetchableIterable) { return (PrefetchableIterable) iterable; } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java index e3464d98526c..8a273127b4fc 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java @@ -2951,17 +2951,18 @@ public void processElement( count.add(1); if (count.read() >= 4) { - // Those should be evaluated only when ReadableState.read is called. + // This should be evaluated only when ReadableState.read is called. ReadableState>> entriesView = state.entries(); - ReadableState containsBadView = state.containsKey("BadKey"); - // Those are evaluated immediately. + // This is evaluated immediately. Iterable> entries = state.entries().read(); state.remove("b"); assertEquals(4, Iterables.size(entries)); state.put("a", 2); + state.put("a", 3); + assertEquals(5, Iterables.size(entriesView.read())); // Note we output the view of state before the modifications in this if statement. for (Entry entry : entries) { r.output(KV.of(entry.getKey(), entry.getValue())); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java index 967cb76af1f3..8e3d76f5fc8f 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java @@ -29,6 +29,7 @@ import java.util.List; import java.util.Map; import java.util.NoSuchElementException; +import java.util.Objects; import java.util.Set; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; @@ -46,8 +47,11 @@ import org.apache.beam.sdk.util.ByteStringOutputStream; import org.apache.beam.sdk.values.KV; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; /** * An implementation of a multimap user state that utilizes the Beam Fn State API to fetch, clear @@ -92,6 +96,8 @@ public MultimapUserState( this.mapKeyCoder = mapKeyCoder; this.valueCoder = valueCoder; + // Note: These StateRequest protos are constructed even if we never try to read the + // corresponding state type. Consider constructing them lazily, as needed. this.keysStateRequest = StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = @@ -219,7 +225,7 @@ public boolean hasNext() { nextKey = persistedKeysIterator.next(); Object nextKeyStructuralValue = mapKeyCoder.structuralValue(nextKey); if (!pendingRemovesNow.contains(nextKeyStructuralValue)) { - // Remove all keys that we will visit when passing over the persistedKeysIterator + // Remove all keys that we will visit when passing over the persistedKeysIterator, // so we do not revisit them when passing over the pendingAddsNowIterator if (pendingAddsNow.containsKey(nextKeyStructuralValue)) { pendingAddsNow.remove(nextKeyStructuralValue); @@ -265,9 +271,11 @@ public PrefetchableIterable> entries() { !isClosed, "Multimap user state is no longer usable because it is closed for %s", keysStateRequest.getStateKey()); + // Make a deep copy of pendingAdds so this iterator represents a snapshot of state at the time + // it was created. + Map>> pendingAddsNow = ImmutableMap.copyOf(pendingAdds); if (isCleared) { - Map>> pendingAddsNow = new HashMap<>(pendingAdds); - return PrefetchableIterables.concat( + return PrefetchableIterables.maybePrefetchable( Iterables.concat( Iterables.transform( pendingAddsNow.entrySet(), @@ -277,93 +285,81 @@ public PrefetchableIterable> entries() { value -> Maps.immutableEntry(entry.getValue().getKey(), value))))); } - Set pendingRemovesNow = new HashSet<>(pendingRemoves.keySet()); - // Make a deep copy of pendingAdds so this iterator represents a snapshot of state at the time - // it was created. - Map> pendingAddsNow = new HashMap<>(); - for (Map.Entry>> entry : pendingAdds.entrySet()) { - pendingAddsNow.put( - entry.getValue().getKey(), new ArrayList<>()); // entry.getValue().getValue()); - for (V value : entry.getValue().getValue()) { - List values = pendingAddsNow.get(entry.getValue().getKey()); - values.add(value); - } - } + Set pendingRemovesNow = ImmutableSet.copyOf(pendingRemoves.keySet()); return new PrefetchableIterables.Default>() { @Override public PrefetchableIterator> createIterator() { return new PrefetchableIterator>() { - PrefetchableIterator> persistedEntriesIterator = + // We can get the same key multiple times from persistedEntries in the case that its + // values are paginated across multiple pages. Keep track of which keys we've seen, so we + // only add in pendingAdds once (with the first page). We'll also use it to return all + // keys not on the backend at the end of the iterator. + Set seenKeys = Sets.newHashSet(); + final PrefetchableIterator> allEntries = PrefetchableIterables.concat( Iterables.concat( - Iterables.transform( - persistedEntries, - entry -> - Iterables.transform( - entry.getValue(), - value -> Maps.immutableEntry(entry.getKey(), value))))) + Iterables.filter( + Iterables.transform( + persistedEntries, + entry -> { + final Object structuralKey = + mapKeyCoder.structuralValue(entry.getKey()); + if (pendingRemovesNow.contains(structuralKey)) { + return null; + } + // add returns true if we haven't seen this key yet. + if (seenKeys.add(structuralKey) + && pendingAddsNow.containsKey(structuralKey)) { + return PrefetchableIterables.concat( + Iterables.transform( + pendingAddsNow.get(structuralKey).getValue(), + pendingAdd -> + Maps.immutableEntry(entry.getKey(), pendingAdd)), + Iterables.transform( + entry.getValue(), + value -> Maps.immutableEntry(entry.getKey(), value))); + } + return Iterables.transform( + entry.getValue(), + value -> Maps.immutableEntry(entry.getKey(), value)); + }), + Objects::nonNull)), + Iterables.concat( + Iterables.filter( + Iterables.transform( + pendingAddsNow.entrySet(), + entry -> { + if (seenKeys.contains(entry.getKey())) { + return null; + } + return Iterables.transform( + entry.getValue().getValue(), + value -> + Maps.immutableEntry(entry.getValue().getKey(), value)); + }), + Objects::nonNull))) .iterator(); - Iterator> pendingAddsNowIterator; - boolean hasNext; - Map.Entry nextEntry; @Override public boolean isReady() { - return persistedEntriesIterator.isReady(); + return allEntries.isReady(); } @Override public void prefetch() { if (!isReady()) { - persistedEntriesIterator.prefetch(); + allEntries.prefetch(); } } @Override public boolean hasNext() { - if (hasNext) { - return true; - } - - while (persistedEntriesIterator.hasNext()) { - Map.Entry nextPersistedEntry = persistedEntriesIterator.next(); - Object nextKeyStructuralValue = - mapKeyCoder.structuralValue(nextPersistedEntry.getKey()); - if (!pendingRemovesNow.contains(nextKeyStructuralValue)) { - nextEntry = - Maps.immutableEntry(nextPersistedEntry.getKey(), nextPersistedEntry.getValue()); - hasNext = true; - return true; - } - } - - if (pendingAddsNowIterator == null) { - pendingAddsNowIterator = - Iterables.concat( - Iterables.transform( - pendingAddsNow.entrySet(), - entry -> - Iterables.transform( - entry.getValue(), - value -> Maps.immutableEntry(entry.getKey(), value)))) - .iterator(); - } - while (pendingAddsNowIterator.hasNext()) { - nextEntry = pendingAddsNowIterator.next(); - hasNext = true; - return true; - } - - return false; + return allEntries.hasNext(); } @Override public Map.Entry next() { - if (!hasNext()) { - throw new NoSuchElementException(); - } - hasNext = false; - return nextEntry; + return allEntries.next(); } }; } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java index 457362815cf6..1d57b33dfa1d 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java @@ -213,6 +213,9 @@ public void testEntries() throws Exception { KV.of(ByteString.copyFrom(A2), "V3"))); userState.put(A1, "V4"); + // Iterable is a snapshot of the entries at this time. + PrefetchableIterable> entriesBeforeOperations = userState.entries(); + assertThat( StreamSupport.stream(userState.entries().spliterator(), false) .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) @@ -240,6 +243,18 @@ public void testEntries() throws Exception { userState.clear(); assertThat(userState.entries(), emptyIterable()); + // Check that after applying all these operations, our original entries Iterable contains a + // snapshot of state from when it was created. + assertThat( + StreamSupport.stream(entriesBeforeOperations.spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A1), "V1"), + KV.of(ByteString.copyFrom(A1), "V2"), + KV.of(ByteString.copyFrom(A1), "V4"), + KV.of(ByteString.copyFrom(A2), "V3"))); + userState.asyncClose(); assertThrows(IllegalStateException.class, () -> userState.entries()); } From cf3c8f4a1e8b723c27b779e4530a8752a6ed0117 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Mon, 20 Oct 2025 19:19:04 +0000 Subject: [PATCH 09/10] Adds a test that prefetching multimap entries results in a StateRequest sent across FnApi. --- .../harness/state/MultimapUserStateTest.java | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java index 1d57b33dfa1d..679307321826 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java @@ -700,6 +700,44 @@ public void testRemoveKeysPrefetch() throws Exception { assertEquals(0, fakeClient.getCallCount()); } + @Test + public void testEntriesPrefetched() throws Exception { + // Use a really large chunk size so all elements get returned in a single page. This makes it + // easier to count how many get calls we should expect. + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapEntriesStateKey(), + KV.of( + KvCoder.of(ByteArrayCoder.of(), IterableCoder.of(StringUtf8Coder.of())), + asList(KV.of(A1, asList("V1", "V2")), KV.of(A2, asList("V3"))))), + 1000000); + MultimapUserState userState = + new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createMultimapKeyStateKey(), + ByteArrayCoder.of(), + StringUtf8Coder.of()); + + userState.put(A1, "V4"); + PrefetchableIterable> entries = userState.entries(); + assertEquals(0, fakeClient.getCallCount()); + entries.prefetch(); + assertEquals(1, fakeClient.getCallCount()); + assertThat( + StreamSupport.stream(entries.spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A1), "V1"), + KV.of(ByteString.copyFrom(A1), "V2"), + KV.of(ByteString.copyFrom(A1), "V4"), + KV.of(ByteString.copyFrom(A2), "V3"))); + assertEquals(1, fakeClient.getCallCount()); + } + @Test public void testClearPrefetch() throws Exception { FakeBeamFnStateClient fakeClient = From 555638acba5f7febdde70f4561fe9d0a0ac5fdb8 Mon Sep 17 00:00:00 2001 From: Andrew Crites Date: Fri, 24 Oct 2025 21:06:16 +0000 Subject: [PATCH 10/10] Adds an environment capability for multimap state and sets in for the java sdk. --- .../beam/model/pipeline/v1/beam_runner_api.proto | 10 +++++++--- .../beam/sdk/util/construction/Environments.java | 1 + .../beam/sdk/util/construction/EnvironmentsTest.java | 3 +++ 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto b/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto index c615b2a5279a..0bdc4f69aab6 100644 --- a/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto +++ b/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto @@ -1621,13 +1621,13 @@ message AnyOfEnvironmentPayload { // environment understands. message StandardProtocols { enum Enum { - // Indicates suport for progress reporting via the legacy Metrics proto. + // Indicates support for progress reporting via the legacy Metrics proto. LEGACY_PROGRESS_REPORTING = 0 [(beam_urn) = "beam:protocol:progress_reporting:v0"]; - // Indicates suport for progress reporting via the new MonitoringInfo proto. + // Indicates support for progress reporting via the new MonitoringInfo proto. PROGRESS_REPORTING = 1 [(beam_urn) = "beam:protocol:progress_reporting:v1"]; - // Indicates suport for worker status protocol defined at + // Indicates support for worker status protocol defined at // https://s.apache.org/beam-fn-api-harness-status. WORKER_STATUS = 2 [(beam_urn) = "beam:protocol:worker_status:v1"]; @@ -1681,6 +1681,10 @@ message StandardProtocols { // Indicates support for reading, writing and propagating Element's metadata ELEMENT_METADATA = 11 [(beam_urn) = "beam:protocol:element_metadata:v1"]; + + // Indicates whether the SDK supports multimap state. + MULTIMAP_STATE = 12 + [(beam_urn) = "beam:protocol:multimap_state:v1"]; } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java index 05ecb21fd956..7a7725cc26d2 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Environments.java @@ -495,6 +495,7 @@ public static Set getJavaCapabilities() { capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.DATA_SAMPLING)); capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.SDK_CONSUMING_RECEIVED_DATA)); capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.ORDERED_LIST_STATE)); + capabilities.add(BeamUrns.getUrn(StandardProtocols.Enum.MULTIMAP_STATE)); return capabilities.build(); } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/EnvironmentsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/EnvironmentsTest.java index 410b52cba23b..46aaeebd1cdf 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/EnvironmentsTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/EnvironmentsTest.java @@ -219,6 +219,9 @@ public void testCapabilities() { assertThat( Environments.getJavaCapabilities(), hasItem(BeamUrns.getUrn(RunnerApi.StandardProtocols.Enum.ORDERED_LIST_STATE))); + assertThat( + Environments.getJavaCapabilities(), + hasItem(BeamUrns.getUrn(RunnerApi.StandardProtocols.Enum.MULTIMAP_STATE))); // Check that SDF truncation is supported assertThat( Environments.getJavaCapabilities(),