Omit state commits in stateful Dofn

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/PTransformRunnerFactory.java

@@ -160,6 +160,16 @@ <T> void addIncomingTimerEndpoint(
      * instant provides the timeout on how long the finalization callback is valid for.
      */
     DoFn.BundleFinalizer getBundleFinalizer();
+
+    /** A supplier containing hasNoState. */
+    default Supplier<Boolean> getHasNoStateSupplier() {
+      return () -> false;
+    }
+
+    /** A supplier containing onlyBundleForKeys. */
+    default Supplier<Boolean> getOnlyBundleForKeysSupplier() {
+      return () -> false;
+    }
   }
 
   /**

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java

@@ -245,7 +245,9 @@ private void addRunnerAndConsumersForPTransformRecursively(
       BundleFinalizer bundleFinalizer,
       Collection<BeamFnDataReadRunner<?>> channelRoots,
       Map<ApiServiceDescriptor, BeamFnDataOutboundAggregator> outboundAggregatorMap,
-      Set<String> runnerCapabilities)
+      Set<String> runnerCapabilities,
+      Supplier<Boolean> hasNoState,
+      Supplier<Boolean> 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<Boolean> getHasNoStateSupplier() {
+                  return hasNoState;
+                }
+
+                @Override
+                public Supplier<Boolean> 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<CacheToken> cacheTokens;
     private ClearableCache<Object, Object> bundleCache;
+    private boolean hasNoState = false;
+    private boolean onlyBundleForKeys = false;
 
     abstract Cache<?, ?> getProcessWideCache();
 
@@ -1145,6 +1163,14 @@ synchronized Cache<Object, Object> getBundleCache() {
       return this.bundleCache;
     }
 
+    synchronized boolean getHasNoState() {
+      return this.hasNoState;
+    }
+
+    synchronized boolean getOnlyBundleForKeys() {
+      return this.onlyBundleForKeys;
+    }
+
     private BeamFnDataInboundObserver inboundObserver2;
 
     BeamFnDataInboundObserver getInboundObserver() {
@@ -1163,13 +1189,17 @@ 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());
     }
 
     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();
         }

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<T> {
   private final BeamFnStateClient beamFnStateClient;
   private final StateRequest request;
   private final Coder<T> valueCoder;
-  private final CachingStateIterable<T> oldValues;
+  private final @Nullable CachingStateIterable<T> oldValues;
   private List<T> 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<T> 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<T> 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<T> 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);
+      }
     }
   }
 }

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java

@@ -114,6 +114,8 @@ public static class Factory<K> {
     private final Map<TupleTag<?>, SideInputSpec> sideInputSpecMap;
     private final Coder<K> keyCoder;
     private final Coder<BoundedWindow> windowCoder;
+    private final Supplier<Boolean> hasNoState;
+    private final Supplier<Boolean> onlyBundleForKeys;
 
     public Factory(
         PipelineOptions pipelineOptions,
@@ -127,6 +129,36 @@ public Factory(
         BeamFnStateClient beamFnStateClient,
         Coder<K> keyCoder,
         Coder<BoundedWindow> windowCoder) {
+      this(
+          pipelineOptions,
+          runnerCapabilities,
+          ptransformId,
+          processBundleInstructionId,
+          cacheTokens,
+          bundleCache,
+          processWideCache,
+          sideInputSpecMap,
+          beamFnStateClient,
+          keyCoder,
+          windowCoder,
+          () -> false,
+          () -> false);
+    }
+
+    public Factory(
+        PipelineOptions pipelineOptions,
+        Set<String> runnerCapabilities,
+        String ptransformId,
+        Supplier<String> processBundleInstructionId,
+        Supplier<List<CacheToken>> cacheTokens,
+        Supplier<Cache<?, ?>> bundleCache,
+        Cache<?, ?> processWideCache,
+        Map<TupleTag<?>, SideInputSpec> sideInputSpecMap,
+        BeamFnStateClient beamFnStateClient,
+        Coder<K> keyCoder,
+        Coder<BoundedWindow> windowCoder,
+        Supplier<Boolean> hasNoState,
+        Supplier<Boolean> 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 <K> Factory<K> factoryForPTransformContext(
@@ -220,7 +254,9 @@ public static <K> Factory<K> factoryForPTransformContext(
           tagToSideInputSpecMap,
           context.getBeamFnStateClient(),
           keyCoder,
-          windowCoder);
+          windowCoder,
+          context.getHasNoStateSupplier(),
+          context.getOnlyBundleForKeysSupplier());
     }
 
     public FnApiStateAccessor<K> create() {
@@ -235,7 +271,9 @@ public FnApiStateAccessor<K> create() {
           sideInputSpecMap,
           beamFnStateClient,
           keyCoder,
-          windowCoder);
+          windowCoder,
+          hasNoState,
+          onlyBundleForKeys);
     }
   }
 
@@ -252,6 +290,8 @@ public FnApiStateAccessor<K> create() {
   private final Collection<ThrowingRunnable> stateFinalizers;
   private final Coder<K> keyCoder;
   private final Coder<BoundedWindow> windowCoder;
+  private final Supplier<Boolean> hasNoState;
+  private final Supplier<Boolean> onlyBundleForKeys;
 
   private @Nullable Supplier<BoundedWindow> currentWindowSupplier;
   private @Nullable Supplier<ByteString> encodedCurrentKeySupplier;
@@ -269,6 +309,36 @@ public FnApiStateAccessor(
       BeamFnStateClient beamFnStateClient,
       Coder<K> keyCoder,
       Coder<BoundedWindow> windowCoder) {
+    this(
+        pipelineOptions,
+        runnerCapabilities,
+        ptransformId,
+        processBundleInstructionId,
+        cacheTokens,
+        bundleCache,
+        processWideCache,
+        sideInputSpecMap,
+        beamFnStateClient,
+        keyCoder,
+        windowCoder,
+        () -> false,
+        () -> false);
+  }
+
+  public FnApiStateAccessor(
+      PipelineOptions pipelineOptions,
+      Set<String> runnerCapabilities,
+      String ptransformId,
+      Supplier<String> processBundleInstructionId,
+      Supplier<List<CacheToken>> cacheTokens,
+      Supplier<Cache<?, ?>> bundleCache,
+      Cache<?, ?> processWideCache,
+      Map<TupleTag<?>, SideInputSpec> sideInputSpecMap,
+      BeamFnStateClient beamFnStateClient,
+      Coder<K> keyCoder,
+      Coder<BoundedWindow> windowCoder,
+      Supplier<Boolean> hasNoState,
+      Supplier<Boolean> 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<K, BoundedWindow> keyAndWindowContext) {
@@ -1201,7 +1273,9 @@ private <T> BagUserState<T> createBagUserState(StateKey stateKey, Coder<T> value
             beamFnStateClient,
             processBundleInstructionId.get(),
             stateKey,
-            valueCoder);
+            valueCoder,
+            hasNoState.get(),
+            onlyBundleForKeys.get());
     stateFinalizers.add(rval::asyncClose);
     return rval;
   }
@@ -1283,7 +1357,9 @@ private <KeyT, ValueT> MultimapUserState<KeyT, ValueT> createMultimapUserState(
             processBundleInstructionId.get(),
             stateKey,
             keyCoder,
-            valueCoder);
+            valueCoder,
+            hasNoState.get(),
+            onlyBundleForKeys.get());
     stateFinalizers.add(rval::asyncClose);
     return rval;
   }
@@ -1318,7 +1394,9 @@ private <T> OrderedListUserState<T> createOrderedListUserState(
             beamFnStateClient,
             processBundleInstructionId.get(),
             stateKey,
-            valueCoder);
+            valueCoder,
+            hasNoState.get(),
+            onlyBundleForKeys.get());
     stateFinalizers.add(rval::asyncClose);
     return rval;
   }