This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new c979b5cfbe2 [Flink] Select bounded-source split assignment by
estimated size (#39874)
c979b5cfbe2 is described below
commit c979b5cfbe26b545c4d742081292d187d58dbd11
Author: Paulius Kuzmickas <[email protected]>
AuthorDate: Wed Sep 2 18:59:13 2026 +0100
[Flink] Select bounded-source split assignment by estimated size (#39874)
* [Flink] Select bounded-source split assignment by size
---
CHANGES.md | 1 +
.../beam/runners/flink/FlinkPipelineOptions.java | 11 +
.../beam/runners/flink/FlinkPipelineOptions.java | 11 +
.../wrappers/streaming/io/source/FlinkSource.java | 64 +--
.../io/source/FlinkSourceEnumeratorState.java | 52 ++
.../FlinkSourceEnumeratorStateSerializer.java | 103 ++++
.../io/source/FlinkSourceSplitEnumerator.java | 141 +++---
.../streaming/io/source/FlinkSourceSplitUtils.java | 84 ++++
.../io/source/LazyFlinkSourceSplitEnumerator.java | 184 ++++----
.../SizeBasedFlinkSourceSplitEnumerator.java | 205 ++++++++
.../io/source/FlinkSourceSplitEnumeratorTest.java | 522 +++++++++++++++++----
.../shortcodes/flink_java_pipeline_options.html | 5 +
.../shortcodes/flink_python_pipeline_options.html | 5 +
13 files changed, 1101 insertions(+), 287 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index b857e09ac29..fc34c155256 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -123,6 +123,7 @@
## New Features / Improvements
+* [Flink Runner] Added opt-in static round-robin split assignment for small
bounded sources via the new `sourceStaticSplitThresholdMb` pipeline option. The
default of 0 keeps the existing lazy pull-based assignment
([#39873](https://github.com/apache/beam/issues/39873)).
* Added automatic caching of bounded, single-pane side-input views for classic
Java Flink DataStream execution
([#39866](https://github.com/apache/beam/issues/39866)).
* Added `GroupIntoBatches` transform and the standard
`beam:coder:sharded_key:v1` coder to the Go SDK, along with
diff --git
a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
index 3fee130d58e..d3c5fefd9cf 100644
---
a/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
+++
b/runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
@@ -365,6 +365,17 @@ public interface FlinkPipelineOptions
void setFileInputSplitMaxSizeMB(Long fileInputSplitMaxSizeMB);
+ @Description(
+ "Static split assignment threshold in MiB per source reader for bounded
sources in Flink "
+ + "DataStream mode. The default of 0 always uses lazy assignment. A
positive value "
+ + "selects static round-robin assignment for sources with a known,
positive estimate "
+ + "below the threshold and lazy assignment otherwise. Any negative
value always uses "
+ + "static assignment.")
+ @Default.Long(0)
+ Long getSourceStaticSplitThresholdMb();
+
+ void setSourceStaticSplitThresholdMb(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.")
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
index 6e4d998ed09..a6729356bdc 100644
---
a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java
@@ -382,6 +382,17 @@ public interface FlinkPipelineOptions
void setFileInputSplitMaxSizeMB(Long fileInputSplitMaxSizeMB);
+ @Description(
+ "Static split assignment threshold in MiB per source reader for bounded
sources in Flink "
+ + "DataStream mode. The default of 0 always uses lazy assignment. A
positive value "
+ + "selects static round-robin assignment for sources with a known,
positive estimate "
+ + "below the threshold and lazy assignment otherwise. Any negative
value always uses "
+ + "static assignment.")
+ @Default.Long(0)
+ Long getSourceStaticSplitThresholdMb();
+
+ void setSourceStaticSplitThresholdMb(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.")
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSource.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSource.java
index 0c1ba73b2a2..f3bbf14f68f 100644
---
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSource.java
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSource.java
@@ -18,12 +18,10 @@
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.FlinkSourceEnumeratorState.AssignmentMode;
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 +41,7 @@ import org.apache.flink.core.io.SimpleVersionedSerializer;
* @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 +100,38 @@ public abstract class FlinkSource<T, OutputT>
}
@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.getSourceStaticSplitThresholdMb();
+ 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() == AssignmentMode.LAZY) {
+ return new LazyFlinkSourceSplitEnumerator<>(
+ enumContext, beamSource, options, numSplits, checkpoint);
+ }
+ if (checkpoint.getAssignmentMode() == AssignmentMode.STATIC) {
+ return new FlinkSourceSplitEnumerator<>(
+ enumContext, beamSource, options, numSplits, checkpoint);
+ }
+ return createEnumerator(enumContext);
}
@Override
@@ -140,9 +140,11 @@ public abstract class FlinkSource<T, OutputT>
}
@Override
- public SimpleVersionedSerializer<Map<Integer, List<FlinkSourceSplit<T>>>>
+ public SimpleVersionedSerializer<FlinkSourceEnumeratorState<T>>
getEnumeratorCheckpointSerializer() {
- return SerdeUtils.getNaiveObjectSerializer();
+ AssignmentMode legacyAssignmentMode =
+ boundedness == Boundedness.BOUNDED ? AssignmentMode.LAZY :
AssignmentMode.STATIC;
+ return new FlinkSourceEnumeratorStateSerializer<>(legacyAssignmentMode);
}
public int getNumSplits() {
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorState.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorState.java
new file mode 100644
index 00000000000..8a825eeb991
--- /dev/null
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorState.java
@@ -0,0 +1,52 @@
+/*
+ * 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 AssignmentMode assignmentMode;
+ private final ArrayList<FlinkSourceSplit<T>> pendingSplits;
+
+ /** Takes ownership of {@code pendingSplits}; the caller must not mutate it
afterwards. */
+ FlinkSourceEnumeratorState(
+ AssignmentMode assignmentMode, ArrayList<FlinkSourceSplit<T>>
pendingSplits) {
+ this.assignmentMode = assignmentMode;
+ this.pendingSplits = pendingSplits;
+ }
+
+ AssignmentMode getAssignmentMode() {
+ return assignmentMode;
+ }
+
+ List<FlinkSourceSplit<T>> getPendingSplits() {
+ return Collections.unmodifiableList(pendingSplits);
+ }
+
+ enum AssignmentMode {
+ UNDECIDED,
+ LAZY,
+ STATIC
+ }
+}
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorStateSerializer.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorStateSerializer.java
new file mode 100644
index 00000000000..3470e60d0bd
--- /dev/null
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorStateSerializer.java
@@ -0,0 +1,103 @@
+/*
+ * 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.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceEnumeratorState.AssignmentMode;
+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 AssignmentMode legacyAssignmentMode;
+
+ FlinkSourceEnumeratorStateSerializer(AssignmentMode 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();
+ }
+}
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java
index 38d2f363939..bb8956f5dae 100644
---
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java
@@ -21,10 +21,13 @@ import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import javax.annotation.Nullable;
+import
org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceEnumeratorState.AssignmentMode;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.Source;
import org.apache.beam.sdk.io.UnboundedSource;
@@ -35,27 +38,17 @@ import
org.apache.flink.api.connector.source.SplitsAssignment;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-/**
- * A Flink {@link org.apache.flink.api.connector.source.SplitEnumerator
SplitEnumerator}
- * implementation that holds a Beam {@link Source} and does the following:
- *
- * <ul>
- * <li>Split the Beam {@link Source} to desired number of splits.
- * <li>Assign the splits to the Flink Source Reader.
- * </ul>
- *
- * <p>Note that at this point, this class has a static round-robin split
assignment strategy.
- *
- * @param <T> The output type of the encapsulated Beam {@link Source}.
- */
+/** Splits a Beam source and assigns its splits to Flink source readers
round-robin. */
public class FlinkSourceSplitEnumerator<T>
- implements SplitEnumerator<FlinkSourceSplit<T>, Map<Integer,
List<FlinkSourceSplit<T>>>> {
+ implements SplitEnumerator<FlinkSourceSplit<T>,
FlinkSourceEnumeratorState<T>> {
private static final Logger LOG =
LoggerFactory.getLogger(FlinkSourceSplitEnumerator.class);
+
private final SplitEnumeratorContext<FlinkSourceSplit<T>> context;
private final Source<T> beamSource;
private final PipelineOptions pipelineOptions;
private final int numSplits;
private final Map<Integer, List<FlinkSourceSplit<T>>> pendingSplits;
+
private boolean splitsInitialized;
public FlinkSourceSplitEnumerator(
@@ -63,8 +56,7 @@ public class FlinkSourceSplitEnumerator<T>
Source<T> beamSource,
PipelineOptions pipelineOptions,
int numSplits) {
-
- this(context, beamSource, pipelineOptions, numSplits, false);
+ this(context, beamSource, pipelineOptions, numSplits, null);
}
public FlinkSourceSplitEnumerator(
@@ -72,17 +64,31 @@ public class FlinkSourceSplitEnumerator<T>
Source<T> beamSource,
PipelineOptions pipelineOptions,
int numSplits,
- boolean splitsInitialized) {
-
+ @Nullable FlinkSourceEnumeratorState<T> restoredState) {
this.context = context;
this.beamSource = beamSource;
this.pipelineOptions = pipelineOptions;
this.numSplits = numSplits;
this.pendingSplits = new HashMap<>(numSplits);
- this.splitsInitialized = splitsInitialized;
+ this.splitsInitialized = restoredState != null;
+
+ if (restoredState != null) {
+ if (restoredState.getAssignmentMode() != AssignmentMode.STATIC) {
+ throw new IllegalArgumentException(
+ "Cannot restore the static source enumerator from "
+ + restoredState.getAssignmentMode()
+ + " state.");
+ }
+ int parallelism = context.currentParallelism();
+ for (FlinkSourceSplit<T> split : restoredState.getPendingSplits()) {
+ int targetSubtask = split.splitIndex() % parallelism;
+ pendingSplits.computeIfAbsent(targetSubtask, ignored -> new
ArrayList<>()).add(split);
+ }
+ }
LOG.info(
- "Created new enumerator with parallelism {}, source {}, numSplits {},
initialized {}",
+ "Created static source enumerator with parallelism {}, source {},
numSplits {}, "
+ + "initialized {}",
context.currentParallelism(),
beamSource,
numSplits,
@@ -93,52 +99,33 @@ public class FlinkSourceSplitEnumerator<T>
public void start() {
if (!splitsInitialized) {
initializeSplits();
+ } else {
+ sendPendingSplitsToSourceReaders();
}
}
private void initializeSplits() {
context.callAsync(
- () -> {
- try {
- LOG.info("Starting source {}", beamSource);
- List<? extends Source<T>> beamSplitSourceList = splitBeamSource();
- Map<Integer, List<FlinkSourceSplit<T>>> flinkSourceSplitsList =
new HashMap<>();
- int i = 0;
- for (Source<T> beamSplitSource : beamSplitSourceList) {
- int targetSubtask = i % context.currentParallelism();
- List<FlinkSourceSplit<T>> splitsForTask =
- flinkSourceSplitsList.computeIfAbsent(
- targetSubtask, ignored -> new ArrayList<>());
- splitsForTask.add(new FlinkSourceSplit<>(i, beamSplitSource));
- i++;
- }
- return flinkSourceSplitsList;
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- },
+ this::splitBeamSource,
(sourceSplits, error) -> {
if (error != null) {
throw new RuntimeException("Failed to start source enumerator.",
error);
- } else {
- pendingSplits.putAll(sourceSplits);
- splitsInitialized = true;
- sendPendingSplitsToSourceReaders();
}
+ prepareAssignments(sourceSplits);
+ splitsInitialized = true;
+ sendPendingSplitsToSourceReaders();
});
}
@Override
public void handleSplitRequest(int subtaskId, @Nullable String
requesterHostname) {
- // Not used.
+ // Static assignment happens when readers register.
}
@Override
public void addSplitsBack(List<FlinkSourceSplit<T>> splits, int subtaskId) {
LOG.info("Adding splits {} back from subtask {}", splits, subtaskId);
- List<FlinkSourceSplit<T>> splitsForSubtask =
- pendingSplits.computeIfAbsent(subtaskId, ignored -> new ArrayList<>());
- splitsForSubtask.addAll(splits);
+ pendingSplits.computeIfAbsent(subtaskId, ignored -> new
ArrayList<>()).addAll(splits);
}
@Override
@@ -146,18 +133,19 @@ public class FlinkSourceSplitEnumerator<T>
List<FlinkSourceSplit<T>> splitsForSubtask =
pendingSplits.remove(subtaskId);
if (splitsForSubtask != null) {
assignSplitsAndLog(splitsForSubtask, subtaskId);
- } else {
- if (splitsInitialized) {
- LOG.info("There is no split for subtask {}. Signaling no more
splits.", subtaskId);
- context.signalNoMoreSplits(subtaskId);
- }
+ } else if (splitsInitialized) {
+ LOG.info("There is no split for subtask {}. Signaling no more splits.",
subtaskId);
+ context.signalNoMoreSplits(subtaskId);
}
}
@Override
- public Map<Integer, List<FlinkSourceSplit<T>>> snapshotState(long
checkpointId) throws Exception {
+ public FlinkSourceEnumeratorState<T> snapshotState(long checkpointId) {
LOG.info("Taking snapshot for checkpoint {}", checkpointId);
- return pendingSplits;
+ ArrayList<FlinkSourceSplit<T>> checkpointSplits = new ArrayList<>();
+ pendingSplits.values().forEach(checkpointSplits::addAll);
+ AssignmentMode mode = splitsInitialized ? AssignmentMode.STATIC :
AssignmentMode.UNDECIDED;
+ return new FlinkSourceEnumeratorState<>(mode, checkpointSplits);
}
@Override
@@ -165,34 +153,49 @@ public class FlinkSourceSplitEnumerator<T>
// NoOp
}
- // -------------- Private helper methods ----------------------
- private List<? extends Source<T>> splitBeamSource() throws Exception {
+ private ArrayList<FlinkSourceSplit<T>> splitBeamSource() throws Exception {
+ LOG.info("Starting source {}", beamSource);
if (beamSource instanceof BoundedSource) {
BoundedSource<T> boundedSource = (BoundedSource<T>) beamSource;
- long desiredSizeBytes =
boundedSource.getEstimatedSizeBytes(pipelineOptions) / numSplits;
- return boundedSource.split(desiredSizeBytes, pipelineOptions);
- } else if (beamSource instanceof UnboundedSource) {
- List<? extends UnboundedSource<T, ?>> splits =
- ((UnboundedSource<T, ?>) beamSource).split(numSplits,
pipelineOptions);
- LOG.info("Split source {} to {} splits", beamSource, splits);
- return splits;
- } else {
- throw new IllegalStateException("Unknown source type " +
beamSource.getClass());
+ long estimatedSizeBytes =
+ FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource,
pipelineOptions);
+ return FlinkSourceSplitUtils.splitBoundedSource(
+ boundedSource, pipelineOptions, numSplits, estimatedSizeBytes);
+ }
+ if (beamSource instanceof UnboundedSource) {
+ return FlinkSourceSplitUtils.splitUnboundedSource(
+ (UnboundedSource<T, ?>) beamSource, pipelineOptions, numSplits);
+ }
+ throw new IllegalStateException("Unknown source type " +
beamSource.getClass());
+ }
+
+ private void prepareAssignments(List<FlinkSourceSplit<T>> sourceSplits) {
+ int parallelism = context.currentParallelism();
+ for (FlinkSourceSplit<T> split : sourceSplits) {
+ int targetSubtask = split.splitIndex() % parallelism;
+ pendingSplits.computeIfAbsent(targetSubtask, ignored -> new
ArrayList<>()).add(split);
}
}
private void sendPendingSplitsToSourceReaders() {
+ Set<Integer> assignedReaders = new HashSet<>();
Iterator<Map.Entry<Integer, List<FlinkSourceSplit<T>>>> splitIter =
pendingSplits.entrySet().iterator();
while (splitIter.hasNext()) {
Map.Entry<Integer, List<FlinkSourceSplit<T>>> entry = splitIter.next();
- int readerIndex = entry.getKey();
- int targetSubtask = readerIndex % context.currentParallelism();
- if (context.registeredReaders().containsKey(targetSubtask)) {
- assignSplitsAndLog(entry.getValue(), targetSubtask);
+ int subtaskId = entry.getKey();
+ if (context.registeredReaders().containsKey(subtaskId)) {
+ assignSplitsAndLog(entry.getValue(), subtaskId);
+ assignedReaders.add(subtaskId);
splitIter.remove();
}
}
+
+ for (int subtaskId : context.registeredReaders().keySet()) {
+ if (!assignedReaders.contains(subtaskId) &&
!pendingSplits.containsKey(subtaskId)) {
+ context.signalNoMoreSplits(subtaskId);
+ }
+ }
}
private void assignSplitsAndLog(List<FlinkSourceSplit<T>> splits, int
subtaskId) {
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitUtils.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitUtils.java
new file mode 100644
index 00000000000..dda99bf53be
--- /dev/null
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitUtils.java
@@ -0,0 +1,84 @@
+/*
+ * 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.util.ArrayList;
+import java.util.List;
+import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.io.FileBasedSource;
+import org.apache.beam.sdk.io.Source;
+import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+
+/** Shared Beam source sizing and splitting helpers. */
+final class FlinkSourceSplitUtils {
+ static final long MEBIBYTE = 1024L * 1024L;
+
+ private FlinkSourceSplitUtils() {}
+
+ static <T> long estimateBoundedSourceSize(
+ BoundedSource<T> boundedSource, PipelineOptions pipelineOptions) throws
Exception {
+ return boundedSource.getEstimatedSizeBytes(pipelineOptions);
+ }
+
+ static <T> ArrayList<FlinkSourceSplit<T>> splitBoundedSource(
+ BoundedSource<T> boundedSource,
+ PipelineOptions pipelineOptions,
+ int numSplits,
+ long estimatedSizeBytes)
+ throws Exception {
+ long desiredSizeBytes =
+ getDesiredSizeBytes(boundedSource, pipelineOptions, numSplits,
estimatedSizeBytes);
+ return toFlinkSplits(boundedSource.split(desiredSizeBytes,
pipelineOptions));
+ }
+
+ static <T> ArrayList<FlinkSourceSplit<T>> splitUnboundedSource(
+ UnboundedSource<T, ?> unboundedSource, PipelineOptions pipelineOptions,
int numSplits)
+ throws Exception {
+ return toFlinkSplits(unboundedSource.split(numSplits, pipelineOptions));
+ }
+
+ static long getDesiredSizeBytes(
+ Source<?> beamSource,
+ PipelineOptions pipelineOptions,
+ int numSplits,
+ long estimatedSizeBytes) {
+ long desiredSizeBytes = estimatedSizeBytes / numSplits;
+
+ long maxSplitSizeMb =
+
pipelineOptions.as(FlinkPipelineOptions.class).getFileInputSplitMaxSizeMB();
+ if (beamSource instanceof FileBasedSource && maxSplitSizeMb > 0) {
+ return Math.min(desiredSizeBytes, mebibytesToBytes(maxSplitSizeMb));
+ }
+ return desiredSizeBytes;
+ }
+
+ static long mebibytesToBytes(long mebibytes) {
+ return mebibytes > Long.MAX_VALUE / MEBIBYTE ? Long.MAX_VALUE : mebibytes
* MEBIBYTE;
+ }
+
+ private static <T> ArrayList<FlinkSourceSplit<T>> toFlinkSplits(
+ List<? extends Source<T>> beamSplits) {
+ ArrayList<FlinkSourceSplit<T>> flinkSplits = new
ArrayList<>(beamSplits.size());
+ for (int i = 0; i < beamSplits.size(); i++) {
+ flinkSplits.add(new FlinkSourceSplit<>(i, beamSplits.get(i)));
+ }
+ return flinkSplits;
+ }
+}
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
index 94c14b2999b..bd07d00cd04 100644
---
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java
@@ -19,123 +19,109 @@ package
org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
-import java.util.concurrent.CountDownLatch;
+import java.util.Optional;
import javax.annotation.Nullable;
-import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import
org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceEnumeratorState.AssignmentMode;
import org.apache.beam.sdk.io.BoundedSource;
-import org.apache.beam.sdk.io.FileBasedSource;
import org.apache.beam.sdk.io.Source;
-import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-/**
- * A Flink {@link org.apache.flink.api.connector.source.SplitEnumerator
SplitEnumerator}
- * implementation that holds a Beam {@link Source} and does the following:
- *
- * <ul>
- * <li>Split the Beam {@link Source} to desired number of splits.
- * <li>Lazily assign the splits to the Flink Source Reader.
- * </ul>
- *
- * @param <T> The output type of the encapsulated Beam {@link Source}.
- */
+/** Splits a bounded Beam source and assigns one split for each reader
request. */
public class LazyFlinkSourceSplitEnumerator<T>
- implements SplitEnumerator<FlinkSourceSplit<T>, Map<Integer,
List<FlinkSourceSplit<T>>>> {
+ implements SplitEnumerator<FlinkSourceSplit<T>,
FlinkSourceEnumeratorState<T>> {
private static final Logger LOG =
LoggerFactory.getLogger(LazyFlinkSourceSplitEnumerator.class);
+
private final SplitEnumeratorContext<FlinkSourceSplit<T>> context;
private final Source<T> beamSource;
private final PipelineOptions pipelineOptions;
private final int numSplits;
private final List<FlinkSourceSplit<T>> pendingSplits;
- private volatile boolean splitsInitialized;
- private final CountDownLatch initializationLatch = new CountDownLatch(1);
+ private final Map<Integer, Optional<String>> pendingSplitRequests;
+
+ private boolean splitsInitialized;
+
+ public LazyFlinkSourceSplitEnumerator(
+ SplitEnumeratorContext<FlinkSourceSplit<T>> context,
+ Source<T> beamSource,
+ PipelineOptions pipelineOptions,
+ int numSplits) {
+ this(context, beamSource, pipelineOptions, numSplits, null);
+ }
public LazyFlinkSourceSplitEnumerator(
SplitEnumeratorContext<FlinkSourceSplit<T>> context,
Source<T> beamSource,
PipelineOptions pipelineOptions,
int numSplits,
- boolean splitInitialized) {
+ @Nullable FlinkSourceEnumeratorState<T> restoredState) {
this.context = context;
this.beamSource = beamSource;
this.pipelineOptions = pipelineOptions;
this.numSplits = numSplits;
this.pendingSplits = new ArrayList<>(numSplits);
- this.splitsInitialized = splitInitialized;
+ this.pendingSplitRequests = new LinkedHashMap<>();
+ this.splitsInitialized = restoredState != null;
+
+ if (restoredState != null) {
+ if (restoredState.getAssignmentMode() != AssignmentMode.LAZY) {
+ throw new IllegalArgumentException(
+ "Cannot restore the lazy source enumerator from "
+ + restoredState.getAssignmentMode()
+ + " state.");
+ }
+ pendingSplits.addAll(restoredState.getPendingSplits());
+ }
+
+ LOG.info(
+ "Created lazy source enumerator with parallelism {}, source {},
numSplits {}, "
+ + "initialized {}",
+ context.currentParallelism(),
+ beamSource,
+ numSplits,
+ splitsInitialized);
}
@Override
public void start() {
if (!splitsInitialized) {
initializeSplits();
+ } else {
+ sendPendingSplitRequests();
}
}
- public void initializeSplits() {
+ private void initializeSplits() {
context.callAsync(
- () -> {
- try {
- LOG.info("Starting source {}", beamSource);
- List<? extends Source<T>> beamSplitSourceList = splitBeamSource();
- int i = 0;
- for (Source<T> beamSplitSource : beamSplitSourceList) {
- pendingSplits.add(new FlinkSourceSplit<>(i, beamSplitSource));
- i++;
- }
- return pendingSplits;
- } catch (Exception e) {
- throw new RuntimeException(e);
- } finally {
- initializationLatch.countDown();
- }
- },
+ this::splitBeamSource,
(sourceSplits, error) -> {
if (error != null) {
- pendingSplits.addAll(sourceSplits);
throw new RuntimeException("Failed to start source enumerator.",
error);
}
+ pendingSplits.addAll(sourceSplits);
splitsInitialized = true;
+ sendPendingSplitRequests();
});
}
@Override
- public void handleSplitRequest(int subtask, @Nullable String hostname) {
- if (!context.registeredReaders().containsKey(subtask)) {
- // reader failed between sending the request and now. skip this request.
+ public void handleSplitRequest(int subtaskId, @Nullable String
requesterHostname) {
+ if (!context.registeredReaders().containsKey(subtaskId)) {
return;
}
- if (LOG.isInfoEnabled()) {
- final String hostInfo =
- hostname == null ? "(no host locality info)" : "(on host '" +
hostname + "')";
- LOG.info("Subtask {} {} is requesting a file source split", subtask,
hostInfo);
- }
-
if (!splitsInitialized) {
- try {
- initializationLatch.await();
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- LOG.warn("Interrupted while waiting for splits initialization", e);
- return;
- }
+ pendingSplitRequests.put(subtaskId,
Optional.ofNullable(requesterHostname));
+ return;
}
- if (!pendingSplits.isEmpty()) {
- final FlinkSourceSplit<T> split =
pendingSplits.remove(pendingSplits.size() - 1);
- context.assignSplit(split, subtask);
- LOG.info("Assigned split to subtask {} : {}", subtask, split);
- } else {
- context.signalNoMoreSplits(subtask);
- LOG.info("No more splits available for subtask {}", subtask);
- }
+ assignNextSplit(subtaskId, requesterHostname);
}
@Override
@@ -146,20 +132,14 @@ public class LazyFlinkSourceSplitEnumerator<T>
@Override
public void addReader(int subtaskId) {
- // this source is purely lazy-pull-based, nothing to do upon registration
+ // Readers request lazy splits when they are ready for work.
}
@Override
- public Map<Integer, List<FlinkSourceSplit<T>>> snapshotState(long
checkpointId) throws Exception {
+ public FlinkSourceEnumeratorState<T> snapshotState(long checkpointId) {
LOG.info("Taking snapshot for checkpoint {}", checkpointId);
- return snapshotState();
- }
-
- public Map<Integer, List<FlinkSourceSplit<T>>> snapshotState() throws
Exception {
- // For type compatibility reasons, we return a Map but we do not actually
care about the key
- Map<Integer, List<FlinkSourceSplit<T>>> state = new HashMap<>(1);
- state.put(1, pendingSplits);
- return state;
+ AssignmentMode mode = splitsInitialized ? AssignmentMode.LAZY :
AssignmentMode.UNDECIDED;
+ return new FlinkSourceEnumeratorState<>(mode, new
ArrayList<>(pendingSplits));
}
@Override
@@ -167,39 +147,37 @@ public class LazyFlinkSourceSplitEnumerator<T>
// NoOp
}
- private long getDesiredSizeBytes(int numSplits, BoundedSource<T>
boundedSource) throws Exception {
- long totalSize = boundedSource.getEstimatedSizeBytes(pipelineOptions);
- long defaultSplitSize = totalSize / numSplits;
- long maxSplitSize = 0;
- if (pipelineOptions != null) {
- maxSplitSize =
pipelineOptions.as(FlinkPipelineOptions.class).getFileInputSplitMaxSizeMB();
- }
- if (beamSource instanceof FileBasedSource && maxSplitSize > 0) {
- // Most of the time parallelism is < number of files in source.
- // Each file becomes a unique split which commonly create skew.
- // This limits the size of splits to reduce skew.
- return Math.min(defaultSplitSize, maxSplitSize * 1024 * 1024);
- } else {
- return defaultSplitSize;
+ private ArrayList<FlinkSourceSplit<T>> splitBeamSource() throws Exception {
+ if (!(beamSource instanceof BoundedSource)) {
+ throw new IllegalStateException("Lazy assignment requires a bounded
source.");
}
+ LOG.info("Starting source {}", beamSource);
+ BoundedSource<T> boundedSource = (BoundedSource<T>) beamSource;
+ long estimatedSizeBytes =
+ FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource,
pipelineOptions);
+ return FlinkSourceSplitUtils.splitBoundedSource(
+ boundedSource, pipelineOptions, numSplits, estimatedSizeBytes);
}
- // -------------- Private helper methods ----------------------
- private List<? extends Source<T>> splitBeamSource() throws Exception {
- if (beamSource instanceof BoundedSource) {
- BoundedSource<T> boundedSource = (BoundedSource<T>) beamSource;
- long desiredSizeBytes = getDesiredSizeBytes(numSplits, boundedSource);
- List<? extends BoundedSource<T>> splits =
- ((BoundedSource<T>) beamSource).split(desiredSizeBytes,
pipelineOptions);
- LOG.info("Split bounded source {} in {} splits", beamSource,
splits.size());
- return splits;
- } else if (beamSource instanceof UnboundedSource) {
- List<? extends UnboundedSource<T, ?>> splits =
- ((UnboundedSource<T, ?>) beamSource).split(numSplits,
pipelineOptions);
- LOG.info("Split source {} to {} splits", beamSource, splits);
- return splits;
- } else {
- throw new IllegalStateException("Unknown source type " +
beamSource.getClass());
+ private void sendPendingSplitRequests() {
+ Map<Integer, Optional<String>> splitRequests = new
LinkedHashMap<>(pendingSplitRequests);
+ pendingSplitRequests.clear();
+ splitRequests.forEach(
+ (subtaskId, hostname) -> assignNextSplit(subtaskId,
hostname.orElse(null)));
+ }
+
+ private void assignNextSplit(int subtaskId, @Nullable String
requesterHostname) {
+ if (!context.registeredReaders().containsKey(subtaskId)) {
+ return;
+ }
+ if (pendingSplits.isEmpty()) {
+ context.signalNoMoreSplits(subtaskId);
+ LOG.info("No more splits available for subtask {}", subtaskId);
+ return;
}
+
+ FlinkSourceSplit<T> split = pendingSplits.remove(pendingSplits.size() - 1);
+ context.assignSplit(split, subtaskId);
+ LOG.info("Assigned split to subtask {} on host {}: {}", subtaskId,
requesterHostname, split);
}
}
diff --git
a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/SizeBasedFlinkSourceSplitEnumerator.java
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/SizeBasedFlinkSourceSplitEnumerator.java
new file mode 100644
index 00000000000..78b0fb59f2c
--- /dev/null
+++
b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/SizeBasedFlinkSourceSplitEnumerator.java
@@ -0,0 +1,205 @@
+/*
+ * 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.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import
org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceEnumeratorState.AssignmentMode;
+import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.flink.api.connector.source.SplitEnumerator;
+import org.apache.flink.api.connector.source.SplitEnumeratorContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Selects static or lazy assignment by estimated bounded-source size. */
+final class SizeBasedFlinkSourceSplitEnumerator<T>
+ implements SplitEnumerator<FlinkSourceSplit<T>,
FlinkSourceEnumeratorState<T>> {
+ private static final Logger LOG =
+ LoggerFactory.getLogger(SizeBasedFlinkSourceSplitEnumerator.class);
+
+ private final SplitEnumeratorContext<FlinkSourceSplit<T>> context;
+ private final BoundedSource<T> boundedSource;
+ private final PipelineOptions pipelineOptions;
+ private final int numSplits;
+ private final Map<Integer, Optional<String>> pendingSplitRequests;
+ private final List<ReturnedSplits<T>> returnedSplits;
+
+ private @Nullable SplitEnumerator<FlinkSourceSplit<T>,
FlinkSourceEnumeratorState<T>> delegate;
+
+ SizeBasedFlinkSourceSplitEnumerator(
+ SplitEnumeratorContext<FlinkSourceSplit<T>> context,
+ BoundedSource<T> boundedSource,
+ PipelineOptions pipelineOptions,
+ int numSplits) {
+ this.context = context;
+ this.boundedSource = boundedSource;
+ this.pipelineOptions = pipelineOptions;
+ this.numSplits = numSplits;
+ this.pendingSplitRequests = new LinkedHashMap<>();
+ this.returnedSplits = new ArrayList<>();
+ }
+
+ @Override
+ public void start() {
+ context.callAsync(
+ this::selectAndSplit,
+ (initialState, error) -> {
+ if (error != null) {
+ throw new RuntimeException("Failed to select a source split
assignment mode.", error);
+ }
+
+ SplitEnumerator<FlinkSourceSplit<T>, FlinkSourceEnumeratorState<T>>
selectedDelegate =
+ createDelegate(initialState);
+ delegate = selectedDelegate;
+ returnedSplits.forEach(
+ returned -> selectedDelegate.addSplitsBack(returned.splits,
returned.subtaskId));
+ returnedSplits.clear();
+ selectedDelegate.start();
+ pendingSplitRequests.forEach(
+ (subtaskId, hostname) ->
+ selectedDelegate.handleSplitRequest(subtaskId,
hostname.orElse(null)));
+ pendingSplitRequests.clear();
+ });
+ }
+
+ @Override
+ public void handleSplitRequest(int subtaskId, @Nullable String
requesterHostname) {
+ if (delegate == null) {
+ pendingSplitRequests.put(subtaskId,
Optional.ofNullable(requesterHostname));
+ } else {
+ delegate.handleSplitRequest(subtaskId, requesterHostname);
+ }
+ }
+
+ @Override
+ public void addSplitsBack(List<FlinkSourceSplit<T>> splits, int subtaskId) {
+ if (delegate == null) {
+ returnedSplits.add(new ReturnedSplits<>(new ArrayList<>(splits),
subtaskId));
+ } else {
+ delegate.addSplitsBack(splits, subtaskId);
+ }
+ }
+
+ @Override
+ public void addReader(int subtaskId) {
+ if (delegate != null) {
+ delegate.addReader(subtaskId);
+ }
+ }
+
+ @Override
+ public FlinkSourceEnumeratorState<T> snapshotState(long checkpointId) throws
Exception {
+ if (delegate == null) {
+ return new FlinkSourceEnumeratorState<>(AssignmentMode.UNDECIDED, new
ArrayList<>());
+ }
+ return delegate.snapshotState(checkpointId);
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (delegate != null) {
+ delegate.close();
+ }
+ }
+
+ private FlinkSourceEnumeratorState<T> selectAndSplit() throws Exception {
+ long estimatedSizeBytes =
+ FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource,
pipelineOptions);
+ AssignmentMode selectedMode = selectAssignmentMode(estimatedSizeBytes);
+ ArrayList<FlinkSourceSplit<T>> splits =
+ FlinkSourceSplitUtils.splitBoundedSource(
+ boundedSource, pipelineOptions, numSplits, estimatedSizeBytes);
+ LOG.info(
+ "Split bounded source {} into {} splits using {} assignment",
+ boundedSource,
+ splits.size(),
+ selectedMode);
+ return new FlinkSourceEnumeratorState<>(selectedMode, splits);
+ }
+
+ private AssignmentMode selectAssignmentMode(long estimatedSizeBytes) {
+ long thresholdMb =
+
pipelineOptions.as(FlinkPipelineOptions.class).getSourceStaticSplitThresholdMb();
+ if (thresholdMb <= 0) {
+ throw new IllegalArgumentException(
+ "Size-based source assignment requires a positive threshold, but
received "
+ + thresholdMb
+ + ".");
+ }
+ if (estimatedSizeBytes <= 0 || estimatedSizeBytes == Long.MAX_VALUE) {
+ LOG.info(
+ "Estimated size of bounded source {} is zero or unknown. Using lazy
split assignment.",
+ boundedSource);
+ return AssignmentMode.LAZY;
+ }
+
+ int sourceParallelism = context.currentParallelism();
+ if (sourceParallelism <= 0) {
+ throw new IllegalStateException(
+ "Source parallelism must be positive, but was " + sourceParallelism
+ ".");
+ }
+ long estimatedBytesPerReader = estimatedSizeBytes / sourceParallelism;
+ long thresholdBytes = FlinkSourceSplitUtils.mebibytesToBytes(thresholdMb);
+ AssignmentMode selectedMode =
+ estimatedBytesPerReader >= thresholdBytes ? AssignmentMode.LAZY :
AssignmentMode.STATIC;
+ LOG.info(
+ "Using {} split assignment for bounded source {}: estimated size {}
bytes, source "
+ + "parallelism {}, estimated bytes per reader {}, static
assignment threshold {} "
+ + "bytes",
+ selectedMode,
+ boundedSource,
+ estimatedSizeBytes,
+ sourceParallelism,
+ estimatedBytesPerReader,
+ thresholdBytes);
+ return selectedMode;
+ }
+
+ private SplitEnumerator<FlinkSourceSplit<T>, FlinkSourceEnumeratorState<T>>
createDelegate(
+ FlinkSourceEnumeratorState<T> initialState) {
+ if (initialState.getAssignmentMode() == AssignmentMode.LAZY) {
+ return new LazyFlinkSourceSplitEnumerator<>(
+ context, boundedSource, pipelineOptions, numSplits, initialState);
+ }
+ if (initialState.getAssignmentMode() == AssignmentMode.STATIC) {
+ return new FlinkSourceSplitEnumerator<>(
+ context, boundedSource, pipelineOptions, numSplits, initialState);
+ }
+ throw new IllegalArgumentException(
+ "Cannot create a source enumerator for "
+ + initialState.getAssignmentMode()
+ + " assignment.");
+ }
+
+ private static final class ReturnedSplits<T> {
+ private final List<FlinkSourceSplit<T>> splits;
+ private final int subtaskId;
+
+ private ReturnedSplits(List<FlinkSourceSplit<T>> splits, int subtaskId) {
+ this.splits = splits;
+ this.subtaskId = subtaskId;
+ }
+ }
+}
diff --git
a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
index e911ee72db5..c4d220d2213 100644
---
a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
+++
b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumeratorTest.java
@@ -18,23 +18,303 @@
package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.stream.Collectors;
+import java.util.concurrent.atomic.AtomicInteger;
+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.TestBoundedCountingSource;
import
org.apache.beam.runners.flink.translation.wrappers.streaming.io.TestCountingSource;
+import
org.apache.beam.runners.flink.translation.wrappers.streaming.io.source.FlinkSourceEnumeratorState.AssignmentMode;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.Source;
+import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.values.KV;
+import org.apache.flink.api.connector.source.SplitEnumerator;
import
org.apache.flink.connector.testutils.source.reader.TestingSplitEnumeratorContext;
+import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.junit.Test;
-/** Unit tests for {@link FlinkSourceSplitEnumerator}. */
+/** Unit tests for the Flink source split enumerators. */
public class FlinkSourceSplitEnumeratorTest {
+ private static final long MEBIBYTE = 1024L * 1024L;
+ private static final long STATIC_SPLIT_THRESHOLD_MB = 6144L;
+ private static final int SOURCE_PARALLELISM = 2;
+ private static final int REQUESTED_SPLITS = 4;
+
+ @Test
+ public void testSelectsAssignmentModeFromEstimatedSizeAndConfiguration()
throws Exception {
+ assertEquals(0L, (long)
FlinkPipelineOptions.defaults().getSourceStaticSplitThresholdMb());
+
+ long thresholdBytes = STATIC_SPLIT_THRESHOLD_MB * MEBIBYTE;
+ long largeEstimate = 1024L * MEBIBYTE;
+ long[] estimatedSizes = {
+ SOURCE_PARALLELISM * thresholdBytes - 1L,
+ SOURCE_PARALLELISM * thresholdBytes,
+ -1L,
+ 0L,
+ Long.MAX_VALUE,
+ largeEstimate,
+ largeEstimate,
+ largeEstimate,
+ largeEstimate
+ };
+ long[] configuredThresholds = {
+ STATIC_SPLIT_THRESHOLD_MB,
+ STATIC_SPLIT_THRESHOLD_MB,
+ STATIC_SPLIT_THRESHOLD_MB,
+ STATIC_SPLIT_THRESHOLD_MB,
+ STATIC_SPLIT_THRESHOLD_MB,
+ -1L,
+ -100L,
+ 0L,
+ Long.MAX_VALUE
+ };
+ AssignmentMode[] expectedModes = {
+ AssignmentMode.STATIC,
+ AssignmentMode.LAZY,
+ AssignmentMode.LAZY,
+ AssignmentMode.LAZY,
+ AssignmentMode.LAZY,
+ AssignmentMode.STATIC,
+ AssignmentMode.STATIC,
+ AssignmentMode.LAZY,
+ AssignmentMode.STATIC
+ };
+
+ for (int i = 0; i < estimatedSizes.length; i++) {
+ assertAssignmentMode(
+ estimatedSizes[i], configuredThresholds[i], expectedModes[i],
REQUESTED_SPLITS);
+ }
+ }
+
+ @Test
+ public void testSizeBasedSelectionReplaysLazyRequestAfterInitialization()
throws Exception {
+ FlinkPipelineOptions options = thresholdOptions();
+ long estimatedSizeBytes = SOURCE_PARALLELISM * STATIC_SPLIT_THRESHOLD_MB *
MEBIBYTE;
+ TestEstimatedSizeBoundedSource testSource =
+ TestEstimatedSizeBoundedSource.create(estimatedSizeBytes,
REQUESTED_SPLITS);
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> context =
+ new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM);
+ FlinkSource<String, ?> flinkSource = createBoundedSource(testSource,
options, REQUESTED_SPLITS);
+
+ try (SplitEnumerator<FlinkSourceSplit<String>,
FlinkSourceEnumeratorState<String>> enumerator =
+ flinkSource.createEnumerator(context)) {
+ enumerator.start();
+ context.registerReader(0, "reader-0");
+ enumerator.addReader(0);
+ enumerator.handleSplitRequest(0, "reader-0");
+
+ context.getExecutorService().triggerAll();
+
+ assertEquals(1,
context.getSplitAssignments().get(0).getAssignedSplits().size());
+ assertEquals(AssignmentMode.LAZY,
enumerator.snapshotState(1L).getAssignmentMode());
+ }
+ }
+
+ @Test
+ public void testSizeBasedSelectionAssignsStaticSplitsToEarlyReaders() throws
Exception {
+ FlinkPipelineOptions options = thresholdOptions();
+ long estimatedSizeBytes = SOURCE_PARALLELISM * STATIC_SPLIT_THRESHOLD_MB *
MEBIBYTE - 1L;
+ TestEstimatedSizeBoundedSource testSource =
+ TestEstimatedSizeBoundedSource.create(estimatedSizeBytes,
REQUESTED_SPLITS);
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> context =
+ new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM);
+ FlinkSource<String, ?> flinkSource = createBoundedSource(testSource,
options, REQUESTED_SPLITS);
+
+ try (SplitEnumerator<FlinkSourceSplit<String>,
FlinkSourceEnumeratorState<String>> enumerator =
+ flinkSource.createEnumerator(context)) {
+ enumerator.start();
+ for (int subtaskId = 0; subtaskId < SOURCE_PARALLELISM; subtaskId++) {
+ context.registerReader(subtaskId, "reader-" + subtaskId);
+ enumerator.addReader(subtaskId);
+ }
+
+ context.getExecutorService().triggerAll();
+
+ assertEquals(REQUESTED_SPLITS, countAssignedSplits(context));
+ context
+ .getSplitAssignments()
+ .values()
+ .forEach(state -> assertTrue(state.hasReceivedNoMoreSplitsSignal()));
+ assertEquals(AssignmentMode.STATIC,
enumerator.snapshotState(1L).getAssignmentMode());
+ }
+ }
+
+ @Test
+ public void testRestoreKeepsLazyAssignmentAcrossRescaleWithoutEstimating()
throws Exception {
+ final int initialParallelism = 4;
+ final int restoredParallelism = 1;
+ final int generatedSplits = 4;
+ FlinkPipelineOptions options = thresholdOptions();
+ long thresholdBytes = STATIC_SPLIT_THRESHOLD_MB * MEBIBYTE;
+ long estimatedSizeBytes = initialParallelism * thresholdBytes;
+ TestEstimatedSizeBoundedSource testSource =
+ TestEstimatedSizeBoundedSource.create(estimatedSizeBytes,
generatedSplits);
+ FlinkSource<String, ?> flinkSource = createBoundedSource(testSource,
options, generatedSplits);
+ FlinkSourceEnumeratorState<String> checkpoint;
+
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> initialContext =
+ new TestingSplitEnumeratorContext<>(initialParallelism);
+ try (SplitEnumerator<FlinkSourceSplit<String>,
FlinkSourceEnumeratorState<String>> enumerator =
+ flinkSource.createEnumerator(initialContext)) {
+ enumerator.start();
+ initialContext.getExecutorService().triggerAll();
+ checkpoint = roundTripState(flinkSource, enumerator.snapshotState(1L));
+ assertEquals(AssignmentMode.LAZY, checkpoint.getAssignmentMode());
+ assertEquals(generatedSplits, checkpoint.getPendingSplits().size());
+ }
+
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> restoredContext =
+ new TestingSplitEnumeratorContext<>(restoredParallelism);
+ try (SplitEnumerator<FlinkSourceSplit<String>,
FlinkSourceEnumeratorState<String>> restored =
+ flinkSource.restoreEnumerator(restoredContext, checkpoint)) {
+ restored.start();
+ restoredContext.registerReader(0, "reader-0");
+ restored.addReader(0);
+ for (int i = 0; i < generatedSplits; i++) {
+ restored.handleSplitRequest(0, "reader-0");
+ }
+ restored.handleSplitRequest(0, "reader-0");
+
+ assertEquals(
+ generatedSplits,
restoredContext.getSplitAssignments().get(0).getAssignedSplits().size());
+
assertTrue(restoredContext.getSplitAssignments().get(0).hasReceivedNoMoreSplitsSignal());
+ assertEquals(
+ "restoring a decided strategy must not estimate the source again",
+ 1,
+ testSource.getEstimationCalls());
+ assertEquals(AssignmentMode.LAZY,
restored.snapshotState(2L).getAssignmentMode());
+ }
+ }
+
+ @Test
+ public void testLegacyLazyCheckpointMapUpgradesToStrategyNeutralState()
throws Exception {
+ FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
+ TestEstimatedSizeBoundedSource testSource =
TestEstimatedSizeBoundedSource.create(1L, 1);
+ FlinkSource<String, ?> flinkSource = createBoundedSource(testSource,
options, 1);
+ FlinkSourceSplit<String> pendingSplit = new FlinkSourceSplit<>(0,
testSource);
+ Map<Integer, List<FlinkSourceSplit<String>>> legacyCheckpoint =
+ Collections.singletonMap(1, Collections.singletonList(pendingSplit));
+ byte[] serialized = SerdeUtils.serializeObject(legacyCheckpoint);
+
+ FlinkSourceEnumeratorState<String> upgraded =
+ flinkSource.getEnumeratorCheckpointSerializer().deserialize(0,
serialized);
+
+ assertEquals(AssignmentMode.LAZY, upgraded.getAssignmentMode());
+ assertEquals(1, upgraded.getPendingSplits().size());
+ assertEquals(0, upgraded.getPendingSplits().get(0).splitIndex());
+ }
+
+ @Test
+ public void testLegacyStaticCheckpointMapKeepsSplitsOnTheirOriginalReaders()
throws Exception {
+ final int parallelism = 2;
+ TestEstimatedSizeBoundedSource testSource =
TestEstimatedSizeBoundedSource.create(1L, 1);
+ Map<Integer, List<FlinkSourceSplit<String>>> legacyCheckpoint = new
HashMap<>();
+ legacyCheckpoint.put(
+ 0,
+ Arrays.asList(
+ new FlinkSourceSplit<>(0, testSource), new FlinkSourceSplit<>(2,
testSource)));
+ legacyCheckpoint.put(1, Collections.singletonList(new
FlinkSourceSplit<>(1, testSource)));
+ byte[] serialized = SerdeUtils.serializeObject(legacyCheckpoint);
+
+ FlinkSourceEnumeratorState<String> upgraded =
+ new FlinkSourceEnumeratorStateSerializer<String>(AssignmentMode.STATIC)
+ .deserialize(0, serialized);
+ assertEquals(AssignmentMode.STATIC, upgraded.getAssignmentMode());
+ assertEquals(3, upgraded.getPendingSplits().size());
+
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> restoredContext =
+ new TestingSplitEnumeratorContext<>(parallelism);
+ try (FlinkSourceSplitEnumerator<String> restored =
+ new FlinkSourceSplitEnumerator<>(
+ restoredContext, testSource, staticOptions(), 3, upgraded)) {
+ restored.start();
+ for (int subtaskId = 0; subtaskId < parallelism; subtaskId++) {
+ restoredContext.registerReader(subtaskId, "reader-" + subtaskId);
+ restored.addReader(subtaskId);
+ }
+
+ assertEquals(Arrays.asList(0, 2),
assignedSplitIndexesForSubtask(restoredContext, 0));
+ assertEquals(
+ Collections.singletonList(1),
assignedSplitIndexesForSubtask(restoredContext, 1));
+ }
+ }
+
+ @Test
+ public void testSerializerRejectsUnknownVersionsAndUnexpectedPayloads()
throws Exception {
+ FlinkSourceEnumeratorStateSerializer<String> serializer =
+ new FlinkSourceEnumeratorStateSerializer<>(AssignmentMode.LAZY);
+ byte[] legacyMapBytes =
+ SerdeUtils.serializeObject(Collections.singletonMap(1, new
ArrayList<>()));
+ byte[] stateBytes =
+ serializer.serialize(
+ new FlinkSourceEnumeratorState<>(AssignmentMode.LAZY, new
ArrayList<>()));
+
+ // The payload type must match the version it was written with.
+ assertThrows(IOException.class, () -> serializer.deserialize(1,
legacyMapBytes));
+ assertThrows(IOException.class, () -> serializer.deserialize(0,
stateBytes));
+ assertThrows(IOException.class, () -> serializer.deserialize(2,
stateBytes));
+ }
+
+ @Test
+ public void
testRestoreRepartitionsStaticSplitsForNewParallelismWithoutEstimating()
+ throws Exception {
+ final int initialParallelism = 2;
+ final int restoredParallelism = 3;
+ final int generatedSplits = 5;
+ FlinkPipelineOptions options = thresholdOptions();
+ long thresholdBytes = STATIC_SPLIT_THRESHOLD_MB * MEBIBYTE;
+ long estimatedSizeBytes = initialParallelism * thresholdBytes - 1L;
+ TestEstimatedSizeBoundedSource testSource =
+ TestEstimatedSizeBoundedSource.create(estimatedSizeBytes,
generatedSplits);
+ FlinkSource<String, ?> flinkSource = createBoundedSource(testSource,
options, generatedSplits);
+ FlinkSourceEnumeratorState<String> checkpoint;
+
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> initialContext =
+ new TestingSplitEnumeratorContext<>(initialParallelism);
+ try (SplitEnumerator<FlinkSourceSplit<String>,
FlinkSourceEnumeratorState<String>> enumerator =
+ flinkSource.createEnumerator(initialContext)) {
+ enumerator.start();
+ initialContext.getExecutorService().triggerAll();
+ checkpoint = roundTripState(flinkSource, enumerator.snapshotState(1L));
+ assertEquals(AssignmentMode.STATIC, checkpoint.getAssignmentMode());
+ assertEquals(generatedSplits, checkpoint.getPendingSplits().size());
+ }
+
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> restoredContext =
+ new TestingSplitEnumeratorContext<>(restoredParallelism);
+ try (SplitEnumerator<FlinkSourceSplit<String>,
FlinkSourceEnumeratorState<String>> restored =
+ flinkSource.restoreEnumerator(restoredContext, checkpoint)) {
+ restored.start();
+ for (int subtaskId = 0; subtaskId < restoredParallelism; subtaskId++) {
+ restoredContext.registerReader(subtaskId, "reader-" + subtaskId);
+ restored.addReader(subtaskId);
+ }
+
+ assertEquals(generatedSplits, countAssignedSplits(restoredContext));
+ assertEquals(2,
restoredContext.getSplitAssignments().get(0).getAssignedSplits().size());
+ assertEquals(2,
restoredContext.getSplitAssignments().get(1).getAssignedSplits().size());
+ assertEquals(1,
restoredContext.getSplitAssignments().get(2).getAssignedSplits().size());
+ restoredContext
+ .getSplitAssignments()
+ .values()
+ .forEach(state -> assertTrue(state.hasReceivedNoMoreSplitsSignal()));
+ assertEquals(1, testSource.getEstimationCalls());
+ }
+ }
@Test
public void testAssignSplitsWithBoundedSource() throws IOException {
@@ -54,13 +334,8 @@ public class FlinkSourceSplitEnumeratorTest {
.forEach(
(subtaskId, state) -> {
int expectedNumSplitsPerSubtask = numSplits / numSubtasks;
- assertEquals(
- "Each subtask should have " + expectedNumSplitsPerSubtask +
" assigned splits",
- expectedNumSplitsPerSubtask,
- state.getAssignedSplits().size());
- assertTrue(
- "Each subtask should have received NoMoreSplits",
- state.hasReceivedNoMoreSplitsSignal());
+ assertEquals(expectedNumSplitsPerSubtask,
state.getAssignedSplits().size());
+ assertTrue(state.hasReceivedNoMoreSplitsSignal());
state
.getAssignedSplits()
.forEach(
@@ -72,8 +347,8 @@ public class FlinkSourceSplitEnumeratorTest {
assertEquals(
expectedSplitSize,
source.getEstimatedSizeBytes(FlinkPipelineOptions.defaults()));
- } catch (Exception e) {
- fail("Received exception" + e);
+ } catch (Exception error) {
+ fail("Received exception " + error);
}
});
});
@@ -88,90 +363,79 @@ public class FlinkSourceSplitEnumeratorTest {
new TestingSplitEnumeratorContext<>(numSubtasks);
TestCountingSource testSource = new TestCountingSource(numRecordsPerSplit);
- assignSplits(testContext, testSource, numSplits);
+ try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> enumerator =
+ new FlinkSourceSplitEnumerator<>(
+ testContext, testSource, FlinkPipelineOptions.defaults(),
numSplits)) {
+ enumerator.start();
+ for (int subtaskId = 0; subtaskId < numSubtasks; subtaskId++) {
+ testContext.registerReader(subtaskId, String.valueOf(subtaskId));
+ enumerator.addReader(subtaskId);
+ }
+ testContext.getExecutorService().triggerAll();
+ }
testContext
.getSplitAssignments()
.forEach(
(subtaskId, state) -> {
- int expectedNumSplitsPerSubtask = numSplits / numSubtasks;
- assertEquals(
- "Each subtask should have " + expectedNumSplitsPerSubtask +
" assigned splits",
- expectedNumSplitsPerSubtask,
- state.getAssignedSplits().size());
- assertTrue(
- "Each subtask should have received NoMoreSplits",
- state.hasReceivedNoMoreSplitsSignal());
+ assertEquals(numSplits / numSubtasks,
state.getAssignedSplits().size());
+ assertTrue(state.hasReceivedNoMoreSplitsSignal());
});
}
@Test
- public void testAddSplitsBack() throws IOException {
+ public void testAddSplitsBackToStaticReader() throws IOException {
final int numSubtasks = 2;
final int numSplits = 10;
- final int totalNumRecords = 10;
TestingSplitEnumeratorContext<FlinkSourceSplit<KV<Integer, Integer>>>
testContext =
new TestingSplitEnumeratorContext<>(numSubtasks);
- TestBoundedCountingSource testSource =
- new TestBoundedCountingSource(numSplits, totalNumRecords);
- try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
- new FlinkSourceSplitEnumerator<>(
- testContext, testSource, FlinkPipelineOptions.defaults(),
numSplits)) {
- splitEnumerator.start();
+ TestBoundedCountingSource testSource = new
TestBoundedCountingSource(numSplits, numSplits);
+
+ try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> enumerator =
+ new FlinkSourceSplitEnumerator<>(testContext, testSource,
staticOptions(), numSplits)) {
+ enumerator.start();
testContext.registerReader(0, "0");
- splitEnumerator.addReader(0);
+ enumerator.addReader(0);
testContext.getExecutorService().triggerAll();
- List<FlinkSourceSplit<KV<Integer, Integer>>> splitsForReader =
- testContext.getSplitAssignments().get(0).getAssignedSplits();
- assertEquals(numSplits / numSubtasks, splitsForReader.size());
+ List<FlinkSourceSplit<KV<Integer, Integer>>> returnedSplits =
+ new
ArrayList<>(testContext.getSplitAssignments().get(0).getAssignedSplits());
+ assertEquals(numSplits / numSubtasks, returnedSplits.size());
- splitEnumerator.addSplitsBack(splitsForReader, 0);
- splitEnumerator.addReader(0);
- assertEquals(2 * numSplits / numSubtasks, splitsForReader.size());
+ enumerator.addSplitsBack(returnedSplits, 0);
+ enumerator.addReader(0);
+ assertEquals(
+ 2 * numSplits / numSubtasks,
+ testContext.getSplitAssignments().get(0).getAssignedSplits().size());
}
}
- @Test
- public void testAddSplitsBackAfterRescale() throws Exception {
- final int numSubtasks = 2;
- final int numSplits = 10;
- final int totalNumRecords = 10;
- TestingSplitEnumeratorContext<FlinkSourceSplit<KV<Integer, Integer>>>
testContext =
- new TestingSplitEnumeratorContext<>(numSubtasks);
- TestBoundedCountingSource testSource =
- new TestBoundedCountingSource(numSplits, totalNumRecords);
- final Map<Integer, List<FlinkSourceSplit<KV<Integer, Integer>>>>
assignment;
- try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
- new FlinkSourceSplitEnumerator<>(
- testContext, testSource, FlinkPipelineOptions.defaults(),
numSplits)) {
- splitEnumerator.start();
- for (int i = 0; i < numSubtasks; i++) {
- testContext.registerReader(i, String.valueOf(i));
- splitEnumerator.addReader(i);
- }
- testContext.getExecutorService().triggerAll();
- assignment =
- testContext.getSplitAssignments().entrySet().stream()
- .map(e -> KV.of(e.getKey(), e.getValue().getAssignedSplits()))
- .collect(Collectors.toMap(KV::getKey, KV::getValue));
- }
+ private void assertAssignmentMode(
+ long estimatedSizeBytes,
+ long configuredThresholdMb,
+ AssignmentMode expectedMode,
+ int generatedSplits)
+ throws Exception {
+ FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
+ options.setSourceStaticSplitThresholdMb(configuredThresholdMb);
+ TestEstimatedSizeBoundedSource testSource =
+ TestEstimatedSizeBoundedSource.create(estimatedSizeBytes,
generatedSplits);
+ TestingSplitEnumeratorContext<FlinkSourceSplit<String>> context =
+ new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM);
+ FlinkSource<String, ?> flinkSource = createBoundedSource(testSource,
options, generatedSplits);
- // add tasks back
- testContext = new TestingSplitEnumeratorContext<>(numSubtasks);
- try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
- new FlinkSourceSplitEnumerator<>(
- testContext, testSource, FlinkPipelineOptions.defaults(),
numSplits, true)) {
- splitEnumerator.start();
- assignment.forEach(
- (splitId, assignedSplits) ->
splitEnumerator.addSplitsBack(assignedSplits, splitId));
- testContext.registerReader(0, "0");
- splitEnumerator.addReader(0);
- testContext.getExecutorService().triggerAll();
-
- List<FlinkSourceSplit<KV<Integer, Integer>>> splitsForReader =
- testContext.getSplitAssignments().get(0).getAssignedSplits();
- assertEquals(numSplits / numSubtasks, splitsForReader.size());
+ try (SplitEnumerator<FlinkSourceSplit<String>,
FlinkSourceEnumeratorState<String>> enumerator =
+ flinkSource.createEnumerator(context)) {
+ if (configuredThresholdMb < 0) {
+ assertTrue(enumerator instanceof FlinkSourceSplitEnumerator);
+ } else if (configuredThresholdMb == 0) {
+ assertTrue(enumerator instanceof LazyFlinkSourceSplitEnumerator);
+ } else {
+ assertTrue(enumerator instanceof SizeBasedFlinkSourceSplitEnumerator);
+ }
+ enumerator.start();
+ context.getExecutorService().triggerAll();
+ assertEquals(expectedMode,
enumerator.snapshotState(1L).getAssignmentMode());
}
}
@@ -180,17 +444,107 @@ public class FlinkSourceSplitEnumeratorTest {
Source<KV<Integer, Integer>> source,
int numSplits)
throws IOException {
- try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> splitEnumerator =
- new FlinkSourceSplitEnumerator<>(
- context, source, FlinkPipelineOptions.defaults(), numSplits)) {
- splitEnumerator.start();
- // Add a reader before splitting the beam source.
- context.registerReader(0, "0");
- splitEnumerator.addReader(0);
+ try (FlinkSourceSplitEnumerator<KV<Integer, Integer>> enumerator =
+ new FlinkSourceSplitEnumerator<>(context, source, staticOptions(),
numSplits)) {
+ enumerator.start();
+ for (int subtaskId = 0; subtaskId < context.currentParallelism();
subtaskId++) {
+ context.registerReader(subtaskId, String.valueOf(subtaskId));
+ enumerator.addReader(subtaskId);
+ }
context.getExecutorService().triggerAll();
- context.registerReader(1, "1");
- // Add another reader after splitting the beam source.
- splitEnumerator.addReader(1);
+ }
+ }
+
+ private static FlinkPipelineOptions staticOptions() {
+ FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
+ options.setSourceStaticSplitThresholdMb(-1L);
+ return options;
+ }
+
+ private static FlinkPipelineOptions thresholdOptions() {
+ FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
+ options.setSourceStaticSplitThresholdMb(STATIC_SPLIT_THRESHOLD_MB);
+ return options;
+ }
+
+ private static FlinkSource<String, ?> createBoundedSource(
+ BoundedSource<String> source, FlinkPipelineOptions options, int
numSplits) {
+ return FlinkSource.bounded(
+ "test-bounded-source", source, new
SerializablePipelineOptions(options), numSplits);
+ }
+
+ private static FlinkSourceEnumeratorState<String> roundTripState(
+ FlinkSource<String, ?> source, FlinkSourceEnumeratorState<String> state)
throws IOException {
+ SimpleVersionedSerializer<FlinkSourceEnumeratorState<String>> serializer =
+ source.getEnumeratorCheckpointSerializer();
+ byte[] serialized = serializer.serialize(state);
+ return serializer.deserialize(serializer.getVersion(), serialized);
+ }
+
+ private static <T> int countAssignedSplits(
+ TestingSplitEnumeratorContext<FlinkSourceSplit<T>> context) {
+ return context.getSplitAssignments().values().stream()
+ .mapToInt(state -> state.getAssignedSplits().size())
+ .sum();
+ }
+
+ private static <T> List<Integer> assignedSplitIndexesForSubtask(
+ TestingSplitEnumeratorContext<FlinkSourceSplit<T>> context, int
subtaskId) {
+ List<Integer> splitIndexes = new ArrayList<>();
+ for (FlinkSourceSplit<T> split :
+ context.getSplitAssignments().get(subtaskId).getAssignedSplits()) {
+ splitIndexes.add(split.splitIndex());
+ }
+ Collections.sort(splitIndexes);
+ return splitIndexes;
+ }
+
+ private static final class TestEstimatedSizeBoundedSource extends
BoundedSource<String> {
+ private final long estimatedSizeBytes;
+ private final int generatedSplits;
+ private final AtomicInteger estimationCalls;
+
+ private TestEstimatedSizeBoundedSource(
+ long estimatedSizeBytes, int generatedSplits, AtomicInteger
estimationCalls) {
+ this.estimatedSizeBytes = estimatedSizeBytes;
+ this.generatedSplits = generatedSplits;
+ this.estimationCalls = estimationCalls;
+ }
+
+ private static TestEstimatedSizeBoundedSource create(
+ long estimatedSizeBytes, int generatedSplits) {
+ return new TestEstimatedSizeBoundedSource(
+ estimatedSizeBytes, generatedSplits, new AtomicInteger());
+ }
+
+ @Override
+ public List<? extends BoundedSource<String>> split(
+ long desiredBundleSizeBytes, PipelineOptions options) {
+ List<TestEstimatedSizeBoundedSource> splits = new
ArrayList<>(generatedSplits);
+ for (int i = 0; i < generatedSplits; i++) {
+ splits.add(new TestEstimatedSizeBoundedSource(1L, 1, estimationCalls));
+ }
+ return splits;
+ }
+
+ @Override
+ public long getEstimatedSizeBytes(PipelineOptions options) {
+ estimationCalls.incrementAndGet();
+ return estimatedSizeBytes;
+ }
+
+ @Override
+ public BoundedReader<String> createReader(PipelineOptions options) {
+ throw new UnsupportedOperationException("This source is only used to
test split assignment");
+ }
+
+ @Override
+ public Coder<String> getOutputCoder() {
+ return StringUtf8Coder.of();
+ }
+
+ private int getEstimationCalls() {
+ return estimationCalls.get();
}
}
}
diff --git
a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html
b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html
index 34d6c524377..8e0fc007f13 100644
--- a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html
+++ b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html
@@ -187,6 +187,11 @@ Should be called before running the tests.
<td>Shuts down sources which have been idle for the configured time of
milliseconds. Once a source has been shut down, checkpointing is not possible
anymore. Shutting down the sources eventually leads to pipeline shutdown
(=Flink job finishes) once all input has been processed. Unless explicitly set,
this will default to Long.MAX_VALUE when checkpointing is enabled and to 0 when
checkpointing is disabled. See https://issues.apache.org/jira/browse/FLINK-2491
for progress on this issue.</td>
<td>Default: <code>-1</code></td>
</tr>
+<tr>
+ <td><code>sourceStaticSplitThresholdMb</code></td>
+ <td>Static split assignment threshold in MiB per source reader for bounded
sources in Flink DataStream mode. The default of 0 always uses lazy assignment.
A positive value selects static round-robin assignment for sources with a
known, positive estimate below the threshold and lazy assignment otherwise. Any
negative value always uses static assignment.</td>
+ <td>Default: <code>0</code></td>
+</tr>
<tr>
<td><code>stateBackend</code></td>
<td>State backend to store Beam's state. Use 'rocksdb' or 'hashmap' (same as
'filesystem').</td>
diff --git
a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html
b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html
index e3fe24216a5..d2ab687c42b 100644
--- a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html
+++ b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html
@@ -187,6 +187,11 @@ Should be called before running the tests.
<td>Shuts down sources which have been idle for the configured time of
milliseconds. Once a source has been shut down, checkpointing is not possible
anymore. Shutting down the sources eventually leads to pipeline shutdown
(=Flink job finishes) once all input has been processed. Unless explicitly set,
this will default to Long.MAX_VALUE when checkpointing is enabled and to 0 when
checkpointing is disabled. See https://issues.apache.org/jira/browse/FLINK-2491
for progress on this issue.</td>
<td>Default: <code>-1</code></td>
</tr>
+<tr>
+ <td><code>source_static_split_threshold_mb</code></td>
+ <td>Static split assignment threshold in MiB per source reader for bounded
sources in Flink DataStream mode. The default of 0 always uses lazy assignment.
A positive value selects static round-robin assignment for sources with a
known, positive estimate below the threshold and lazy assignment otherwise. Any
negative value always uses static assignment.</td>
+ <td>Default: <code>0</code></td>
+</tr>
<tr>
<td><code>state_backend</code></td>
<td>State backend to store Beam's state. Use 'rocksdb' or 'hashmap' (same as
'filesystem').</td>
