[Flink] Add a dedicated operator for batch combines in the global window

CHANGES.md

@@ -68,6 +68,7 @@
 
 ## New Features / Improvements
 
+* (Java) Flink runner: the batch-only `--sortedBatchCombine` option is enabled by default (disable it with `--sortedBatchCombine=false`). GroupByKey and Combine in the global window with the default trigger then finish with a dedicated operator that relies on Flink sorting its input by key, instead of keyed state and timers ([#40607](https://github.com/apache/beam/issues/40607)).
 * (Go) Added `wait.On`, which delays each input window until the corresponding windows in its signal PCollections have closed ([#39909](https://github.com/apache/beam/issues/39909)).
 * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)).
 * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)).

runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java

@@ -350,6 +350,16 @@ public Long create(PipelineOptions options) {
 
   void setFasterCopy(Boolean fasterCopy);
 
+  @Description(
+      "Batch only. Use a dedicated operator for the final step of GroupByKey and Combine in the "
+          + "global window with the default trigger. It relies on Flink sorting the input by key "
+          + "and does not use keyed state or timers. Has no effect in streaming mode. "
+          + "Enabled by default; set to false to disable.")
+  @Default.Boolean(true)
+  Boolean getSortedBatchCombine();
+
+  void setSortedBatchCombine(Boolean sortedBatchCombine);
+
   @Description(
       "Directory containing Flink YAML configuration files. "
           + "These properties will be set to all jobs submitted to Flink and take precedence "

runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java

@@ -367,6 +367,16 @@ public Long create(PipelineOptions options) {
 
   void setFasterCopy(Boolean fasterCopy);
 
+  @Description(
+      "Batch only. Use a dedicated operator for the final step of GroupByKey and Combine in the "
+          + "global window with the default trigger. It relies on Flink sorting the input by key "
+          + "and does not use keyed state or timers. Has no effect in streaming mode. "
+          + "Enabled by default; set to false to disable.")
+  @Default.Boolean(true)
+  Boolean getSortedBatchCombine();
+
+  void setSortedBatchCombine(Boolean sortedBatchCombine);
+
   @Description(
       "Directory containing Flink YAML configuration files. "
           + "These properties will be set to all jobs submitted to Flink and take precedence "

runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingAggregationsTranslators.java

@@ -22,6 +22,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 import org.apache.beam.runners.core.KeyedWorkItem;
 import org.apache.beam.runners.core.SystemReduceFn;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
@@ -31,19 +32,23 @@
 import org.apache.beam.runners.flink.translation.wrappers.streaming.KvToFlinkKeyKeySelector;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.PartialReduceBundleOperator;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.SingletonKeyedWorkItemCoder;
+import org.apache.beam.runners.flink.translation.wrappers.streaming.SortedBatchCombineOperator;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.WindowDoFnOperator;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.WorkItemKeySelector;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderRegistry;
 import org.apache.beam.sdk.coders.IterableCoder;
 import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.transforms.Combine;
 import org.apache.beam.sdk.transforms.CombineFnBase;
 import org.apache.beam.sdk.transforms.CombineWithContext;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.join.RawUnionValue;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.DefaultTrigger;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
 import org.apache.beam.sdk.util.AppliedCombineFn;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
@@ -61,6 +66,7 @@
 import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.datastream.KeyedStream;
 import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
 import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
 import org.apache.flink.util.Collector;
 
@@ -279,7 +285,8 @@ SingleOutputStreamOperator<WindowedValue<KV<K, OutputT>>> getBatchCombinePerKeyO
           List<PCollectionView<?>> sideInputs,
           Coder<WindowedValue<KV<K, AccumT>>> windowedAccumCoder,
           CombineFnBase.GlobalCombineFn<InputT, AccumT, ?> combineFn,
-          WindowDoFnOperator<K, AccumT, OutputT> finalDoFnOperator,
+          OneInputStreamOperator<WindowedValue<KV<K, AccumT>>, WindowedValue<KV<K, OutputT>>>
+              finalDoFnOperator,
           TypeInformation<WindowedValue<KV<K, OutputT>>> outputTypeInfo) {
 
     String fullName = FlinkStreamingTransformTranslators.getCurrentTransformName(context);
@@ -348,7 +355,7 @@ SingleOutputStreamOperator<WindowedValue<KV<K, OutputT>>> getBatchCombinePerKeyO
           partialyCombinedStream.keyBy(accumKeySelector),
           transformSideInputs.f1,
           fullName,
-          finalDoFnOperator,
+          (WindowDoFnOperator<K, AccumT, OutputT>) finalDoFnOperator,
           outputTypeInfo);
     }
   }
@@ -400,9 +407,23 @@ SingleOutputStreamOperator<WindowedValue<KV<K, Iterable<InputT>>>> batchGroupByK
             serializablePipelineOptions);
 
     // final aggregation
