From 90d6c96d594a86d41ecc15bd4b31a2596fa5c097 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 23 Feb 2026 14:09:49 -0500 Subject: [PATCH 01/33] Propagate 'has_no_state' and 'only_bundle_for_keys' bits from ProcessBundleHandler to where StateAppendRequest and StateGetRequest are constructed --- .../fn/harness/PTransformRunnerFactory.java | 12 ++++++ .../harness/control/ProcessBundleHandler.java | 22 +++++++++-- .../beam/fn/harness/state/BagUserState.java | 18 ++++++--- .../fn/harness/state/FnApiStateAccessor.java | 39 +++++++++++++++---- .../fn/harness/state/IterableSideInput.java | 6 ++- .../fn/harness/state/MultimapSideInput.java | 16 ++++++-- .../fn/harness/state/MultimapUserState.java | 16 ++++++-- .../harness/state/OrderedListUserState.java | 12 ++++-- .../harness/state/StateFetchingIterators.java | 22 +++++++++-- 9 files changed, 130 insertions(+), 33 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java index 8e7c670d5c26..8752408770ac 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java @@ -160,6 +160,18 @@ void addIncomingTimerEndpoint( * instant provides the timeout on how long the finalization callback is valid for. */ DoFn.BundleFinalizer getBundleFinalizer(); + + /** + * Returns true if the runner has no state for the keys in the ProcessBundleRequest. If true, + * the SDK can begin stateful processing with an initial empty state. + */ + boolean getHasNoState(); + + /** + * Returns true if the runner will never process another bundle for the keys it contains. + * Therefore, the generated state need not be included in the bundle commit. + */ + boolean getOnlyBundleForKeys(); } /** diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index b8ad51816a7a..e4cffeed0b53 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -246,7 +246,9 @@ private void addRunnerAndConsumersForPTransformRecursively( BundleFinalizer bundleFinalizer, Collection> channelRoots, Map outboundAggregatorMap, - Set runnerCapabilities) + Set runnerCapabilities, + boolean hasNoState, + boolean onlyBundleForKeys) throws IOException { // Recursively ensure that all consumers of the output PCollection have been created. @@ -279,7 +281,9 @@ private void addRunnerAndConsumersForPTransformRecursively( bundleFinalizer, channelRoots, outboundAggregatorMap, - runnerCapabilities); + runnerCapabilities, + hasNoState, + onlyBundleForKeys); } } @@ -488,6 +492,16 @@ public BundleSplitListener getSplitListener() { public BundleFinalizer getBundleFinalizer() { return bundleFinalizer; } + + @Override + public boolean getHasNoState() { + return hasNoState; + } + + @Override + public boolean getOnlyBundleForKeys() { + return onlyBundleForKeys; + } }); processedPTransformIds.add(pTransformId); } @@ -913,7 +927,9 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { bundleFinalizer, bundleProcessor.getChannelRoots(), bundleProcessor.getOutboundAggregators(), - bundleProcessor.getRunnerCapabilities()); + bundleProcessor.getRunnerCapabilities(), + processBundleRequest.getHasNoState(), + processBundleRequest.getOnlyBundleForKeys()); } bundleProcessor.finish(); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index ba56c6d656ca..b579423d9446 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -54,27 +54,33 @@ public class BagUserState { private List newValues; private boolean isCleared; private boolean isClosed; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; static final int BAG_APPEND_BATCHING_LIMIT = 10 * 1024 * 1024; /** The cache must be namespaced for this state object accordingly. */ public BagUserState( - Cache cache, - BeamFnStateClient beamFnStateClient, - String instructionId, - StateKey stateKey, - Coder valueCoder) { + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder valueCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasBagUserState(), "Expected BagUserState StateKey but received %s.", stateKey); this.cache = cache; this.beamFnStateClient = beamFnStateClient; this.valueCoder = valueCoder; + this.hasNoState = hasNoState; + this.onlyBundleForKeys = onlyBundleForKeys; this.request = StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.oldValues = StateFetchingIterators.readAllAndDecodeStartingFrom( - this.cache, beamFnStateClient, request, valueCoder); + this.cache, beamFnStateClient, request, valueCoder, hasNoState); this.newValues = new ArrayList<>(); } 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 6913c75a5f2d..c14d9745bc30 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 @@ -114,6 +114,8 @@ public static class Factory { private final Map, SideInputSpec> sideInputSpecMap; private final Coder keyCoder; private final Coder windowCoder; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; public Factory( PipelineOptions pipelineOptions, @@ -126,7 +128,9 @@ public Factory( Map, SideInputSpec> sideInputSpecMap, BeamFnStateClient beamFnStateClient, Coder keyCoder, - Coder windowCoder) { + Coder windowCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.ptransformId = ptransformId; @@ -138,6 +142,8 @@ public Factory( this.beamFnStateClient = beamFnStateClient; this.keyCoder = keyCoder; this.windowCoder = windowCoder; + this.hasNoState = hasNoState; + this.onlyBundleForKeys = onlyBundleForKeys; } public static Factory factoryForPTransformContext( @@ -220,7 +226,9 @@ public static Factory factoryForPTransformContext( tagToSideInputSpecMap, context.getBeamFnStateClient(), keyCoder, - windowCoder); + windowCoder, + context.getHasNoState(), + context.getOnlyBundleForKeys()); } public FnApiStateAccessor create() { @@ -235,7 +243,9 @@ public FnApiStateAccessor create() { sideInputSpecMap, beamFnStateClient, keyCoder, - windowCoder); + windowCoder, + hasNoState, + onlyBundleForKeys); } } @@ -252,6 +262,8 @@ public FnApiStateAccessor create() { private final Collection stateFinalizers; private final Coder keyCoder; private final Coder windowCoder; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; private @Nullable Supplier currentWindowSupplier; private @Nullable Supplier encodedCurrentKeySupplier; @@ -268,7 +280,9 @@ public FnApiStateAccessor( Map, SideInputSpec> sideInputSpecMap, BeamFnStateClient beamFnStateClient, Coder keyCoder, - Coder windowCoder) { + Coder windowCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.stateKeyObjectCache = Maps.newHashMap(); @@ -282,6 +296,8 @@ public FnApiStateAccessor( this.keyCoder = keyCoder; this.windowCoder = windowCoder; this.stateFinalizers = new ArrayList<>(); + this.hasNoState = hasNoState; + this.onlyBundleForKeys = onlyBundleForKeys; } public void setKeyAndWindowContext(MutatingStateContext keyAndWindowContext) { @@ -417,7 +433,8 @@ public T get(PCollectionView view, BoundedWindow window) { runnerCapabilities.contains( BeamUrns.getUrn( RunnerApi.StandardRunnerProtocols.Enum - .MULTIMAP_KEYS_VALUES_SIDE_INPUT)))); + .MULTIMAP_KEYS_VALUES_SIDE_INPUT)), + hasNoState)); default: throw new IllegalStateException( String.format( @@ -1201,7 +1218,9 @@ private BagUserState createBagUserState(StateKey stateKey, Coder value beamFnStateClient, processBundleInstructionId.get(), stateKey, - valueCoder); + valueCoder, + hasNoState, + onlyBundleForKeys); stateFinalizers.add(rval::asyncClose); return rval; } @@ -1283,7 +1302,9 @@ private MultimapUserState createMultimapUserState( processBundleInstructionId.get(), stateKey, keyCoder, - valueCoder); + valueCoder, + hasNoState, + onlyBundleForKeys); stateFinalizers.add(rval::asyncClose); return rval; } @@ -1318,7 +1339,9 @@ private OrderedListUserState createOrderedListUserState( beamFnStateClient, processBundleInstructionId.get(), stateKey, - valueCoder); + valueCoder, + hasNoState, + onlyBundleForKeys); stateFinalizers.add(rval::asyncClose); return rval; } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java index 7f87cb4d4e41..b5492bfa8565 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java @@ -40,7 +40,8 @@ public IterableSideInput( BeamFnStateClient beamFnStateClient, String instructionId, StateKey stateKey, - Coder valueCoder) { + Coder valueCoder, + boolean hasNoState) { checkArgument( stateKey.hasIterableSideInput(), "Expected IterableSideInput StateKey but received %s.", @@ -50,7 +51,8 @@ public IterableSideInput( cache, beamFnStateClient, StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(), - valueCoder); + valueCoder, + hasNoState); } @Override diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java index 8e38a57ff9fe..e403332d8f13 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java @@ -55,6 +55,7 @@ public class MultimapSideInput implements MultimapView { private final Coder valueCoder; private volatile Function> bulkReadResult; private final boolean useBulkRead; + private final boolean hasNoState; public MultimapSideInput( Cache cache, @@ -63,7 +64,8 @@ public MultimapSideInput( StateKey stateKey, Coder keyCoder, Coder valueCoder, - boolean useBulkRead) { + boolean useBulkRead, + boolean hasNoState) { checkArgument( stateKey.hasMultimapKeysSideInput(), "Expected MultimapKeysSideInput StateKey but received %s.", @@ -75,12 +77,13 @@ public MultimapSideInput( this.keyCoder = keyCoder; this.valueCoder = valueCoder; this.useBulkRead = useBulkRead; + this.hasNoState = hasNoState; } @Override public Iterable get() { return StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, beamFnStateClient, keysRequest, keyCoder); + cache, beamFnStateClient, keysRequest, keyCoder, hasNoState); } @Override @@ -120,7 +123,8 @@ public Iterable get(K k) { Caches.noop(), beamFnStateClient, bulkReadRequest, - KvCoder.of(keyCoder, IterableCoder.of(valueCoder))) + KvCoder.of(keyCoder, IterableCoder.of(valueCoder)), + hasNoState) .iterator(); while (bulkRead.size() < BULK_READ_SIZE && entries.hasNext()) { KV> entry = entries.next(); @@ -169,7 +173,11 @@ public Iterable get(K k) { StateRequest request = keysRequest.toBuilder().setStateKey(stateKey).build(); return StateFetchingIterators.readAllAndDecodeStartingFrom( - Caches.subCache(cache, "ValuesForKey", encodedKey), beamFnStateClient, request, valueCoder); + Caches.subCache(cache, "ValuesForKey", encodedKey), + beamFnStateClient, + request, + valueCoder, + hasNoState); } private ByteString encodeKey(K k) { 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 83d78ff836c7..c4a1482444fb 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 @@ -72,6 +72,8 @@ public class MultimapUserState { private boolean isClosed; private boolean isCleared; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; // Pending updates to persistent storage private HashMap pendingRemoves = Maps.newHashMap(); private HashMap>> pendingAdds = Maps.newHashMap(); @@ -84,7 +86,9 @@ public MultimapUserState( String instructionId, StateKey stateKey, Coder mapKeyCoder, - Coder valueCoder) { + Coder valueCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasMultimapKeysUserState(), "Expected MultimapKeysUserState StateKey but received %s.", @@ -93,6 +97,8 @@ public MultimapUserState( this.beamFnStateClient = beamFnStateClient; this.mapKeyCoder = mapKeyCoder; this.valueCoder = valueCoder; + this.hasNoState = hasNoState; + this.onlyBundleForKeys = onlyBundleForKeys; // Note: These StateRequest protos are constructed even if we never try to read the // corresponding state type. Consider constructing them lazily, as needed. @@ -100,7 +106,7 @@ public MultimapUserState( StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, beamFnStateClient, keysStateRequest, mapKeyCoder); + cache, beamFnStateClient, keysStateRequest, mapKeyCoder, hasNoState); StateRequest.Builder userStateRequestBuilder = StateRequest.newBuilder(); userStateRequestBuilder @@ -128,7 +134,8 @@ public MultimapUserState( Caches.subCache(this.cache, "AllEntries"), beamFnStateClient, entriesStateRequest, - KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder))); + KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder)), + hasNoState); } public void clear() { @@ -542,7 +549,8 @@ private CachingStateIterable getPersistedValues(Object structuralKey, K key) request.getStateKey().getMultimapUserState().getMapKey()), beamFnStateClient, request, - valueCoder)); + valueCoder, + hasNoState)); }) .getValue(); } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index 47b5057880b9..c841e1b7eec0 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -88,7 +88,8 @@ public class OrderedListUserState { private boolean isCleared = false; private boolean isClosed = false; - + private final boolean hasNoState; + private final boolean onlyBundleForKeys; public static class TimestampedValueCoder extends StructuredCoder> { private final Coder valueCoder; @@ -161,7 +162,9 @@ public OrderedListUserState( BeamFnStateClient beamFnStateClient, String instructionId, StateKey stateKey, - Coder valueCoder) { + Coder valueCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasOrderedListUserState(), "Expected OrderedListUserState StateKey but received %s.", @@ -170,6 +173,8 @@ public OrderedListUserState( this.timestampedValueCoder = TimestampedValueCoder.of(valueCoder); this.requestTemplate = StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); + this.hasNoState = hasNoState; + this.onlyBundleForKeys = onlyBundleForKeys; } public void add(TimestampedValue value) { @@ -218,7 +223,8 @@ public Iterable> readRange(Instant minTimestamp, Instant lim Caches.noop(), this.beamFnStateClient, getRequestBuilder.build(), - this.timestampedValueCoder); + this.timestampedValueCoder, + hasNoState); // Make a snapshot of the current pendingRemoves and use them to filter persistent values. // The values of pendingRemoves are copied, so that they will still be accessible in diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 1e06c98f2e31..6210aba2c801 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -86,12 +86,14 @@ public static CachingStateIterable readAllAndDecodeStartingFrom( Cache cache, BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk, - Coder valueCoder) { + Coder valueCoder, + boolean hasNoState) { return new CachingStateIterable<>( (Cache>) cache, beamFnStateClient, stateRequestForFirstChunk, - valueCoder); + valueCoder, + hasNoState); } /** @@ -328,16 +330,19 @@ public static Block fromValues( private final BeamFnStateClient beamFnStateClient; private final StateRequest stateRequestForFirstChunk; private final Coder valueCoder; + private final boolean hasNoState; public CachingStateIterable( Cache> cache, BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk, - Coder valueCoder) { + Coder valueCoder, + boolean hasNoState) { this.cache = cache; this.beamFnStateClient = beamFnStateClient; this.stateRequestForFirstChunk = stateRequestForFirstChunk; this.valueCoder = valueCoder; + this.hasNoState = hasNoState; } /** @@ -669,12 +674,23 @@ static class LazyBlockingStateFetchingIterator implements PrefetchableIterator prefetchedResponse; + private boolean hasNoState; LazyBlockingStateFetchingIterator( BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk) { this.beamFnStateClient = beamFnStateClient; this.stateRequestForFirstChunk = stateRequestForFirstChunk; this.continuationToken = stateRequestForFirstChunk.getGet().getContinuationToken(); + this.hasNoState = false; + } + + LazyBlockingStateFetchingIterator( + BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk, + boolean hasNoState) { + this.beamFnStateClient = beamFnStateClient; + this.stateRequestForFirstChunk = stateRequestForFirstChunk; + this.continuationToken = stateRequestForFirstChunk.getGet().getContinuationToken(); + this.hasNoState = hasNoState; } /** From 4b6d58495f94bd96c6e50927f3c5026c58dcd968 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 23 Feb 2026 14:10:17 -0500 Subject: [PATCH 02/33] Skip StateRequest construction if the bits are true --- .../java/org/apache/beam/fn/harness/state/BagUserState.java | 6 +++--- .../beam/fn/harness/state/StateFetchingIterators.java | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index b579423d9446..d39bf0b5347b 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -141,7 +141,7 @@ public void asyncClose() throws Exception { for (T newValue : newValues) { int previousSize = out.size(); valueCoder.encode(newValue, out); - if (out.size() > BAG_APPEND_BATCHING_LIMIT && previousSize > 0) { + if (out.size() > BAG_APPEND_BATCHING_LIMIT && previousSize > 0 && !onlyBundleForKeys) { // Respect the batching limit by outputting the previous batch of // elements. beamFnStateClient.handle( @@ -151,7 +151,7 @@ public void asyncClose() throws Exception { StateAppendRequest.newBuilder() .setData(out.consumePrefixToByteString(previousSize)))); } - if (out.size() > BAG_APPEND_BATCHING_LIMIT) { + if (out.size() > BAG_APPEND_BATCHING_LIMIT && !onlyBundleForKeys) { // The last element was over the batching limit by itself. To avoid // exceeding runner state limits due to large elements, we output // without additional batching. @@ -161,7 +161,7 @@ public void asyncClose() throws Exception { .setAppend(StateAppendRequest.newBuilder().setData(out.toByteStringAndReset()))); } } - if (out.size() > 0) { + if (out.size() > 0 && !onlyBundleForKeys) { beamFnStateClient.handle( request .toBuilder() diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 6210aba2c801..0bed6a5f6e0b 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -729,7 +729,7 @@ public boolean isReady() { @Override public void prefetch() { - if (continuationToken != null && prefetchedResponse == null) { + if (continuationToken != null && prefetchedResponse == null && !hasNoState) { prefetchedResponse = loadPrefetchedResponse(continuationToken); } } From 708247a8c0a07b5f5dc95c391332a2d1fe9f836c Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 23 Feb 2026 17:28:38 -0500 Subject: [PATCH 03/33] minor fix --- .../org/apache/beam/fn/harness/state/MultimapUserState.java | 2 +- .../apache/beam/fn/harness/state/StateFetchingIterators.java | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) 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 c4a1482444fb..da858240e569 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 @@ -445,7 +445,7 @@ private void startStateApiWrites() { } // Persist pending key-values - if (!pendingAdds.isEmpty()) { + if (!pendingAdds.isEmpty() && !onlyBundleForKeys) { for (KV> entry : pendingAdds.values()) { StateRequest request = createUserStateRequest(entry.getKey()); beamFnStateClient.handle( diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 0bed6a5f6e0b..37a71f1630e4 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -515,7 +515,8 @@ class CachingStateIterator implements PrefetchableIterator { public CachingStateIterator() { this.underlyingStateFetchingIterator = - new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk); + new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk, + hasNoState); this.dataStreamDecoder = new DataStreamDecoder<>(valueCoder, underlyingStateFetchingIterator); this.currentBlock = From 6ef00e8073202bffddb882c5cbf48a20b90c18cd Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 23 Feb 2026 17:50:06 -0500 Subject: [PATCH 04/33] minor fix --- .../apache/beam/fn/harness/state/OrderedListUserState.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index c841e1b7eec0..dc8b92002213 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -309,6 +309,11 @@ public void asyncClose() throws Exception { pendingRemoves.clear(); } + if (hasNoState) { + pendingAdds.clear(); + return; + } + if (!pendingAdds.isEmpty()) { ByteStringOutputStream outStream = new ByteStringOutputStream(); From da95d800ff172835d585e6c5da76cf3d47e2be8e Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 23 Feb 2026 17:55:46 -0500 Subject: [PATCH 05/33] minor fix --- .../org/apache/beam/fn/harness/state/OrderedListUserState.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index dc8b92002213..008ecb77cb37 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -309,7 +309,7 @@ public void asyncClose() throws Exception { pendingRemoves.clear(); } - if (hasNoState) { + if (onlyBundleForKeys) { pendingAdds.clear(); return; } From b4836e612a6e6a424953dbfc3ff42cf46b70bb44 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Tue, 24 Feb 2026 03:39:46 -0500 Subject: [PATCH 06/33] if has_no_state is true, will not load prefetched response --- .../apache/beam/fn/harness/state/StateFetchingIterators.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 37a71f1630e4..639b1030a0cd 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -766,6 +766,8 @@ public ByteString next() { } Throwables.throwIfUnchecked(e.getCause()); throw new IllegalStateException(e.getCause()); + } catch (NullPointerException e) { + throw new NoSuchElementException(); } prefetchedResponse = null; From cfe916cd45fd80561ec318d1504fbd7ca74a24dd Mon Sep 17 00:00:00 2001 From: Kai Li Date: Tue, 24 Feb 2026 03:57:12 -0500 Subject: [PATCH 07/33] when has_no_state is true, continuationToken set to null --- .../beam/fn/harness/state/StateFetchingIterators.java | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 639b1030a0cd..9280c3191308 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -675,14 +675,12 @@ static class LazyBlockingStateFetchingIterator implements PrefetchableIterator prefetchedResponse; - private boolean hasNoState; LazyBlockingStateFetchingIterator( BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk) { this.beamFnStateClient = beamFnStateClient; this.stateRequestForFirstChunk = stateRequestForFirstChunk; this.continuationToken = stateRequestForFirstChunk.getGet().getContinuationToken(); - this.hasNoState = false; } LazyBlockingStateFetchingIterator( @@ -690,8 +688,7 @@ static class LazyBlockingStateFetchingIterator implements PrefetchableIterator Date: Tue, 24 Feb 2026 04:12:29 -0500 Subject: [PATCH 08/33] simplify --- .../org/apache/beam/fn/harness/state/BagUserState.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index d39bf0b5347b..c4b98153709f 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -133,7 +133,7 @@ public void asyncClose() throws Exception { beamFnStateClient.handle( request.toBuilder().setClear(StateClearRequest.getDefaultInstance())); } - if (!newValues.isEmpty()) { + if (!onlyBundleForKeys && !newValues.isEmpty()) { // Batch values up to a arbitrary limit to reduce overhead of write // requests. We treat this limit as strict to ensure that large elements // are not batched as they may otherwise exceed runner limits. @@ -141,7 +141,7 @@ public void asyncClose() throws Exception { for (T newValue : newValues) { int previousSize = out.size(); valueCoder.encode(newValue, out); - if (out.size() > BAG_APPEND_BATCHING_LIMIT && previousSize > 0 && !onlyBundleForKeys) { + if (out.size() > BAG_APPEND_BATCHING_LIMIT && previousSize > 0) { // Respect the batching limit by outputting the previous batch of // elements. beamFnStateClient.handle( @@ -151,7 +151,7 @@ public void asyncClose() throws Exception { StateAppendRequest.newBuilder() .setData(out.consumePrefixToByteString(previousSize)))); } - if (out.size() > BAG_APPEND_BATCHING_LIMIT && !onlyBundleForKeys) { + if (out.size() > BAG_APPEND_BATCHING_LIMIT) { // The last element was over the batching limit by itself. To avoid // exceeding runner state limits due to large elements, we output // without additional batching. @@ -161,7 +161,7 @@ public void asyncClose() throws Exception { .setAppend(StateAppendRequest.newBuilder().setData(out.toByteStringAndReset()))); } } - if (out.size() > 0 && !onlyBundleForKeys) { + if (out.size() > 0) { beamFnStateClient.handle( request .toBuilder() From d5acb6b3f2c452519050ec060c6de0a7153f2a2d Mon Sep 17 00:00:00 2001 From: LEEKYE <50568168+LEEKYE@users.noreply.github.com> Date: Tue, 24 Feb 2026 04:19:03 -0500 Subject: [PATCH 09/33] Update sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../apache/beam/fn/harness/state/StateFetchingIterators.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 9280c3191308..0f53a472ba4c 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -684,8 +684,9 @@ static class LazyBlockingStateFetchingIterator implements PrefetchableIterator Date: Sun, 8 Mar 2026 21:30:00 -0400 Subject: [PATCH 10/33] Apply suggestion from @ryan-mbuashundip Co-authored-by: Ryan Mbuashu-Ndip --- .../apache/beam/fn/harness/state/BagUserState.java | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index c4b98153709f..196bb3095cf2 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -61,13 +61,13 @@ public class BagUserState { /** The cache must be namespaced for this state object accordingly. */ public BagUserState( - Cache cache, - BeamFnStateClient beamFnStateClient, - String instructionId, - StateKey stateKey, - Coder valueCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder valueCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasBagUserState(), "Expected BagUserState StateKey but received %s.", stateKey); this.cache = cache; From 5f030112f7e0124dc265905709e5a92bb6a5cc0f Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 9 Mar 2026 00:58:08 -0400 Subject: [PATCH 11/33] fix based on comments --- .../beam/fn/harness/state/BagUserState.java | 4 +-- .../fn/harness/state/FnApiStateAccessor.java | 3 ++- .../fn/harness/state/IterableSideInput.java | 5 ++-- .../fn/harness/state/MultimapSideInput.java | 13 ++++++--- .../fn/harness/state/MultimapUserState.java | 27 ++++++++++--------- .../harness/state/OrderedListUserState.java | 11 ++++---- .../harness/state/StateFetchingIterators.java | 19 +++++-------- 7 files changed, 42 insertions(+), 40 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index c4b98153709f..43f8d577a6b4 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -126,14 +126,14 @@ public void asyncClose() throws Exception { "Bag user state is no longer usable because it is closed for %s", request.getStateKey()); isClosed = true; - if (!isCleared && newValues.isEmpty()) { + if (onlyBundleForKeys || (!isCleared && newValues.isEmpty())) { return; } if (isCleared) { beamFnStateClient.handle( request.toBuilder().setClear(StateClearRequest.getDefaultInstance())); } - if (!onlyBundleForKeys && !newValues.isEmpty()) { + if (!newValues.isEmpty()) { // Batch values up to a arbitrary limit to reduce overhead of write // requests. We treat this limit as strict to ensure that large elements // are not batched as they may otherwise exceed runner limits. 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 c14d9745bc30..70b6f5284ad7 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 @@ -418,7 +418,8 @@ public T get(PCollectionView view, BoundedWindow window) { beamFnStateClient, processBundleInstructionId.get(), key, - sideInputSpec.getCoder())); + sideInputSpec.getCoder(), + hasNoState)); case Materializations.MULTIMAP_MATERIALIZATION_URN: return sideInputSpec .getViewFn() diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java index b5492bfa8565..d7ce194d1164 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java @@ -46,13 +46,12 @@ public IterableSideInput( stateKey.hasIterableSideInput(), "Expected IterableSideInput StateKey but received %s.", stateKey); - this.values = + this.values = hasNoState ? PrefetchableIterables.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( cache, beamFnStateClient, StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(), - valueCoder, - hasNoState); + valueCoder); } @Override diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java index e403332d8f13..46ce7db17857 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java @@ -82,12 +82,18 @@ public MultimapSideInput( @Override public Iterable get() { - return StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, beamFnStateClient, keysRequest, keyCoder, hasNoState); + return this.hasNoState + ? PrefetchableIterables.emptyIterable() + : StateFetchingIterators.readAllAndDecodeStartingFrom( + cache, beamFnStateClient, keysRequest, keyCoder); } @Override public Iterable get(K k) { + if (this.hasNoState) { + return PrefetchableIterables.emptyIterable(); + } + ByteString encodedKey = encodeKey(k); if (useBulkRead) { @@ -176,8 +182,7 @@ public Iterable get(K k) { Caches.subCache(cache, "ValuesForKey", encodedKey), beamFnStateClient, request, - valueCoder, - hasNoState); + valueCoder); } private ByteString encodeKey(K k) { 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 da858240e569..b44e403ab7e2 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 @@ -105,8 +105,10 @@ public MultimapUserState( this.keysStateRequest = StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = - StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, beamFnStateClient, keysStateRequest, mapKeyCoder, hasNoState); + hasNoState + ? PrefetchableIterables.emptyIterable() + : StateFetchingIterators.readAllAndDecodeStartingFrom( + cache, beamFnStateClient, keysStateRequest, mapKeyCoder, hasNoState); StateRequest.Builder userStateRequestBuilder = StateRequest.newBuilder(); userStateRequestBuilder @@ -130,12 +132,13 @@ public MultimapUserState( .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)), - hasNoState); + hasNoState + ? PrefetchableIterables.emptyIterable() + : StateFetchingIterators.readAllAndDecodeStartingFrom( + Caches.subCache(this.cache, "AllEntries"), + beamFnStateClient, + entriesStateRequest, + KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder))); } public void clear() { @@ -186,7 +189,7 @@ public PrefetchableIterable keys() { !isClosed, "Multimap user state is no longer usable because it is closed for %s", keysStateRequest.getStateKey()); - if (isCleared) { + if (isCleared || hasNoState) { List keys = new ArrayList<>(pendingAdds.size()); for (Map.Entry>> entry : pendingAdds.entrySet()) { keys.add(entry.getValue().getKey()); @@ -284,7 +287,7 @@ public PrefetchableIterable> entries() { entry.getKey(), KV.of(entry.getValue().getKey(), new ArrayList<>(entry.getValue().getValue()))); } - if (isCleared) { + if (isCleared || hasNoState) { return PrefetchableIterables.maybePrefetchable( Iterables.concat( Iterables.transform( @@ -422,7 +425,7 @@ public void asyncClose() throws Exception { keysStateRequest.getStateKey()); isClosed = true; // No mutations necessary - if (!isCleared && pendingRemoves.isEmpty() && pendingAdds.isEmpty()) { + if (onlyBundleForKeys || (!isCleared && pendingRemoves.isEmpty() && pendingAdds.isEmpty())) { return; } @@ -445,7 +448,7 @@ private void startStateApiWrites() { } // Persist pending key-values - if (!pendingAdds.isEmpty() && !onlyBundleForKeys) { + if (!pendingAdds.isEmpty()) { for (KV> entry : pendingAdds.values()) { StateRequest request = createUserStateRequest(entry.getKey()); beamFnStateClient.handle( diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index 008ecb77cb37..ae0c204fb48a 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -90,6 +90,7 @@ public class OrderedListUserState { private boolean isClosed = false; private final boolean hasNoState; private final boolean onlyBundleForKeys; + public static class TimestampedValueCoder extends StructuredCoder> { private final Coder valueCoder; @@ -289,6 +290,11 @@ public void clear() { public void asyncClose() throws Exception { isClosed = true; + if (onlyBundleForKeys) { + pendingRemoves.clear(); + pendingAdds.clear(); + return; + } if (!pendingRemoves.isEmpty()) { for (Range r : pendingRemoves.asRanges()) { @@ -309,11 +315,6 @@ public void asyncClose() throws Exception { pendingRemoves.clear(); } - if (onlyBundleForKeys) { - pendingAdds.clear(); - return; - } - if (!pendingAdds.isEmpty()) { ByteStringOutputStream outStream = new ByteStringOutputStream(); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 9280c3191308..ee8087fabc65 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -86,14 +86,12 @@ public static CachingStateIterable readAllAndDecodeStartingFrom( Cache cache, BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk, - Coder valueCoder, - boolean hasNoState) { + Coder valueCoder) { return new CachingStateIterable<>( (Cache>) cache, beamFnStateClient, stateRequestForFirstChunk, - valueCoder, - hasNoState); + valueCoder); } /** @@ -330,19 +328,16 @@ public static Block fromValues( private final BeamFnStateClient beamFnStateClient; private final StateRequest stateRequestForFirstChunk; private final Coder valueCoder; - private final boolean hasNoState; public CachingStateIterable( Cache> cache, BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk, - Coder valueCoder, - boolean hasNoState) { + Coder valueCoder) { this.cache = cache; this.beamFnStateClient = beamFnStateClient; this.stateRequestForFirstChunk = stateRequestForFirstChunk; this.valueCoder = valueCoder; - this.hasNoState = hasNoState; } /** @@ -515,8 +510,7 @@ class CachingStateIterator implements PrefetchableIterator { public CachingStateIterator() { this.underlyingStateFetchingIterator = - new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk, - hasNoState); + new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk); this.dataStreamDecoder = new DataStreamDecoder<>(valueCoder, underlyingStateFetchingIterator); this.currentBlock = @@ -684,11 +678,10 @@ static class LazyBlockingStateFetchingIterator implements PrefetchableIterator Date: Mon, 9 Mar 2026 00:58:34 -0400 Subject: [PATCH 12/33] unit tests --- .../PTransformRunnerFactoryTestContext.java | 14 +- .../fn/harness/state/BagUserStateTest.java | 57 ++++-- .../harness/state/IterableSideInputTest.java | 12 +- .../harness/state/MultimapSideInputTest.java | 12 +- .../harness/state/MultimapUserStateTest.java | 168 +++++++++++++----- .../state/OrderedListUserStateTest.java | 76 ++++++-- 6 files changed, 259 insertions(+), 80 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java index 7b4387738a4c..c3aa19d820e0 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java @@ -64,6 +64,12 @@ public abstract class PTransformRunnerFactoryTestContext implements PTransformRunnerFactory.Context { + @Override + public abstract boolean getHasNoState(); + + @Override + public abstract boolean getOnlyBundleForKeys(); + /** Returns a builder for the specified PTransform id and PTransform definition. */ public static Builder builder(String pTransformId, RunnerApi.PTransform pTransform) { return new AutoValue_PTransformRunnerFactoryTestContext.Builder() @@ -149,7 +155,9 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { throw new UnsupportedOperationException("Unexpected call during test."); } }) - .runnerCapabilities(new HashSet<>()); + .runnerCapabilities(new HashSet<>()) + .hasNoState(false) + .onlyBundleForKeys(false); } /** A builder to create a context for tests. */ @@ -217,6 +225,10 @@ default Builder processBundleInstructionId(String value) { Builder timerApiServiceDescriptor(ApiServiceDescriptor value); + Builder hasNoState(boolean value); + + Builder onlyBundleForKeys(boolean value); + PTransformRunnerFactoryTestContext build(); } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java index 05a280245d13..a35531149fa8 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java @@ -49,7 +49,13 @@ public void testGet() throws Exception { StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList("A1", "A2", "A3"))); BagUserState userState = new BagUserState<>( - Caches.noop(), fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + Caches.noop(), + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + false, + false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); @@ -67,7 +73,8 @@ public void testGetCached() throws Exception { { // First user state populates the cache. BagUserState userState = - new BagUserState<>(cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + new BagUserState<>( + cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of(), false, false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -83,7 +90,9 @@ public void testGetCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -96,7 +105,13 @@ public void testAppend() throws Exception { new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList("A1"))); BagUserState userState = new BagUserState<>( - Caches.noop(), fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + Caches.noop(), + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + false, + false); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -118,7 +133,13 @@ public void testAppendBatchingLimit() throws Exception { new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList(a1))); BagUserState userState = new BagUserState<>( - Caches.noop(), fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + Caches.noop(), + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + false, + false); String a2 = Strings.repeat("A2", 2 * 1024 * 1024); userState.append(a2); String a3 = Strings.repeat("A3", 4 * 1024 * 1024); @@ -155,7 +176,8 @@ public void testAppendCached() throws Exception { { // First user state populates the cache. BagUserState userState = - new BagUserState<>(cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + new BagUserState<>( + cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of(), false, false); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -180,7 +202,9 @@ public void testAppendCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.append("A4"); Iterable stateBeforeA5 = userState.get(); assertArrayEquals( @@ -203,7 +227,13 @@ public void testClear() throws Exception { StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList("A1", "A2", "A3"))); BagUserState userState = new BagUserState<>( - Caches.noop(), fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + Caches.noop(), + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + false, + false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -228,7 +258,8 @@ public void testClearCached() throws Exception { { // First user state populates the cache. BagUserState userState = - new BagUserState<>(cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of()); + new BagUserState<>( + cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of(), false, false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -252,7 +283,9 @@ public void testClearCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals(new String[] {"A4"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); assertFalse(userState.get().iterator().hasNext()); @@ -273,7 +306,9 @@ public void testClearCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java index aa6d9115c1ee..cbeed603a96e 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java @@ -43,7 +43,12 @@ public void testGet() throws Exception { IterableSideInput iterableSideInput = new IterableSideInput<>( - Caches.noop(), fakeBeamFnStateClient, "instructionId", key(), StringUtf8Coder.of()); + Caches.noop(), + fakeBeamFnStateClient, + "instructionId", + key(), + StringUtf8Coder.of(), + false); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); @@ -61,7 +66,7 @@ public void testGetCached() throws Exception { // The first side input will populate the cache. IterableSideInput iterableSideInput = new IterableSideInput<>( - cache, fakeBeamFnStateClient, "instructionId", key(), StringUtf8Coder.of()); + cache, fakeBeamFnStateClient, "instructionId", key(), StringUtf8Coder.of(), false); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); @@ -77,7 +82,8 @@ public void testGetCached() throws Exception { }, "instructionId", key(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java index 88982bf1b6f6..c74ae9d0edca 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java @@ -70,7 +70,8 @@ public void testGetWithBulkRead() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true); + true, + false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); assertArrayEquals( @@ -96,7 +97,8 @@ public void testGet() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true); + true, + false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); assertArrayEquals( @@ -127,7 +129,8 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true); + true, + false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); @@ -151,7 +154,8 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true); + true, + false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); 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 679307321826..b921c8ba43ec 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 @@ -84,7 +84,9 @@ public void testNoPersistedValues() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.keys(), is(emptyIterable())); } @@ -104,7 +106,9 @@ public void testGet() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -132,7 +136,9 @@ public void testClear() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable initValues = userState.get(A1); userState.clear(); @@ -166,7 +172,9 @@ public void testKeys() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A2, "V1"); Iterable initKeys = userState.keys(); @@ -200,7 +208,9 @@ public void testEntries() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals(A1, userState.entries().iterator().next().getKey()); assertThat( @@ -275,7 +285,9 @@ public void testPut() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -302,7 +314,9 @@ public void testPutAfterRemove() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.remove(A0); userState.put(A0, "V2"); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); @@ -326,7 +340,9 @@ public void testPutAfterClear() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); userState.put(A0, "V2"); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); @@ -348,7 +364,9 @@ public void testRemoveBeforeClear() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.remove(A0); userState.clear(); userState.asyncClose(); @@ -366,7 +384,9 @@ public void testPutBeforeClear() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A0, "V0"); userState.put(A1, "V1"); Iterable values = userState.get(A1); // fakeClient call = 1 @@ -387,7 +407,9 @@ public void testPutBeforeRemove() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A0, "V0"); userState.put(A1, "V1"); Iterable values = userState.get(A1); // fakeClient call = 1 @@ -416,7 +438,9 @@ public void testRemove() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -444,7 +468,9 @@ public void testImmutableKeys() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable keys = userState.keys(); Iterator keysIterator = keys.iterator(); keysIterator.next(); @@ -467,7 +493,9 @@ public void testImmutableValues() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable values = userState.get(A1); assertThrows( UnsupportedOperationException.class, @@ -490,7 +518,9 @@ public void testClearAsyncClose() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); userState.asyncClose(); Map data = fakeClient.getData(); @@ -514,7 +544,9 @@ public void testNoopAsyncClose() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.asyncClose(); assertThrows(IllegalStateException.class, () -> userState.keys()); assertEquals(0, fakeClient.getCallCount()); @@ -538,7 +570,9 @@ public void testAsyncClose() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.remove(A0); userState.put(A1, "V3"); userState.put(A2, "V1"); @@ -567,7 +601,9 @@ public void testNullKeysAndValues() throws Exception { "instructionId", createMultimapKeyStateKey(), NullableCoder.of(ByteArrayCoder.of()), - NullableCoder.of(StringUtf8Coder.of())); + NullableCoder.of(StringUtf8Coder.of()), + false, + false); userState.put(null, null); userState.put(null, null); userState.put(null, "V1"); @@ -585,7 +621,9 @@ public void testNegativeCache() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(A1), String.class)); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(A1), String.class)); assertThat(fakeClient.getCallCount(), is(1)); @@ -607,7 +645,9 @@ public void testGetValuesPrefetch() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); PrefetchableIterable values = userState.get(A1); assertEquals(0, fakeClient.getCallCount()); @@ -633,7 +673,9 @@ public void testGetKeysPrefetch() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); PrefetchableIterable keys = userState.keys(); assertEquals(0, fakeClient.getCallCount()); @@ -659,7 +701,9 @@ public void testPutKeysPrefetch() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A2, "V3"); PrefetchableIterable keys = userState.keys(); @@ -686,7 +730,9 @@ public void testRemoveKeysPrefetch() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.remove(A1); userState.put(A1, "V3"); @@ -719,7 +765,9 @@ public void testEntriesPrefetched() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A1, "V4"); PrefetchableIterable> entries = userState.entries(); @@ -754,7 +802,9 @@ public void testClearPrefetch() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); userState.put(A2, "V3"); @@ -784,7 +834,9 @@ public void testAppendValuesPrefetch() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A1, "V3"); PrefetchableIterable values = userState.get(A1); @@ -808,7 +860,9 @@ public void testNoPersistedValuesCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.keys(), is(emptyIterable())); assertThat(userState.get(A1), is(emptyIterable())); } @@ -824,7 +878,9 @@ public void testNoPersistedValuesCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.keys(), is(emptyIterable())); assertThat(userState.get(A1), is(emptyIterable())); } @@ -849,7 +905,9 @@ public void testGetCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); @@ -867,7 +925,9 @@ public void testGetCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); @@ -894,7 +954,9 @@ public void testClearCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); assertThat(userState.keys(), is(emptyIterable())); @@ -913,7 +975,9 @@ public void testClearCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.keys(), is(emptyIterable())); userState.asyncClose(); @@ -939,7 +1003,9 @@ public void testKeysCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A2, "V1"); userState.put(A3, "V1"); @@ -961,7 +1027,9 @@ public void testKeysCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat( Iterables.toArray(userState.keys(), byte[].class), @@ -989,7 +1057,9 @@ public void testPutCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.put(A1, "V3"); userState.put(A2, "V1"); @@ -1010,7 +1080,9 @@ public void testPutCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( new String[] {"V1", "V2", "V3"}, Iterables.toArray(userState.get(A1), String.class)); @@ -1042,7 +1114,9 @@ public void testPutAfterRemoveCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.remove(A0); userState.put(A0, "V2"); @@ -1061,7 +1135,9 @@ public void testPutAfterRemoveCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); userState.asyncClose(); } @@ -1086,7 +1162,9 @@ public void testPutAfterClearCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); userState.put(A0, "V2"); userState.asyncClose(); @@ -1104,7 +1182,9 @@ public void testPutAfterClearCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); // Even though we never load assertArrayEquals(new byte[][] {A0}, Iterables.toArray(userState.keys(), byte[].class)); @@ -1131,7 +1211,9 @@ public void testRemoveCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); userState.remove(A1); @@ -1152,7 +1234,9 @@ public void testRemoveCached() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.get(A1), is(emptyIterable())); assertThat(userState.get(A2), is(emptyIterable())); assertThat(userState.keys(), is(emptyIterable())); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java index d831692b443b..c375f17a1c2e 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java @@ -76,7 +76,9 @@ public void testNoPersistedValues() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.read(), is(emptyIterable())); } @@ -92,7 +94,9 @@ public void testRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( asList(A1, B1).toArray(), Iterables.toArray(userState.read(), TimestampedValue.class)); @@ -116,7 +120,9 @@ public void testReadRange() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable> stateBeforeB2 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -171,7 +177,9 @@ public void testAdd() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); // add to an existing timestamp userState.add(B1); @@ -216,7 +224,9 @@ public void testClearRange() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable> initStateFrom2To3 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -284,7 +294,9 @@ public void testClear() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterable> stateBeforeClear = userState.read(); userState.clear(); @@ -316,7 +328,9 @@ public void testAddAndClearRange() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); // add to a non-existing timestamp, clear, and then add userState.add(A2); @@ -387,7 +401,9 @@ public void testAddAndClearRangeAfterClear() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); userState.clearRange(Instant.ofEpochMilli(0), Instant.ofEpochMilli(5)); @@ -427,7 +443,9 @@ public void testNoopAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.asyncClose(); } @@ -439,7 +457,9 @@ public void testNoopAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( asList(A1, A3, A4).toArray(), @@ -466,7 +486,9 @@ public void testAddAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.add(B1); userState.add(A2); @@ -479,7 +501,9 @@ public void testAddAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( asList(A1, B1, A2, A3, A4).toArray(), @@ -508,7 +532,9 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clearRange(Instant.ofEpochMilli(1), Instant.ofEpochMilli(3)); userState.clearRange(Instant.ofEpochMilli(4), Instant.ofEpochMilli(5)); @@ -521,7 +547,9 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -546,7 +574,9 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.add(B1); userState.add(A2); @@ -562,7 +592,9 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -591,7 +623,9 @@ public void testClearAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); userState.asyncClose(); @@ -603,7 +637,9 @@ public void testClearAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.read(), is(emptyIterable())); } @@ -630,7 +666,9 @@ public void testOperationsDuringNavigatingIterable() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); Iterator> iter = userState.read().iterator(); assertEquals(iter.next(), A1); From 208e3a2b15cd6dbaaf34e79491abaa2a2b2f266e Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 9 Mar 2026 01:31:59 -0400 Subject: [PATCH 13/33] fix --- .../beam/fn/harness/state/BagUserState.java | 24 ++++++++++--------- .../fn/harness/state/IterableSideInput.java | 18 +++++++++----- .../fn/harness/state/MultimapSideInput.java | 9 +++---- .../fn/harness/state/MultimapUserState.java | 7 +++--- .../harness/state/OrderedListUserState.java | 7 +++--- 5 files changed, 34 insertions(+), 31 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index 43f8d577a6b4..2e2b7d97c253 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -61,13 +61,13 @@ public class BagUserState { /** The cache must be namespaced for this state object accordingly. */ public BagUserState( - Cache cache, - BeamFnStateClient beamFnStateClient, - String instructionId, - StateKey stateKey, - Coder valueCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder valueCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasBagUserState(), "Expected BagUserState StateKey but received %s.", stateKey); this.cache = cache; @@ -79,8 +79,10 @@ public BagUserState( StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.oldValues = - StateFetchingIterators.readAllAndDecodeStartingFrom( - this.cache, beamFnStateClient, request, valueCoder, hasNoState); + hasNoState + ? PrefetchableIterables.emptyIterable() + : StateFetchingIterators.readAllAndDecodeStartingFrom( + this.cache, beamFnStateClient, request, valueCoder); this.newValues = new ArrayList<>(); } @@ -89,8 +91,8 @@ public PrefetchableIterable get() { !isClosed, "Bag user state is no longer usable because it is closed for %s", request.getStateKey()); - if (isCleared) { - // If we were cleared we should disregard old values. + if (isCleared || hasNoState) { + // If we were cleared or have no state we should disregard old values. return PrefetchableIterables.limit(Collections.unmodifiableList(newValues), newValues.size()); } else if (newValues.isEmpty()) { // If we have no new values then just return the old values. diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java index d7ce194d1164..ab2a97e838fc 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java @@ -23,6 +23,7 @@ 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.fn.stream.PrefetchableIterables; import org.apache.beam.sdk.transforms.Materializations.IterableView; /** @@ -46,12 +47,17 @@ public IterableSideInput( stateKey.hasIterableSideInput(), "Expected IterableSideInput StateKey but received %s.", stateKey); - this.values = hasNoState ? PrefetchableIterables.emptyIterable() : - StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, - beamFnStateClient, - StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(), - valueCoder); + this.values = + hasNoState + ? PrefetchableIterables.emptyIterable() + : StateFetchingIterators.readAllAndDecodeStartingFrom( + cache, + beamFnStateClient, + StateRequest.newBuilder() + .setInstructionId(instructionId) + .setStateKey(stateKey) + .build(), + valueCoder); } @Override diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java index 46ce7db17857..25e7868480ea 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java @@ -32,6 +32,7 @@ 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.PrefetchableIterables; import org.apache.beam.sdk.transforms.Materializations.MultimapView; import org.apache.beam.sdk.util.ByteStringOutputStream; import org.apache.beam.sdk.values.KV; @@ -129,8 +130,7 @@ public Iterable get(K k) { Caches.noop(), beamFnStateClient, bulkReadRequest, - KvCoder.of(keyCoder, IterableCoder.of(valueCoder)), - hasNoState) + KvCoder.of(keyCoder, IterableCoder.of(valueCoder))) .iterator(); while (bulkRead.size() < BULK_READ_SIZE && entries.hasNext()) { KV> entry = entries.next(); @@ -179,10 +179,7 @@ public Iterable get(K k) { StateRequest request = keysRequest.toBuilder().setStateKey(stateKey).build(); return StateFetchingIterators.readAllAndDecodeStartingFrom( - Caches.subCache(cache, "ValuesForKey", encodedKey), - beamFnStateClient, - request, - valueCoder); + Caches.subCache(cache, "ValuesForKey", encodedKey), beamFnStateClient, request, valueCoder); } private ByteString encodeKey(K k) { 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 b44e403ab7e2..e0f638adeb36 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 @@ -108,7 +108,7 @@ public MultimapUserState( hasNoState ? PrefetchableIterables.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, beamFnStateClient, keysStateRequest, mapKeyCoder, hasNoState); + cache, beamFnStateClient, keysStateRequest, mapKeyCoder); StateRequest.Builder userStateRequestBuilder = StateRequest.newBuilder(); userStateRequestBuilder @@ -171,7 +171,7 @@ public PrefetchableIterable get(K key) { ? PrefetchableIterables.fromArray() : PrefetchableIterables.limit( pendingAddValues.getValue(), pendingAddValues.getValue().size()); - if (isCleared || pendingRemoves.containsKey(structuralKey)) { + if (isCleared || hasNoState || pendingRemoves.containsKey(structuralKey)) { return pendingValues; } @@ -552,8 +552,7 @@ private CachingStateIterable getPersistedValues(Object structuralKey, K key) request.getStateKey().getMultimapUserState().getMapKey()), beamFnStateClient, request, - valueCoder, - hasNoState)); + valueCoder)); }) .getValue(); } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index ae0c204fb48a..49c25ee449a5 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -90,7 +90,7 @@ public class OrderedListUserState { private boolean isClosed = false; private final boolean hasNoState; private final boolean onlyBundleForKeys; - + public static class TimestampedValueCoder extends StructuredCoder> { private final Coder valueCoder; @@ -209,7 +209,7 @@ public Iterable> readRange(Instant minTimestamp, Instant lim } Iterable> valuesInRange = Iterables.concat(pendingAddsInRange); - if (!isCleared) { + if (!isCleared && !hasNoState) { StateRequest.Builder getRequestBuilder = this.requestTemplate.toBuilder(); getRequestBuilder .getStateKeyBuilder() @@ -224,8 +224,7 @@ public Iterable> readRange(Instant minTimestamp, Instant lim Caches.noop(), this.beamFnStateClient, getRequestBuilder.build(), - this.timestampedValueCoder, - hasNoState); + this.timestampedValueCoder); // Make a snapshot of the current pendingRemoves and use them to filter persistent values. // The values of pendingRemoves are copied, so that they will still be accessible in From 657ac0871811c3ece7322533257bdb2fd7387b94 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 9 Mar 2026 02:25:48 -0400 Subject: [PATCH 14/33] fix --- .../beam/fn/harness/state/BagUserState.java | 3 +- .../fn/harness/state/MultimapUserState.java | 8 ++- .../harness/state/StateFetchingIterators.java | 21 ++++-- .../fn/harness/state/BagUserStateTest.java | 56 +++++++++++++++ .../state/OrderedListUserStateTest.java | 70 +++++++++++++++++++ 5 files changed, 148 insertions(+), 10 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index 2e2b7d97c253..23fb8aa99fd8 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -80,7 +80,8 @@ public BagUserState( this.oldValues = hasNoState - ? PrefetchableIterables.emptyIterable() + ? StateFetchingIterators.emptyCachingStateIterable( + beamFnStateClient, request, valueCoder) : StateFetchingIterators.readAllAndDecodeStartingFrom( this.cache, beamFnStateClient, request, valueCoder); this.newValues = new ArrayList<>(); 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 e0f638adeb36..1fd06935d2c3 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 @@ -106,7 +106,8 @@ public MultimapUserState( StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = hasNoState - ? PrefetchableIterables.emptyIterable() + ? StateFetchingIterators.emptyCachingStateIterable( + beamFnStateClient, keysStateRequest, mapKeyCoder) : StateFetchingIterators.readAllAndDecodeStartingFrom( cache, beamFnStateClient, keysStateRequest, mapKeyCoder); @@ -133,7 +134,10 @@ public MultimapUserState( this.entriesStateRequest = entriesStateRequestBuilder.build(); this.persistedEntries = hasNoState - ? PrefetchableIterables.emptyIterable() + ? StateFetchingIterators.emptyCachingStateIterable( + beamFnStateClient, + entriesStateRequest, + KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder))) : StateFetchingIterators.readAllAndDecodeStartingFrom( Caches.subCache(this.cache, "AllEntries"), beamFnStateClient, diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index ee8087fabc65..ed2ade1809c1 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -94,6 +94,20 @@ public static CachingStateIterable readAllAndDecodeStartingFrom( valueCoder); } + public static CachingStateIterable emptyCachingStateIterable( + BeamFnStateClient beamFnStateClient, + StateRequest stateRequestForFirstChunk, + Coder valueCoder) { + CachingStateIterable result = + readAllAndDecodeStartingFrom( + org.apache.beam.fn.harness.Caches.noop(), + beamFnStateClient, + stateRequestForFirstChunk, + valueCoder); + result.clearAndAppend(java.util.Collections.emptyList()); + return result; + } + /** * This adapter handles using the continuation token to provide iteration over all the elements * returned by the Beam Fn State API using the supplied state client, state request for the first @@ -677,13 +691,6 @@ static class LazyBlockingStateFetchingIterator implements PrefetchableIterator userState = + new BagUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + true, /* hasNoState */ + false /* onlyBundleForKeys */); + + // Iterating should be empty since hasNoState is true + assertFalse(userState.get().iterator().hasNext()); + + // We can append new values + userState.append("A4"); + assertArrayEquals(new String[] {"A4"}, Iterables.toArray(userState.get(), String.class)); + userState.asyncClose(); + + // The appended string will be persisted, appended to the old values (since we didn't clear) + assertEquals(encode("A1", "A2", "A3", "A4"), fakeClient.getData().get(key("A"))); + } + + @Test + public void testOnlyBundleForKeys() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList("A1"))); + BagUserState userState = + new BagUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + false /* hasNoState */, + true /* onlyBundleForKeys */); + + // Reads still work + assertArrayEquals(new String[] {"A1"}, Iterables.toArray(userState.get(), String.class)); + + // Append something + userState.append("A2"); + + // We can observe the appended value locally + assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(userState.get(), String.class)); + + userState.asyncClose(); + + // But when we close, the data is NOT persisted because onlyBundleForKeys = true + assertEquals(encode("A1"), fakeClient.getData().get(key("A"))); + } } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java index c375f17a1c2e..2f0c31cef7d0 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java @@ -23,6 +23,7 @@ import static org.hamcrest.core.Is.is; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; import java.io.IOException; @@ -688,6 +689,75 @@ public void testOperationsDuringNavigatingIterable() throws Exception { assertEquals(iter.next(), A4); } + @Test + public void testHasNoState() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + timestampedValueCoder, + ImmutableMap.of(createOrderedListStateKey("A", 1), Collections.singletonList(A1))); + OrderedListUserState userState = + new OrderedListUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createOrderedListStateKey("A"), + StringUtf8Coder.of(), + true, /* hasNoState */ + false /* onlyBundleForKeys */); + + // Iterating should be empty since hasNoState is true + assertThat(userState.read(), is(emptyIterable())); + + // We can add new values + userState.add(A2); + assertArrayEquals( + Collections.singletonList(A2).toArray(), + Iterables.toArray(userState.read(), TimestampedValue.class)); + userState.asyncClose(); + + // The added value will be persisted, appended to the old values (since we didn't clear) + ByteStringOutputStream out = new ByteStringOutputStream(); + timestampedValueCoder.encode(A2, out); + assertEquals(out.toByteString(), fakeClient.getData().get(createOrderedListStateKey("A", 2))); + } + + @Test + public void testOnlyBundleForKeys() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + timestampedValueCoder, + ImmutableMap.of(createOrderedListStateKey("A", 1), Collections.singletonList(A1))); + OrderedListUserState userState = + new OrderedListUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createOrderedListStateKey("A"), + StringUtf8Coder.of(), + false /* hasNoState */, + true /* onlyBundleForKeys */); + + // Reads still work + assertArrayEquals( + Collections.singletonList(A1).toArray(), + Iterables.toArray(userState.read(), TimestampedValue.class)); + + // Add something + userState.add(A2); + + // We can observe the added value locally + assertArrayEquals( + asList(A1, A2).toArray(), Iterables.toArray(userState.read(), TimestampedValue.class)); + + userState.asyncClose(); + + // But when we close, the data is NOT persisted because onlyBundleForKeys = true + ByteStringOutputStream out = new ByteStringOutputStream(); + timestampedValueCoder.encode(A1, out); + assertEquals(out.toByteString(), fakeClient.getData().get(createOrderedListStateKey("A", 1))); + assertFalse(fakeClient.getData().containsKey(createOrderedListStateKey("A", 2))); + } + private ByteString encode(String... values) throws IOException { ByteStringOutputStream out = new ByteStringOutputStream(); for (String value : values) { From 1bb60df561386191f3c01c115e37b16b973e1206 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 18 Mar 2026 01:46:05 -0400 Subject: [PATCH 15/33] hasNoState, onlyBundleForKeys boolean -> Supplier --- .../fn/harness/PTransformRunnerFactory.java | 4 +- .../harness/control/ProcessBundleHandler.java | 12 +- .../beam/fn/harness/state/BagUserState.java | 15 +- .../fn/harness/state/FnApiStateAccessor.java | 16 +- .../fn/harness/state/IterableSideInput.java | 5 +- .../fn/harness/state/MultimapSideInput.java | 9 +- .../fn/harness/state/MultimapUserState.java | 22 +-- .../harness/state/OrderedListUserState.java | 13 +- .../PTransformRunnerFactoryTestContext.java | 12 +- .../fn/harness/state/BagUserStateTest.java | 64 ++++--- .../harness/state/IterableSideInputTest.java | 11 +- .../harness/state/MultimapSideInputTest.java | 16 +- .../harness/state/MultimapUserStateTest.java | 168 +++++++++--------- .../state/OrderedListUserStateTest.java | 84 ++++----- 14 files changed, 240 insertions(+), 211 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java index 8752408770ac..744e834b8942 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java @@ -165,13 +165,13 @@ void addIncomingTimerEndpoint( * Returns true if the runner has no state for the keys in the ProcessBundleRequest. If true, * the SDK can begin stateful processing with an initial empty state. */ - boolean getHasNoState(); + Supplier getHasNoState(); /** * Returns true if the runner will never process another bundle for the keys it contains. * Therefore, the generated state need not be included in the bundle commit. */ - boolean getOnlyBundleForKeys(); + Supplier getOnlyBundleForKeys(); } /** diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index e4cffeed0b53..7146419429d9 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -247,8 +247,8 @@ private void addRunnerAndConsumersForPTransformRecursively( Collection> channelRoots, Map outboundAggregatorMap, Set runnerCapabilities, - boolean hasNoState, - boolean onlyBundleForKeys) + Supplier hasNoState, + Supplier onlyBundleForKeys) throws IOException { // Recursively ensure that all consumers of the output PCollection have been created. @@ -494,12 +494,12 @@ public BundleFinalizer getBundleFinalizer() { } @Override - public boolean getHasNoState() { + public Supplier getHasNoState() { return hasNoState; } @Override - public boolean getOnlyBundleForKeys() { + public Supplier getOnlyBundleForKeys() { return onlyBundleForKeys; } }); @@ -928,8 +928,8 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { bundleProcessor.getChannelRoots(), bundleProcessor.getOutboundAggregators(), bundleProcessor.getRunnerCapabilities(), - processBundleRequest.getHasNoState(), - processBundleRequest.getOnlyBundleForKeys()); + processBundleRequest::getHasNoState, + processBundleRequest::getOnlyBundleForKeys); } bundleProcessor.finish(); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index 23fb8aa99fd8..c8ab95aa6945 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -23,6 +23,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.state.StateFetchingIterators.CachingStateIterable; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateAppendRequest; @@ -54,8 +55,8 @@ public class BagUserState { private List newValues; private boolean isCleared; private boolean isClosed; - private final boolean hasNoState; - private final boolean onlyBundleForKeys; + private final Supplier hasNoState; + private final Supplier onlyBundleForKeys; static final int BAG_APPEND_BATCHING_LIMIT = 10 * 1024 * 1024; @@ -66,8 +67,8 @@ public BagUserState( String instructionId, StateKey stateKey, Coder valueCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Supplier hasNoState, + Supplier onlyBundleForKeys) { checkArgument( stateKey.hasBagUserState(), "Expected BagUserState StateKey but received %s.", stateKey); this.cache = cache; @@ -79,7 +80,7 @@ public BagUserState( StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.oldValues = - hasNoState + hasNoState.get() ? StateFetchingIterators.emptyCachingStateIterable( beamFnStateClient, request, valueCoder) : StateFetchingIterators.readAllAndDecodeStartingFrom( @@ -92,7 +93,7 @@ public PrefetchableIterable get() { !isClosed, "Bag user state is no longer usable because it is closed for %s", request.getStateKey()); - if (isCleared || hasNoState) { + if (isCleared || hasNoState.get()) { // If we were cleared or have no state we should disregard old values. return PrefetchableIterables.limit(Collections.unmodifiableList(newValues), newValues.size()); } else if (newValues.isEmpty()) { @@ -129,7 +130,7 @@ public void asyncClose() throws Exception { "Bag user state is no longer usable because it is closed for %s", request.getStateKey()); isClosed = true; - if (onlyBundleForKeys || (!isCleared && newValues.isEmpty())) { + if (onlyBundleForKeys.get() || (!isCleared && newValues.isEmpty())) { return; } if (isCleared) { 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 70b6f5284ad7..c81042618b38 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 @@ -114,8 +114,8 @@ public static class Factory { private final Map, SideInputSpec> sideInputSpecMap; private final Coder keyCoder; private final Coder windowCoder; - private final boolean hasNoState; - private final boolean onlyBundleForKeys; + private final Supplier hasNoState; + private final Supplier onlyBundleForKeys; public Factory( PipelineOptions pipelineOptions, @@ -129,8 +129,8 @@ public Factory( BeamFnStateClient beamFnStateClient, Coder keyCoder, Coder windowCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Supplier hasNoState, + Supplier onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.ptransformId = ptransformId; @@ -262,8 +262,8 @@ public FnApiStateAccessor create() { private final Collection stateFinalizers; private final Coder keyCoder; private final Coder windowCoder; - private final boolean hasNoState; - private final boolean onlyBundleForKeys; + private final Supplier hasNoState; + private final Supplier onlyBundleForKeys; private @Nullable Supplier currentWindowSupplier; private @Nullable Supplier encodedCurrentKeySupplier; @@ -281,8 +281,8 @@ public FnApiStateAccessor( BeamFnStateClient beamFnStateClient, Coder keyCoder, Coder windowCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Supplier hasNoState, + Supplier onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.stateKeyObjectCache = Maps.newHashMap(); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java index ab2a97e838fc..4fab87c0b3aa 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java @@ -19,6 +19,7 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest; @@ -42,13 +43,13 @@ public IterableSideInput( String instructionId, StateKey stateKey, Coder valueCoder, - boolean hasNoState) { + Supplier hasNoState) { checkArgument( stateKey.hasIterableSideInput(), "Expected IterableSideInput StateKey but received %s.", stateKey); this.values = - hasNoState + hasNoState.get() ? PrefetchableIterables.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( cache, diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java index 25e7868480ea..4610fdd797c3 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java @@ -25,6 +25,7 @@ import java.util.Iterator; import java.util.Map; import java.util.function.Function; +import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; @@ -56,7 +57,7 @@ public class MultimapSideInput implements MultimapView { private final Coder valueCoder; private volatile Function> bulkReadResult; private final boolean useBulkRead; - private final boolean hasNoState; + private final Supplier hasNoState; public MultimapSideInput( Cache cache, @@ -66,7 +67,7 @@ public MultimapSideInput( Coder keyCoder, Coder valueCoder, boolean useBulkRead, - boolean hasNoState) { + Supplier hasNoState) { checkArgument( stateKey.hasMultimapKeysSideInput(), "Expected MultimapKeysSideInput StateKey but received %s.", @@ -83,7 +84,7 @@ public MultimapSideInput( @Override public Iterable get() { - return this.hasNoState + return this.hasNoState.get() ? PrefetchableIterables.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( cache, beamFnStateClient, keysRequest, keyCoder); @@ -91,7 +92,7 @@ public Iterable get() { @Override public Iterable get(K k) { - if (this.hasNoState) { + if (this.hasNoState.get()) { return PrefetchableIterables.emptyIterable(); } 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 1fd06935d2c3..9a31b01c4df3 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 @@ -31,6 +31,7 @@ import java.util.NoSuchElementException; import java.util.Objects; import java.util.Set; +import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.fn.harness.state.StateFetchingIterators.CachingStateIterable; @@ -72,8 +73,8 @@ public class MultimapUserState { private boolean isClosed; private boolean isCleared; - private final boolean hasNoState; - private final boolean onlyBundleForKeys; + private final Supplier hasNoState; + private final Supplier onlyBundleForKeys; // Pending updates to persistent storage private HashMap pendingRemoves = Maps.newHashMap(); private HashMap>> pendingAdds = Maps.newHashMap(); @@ -87,8 +88,8 @@ public MultimapUserState( StateKey stateKey, Coder mapKeyCoder, Coder valueCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Supplier hasNoState, + Supplier onlyBundleForKeys) { checkArgument( stateKey.hasMultimapKeysUserState(), "Expected MultimapKeysUserState StateKey but received %s.", @@ -105,7 +106,7 @@ public MultimapUserState( this.keysStateRequest = StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = - hasNoState + hasNoState.get() ? StateFetchingIterators.emptyCachingStateIterable( beamFnStateClient, keysStateRequest, mapKeyCoder) : StateFetchingIterators.readAllAndDecodeStartingFrom( @@ -133,7 +134,7 @@ public MultimapUserState( .setKey(stateKey.getMultimapKeysUserState().getKey()); this.entriesStateRequest = entriesStateRequestBuilder.build(); this.persistedEntries = - hasNoState + hasNoState.get() ? StateFetchingIterators.emptyCachingStateIterable( beamFnStateClient, entriesStateRequest, @@ -175,7 +176,7 @@ public PrefetchableIterable get(K key) { ? PrefetchableIterables.fromArray() : PrefetchableIterables.limit( pendingAddValues.getValue(), pendingAddValues.getValue().size()); - if (isCleared || hasNoState || pendingRemoves.containsKey(structuralKey)) { + if (isCleared || hasNoState.get() || pendingRemoves.containsKey(structuralKey)) { return pendingValues; } @@ -193,7 +194,7 @@ public PrefetchableIterable keys() { !isClosed, "Multimap user state is no longer usable because it is closed for %s", keysStateRequest.getStateKey()); - if (isCleared || hasNoState) { + if (isCleared || hasNoState.get()) { List keys = new ArrayList<>(pendingAdds.size()); for (Map.Entry>> entry : pendingAdds.entrySet()) { keys.add(entry.getValue().getKey()); @@ -291,7 +292,7 @@ public PrefetchableIterable> entries() { entry.getKey(), KV.of(entry.getValue().getKey(), new ArrayList<>(entry.getValue().getValue()))); } - if (isCleared || hasNoState) { + if (isCleared || hasNoState.get()) { return PrefetchableIterables.maybePrefetchable( Iterables.concat( Iterables.transform( @@ -429,7 +430,8 @@ public void asyncClose() throws Exception { keysStateRequest.getStateKey()); isClosed = true; // No mutations necessary - if (onlyBundleForKeys || (!isCleared && pendingRemoves.isEmpty() && pendingAdds.isEmpty())) { + if (onlyBundleForKeys.get() + || (!isCleared && pendingRemoves.isEmpty() && pendingAdds.isEmpty())) { return; } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index 49c25ee449a5..8c002a5ecc4c 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -33,6 +33,7 @@ import java.util.Map.Entry; import java.util.NavigableMap; import java.util.concurrent.CompletableFuture; +import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.fn.harness.state.StateFetchingIterators.CachingStateIterable; @@ -88,8 +89,8 @@ public class OrderedListUserState { private boolean isCleared = false; private boolean isClosed = false; - private final boolean hasNoState; - private final boolean onlyBundleForKeys; + private final Supplier hasNoState; + private final Supplier onlyBundleForKeys; public static class TimestampedValueCoder extends StructuredCoder> { @@ -164,8 +165,8 @@ public OrderedListUserState( String instructionId, StateKey stateKey, Coder valueCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Supplier hasNoState, + Supplier onlyBundleForKeys) { checkArgument( stateKey.hasOrderedListUserState(), "Expected OrderedListUserState StateKey but received %s.", @@ -209,7 +210,7 @@ public Iterable> readRange(Instant minTimestamp, Instant lim } Iterable> valuesInRange = Iterables.concat(pendingAddsInRange); - if (!isCleared && !hasNoState) { + if (!isCleared && !hasNoState.get()) { StateRequest.Builder getRequestBuilder = this.requestTemplate.toBuilder(); getRequestBuilder .getStateKeyBuilder() @@ -289,7 +290,7 @@ public void clear() { public void asyncClose() throws Exception { isClosed = true; - if (onlyBundleForKeys) { + if (onlyBundleForKeys.get()) { pendingRemoves.clear(); pendingAdds.clear(); return; diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java index c3aa19d820e0..a5f335bd56eb 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java @@ -65,10 +65,10 @@ public abstract class PTransformRunnerFactoryTestContext implements PTransformRunnerFactory.Context { @Override - public abstract boolean getHasNoState(); + public abstract Supplier getHasNoState(); @Override - public abstract boolean getOnlyBundleForKeys(); + public abstract Supplier getOnlyBundleForKeys(); /** Returns a builder for the specified PTransform id and PTransform definition. */ public static Builder builder(String pTransformId, RunnerApi.PTransform pTransform) { @@ -156,8 +156,8 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { } }) .runnerCapabilities(new HashSet<>()) - .hasNoState(false) - .onlyBundleForKeys(false); + .hasNoState(() -> false) + .onlyBundleForKeys(() -> false); } /** A builder to create a context for tests. */ @@ -225,9 +225,9 @@ default Builder processBundleInstructionId(String value) { Builder timerApiServiceDescriptor(ApiServiceDescriptor value); - Builder hasNoState(boolean value); + Builder hasNoState(Supplier value); - Builder onlyBundleForKeys(boolean value); + Builder onlyBundleForKeys(Supplier value); PTransformRunnerFactoryTestContext build(); } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java index c49674df52f3..f4311fb65d3d 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java @@ -54,8 +54,8 @@ public void testGet() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); @@ -74,7 +74,13 @@ public void testGetCached() throws Exception { // First user state populates the cache. BagUserState userState = new BagUserState<>( - cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of(), false, false); + cache, + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + () -> false, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -91,8 +97,8 @@ public void testGetCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -110,8 +116,8 @@ public void testAppend() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -138,8 +144,8 @@ public void testAppendBatchingLimit() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); String a2 = Strings.repeat("A2", 2 * 1024 * 1024); userState.append(a2); String a3 = Strings.repeat("A3", 4 * 1024 * 1024); @@ -177,7 +183,13 @@ public void testAppendCached() throws Exception { // First user state populates the cache. BagUserState userState = new BagUserState<>( - cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of(), false, false); + cache, + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + () -> false, + () -> false); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -203,8 +215,8 @@ public void testAppendCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.append("A4"); Iterable stateBeforeA5 = userState.get(); assertArrayEquals( @@ -232,8 +244,8 @@ public void testClear() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -259,7 +271,13 @@ public void testClearCached() throws Exception { // First user state populates the cache. BagUserState userState = new BagUserState<>( - cache, fakeClient, "instructionId", key("A"), StringUtf8Coder.of(), false, false); + cache, + fakeClient, + "instructionId", + key("A"), + StringUtf8Coder.of(), + () -> false, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -284,8 +302,8 @@ public void testClearCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals(new String[] {"A4"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); assertFalse(userState.get().iterator().hasNext()); @@ -307,8 +325,8 @@ public void testClearCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); } @@ -346,8 +364,8 @@ public void testHasNoState() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - true, /* hasNoState */ - false /* onlyBundleForKeys */); + () -> true, /* hasNoState */ + () -> false /* onlyBundleForKeys */); // Iterating should be empty since hasNoState is true assertFalse(userState.get().iterator().hasNext()); @@ -372,8 +390,8 @@ public void testOnlyBundleForKeys() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - false /* hasNoState */, - true /* onlyBundleForKeys */); + () -> false /* hasNoState */, + () -> true /* onlyBundleForKeys */); // Reads still work assertArrayEquals(new String[] {"A1"}, Iterables.toArray(userState.get(), String.class)); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java index cbeed603a96e..cd630cbf02f1 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java @@ -48,7 +48,7 @@ public void testGet() throws Exception { "instructionId", key(), StringUtf8Coder.of(), - false); + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); @@ -66,7 +66,12 @@ public void testGetCached() throws Exception { // The first side input will populate the cache. IterableSideInput iterableSideInput = new IterableSideInput<>( - cache, fakeBeamFnStateClient, "instructionId", key(), StringUtf8Coder.of(), false); + cache, + fakeBeamFnStateClient, + "instructionId", + key(), + StringUtf8Coder.of(), + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); @@ -83,7 +88,7 @@ public void testGetCached() throws Exception { "instructionId", key(), StringUtf8Coder.of(), - false); + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java index c74ae9d0edca..595aa71a6666 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java @@ -70,8 +70,8 @@ public void testGetWithBulkRead() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - false); + () -> true, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); assertArrayEquals( @@ -97,8 +97,8 @@ public void testGet() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - false); + () -> true, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); assertArrayEquals( @@ -129,8 +129,8 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - false); + () -> true, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); @@ -154,8 +154,8 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - false); + () -> true, + () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); 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 b921c8ba43ec..5675e9529d11 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 @@ -85,8 +85,8 @@ public void testNoPersistedValues() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat(userState.keys(), is(emptyIterable())); } @@ -107,8 +107,8 @@ public void testGet() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -137,8 +137,8 @@ public void testClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable initValues = userState.get(A1); userState.clear(); @@ -173,8 +173,8 @@ public void testKeys() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A2, "V1"); Iterable initKeys = userState.keys(); @@ -209,8 +209,8 @@ public void testEntries() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals(A1, userState.entries().iterator().next().getKey()); assertThat( @@ -286,8 +286,8 @@ public void testPut() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -315,8 +315,8 @@ public void testPutAfterRemove() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.remove(A0); userState.put(A0, "V2"); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); @@ -341,8 +341,8 @@ public void testPutAfterClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clear(); userState.put(A0, "V2"); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); @@ -365,8 +365,8 @@ public void testRemoveBeforeClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.remove(A0); userState.clear(); userState.asyncClose(); @@ -385,8 +385,8 @@ public void testPutBeforeClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A0, "V0"); userState.put(A1, "V1"); Iterable values = userState.get(A1); // fakeClient call = 1 @@ -408,8 +408,8 @@ public void testPutBeforeRemove() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A0, "V0"); userState.put(A1, "V1"); Iterable values = userState.get(A1); // fakeClient call = 1 @@ -439,8 +439,8 @@ public void testRemove() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -469,8 +469,8 @@ public void testImmutableKeys() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable keys = userState.keys(); Iterator keysIterator = keys.iterator(); keysIterator.next(); @@ -494,8 +494,8 @@ public void testImmutableValues() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable values = userState.get(A1); assertThrows( UnsupportedOperationException.class, @@ -519,8 +519,8 @@ public void testClearAsyncClose() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clear(); userState.asyncClose(); Map data = fakeClient.getData(); @@ -545,8 +545,8 @@ public void testNoopAsyncClose() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.asyncClose(); assertThrows(IllegalStateException.class, () -> userState.keys()); assertEquals(0, fakeClient.getCallCount()); @@ -571,8 +571,8 @@ public void testAsyncClose() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.remove(A0); userState.put(A1, "V3"); userState.put(A2, "V1"); @@ -602,8 +602,8 @@ public void testNullKeysAndValues() throws Exception { createMultimapKeyStateKey(), NullableCoder.of(ByteArrayCoder.of()), NullableCoder.of(StringUtf8Coder.of()), - false, - false); + () -> false, + () -> false); userState.put(null, null); userState.put(null, null); userState.put(null, "V1"); @@ -622,8 +622,8 @@ public void testNegativeCache() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(A1), String.class)); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(A1), String.class)); assertThat(fakeClient.getCallCount(), is(1)); @@ -646,8 +646,8 @@ public void testGetValuesPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); PrefetchableIterable values = userState.get(A1); assertEquals(0, fakeClient.getCallCount()); @@ -674,8 +674,8 @@ public void testGetKeysPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); PrefetchableIterable keys = userState.keys(); assertEquals(0, fakeClient.getCallCount()); @@ -702,8 +702,8 @@ public void testPutKeysPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A2, "V3"); PrefetchableIterable keys = userState.keys(); @@ -731,8 +731,8 @@ public void testRemoveKeysPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.remove(A1); userState.put(A1, "V3"); @@ -766,8 +766,8 @@ public void testEntriesPrefetched() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A1, "V4"); PrefetchableIterable> entries = userState.entries(); @@ -803,8 +803,8 @@ public void testClearPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clear(); userState.put(A2, "V3"); @@ -835,8 +835,8 @@ public void testAppendValuesPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A1, "V3"); PrefetchableIterable values = userState.get(A1); @@ -861,8 +861,8 @@ public void testNoPersistedValuesCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat(userState.keys(), is(emptyIterable())); assertThat(userState.get(A1), is(emptyIterable())); } @@ -879,8 +879,8 @@ public void testNoPersistedValuesCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat(userState.keys(), is(emptyIterable())); assertThat(userState.get(A1), is(emptyIterable())); } @@ -906,8 +906,8 @@ public void testGetCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); @@ -926,8 +926,8 @@ public void testGetCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); @@ -955,8 +955,8 @@ public void testClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clear(); assertThat(userState.keys(), is(emptyIterable())); @@ -976,8 +976,8 @@ public void testClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat(userState.keys(), is(emptyIterable())); userState.asyncClose(); @@ -1004,8 +1004,8 @@ public void testKeysCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A2, "V1"); userState.put(A3, "V1"); @@ -1028,8 +1028,8 @@ public void testKeysCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat( Iterables.toArray(userState.keys(), byte[].class), @@ -1058,8 +1058,8 @@ public void testPutCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.put(A1, "V3"); userState.put(A2, "V1"); @@ -1081,8 +1081,8 @@ public void testPutCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( new String[] {"V1", "V2", "V3"}, Iterables.toArray(userState.get(A1), String.class)); @@ -1115,8 +1115,8 @@ public void testPutAfterRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.remove(A0); userState.put(A0, "V2"); @@ -1136,8 +1136,8 @@ public void testPutAfterRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); userState.asyncClose(); } @@ -1163,8 +1163,8 @@ public void testPutAfterClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clear(); userState.put(A0, "V2"); userState.asyncClose(); @@ -1183,8 +1183,8 @@ public void testPutAfterClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); // Even though we never load assertArrayEquals(new byte[][] {A0}, Iterables.toArray(userState.keys(), byte[].class)); @@ -1212,8 +1212,8 @@ public void testRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); userState.remove(A1); @@ -1235,8 +1235,8 @@ public void testRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat(userState.get(A1), is(emptyIterable())); assertThat(userState.get(A2), is(emptyIterable())); assertThat(userState.keys(), is(emptyIterable())); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java index 2f0c31cef7d0..23b1c674d502 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java @@ -78,8 +78,8 @@ public void testNoPersistedValues() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat(userState.read(), is(emptyIterable())); } @@ -96,8 +96,8 @@ public void testRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( asList(A1, B1).toArray(), Iterables.toArray(userState.read(), TimestampedValue.class)); @@ -122,8 +122,8 @@ public void testReadRange() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable> stateBeforeB2 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -179,8 +179,8 @@ public void testAdd() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); // add to an existing timestamp userState.add(B1); @@ -226,8 +226,8 @@ public void testClearRange() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable> initStateFrom2To3 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -296,8 +296,8 @@ public void testClear() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterable> stateBeforeClear = userState.read(); userState.clear(); @@ -330,8 +330,8 @@ public void testAddAndClearRange() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); // add to a non-existing timestamp, clear, and then add userState.add(A2); @@ -403,8 +403,8 @@ public void testAddAndClearRangeAfterClear() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clear(); userState.clearRange(Instant.ofEpochMilli(0), Instant.ofEpochMilli(5)); @@ -445,8 +445,8 @@ public void testNoopAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.asyncClose(); } @@ -459,8 +459,8 @@ public void testNoopAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( asList(A1, A3, A4).toArray(), @@ -488,8 +488,8 @@ public void testAddAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.add(B1); userState.add(A2); @@ -503,8 +503,8 @@ public void testAddAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( asList(A1, B1, A2, A3, A4).toArray(), @@ -534,8 +534,8 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clearRange(Instant.ofEpochMilli(1), Instant.ofEpochMilli(3)); userState.clearRange(Instant.ofEpochMilli(4), Instant.ofEpochMilli(5)); @@ -549,8 +549,8 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -576,8 +576,8 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.add(B1); userState.add(A2); @@ -594,8 +594,8 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -625,8 +625,8 @@ public void testClearAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); userState.clear(); userState.asyncClose(); @@ -639,8 +639,8 @@ public void testClearAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); assertThat(userState.read(), is(emptyIterable())); } @@ -668,8 +668,8 @@ public void testOperationsDuringNavigatingIterable() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false, - false); + () -> false, + () -> false); Iterator> iter = userState.read().iterator(); assertEquals(iter.next(), A1); @@ -702,8 +702,8 @@ public void testHasNoState() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - true, /* hasNoState */ - false /* onlyBundleForKeys */); + () -> true, /* hasNoState */ + () -> false /* onlyBundleForKeys */); // Iterating should be empty since hasNoState is true assertThat(userState.read(), is(emptyIterable())); @@ -734,8 +734,8 @@ public void testOnlyBundleForKeys() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - false /* hasNoState */, - true /* onlyBundleForKeys */); + () -> false /* hasNoState */, + () -> true /* onlyBundleForKeys */); // Reads still work assertArrayEquals( From bdc6c2d5a1f09f6ce391db1e98a0d92e161086a3 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 18 Mar 2026 02:20:25 -0400 Subject: [PATCH 16/33] minor fix --- .../beam/fn/harness/state/MultimapSideInputTest.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java index 595aa71a6666..8e7fb68a4261 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java @@ -70,7 +70,7 @@ public void testGetWithBulkRead() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> true, + true, () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); @@ -97,7 +97,7 @@ public void testGet() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> true, + true, () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); @@ -129,7 +129,7 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> true, + true, () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, @@ -154,7 +154,7 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> true, + true, () -> false); assertArrayEquals( new String[] {"A1", "A2", "A3"}, From c1d41ce922d3e0cc5b71cb3ec252f6cc93eaf210 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 18 Mar 2026 23:22:14 -0400 Subject: [PATCH 17/33] define emptyCachingStateIterable method --- .../harness/state/StateFetchingIterators.java | 51 ++++++++++++++----- 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index ed2ade1809c1..fa821a248822 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -94,20 +94,6 @@ public static CachingStateIterable readAllAndDecodeStartingFrom( valueCoder); } - public static CachingStateIterable emptyCachingStateIterable( - BeamFnStateClient beamFnStateClient, - StateRequest stateRequestForFirstChunk, - Coder valueCoder) { - CachingStateIterable result = - readAllAndDecodeStartingFrom( - org.apache.beam.fn.harness.Caches.noop(), - beamFnStateClient, - stateRequestForFirstChunk, - valueCoder); - result.clearAndAppend(java.util.Collections.emptyList()); - return result; - } - /** * This adapter handles using the continuation token to provide iteration over all the elements * returned by the Beam Fn State API using the supplied state client, state request for the first @@ -515,6 +501,18 @@ private void appendHelper(List newValues, long newWeight) { cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues))); } + private static final CachingStateIterable EMPTY_ITERABLE = new CachingStateIterable() { + @Override + protected CachingStateIterator createIterator() { + return CachingStateIterator.emptyIterator(); + } + }; + + /** Returns an empty {@link CachingStateIterable}. */ + public static CachingStateIterable emptyIterable() { + return (CachingStateIterable) EMPTY_ITERABLE; + } + class CachingStateIterator implements PrefetchableIterator { private final LazyBlockingStateFetchingIterator underlyingStateFetchingIterator; @@ -667,6 +665,31 @@ public T next() { } return currentBlock.getValues().get(currentCachedBlockValueIndex++); } + + public static CachingStateIterator emptyIterator() { + return (CachingStateIterator) EMPTY_ITERATOR; + } + + private static final CachingStateIterator EMPTY_ITERATOR = new CachingStateIterator() { + @Override + public boolean isReady() { + return true; + } + + @Override + public void prefetch() { + } + + @Override + public boolean hasNext() { + return false; + } + + @Override + public Object next() { + throw new NoSuchElementException(); + } + }; } } From 07608d3b15f9df5fa01d7724da009255cfb19c98 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 18 Mar 2026 23:41:00 -0400 Subject: [PATCH 18/33] hasNoState does not affect side input --- .../fn/harness/state/IterableSideInput.java | 24 +++++++------------ .../fn/harness/state/MultimapSideInput.java | 17 +++---------- 2 files changed, 12 insertions(+), 29 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java index 4fab87c0b3aa..cdf3098607c3 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java @@ -19,12 +19,10 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; 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.fn.stream.PrefetchableIterables; import org.apache.beam.sdk.transforms.Materializations.IterableView; /** @@ -42,23 +40,19 @@ public IterableSideInput( BeamFnStateClient beamFnStateClient, String instructionId, StateKey stateKey, - Coder valueCoder, - Supplier hasNoState) { + Coder valueCoder) { checkArgument( stateKey.hasIterableSideInput(), "Expected IterableSideInput StateKey but received %s.", stateKey); - this.values = - hasNoState.get() - ? PrefetchableIterables.emptyIterable() - : StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, - beamFnStateClient, - StateRequest.newBuilder() - .setInstructionId(instructionId) - .setStateKey(stateKey) - .build(), - valueCoder); + this.values = StateFetchingIterators.readAllAndDecodeStartingFrom( + cache, + beamFnStateClient, + StateRequest.newBuilder() + .setInstructionId(instructionId) + .setStateKey(stateKey) + .build(), + valueCoder); } @Override diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java index 4610fdd797c3..8e38a57ff9fe 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapSideInput.java @@ -25,7 +25,6 @@ import java.util.Iterator; import java.util.Map; import java.util.function.Function; -import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; @@ -33,7 +32,6 @@ 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.PrefetchableIterables; import org.apache.beam.sdk.transforms.Materializations.MultimapView; import org.apache.beam.sdk.util.ByteStringOutputStream; import org.apache.beam.sdk.values.KV; @@ -57,7 +55,6 @@ public class MultimapSideInput implements MultimapView { private final Coder valueCoder; private volatile Function> bulkReadResult; private final boolean useBulkRead; - private final Supplier hasNoState; public MultimapSideInput( Cache cache, @@ -66,8 +63,7 @@ public MultimapSideInput( StateKey stateKey, Coder keyCoder, Coder valueCoder, - boolean useBulkRead, - Supplier hasNoState) { + boolean useBulkRead) { checkArgument( stateKey.hasMultimapKeysSideInput(), "Expected MultimapKeysSideInput StateKey but received %s.", @@ -79,23 +75,16 @@ public MultimapSideInput( this.keyCoder = keyCoder; this.valueCoder = valueCoder; this.useBulkRead = useBulkRead; - this.hasNoState = hasNoState; } @Override public Iterable get() { - return this.hasNoState.get() - ? PrefetchableIterables.emptyIterable() - : StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, beamFnStateClient, keysRequest, keyCoder); + return StateFetchingIterators.readAllAndDecodeStartingFrom( + cache, beamFnStateClient, keysRequest, keyCoder); } @Override public Iterable get(K k) { - if (this.hasNoState.get()) { - return PrefetchableIterables.emptyIterable(); - } - ByteString encodedKey = encodeKey(k); if (useBulkRead) { From 56939091c5ba879eeef8006fe25cc1c23ef1c04d Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 18 Mar 2026 23:43:42 -0400 Subject: [PATCH 19/33] refactor --- .../org/apache/beam/fn/harness/state/BagUserState.java | 3 +-- .../apache/beam/fn/harness/state/FnApiStateAccessor.java | 6 ++---- .../apache/beam/fn/harness/state/MultimapUserState.java | 8 ++------ 3 files changed, 5 insertions(+), 12 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index c8ab95aa6945..f62db532e651 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -81,8 +81,7 @@ public BagUserState( this.oldValues = hasNoState.get() - ? StateFetchingIterators.emptyCachingStateIterable( - beamFnStateClient, request, valueCoder) + ? CachingStateIterable.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( this.cache, beamFnStateClient, request, valueCoder); this.newValues = new ArrayList<>(); 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 c81042618b38..eb91b6e24f0e 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 @@ -418,8 +418,7 @@ public T get(PCollectionView view, BoundedWindow window) { beamFnStateClient, processBundleInstructionId.get(), key, - sideInputSpec.getCoder(), - hasNoState)); + sideInputSpec.getCoder())); case Materializations.MULTIMAP_MATERIALIZATION_URN: return sideInputSpec .getViewFn() @@ -434,8 +433,7 @@ public T get(PCollectionView view, BoundedWindow window) { runnerCapabilities.contains( BeamUrns.getUrn( RunnerApi.StandardRunnerProtocols.Enum - .MULTIMAP_KEYS_VALUES_SIDE_INPUT)), - hasNoState)); + .MULTIMAP_KEYS_VALUES_SIDE_INPUT)))); default: throw new IllegalStateException( String.format( 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 9a31b01c4df3..31afdf3678c3 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 @@ -107,8 +107,7 @@ public MultimapUserState( StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = hasNoState.get() - ? StateFetchingIterators.emptyCachingStateIterable( - beamFnStateClient, keysStateRequest, mapKeyCoder) + ? CachingStateIterable.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( cache, beamFnStateClient, keysStateRequest, mapKeyCoder); @@ -135,10 +134,7 @@ public MultimapUserState( this.entriesStateRequest = entriesStateRequestBuilder.build(); this.persistedEntries = hasNoState.get() - ? StateFetchingIterators.emptyCachingStateIterable( - beamFnStateClient, - entriesStateRequest, - KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder))) + ? CachingStateIterable.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( Caches.subCache(this.cache, "AllEntries"), beamFnStateClient, From e9917d1c43a87a993599ce45c6b8b3608cdb37df Mon Sep 17 00:00:00 2001 From: Kai Li Date: Thu, 19 Mar 2026 00:07:37 -0400 Subject: [PATCH 20/33] fix sideinput test --- .../beam/fn/harness/state/IterableSideInputTest.java | 9 +++------ .../beam/fn/harness/state/MultimapSideInputTest.java | 12 ++++-------- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java index cd630cbf02f1..24de55fad976 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java @@ -47,8 +47,7 @@ public void testGet() throws Exception { fakeBeamFnStateClient, "instructionId", key(), - StringUtf8Coder.of(), - () -> false); + StringUtf8Coder.of()); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); @@ -70,8 +69,7 @@ public void testGetCached() throws Exception { fakeBeamFnStateClient, "instructionId", key(), - StringUtf8Coder.of(), - () -> false); + StringUtf8Coder.of()); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); @@ -87,8 +85,7 @@ public void testGetCached() throws Exception { }, "instructionId", key(), - StringUtf8Coder.of(), - () -> false); + StringUtf8Coder.of()); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java index 8e7fb68a4261..88982bf1b6f6 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapSideInputTest.java @@ -70,8 +70,7 @@ public void testGetWithBulkRead() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - () -> false); + true); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); assertArrayEquals( @@ -97,8 +96,7 @@ public void testGet() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - () -> false); + true); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); assertArrayEquals( @@ -129,8 +127,7 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - () -> false); + true); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); @@ -154,8 +151,7 @@ public void testGetCached() throws Exception { keysStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - true, - () -> false); + true); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(multimapSideInput.get(A), String.class)); From e27bdffc36800ce884ab1b8766e46badfa2957f2 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Thu, 19 Mar 2026 00:13:23 -0400 Subject: [PATCH 21/33] spotless check fix --- .../fn/harness/state/IterableSideInput.java | 14 +++--- .../harness/state/StateFetchingIterators.java | 47 ++++++++++--------- .../harness/state/IterableSideInputTest.java | 12 +---- 3 files changed, 32 insertions(+), 41 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java index cdf3098607c3..7f87cb4d4e41 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/IterableSideInput.java @@ -45,14 +45,12 @@ public IterableSideInput( stateKey.hasIterableSideInput(), "Expected IterableSideInput StateKey but received %s.", stateKey); - this.values = StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, - beamFnStateClient, - StateRequest.newBuilder() - .setInstructionId(instructionId) - .setStateKey(stateKey) - .build(), - valueCoder); + this.values = + StateFetchingIterators.readAllAndDecodeStartingFrom( + cache, + beamFnStateClient, + StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(), + valueCoder); } @Override diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index fa821a248822..e3b23d090d59 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -501,12 +501,13 @@ private void appendHelper(List newValues, long newWeight) { cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues))); } - private static final CachingStateIterable EMPTY_ITERABLE = new CachingStateIterable() { - @Override - protected CachingStateIterator createIterator() { - return CachingStateIterator.emptyIterator(); - } - }; + private static final CachingStateIterable EMPTY_ITERABLE = + new CachingStateIterable() { + @Override + protected CachingStateIterator createIterator() { + return CachingStateIterator.emptyIterator(); + } + }; /** Returns an empty {@link CachingStateIterable}. */ public static CachingStateIterable emptyIterable() { @@ -670,26 +671,26 @@ public static CachingStateIterator emptyIterator() { return (CachingStateIterator) EMPTY_ITERATOR; } - private static final CachingStateIterator EMPTY_ITERATOR = new CachingStateIterator() { - @Override - public boolean isReady() { - return true; - } + private static final CachingStateIterator EMPTY_ITERATOR = + new CachingStateIterator() { + @Override + public boolean isReady() { + return true; + } - @Override - public void prefetch() { - } + @Override + public void prefetch() {} - @Override - public boolean hasNext() { - return false; - } + @Override + public boolean hasNext() { + return false; + } - @Override - public Object next() { - throw new NoSuchElementException(); - } - }; + @Override + public Object next() { + throw new NoSuchElementException(); + } + }; } } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java index 24de55fad976..aa6d9115c1ee 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/IterableSideInputTest.java @@ -43,11 +43,7 @@ public void testGet() throws Exception { IterableSideInput iterableSideInput = new IterableSideInput<>( - Caches.noop(), - fakeBeamFnStateClient, - "instructionId", - key(), - StringUtf8Coder.of()); + Caches.noop(), fakeBeamFnStateClient, "instructionId", key(), StringUtf8Coder.of()); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); @@ -65,11 +61,7 @@ public void testGetCached() throws Exception { // The first side input will populate the cache. IterableSideInput iterableSideInput = new IterableSideInput<>( - cache, - fakeBeamFnStateClient, - "instructionId", - key(), - StringUtf8Coder.of()); + cache, fakeBeamFnStateClient, "instructionId", key(), StringUtf8Coder.of()); assertArrayEquals( new String[] {"A1", "A2", "A3", "A4", "A5", "A6"}, Iterables.toArray(iterableSideInput.get(), String.class)); From be0f73e4d4fc3198792a78d2318b230f3c84d723 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Thu, 19 Mar 2026 01:12:12 -0400 Subject: [PATCH 22/33] fixing emptyIterable --- .../harness/state/StateFetchingIterators.java | 58 ++++++++++--------- 1 file changed, 30 insertions(+), 28 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index e3b23d090d59..8143db9ba467 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -501,15 +501,42 @@ private void appendHelper(List newValues, long newWeight) { cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues))); } + private static final PrefetchableIterator EMPTY_ITERATOR = new PrefetchableIterator() { + @Override + public boolean isReady() { + return true; + } + + @Override + public void prefetch() { + } + + @Override + public boolean hasNext() { + return false; + } + + @Override + public Object next() { + throw new NoSuchElementException(); + } + }; + + @SuppressWarnings("unchecked") + public static PrefetchableIterator emptyIterator() { + return (PrefetchableIterator) EMPTY_ITERATOR; + } + private static final CachingStateIterable EMPTY_ITERABLE = - new CachingStateIterable() { + new CachingStateIterable(null, null, null, null) { @Override - protected CachingStateIterator createIterator() { - return CachingStateIterator.emptyIterator(); + public PrefetchableIterator createIterator() { + return emptyIterator(); } }; /** Returns an empty {@link CachingStateIterable}. */ + @SuppressWarnings("unchecked") public static CachingStateIterable emptyIterable() { return (CachingStateIterable) EMPTY_ITERABLE; } @@ -666,31 +693,6 @@ public T next() { } return currentBlock.getValues().get(currentCachedBlockValueIndex++); } - - public static CachingStateIterator emptyIterator() { - return (CachingStateIterator) EMPTY_ITERATOR; - } - - private static final CachingStateIterator EMPTY_ITERATOR = - new CachingStateIterator() { - @Override - public boolean isReady() { - return true; - } - - @Override - public void prefetch() {} - - @Override - public boolean hasNext() { - return false; - } - - @Override - public Object next() { - throw new NoSuchElementException(); - } - }; } } From 306ccbc1337c7f17dbe85d11532a4480df0afd4c Mon Sep 17 00:00:00 2001 From: Kai Li Date: Thu, 19 Mar 2026 01:43:44 -0400 Subject: [PATCH 23/33] fix unit test --- .../harness/state/StateFetchingIterators.java | 54 +++++++++++++------ 1 file changed, 37 insertions(+), 17 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 8143db9ba467..1642897a468c 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -501,26 +501,26 @@ private void appendHelper(List newValues, long newWeight) { cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues))); } - private static final PrefetchableIterator EMPTY_ITERATOR = new PrefetchableIterator() { - @Override - public boolean isReady() { - return true; - } + private static final PrefetchableIterator EMPTY_ITERATOR = + new PrefetchableIterator() { + @Override + public boolean isReady() { + return true; + } - @Override - public void prefetch() { - } + @Override + public void prefetch() {} - @Override - public boolean hasNext() { - return false; - } + @Override + public boolean hasNext() { + return false; + } - @Override - public Object next() { - throw new NoSuchElementException(); - } - }; + @Override + public Object next() { + throw new NoSuchElementException(); + } + }; @SuppressWarnings("unchecked") public static PrefetchableIterator emptyIterator() { @@ -533,6 +533,26 @@ public static PrefetchableIterator emptyIterator() { public PrefetchableIterator createIterator() { return emptyIterator(); } + + @Override + public void remove(Set toRemoveStructuralValues) { + } + + @Override + public void clearAndAppend(List values) { + } + + @Override + public void clearAndAppend(WeightedList values) { + } + + @Override + public void append(List values) { + } + + @Override + public void append(WeightedList values) { + } }; /** Returns an empty {@link CachingStateIterable}. */ From 5da21f2d35ecdde964e213cb1187278e9930f00d Mon Sep 17 00:00:00 2001 From: Kai Li Date: Thu, 19 Mar 2026 01:57:10 -0400 Subject: [PATCH 24/33] fix --- .../fn/harness/state/StateFetchingIterators.java | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 1642897a468c..f9166e0c0742 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -535,24 +535,19 @@ public PrefetchableIterator createIterator() { } @Override - public void remove(Set toRemoveStructuralValues) { - } + public void remove(Set toRemoveStructuralValues) {} @Override - public void clearAndAppend(List values) { - } + public void clearAndAppend(List values) {} @Override - public void clearAndAppend(WeightedList values) { - } + public void clearAndAppend(WeightedList values) {} @Override - public void append(List values) { - } + public void append(List values) {} @Override - public void append(WeightedList values) { - } + public void append(WeightedList values) {} }; /** Returns an empty {@link CachingStateIterable}. */ From a85798239e5f8f7c3dba44d16bde554a7fbeb081 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Thu, 19 Mar 2026 02:13:46 -0400 Subject: [PATCH 25/33] add unit tests --- .../state/OrderedListUserStateTest.java | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java index 23b1c674d502..7280b016b7bd 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java @@ -721,6 +721,41 @@ public void testHasNoState() throws Exception { assertEquals(out.toByteString(), fakeClient.getData().get(createOrderedListStateKey("A", 2))); } + @Test + public void testReadRangeHasNoState() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + timestampedValueCoder, + ImmutableMap.of(createOrderedListStateKey("A", 1), Collections.singletonList(A1))); + OrderedListUserState userState = + new OrderedListUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createOrderedListStateKey("A"), + StringUtf8Coder.of(), + () -> true, /* hasNoState */ + () -> false /* onlyBundleForKeys */); + + // Iterating should be empty since hasNoState is true + assertThat( + userState.readRange(Instant.ofEpochMilli(0), Instant.ofEpochMilli(2)), is(emptyIterable())); + + // We can add new values + userState.add(A2); + assertArrayEquals( + Collections.singletonList(A2).toArray(), + Iterables.toArray( + userState.readRange(Instant.ofEpochMilli(1), Instant.ofEpochMilli(3)), + TimestampedValue.class)); + userState.asyncClose(); + + // The added value will be persisted, appended to the old values (since we didn't clear) + ByteStringOutputStream out = new ByteStringOutputStream(); + timestampedValueCoder.encode(A2, out); + assertEquals(out.toByteString(), fakeClient.getData().get(createOrderedListStateKey("A", 2))); + } + @Test public void testOnlyBundleForKeys() throws Exception { FakeBeamFnStateClient fakeClient = @@ -758,6 +793,43 @@ public void testOnlyBundleForKeys() throws Exception { assertFalse(fakeClient.getData().containsKey(createOrderedListStateKey("A", 2))); } + @Test + public void testHasNoStateAndOnlyBundleForKeys() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + timestampedValueCoder, + ImmutableMap.of(createOrderedListStateKey("A", 1), Collections.singletonList(A1))); + OrderedListUserState userState = + new OrderedListUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createOrderedListStateKey("A"), + StringUtf8Coder.of(), + () -> true /* hasNoState */, + () -> true /* onlyBundleForKeys */); + + // Iterating should be empty since hasNoState is true + assertThat(userState.read(), is(emptyIterable())); + + // We can add new values + userState.add(A2); + + // We can observe the added value locally + assertArrayEquals( + Collections.singletonList(A2).toArray(), + Iterables.toArray(userState.read(), TimestampedValue.class)); + + userState.asyncClose(); + + // The data is NOT persisted because onlyBundleForKeys = true, and the old data remains + // untouched + ByteStringOutputStream out = new ByteStringOutputStream(); + timestampedValueCoder.encode(A1, out); + assertEquals(out.toByteString(), fakeClient.getData().get(createOrderedListStateKey("A", 1))); + assertFalse(fakeClient.getData().containsKey(createOrderedListStateKey("A", 2))); + } + private ByteString encode(String... values) throws IOException { ByteStringOutputStream out = new ByteStringOutputStream(); for (String value : values) { From 20f7ab3e89d54b957e2c921a108936ac1fe89cd7 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Fri, 20 Mar 2026 01:38:21 -0400 Subject: [PATCH 26/33] pass only boolean hasNoState/onlyBundleForKeys to non-factory class --- .../harness/control/ProcessBundleHandler.java | 8 +++---- .../beam/fn/harness/state/BagUserState.java | 15 ++++++------- .../fn/harness/state/FnApiStateAccessor.java | 12 +++++----- .../fn/harness/state/MultimapUserState.java | 22 +++++++++---------- .../harness/state/OrderedListUserState.java | 13 +++++------ .../fn/harness/state/BagUserStateTest.java | 8 +++---- .../state/OrderedListUserStateTest.java | 16 +++++++------- 7 files changed, 45 insertions(+), 49 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 7146419429d9..045d3f430d19 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -247,8 +247,8 @@ private void addRunnerAndConsumersForPTransformRecursively( Collection> channelRoots, Map outboundAggregatorMap, Set runnerCapabilities, - Supplier hasNoState, - Supplier onlyBundleForKeys) + boolean hasNoState, + boolean onlyBundleForKeys) throws IOException { // Recursively ensure that all consumers of the output PCollection have been created. @@ -495,12 +495,12 @@ public BundleFinalizer getBundleFinalizer() { @Override public Supplier getHasNoState() { - return hasNoState; + return () -> hasNoState; } @Override public Supplier getOnlyBundleForKeys() { - return onlyBundleForKeys; + return () -> onlyBundleForKeys; } }); processedPTransformIds.add(pTransformId); diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index f62db532e651..b945ebe7d0f3 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -23,7 +23,6 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.state.StateFetchingIterators.CachingStateIterable; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateAppendRequest; @@ -55,8 +54,8 @@ public class BagUserState { private List newValues; private boolean isCleared; private boolean isClosed; - private final Supplier hasNoState; - private final Supplier onlyBundleForKeys; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; static final int BAG_APPEND_BATCHING_LIMIT = 10 * 1024 * 1024; @@ -67,8 +66,8 @@ public BagUserState( String instructionId, StateKey stateKey, Coder valueCoder, - Supplier hasNoState, - Supplier onlyBundleForKeys) { + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasBagUserState(), "Expected BagUserState StateKey but received %s.", stateKey); this.cache = cache; @@ -80,7 +79,7 @@ public BagUserState( StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.oldValues = - hasNoState.get() + hasNoState ? CachingStateIterable.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( this.cache, beamFnStateClient, request, valueCoder); @@ -92,7 +91,7 @@ public PrefetchableIterable get() { !isClosed, "Bag user state is no longer usable because it is closed for %s", request.getStateKey()); - if (isCleared || hasNoState.get()) { + if (isCleared || hasNoState) { // If we were cleared or have no state we should disregard old values. return PrefetchableIterables.limit(Collections.unmodifiableList(newValues), newValues.size()); } else if (newValues.isEmpty()) { @@ -129,7 +128,7 @@ public void asyncClose() throws Exception { "Bag user state is no longer usable because it is closed for %s", request.getStateKey()); isClosed = true; - if (onlyBundleForKeys.get() || (!isCleared && newValues.isEmpty())) { + if (onlyBundleForKeys || (!isCleared && newValues.isEmpty())) { return; } if (isCleared) { 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 eb91b6e24f0e..e5c686d3f23f 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 @@ -244,8 +244,8 @@ public FnApiStateAccessor create() { beamFnStateClient, keyCoder, windowCoder, - hasNoState, - onlyBundleForKeys); + hasNoState.get(), + onlyBundleForKeys.get()); } } @@ -262,8 +262,8 @@ public FnApiStateAccessor create() { private final Collection stateFinalizers; private final Coder keyCoder; private final Coder windowCoder; - private final Supplier hasNoState; - private final Supplier onlyBundleForKeys; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; private @Nullable Supplier currentWindowSupplier; private @Nullable Supplier encodedCurrentKeySupplier; @@ -281,8 +281,8 @@ public FnApiStateAccessor( BeamFnStateClient beamFnStateClient, Coder keyCoder, Coder windowCoder, - Supplier hasNoState, - Supplier onlyBundleForKeys) { + boolean hasNoState, + boolean onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.stateKeyObjectCache = Maps.newHashMap(); 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 31afdf3678c3..567f7fcabbac 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 @@ -31,7 +31,6 @@ import java.util.NoSuchElementException; import java.util.Objects; import java.util.Set; -import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.fn.harness.state.StateFetchingIterators.CachingStateIterable; @@ -73,8 +72,8 @@ public class MultimapUserState { private boolean isClosed; private boolean isCleared; - private final Supplier hasNoState; - private final Supplier onlyBundleForKeys; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; // Pending updates to persistent storage private HashMap pendingRemoves = Maps.newHashMap(); private HashMap>> pendingAdds = Maps.newHashMap(); @@ -88,8 +87,8 @@ public MultimapUserState( StateKey stateKey, Coder mapKeyCoder, Coder valueCoder, - Supplier hasNoState, - Supplier onlyBundleForKeys) { + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasMultimapKeysUserState(), "Expected MultimapKeysUserState StateKey but received %s.", @@ -106,7 +105,7 @@ public MultimapUserState( this.keysStateRequest = StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = - hasNoState.get() + hasNoState ? CachingStateIterable.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( cache, beamFnStateClient, keysStateRequest, mapKeyCoder); @@ -133,7 +132,7 @@ public MultimapUserState( .setKey(stateKey.getMultimapKeysUserState().getKey()); this.entriesStateRequest = entriesStateRequestBuilder.build(); this.persistedEntries = - hasNoState.get() + hasNoState ? CachingStateIterable.emptyIterable() : StateFetchingIterators.readAllAndDecodeStartingFrom( Caches.subCache(this.cache, "AllEntries"), @@ -172,7 +171,7 @@ public PrefetchableIterable get(K key) { ? PrefetchableIterables.fromArray() : PrefetchableIterables.limit( pendingAddValues.getValue(), pendingAddValues.getValue().size()); - if (isCleared || hasNoState.get() || pendingRemoves.containsKey(structuralKey)) { + if (isCleared || hasNoState || pendingRemoves.containsKey(structuralKey)) { return pendingValues; } @@ -190,7 +189,7 @@ public PrefetchableIterable keys() { !isClosed, "Multimap user state is no longer usable because it is closed for %s", keysStateRequest.getStateKey()); - if (isCleared || hasNoState.get()) { + if (isCleared || hasNoState) { List keys = new ArrayList<>(pendingAdds.size()); for (Map.Entry>> entry : pendingAdds.entrySet()) { keys.add(entry.getValue().getKey()); @@ -288,7 +287,7 @@ public PrefetchableIterable> entries() { entry.getKey(), KV.of(entry.getValue().getKey(), new ArrayList<>(entry.getValue().getValue()))); } - if (isCleared || hasNoState.get()) { + if (isCleared || hasNoState) { return PrefetchableIterables.maybePrefetchable( Iterables.concat( Iterables.transform( @@ -426,8 +425,7 @@ public void asyncClose() throws Exception { keysStateRequest.getStateKey()); isClosed = true; // No mutations necessary - if (onlyBundleForKeys.get() - || (!isCleared && pendingRemoves.isEmpty() && pendingAdds.isEmpty())) { + if (onlyBundleForKeys || (!isCleared && pendingRemoves.isEmpty() && pendingAdds.isEmpty())) { return; } diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index 8c002a5ecc4c..49c25ee449a5 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -33,7 +33,6 @@ import java.util.Map.Entry; import java.util.NavigableMap; import java.util.concurrent.CompletableFuture; -import java.util.function.Supplier; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.fn.harness.state.StateFetchingIterators.CachingStateIterable; @@ -89,8 +88,8 @@ public class OrderedListUserState { private boolean isCleared = false; private boolean isClosed = false; - private final Supplier hasNoState; - private final Supplier onlyBundleForKeys; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; public static class TimestampedValueCoder extends StructuredCoder> { @@ -165,8 +164,8 @@ public OrderedListUserState( String instructionId, StateKey stateKey, Coder valueCoder, - Supplier hasNoState, - Supplier onlyBundleForKeys) { + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasOrderedListUserState(), "Expected OrderedListUserState StateKey but received %s.", @@ -210,7 +209,7 @@ public Iterable> readRange(Instant minTimestamp, Instant lim } Iterable> valuesInRange = Iterables.concat(pendingAddsInRange); - if (!isCleared && !hasNoState.get()) { + if (!isCleared && !hasNoState) { StateRequest.Builder getRequestBuilder = this.requestTemplate.toBuilder(); getRequestBuilder .getStateKeyBuilder() @@ -290,7 +289,7 @@ public void clear() { public void asyncClose() throws Exception { isClosed = true; - if (onlyBundleForKeys.get()) { + if (onlyBundleForKeys) { pendingRemoves.clear(); pendingAdds.clear(); return; diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java index f4311fb65d3d..a57f0d256e15 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java @@ -364,8 +364,8 @@ public void testHasNoState() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> true, /* hasNoState */ - () -> false /* onlyBundleForKeys */); + true, /* hasNoState */ + false /* onlyBundleForKeys */); // Iterating should be empty since hasNoState is true assertFalse(userState.get().iterator().hasNext()); @@ -390,8 +390,8 @@ public void testOnlyBundleForKeys() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false /* hasNoState */, - () -> true /* onlyBundleForKeys */); + false /* hasNoState */, + true /* onlyBundleForKeys */); // Reads still work assertArrayEquals(new String[] {"A1"}, Iterables.toArray(userState.get(), String.class)); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java index 7280b016b7bd..9488e8d5eb4d 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java @@ -702,8 +702,8 @@ public void testHasNoState() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> true, /* hasNoState */ - () -> false /* onlyBundleForKeys */); + true, /* hasNoState */ + false /* onlyBundleForKeys */); // Iterating should be empty since hasNoState is true assertThat(userState.read(), is(emptyIterable())); @@ -734,8 +734,8 @@ public void testReadRangeHasNoState() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> true, /* hasNoState */ - () -> false /* onlyBundleForKeys */); + true, /* hasNoState */ + false /* onlyBundleForKeys */); // Iterating should be empty since hasNoState is true assertThat( @@ -769,8 +769,8 @@ public void testOnlyBundleForKeys() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false /* hasNoState */, - () -> true /* onlyBundleForKeys */); + false /* hasNoState */, + true /* onlyBundleForKeys */); // Reads still work assertArrayEquals( @@ -806,8 +806,8 @@ public void testHasNoStateAndOnlyBundleForKeys() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> true /* hasNoState */, - () -> true /* onlyBundleForKeys */); + true /* hasNoState */, + true /* onlyBundleForKeys */); // Iterating should be empty since hasNoState is true assertThat(userState.read(), is(emptyIterable())); From 213e6ce7d98c197cf90aafca60ea25dc4004bb3a Mon Sep 17 00:00:00 2001 From: Kai Li Date: Fri, 20 Mar 2026 03:44:42 -0400 Subject: [PATCH 27/33] fix unit test --- .../harness/control/ProcessBundleHandler.java | 4 +- .../harness/state/MultimapUserStateTest.java | 269 ++++++++++++------ 2 files changed, 187 insertions(+), 86 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 045d3f430d19..d942793826c2 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -928,8 +928,8 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { bundleProcessor.getChannelRoots(), bundleProcessor.getOutboundAggregators(), bundleProcessor.getRunnerCapabilities(), - processBundleRequest::getHasNoState, - processBundleRequest::getOnlyBundleForKeys); + processBundleRequest.getHasNoState(), + processBundleRequest.getOnlyBundleForKeys()); } bundleProcessor.finish(); 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 5675e9529d11..1ae922e4b820 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 @@ -26,6 +26,7 @@ import static org.hamcrest.core.Is.is; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; @@ -85,8 +86,8 @@ public void testNoPersistedValues() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertThat(userState.keys(), is(emptyIterable())); } @@ -107,8 +108,8 @@ public void testGet() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -137,8 +138,8 @@ public void testClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); Iterable initValues = userState.get(A1); userState.clear(); @@ -173,8 +174,8 @@ public void testKeys() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A2, "V1"); Iterable initKeys = userState.keys(); @@ -209,8 +210,8 @@ public void testEntries() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals(A1, userState.entries().iterator().next().getKey()); assertThat( @@ -286,8 +287,8 @@ public void testPut() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -315,8 +316,8 @@ public void testPutAfterRemove() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.remove(A0); userState.put(A0, "V2"); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); @@ -341,8 +342,8 @@ public void testPutAfterClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.clear(); userState.put(A0, "V2"); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); @@ -365,8 +366,8 @@ public void testRemoveBeforeClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.remove(A0); userState.clear(); userState.asyncClose(); @@ -385,8 +386,8 @@ public void testPutBeforeClear() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A0, "V0"); userState.put(A1, "V1"); Iterable values = userState.get(A1); // fakeClient call = 1 @@ -408,8 +409,8 @@ public void testPutBeforeRemove() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A0, "V0"); userState.put(A1, "V1"); Iterable values = userState.get(A1); // fakeClient call = 1 @@ -439,8 +440,8 @@ public void testRemove() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); Iterable initValues = userState.get(A1); userState.put(A1, "V3"); @@ -469,8 +470,8 @@ public void testImmutableKeys() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); Iterable keys = userState.keys(); Iterator keysIterator = keys.iterator(); keysIterator.next(); @@ -494,8 +495,8 @@ public void testImmutableValues() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); Iterable values = userState.get(A1); assertThrows( UnsupportedOperationException.class, @@ -519,8 +520,8 @@ public void testClearAsyncClose() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.clear(); userState.asyncClose(); Map data = fakeClient.getData(); @@ -545,8 +546,8 @@ public void testNoopAsyncClose() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.asyncClose(); assertThrows(IllegalStateException.class, () -> userState.keys()); assertEquals(0, fakeClient.getCallCount()); @@ -571,8 +572,8 @@ public void testAsyncClose() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.remove(A0); userState.put(A1, "V3"); userState.put(A2, "V1"); @@ -602,8 +603,8 @@ public void testNullKeysAndValues() throws Exception { createMultimapKeyStateKey(), NullableCoder.of(ByteArrayCoder.of()), NullableCoder.of(StringUtf8Coder.of()), - () -> false, - () -> false); + false, + false); userState.put(null, null); userState.put(null, null); userState.put(null, "V1"); @@ -622,8 +623,8 @@ public void testNegativeCache() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(A1), String.class)); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(A1), String.class)); assertThat(fakeClient.getCallCount(), is(1)); @@ -646,8 +647,8 @@ public void testGetValuesPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); PrefetchableIterable values = userState.get(A1); assertEquals(0, fakeClient.getCallCount()); @@ -674,8 +675,8 @@ public void testGetKeysPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); PrefetchableIterable keys = userState.keys(); assertEquals(0, fakeClient.getCallCount()); @@ -702,8 +703,8 @@ public void testPutKeysPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A2, "V3"); PrefetchableIterable keys = userState.keys(); @@ -731,8 +732,8 @@ public void testRemoveKeysPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.remove(A1); userState.put(A1, "V3"); @@ -766,8 +767,8 @@ public void testEntriesPrefetched() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A1, "V4"); PrefetchableIterable> entries = userState.entries(); @@ -803,8 +804,8 @@ public void testClearPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.clear(); userState.put(A2, "V3"); @@ -835,8 +836,8 @@ public void testAppendValuesPrefetch() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A1, "V3"); PrefetchableIterable values = userState.get(A1); @@ -861,8 +862,8 @@ public void testNoPersistedValuesCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertThat(userState.keys(), is(emptyIterable())); assertThat(userState.get(A1), is(emptyIterable())); } @@ -879,8 +880,8 @@ public void testNoPersistedValuesCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertThat(userState.keys(), is(emptyIterable())); assertThat(userState.get(A1), is(emptyIterable())); } @@ -906,8 +907,8 @@ public void testGetCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); @@ -926,8 +927,8 @@ public void testGetCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); @@ -955,8 +956,8 @@ public void testClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.clear(); assertThat(userState.keys(), is(emptyIterable())); @@ -976,8 +977,8 @@ public void testClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertThat(userState.keys(), is(emptyIterable())); userState.asyncClose(); @@ -1004,8 +1005,8 @@ public void testKeysCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A2, "V1"); userState.put(A3, "V1"); @@ -1028,8 +1029,8 @@ public void testKeysCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertThat( Iterables.toArray(userState.keys(), byte[].class), @@ -1058,8 +1059,8 @@ public void testPutCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.put(A1, "V3"); userState.put(A2, "V1"); @@ -1081,8 +1082,8 @@ public void testPutCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals( new String[] {"V1", "V2", "V3"}, Iterables.toArray(userState.get(A1), String.class)); @@ -1115,8 +1116,8 @@ public void testPutAfterRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.remove(A0); userState.put(A0, "V2"); @@ -1136,8 +1137,8 @@ public void testPutAfterRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); userState.asyncClose(); } @@ -1163,8 +1164,8 @@ public void testPutAfterClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); userState.clear(); userState.put(A0, "V2"); userState.asyncClose(); @@ -1183,8 +1184,8 @@ public void testPutAfterClearCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A0), String.class)); // Even though we never load assertArrayEquals(new byte[][] {A0}, Iterables.toArray(userState.keys(), byte[].class)); @@ -1212,8 +1213,8 @@ public void testRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertArrayEquals( new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); userState.remove(A1); @@ -1235,8 +1236,8 @@ public void testRemoveCached() throws Exception { createMultimapKeyStateKey(), ByteArrayCoder.of(), StringUtf8Coder.of(), - () -> false, - () -> false); + false, + false); assertThat(userState.get(A1), is(emptyIterable())); assertThat(userState.get(A2), is(emptyIterable())); assertThat(userState.keys(), is(emptyIterable())); @@ -1244,6 +1245,106 @@ public void testRemoveCached() throws Exception { } } + @Test + public void testHasNoState() throws Exception { + FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapKeyStateKey(), + KV.of(ByteArrayCoder.of(), singletonList(A1)), + createMultimapValueStateKey(A1), + KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); + MultimapUserState userState = new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createMultimapKeyStateKey(), + ByteArrayCoder.of(), + StringUtf8Coder.of(), + true /* hasNoState */, + false /* onlyBundleForKeys */); + + // Iterating should be empty since hasNoState is true + assertThat(userState.keys(), is(emptyIterable())); + assertThat(userState.get(A1), is(emptyIterable())); + + // We can add new values + userState.put(A2, "V2"); + + // We can observe the added value locally + assertArrayEquals(new byte[][] { A2 }, Iterables.toArray(userState.keys(), byte[].class)); + assertArrayEquals(new String[] { "V2" }, Iterables.toArray(userState.get(A2), String.class)); + + userState.asyncClose(); + + // The data is persisted + assertEquals(encode("V1", "V2"), fakeClient.getData().get(createMultimapValueStateKey(A1))); + assertEquals(encode("V2"), fakeClient.getData().get(createMultimapValueStateKey(A2))); + } + + @Test + public void testOnlyBundleForKeys() throws Exception { + FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapKeyStateKey(), + KV.of(ByteArrayCoder.of(), singletonList(A1)), + createMultimapValueStateKey(A1), + KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); + MultimapUserState userState = new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createMultimapKeyStateKey(), + ByteArrayCoder.of(), + StringUtf8Coder.of(), + false /* hasNoState */, + true /* onlyBundleForKeys */); + + assertArrayEquals(new byte[][] { A1 }, Iterables.toArray(userState.keys(), byte[].class)); + assertArrayEquals( + new String[] { "V1", "V2" }, Iterables.toArray(userState.get(A1), String.class)); + + userState.put(A2, "V2"); + + assertArrayEquals(new byte[][] { A1, A2 }, Iterables.toArray(userState.keys(), byte[].class)); + assertArrayEquals(new String[] { "V2" }, Iterables.toArray(userState.get(A2), String.class)); + + userState.asyncClose(); + + // The new data is NOT persisted + assertFalse(fakeClient.getData().containsKey(createMultimapValueStateKey(A2))); + } + + @Test + public void testHasNoStateAndOnlyBundleForKeys() throws Exception { + FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapKeyStateKey(), + KV.of(ByteArrayCoder.of(), singletonList(A1)), + createMultimapValueStateKey(A1), + KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); + MultimapUserState userState = new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createMultimapKeyStateKey(), + ByteArrayCoder.of(), + StringUtf8Coder.of(), + true /* hasNoState */, + true /* onlyBundleForKeys */); + + assertThat(userState.keys(), is(emptyIterable())); + assertThat(userState.get(A1), is(emptyIterable())); + + userState.put(A2, "V2"); + + assertArrayEquals(new byte[][] { A2 }, Iterables.toArray(userState.keys(), byte[].class)); + assertArrayEquals(new String[] { "V2" }, Iterables.toArray(userState.get(A2), String.class)); + + userState.asyncClose(); + + assertFalse(fakeClient.getData().containsKey(createMultimapValueStateKey(A2))); + } + private StateKey createMultimapKeyStateKey() throws IOException { return StateKey.newBuilder() .setMultimapKeysUserState( From dbb3dced49dd91f483c00aba935d5410bc641c10 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Mon, 23 Mar 2026 17:47:33 -0400 Subject: [PATCH 28/33] fix unit tests --- .../fn/harness/state/BagUserStateTest.java | 44 +++---- .../harness/state/MultimapUserStateTest.java | 108 +++++++++--------- .../state/OrderedListUserStateTest.java | 76 ++++++------ 3 files changed, 117 insertions(+), 111 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java index a57f0d256e15..5d3d849fed5f 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/BagUserStateTest.java @@ -54,8 +54,8 @@ public void testGet() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); @@ -79,8 +79,8 @@ public void testGetCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -97,8 +97,8 @@ public void testGetCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -116,8 +116,8 @@ public void testAppend() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -144,8 +144,8 @@ public void testAppendBatchingLimit() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); String a2 = Strings.repeat("A2", 2 * 1024 * 1024); userState.append(a2); String a3 = Strings.repeat("A3", 4 * 1024 * 1024); @@ -188,8 +188,8 @@ public void testAppendCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -215,8 +215,8 @@ public void testAppendCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.append("A4"); Iterable stateBeforeA5 = userState.get(); assertArrayEquals( @@ -244,8 +244,8 @@ public void testClear() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -276,8 +276,8 @@ public void testClearCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -302,8 +302,8 @@ public void testClearCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals(new String[] {"A4"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); assertFalse(userState.get().iterator().hasNext()); @@ -325,8 +325,8 @@ public void testClearCached() throws Exception { "instructionId", key("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); } 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 1ae922e4b820..85c53278c73e 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 @@ -1247,21 +1247,23 @@ public void testRemoveCached() throws Exception { @Test public void testHasNoState() throws Exception { - FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient( - ImmutableMap.of( + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapKeyStateKey(), + KV.of(ByteArrayCoder.of(), singletonList(A1)), + createMultimapValueStateKey(A1), + KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); + MultimapUserState userState = + new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", createMultimapKeyStateKey(), - KV.of(ByteArrayCoder.of(), singletonList(A1)), - createMultimapValueStateKey(A1), - KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); - MultimapUserState userState = new MultimapUserState<>( - Caches.noop(), - fakeClient, - "instructionId", - createMultimapKeyStateKey(), - ByteArrayCoder.of(), - StringUtf8Coder.of(), - true /* hasNoState */, - false /* onlyBundleForKeys */); + ByteArrayCoder.of(), + StringUtf8Coder.of(), + true /* hasNoState */, + false /* onlyBundleForKeys */); // Iterating should be empty since hasNoState is true assertThat(userState.keys(), is(emptyIterable())); @@ -1271,8 +1273,8 @@ public void testHasNoState() throws Exception { userState.put(A2, "V2"); // We can observe the added value locally - assertArrayEquals(new byte[][] { A2 }, Iterables.toArray(userState.keys(), byte[].class)); - assertArrayEquals(new String[] { "V2" }, Iterables.toArray(userState.get(A2), String.class)); + assertArrayEquals(new byte[][] {A2}, Iterables.toArray(userState.keys(), byte[].class)); + assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A2), String.class)); userState.asyncClose(); @@ -1283,30 +1285,32 @@ public void testHasNoState() throws Exception { @Test public void testOnlyBundleForKeys() throws Exception { - FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient( - ImmutableMap.of( + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapKeyStateKey(), + KV.of(ByteArrayCoder.of(), singletonList(A1)), + createMultimapValueStateKey(A1), + KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); + MultimapUserState userState = + new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", createMultimapKeyStateKey(), - KV.of(ByteArrayCoder.of(), singletonList(A1)), - createMultimapValueStateKey(A1), - KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); - MultimapUserState userState = new MultimapUserState<>( - Caches.noop(), - fakeClient, - "instructionId", - createMultimapKeyStateKey(), - ByteArrayCoder.of(), - StringUtf8Coder.of(), - false /* hasNoState */, - true /* onlyBundleForKeys */); - - assertArrayEquals(new byte[][] { A1 }, Iterables.toArray(userState.keys(), byte[].class)); + ByteArrayCoder.of(), + StringUtf8Coder.of(), + false /* hasNoState */, + true /* onlyBundleForKeys */); + + assertArrayEquals(new byte[][] {A1}, Iterables.toArray(userState.keys(), byte[].class)); assertArrayEquals( - new String[] { "V1", "V2" }, Iterables.toArray(userState.get(A1), String.class)); + new String[] {"V1", "V2"}, Iterables.toArray(userState.get(A1), String.class)); userState.put(A2, "V2"); - assertArrayEquals(new byte[][] { A1, A2 }, Iterables.toArray(userState.keys(), byte[].class)); - assertArrayEquals(new String[] { "V2" }, Iterables.toArray(userState.get(A2), String.class)); + assertArrayEquals(new byte[][] {A1, A2}, Iterables.toArray(userState.keys(), byte[].class)); + assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A2), String.class)); userState.asyncClose(); @@ -1316,29 +1320,31 @@ public void testOnlyBundleForKeys() throws Exception { @Test public void testHasNoStateAndOnlyBundleForKeys() throws Exception { - FakeBeamFnStateClient fakeClient = new FakeBeamFnStateClient( - ImmutableMap.of( + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapKeyStateKey(), + KV.of(ByteArrayCoder.of(), singletonList(A1)), + createMultimapValueStateKey(A1), + KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); + MultimapUserState userState = + new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", createMultimapKeyStateKey(), - KV.of(ByteArrayCoder.of(), singletonList(A1)), - createMultimapValueStateKey(A1), - KV.of(StringUtf8Coder.of(), asList("V1", "V2")))); - MultimapUserState userState = new MultimapUserState<>( - Caches.noop(), - fakeClient, - "instructionId", - createMultimapKeyStateKey(), - ByteArrayCoder.of(), - StringUtf8Coder.of(), - true /* hasNoState */, - true /* onlyBundleForKeys */); + ByteArrayCoder.of(), + StringUtf8Coder.of(), + true /* hasNoState */, + true /* onlyBundleForKeys */); assertThat(userState.keys(), is(emptyIterable())); assertThat(userState.get(A1), is(emptyIterable())); userState.put(A2, "V2"); - assertArrayEquals(new byte[][] { A2 }, Iterables.toArray(userState.keys(), byte[].class)); - assertArrayEquals(new String[] { "V2" }, Iterables.toArray(userState.get(A2), String.class)); + assertArrayEquals(new byte[][] {A2}, Iterables.toArray(userState.keys(), byte[].class)); + assertArrayEquals(new String[] {"V2"}, Iterables.toArray(userState.get(A2), String.class)); userState.asyncClose(); diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java index 9488e8d5eb4d..759a83afca90 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/OrderedListUserStateTest.java @@ -78,8 +78,8 @@ public void testNoPersistedValues() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertThat(userState.read(), is(emptyIterable())); } @@ -96,8 +96,8 @@ public void testRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( asList(A1, B1).toArray(), Iterables.toArray(userState.read(), TimestampedValue.class)); @@ -122,8 +122,8 @@ public void testReadRange() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterable> stateBeforeB2 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -179,8 +179,8 @@ public void testAdd() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); // add to an existing timestamp userState.add(B1); @@ -226,8 +226,8 @@ public void testClearRange() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterable> initStateFrom2To3 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -296,8 +296,8 @@ public void testClear() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterable> stateBeforeClear = userState.read(); userState.clear(); @@ -330,8 +330,8 @@ public void testAddAndClearRange() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); // add to a non-existing timestamp, clear, and then add userState.add(A2); @@ -403,8 +403,8 @@ public void testAddAndClearRangeAfterClear() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.clear(); userState.clearRange(Instant.ofEpochMilli(0), Instant.ofEpochMilli(5)); @@ -445,8 +445,8 @@ public void testNoopAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.asyncClose(); } @@ -459,8 +459,8 @@ public void testNoopAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( asList(A1, A3, A4).toArray(), @@ -488,8 +488,8 @@ public void testAddAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.add(B1); userState.add(A2); @@ -503,8 +503,8 @@ public void testAddAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( asList(A1, B1, A2, A3, A4).toArray(), @@ -534,8 +534,8 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.clearRange(Instant.ofEpochMilli(1), Instant.ofEpochMilli(3)); userState.clearRange(Instant.ofEpochMilli(4), Instant.ofEpochMilli(5)); @@ -549,8 +549,8 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -576,8 +576,8 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.add(B1); userState.add(A2); @@ -594,8 +594,8 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -625,8 +625,8 @@ public void testClearAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.clear(); userState.asyncClose(); @@ -639,8 +639,8 @@ public void testClearAsyncCloseAndRead() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); assertThat(userState.read(), is(emptyIterable())); } @@ -668,8 +668,8 @@ public void testOperationsDuringNavigatingIterable() throws Exception { "instructionId", createOrderedListStateKey("A"), StringUtf8Coder.of(), - () -> false, - () -> false); + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterator> iter = userState.read().iterator(); assertEquals(iter.next(), A1); From 2574cd9028c9be8b83230f8a1832a370fde39a6f Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 25 Mar 2026 03:02:10 -0400 Subject: [PATCH 29/33] add hasNoState and onlyBundleForKeys to BundleProcess --- .../fn/harness/PTransformRunnerFactory.java | 14 +++------- .../harness/control/ProcessBundleHandler.java | 28 +++++++++++++++---- .../fn/harness/state/FnApiStateAccessor.java | 28 +++++++++---------- 3 files changed, 40 insertions(+), 30 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java index 744e834b8942..db0c69d0abd9 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java @@ -161,17 +161,11 @@ void addIncomingTimerEndpoint( */ DoFn.BundleFinalizer getBundleFinalizer(); - /** - * Returns true if the runner has no state for the keys in the ProcessBundleRequest. If true, - * the SDK can begin stateful processing with an initial empty state. - */ - Supplier getHasNoState(); + /** A supplier containing hasNoState. */ + Supplier getHasNoStateSupplier(); - /** - * Returns true if the runner will never process another bundle for the keys it contains. - * Therefore, the generated state need not be included in the bundle commit. - */ - Supplier getOnlyBundleForKeys(); + /** A supplier containing onlyBundleForKeys. */ + Supplier getOnlyBundleForKeysSupplier(); } /** diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index d942793826c2..262a9db64365 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -247,8 +247,8 @@ private void addRunnerAndConsumersForPTransformRecursively( Collection> channelRoots, Map outboundAggregatorMap, Set runnerCapabilities, - boolean hasNoState, - boolean onlyBundleForKeys) + Supplier hasNoState, + Supplier onlyBundleForKeys) throws IOException { // Recursively ensure that all consumers of the output PCollection have been created. @@ -495,12 +495,12 @@ public BundleFinalizer getBundleFinalizer() { @Override public Supplier getHasNoState() { - return () -> hasNoState; + return hasNoState; } @Override public Supplier getOnlyBundleForKeys() { - return () -> onlyBundleForKeys; + return onlyBundleForKeys; } }); processedPTransformIds.add(pTransformId); @@ -928,8 +928,8 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { bundleProcessor.getChannelRoots(), bundleProcessor.getOutboundAggregators(), bundleProcessor.getRunnerCapabilities(), - processBundleRequest.getHasNoState(), - processBundleRequest.getOnlyBundleForKeys()); + bundleProcessor::getHasNoState, + bundleProcessor::getOnlyBundleForKeys); } bundleProcessor.finish(); @@ -1092,6 +1092,8 @@ public static BundleProcessor create( private String instructionId; private List cacheTokens; private ClearableCache bundleCache; + private boolean hasNoState = false; + private boolean onlyBundleForKeys = false; abstract Cache getProcessWideCache(); @@ -1150,6 +1152,14 @@ synchronized Cache getBundleCache() { return this.bundleCache; } + synchronized boolean getHasNoState() { + return this.hasNoState; + } + + synchronized boolean getOnlyBundleForKeys() { + return this.onlyBundleForKeys; + } + private BeamFnDataInboundObserver inboundObserver2; BeamFnDataInboundObserver getInboundObserver() { @@ -1168,6 +1178,8 @@ void finish() { synchronized void setupForProcessBundleRequest(InstructionRequest request) { this.instructionId = request.getInstructionId(); this.cacheTokens = request.getProcessBundle().getCacheTokensList(); + this.hasNoState = request.getProcessBundle().getHasNoState(); + this.onlyBundleForKeys = request.getProcessBundle().getOnlyBundleForKeys(); getMetricsEnvironmentStateForBundle().start(getStateTracker().getMetricsContainer()); } @@ -1175,6 +1187,8 @@ void reset() throws Exception { synchronized (this) { this.instructionId = null; this.cacheTokens = null; + this.hasNoState = false; + this.onlyBundleForKeys = false; if (this.bundleCache != null) { this.bundleCache.clear(); this.bundleCache = null; @@ -1196,6 +1210,8 @@ void discard() { synchronized (this) { this.instructionId = null; this.cacheTokens = null; + this.hasNoState = false; + this.onlyBundleForKeys = false; if (this.bundleCache != null) { this.bundleCache.clear(); } 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 e5c686d3f23f..aa6216b1bd1a 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 @@ -227,8 +227,8 @@ public static Factory factoryForPTransformContext( context.getBeamFnStateClient(), keyCoder, windowCoder, - context.getHasNoState(), - context.getOnlyBundleForKeys()); + context.getHasNoStateSupplier(), + context.getOnlyBundleForKeysSupplier()); } public FnApiStateAccessor create() { @@ -244,8 +244,8 @@ public FnApiStateAccessor create() { beamFnStateClient, keyCoder, windowCoder, - hasNoState.get(), - onlyBundleForKeys.get()); + hasNoState, + onlyBundleForKeys); } } @@ -262,8 +262,8 @@ public FnApiStateAccessor create() { private final Collection stateFinalizers; private final Coder keyCoder; private final Coder windowCoder; - private final boolean hasNoState; - private final boolean onlyBundleForKeys; + private final Supplier hasNoState; + private final Supplier onlyBundleForKeys; private @Nullable Supplier currentWindowSupplier; private @Nullable Supplier encodedCurrentKeySupplier; @@ -281,8 +281,8 @@ public FnApiStateAccessor( BeamFnStateClient beamFnStateClient, Coder keyCoder, Coder windowCoder, - boolean hasNoState, - boolean onlyBundleForKeys) { + Supplier hasNoState, + Supplier onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.stateKeyObjectCache = Maps.newHashMap(); @@ -1218,8 +1218,8 @@ private BagUserState createBagUserState(StateKey stateKey, Coder value processBundleInstructionId.get(), stateKey, valueCoder, - hasNoState, - onlyBundleForKeys); + hasNoState.get(), + onlyBundleForKeys.get()); stateFinalizers.add(rval::asyncClose); return rval; } @@ -1302,8 +1302,8 @@ private MultimapUserState createMultimapUserState( stateKey, keyCoder, valueCoder, - hasNoState, - onlyBundleForKeys); + hasNoState.get(), + onlyBundleForKeys.get()); stateFinalizers.add(rval::asyncClose); return rval; } @@ -1339,8 +1339,8 @@ private OrderedListUserState createOrderedListUserState( processBundleInstructionId.get(), stateKey, valueCoder, - hasNoState, - onlyBundleForKeys); + hasNoState.get(), + onlyBundleForKeys.get()); stateFinalizers.add(rval::asyncClose); return rval; } From 4e4ee12f699cfde8e52772ceb8c74f8c309efac3 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 25 Mar 2026 03:36:36 -0400 Subject: [PATCH 30/33] unit test --- .../harness/control/ProcessBundleHandler.java | 4 +- .../control/ProcessBundleHandlerTest.java | 66 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java index 262a9db64365..39cc5e7a47f7 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java @@ -494,12 +494,12 @@ public BundleFinalizer getBundleFinalizer() { } @Override - public Supplier getHasNoState() { + public Supplier getHasNoStateSupplier() { return hasNoState; } @Override - public Supplier getOnlyBundleForKeys() { + public Supplier getOnlyBundleForKeysSupplier() { return onlyBundleForKeys; } }); 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 52b6c87a5c05..211ce574cf14 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 @@ -2060,6 +2060,72 @@ public void testStartFinishBundleMetrics() throws Exception { assertEquals(1, finishCounterFound); } + @Test + public void testStateFlagsUpdatedOnReuse() throws Exception { + BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = + BeamFnApi.ProcessBundleDescriptor.newBuilder() + .putTransforms( + "2L", + RunnerApi.PTransform.newBuilder() + .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()) + .build()) + .build(); + Map fnApiRegistry = + ImmutableMap.of("1L", processBundleDescriptor); + + List observedHasNoState = new ArrayList<>(); + List observedOnlyBundleForKeys = new ArrayList<>(); + + PTransformRunnerFactory flagRecorder = + (context) -> { + Supplier hasNoStateSupplier = context.getHasNoStateSupplier(); + Supplier onlyBundleForKeysSupplier = context.getOnlyBundleForKeysSupplier(); + context.addStartBundleFunction( + () -> { + observedHasNoState.add(hasNoStateSupplier.get()); + observedOnlyBundleForKeys.add(onlyBundleForKeysSupplier.get()); + }); + }; + + ProcessBundleHandler handler = + new ProcessBundleHandler( + PipelineOptionsFactory.create(), + Collections.emptySet(), + fnApiRegistry::get, + beamFnDataClient, + null, + null, + new ShortIdMap(), + executionStateSampler, + ImmutableMap.of(DATA_INPUT_URN, flagRecorder), + Caches.noop(), + new BundleProcessorCache(Duration.ZERO), + null); + + handler.processBundle( + BeamFnApi.InstructionRequest.newBuilder() + .setInstructionId("100L") + .setProcessBundle( + BeamFnApi.ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("1L") + .setHasNoState(true) + .setOnlyBundleForKeys(false)) + .build()); + + handler.processBundle( + BeamFnApi.InstructionRequest.newBuilder() + .setInstructionId("101L") + .setProcessBundle( + BeamFnApi.ProcessBundleRequest.newBuilder() + .setProcessBundleDescriptorId("1L") + .setHasNoState(false) + .setOnlyBundleForKeys(true)) + .build()); + + assertThat(observedHasNoState, contains(true, false)); + assertThat(observedOnlyBundleForKeys, contains(false, true)); + } + private static void throwException() { throw new IllegalStateException("TestException"); } From 85981883888455d0cfe9615ba69090cc7d16e556 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 25 Mar 2026 03:56:58 -0400 Subject: [PATCH 31/33] test context --- .../fn/harness/PTransformRunnerFactoryTestContext.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java index a5f335bd56eb..dddbda758bfe 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java @@ -156,8 +156,8 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { } }) .runnerCapabilities(new HashSet<>()) - .hasNoState(() -> false) - .onlyBundleForKeys(() -> false); + .hasNoStateSupplier(() -> false) + .onlyBundleForKeysSupplier(() -> false); } /** A builder to create a context for tests. */ @@ -225,9 +225,9 @@ default Builder processBundleInstructionId(String value) { Builder timerApiServiceDescriptor(ApiServiceDescriptor value); - Builder hasNoState(Supplier value); + Builder hasNoStateSupplier(Supplier value); - Builder onlyBundleForKeys(Supplier value); + Builder onlyBundleForKeysSupplier(Supplier value); PTransformRunnerFactoryTestContext build(); } From 6e8b2c4f7916d857c4482f9c751542759bc59c7f Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 25 Mar 2026 04:07:24 -0400 Subject: [PATCH 32/33] fix --- .../beam/fn/harness/PTransformRunnerFactoryTestContext.java | 6 ------ 1 file changed, 6 deletions(-) diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java index dddbda758bfe..e38679d9ee26 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java @@ -64,12 +64,6 @@ public abstract class PTransformRunnerFactoryTestContext implements PTransformRunnerFactory.Context { - @Override - public abstract Supplier getHasNoState(); - - @Override - public abstract Supplier getOnlyBundleForKeys(); - /** Returns a builder for the specified PTransform id and PTransform definition. */ public static Builder builder(String pTransformId, RunnerApi.PTransform pTransform) { return new AutoValue_PTransformRunnerFactoryTestContext.Builder() From 06a1365aacc84f242d0269b0dac65d60489f3d89 Mon Sep 17 00:00:00 2001 From: Kai Li Date: Wed, 14 Oct 2026 00:57:03 -0400 Subject: [PATCH 33/33] Fix backwards compatibility of constructors and remove CachingStateIterable empty stub --- .../fn/harness/PTransformRunnerFactory.java | 8 ++- .../beam/fn/harness/state/BagUserState.java | 35 ++++++++---- .../fn/harness/state/FnApiStateAccessor.java | 56 +++++++++++++++++++ .../fn/harness/state/MultimapUserState.java | 29 +++++++--- .../harness/state/OrderedListUserState.java | 9 +++ .../harness/state/StateFetchingIterators.java | 55 ------------------ .../PTransformRunnerFactoryTestContext.java | 6 ++ 7 files changed, 124 insertions(+), 74 deletions(-) diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java index db0c69d0abd9..840973d5bc04 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java @@ -162,10 +162,14 @@ void addIncomingTimerEndpoint( DoFn.BundleFinalizer getBundleFinalizer(); /** A supplier containing hasNoState. */ - Supplier getHasNoStateSupplier(); + default Supplier getHasNoStateSupplier() { + return () -> false; + } /** A supplier containing onlyBundleForKeys. */ - Supplier getOnlyBundleForKeysSupplier(); + default Supplier getOnlyBundleForKeysSupplier() { + return () -> false; + } } /** diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index 98258de6a544..ef3adbed5b0d 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java @@ -18,6 +18,7 @@ package org.apache.beam.fn.harness.state; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.util.ArrayList; @@ -34,6 +35,7 @@ import org.apache.beam.sdk.fn.stream.PrefetchableIterables; import org.apache.beam.sdk.util.ByteStringOutputStream; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; /** * An implementation of a bag user state that utilizes the Beam Fn State API to fetch, clear and @@ -50,7 +52,7 @@ public class BagUserState { private final BeamFnStateClient beamFnStateClient; private final StateRequest request; private final Coder valueCoder; - private final CachingStateIterable oldValues; + private final @Nullable CachingStateIterable oldValues; private List newValues; private boolean isCleared; private boolean isClosed; @@ -59,6 +61,16 @@ public class BagUserState { static final int BAG_APPEND_BATCHING_LIMIT = 10 * 1024 * 1024; + /** The cache must be namespaced for this state object accordingly. */ + public BagUserState( + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder valueCoder) { + this(cache, beamFnStateClient, instructionId, stateKey, valueCoder, false, false); + } + /** The cache must be namespaced for this state object accordingly. */ public BagUserState( Cache cache, @@ -80,7 +92,7 @@ public BagUserState( this.oldValues = hasNoState - ? CachingStateIterable.emptyIterable() + ? null : StateFetchingIterators.readAllAndDecodeStartingFrom( this.cache, beamFnStateClient, request, valueCoder); this.newValues = new ArrayList<>(); @@ -96,10 +108,11 @@ public PrefetchableIterable get() { return PrefetchableIterables.limit(Collections.unmodifiableList(newValues), newValues.size()); } else if (newValues.isEmpty()) { // If we have no new values then just return the old values. - return oldValues; + return checkNotNull(oldValues); } return PrefetchableIterables.concat( - oldValues, Iterables.limit(Collections.unmodifiableList(newValues), newValues.size())); + checkNotNull(oldValues), + Iterables.limit(Collections.unmodifiableList(newValues), newValues.size())); } public void append(T t) { @@ -169,12 +182,14 @@ public void asyncClose() throws Exception { } // Modify the underlying cached state depending on the mutations performed - if (isCleared) { - // Note this takes ownership of newValues. This object is no longer used after it has been - // closed. - oldValues.clearAndAppend(newValues); - } else { - oldValues.append(newValues); + if (oldValues != null) { + if (isCleared) { + // Note this takes ownership of newValues. This object is no longer used after it has been + // closed. + oldValues.clearAndAppend(newValues); + } else { + oldValues.append(newValues); + } } } } 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 e4d074540245..8205bfaaefd4 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 @@ -117,6 +117,34 @@ public static class Factory { private final Supplier hasNoState; private final Supplier onlyBundleForKeys; + public Factory( + PipelineOptions pipelineOptions, + Set runnerCapabilities, + String ptransformId, + Supplier processBundleInstructionId, + Supplier> cacheTokens, + Supplier> bundleCache, + Cache processWideCache, + Map, SideInputSpec> sideInputSpecMap, + BeamFnStateClient beamFnStateClient, + Coder keyCoder, + Coder windowCoder) { + this( + pipelineOptions, + runnerCapabilities, + ptransformId, + processBundleInstructionId, + cacheTokens, + bundleCache, + processWideCache, + sideInputSpecMap, + beamFnStateClient, + keyCoder, + windowCoder, + () -> false, + () -> false); + } + public Factory( PipelineOptions pipelineOptions, Set runnerCapabilities, @@ -269,6 +297,34 @@ public FnApiStateAccessor create() { private @Nullable Supplier encodedCurrentKeySupplier; private @Nullable Supplier encodedCurrentWindowSupplier; + public FnApiStateAccessor( + PipelineOptions pipelineOptions, + Set runnerCapabilities, + String ptransformId, + Supplier processBundleInstructionId, + Supplier> cacheTokens, + Supplier> bundleCache, + Cache processWideCache, + Map, SideInputSpec> sideInputSpecMap, + BeamFnStateClient beamFnStateClient, + Coder keyCoder, + Coder windowCoder) { + this( + pipelineOptions, + runnerCapabilities, + ptransformId, + processBundleInstructionId, + cacheTokens, + bundleCache, + processWideCache, + sideInputSpecMap, + beamFnStateClient, + keyCoder, + windowCoder, + () -> false, + () -> false); + } + public FnApiStateAccessor( PipelineOptions pipelineOptions, Set runnerCapabilities, 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 d639cdc66d53..e278ec58ce20 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 @@ -50,6 +50,7 @@ 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.Nullable; /** * An implementation of a multimap user state that utilizes the Beam Fn State API to fetch, clear @@ -67,8 +68,8 @@ public class MultimapUserState { private final StateRequest keysStateRequest; private final StateRequest entriesStateRequest; private final StateRequest userStateRequest; - private final CachingStateIterable persistedKeys; - private final CachingStateIterable>> persistedEntries; + private final @Nullable CachingStateIterable persistedKeys; + private final @Nullable CachingStateIterable>> persistedEntries; private boolean isClosed; private boolean isCleared; @@ -80,6 +81,16 @@ public class MultimapUserState { // Values retrieved from persistent storage private HashMap>> persistedValues = Maps.newHashMap(); + public MultimapUserState( + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder mapKeyCoder, + Coder valueCoder) { + this(cache, beamFnStateClient, instructionId, stateKey, mapKeyCoder, valueCoder, false, false); + } + public MultimapUserState( Cache cache, BeamFnStateClient beamFnStateClient, @@ -106,7 +117,7 @@ public MultimapUserState( StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = hasNoState - ? CachingStateIterable.emptyIterable() + ? null : StateFetchingIterators.readAllAndDecodeStartingFrom( cache, beamFnStateClient, keysStateRequest, mapKeyCoder); @@ -133,7 +144,7 @@ public MultimapUserState( this.entriesStateRequest = entriesStateRequestBuilder.build(); this.persistedEntries = hasNoState - ? CachingStateIterable.emptyIterable() + ? null : StateFetchingIterators.readAllAndDecodeStartingFrom( Caches.subCache(this.cache, "AllEntries"), beamFnStateClient, @@ -468,7 +479,9 @@ private void updateCache() { if (isCleared) { // This will clear all keys and values since values is a sub-cache of keys. Note this // takes ownership of pendingAddKeys. This object is no longer used after it has been closed. - persistedKeys.clearAndAppend(pendingAddsKeys); + if (persistedKeys != null) { + persistedKeys.clearAndAppend(pendingAddsKeys); + } // Since the map was cleared we can add all the values that are pending since we know // that they must have been cleared. @@ -483,8 +496,10 @@ private void updateCache() { // The cast to Set is necessary since the checker framework would like to further // limit the type to Set<@KeyFor("this.pendingRemoves") Object> which is incompatible with // the API being remove(Set). We don't want to limit the API for remove either. - persistedKeys.remove((Set) pendingRemoves.keySet()); - persistedKeys.append(pendingAddsKeys); + if (persistedKeys != null) { + persistedKeys.remove((Set) pendingRemoves.keySet()); + persistedKeys.append(pendingAddsKeys); + } // For each removed key, we want to update the internal cache to clear its set of values for (Map.Entry entry : pendingRemoves.entrySet()) { diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java index 49c25ee449a5..b305c067a356 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/OrderedListUserState.java @@ -158,6 +158,15 @@ public List> getComponents() { } } + public OrderedListUserState( + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder valueCoder) { + this(cache, beamFnStateClient, instructionId, stateKey, valueCoder, false, false); + } + public OrderedListUserState( Cache cache, BeamFnStateClient beamFnStateClient, diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java index 0c19a79435e0..1b4d7ec371ba 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java @@ -572,61 +572,6 @@ private void appendHelper(List newValues, long newWeight) { cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(builder.build())); } - private static final PrefetchableIterator EMPTY_ITERATOR = - new PrefetchableIterator() { - @Override - public boolean isReady() { - return true; - } - - @Override - public void prefetch() {} - - @Override - public boolean hasNext() { - return false; - } - - @Override - public Object next() { - throw new NoSuchElementException(); - } - }; - - @SuppressWarnings("unchecked") - public static PrefetchableIterator emptyIterator() { - return (PrefetchableIterator) EMPTY_ITERATOR; - } - - private static final CachingStateIterable EMPTY_ITERABLE = - new CachingStateIterable(null, null, null, null) { - @Override - public PrefetchableIterator createIterator() { - return emptyIterator(); - } - - @Override - public void remove(Set toRemoveStructuralValues) {} - - @Override - public void clearAndAppend(List values) {} - - @Override - public void clearAndAppend(WeightedList values) {} - - @Override - public void append(List values) {} - - @Override - public void append(WeightedList values) {} - }; - - /** Returns an empty {@link CachingStateIterable}. */ - @SuppressWarnings("unchecked") - public static CachingStateIterable emptyIterable() { - return (CachingStateIterable) EMPTY_ITERABLE; - } - class CachingStateIterator implements PrefetchableIterator { private final LazyBlockingStateFetchingIterator underlyingStateFetchingIterator; diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java index b7d6cf37c4d3..a6af66761e96 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/PTransformRunnerFactoryTestContext.java @@ -65,6 +65,12 @@ public abstract class PTransformRunnerFactoryTestContext implements PTransformRunnerFactory.Context { + @Override + public abstract Supplier getHasNoStateSupplier(); + + @Override + public abstract Supplier getOnlyBundleForKeysSupplier(); + /** Returns a builder for the specified PTransform id and PTransform definition. */ public static Builder builder(String pTransformId, RunnerApi.PTransform pTransform) { return new AutoValue_PTransformRunnerFactoryTestContext.Builder()