From f110427f6c295f4dfccfdf1331b959d25732d80b Mon Sep 17 00:00:00 2001 From: Paulius Kuzmickas Date: Mon, 24 Aug 2026 11:06:41 +0100 Subject: [PATCH] [Flink] Select bounded-source split assignment by size --- CHANGES.md | 1 + .../runners/flink/FlinkPipelineOptions.java | 11 + .../runners/flink/FlinkPipelineOptions.java | 11 + .../streaming/io/source/FlinkSource.java | 65 +- .../io/source/FlinkSourceEnumeratorState.java | 46 ++ .../FlinkSourceEnumeratorStateSerializer.java | 102 +++ .../FlinkSourceSplitAssignmentMode.java | 24 + .../io/source/FlinkSourceSplitEnumerator.java | 143 ++-- .../io/source/FlinkSourceSplitUtils.java | 84 ++ .../LazyFlinkSourceSplitEnumerator.java | 186 ++--- .../SizeBasedFlinkSourceSplitEnumerator.java | 208 +++++ .../FlinkSourceSplitEnumeratorTest.java | 771 ++++++++++++++++-- .../flink_java_pipeline_options.html | 5 + .../flink_python_pipeline_options.html | 5 + 14 files changed, 1375 insertions(+), 287 deletions(-) create mode 100644 runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorState.java create mode 100644 runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceEnumeratorStateSerializer.java create mode 100644 runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitAssignmentMode.java create mode 100644 runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitUtils.java create mode 100644 runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/SizeBasedFlinkSourceSplitEnumerator.java diff --git a/CHANGES.md b/CHANGES.md index 73966a48313c..a351c0b1f1ed 100644 --- a/CHANGES.md +++ b/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`, 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 3fee130d58ee..e6dfe5045712 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 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.") 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 f0724b4d031f..c26beac043a1 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 @@ -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.") 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 0c1ba73b2a2d..fc8dd359bee2 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,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 The data type of the records emitted by the Flink Source. */ public abstract class FlinkSource - implements Source, Map>>> { + implements Source, FlinkSourceEnumeratorState> { protected final String stepName; protected final org.apache.beam.sdk.io.Source beamSource; @@ -102,36 +99,38 @@ public Boundedness getBoundedness() { } @Override - public SplitEnumerator, Map>>> - createEnumerator(SplitEnumeratorContext> enumContext) throws Exception { - return createEnumerator(enumContext, false); - } - - public SplitEnumerator, Map>>> - createEnumerator( - SplitEnumeratorContext> enumContext, boolean splitInitialized) - throws Exception { - + public SplitEnumerator, FlinkSourceEnumeratorState> createEnumerator( + SplitEnumeratorContext> 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) beamSource, options, numSplits); + } + return new LazyFlinkSourceSplitEnumerator<>(enumContext, beamSource, options, numSplits); } + return new FlinkSourceSplitEnumerator<>(enumContext, beamSource, options, numSplits); } @Override - public SplitEnumerator, Map>>> - restoreEnumerator( - SplitEnumeratorContext> enumContext, - Map>> checkpoint) - throws Exception { - SplitEnumerator, Map>>> enumerator = - createEnumerator(enumContext, true); - checkpoint.forEach( - (subtaskId, splitsForSubtask) -> enumerator.addSplitsBack(splitsForSubtask, subtaskId)); - return enumerator; + public SplitEnumerator, FlinkSourceEnumeratorState> restoreEnumerator( + SplitEnumeratorContext> enumContext, + FlinkSourceEnumeratorState 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> getSplitSerializer() { } @Override - public SimpleVersionedSerializer>>> + public SimpleVersionedSerializer> getEnumeratorCheckpointSerializer() { - return SerdeUtils.getNaiveObjectSerializer(); + FlinkSourceSplitAssignmentMode legacyAssignmentMode = + boundedness == Boundedness.BOUNDED + ? FlinkSourceSplitAssignmentMode.LAZY + : FlinkSourceSplitAssignmentMode.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 000000000000..1c2ed31ee444 --- /dev/null +++ b/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 implements Serializable { + private static final long serialVersionUID = 1L; + + private final FlinkSourceSplitAssignmentMode assignmentMode; + private final ArrayList> pendingSplits; + + /** Takes ownership of {@code pendingSplits}; the caller must not mutate it afterwards. */ + FlinkSourceEnumeratorState( + FlinkSourceSplitAssignmentMode assignmentMode, ArrayList> pendingSplits) { + this.assignmentMode = assignmentMode; + this.pendingSplits = pendingSplits; + } + + FlinkSourceSplitAssignmentMode getAssignmentMode() { + return assignmentMode; + } + + List> getPendingSplits() { + return Collections.unmodifiableList(pendingSplits); + } +} 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 000000000000..fb4da57ace8b --- /dev/null +++ b/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 + implements SimpleVersionedSerializer> { + // Version written by SerdeUtils.getNaiveObjectSerializer, which serialized the + // Map>> 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 state) throws IOException { + return SerdeUtils.serializeObject(state); + } + + @Override + @SuppressWarnings("unchecked") + public FlinkSourceEnumeratorState deserialize(int version, byte[] serialized) + throws IOException { + if (version == VERSION) { + Object deserialized = SerdeUtils.deserializeObject(serialized); + if (deserialized instanceof FlinkSourceEnumeratorState) { + return (FlinkSourceEnumeratorState) 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 upgradeLegacyState(Map legacyState) + throws IOException { + ArrayList> pendingSplits = new ArrayList<>(); + for (Object value : legacyState.values()) { + List> legacySplits = (List>) 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/FlinkSourceSplitAssignmentMode.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitAssignmentMode.java new file mode 100644 index 000000000000..25e6b089a1ad --- /dev/null +++ b/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 +} 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 38d2f3639394..90af417adce4 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,9 +21,11 @@ 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.sdk.io.BoundedSource; import org.apache.beam.sdk.io.Source; @@ -35,27 +37,17 @@ 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: - * - *
    - *
  • Split the Beam {@link Source} to desired number of splits. - *
  • Assign the splits to the Flink Source Reader. - *
- * - *

Note that at this point, this class has a static round-robin split assignment strategy. - * - * @param 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 - implements SplitEnumerator, Map>>> { + implements SplitEnumerator, FlinkSourceEnumeratorState> { private static final Logger LOG = LoggerFactory.getLogger(FlinkSourceSplitEnumerator.class); + private final SplitEnumeratorContext> context; private final Source beamSource; private final PipelineOptions pipelineOptions; private final int numSplits; private final Map>> pendingSplits; + private boolean splitsInitialized; public FlinkSourceSplitEnumerator( @@ -63,8 +55,7 @@ public FlinkSourceSplitEnumerator( Source beamSource, PipelineOptions pipelineOptions, int numSplits) { - - this(context, beamSource, pipelineOptions, numSplits, false); + this(context, beamSource, pipelineOptions, numSplits, null); } public FlinkSourceSplitEnumerator( @@ -72,17 +63,31 @@ public FlinkSourceSplitEnumerator( Source beamSource, PipelineOptions pipelineOptions, int numSplits, - boolean splitsInitialized) { - + @Nullable FlinkSourceEnumeratorState 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() != FlinkSourceSplitAssignmentMode.STATIC) { + throw new IllegalArgumentException( + "Cannot restore the static source enumerator from " + + restoredState.getAssignmentMode() + + " state."); + } + int parallelism = context.currentParallelism(); + for (FlinkSourceSplit 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 +98,33 @@ public FlinkSourceSplitEnumerator( public void start() { if (!splitsInitialized) { initializeSplits(); + } else { + sendPendingSplitsToSourceReaders(); } } private void initializeSplits() { context.callAsync( - () -> { - try { - LOG.info("Starting source {}", beamSource); - List> beamSplitSourceList = splitBeamSource(); - Map>> flinkSourceSplitsList = new HashMap<>(); - int i = 0; - for (Source beamSplitSource : beamSplitSourceList) { - int targetSubtask = i % context.currentParallelism(); - List> 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> splits, int subtaskId) { LOG.info("Adding splits {} back from subtask {}", splits, subtaskId); - List> splitsForSubtask = - pendingSplits.computeIfAbsent(subtaskId, ignored -> new ArrayList<>()); - splitsForSubtask.addAll(splits); + pendingSplits.computeIfAbsent(subtaskId, ignored -> new ArrayList<>()).addAll(splits); } @Override @@ -146,18 +132,22 @@ public void addReader(int subtaskId) { List> 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>> snapshotState(long checkpointId) throws Exception { + public FlinkSourceEnumeratorState snapshotState(long checkpointId) { LOG.info("Taking snapshot for checkpoint {}", checkpointId); - return pendingSplits; + ArrayList> checkpointSplits = new ArrayList<>(); + pendingSplits.values().forEach(checkpointSplits::addAll); + FlinkSourceSplitAssignmentMode mode = + splitsInitialized + ? FlinkSourceSplitAssignmentMode.STATIC + : FlinkSourceSplitAssignmentMode.UNDECIDED; + return new FlinkSourceEnumeratorState<>(mode, checkpointSplits); } @Override @@ -165,34 +155,49 @@ public void close() throws IOException { // NoOp } - // -------------- Private helper methods ---------------------- - private List> splitBeamSource() throws Exception { + private ArrayList> splitBeamSource() throws Exception { + LOG.info("Starting source {}", beamSource); if (beamSource instanceof BoundedSource) { BoundedSource boundedSource = (BoundedSource) beamSource; - long desiredSizeBytes = boundedSource.getEstimatedSizeBytes(pipelineOptions) / numSplits; - return boundedSource.split(desiredSizeBytes, pipelineOptions); - } else if (beamSource instanceof UnboundedSource) { - List> splits = - ((UnboundedSource) 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) beamSource, pipelineOptions, numSplits); + } + throw new IllegalStateException("Unknown source type " + beamSource.getClass()); + } + + private void prepareAssignments(List> sourceSplits) { + int parallelism = context.currentParallelism(); + for (FlinkSourceSplit split : sourceSplits) { + int targetSubtask = split.splitIndex() % parallelism; + pendingSplits.computeIfAbsent(targetSubtask, ignored -> new ArrayList<>()).add(split); } } private void sendPendingSplitsToSourceReaders() { + Set assignedReaders = new HashSet<>(); Iterator>>> splitIter = pendingSplits.entrySet().iterator(); while (splitIter.hasNext()) { Map.Entry>> 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> 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 000000000000..dda99bf53bef --- /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 long estimateBoundedSourceSize( + BoundedSource boundedSource, PipelineOptions pipelineOptions) throws Exception { + return boundedSource.getEstimatedSizeBytes(pipelineOptions); + } + + static ArrayList> splitBoundedSource( + BoundedSource boundedSource, + PipelineOptions pipelineOptions, + int numSplits, + long estimatedSizeBytes) + throws Exception { + long desiredSizeBytes = + getDesiredSizeBytes(boundedSource, pipelineOptions, numSplits, estimatedSizeBytes); + return toFlinkSplits(boundedSource.split(desiredSizeBytes, pipelineOptions)); + } + + static ArrayList> splitUnboundedSource( + UnboundedSource 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 ArrayList> toFlinkSplits( + List> beamSplits) { + ArrayList> 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 94c14b2999b9..b4336cf84d65 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,108 @@ 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.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: - * - *

    - *
  • Split the Beam {@link Source} to desired number of splits. - *
  • Lazily assign the splits to the Flink Source Reader. - *
- * - * @param 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 - implements SplitEnumerator, Map>>> { + implements SplitEnumerator, FlinkSourceEnumeratorState> { private static final Logger LOG = LoggerFactory.getLogger(LazyFlinkSourceSplitEnumerator.class); + private final SplitEnumeratorContext> context; private final Source beamSource; private final PipelineOptions pipelineOptions; private final int numSplits; private final List> pendingSplits; - private volatile boolean splitsInitialized; - private final CountDownLatch initializationLatch = new CountDownLatch(1); + private final Map> pendingSplitRequests; + + private boolean splitsInitialized; + + public LazyFlinkSourceSplitEnumerator( + SplitEnumeratorContext> context, + Source beamSource, + PipelineOptions pipelineOptions, + int numSplits) { + this(context, beamSource, pipelineOptions, numSplits, null); + } public LazyFlinkSourceSplitEnumerator( SplitEnumeratorContext> context, Source beamSource, PipelineOptions pipelineOptions, int numSplits, - boolean splitInitialized) { + @Nullable FlinkSourceEnumeratorState 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() != FlinkSourceSplitAssignmentMode.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> beamSplitSourceList = splitBeamSource(); - int i = 0; - for (Source 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 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 +131,17 @@ public void addSplitsBack(List> splits, int subtaskId) { @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>> snapshotState(long checkpointId) throws Exception { + public FlinkSourceEnumeratorState snapshotState(long checkpointId) { LOG.info("Taking snapshot for checkpoint {}", checkpointId); - return snapshotState(); - } - - public Map>> snapshotState() throws Exception { - // For type compatibility reasons, we return a Map but we do not actually care about the key - Map>> state = new HashMap<>(1); - state.put(1, pendingSplits); - return state; + FlinkSourceSplitAssignmentMode mode = + splitsInitialized + ? FlinkSourceSplitAssignmentMode.LAZY + : FlinkSourceSplitAssignmentMode.UNDECIDED; + return new FlinkSourceEnumeratorState<>(mode, new ArrayList<>(pendingSplits)); } @Override @@ -167,39 +149,37 @@ public void close() throws IOException { // NoOp } - private long getDesiredSizeBytes(int numSplits, BoundedSource 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> splitBeamSource() throws Exception { + if (!(beamSource instanceof BoundedSource)) { + throw new IllegalStateException("Lazy assignment requires a bounded source."); } + LOG.info("Starting source {}", beamSource); + BoundedSource boundedSource = (BoundedSource) beamSource; + long estimatedSizeBytes = + FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource, pipelineOptions); + return FlinkSourceSplitUtils.splitBoundedSource( + boundedSource, pipelineOptions, numSplits, estimatedSizeBytes); } - // -------------- Private helper methods ---------------------- - private List> splitBeamSource() throws Exception { - if (beamSource instanceof BoundedSource) { - BoundedSource boundedSource = (BoundedSource) beamSource; - long desiredSizeBytes = getDesiredSizeBytes(numSplits, boundedSource); - List> splits = - ((BoundedSource) beamSource).split(desiredSizeBytes, pipelineOptions); - LOG.info("Split bounded source {} in {} splits", beamSource, splits.size()); - return splits; - } else if (beamSource instanceof UnboundedSource) { - List> splits = - ((UnboundedSource) 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> 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 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 000000000000..b74ee2d7f7d7 --- /dev/null +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/SizeBasedFlinkSourceSplitEnumerator.java @@ -0,0 +1,208 @@ +/* + * 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.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 + implements SplitEnumerator, FlinkSourceEnumeratorState> { + private static final Logger LOG = + LoggerFactory.getLogger(SizeBasedFlinkSourceSplitEnumerator.class); + + private final SplitEnumeratorContext> context; + private final BoundedSource boundedSource; + private final PipelineOptions pipelineOptions; + private final int numSplits; + private final Map> pendingSplitRequests; + private final List> returnedSplits; + + private @Nullable SplitEnumerator, FlinkSourceEnumeratorState> delegate; + + SizeBasedFlinkSourceSplitEnumerator( + SplitEnumeratorContext> context, + BoundedSource 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, FlinkSourceEnumeratorState> 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> 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 snapshotState(long checkpointId) throws Exception { + if (delegate == null) { + return new FlinkSourceEnumeratorState<>( + FlinkSourceSplitAssignmentMode.UNDECIDED, new ArrayList<>()); + } + return delegate.snapshotState(checkpointId); + } + + @Override + public void close() throws IOException { + if (delegate != null) { + delegate.close(); + } + } + + private FlinkSourceEnumeratorState selectAndSplit() throws Exception { + long estimatedSizeBytes = + FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource, pipelineOptions); + FlinkSourceSplitAssignmentMode selectedMode = selectAssignmentMode(estimatedSizeBytes); + ArrayList> 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 FlinkSourceSplitAssignmentMode selectAssignmentMode(long estimatedSizeBytes) { + long thresholdMb = + pipelineOptions + .as(FlinkPipelineOptions.class) + .getLazySourceSplitAssignmentMinSizeMbPerReader(); + 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 unknown. Using lazy split assignment.", + boundedSource); + return FlinkSourceSplitAssignmentMode.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); + FlinkSourceSplitAssignmentMode selectedMode = + estimatedBytesPerReader >= thresholdBytes + ? FlinkSourceSplitAssignmentMode.LAZY + : FlinkSourceSplitAssignmentMode.STATIC; + LOG.info( + "Using {} split assignment for bounded source {}: estimated size {} bytes, source " + + "parallelism {}, estimated bytes per reader {}, lazy assignment threshold {} bytes", + selectedMode, + boundedSource, + estimatedSizeBytes, + sourceParallelism, + estimatedBytesPerReader, + thresholdBytes); + return selectedMode; + } + + private SplitEnumerator, FlinkSourceEnumeratorState> createDelegate( + FlinkSourceEnumeratorState initialState) { + if (initialState.getAssignmentMode() == FlinkSourceSplitAssignmentMode.LAZY) { + return new LazyFlinkSourceSplitEnumerator<>( + context, boundedSource, pipelineOptions, numSplits, initialState); + } + if (initialState.getAssignmentMode() == FlinkSourceSplitAssignmentMode.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 { + private final List> splits; + private final int subtaskId; + + private ReturnedSplits(List> 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 e911ee72db5e..07d9c5a2f098 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,437 @@ 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.io.Serializable; +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 java.util.concurrent.atomic.AtomicLong; +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.sdk.coders.Coder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.io.FileBasedSource; +import org.apache.beam.sdk.io.FileBasedSource.FileBasedReader; +import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.io.Source; +import org.apache.beam.sdk.io.fs.MatchResult.Metadata; +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 AUTO_THRESHOLD_MB = 6144L; + private static final int SOURCE_PARALLELISM = 2; + private static final int REQUESTED_SPLITS = 4; + + @Test + public void testSmallBoundedSourceUsesStaticAssignmentAsynchronously() throws Exception { + FlinkPipelineOptions options = autoOptions(); + long thresholdBytes = AUTO_THRESHOLD_MB * MEBIBYTE; + long estimatedSizeBytes = SOURCE_PARALLELISM * thresholdBytes - 1L; + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(estimatedSizeBytes, REQUESTED_SPLITS); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = createBoundedSource(testSource, options, REQUESTED_SPLITS); + + try (SplitEnumerator, FlinkSourceEnumeratorState> enumerator = + flinkSource.createEnumerator(context)) { + assertTrue(enumerator instanceof SizeBasedFlinkSourceSplitEnumerator); + assertEquals(0, testSource.getEstimationCalls()); + enumerator.start(); + assertEquals( + "start must schedule estimation instead of blocking the coordinator thread", + 0, + testSource.getEstimationCalls()); + + context.getExecutorService().triggerAll(); + + FlinkSourceEnumeratorState state = enumerator.snapshotState(1L); + assertEquals(FlinkSourceSplitAssignmentMode.STATIC, state.getAssignmentMode()); + assertEquals(1, testSource.getEstimationCalls()); + assertEquals(estimatedSizeBytes / REQUESTED_SPLITS, testSource.getDesiredBundleSizeBytes()); + } + } + + @Test + public void testLargeBoundedSourceUsesLazyAssignment() throws Exception { + FlinkPipelineOptions options = autoOptions(); + long thresholdBytes = AUTO_THRESHOLD_MB * MEBIBYTE; + long estimatedSizeBytes = SOURCE_PARALLELISM * thresholdBytes; + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(estimatedSizeBytes, REQUESTED_SPLITS); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = createBoundedSource(testSource, options, REQUESTED_SPLITS); + + try (SplitEnumerator, FlinkSourceEnumeratorState> enumerator = + flinkSource.createEnumerator(context)) { + assertTrue(enumerator instanceof SizeBasedFlinkSourceSplitEnumerator); + enumerator.start(); + context.getExecutorService().triggerAll(); + + FlinkSourceEnumeratorState state = enumerator.snapshotState(1L); + assertEquals(FlinkSourceSplitAssignmentMode.LAZY, state.getAssignmentMode()); + assertEquals(1, testSource.getEstimationCalls()); + assertEquals(estimatedSizeBytes / REQUESTED_SPLITS, testSource.getDesiredBundleSizeBytes()); + } + } + + @Test + public void testSizeBasedSelectionReplaysLazyRequestAfterInitialization() throws Exception { + FlinkPipelineOptions options = autoOptions(); + long estimatedSizeBytes = SOURCE_PARALLELISM * AUTO_THRESHOLD_MB * MEBIBYTE; + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(estimatedSizeBytes, REQUESTED_SPLITS); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = createBoundedSource(testSource, options, REQUESTED_SPLITS); + + try (SplitEnumerator, FlinkSourceEnumeratorState> 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( + FlinkSourceSplitAssignmentMode.LAZY, enumerator.snapshotState(1L).getAssignmentMode()); + } + } + + @Test + public void testSizeBasedSelectionAssignsStaticSplitsToEarlyReaders() throws Exception { + FlinkPipelineOptions options = autoOptions(); + long estimatedSizeBytes = SOURCE_PARALLELISM * AUTO_THRESHOLD_MB * MEBIBYTE - 1L; + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(estimatedSizeBytes, REQUESTED_SPLITS); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = createBoundedSource(testSource, options, REQUESTED_SPLITS); + + try (SplitEnumerator, FlinkSourceEnumeratorState> 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( + FlinkSourceSplitAssignmentMode.STATIC, enumerator.snapshotState(1L).getAssignmentMode()); + } + } + + @Test + public void testEstimationFailureFailsSplitInitialization() throws Exception { + FlinkPipelineOptions options = autoOptions(); + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.createFailing(REQUESTED_SPLITS); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = createBoundedSource(testSource, options, REQUESTED_SPLITS); + + try (SplitEnumerator, FlinkSourceEnumeratorState> enumerator = + flinkSource.createEnumerator(context)) { + enumerator.start(); + assertThrows(RuntimeException.class, () -> context.getExecutorService().triggerAll()); + assertEquals(1, testSource.getEstimationCalls()); + } + } + + @Test + public void testUnknownEstimatesUseLazyAssignmentAndPreserveBundleSize() throws Exception { + long[] unknownEstimates = {-1L, Long.MAX_VALUE}; + for (long unknownEstimate : unknownEstimates) { + FlinkPipelineOptions options = autoOptions(); + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(unknownEstimate, REQUESTED_SPLITS); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = + createBoundedSource(testSource, options, REQUESTED_SPLITS); + + try (SplitEnumerator, FlinkSourceEnumeratorState> + enumerator = flinkSource.createEnumerator(context)) { + enumerator.start(); + context.getExecutorService().triggerAll(); + + FlinkSourceEnumeratorState state = enumerator.snapshotState(1L); + assertEquals(FlinkSourceSplitAssignmentMode.LAZY, state.getAssignmentMode()); + assertEquals(1, testSource.getEstimationCalls()); + assertEquals(unknownEstimate / REQUESTED_SPLITS, testSource.getDesiredBundleSizeBytes()); + } + } + } + + @Test + public void testConfigurationCanForceLazyOrStaticAssignment() throws Exception { + long largeEstimate = 1024L * MEBIBYTE; + assertAssignmentMode( + largeEstimate, -1L, FlinkSourceSplitAssignmentMode.STATIC, REQUESTED_SPLITS); + assertAssignmentMode( + largeEstimate, -100L, FlinkSourceSplitAssignmentMode.STATIC, REQUESTED_SPLITS); + assertAssignmentMode(0L, 0L, FlinkSourceSplitAssignmentMode.LAZY, REQUESTED_SPLITS); + assertAssignmentMode( + largeEstimate, Long.MAX_VALUE, FlinkSourceSplitAssignmentMode.STATIC, REQUESTED_SPLITS); + } + + @Test + public void testDefaultConfigurationKeepsLazyAssignmentWithoutThresholdComparison() + throws Exception { + FlinkPipelineOptions options = FlinkPipelineOptions.defaults(); + assertEquals(0L, (long) options.getLazySourceSplitAssignmentMinSizeMbPerReader()); + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(MEBIBYTE, REQUESTED_SPLITS); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = createBoundedSource(testSource, options, REQUESTED_SPLITS); + + try (SplitEnumerator, FlinkSourceEnumeratorState> enumerator = + flinkSource.createEnumerator(context)) { + assertTrue(enumerator instanceof LazyFlinkSourceSplitEnumerator); + enumerator.start(); + context.getExecutorService().triggerAll(); + assertEquals( + FlinkSourceSplitAssignmentMode.LAZY, enumerator.snapshotState(1L).getAssignmentMode()); + } + } + + @Test + public void testEmptyBoundedSourceUsesStaticAssignment() throws Exception { + // An estimated size of exactly 0 is a valid answer (e.g. an empty file glob), not an + // unknown estimate, and sits below any positive threshold. + assertAssignmentMode( + 0L, AUTO_THRESHOLD_MB, FlinkSourceSplitAssignmentMode.STATIC, REQUESTED_SPLITS); + } + + @Test + public void testRestoreKeepsLazyAssignmentAcrossRescaleWithoutEstimating() throws Exception { + final int initialParallelism = 4; + final int restoredParallelism = 1; + final int generatedSplits = 4; + FlinkPipelineOptions options = autoOptions(); + long thresholdBytes = AUTO_THRESHOLD_MB * MEBIBYTE; + long estimatedSizeBytes = initialParallelism * thresholdBytes; + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(estimatedSizeBytes, generatedSplits); + FlinkSource flinkSource = createBoundedSource(testSource, options, generatedSplits); + FlinkSourceEnumeratorState checkpoint; + + TestingSplitEnumeratorContext> initialContext = + new TestingSplitEnumeratorContext<>(initialParallelism); + try (SplitEnumerator, FlinkSourceEnumeratorState> enumerator = + flinkSource.createEnumerator(initialContext)) { + enumerator.start(); + initialContext.getExecutorService().triggerAll(); + checkpoint = roundTripState(flinkSource, enumerator.snapshotState(1L)); + assertEquals(FlinkSourceSplitAssignmentMode.LAZY, checkpoint.getAssignmentMode()); + assertEquals(generatedSplits, checkpoint.getPendingSplits().size()); + } + + TestingSplitEnumeratorContext> restoredContext = + new TestingSplitEnumeratorContext<>(restoredParallelism); + try (SplitEnumerator, FlinkSourceEnumeratorState> 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( + FlinkSourceSplitAssignmentMode.LAZY, restored.snapshotState(2L).getAssignmentMode()); + } + } + + @Test + public void testLegacyLazyCheckpointMapUpgradesToStrategyNeutralState() throws Exception { + FlinkPipelineOptions options = FlinkPipelineOptions.defaults(); + TestEstimatedSizeBoundedSource testSource = TestEstimatedSizeBoundedSource.create(1L, 1); + FlinkSource flinkSource = createBoundedSource(testSource, options, 1); + FlinkSourceSplit pendingSplit = new FlinkSourceSplit<>(0, testSource); + Map>> legacyCheckpoint = + Collections.singletonMap(1, Collections.singletonList(pendingSplit)); + byte[] serialized = SerdeUtils.serializeObject(legacyCheckpoint); + + FlinkSourceEnumeratorState upgraded = + flinkSource.getEnumeratorCheckpointSerializer().deserialize(0, serialized); + + assertEquals(FlinkSourceSplitAssignmentMode.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>> 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 upgraded = + new FlinkSourceEnumeratorStateSerializer(FlinkSourceSplitAssignmentMode.STATIC) + .deserialize(0, serialized); + assertEquals(FlinkSourceSplitAssignmentMode.STATIC, upgraded.getAssignmentMode()); + assertEquals(3, upgraded.getPendingSplits().size()); + + TestingSplitEnumeratorContext> restoredContext = + new TestingSplitEnumeratorContext<>(parallelism); + try (FlinkSourceSplitEnumerator 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 testStaticRestoreReturnsPendingSplitsToOriginalOwnersAtSameParallelism() + throws Exception { + final int parallelism = 4; + TestEstimatedSizeBoundedSource testSource = TestEstimatedSizeBoundedSource.create(1L, 1); + ArrayList> pendingSplits = new ArrayList<>(); + pendingSplits.add(new FlinkSourceSplit<>(3, testSource)); + FlinkSourceEnumeratorState checkpoint = + new FlinkSourceEnumeratorState<>(FlinkSourceSplitAssignmentMode.STATIC, pendingSplits); + + TestingSplitEnumeratorContext> restoredContext = + new TestingSplitEnumeratorContext<>(parallelism); + try (FlinkSourceSplitEnumerator restored = + new FlinkSourceSplitEnumerator<>( + restoredContext, testSource, staticOptions(), parallelism, checkpoint)) { + restored.start(); + for (int subtaskId = 0; subtaskId < parallelism; subtaskId++) { + restoredContext.registerReader(subtaskId, "reader-" + subtaskId); + restored.addReader(subtaskId); + } + + for (int subtaskId = 0; subtaskId < parallelism - 1; subtaskId++) { + assertEquals( + 0, restoredContext.getSplitAssignments().get(subtaskId).getAssignedSplits().size()); + } + assertEquals( + Collections.singletonList(3), assignedSplitIndexesForSubtask(restoredContext, 3)); + restoredContext + .getSplitAssignments() + .values() + .forEach(state -> assertTrue(state.hasReceivedNoMoreSplitsSignal())); + } + } + + @Test + public void testSerializerRejectsUnknownVersionsAndUnexpectedPayloads() throws Exception { + FlinkSourceEnumeratorStateSerializer serializer = + new FlinkSourceEnumeratorStateSerializer<>(FlinkSourceSplitAssignmentMode.LAZY); + byte[] legacyMapBytes = + SerdeUtils.serializeObject(Collections.singletonMap(1, new ArrayList<>())); + byte[] stateBytes = + serializer.serialize( + new FlinkSourceEnumeratorState<>( + FlinkSourceSplitAssignmentMode.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 = autoOptions(); + long thresholdBytes = AUTO_THRESHOLD_MB * MEBIBYTE; + long estimatedSizeBytes = initialParallelism * thresholdBytes - 1L; + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(estimatedSizeBytes, generatedSplits); + FlinkSource flinkSource = createBoundedSource(testSource, options, generatedSplits); + FlinkSourceEnumeratorState checkpoint; + + TestingSplitEnumeratorContext> initialContext = + new TestingSplitEnumeratorContext<>(initialParallelism); + try (SplitEnumerator, FlinkSourceEnumeratorState> enumerator = + flinkSource.createEnumerator(initialContext)) { + enumerator.start(); + initialContext.getExecutorService().triggerAll(); + checkpoint = roundTripState(flinkSource, enumerator.snapshotState(1L)); + assertEquals(FlinkSourceSplitAssignmentMode.STATIC, checkpoint.getAssignmentMode()); + assertEquals(generatedSplits, checkpoint.getPendingSplits().size()); + } + + TestingSplitEnumeratorContext> restoredContext = + new TestingSplitEnumeratorContext<>(restoredParallelism); + try (SplitEnumerator, FlinkSourceEnumeratorState> 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 +468,8 @@ public void testAssignSplitsWithBoundedSource() throws IOException { .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,13 +481,72 @@ public void testAssignSplitsWithBoundedSource() throws IOException { assertEquals( expectedSplitSize, source.getEstimatedSizeBytes(FlinkPipelineOptions.defaults())); - } catch (Exception e) { - fail("Received exception" + e); + } catch (Exception error) { + fail("Received exception " + error); } }); }); } + @Test + public void testSignalsNoMoreSplitsToEarlyReaderWithoutAssignment() throws IOException { + final int numSubtasks = 2; + final int numSplits = 1; + TestingSplitEnumeratorContext>> testContext = + new TestingSplitEnumeratorContext<>(numSubtasks); + TestBoundedCountingSource testSource = new TestBoundedCountingSource(numSplits, numSplits); + + try (FlinkSourceSplitEnumerator> enumerator = + new FlinkSourceSplitEnumerator<>(testContext, testSource, staticOptions(), numSplits)) { + enumerator.start(); + for (int subtaskId = 0; subtaskId < numSubtasks; subtaskId++) { + testContext.registerReader(subtaskId, String.valueOf(subtaskId)); + enumerator.addReader(subtaskId); + } + testContext.getExecutorService().triggerAll(); + + assertEquals(numSubtasks, testContext.getSplitAssignments().size()); + assertEquals(1, testContext.getSplitAssignments().get(0).getAssignedSplits().size()); + assertEquals(0, testContext.getSplitAssignments().get(1).getAssignedSplits().size()); + assertTrue(testContext.getSplitAssignments().get(0).hasReceivedNoMoreSplitsSignal()); + assertTrue(testContext.getSplitAssignments().get(1).hasReceivedNoMoreSplitsSignal()); + } + } + + @Test + public void testStaticAssignmentRespectsFileInputSplitMaxSize() throws IOException { + final int numSubtasks = 2; + final int requestedSplits = 2; + final long fileSize = 100L * MEBIBYTE; + final long maxSplitSizeMb = 10L; + final int expectedSplits = 10; + FlinkPipelineOptions options = staticOptions(); + options.setFileInputSplitMaxSizeMB(maxSplitSizeMb); + TestingSplitEnumeratorContext> testContext = + new TestingSplitEnumeratorContext<>(numSubtasks); + TestFileBasedSource testSource = TestFileBasedSource.create(fileSize); + + try (FlinkSourceSplitEnumerator enumerator = + new FlinkSourceSplitEnumerator<>(testContext, testSource, options, requestedSplits)) { + enumerator.start(); + for (int subtaskId = 0; subtaskId < numSubtasks; subtaskId++) { + testContext.registerReader(subtaskId, String.valueOf(subtaskId)); + enumerator.addReader(subtaskId); + } + testContext.getExecutorService().triggerAll(); + + assertEquals(expectedSplits, countAssignedSplits(testContext)); + testContext + .getSplitAssignments() + .values() + .forEach( + state -> { + assertEquals(expectedSplits / numSubtasks, state.getAssignedSplits().size()); + assertTrue(state.hasReceivedNoMoreSplitsSignal()); + }); + } + } + @Test public void testAssignSplitsWithUnboundedSource() throws IOException { final int numSplits = 10; @@ -88,90 +556,79 @@ public void testAssignSplitsWithUnboundedSource() throws IOException { new TestingSplitEnumeratorContext<>(numSubtasks); TestCountingSource testSource = new TestCountingSource(numRecordsPerSplit); - assignSplits(testContext, testSource, numSplits); + try (FlinkSourceSplitEnumerator> 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>> testContext = new TestingSplitEnumeratorContext<>(numSubtasks); - TestBoundedCountingSource testSource = - new TestBoundedCountingSource(numSplits, totalNumRecords); - try (FlinkSourceSplitEnumerator> splitEnumerator = - new FlinkSourceSplitEnumerator<>( - testContext, testSource, FlinkPipelineOptions.defaults(), numSplits)) { - splitEnumerator.start(); + TestBoundedCountingSource testSource = new TestBoundedCountingSource(numSplits, numSplits); + + try (FlinkSourceSplitEnumerator> enumerator = + new FlinkSourceSplitEnumerator<>(testContext, testSource, staticOptions(), numSplits)) { + enumerator.start(); testContext.registerReader(0, "0"); - splitEnumerator.addReader(0); + enumerator.addReader(0); testContext.getExecutorService().triggerAll(); - List>> splitsForReader = - testContext.getSplitAssignments().get(0).getAssignedSplits(); - assertEquals(numSplits / numSubtasks, splitsForReader.size()); + List>> 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>> testContext = - new TestingSplitEnumeratorContext<>(numSubtasks); - TestBoundedCountingSource testSource = - new TestBoundedCountingSource(numSplits, totalNumRecords); - final Map>>> assignment; - try (FlinkSourceSplitEnumerator> 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, + FlinkSourceSplitAssignmentMode expectedMode, + int generatedSplits) + throws Exception { + FlinkPipelineOptions options = FlinkPipelineOptions.defaults(); + options.setLazySourceSplitAssignmentMinSizeMbPerReader(configuredThresholdMb); + TestEstimatedSizeBoundedSource testSource = + TestEstimatedSizeBoundedSource.create(estimatedSizeBytes, generatedSplits); + TestingSplitEnumeratorContext> context = + new TestingSplitEnumeratorContext<>(SOURCE_PARALLELISM); + FlinkSource flinkSource = createBoundedSource(testSource, options, generatedSplits); - // add tasks back - testContext = new TestingSplitEnumeratorContext<>(numSubtasks); - try (FlinkSourceSplitEnumerator> 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>> splitsForReader = - testContext.getSplitAssignments().get(0).getAssignedSplits(); - assertEquals(numSplits / numSubtasks, splitsForReader.size()); + try (SplitEnumerator, FlinkSourceEnumeratorState> 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 +637,163 @@ private void assignSplits( Source> source, int numSplits) throws IOException { - try (FlinkSourceSplitEnumerator> 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> 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.setLazySourceSplitAssignmentMinSizeMbPerReader(-1L); + return options; + } + + private static FlinkPipelineOptions autoOptions() { + FlinkPipelineOptions options = FlinkPipelineOptions.defaults(); + options.setLazySourceSplitAssignmentMinSizeMbPerReader(AUTO_THRESHOLD_MB); + return options; + } + + private static FlinkSource createBoundedSource( + BoundedSource source, FlinkPipelineOptions options, int numSplits) { + return FlinkSource.bounded( + "test-bounded-source", source, new SerializablePipelineOptions(options), numSplits); + } + + private static FlinkSourceEnumeratorState roundTripState( + FlinkSource source, FlinkSourceEnumeratorState state) throws IOException { + SimpleVersionedSerializer> serializer = + source.getEnumeratorCheckpointSerializer(); + byte[] serialized = serializer.serialize(state); + return serializer.deserialize(serializer.getVersion(), serialized); + } + + private static int countAssignedSplits( + TestingSplitEnumeratorContext> context) { + return context.getSplitAssignments().values().stream() + .mapToInt(state -> state.getAssignedSplits().size()) + .sum(); + } + + private static List assignedSplitIndexesForSubtask( + TestingSplitEnumeratorContext> context, int subtaskId) { + List splitIndexes = new ArrayList<>(); + for (FlinkSourceSplit split : + context.getSplitAssignments().get(subtaskId).getAssignedSplits()) { + splitIndexes.add(split.splitIndex()); + } + Collections.sort(splitIndexes); + return splitIndexes; + } + + private static final class TestEstimatedSizeBoundedSource extends BoundedSource { + private final long estimatedSizeBytes; + private final int generatedSplits; + private final boolean failEstimation; + private final EstimationTracker tracker; + + private TestEstimatedSizeBoundedSource( + long estimatedSizeBytes, + int generatedSplits, + boolean failEstimation, + EstimationTracker tracker) { + this.estimatedSizeBytes = estimatedSizeBytes; + this.generatedSplits = generatedSplits; + this.failEstimation = failEstimation; + this.tracker = tracker; + } + + private static TestEstimatedSizeBoundedSource create( + long estimatedSizeBytes, int generatedSplits) { + return new TestEstimatedSizeBoundedSource( + estimatedSizeBytes, generatedSplits, false, new EstimationTracker()); + } + + private static TestEstimatedSizeBoundedSource createFailing(int generatedSplits) { + return new TestEstimatedSizeBoundedSource(0L, generatedSplits, true, new EstimationTracker()); + } + + @Override + public List> split( + long desiredBundleSizeBytes, PipelineOptions options) { + tracker.desiredBundleSizeBytes.set(desiredBundleSizeBytes); + List splits = new ArrayList<>(generatedSplits); + for (int i = 0; i < generatedSplits; i++) { + splits.add(new TestEstimatedSizeBoundedSource(1L, 1, false, tracker)); + } + return splits; + } + + @Override + public long getEstimatedSizeBytes(PipelineOptions options) throws IOException { + tracker.estimationCalls.incrementAndGet(); + if (failEstimation) { + throw new IOException("Expected test estimation failure"); + } + return estimatedSizeBytes; + } + + @Override + public BoundedReader createReader(PipelineOptions options) { + throw new UnsupportedOperationException("This source is only used to test split assignment"); + } + + @Override + public Coder getOutputCoder() { + return StringUtf8Coder.of(); + } + + private int getEstimationCalls() { + return tracker.estimationCalls.get(); + } + + private long getDesiredBundleSizeBytes() { + return tracker.desiredBundleSizeBytes.get(); + } + } + + private static final class EstimationTracker implements Serializable { + private static final long serialVersionUID = 1L; + + private final AtomicInteger estimationCalls = new AtomicInteger(); + private final AtomicLong desiredBundleSizeBytes = new AtomicLong(-1L); + } + + private static final class TestFileBasedSource extends FileBasedSource { + private TestFileBasedSource(Metadata metadata, long startOffset, long endOffset) { + super(metadata, 1L, startOffset, endOffset); + } + + private static TestFileBasedSource create(long sizeBytes) { + Metadata metadata = + Metadata.builder() + .setResourceId(FileSystems.matchNewResource("static-split-size-test", false)) + .setSizeBytes(sizeBytes) + .setIsReadSeekEfficient(true) + .build(); + return new TestFileBasedSource(metadata, 0L, sizeBytes); + } + + @Override + public Coder getOutputCoder() { + return StringUtf8Coder.of(); + } + + @Override + protected FileBasedSource createForSubrangeOfFile( + Metadata metadata, long start, long end) { + return new TestFileBasedSource(metadata, start, end); + } + + @Override + protected FileBasedReader createSingleFileReader(PipelineOptions options) { + throw new UnsupportedOperationException("This source is only used to test split sizing"); } } } 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 34d6c5243776..eeca59a65576 100644 --- a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html @@ -122,6 +122,11 @@ Interval in milliseconds for sending latency tracking marks from the sources to the sinks. Interval value <= 0 disables the feature. Default: 0 + + lazySourceSplitAssignmentMinSizeMbPerReader + 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: 0 + maxBundleSize The maximum number of elements in a bundle. Default values are 1000 for a streaming job and 1,000,000 for batch 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 e3fe24216a54..5cb9fe7a049b 100644 --- a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html @@ -122,6 +122,11 @@ Interval in milliseconds for sending latency tracking marks from the sources to the sinks. Interval value <= 0 disables the feature. Default: 0 + + lazy_source_split_assignment_min_size_mb_per_reader + 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: 0 + max_bundle_size The maximum number of elements in a bundle. Default values are 1000 for a streaming job and 1,000,000 for batch