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..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 @@ -160,6 +160,16 @@ void addIncomingTimerEndpoint( * instant provides the timeout on how long the finalization callback is valid for. */ DoFn.BundleFinalizer getBundleFinalizer(); + + /** A supplier containing hasNoState. */ + default Supplier getHasNoStateSupplier() { + return () -> false; + } + + /** A supplier containing onlyBundleForKeys. */ + default Supplier getOnlyBundleForKeysSupplier() { + return () -> false; + } } /** 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 f12340552fd0..9197d8359a2a 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 @@ -245,7 +245,9 @@ private void addRunnerAndConsumersForPTransformRecursively( BundleFinalizer bundleFinalizer, Collection> channelRoots, Map outboundAggregatorMap, - Set runnerCapabilities) + Set runnerCapabilities, + Supplier hasNoState, + Supplier onlyBundleForKeys) throws IOException { // Recursively ensure that all consumers of the output PCollection have been created. @@ -277,7 +279,9 @@ private void addRunnerAndConsumersForPTransformRecursively( bundleFinalizer, channelRoots, outboundAggregatorMap, - runnerCapabilities); + runnerCapabilities, + hasNoState, + onlyBundleForKeys); } } @@ -483,6 +487,16 @@ public BundleSplitListener getSplitListener() { public BundleFinalizer getBundleFinalizer() { return bundleFinalizer; } + + @Override + public Supplier getHasNoStateSupplier() { + return hasNoState; + } + + @Override + public Supplier getOnlyBundleForKeysSupplier() { + return onlyBundleForKeys; + } }); processedPTransformIds.add(pTransformId); } @@ -924,7 +938,9 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { bundleFinalizer, bundleProcessor.getChannelRoots(), bundleProcessor.getOutboundAggregators(), - bundleProcessor.getRunnerCapabilities()); + bundleProcessor.getRunnerCapabilities(), + bundleProcessor::getHasNoState, + bundleProcessor::getOnlyBundleForKeys); } bundleProcessor.finish(); @@ -1087,6 +1103,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(); @@ -1145,6 +1163,14 @@ synchronized Cache getBundleCache() { return this.bundleCache; } + synchronized boolean getHasNoState() { + return this.hasNoState; + } + + synchronized boolean getOnlyBundleForKeys() { + return this.onlyBundleForKeys; + } + private BeamFnDataInboundObserver inboundObserver2; BeamFnDataInboundObserver getInboundObserver() { @@ -1163,6 +1189,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()); } @@ -1170,6 +1198,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; @@ -1191,6 +1221,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/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java index 7f61008bc66b..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,10 +52,12 @@ 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; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; static final int BAG_APPEND_BATCHING_LIMIT = 10 * 1024 * 1024; @@ -64,17 +68,33 @@ public BagUserState( 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, + 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); + hasNoState + ? null + : StateFetchingIterators.readAllAndDecodeStartingFrom( + this.cache, beamFnStateClient, request, valueCoder); this.newValues = new ArrayList<>(); } @@ -83,15 +103,16 @@ 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. - 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) { @@ -120,7 +141,7 @@ 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) { @@ -161,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 8d02e424adc5..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 @@ -114,6 +114,8 @@ public static class Factory { private final Map, SideInputSpec> sideInputSpecMap; private final Coder keyCoder; private final Coder windowCoder; + private final Supplier hasNoState; + private final Supplier onlyBundleForKeys; public Factory( PipelineOptions pipelineOptions, @@ -127,6 +129,36 @@ public Factory( 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, + String ptransformId, + Supplier processBundleInstructionId, + Supplier> cacheTokens, + Supplier> bundleCache, + Cache processWideCache, + Map, SideInputSpec> sideInputSpecMap, + BeamFnStateClient beamFnStateClient, + Coder keyCoder, + Coder windowCoder, + Supplier hasNoState, + Supplier onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.ptransformId = ptransformId; @@ -138,6 +170,8 @@ public Factory( this.beamFnStateClient = beamFnStateClient; this.keyCoder = keyCoder; this.windowCoder = windowCoder; + this.hasNoState = hasNoState; + this.onlyBundleForKeys = onlyBundleForKeys; } public static Factory factoryForPTransformContext( @@ -220,7 +254,9 @@ public static Factory factoryForPTransformContext( tagToSideInputSpecMap, context.getBeamFnStateClient(), keyCoder, - windowCoder); + windowCoder, + context.getHasNoStateSupplier(), + context.getOnlyBundleForKeysSupplier()); } public FnApiStateAccessor create() { @@ -235,7 +271,9 @@ public FnApiStateAccessor create() { sideInputSpecMap, beamFnStateClient, keyCoder, - windowCoder); + windowCoder, + hasNoState, + onlyBundleForKeys); } } @@ -252,6 +290,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 @Nullable Supplier currentWindowSupplier; private @Nullable Supplier encodedCurrentKeySupplier; @@ -269,6 +309,36 @@ public FnApiStateAccessor( 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, + String ptransformId, + Supplier processBundleInstructionId, + Supplier> cacheTokens, + Supplier> bundleCache, + Cache processWideCache, + Map, SideInputSpec> sideInputSpecMap, + BeamFnStateClient beamFnStateClient, + Coder keyCoder, + Coder windowCoder, + Supplier hasNoState, + Supplier onlyBundleForKeys) { this.pipelineOptions = pipelineOptions; this.runnerCapabilities = runnerCapabilities; this.stateKeyObjectCache = Maps.newHashMap(); @@ -282,6 +352,8 @@ public FnApiStateAccessor( this.keyCoder = keyCoder; this.windowCoder = windowCoder; this.stateFinalizers = new ArrayList<>(); + this.hasNoState = hasNoState; + this.onlyBundleForKeys = onlyBundleForKeys; } public void setKeyAndWindowContext(MutatingStateContext keyAndWindowContext) { @@ -1201,7 +1273,9 @@ private BagUserState createBagUserState(StateKey stateKey, Coder value beamFnStateClient, processBundleInstructionId.get(), stateKey, - valueCoder); + valueCoder, + hasNoState.get(), + onlyBundleForKeys.get()); stateFinalizers.add(rval::asyncClose); return rval; } @@ -1283,7 +1357,9 @@ private MultimapUserState createMultimapUserState( processBundleInstructionId.get(), stateKey, keyCoder, - valueCoder); + valueCoder, + hasNoState.get(), + onlyBundleForKeys.get()); stateFinalizers.add(rval::asyncClose); return rval; } @@ -1318,7 +1394,9 @@ private OrderedListUserState createOrderedListUserState( beamFnStateClient, processBundleInstructionId.get(), stateKey, - valueCoder); + valueCoder, + hasNoState.get(), + onlyBundleForKeys.get()); stateFinalizers.add(rval::asyncClose); return rval; } 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 c62b7ab9c8ad..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,11 +68,13 @@ 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; + private final boolean hasNoState; + private final boolean onlyBundleForKeys; // Pending updates to persistent storage private HashMap pendingRemoves = Maps.newHashMap(); private HashMap>> pendingAdds = Maps.newHashMap(); @@ -85,6 +88,18 @@ public MultimapUserState( StateKey stateKey, Coder mapKeyCoder, Coder valueCoder) { + this(cache, beamFnStateClient, instructionId, stateKey, mapKeyCoder, valueCoder, false, false); + } + + public MultimapUserState( + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder mapKeyCoder, + Coder valueCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasMultimapKeysUserState(), "Expected MultimapKeysUserState StateKey but received %s.", @@ -93,14 +108,18 @@ 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. this.keysStateRequest = StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build(); this.persistedKeys = - StateFetchingIterators.readAllAndDecodeStartingFrom( - cache, beamFnStateClient, keysStateRequest, mapKeyCoder); + hasNoState + ? null + : StateFetchingIterators.readAllAndDecodeStartingFrom( + cache, beamFnStateClient, keysStateRequest, mapKeyCoder); StateRequest.Builder userStateRequestBuilder = StateRequest.newBuilder(); userStateRequestBuilder @@ -124,11 +143,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 + ? null + : StateFetchingIterators.readAllAndDecodeStartingFrom( + Caches.subCache(this.cache, "AllEntries"), + beamFnStateClient, + entriesStateRequest, + KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder))); } public void clear() { @@ -161,7 +182,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; } @@ -179,7 +200,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()); @@ -277,7 +298,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( @@ -415,7 +436,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; } @@ -458,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. @@ -473,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 47b5057880b9..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 @@ -88,6 +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> { @@ -162,6 +164,17 @@ public OrderedListUserState( String instructionId, StateKey stateKey, Coder valueCoder) { + this(cache, beamFnStateClient, instructionId, stateKey, valueCoder, false, false); + } + + public OrderedListUserState( + Cache cache, + BeamFnStateClient beamFnStateClient, + String instructionId, + StateKey stateKey, + Coder valueCoder, + boolean hasNoState, + boolean onlyBundleForKeys) { checkArgument( stateKey.hasOrderedListUserState(), "Expected OrderedListUserState StateKey but received %s.", @@ -170,6 +183,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) { @@ -203,7 +218,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() @@ -283,6 +298,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()) { 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 51e49953b406..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() @@ -151,7 +157,9 @@ public void afterBundleCommit(Instant callbackExpiry, Callback callback) { throw new UnsupportedOperationException("Unexpected call during test."); } }) - .runnerCapabilities(new HashSet<>()); + .runnerCapabilities(new HashSet<>()) + .hasNoStateSupplier(() -> false) + .onlyBundleForKeysSupplier(() -> false); } /** A builder to create a context for tests. */ @@ -219,6 +227,10 @@ default Builder processBundleInstructionId(String value) { Builder timerApiServiceDescriptor(ApiServiceDescriptor value); + Builder hasNoStateSupplier(Supplier value); + + Builder onlyBundleForKeysSupplier(Supplier value); + PTransformRunnerFactoryTestContext build(); } 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 c03f82726740..ad0bb68e7872 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 @@ -2056,6 +2056,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"); } 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..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 @@ -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 /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); @@ -67,7 +73,14 @@ 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 /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -83,7 +96,9 @@ public void testGetCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); @@ -96,7 +111,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 /* hasNoState */, + false /* onlyBundleForKeys */); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -118,7 +139,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 /* hasNoState */, + false /* onlyBundleForKeys */); String a2 = Strings.repeat("A2", 2 * 1024 * 1024); userState.append(a2); String a3 = Strings.repeat("A3", 4 * 1024 * 1024); @@ -155,7 +182,14 @@ 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 /* hasNoState */, + false /* onlyBundleForKeys */); userState.append("A2"); Iterable stateBeforeA3 = userState.get(); assertArrayEquals(new String[] {"A1", "A2"}, Iterables.toArray(stateBeforeA3, String.class)); @@ -180,7 +214,9 @@ public void testAppendCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.append("A4"); Iterable stateBeforeA5 = userState.get(); assertArrayEquals( @@ -203,7 +239,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 /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -228,7 +270,14 @@ 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 /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( new String[] {"A1", "A2", "A3"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); @@ -252,7 +301,9 @@ public void testClearCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals(new String[] {"A4"}, Iterables.toArray(userState.get(), String.class)); userState.clear(); assertFalse(userState.get().iterator().hasNext()); @@ -273,7 +324,9 @@ public void testClearCached() throws Exception { }, "instructionId", key("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals(new String[] {}, Iterables.toArray(userState.get(), String.class)); userState.asyncClose(); } @@ -298,4 +351,60 @@ private ByteString encode(String... values) throws IOException { } return out.toByteString(); } + + @Test + public void testHasNoState() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + StringUtf8Coder.of(), ImmutableMap.of(key("A"), asList("A1", "A2", "A3"))); + BagUserState 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/MultimapUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java index 679307321826..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 @@ -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; @@ -84,7 +85,9 @@ public void testNoPersistedValues() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); assertThat(userState.keys(), is(emptyIterable())); } @@ -104,7 +107,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 +137,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 +173,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 +209,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 +286,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 +315,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 +341,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 +365,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 +385,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 +408,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 +439,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 +469,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 +494,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 +519,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 +545,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 +571,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 +602,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 +622,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 +646,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 +674,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 +702,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 +731,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 +766,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 +803,9 @@ public void testClearPrefetch() throws Exception { "instructionId", createMultimapKeyStateKey(), ByteArrayCoder.of(), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false, + false); userState.clear(); userState.put(A2, "V3"); @@ -784,7 +835,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 +861,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 +879,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 +906,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 +926,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 +955,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 +976,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 +1004,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 +1028,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 +1058,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 +1081,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 +1115,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 +1136,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 +1163,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 +1183,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 +1212,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 +1235,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())); @@ -1160,6 +1245,112 @@ 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( 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..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 @@ -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; @@ -76,7 +77,9 @@ public void testNoPersistedValues() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertThat(userState.read(), is(emptyIterable())); } @@ -92,7 +95,9 @@ public void testRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( asList(A1, B1).toArray(), Iterables.toArray(userState.read(), TimestampedValue.class)); @@ -116,7 +121,9 @@ public void testReadRange() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterable> stateBeforeB2 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -171,7 +178,9 @@ public void testAdd() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); // add to an existing timestamp userState.add(B1); @@ -216,7 +225,9 @@ public void testClearRange() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterable> initStateFrom2To3 = userState.readRange(Instant.ofEpochMilli(2), Instant.ofEpochMilli(4)); @@ -284,7 +295,9 @@ public void testClear() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterable> stateBeforeClear = userState.read(); userState.clear(); @@ -316,7 +329,9 @@ public void testAddAndClearRange() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); // add to a non-existing timestamp, clear, and then add userState.add(A2); @@ -387,7 +402,9 @@ public void testAddAndClearRangeAfterClear() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.clear(); userState.clearRange(Instant.ofEpochMilli(0), Instant.ofEpochMilli(5)); @@ -427,7 +444,9 @@ public void testNoopAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.asyncClose(); } @@ -439,7 +458,9 @@ public void testNoopAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( asList(A1, A3, A4).toArray(), @@ -466,7 +487,9 @@ public void testAddAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.add(B1); userState.add(A2); @@ -479,7 +502,9 @@ public void testAddAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( asList(A1, B1, A2, A3, A4).toArray(), @@ -508,7 +533,9 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.clearRange(Instant.ofEpochMilli(1), Instant.ofEpochMilli(3)); userState.clearRange(Instant.ofEpochMilli(4), Instant.ofEpochMilli(5)); @@ -521,7 +548,9 @@ public void testClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -546,7 +575,9 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.add(B1); userState.add(A2); @@ -562,7 +593,9 @@ public void testAddClearRangeAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertArrayEquals( Collections.singletonList(A3).toArray(), @@ -591,7 +624,9 @@ public void testClearAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); userState.clear(); userState.asyncClose(); @@ -603,7 +638,9 @@ public void testClearAsyncCloseAndRead() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); assertThat(userState.read(), is(emptyIterable())); } @@ -630,7 +667,9 @@ public void testOperationsDuringNavigatingIterable() throws Exception { fakeClient, "instructionId", createOrderedListStateKey("A"), - StringUtf8Coder.of()); + StringUtf8Coder.of(), + false /* hasNoState */, + false /* onlyBundleForKeys */); Iterator> iter = userState.read().iterator(); assertEquals(iter.next(), A1); @@ -650,6 +689,147 @@ 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 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 = + 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))); + } + + @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) {