Add DoFnRunner::finishKey() method (#38454)

Currently the method is a no-op.

In upcoming changes there'll be multiple streaming work items
in a single beam bundle. With multiple work items, we've to process
elements and timers of each work item before moving to the next work
items.

finishKey() will be called by the NativeIterator classes after iterating
through all elements from a work item.

Author: arunpandianp

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;