Add DoFnRunner::finishKey() method

runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java

@@ -63,6 +63,18 @@ public interface DoFnRunner<InputT extends @Nullable Object, OutputT extends @Nu
   <KeyT extends @Nullable Object> void onWindowExpiration(
       BoundedWindow window, Instant timestamp, KeyT key);
 
+  /**
+   * Performs per-key cleanup or processing after all elements, timers for a key have been processed
+   * and before moving to the next key or before finishBundle for the last key.
+   *
+   * <p>This is an optional method that can be used by runners as a hook to reset any per key state
+   * before moving to a different key in the same bundle. Currently used only by the Dataflow
+   * Streaming runner.
+   *
+   * @param key current key to clean up or finish processing
+   */
+  <KeyT extends @Nullable Object> void finishKey(KeyT key);
+
   /**
    * Returns the underlying fn instance.
    *

runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -101,6 +102,11 @@ public void finishBundle() {
     doFnRunner.finishBundle();
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {
+    doFnRunner.finishKey(key);
+  }
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     doFnRunner.onWindowExpiration(window, timestamp, key);

runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java

@@ -230,6 +230,9 @@ public void finishBundle() {
     }
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     invoker.invokeOnWindowExpiration(

runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java

@@ -43,6 +43,7 @@
 import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 
@@ -131,6 +132,11 @@ public void finishBundle() {
     doFnRunner.finishBundle();
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {
+    doFnRunner.finishKey(key);
+  }
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     doFnRunner.onWindowExpiration(window, timestamp, key);

runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java

@@ -65,6 +65,7 @@
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 import org.junit.Before;
@@ -379,6 +380,9 @@ public void finishBundle() {
       finished = true;
     }
 
+    @Override
+    public <KeyT extends @Nullable Object> void finishKey(KeyT key) {}
+
     @Override
     public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
   }

runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java

@@ -27,6 +27,7 @@
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.WindowedValue;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -105,6 +106,9 @@ public void finishBundle() {
     container.updateMetrics(stepName);
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     delegate.onWindowExpiration(window, timestamp, key);

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java

@@ -1064,6 +1064,9 @@ public void finishBundle() {
       }
     }
 
+    @Override
+    public <KeyT extends @Nullable Object> void finishKey(KeyT key) {}
+
     @Override
     public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
 

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java

@@ -255,6 +255,9 @@ public void finishBundle() {
     Optional.ofNullable(finishBundleCallback).ifPresent(Runnable::run);
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
 

runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java

@@ -523,6 +523,9 @@ public void finishBundle() {
                 wrappedRunner.finishBundle();
               }
 
+              @Override
+              public <KeyT extends @Nullable Object> void finishKey(KeyT key) {}
+
               @Override
               public <KeyT> void onWindowExpiration(
                   BoundedWindow window, Instant timestamp, KeyT key) {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java

@@ -119,6 +119,9 @@ public void processTimers() throws Exception {
       // Nothing.
     }
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() throws Exception {
       receiver = null;

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeUngroupingParDoFn.java

@@ -73,6 +73,9 @@ public void processTimers() throws Exception {
     // The timers for the underlying ParDoFn are processed at the end of each element
   }
 
+  @Override
+  public void finishKey(Object key) throws Exception {}
+
   @Override
   public void finishBundle() throws Exception {
     underlyingParDoFn.finishBundle();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/CreateIsmShardKeyAndSortKeyDoFnFactory.java

@@ -114,6 +114,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() {}
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java

@@ -35,6 +35,7 @@
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -128,6 +129,11 @@ public void finishBundle() {
     simpleRunner.finishBundle();
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {
+    simpleRunner.finishKey(key);
+  }
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     simpleRunner.onWindowExpiration(window, timestamp, key);

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ForwardingParDoFn.java

@@ -19,6 +19,7 @@
 
 import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * A base class for {@link ParDoFn} implementations for overriding particular methods while
@@ -47,6 +48,11 @@ public void processTimers() throws Exception {
     delegate.processTimers();
   }
 
+  @Override
+  public void finishKey(@Nullable Object key) throws Exception {
+    delegate.finishKey(key);
+  }
+
   @Override
   public void finishBundle() throws Exception {
     delegate.finishBundle();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java

@@ -27,6 +27,7 @@
 import org.apache.beam.sdk.util.WindowedValueReceiver;
 import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.WindowedValue;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -102,6 +103,9 @@ private void invokeProcessElement(WindowedValue<InputT> elem) {
   @Override
   public void finishBundle() {}
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java

@@ -142,6 +142,12 @@ public void processTimers() throws Exception {
     // it here to build a KeyedWorkItem
   }
 
+  @Override
+  public void finishKey(Object key) throws Exception {
+    checkState(fnRunner != null);
+    fnRunner.finishKey(key);
+  }
+
   @Override
   public void finishBundle() throws Exception {
     checkState(fnRunner != null);

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PairWithConstantKeyDoFnFactory.java

@@ -98,6 +98,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() {}
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java

@@ -317,6 +317,9 @@ public void processElement(Object elem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() throws Exception {
       groupingTable.flush(receiver);
@@ -377,6 +380,9 @@ public void processElement(Object elem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() throws Exception {
       groupingTable.flush(receiver);

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ReifyTimestampAndWindowsParDoFnFactory.java

@@ -86,6 +86,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() throws Exception {
       this.receiver = null;

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java

@@ -195,6 +195,9 @@ public void processTimers() throws Exception {
         () -> sideInputProcessor);
   }
 
+  @Override
+  public void finishKey(Object key) throws Exception {}
+
   @Override
   public void finishBundle() throws Exception {
     helpers.finishBundle(sideInputProcessor);

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java

@@ -39,6 +39,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.Lists;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -155,6 +156,11 @@ public void finishBundle() {
     sideInputFetcher.persist();
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {
+    simpleDoFnRunner.finishKey(key);
+  }
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     simpleDoFnRunner.onWindowExpiration(window, timestamp, key);

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputParDoFn.java

@@ -193,6 +193,9 @@ public void processTimers() throws Exception {
         () -> sideInputProcessor);
   }
 
+  @Override
+  public void finishKey(Object key) throws Exception {}
+
   @Override
   public void finishBundle() throws Exception {
     helpers.finishBundle(sideInputProcessor);

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java

@@ -282,7 +282,7 @@ public void finishKey() {
     checkState(!finishKeyCalled, "finishKey was already called");
     checkStateNotNull(workExecutor, "workExecutor must be set before calling finishKey()");
     try {
-      workExecutor.finishKey();
+      workExecutor.finishKey(key);
     } catch (Exception e) {
       throw new RuntimeException(e);
     }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingPCollectionViewWriterParDoFn.java

@@ -73,6 +73,9 @@ public void processElement(Object element) throws Exception {
   @Override
   public void processTimers() {}
 
+  @Override
+  public void finishKey(Object key) throws Exception {}
+
   @Override
   public void finishBundle() throws Exception {}
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java

@@ -24,6 +24,7 @@
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.WindowedValue;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -80,6 +81,11 @@ public <KeyT> void onTimer(
         "Attempt to deliver a timer to a DoFn, but timers are not supported in Dataflow.");
   }
 
+  @Override
+  public <KeyT extends @Nullable Object> void finishKey(KeyT key) {
+    simpleDoFnRunner.finishKey(key);
+  }
+
   @Override
   public void finishBundle() {
     simpleDoFnRunner.finishBundle();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ToIsmRecordForMultimapDoFnFactory.java

@@ -149,6 +149,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() {}
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ValuesDoFnFactory.java

@@ -80,6 +80,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey(Object key) throws Exception {}
+
     @Override
     public void finishBundle() {}
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourceOperationExecutor.java

@@ -90,7 +90,7 @@ public void execute() throws Exception {
   }
 
   @Override
-  public void finishKey() throws Exception {}
+  public void finishKey(Object key) throws Exception {}
 
   @Override
   public SourceOperationResponse getResponse() {