-    WindowDoFnOperator<K, Iterable<InputT>, Iterable<Iterable<InputT>>> finalDoFnOperator =
-        getWindowedAccumulateDoFnOperator(
-            context, transform, accumKvCoder, outputCoder, sideInputTagMapping, sideInputs);
+    OneInputStreamOperator<
+            WindowedValue<KV<K, Iterable<InputT>>>,
+            WindowedValue<KV<K, Iterable<Iterable<InputT>>>>>
+        finalDoFnOperator =
+            getFinalAggregationOperator(
+                context,
+                input,
+                sideInputs,
+                new SortedBatchCombineOperator.BufferingCombiner<>(),
+                () ->
+                    getWindowedAccumulateDoFnOperator(
+                        context,
+                        transform,
+                        accumKvCoder,
+                        outputCoder,
+                        sideInputTagMapping,
+                        sideInputs));
 
     return getBatchCombinePerKeyOperator(
             context,
@@ -441,6 +462,47 @@ SingleOutputStreamOperator<WindowedValue<KV<K, Iterable<InputT>>>> batchGroupByK
         context, transform, accumKvCoder, outputCoder, reduceFn, sideInputTagMapping, sideInputs);
   }
 
+  /**
+   * The operator for the final step of a batch GroupByKey or Combine: a {@link
+   * SortedBatchCombineOperator} using {@code sortedCombiner} when {@link #useSortedBatchCombine}
+   * allows it, otherwise the {@link WindowDoFnOperator} from {@code windowDoFnOperator}.
+   */
+  private static <K, AccumT, OutputT>
+      OneInputStreamOperator<WindowedValue<KV<K, AccumT>>, WindowedValue<KV<K, OutputT>>>
+          getFinalAggregationOperator(
+              FlinkStreamingTranslationContext context,
+              PCollection<?> input,
+              List<PCollectionView<?>> sideInputs,
+              SortedBatchCombineOperator.KeyCombiner<AccumT, OutputT> sortedCombiner,
+              Supplier<WindowDoFnOperator<K, AccumT, OutputT>> windowDoFnOperator) {
+    WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();
+    if (useSortedBatchCombine(context.getPipelineOptions(), windowingStrategy, sideInputs)) {
+      return new SortedBatchCombineOperator<>(
+          FlinkStreamingTransformTranslators.getCurrentTransformName(context),
+          sortedCombiner,
+          windowingStrategy.getTimestampCombiner(),
+          new SerializablePipelineOptions(context.getPipelineOptions()));
+    } else {
+      return windowDoFnOperator.get();
+    }
+  }
+
+  /**
+   * Whether the final step of a batch GroupByKey or Combine can use a {@link
+   * SortedBatchCombineOperator} instead of a {@link WindowDoFnOperator}.
+   */
+  static boolean useSortedBatchCombine(
+      PipelineOptions pipelineOptions,
+      WindowingStrategy<?, ?> windowingStrategy,
+      List<PCollectionView<?>> sideInputs) {
+    FlinkPipelineOptions options = pipelineOptions.as(FlinkPipelineOptions.class);
+    return !options.isStreaming()
+        && options.getSortedBatchCombine()
+        && sideInputs.isEmpty()
+        && windowingStrategy.getWindowFn() instanceof GlobalWindows
+        && windowingStrategy.getTrigger() instanceof DefaultTrigger;
+  }
+
   public static <K, InputT, AccumT, OutputT>
       SingleOutputStreamOperator<WindowedValue<KV<K, OutputT>>> batchCombinePerKey(
           FlinkStreamingTranslationContext context,
@@ -476,15 +538,22 @@ SingleOutputStreamOperator<WindowedValue<KV<K, OutputT>>> batchCombinePerKey(
     }
 
     // final aggregation from AccumT to OutputT
-    WindowDoFnOperator<K, AccumT, OutputT> finalDoFnOperator =
-        getWindowedAggregateDoFnOperator(
-            context,
-            transform,
-            accumKvCoder,
-            outputCoder,
-            toFinalFlinkCombineFn(combineFn, inputKvCoder.getValueCoder()),
-            sideInputTagMapping,
-            sideInputs);
+    OneInputStreamOperator<WindowedValue<KV<K, AccumT>>, WindowedValue<KV<K, OutputT>>>
+        finalDoFnOperator =
+            getFinalAggregationOperator(
+                context,
+                input,
+                sideInputs,
+                new SortedBatchCombineOperator.MergingCombiner<>(combineFn),
+                () ->
+                    getWindowedAggregateDoFnOperator(
+                        context,
+                        transform,
+                        accumKvCoder,
+                        outputCoder,
+                        toFinalFlinkCombineFn(combineFn, inputKvCoder.getValueCoder()),
+                        sideInputTagMapping,
+                        sideInputs));
 
     return getBatchCombinePerKeyOperator(
         context,

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

@@ -89,7 +89,7 @@ public void updateMetrics(String stepName, List<MetricsApi.MonitoringInfo> monit
    * Update Flink's internal metrics ({@link this#flinkCounterCache}) with the latest metrics for a
    * given step.
    */
-  void updateMetrics(String stepName) {
+  public void updateMetrics(String stepName) {
     MetricResults metricResults = asAttemptedOnlyMetricResults(metricsContainers);
     MetricQueryResults metricQueryResults =
         metricResults.queryMetrics(MetricsFilter.builder().addStep(stepName).build());