[Flink] Select bounded-source split assignment by estimated size

CHANGES.md

@@ -104,6 +104,7 @@
 
 ## New Features / Improvements
 
+* [Flink Runner] Added opt-in static round-robin split assignment for small bounded sources via the new `lazySourceSplitAssignmentMinSizeMbPerReader` pipeline option. The default of 0 keeps the existing lazy pull-based assignment ([#39873](https://github.com/apache/beam/issues/39873)).
 * Added `GroupIntoBatches` transform and the standard
   `beam:coder:sharded_key:v1` coder to the Go SDK, along with
   `beam.Coder.IsDeterministic`, `beam.PCollection.WindowingStrategy`,

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

@@ -365,6 +365,17 @@ public Long create(PipelineOptions options) {
 
   void setFileInputSplitMaxSizeMB(Long fileInputSplitMaxSizeMB);
 
+  @Description(
+      "Minimum estimated input size in MiB per source reader for lazy split assignment of "
+          + "bounded sources. The default of 0 always uses lazy assignment. A positive value "
+          + "selects static round-robin assignment for sources estimated below the threshold "
+          + "and lazy assignment for sources at or above it. Any negative value always uses "
+          + "static assignment.")
+  @Default.Long(0)
+  Long getLazySourceSplitAssignmentMinSizeMbPerReader();
+
+  void setLazySourceSplitAssignmentMinSizeMbPerReader(Long thresholdMb);
+
   @Description(
       "Allow drain operation for flink pipelines that contain RequiresStableInput operator. Note that at time of draining,"
           + "the RequiresStableInput contract might be violated if there any processing related failures in the DoFn operator.")

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

@@ -380,6 +380,17 @@ public Long create(PipelineOptions options) {
 
   void setFileInputSplitMaxSizeMB(Long fileInputSplitMaxSizeMB);
 
+  @Description(
+      "Minimum estimated input size in MiB per source reader for lazy split assignment of "
+          + "bounded sources. The default of 0 always uses lazy assignment. A positive value "
+          + "selects static round-robin assignment for sources estimated below the threshold "
+          + "and lazy assignment for sources at or above it. Any negative value always uses "
+          + "static assignment.")
+  @Default.Long(0)
+  Long getLazySourceSplitAssignmentMinSizeMbPerReader();
+
+  void setLazySourceSplitAssignmentMinSizeMbPerReader(Long thresholdMb);
+
   @Description(
       "Allow drain operation for flink pipelines that contain RequiresStableInput operator. Note that at time of draining,"
           + "the RequiresStableInput contract might be violated if there any processing related failures in the DoFn operator.")

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSource.java

@@ -18,12 +18,9 @@
 package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
 
 import java.io.Serializable;
-import java.util.List;
-import java.util.Map;
 import java.util.function.Function;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
 import org.apache.beam.runners.flink.FlinkPipelineOptions;
-import org.apache.beam.runners.flink.translation.utils.SerdeUtils;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.bounded.FlinkBoundedSource;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.impulse.BeamImpulseSource;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.unbounded.FlinkUnboundedSource;
@@ -43,7 +40,7 @@
  * @param <OutputT> The data type of the records emitted by the Flink Source.
  */
 public abstract class FlinkSource<T, OutputT>
-    implements Source<OutputT, FlinkSourceSplit<T>, Map<Integer, List<FlinkSourceSplit<T>>>> {
+    implements Source<OutputT, FlinkSourceSplit<T>, FlinkSourceEnumeratorState<T>> {
 
   protected final String stepName;
   protected final org.apache.beam.sdk.io.Source<T> beamSource;
@@ -102,36 +99,38 @@ public Boundedness getBoundedness() {
   }
 
   @Override
-  public SplitEnumerator<FlinkSourceSplit<T>, Map<Integer, List<FlinkSourceSplit<T>>>>
-      createEnumerator(SplitEnumeratorContext<FlinkSourceSplit<T>> enumContext) throws Exception {
-    return createEnumerator(enumContext, false);
-  }
-
-  public SplitEnumerator<FlinkSourceSplit<T>, Map<Integer, List<FlinkSourceSplit<T>>>>
-      createEnumerator(
-          SplitEnumeratorContext<FlinkSourceSplit<T>> enumContext, boolean splitInitialized)
-          throws Exception {
-
+  public SplitEnumerator<FlinkSourceSplit<T>, FlinkSourceEnumeratorState<T>> createEnumerator(
+      SplitEnumeratorContext<FlinkSourceSplit<T>> enumContext) throws Exception {
+    FlinkPipelineOptions options = serializablePipelineOptions.get().as(FlinkPipelineOptions.class);
     if (boundedness == Boundedness.BOUNDED) {
-      return new LazyFlinkSourceSplitEnumerator<>(
-          enumContext, beamSource, serializablePipelineOptions.get(), numSplits, splitInitialized);
-    } else {
-      return new FlinkSourceSplitEnumerator<>(
-          enumContext, beamSource, serializablePipelineOptions.get(), numSplits, splitInitialized);
+      long thresholdMb = options.getLazySourceSplitAssignmentMinSizeMbPerReader();
+      if (thresholdMb < 0) {
+        return new FlinkSourceSplitEnumerator<>(enumContext, beamSource, options, numSplits);
+      }
+      if (thresholdMb > 0) {
+        return new SizeBasedFlinkSourceSplitEnumerator<>(
+            enumContext, (BoundedSource<T>) beamSource, options, numSplits);
+      }
+      return new LazyFlinkSourceSplitEnumerator<>(enumContext, beamSource, options, numSplits);
     }
+    return new FlinkSourceSplitEnumerator<>(enumContext, beamSource, options, numSplits);
   }
 
   @Override
-  public SplitEnumerator<FlinkSourceSplit<T>, Map<Integer, List<FlinkSourceSplit<T>>>>
-      restoreEnumerator(
-          SplitEnumeratorContext<FlinkSourceSplit<T>> enumContext,
-          Map<Integer, List<FlinkSourceSplit<T>>> checkpoint)
-          throws Exception {
-    SplitEnumerator<FlinkSourceSplit<T>, Map<Integer, List<FlinkSourceSplit<T>>>> enumerator =
-        createEnumerator(enumContext, true);
-    checkpoint.forEach(
-        (subtaskId, splitsForSubtask) -> enumerator.addSplitsBack(splitsForSubtask, subtaskId));
-    return enumerator;
+  public SplitEnumerator<FlinkSourceSplit<T>, FlinkSourceEnumeratorState<T>> restoreEnumerator(
+      SplitEnumeratorContext<FlinkSourceSplit<T>> enumContext,
+      FlinkSourceEnumeratorState<T> checkpoint)
+      throws Exception {
+    FlinkPipelineOptions options = serializablePipelineOptions.get().as(FlinkPipelineOptions.class);
+    if (checkpoint.getAssignmentMode() == FlinkSourceSplitAssignmentMode.LAZY) {
+      return new LazyFlinkSourceSplitEnumerator<>(
+          enumContext, beamSource, options, numSplits, checkpoint);
+    }
+    if (checkpoint.getAssignmentMode() == FlinkSourceSplitAssignmentMode.STATIC) {
+      return new FlinkSourceSplitEnumerator<>(
+          enumContext, beamSource, options, numSplits, checkpoint);
+    }
+    return createEnumerator(enumContext);
   }
 
   @Override
@@ -140,9 +139,13 @@ public SimpleVersionedSerializer<FlinkSourceSplit<T>> getSplitSerializer() {
   }
 
   @Override
-  public SimpleVersionedSerializer<Map<Integer, List<FlinkSourceSplit<T>>>>
+  public SimpleVersionedSerializer<FlinkSourceEnumeratorState<T>>
       getEnumeratorCheckpointSerializer() {
-    return SerdeUtils.getNaiveObjectSerializer();
+    FlinkSourceSplitAssignmentMode legacyAssignmentMode =
+        boundedness == Boundedness.BOUNDED
+            ? FlinkSourceSplitAssignmentMode.LAZY
+            : FlinkSourceSplitAssignmentMode.STATIC;
+    return new FlinkSourceEnumeratorStateSerializer<>(legacyAssignmentMode);
   }
 
   public int getNumSplits() {

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorState.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/** Checkpoint state shared by the lazy and static source split assignment strategies. */
+public final class FlinkSourceEnumeratorState<T> implements Serializable {
+  private static final long serialVersionUID = 1L;
+
+  private final FlinkSourceSplitAssignmentMode assignmentMode;
+  private final ArrayList<FlinkSourceSplit<T>> pendingSplits;
+
+  /** Takes ownership of {@code pendingSplits}; the caller must not mutate it afterwards. */
+  FlinkSourceEnumeratorState(
+      FlinkSourceSplitAssignmentMode assignmentMode, ArrayList<FlinkSourceSplit<T>> pendingSplits) {
+    this.assignmentMode = assignmentMode;
+    this.pendingSplits = pendingSplits;
+  }
+
+  FlinkSourceSplitAssignmentMode getAssignmentMode() {
+    return assignmentMode;
+  }
+
+  List<FlinkSourceSplit<T>> getPendingSplits() {
+    return Collections.unmodifiableList(pendingSplits);
+  }
+}

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorStateSerializer.java

@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.flink.translation.utils.SerdeUtils;
+import org.apache.flink.core.io.SimpleVersionedSerializer;
+
+/** Serializes source enumerator state and upgrades the map used by earlier runner versions. */
+final class FlinkSourceEnumeratorStateSerializer<T>
+    implements SimpleVersionedSerializer<FlinkSourceEnumeratorState<T>> {
+  // Version written by SerdeUtils.getNaiveObjectSerializer, which serialized the
+  // Map<Integer, List<FlinkSourceSplit<T>>> state used by earlier runner versions.
+  static final int LEGACY_MAP_VERSION = 0;
+  static final int VERSION = 1;
+
+  private final FlinkSourceSplitAssignmentMode legacyAssignmentMode;
+
+  FlinkSourceEnumeratorStateSerializer(FlinkSourceSplitAssignmentMode legacyAssignmentMode) {
+    this.legacyAssignmentMode = legacyAssignmentMode;
+  }
+
+  @Override
+  public int getVersion() {
+    return VERSION;
+  }
+
+  @Override
+  public byte[] serialize(FlinkSourceEnumeratorState<T> state) throws IOException {
+    return SerdeUtils.serializeObject(state);
+  }
+
+  @Override
+  @SuppressWarnings("unchecked")
+  public FlinkSourceEnumeratorState<T> deserialize(int version, byte[] serialized)
+      throws IOException {
+    if (version == VERSION) {
+      Object deserialized = SerdeUtils.deserializeObject(serialized);
+      if (deserialized instanceof FlinkSourceEnumeratorState) {
+        return (FlinkSourceEnumeratorState<T>) deserialized;
+      }
+      throw new IOException(
+          "Expected source enumerator state of type FlinkSourceEnumeratorState for version "
+              + version
+              + ", but got: "
+              + describe(deserialized));
+    }
+    if (version == LEGACY_MAP_VERSION) {
+      Object deserialized = SerdeUtils.deserializeObject(serialized);
+      if (deserialized instanceof Map) {
+        return upgradeLegacyState((Map<?, ?>) deserialized);
+      }
+      throw new IOException(
+          "Expected legacy source enumerator state of type Map for version "
+              + version
+              + ", but got: "
+              + describe(deserialized));
+    }
+    throw new IOException(
+        String.format(
+            "Received source enumerator state version %d, but the highest supported version "
+                + "is %d.",
+            version, VERSION));
+  }
+
+  @SuppressWarnings("unchecked")
+  private FlinkSourceEnumeratorState<T> upgradeLegacyState(Map<?, ?> legacyState)
+      throws IOException {
+    ArrayList<FlinkSourceSplit<T>> pendingSplits = new ArrayList<>();
+    for (Object value : legacyState.values()) {
+      List<FlinkSourceSplit<T>> legacySplits = (List<FlinkSourceSplit<T>>) value;
+      if (legacySplits == null) {
+        throw new IOException("Legacy source enumerator state contains a null split list.");
+      }
+      pendingSplits.addAll(legacySplits);
+    }
+    return new FlinkSourceEnumeratorState<>(legacyAssignmentMode, pendingSplits);
+  }
+
+  private static String describe(@Nullable Object obj) {
+    return obj == null ? "null" : obj.getClass().getName();
+  }
+}

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitAssignmentMode.java

@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
+
+enum FlinkSourceSplitAssignmentMode {
+  UNDECIDED,
+  LAZY,
+  STATIC
+}