From a19889c65ed98d68d846786af83a9a5dca1ecffe Mon Sep 17 00:00:00 2001 From: jto Date: Fri, 2 Oct 2026 16:21:50 +0200 Subject: [PATCH 01/11] [Flink] Re-split bounded sources for even work distribution Share bounded source splitting between the DataSet and DataStream translations in BoundedSourceSplitter. After the initial split, the splitter re-splits when there are fewer splits than readers, re-splits oversized splits when split sizes vary a lot, rounds the split count up to a multiple of the parallelism, and orders splits so round-robin assignment balances readers. The file input max split size is preserved. Batch pipelines now use static split assignment by default, in both DataSet (new StaticInputSplitAssigner) and DataStream. Lazy assignment is used when explicitly requested via sourceStaticSplitThresholdMb, and stays the default for streaming pipelines. Co-Authored-By: Claude Opus 5.5 --- .../runners/flink/FlinkPipelineOptions.java | 19 +- .../runners/flink/FlinkPipelineOptions.java | 19 +- .../wrappers/BoundedSourceSplitter.java | 324 ++++++++++++++++++ .../wrappers/SourceInputFormat.java | 88 +++-- .../streaming/io/source/FlinkSource.java | 17 +- .../io/source/FlinkSourceSplitUtils.java | 38 +- .../LazyFlinkSourceSplitEnumerator.java | 3 +- .../SizeBasedFlinkSourceSplitEnumerator.java | 11 +- .../wrappers/BoundedSourceSplitterTest.java | 213 ++++++++++++ .../wrappers/SourceInputFormatTest.java | 40 +++ .../FlinkSourceSplitEnumeratorTest.java | 5 +- .../flink_java_pipeline_options.html | 4 +- .../flink_python_pipeline_options.html | 4 +- 13 files changed, 689 insertions(+), 96 deletions(-) create mode 100644 runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java create mode 100644 runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java 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 d3c5fefd9cf7..e2ecf2c879c2 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 @@ -26,6 +26,7 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.StreamingOptions; import org.apache.flink.runtime.state.StateBackendFactory; +import org.checkerframework.checker.nullness.qual.Nullable; /** * Options which can be used to configure the Flink Runner. @@ -366,15 +367,15 @@ public Long create(PipelineOptions options) { 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); + "Split assignment for bounded sources. When unset, batch pipelines assign splits statically " + + "(round-robin, after re-splitting the source so that splits are numerous and evenly " + + "sized) and streaming pipelines assign them lazily. 0 always uses lazy assignment. A " + + "positive value is a threshold in MiB per source reader: static assignment is used for " + + "sources with a known, positive estimate below the threshold and lazy assignment " + + "otherwise. Any negative value always uses static assignment.") + @Nullable Long getSourceStaticSplitThresholdMb(); + + void setSourceStaticSplitThresholdMb(@Nullable Long thresholdMb); @Description( "Allow drain operation for flink pipelines that contain RequiresStableInput operator. Note that at time of draining," 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 a6729356bdcb..3b1df33c258b 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 @@ -25,6 +25,7 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.StreamingOptions; +import org.checkerframework.checker.nullness.qual.Nullable; /** * Options which can be used to configure the Flink Runner. @@ -383,15 +384,15 @@ public Long create(PipelineOptions options) { 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); + "Split assignment for bounded sources. When unset, batch pipelines assign splits statically " + + "(round-robin, after re-splitting the source so that splits are numerous and evenly " + + "sized) and streaming pipelines assign them lazily. 0 always uses lazy assignment. A " + + "positive value is a threshold in MiB per source reader: static assignment is used for " + + "sources with a known, positive estimate below the threshold and lazy assignment " + + "otherwise. Any negative value always uses static assignment.") + @Nullable Long getSourceStaticSplitThresholdMb(); + + void setSourceStaticSplitThresholdMb(@Nullable Long thresholdMb); @Description( "Allow drain operation for flink pipelines that contain RequiresStableInput operator. Note that at time of draining," diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java new file mode 100644 index 000000000000..1a6ac3f46068 --- /dev/null +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -0,0 +1,324 @@ +/* + * 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; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +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.options.PipelineOptions; +import org.apache.beam.sdk.options.StreamingOptions; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Splits a {@link BoundedSource} into splits that can be evenly distributed over the source + * readers. Shared by the DataSet and DataStream translations. + * + *

Beam does not require a source to be fully split in one pass, so the initial splits are + * refined in several bounded passes: + * + *

    + *
  1. Split the source using {@code estimatedSize / parallelism} as the desired split size, + * capped by {@link FlinkPipelineOptions#getFileInputSplitMaxSizeMB()} for file sources. + *
  2. While there are fewer splits than readers, re-split each split. + *
  3. While the split sizes vary a lot (high coefficient of variation), re-split the larger + * splits. + *
  4. Halve the largest splits until the number of splits is a multiple of the parallelism. + *
  5. Order the splits so that round-robin assignment gives each reader a similar amount of data. + *
+ * + *

Steps 3 to 5 need per-split size estimates and are skipped when they are not available. + */ +public final class BoundedSourceSplitter { + private static final Logger LOG = LoggerFactory.getLogger(BoundedSourceSplitter.class); + + static final long MEBIBYTE = 1024L * 1024L; + + /** Maximum number of re-split passes per refinement step. */ + static final int MAX_RESPLIT_ROUNDS = 3; + + /** Coefficient of variation of split sizes above which larger splits are re-split. */ + static final double MAX_SIZE_COEFFICIENT_OF_VARIATION = 0.5; + + /** Splits larger than this factor times the target size are re-split during balancing. */ + static final double OVERSIZED_SPLIT_FACTOR = 1.5; + + /** Lower bound of the balancing target size, relative to {@code estimatedSize/parallelism}. */ + static final int MAX_SPLITS_PER_READER_FOR_BALANCING = 4; + + /** How bounded source splits are handed out to readers. */ + public enum Assignment { + /** Splits are assigned to readers up front, round-robin. */ + STATIC, + /** Readers pull one split at a time. */ + LAZY, + /** Static or lazy, depending on the estimated size per reader. */ + SIZE_BASED + } + + private BoundedSourceSplitter() {} + + /** + * Resolves the split assignment from {@link + * FlinkPipelineOptions#getSourceStaticSplitThresholdMb}. When unset, batch pipelines use static + * assignment and streaming pipelines lazy assignment. + */ + public static Assignment assignment(PipelineOptions options) { + @Nullable Long thresholdMb = + options.as(FlinkPipelineOptions.class).getSourceStaticSplitThresholdMb(); + if (thresholdMb == null) { + return options.as(StreamingOptions.class).isStreaming() ? Assignment.LAZY : Assignment.STATIC; + } + if (thresholdMb < 0) { + return Assignment.STATIC; + } + return thresholdMb == 0 ? Assignment.LAZY : Assignment.SIZE_BASED; + } + + /** Returns the configured static split threshold in bytes, or 0 when unset. */ + public static long staticSplitThresholdBytes(PipelineOptions options) { + @Nullable Long thresholdMb = + options.as(FlinkPipelineOptions.class).getSourceStaticSplitThresholdMb(); + return thresholdMb == null ? 0L : mebibytesToBytes(thresholdMb); + } + + public static long mebibytesToBytes(long mebibytes) { + return mebibytes > Long.MAX_VALUE / MEBIBYTE ? Long.MAX_VALUE : mebibytes * MEBIBYTE; + } + + /** Splits {@code source} for {@code parallelism} readers. */ + public static List> split( + BoundedSource source, PipelineOptions options, int parallelism, long estimatedSizeBytes) + throws Exception { + int readers = Math.max(1, parallelism); + long maxSplitSizeBytes = maxSplitSizeBytes(source, options); + long desiredSizeBytes = Math.min(Math.max(1L, estimatedSizeBytes / readers), maxSplitSizeBytes); + + List> splits = sized(source.split(desiredSizeBytes, options), options); + int initialSplits = splits.size(); + + splits = ensureEnoughSplits(splits, options, readers, desiredSizeBytes); + if (allSizesKnown(splits)) { + splits = balanceSizes(splits, options, readers, estimatedSizeBytes, maxSplitSizeBytes); + splits = roundToMultipleOfParallelism(splits, options, readers); + splits = orderForRoundRobin(splits, readers); + } + + LOG.info( + "Split bounded source {} in {} splits (initially {}; estimated size {} bytes, desired " + + "split size {} bytes, parallelism {}, split size coefficient of variation {})", + source, + splits.size(), + initialSplits, + estimatedSizeBytes, + desiredSizeBytes, + readers, + allSizesKnown(splits) ? coefficientOfVariation(splits) : "unknown"); + + List> result = new ArrayList<>(splits.size()); + for (SizedSource split : splits) { + result.add(split.source); + } + return result; + } + + private static long maxSplitSizeBytes(BoundedSource source, PipelineOptions options) { + @Nullable Long maxSplitSizeMb = + options.as(FlinkPipelineOptions.class).getFileInputSplitMaxSizeMB(); + if (source instanceof FileBasedSource && maxSplitSizeMb != null && maxSplitSizeMb > 0) { + return mebibytesToBytes(maxSplitSizeMb); + } + return Long.MAX_VALUE; + } + + /** Re-splits every split while there are fewer splits than readers. */ + private static List> ensureEnoughSplits( + List> splits, PipelineOptions options, int readers, long desiredSizeBytes) + throws Exception { + for (int round = 0; round < MAX_RESPLIT_ROUNDS && splits.size() < readers; round++) { + int factor = (readers + splits.size() - 1) / Math.max(1, splits.size()); + List> next = new ArrayList<>(); + for (SizedSource split : splits) { + long base = split.sizeBytes > 0 ? split.sizeBytes : desiredSizeBytes; + next.addAll(resplit(split, options, Math.max(1L, base / factor))); + } + if (next.size() <= splits.size()) { + break; + } + splits = next; + } + return splits; + } + + /** Re-splits the larger splits while the split sizes vary a lot. */ + private static List> balanceSizes( + List> splits, + PipelineOptions options, + int readers, + long estimatedSizeBytes, + long maxSplitSizeBytes) + throws Exception { + long minTargetBytes = + Math.max(1L, estimatedSizeBytes / ((long) readers * MAX_SPLITS_PER_READER_FOR_BALANCING)); + for (int round = 0; round < MAX_RESPLIT_ROUNDS; round++) { + if (coefficientOfVariation(splits) <= MAX_SIZE_COEFFICIENT_OF_VARIATION) { + break; + } + long targetBytes = Math.min(Math.max((long) mean(splits), minTargetBytes), maxSplitSizeBytes); + List> next = new ArrayList<>(); + boolean progress = false; + for (SizedSource split : splits) { + if (split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes) { + List> parts = resplit(split, options, targetBytes); + progress |= parts.size() > 1; + next.addAll(parts); + } else { + next.add(split); + } + } + if (!progress || !allSizesKnown(next)) { + break; + } + splits = next; + } + return splits; + } + + /** Halves the largest splits until the number of splits is a multiple of the parallelism. */ + private static List> roundToMultipleOfParallelism( + List> splits, PipelineOptions options, int readers) throws Exception { + List> result = new ArrayList<>(splits); + result.sort(Comparator.comparingLong((SizedSource s) -> s.sizeBytes).reversed()); + // Splits that could not be split further are moved here so they are not retried. + List> unsplittable = new ArrayList<>(); + for (int attempt = 0; + attempt < 2 * readers && (result.size() + unsplittable.size()) % readers != 0; + attempt++) { + if (result.isEmpty()) { + break; + } + SizedSource largest = result.remove(0); + List> parts = + largest.sizeBytes > 1 + ? resplit(largest, options, Math.max(1L, largest.sizeBytes / 2)) + : Collections.singletonList(largest); + if (parts.size() <= 1 || !allSizesKnown(parts)) { + unsplittable.add(largest); + continue; + } + result.addAll(parts); + result.sort(Comparator.comparingLong((SizedSource s) -> s.sizeBytes).reversed()); + } + result.addAll(unsplittable); + return result; + } + + /** + * Sorts splits by decreasing size in a "snake" order (0..p-1, p-1..0, ...) so that assigning + * split {@code i} to reader {@code i % p} spreads the large splits across readers. + */ + private static List> orderForRoundRobin( + List> splits, int readers) { + List> sorted = new ArrayList<>(splits); + sorted.sort(Comparator.comparingLong((SizedSource s) -> s.sizeBytes).reversed()); + List> ordered = new ArrayList<>(sorted.size()); + for (int start = 0; start < sorted.size(); start += readers) { + int end = Math.min(start + readers, sorted.size()); + List> row = new ArrayList<>(sorted.subList(start, end)); + if ((start / readers) % 2 == 1) { + Collections.reverse(row); + } + ordered.addAll(row); + } + return ordered; + } + + private static List> resplit( + SizedSource split, PipelineOptions options, long desiredSizeBytes) { + try { + List> parts = split.source.split(desiredSizeBytes, options); + if (parts.isEmpty()) { + return Collections.singletonList(split); + } + return sized(parts, options); + } catch (Exception e) { + LOG.debug("Could not re-split source {}. Keeping it as is.", split.source, e); + return Collections.singletonList(split); + } + } + + private static List> sized( + List> sources, PipelineOptions options) { + List> result = new ArrayList<>(sources.size()); + for (BoundedSource source : sources) { + long size; + try { + size = source.getEstimatedSizeBytes(options); + } catch (Exception e) { + size = -1L; + } + result.add(new SizedSource<>(source, size)); + } + return result; + } + + private static boolean allSizesKnown(List> splits) { + for (SizedSource split : splits) { + if (split.sizeBytes < 0 || split.sizeBytes == Long.MAX_VALUE) { + return false; + } + } + return !splits.isEmpty(); + } + + private static double mean(List> splits) { + double sum = 0; + for (SizedSource split : splits) { + sum += split.sizeBytes; + } + return sum / splits.size(); + } + + static double coefficientOfVariation(List> splits) { + double mean = mean(splits); + if (mean <= 0) { + return 0; + } + double variance = 0; + for (SizedSource split : splits) { + double delta = split.sizeBytes - mean; + variance += delta * delta; + } + return Math.sqrt(variance / splits.size()) / mean; + } + + private static final class SizedSource { + private final BoundedSource source; + private final long sizeBytes; + + private SizedSource(BoundedSource source, long sizeBytes) { + this.source = source; + this.sizeBytes = sizeBytes; + } + } +} diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormat.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormat.java index 8a8d4cf3e9a5..3e6f5f5619f1 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormat.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormat.java @@ -18,13 +18,15 @@ package org.apache.beam.runners.flink.translation.wrappers; import java.io.IOException; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.HashMap; import java.util.List; +import java.util.Map; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; -import org.apache.beam.runners.flink.FlinkPipelineOptions; import org.apache.beam.runners.flink.metrics.FlinkMetricContainer; import org.apache.beam.runners.flink.metrics.ReaderInvocationUtil; 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.options.PipelineOptions; import org.apache.beam.sdk.transforms.windowing.GlobalWindow; @@ -36,6 +38,7 @@ import org.apache.flink.api.common.io.RichInputFormat; import org.apache.flink.api.common.io.statistics.BaseStatistics; import org.apache.flink.configuration.Configuration; +import org.apache.flink.core.io.InputSplit; import org.apache.flink.core.io.InputSplitAssigner; import org.joda.time.Instant; import org.slf4j.Logger; @@ -58,6 +61,10 @@ public class SourceInputFormat extends RichInputFormat, Sour private transient BoundedSource.BoundedReader reader; private boolean inputAvailable = false; + // Set by createInputSplits on the JobManager and read by getInputSplitAssigner. + private transient int numTasks = 1; + private transient boolean lazyAssignment; + private transient ReaderInvocationUtil> readerInvoker; private transient FlinkMetricContainer metricContainer; @@ -111,29 +118,18 @@ public float getAverageRecordWidth() { return null; } - private long getDesiredSizeBytes(int numSplits) throws Exception { - long totalSize = initialSource.getEstimatedSizeBytes(options); - long defaultSplitSize = totalSize / numSplits; - long maxSplitSize = 0; - if (options != null) { - maxSplitSize = options.as(FlinkPipelineOptions.class).getFileInputSplitMaxSizeMB(); - } - if (initialSource 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; - } - } - @Override @SuppressWarnings("unchecked") public SourceInputSplit[] createInputSplits(int numSplits) throws IOException { + if (options == null) { + options = serializedOptions.get(); + } try { - long desiredSizeBytes = getDesiredSizeBytes(numSplits); - List> shards = initialSource.split(desiredSizeBytes, options); + long estimatedSizeBytes = initialSource.getEstimatedSizeBytes(options); + List> shards = + BoundedSourceSplitter.split(initialSource, options, numSplits, estimatedSizeBytes); + numTasks = Math.max(1, numSplits); + lazyAssignment = useLazyAssignment(estimatedSizeBytes, numTasks); int numShards = shards.size(); SourceInputSplit[] sourceInputSplits = new SourceInputSplit[numShards]; @@ -146,9 +142,57 @@ public SourceInputSplit[] createInputSplits(int numSplits) throws IOException } } + private boolean useLazyAssignment(long estimatedSizeBytes, int parallelism) { + switch (BoundedSourceSplitter.assignment(options)) { + case LAZY: + return true; + case SIZE_BASED: + if (estimatedSizeBytes <= 0 || estimatedSizeBytes == Long.MAX_VALUE) { + return true; + } + long thresholdBytes = BoundedSourceSplitter.staticSplitThresholdBytes(options); + return estimatedSizeBytes / parallelism >= thresholdBytes; + case STATIC: + default: + return false; + } + } + @Override public InputSplitAssigner getInputSplitAssigner(final SourceInputSplit[] sourceInputSplits) { - return new DefaultInputSplitAssigner(sourceInputSplits); + LOG.info( + "Using {} assignment for {} splits of source {}", + lazyAssignment ? "lazy" : "static", + sourceInputSplits.length, + initialSource); + if (lazyAssignment) { + return new DefaultInputSplitAssigner(sourceInputSplits); + } + return new StaticInputSplitAssigner(sourceInputSplits, numTasks); + } + + /** Assigns split {@code i} to task {@code i % numTasks}, mirroring the DataStream enumerator. */ + static class StaticInputSplitAssigner implements InputSplitAssigner { + private final Map> pendingSplits = new HashMap<>(); + + StaticInputSplitAssigner(InputSplit[] splits, int numTasks) { + for (InputSplit split : splits) { + pendingSplits + .computeIfAbsent(split.getSplitNumber() % numTasks, ignored -> new ArrayDeque<>()) + .add(split); + } + } + + @Override + public synchronized InputSplit getNextInputSplit(String host, int taskId) { + Deque splits = pendingSplits.get(taskId); + return splits == null ? null : splits.poll(); + } + + @Override + public synchronized void returnInputSplit(List splits, int taskId) { + pendingSplits.computeIfAbsent(taskId, ignored -> new ArrayDeque<>()).addAll(splits); + } } @Override 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 f3bbf14f68fd..b969956cdcda 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 @@ -21,6 +21,7 @@ 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.wrappers.BoundedSourceSplitter; 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; @@ -104,15 +105,15 @@ public SplitEnumerator, FlinkSourceEnumeratorState> creat SplitEnumeratorContext> enumContext) throws Exception { FlinkPipelineOptions options = serializablePipelineOptions.get().as(FlinkPipelineOptions.class); if (boundedness == Boundedness.BOUNDED) { - long thresholdMb = options.getSourceStaticSplitThresholdMb(); - if (thresholdMb < 0) { - return new FlinkSourceSplitEnumerator<>(enumContext, beamSource, options, numSplits); + switch (BoundedSourceSplitter.assignment(options)) { + case STATIC: + return new FlinkSourceSplitEnumerator<>(enumContext, beamSource, options, numSplits); + case SIZE_BASED: + return new SizeBasedFlinkSourceSplitEnumerator<>( + enumContext, (BoundedSource) beamSource, options, numSplits); + case LAZY: + return new LazyFlinkSourceSplitEnumerator<>(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); } 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 index 26a5dd4d9349..481fdc9d8af6 100644 --- 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 @@ -19,9 +19,8 @@ import java.util.ArrayList; import java.util.List; -import org.apache.beam.runners.flink.FlinkPipelineOptions; +import org.apache.beam.runners.flink.translation.wrappers.BoundedSourceSplitter; 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; @@ -30,8 +29,6 @@ /** Shared Beam source sizing and splitting helpers. */ final class FlinkSourceSplitUtils { - static final long MEBIBYTE = 1024L * 1024L; - private static final Logger LOG = LoggerFactory.getLogger(FlinkSourceSplitUtils.class); private FlinkSourceSplitUtils() {} @@ -47,18 +44,8 @@ static ArrayList> splitBoundedSource( int numSplits, long estimatedSizeBytes) throws Exception { - long desiredSizeBytes = - getDesiredSizeBytes(boundedSource, pipelineOptions, numSplits, estimatedSizeBytes); - List> splits = - boundedSource.split(desiredSizeBytes, pipelineOptions); - LOG.info( - "Split bounded source {} in {} splits (estimated size {} bytes, " - + "desired split size {} bytes)", - boundedSource, - splits.size(), - estimatedSizeBytes, - desiredSizeBytes); - return toFlinkSplits(splits); + return toFlinkSplits( + BoundedSourceSplitter.split(boundedSource, pipelineOptions, numSplits, estimatedSizeBytes)); } static ArrayList> splitUnboundedSource( @@ -70,25 +57,6 @@ static ArrayList> splitUnboundedSource( return toFlinkSplits(splits); } - 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()); 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 bd07d00cd049..c4c9190c7540 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 @@ -176,7 +176,8 @@ private void assignNextSplit(int subtaskId, @Nullable String requesterHostname) return; } - FlinkSourceSplit split = pendingSplits.remove(pendingSplits.size() - 1); + // Splits are ordered largest first; handing those out first evens out the reader load. + FlinkSourceSplit split = pendingSplits.remove(0); 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 index ff1a79e3c30a..ef152ee054a3 100644 --- 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 @@ -25,6 +25,7 @@ import java.util.Optional; import javax.annotation.Nullable; import org.apache.beam.runners.flink.FlinkPipelineOptions; +import org.apache.beam.runners.flink.translation.wrappers.BoundedSourceSplitter; 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; @@ -135,13 +136,12 @@ private FlinkSourceEnumeratorState selectAndSplit() throws Exception { } private AssignmentMode selectAssignmentMode(long estimatedSizeBytes) { - long thresholdMb = - pipelineOptions.as(FlinkPipelineOptions.class).getSourceStaticSplitThresholdMb(); - if (thresholdMb <= 0) { + long thresholdBytes = BoundedSourceSplitter.staticSplitThresholdBytes(pipelineOptions); + if (thresholdBytes <= 0) { throw new IllegalArgumentException( "Size-based source assignment requires a positive threshold, but received " - + thresholdMb - + "."); + + pipelineOptions.as(FlinkPipelineOptions.class).getSourceStaticSplitThresholdMb() + + " MiB."); } if (estimatedSizeBytes <= 0 || estimatedSizeBytes == Long.MAX_VALUE) { LOG.info( @@ -156,7 +156,6 @@ private AssignmentMode selectAssignmentMode(long estimatedSizeBytes) { "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( diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java new file mode 100644 index 000000000000..9eb51491c039 --- /dev/null +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -0,0 +1,213 @@ +/* + * 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; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.flink.FlinkPipelineOptions; +import org.apache.beam.runners.flink.translation.wrappers.BoundedSourceSplitter.Assignment; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.VarLongCoder; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.junit.Test; + +/** Tests for {@link BoundedSourceSplitter}. */ +public class BoundedSourceSplitterTest { + + private final FlinkPipelineOptions options = FlinkPipelineOptions.defaults(); + + @Test + public void testResolvesAssignment() { + FlinkPipelineOptions batch = FlinkPipelineOptions.defaults(); + assertEquals(Assignment.STATIC, BoundedSourceSplitter.assignment(batch)); + + FlinkPipelineOptions streaming = FlinkPipelineOptions.defaults(); + streaming.setStreaming(true); + assertEquals(Assignment.LAZY, BoundedSourceSplitter.assignment(streaming)); + + batch.setSourceStaticSplitThresholdMb(0L); + assertEquals(Assignment.LAZY, BoundedSourceSplitter.assignment(batch)); + batch.setSourceStaticSplitThresholdMb(-1L); + assertEquals(Assignment.STATIC, BoundedSourceSplitter.assignment(batch)); + batch.setSourceStaticSplitThresholdMb(10L); + assertEquals(Assignment.SIZE_BASED, BoundedSourceSplitter.assignment(batch)); + } + + @Test + public void testResplitsWhenTooFewSplits() throws Exception { + // The first split ignores the desired size and only produces two splits. + BoundedSource source = new CoarseSource(Collections.singletonList(range(0, 8000))); + + List> splits = BoundedSourceSplitter.split(source, options, 8, 8000); + + assertEquals(8, splits.size()); + assertEquals(8000, totalSize(splits)); + for (BoundedSource split : splits) { + assertEquals(1000, split.getEstimatedSizeBytes(options)); + } + } + + @Test + public void testResplitsLargeSplits() throws Exception { + // One large file and many small ones. + List files = new ArrayList<>(); + files.add(range(0, 10_000)); + for (int i = 0; i < 12; i++) { + files.add(range(10_000 + i * 100, 10_100 + i * 100)); + } + BoundedSource source = new CoarseSource(files); + + List> splits = BoundedSourceSplitter.split(source, options, 4, 11_200); + + assertEquals(11_200, totalSize(splits)); + assertEquals(0, splits.size() % 4); + long[] perReader = perReaderSizes(splits, 4); + long max = Arrays.stream(perReader).max().getAsLong(); + long min = Arrays.stream(perReader).min().getAsLong(); + assertTrue("Unbalanced readers: " + Arrays.toString(perReader), max - min <= 11_200 / 4 / 2); + } + + @Test + public void testRoundsSplitCountToMultipleOfParallelism() throws Exception { + List files = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + files.add(range(i * 1000, (i + 1) * 1000)); + } + BoundedSource source = new CoarseSource(files); + + List> splits = BoundedSourceSplitter.split(source, options, 4, 5000); + + assertEquals(0, splits.size() % 4); + assertEquals(5000, totalSize(splits)); + } + + @Test + public void testKeepsUnsplittableSources() throws Exception { + BoundedSource source = + new CoarseSource(Arrays.asList(new UnsplittableSource(), new UnsplittableSource())); + + List> splits = BoundedSourceSplitter.split(source, options, 4, 2); + + assertEquals(2, splits.size()); + } + + private static RangeSource range(long start, long end) { + return new RangeSource(start, end); + } + + private long totalSize(List> splits) throws Exception { + long total = 0; + for (BoundedSource split : splits) { + total += split.getEstimatedSizeBytes(options); + } + return total; + } + + private long[] perReaderSizes(List> splits, int readers) throws Exception { + long[] sizes = new long[readers]; + for (int i = 0; i < splits.size(); i++) { + sizes[i % readers] += splits.get(i).getEstimatedSizeBytes(options); + } + return sizes; + } + + private abstract static class TestSource extends BoundedSource { + @Override + public BoundedReader createReader(PipelineOptions options) { + throw new UnsupportedOperationException(); + } + + @Override + public Coder getOutputCoder() { + return VarLongCoder.of(); + } + } + + /** Splits into its children regardless of the desired size, like a multi-file source. */ + private static class CoarseSource extends TestSource { + private final List> children; + + CoarseSource(List> children) { + this.children = children; + } + + @Override + public List> split(long desiredSize, PipelineOptions options) + throws Exception { + if (children.size() == 1) { + RangeSource only = (RangeSource) children.get(0); + long middle = (only.start + only.end) / 2; + return Arrays.asList(range(only.start, middle), range(middle, only.end)); + } + return children; + } + + @Override + public long getEstimatedSizeBytes(PipelineOptions options) throws Exception { + long size = 0; + for (BoundedSource child : children) { + size += child.getEstimatedSizeBytes(options); + } + return size; + } + } + + /** One byte per element, splits honoring the desired size. */ + private static class RangeSource extends TestSource { + private final long start; + private final long end; + + RangeSource(long start, long end) { + this.start = start; + this.end = end; + } + + @Override + public List split(long desiredSize, PipelineOptions options) { + List splits = new ArrayList<>(); + long step = Math.max(1, desiredSize); + for (long s = start; s < end; s += step) { + splits.add(range(s, Math.min(end, s + step))); + } + return splits; + } + + @Override + public long getEstimatedSizeBytes(PipelineOptions options) { + return end - start; + } + } + + private static class UnsplittableSource extends TestSource { + @Override + public List split(long desiredSize, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public long getEstimatedSizeBytes(PipelineOptions options) { + return 1; + } + } +} diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormatTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormatTest.java index f070a1c4943e..ffa1f8f68af3 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormatTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/SourceInputFormatTest.java @@ -17,12 +17,19 @@ */ package org.apache.beam.runners.flink.translation.wrappers; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import org.apache.beam.runners.flink.FlinkPipelineOptions; import org.apache.beam.runners.flink.metrics.FlinkMetricContainer; import org.apache.beam.sdk.io.BoundedSource; import org.apache.beam.sdk.io.CountingSource; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.flink.api.common.functions.RuntimeContext; +import org.apache.flink.api.common.io.DefaultInputSplitAssigner; +import org.apache.flink.core.io.InputSplit; +import org.apache.flink.core.io.InputSplitAssigner; import org.junit.Test; import org.mockito.Mockito; import org.powermock.reflect.Whitebox; @@ -52,6 +59,39 @@ public void testAccumulatorRegistrationOnOperatorClose() throws Exception { Mockito.verify(monitoredContainer).registerMetricsForPipelineResult(); } + @Test + public void testStaticSplitAssignmentByDefault() throws Exception { + SourceInputFormat sourceInputFormat = + new TestSourceInputFormat<>( + "step", CountingSource.upTo(1000), PipelineOptionsFactory.create()); + + SourceInputSplit[] splits = sourceInputFormat.createInputSplits(2); + InputSplitAssigner assigner = sourceInputFormat.getInputSplitAssigner(splits); + assertTrue(assigner instanceof SourceInputFormat.StaticInputSplitAssigner); + + int assigned = 0; + for (int taskId = 0; taskId < 2; taskId++) { + InputSplit split; + while ((split = assigner.getNextInputSplit("host", taskId)) != null) { + assertEquals(taskId, split.getSplitNumber() % 2); + assigned++; + } + } + assertEquals(splits.length, assigned); + } + + @Test + public void testLazySplitAssignmentWhenRequested() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.as(FlinkPipelineOptions.class).setSourceStaticSplitThresholdMb(0L); + SourceInputFormat sourceInputFormat = + new TestSourceInputFormat<>("step", CountingSource.upTo(1000), options); + + SourceInputSplit[] splits = sourceInputFormat.createInputSplits(2); + assertTrue( + sourceInputFormat.getInputSplitAssigner(splits) instanceof DefaultInputSplitAssigner); + } + private static class TestSourceInputFormat extends SourceInputFormat { public TestSourceInputFormat( 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 c4d220d22138..68a69a7e95ea 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,6 +18,7 @@ package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -56,7 +57,7 @@ public class FlinkSourceSplitEnumeratorTest { @Test public void testSelectsAssignmentModeFromEstimatedSizeAndConfiguration() throws Exception { - assertEquals(0L, (long) FlinkPipelineOptions.defaults().getSourceStaticSplitThresholdMb()); + assertNull(FlinkPipelineOptions.defaults().getSourceStaticSplitThresholdMb()); long thresholdBytes = STATIC_SPLIT_THRESHOLD_MB * MEBIBYTE; long largeEstimate = 1024L * MEBIBYTE; @@ -522,7 +523,7 @@ public List> split( long desiredBundleSizeBytes, PipelineOptions options) { List splits = new ArrayList<>(generatedSplits); for (int i = 0; i < generatedSplits; i++) { - splits.add(new TestEstimatedSizeBoundedSource(1L, 1, estimationCalls)); + splits.add(new TestEstimatedSizeBoundedSource(1L, 1, new AtomicInteger())); } return splits; } 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 8e0fc007f13b..a5c6c0f64b33 100644 --- a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html @@ -189,8 +189,8 @@ sourceStaticSplitThresholdMb - 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: 0 + Split assignment for bounded sources. When unset, batch pipelines assign splits statically (round-robin, after re-splitting the source so that splits are numerous and evenly sized) and streaming pipelines assign them lazily. 0 always uses lazy assignment. A positive value is a threshold in MiB per source reader: static assignment is used for sources with a known, positive estimate below the threshold and lazy assignment otherwise. Any negative value always uses static assignment. + stateBackend 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 d2ab687c42be..342af6853809 100644 --- a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html @@ -189,8 +189,8 @@ source_static_split_threshold_mb - 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: 0 + Split assignment for bounded sources. When unset, batch pipelines assign splits statically (round-robin, after re-splitting the source so that splits are numerous and evenly sized) and streaming pipelines assign them lazily. 0 always uses lazy assignment. A positive value is a threshold in MiB per source reader: static assignment is used for sources with a known, positive estimate below the threshold and lazy assignment otherwise. Any negative value always uses static assignment. + state_backend From 238aa06586cd5678f7306bbf2cc9361aaf762a9c Mon Sep 17 00:00:00 2001 From: jto Date: Fri, 2 Oct 2026 16:41:14 +0200 Subject: [PATCH 02/11] [Flink] Simplify re-splitting of oversized bounded source splits Replace the count-based re-split (uniform factor per split) with a single pass that re-splits only the splits much larger than the desired split size, as some sources (e.g. BigQuery) return splits way larger than requested. Share the per-split re-split logic with size balancing and express balancing without loops and breaks. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 101 ++++++++---------- 1 file changed, 47 insertions(+), 54 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index 1a6ac3f46068..0516eda9c117 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -21,6 +21,8 @@ import java.util.Collections; import java.util.Comparator; import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; import org.apache.beam.runners.flink.FlinkPipelineOptions; import org.apache.beam.sdk.io.BoundedSource; import org.apache.beam.sdk.io.FileBasedSource; @@ -40,7 +42,8 @@ *

    *
  1. Split the source using {@code estimatedSize / parallelism} as the desired split size, * capped by {@link FlinkPipelineOptions#getFileInputSplitMaxSizeMB()} for file sources. - *
  2. While there are fewer splits than readers, re-split each split. + *
  3. Re-split the splits that are much larger than the desired split size, as some sources + * return splits way larger than requested. *
  4. While the split sizes vary a lot (high coefficient of variation), re-split the larger * splits. *
  5. Halve the largest splits until the number of splits is a multiple of the parallelism. @@ -117,9 +120,13 @@ public static List> split( List> splits = sized(source.split(desiredSizeBytes, options), options); int initialSplits = splits.size(); - splits = ensureEnoughSplits(splits, options, readers, desiredSizeBytes); + // Some sources (e.g. BigQuery) return splits way larger than requested that can still be + // split further. + splits = resplitLargerThan(splits, options, desiredSizeBytes); if (allSizesKnown(splits)) { - splits = balanceSizes(splits, options, readers, estimatedSizeBytes, maxSplitSizeBytes); + long minTargetBytes = + Math.max(1L, estimatedSizeBytes / ((long) readers * MAX_SPLITS_PER_READER_FOR_BALANCING)); + splits = balanceSizes(splits, options, minTargetBytes, maxSplitSizeBytes, MAX_RESPLIT_ROUNDS); splits = roundToMultipleOfParallelism(splits, options, readers); splits = orderForRoundRobin(splits, readers); } @@ -151,57 +158,41 @@ private static long maxSplitSizeBytes(BoundedSource source, PipelineOptions o return Long.MAX_VALUE; } - /** Re-splits every split while there are fewer splits than readers. */ - private static List> ensureEnoughSplits( - List> splits, PipelineOptions options, int readers, long desiredSizeBytes) - throws Exception { - for (int round = 0; round < MAX_RESPLIT_ROUNDS && splits.size() < readers; round++) { - int factor = (readers + splits.size() - 1) / Math.max(1, splits.size()); - List> next = new ArrayList<>(); - for (SizedSource split : splits) { - long base = split.sizeBytes > 0 ? split.sizeBytes : desiredSizeBytes; - next.addAll(resplit(split, options, Math.max(1L, base / factor))); - } - if (next.size() <= splits.size()) { - break; - } - splits = next; - } - return splits; - } - - /** Re-splits the larger splits while the split sizes vary a lot. */ + /** + * Re-splits the splits larger than the mean size (bounded by {@code minTargetBytes} and {@code + * maxSplitSizeBytes}) while the split sizes vary a lot, at most {@code roundsLeft} times. + */ private static List> balanceSizes( List> splits, PipelineOptions options, - int readers, - long estimatedSizeBytes, - long maxSplitSizeBytes) - throws Exception { - long minTargetBytes = - Math.max(1L, estimatedSizeBytes / ((long) readers * MAX_SPLITS_PER_READER_FOR_BALANCING)); - for (int round = 0; round < MAX_RESPLIT_ROUNDS; round++) { - if (coefficientOfVariation(splits) <= MAX_SIZE_COEFFICIENT_OF_VARIATION) { - break; - } - long targetBytes = Math.min(Math.max((long) mean(splits), minTargetBytes), maxSplitSizeBytes); - List> next = new ArrayList<>(); - boolean progress = false; - for (SizedSource split : splits) { - if (split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes) { - List> parts = resplit(split, options, targetBytes); - progress |= parts.size() > 1; - next.addAll(parts); - } else { - next.add(split); - } - } - if (!progress || !allSizesKnown(next)) { - break; - } - splits = next; + long minTargetBytes, + long maxSplitSizeBytes, + int roundsLeft) { + if (roundsLeft == 0 || coefficientOfVariation(splits) <= MAX_SIZE_COEFFICIENT_OF_VARIATION) { + return splits; } - return splits; + long targetBytes = Math.min(Math.max((long) mean(splits), minTargetBytes), maxSplitSizeBytes); + List> next = resplitLargerThan(splits, options, targetBytes); + if (next.size() > splits.size() && allSizesKnown(next)) { + return balanceSizes(next, options, minTargetBytes, maxSplitSizeBytes, roundsLeft - 1); + } else { + return splits; + } + } + + /** Re-splits, once, each split larger than {@link #OVERSIZED_SPLIT_FACTOR} x the target. */ + private static List> resplitLargerThan( + List> splits, PipelineOptions options, long targetBytes) { + return splits.stream() + .flatMap( + split -> { + if (split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes) { + return resplit(split, options, targetBytes).stream(); + } else { + return Stream.of(split); + } + }) + .collect(Collectors.toList()); } /** Halves the largest splits until the number of splits is a multiple of the parallelism. */ @@ -218,10 +209,12 @@ private static List> roundToMultipleOfParallelism( break; } SizedSource largest = result.remove(0); - List> parts = - largest.sizeBytes > 1 - ? resplit(largest, options, Math.max(1L, largest.sizeBytes / 2)) - : Collections.singletonList(largest); + List> parts; + if (largest.sizeBytes > 1) { + parts = resplit(largest, options, Math.max(1L, largest.sizeBytes / 2)); + } else { + parts = Collections.singletonList(largest); + } if (parts.size() <= 1 || !allSizesKnown(parts)) { unsplittable.add(largest); continue; From 4f385a1217f7cd6410920ca523e144676a6e14b9 Mon Sep 17 00:00:00 2001 From: jto Date: Fri, 2 Oct 2026 17:18:08 +0200 Subject: [PATCH 03/11] [Flink] Relax rounding of bounded source split count to parallelism Rewrite roundToMultipleOfParallelism without break/continue, using a priority queue and a halve() helper. Stop halving once the split count is within 10% of the parallelism of a multiple of it, since re-splitting introduces some skew anyway. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 72 +++++++++++++------ .../wrappers/BoundedSourceSplitterTest.java | 14 ++++ 2 files changed, 65 insertions(+), 21 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index 0516eda9c117..86b96cc8c7eb 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -21,6 +21,7 @@ import java.util.Collections; import java.util.Comparator; import java.util.List; +import java.util.PriorityQueue; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.beam.runners.flink.FlinkPipelineOptions; @@ -46,7 +47,8 @@ * return splits way larger than requested. *
  6. While the split sizes vary a lot (high coefficient of variation), re-split the larger * splits. - *
  7. Halve the largest splits until the number of splits is a multiple of the parallelism. + *
  8. Halve the largest splits until the number of splits is close to a multiple of the + * parallelism. *
  9. Order the splits so that round-robin assignment gives each reader a similar amount of data. *
* @@ -66,6 +68,12 @@ public final class BoundedSourceSplitter { /** Splits larger than this factor times the target size are re-split during balancing. */ static final double OVERSIZED_SPLIT_FACTOR = 1.5; + /** + * How far, as a fraction of the parallelism, the split count may be from a multiple of the + * parallelism. Re-splitting introduces some skew anyway, so an exact multiple is not required. + */ + static final double MULTIPLE_OF_PARALLELISM_TOLERANCE = 0.1; + /** Lower bound of the balancing target size, relative to {@code estimatedSize/parallelism}. */ static final int MAX_SPLITS_PER_READER_FOR_BALANCING = 4; @@ -195,37 +203,59 @@ private static List> resplitLargerThan( .collect(Collectors.toList()); } - /** Halves the largest splits until the number of splits is a multiple of the parallelism. */ + /** + * Halves the largest splits until the number of splits is close to a multiple of the parallelism. + * Splits that cannot be halved are set aside so they are not retried. The result is not ordered. + */ private static List> roundToMultipleOfParallelism( - List> splits, PipelineOptions options, int readers) throws Exception { - List> result = new ArrayList<>(splits); - result.sort(Comparator.comparingLong((SizedSource s) -> s.sizeBytes).reversed()); - // Splits that could not be split further are moved here so they are not retried. + List> splits, PipelineOptions options, int readers) { + PriorityQueue> candidates = new PriorityQueue<>(bySizeDescending()); + candidates.addAll(splits); List> unsplittable = new ArrayList<>(); for (int attempt = 0; - attempt < 2 * readers && (result.size() + unsplittable.size()) % readers != 0; + attempt < 2 * readers + && !candidates.isEmpty() + && !isCloseToMultipleOf(candidates.size() + unsplittable.size(), readers); attempt++) { - if (result.isEmpty()) { - break; - } - SizedSource largest = result.remove(0); - List> parts; - if (largest.sizeBytes > 1) { - parts = resplit(largest, options, Math.max(1L, largest.sizeBytes / 2)); + SizedSource largest = candidates.remove(); + List> halves = halve(largest, options); + if (halves.size() > 1) { + candidates.addAll(halves); } else { - parts = Collections.singletonList(largest); - } - if (parts.size() <= 1 || !allSizesKnown(parts)) { unsplittable.add(largest); - continue; } - result.addAll(parts); - result.sort(Comparator.comparingLong((SizedSource s) -> s.sizeBytes).reversed()); } + List> result = new ArrayList<>(candidates); result.addAll(unsplittable); return result; } + /** + * Whether {@code count} is within {@link #MULTIPLE_OF_PARALLELISM_TOLERANCE} x {@code readers} of + * a multiple of {@code readers}, e.g. 1.1x or 1.9x the parallelism with a 10% tolerance. + */ + private static boolean isCloseToMultipleOf(int count, int readers) { + int remainder = count % readers; + return Math.min(remainder, readers - remainder) <= MULTIPLE_OF_PARALLELISM_TOLERANCE * readers; + } + + /** Splits {@code split} in two, or returns it alone if that is not possible. */ + private static List> halve(SizedSource split, PipelineOptions options) { + if (split.sizeBytes <= 1) { + return Collections.singletonList(split); + } + List> parts = resplit(split, options, split.sizeBytes / 2); + if (allSizesKnown(parts)) { + return parts; + } else { + return Collections.singletonList(split); + } + } + + private static Comparator> bySizeDescending() { + return Comparator.comparingLong((SizedSource s) -> s.sizeBytes).reversed(); + } + /** * Sorts splits by decreasing size in a "snake" order (0..p-1, p-1..0, ...) so that assigning * split {@code i} to reader {@code i % p} spreads the large splits across readers. @@ -233,7 +263,7 @@ private static List> roundToMultipleOfParallelism( private static List> orderForRoundRobin( List> splits, int readers) { List> sorted = new ArrayList<>(splits); - sorted.sort(Comparator.comparingLong((SizedSource s) -> s.sizeBytes).reversed()); + sorted.sort(bySizeDescending()); List> ordered = new ArrayList<>(sorted.size()); for (int start = 0; start < sorted.size(); start += readers) { int end = Math.min(start + readers, sorted.size()); diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java index 9eb51491c039..ecfe1f145335 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -102,6 +102,20 @@ public void testRoundsSplitCountToMultipleOfParallelism() throws Exception { assertEquals(5000, totalSize(splits)); } + @Test + public void testToleratesSplitCountCloseToMultipleOfParallelism() throws Exception { + List files = new ArrayList<>(); + for (int i = 0; i < 21; i++) { + files.add(range(i * 1000, (i + 1) * 1000)); + } + BoundedSource source = new CoarseSource(files); + + // 21 splits for 20 readers is within 10% of a multiple: nothing is halved. + List> splits = BoundedSourceSplitter.split(source, options, 20, 21_000); + + assertEquals(21, splits.size()); + } + @Test public void testKeepsUnsplittableSources() throws Exception { BoundedSource source = From 10a7fa01d62190d033d781356650eebd5669f36e Mon Sep 17 00:00:00 2001 From: jto Date: Mon, 5 Oct 2026 15:08:14 +0200 Subject: [PATCH 04/11] [Flink] Order bounded source splits for the actual reader count The DataStream translation passes maxParallelism as numSplits when it is set, while the static enumerator assigns split i to reader i % currentParallelism. Pass the enumerator's current parallelism to the splitter separately, so rounding and round-robin ordering match the actual assignment. numSplits still drives the desired split size. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 19 ++++++++++++- .../io/source/FlinkSourceSplitEnumerator.java | 6 ++++- .../io/source/FlinkSourceSplitUtils.java | 4 ++- .../LazyFlinkSourceSplitEnumerator.java | 6 ++++- .../SizeBasedFlinkSourceSplitEnumerator.java | 6 ++++- .../wrappers/BoundedSourceSplitterTest.java | 27 +++++++++++++++++-- 6 files changed, 61 insertions(+), 7 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index 86b96cc8c7eb..e183e02035e7 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -121,9 +121,26 @@ public static long mebibytesToBytes(long mebibytes) { public static List> split( BoundedSource source, PipelineOptions options, int parallelism, long estimatedSizeBytes) throws Exception { + return split(source, options, parallelism, parallelism, estimatedSizeBytes); + } + + /** + * Splits {@code source} into about {@code numSplits} splits for {@code parallelism} readers. + * {@code numSplits} sets the desired split size, while {@code parallelism} must be the number of + * readers splits are assigned to, as split {@code i} is assumed to go to reader {@code i % + * parallelism}. + */ + public static List> split( + BoundedSource source, + PipelineOptions options, + int numSplits, + int parallelism, + long estimatedSizeBytes) + throws Exception { int readers = Math.max(1, parallelism); long maxSplitSizeBytes = maxSplitSizeBytes(source, options); - long desiredSizeBytes = Math.min(Math.max(1L, estimatedSizeBytes / readers), maxSplitSizeBytes); + long desiredSizeBytes = + Math.min(Math.max(1L, estimatedSizeBytes / Math.max(1, numSplits)), maxSplitSizeBytes); List> splits = sized(source.split(desiredSizeBytes, options), options); int initialSplits = splits.size(); 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 bb8956f5daea..413b5b7e867e 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 @@ -160,7 +160,11 @@ private ArrayList> splitBeamSource() throws Exception { long estimatedSizeBytes = FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource, pipelineOptions); return FlinkSourceSplitUtils.splitBoundedSource( - boundedSource, pipelineOptions, numSplits, estimatedSizeBytes); + boundedSource, + pipelineOptions, + numSplits, + context.currentParallelism(), + estimatedSizeBytes); } if (beamSource instanceof UnboundedSource) { return FlinkSourceSplitUtils.splitUnboundedSource( 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 index 481fdc9d8af6..1962558bbb02 100644 --- 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 @@ -42,10 +42,12 @@ static ArrayList> splitBoundedSource( BoundedSource boundedSource, PipelineOptions pipelineOptions, int numSplits, + int parallelism, long estimatedSizeBytes) throws Exception { return toFlinkSplits( - BoundedSourceSplitter.split(boundedSource, pipelineOptions, numSplits, estimatedSizeBytes)); + BoundedSourceSplitter.split( + boundedSource, pipelineOptions, numSplits, parallelism, estimatedSizeBytes)); } static ArrayList> splitUnboundedSource( 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 c4c9190c7540..e0a64d7dc867 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 @@ -156,7 +156,11 @@ private ArrayList> splitBeamSource() throws Exception { long estimatedSizeBytes = FlinkSourceSplitUtils.estimateBoundedSourceSize(boundedSource, pipelineOptions); return FlinkSourceSplitUtils.splitBoundedSource( - boundedSource, pipelineOptions, numSplits, estimatedSizeBytes); + boundedSource, + pipelineOptions, + numSplits, + context.currentParallelism(), + estimatedSizeBytes); } private void sendPendingSplitRequests() { 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 index ef152ee054a3..bcf4c8ff5d41 100644 --- 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 @@ -131,7 +131,11 @@ private FlinkSourceEnumeratorState selectAndSplit() throws Exception { AssignmentMode selectedMode = selectAssignmentMode(estimatedSizeBytes); ArrayList> splits = FlinkSourceSplitUtils.splitBoundedSource( - boundedSource, pipelineOptions, numSplits, estimatedSizeBytes); + boundedSource, + pipelineOptions, + numSplits, + context.currentParallelism(), + estimatedSizeBytes); return new FlinkSourceEnumeratorState<>(selectedMode, splits); } diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java index ecfe1f145335..38b7bdeba9f3 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.flink.translation.wrappers; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; @@ -116,10 +117,26 @@ public void testToleratesSplitCountCloseToMultipleOfParallelism() throws Excepti assertEquals(21, splits.size()); } + @Test + public void testOrdersSplitsForParallelismRatherThanNumSplits() throws Exception { + BoundedSource source = + new CoarseSource( + Arrays.asList( + new UnsplittableSource(4), + new UnsplittableSource(3), + new UnsplittableSource(2), + new UnsplittableSource(1))); + + // 4 splits requested, but assigned round-robin to 2 readers. + List> splits = BoundedSourceSplitter.split(source, options, 4, 2, 10); + + assertArrayEquals(new long[] {5, 5}, perReaderSizes(splits, 2)); + } + @Test public void testKeepsUnsplittableSources() throws Exception { BoundedSource source = - new CoarseSource(Arrays.asList(new UnsplittableSource(), new UnsplittableSource())); + new CoarseSource(Arrays.asList(new UnsplittableSource(1), new UnsplittableSource(1))); List> splits = BoundedSourceSplitter.split(source, options, 4, 2); @@ -214,6 +231,12 @@ public long getEstimatedSizeBytes(PipelineOptions options) { } private static class UnsplittableSource extends TestSource { + private final long sizeBytes; + + UnsplittableSource(long sizeBytes) { + this.sizeBytes = sizeBytes; + } + @Override public List split(long desiredSize, PipelineOptions options) { return Collections.singletonList(this); @@ -221,7 +244,7 @@ public List split(long desiredSize, PipelineOptions options) @Override public long getEstimatedSizeBytes(PipelineOptions options) { - return 1; + return sizeBytes; } } } From 49d6f69e45bd911b73d838875cf358c9dbfeabde Mon Sep 17 00:00:00 2001 From: jto Date: Mon, 5 Oct 2026 15:12:41 +0200 Subject: [PATCH 05/11] [Flink] Balance bounded source split count by per-reader imbalance Instead of accepting split counts within 10% of the parallelism of a multiple of it, halve splits until round-robin assignment gives the busiest reader at most 10% more splits than the average. This rejects counts like 1.1x the parallelism (busiest reader does ~1.8x the average) while accepting any count once readers get many splits. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 33 ++++++++++--------- .../wrappers/BoundedSourceSplitterTest.java | 33 +++++++++++++------ 2 files changed, 41 insertions(+), 25 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index e183e02035e7..324d1358b826 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -47,8 +47,8 @@ * return splits way larger than requested. *
  • While the split sizes vary a lot (high coefficient of variation), re-split the larger * splits. - *
  • Halve the largest splits until the number of splits is close to a multiple of the - * parallelism. + *
  • Halve the largest splits until round-robin assignment gives each reader about the same + * number of splits. *
  • Order the splits so that round-robin assignment gives each reader a similar amount of data. * * @@ -69,10 +69,11 @@ public final class BoundedSourceSplitter { static final double OVERSIZED_SPLIT_FACTOR = 1.5; /** - * How far, as a fraction of the parallelism, the split count may be from a multiple of the - * parallelism. Re-splitting introduces some skew anyway, so an exact multiple is not required. + * How many more splits than the average the busiest reader may get with round-robin assignment, + * as a fraction of the average. Re-splitting introduces some skew anyway, so the split count does + * not need to be an exact multiple of the parallelism. */ - static final double MULTIPLE_OF_PARALLELISM_TOLERANCE = 0.1; + static final double SPLIT_COUNT_IMBALANCE_TOLERANCE = 0.1; /** Lower bound of the balancing target size, relative to {@code estimatedSize/parallelism}. */ static final int MAX_SPLITS_PER_READER_FOR_BALANCING = 4; @@ -152,7 +153,7 @@ public static List> split( long minTargetBytes = Math.max(1L, estimatedSizeBytes / ((long) readers * MAX_SPLITS_PER_READER_FOR_BALANCING)); splits = balanceSizes(splits, options, minTargetBytes, maxSplitSizeBytes, MAX_RESPLIT_ROUNDS); - splits = roundToMultipleOfParallelism(splits, options, readers); + splits = balanceSplitCount(splits, options, readers); splits = orderForRoundRobin(splits, readers); } @@ -221,10 +222,11 @@ private static List> resplitLargerThan( } /** - * Halves the largest splits until the number of splits is close to a multiple of the parallelism. - * Splits that cannot be halved are set aside so they are not retried. The result is not ordered. + * Halves the largest splits until round-robin assignment gives each reader about the same number + * of splits. Splits that cannot be halved are set aside so they are not retried. The result is + * not ordered. */ - private static List> roundToMultipleOfParallelism( + private static List> balanceSplitCount( List> splits, PipelineOptions options, int readers) { PriorityQueue> candidates = new PriorityQueue<>(bySizeDescending()); candidates.addAll(splits); @@ -232,7 +234,7 @@ private static List> roundToMultipleOfParallelism( for (int attempt = 0; attempt < 2 * readers && !candidates.isEmpty() - && !isCloseToMultipleOf(candidates.size() + unsplittable.size(), readers); + && !isBalancedEnough(candidates.size() + unsplittable.size(), readers); attempt++) { SizedSource largest = candidates.remove(); List> halves = halve(largest, options); @@ -248,12 +250,13 @@ private static List> roundToMultipleOfParallelism( } /** - * Whether {@code count} is within {@link #MULTIPLE_OF_PARALLELISM_TOLERANCE} x {@code readers} of - * a multiple of {@code readers}, e.g. 1.1x or 1.9x the parallelism with a 10% tolerance. + * Whether round-robin assignment of {@code count} splits gives the busiest reader at most {@link + * #SPLIT_COUNT_IMBALANCE_TOLERANCE} more splits than the average. For example 1.9x the + * parallelism is fine (2 / 1.9 = 1.05), but 1.1x is not (2 / 1.1 = 1.82). */ - private static boolean isCloseToMultipleOf(int count, int readers) { - int remainder = count % readers; - return Math.min(remainder, readers - remainder) <= MULTIPLE_OF_PARALLELISM_TOLERANCE * readers; + private static boolean isBalancedEnough(int count, int readers) { + long busiest = ((long) count + readers - 1) / readers; + return busiest * readers <= (1 + SPLIT_COUNT_IMBALANCE_TOLERANCE) * count; } /** Splits {@code split} in two, or returns it alone if that is not possible. */ diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java index 38b7bdeba9f3..a72b800eeeb6 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -82,7 +82,6 @@ public void testResplitsLargeSplits() throws Exception { List> splits = BoundedSourceSplitter.split(source, options, 4, 11_200); assertEquals(11_200, totalSize(splits)); - assertEquals(0, splits.size() % 4); long[] perReader = perReaderSizes(splits, 4); long max = Arrays.stream(perReader).max().getAsLong(); long min = Arrays.stream(perReader).min().getAsLong(); @@ -104,17 +103,31 @@ public void testRoundsSplitCountToMultipleOfParallelism() throws Exception { } @Test - public void testToleratesSplitCountCloseToMultipleOfParallelism() throws Exception { - List files = new ArrayList<>(); - for (int i = 0; i < 21; i++) { - files.add(range(i * 1000, (i + 1) * 1000)); - } - BoundedSource source = new CoarseSource(files); + public void testAcceptsSplitCountWithManySplitsPerReader() throws Exception { + // 39 splits for 4 readers: the busiest reader gets 10 splits for an average of 9.75. + List> splits = + BoundedSourceSplitter.split(equalFiles(39, 1000), options, 4, 39_000); - // 21 splits for 20 readers is within 10% of a multiple: nothing is halved. - List> splits = BoundedSourceSplitter.split(source, options, 20, 21_000); + assertEquals(39, splits.size()); + } - assertEquals(21, splits.size()); + @Test + public void testHalvesSplitsWithFewSplitsPerReader() throws Exception { + // 11 splits for 10 readers: the busiest reader would get 2 splits for an average of 1.1. + // Halving stops at 19 splits: 2 splits for an average of 1.9. + List> splits = + BoundedSourceSplitter.split(equalFiles(11, 1000), options, 10, 11_000); + + assertEquals(19, splits.size()); + assertEquals(11_000, totalSize(splits)); + } + + private static BoundedSource equalFiles(int count, long sizeBytes) { + List files = new ArrayList<>(); + for (int i = 0; i < count; i++) { + files.add(range(i * sizeBytes, (i + 1) * sizeBytes)); + } + return new CoarseSource(files); } @Test From 1a0da0ded758e59ac11bab7b5a58abf6b8112289 Mon Sep 17 00:00:00 2001 From: jto Date: Mon, 5 Oct 2026 15:41:03 +0200 Subject: [PATCH 06/11] [Flink] Assign lazy source splits in constant time The lazy enumerator hands out the largest splits first by taking them from the front of its pending splits. Use an ArrayDeque instead of ArrayList.remove(0), which made assigning n splits O(n^2). Co-Authored-By: Claude Opus 5.5 --- .../io/source/LazyFlinkSourceSplitEnumerator.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) 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 e0a64d7dc867..eed0c15b3e12 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 @@ -18,7 +18,9 @@ package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source; import java.io.IOException; +import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Deque; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -42,7 +44,7 @@ public class LazyFlinkSourceSplitEnumerator private final Source beamSource; private final PipelineOptions pipelineOptions; private final int numSplits; - private final List> pendingSplits; + private final Deque> pendingSplits; private final Map> pendingSplitRequests; private boolean splitsInitialized; @@ -65,7 +67,7 @@ public LazyFlinkSourceSplitEnumerator( this.beamSource = beamSource; this.pipelineOptions = pipelineOptions; this.numSplits = numSplits; - this.pendingSplits = new ArrayList<>(numSplits); + this.pendingSplits = new ArrayDeque<>(numSplits); this.pendingSplitRequests = new LinkedHashMap<>(); this.splitsInitialized = restoredState != null; @@ -181,7 +183,7 @@ private void assignNextSplit(int subtaskId, @Nullable String requesterHostname) } // Splits are ordered largest first; handing those out first evens out the reader load. - FlinkSourceSplit split = pendingSplits.remove(0); + FlinkSourceSplit split = pendingSplits.removeFirst(); context.assignSplit(split, subtaskId); LOG.info("Assigned split to subtask {} on host {}: {}", subtaskId, requesterHostname, split); } From 82f68c34b5fb7e89dcc75a89a482d558329a71f9 Mon Sep 17 00:00:00 2001 From: jto Date: Mon, 5 Oct 2026 15:41:03 +0200 Subject: [PATCH 07/11] [Flink] Do not re-split bounded source splits below 64 MiB The fixed cost of opening and reading a split dominates for small splits, and byte ranges smaller than a file's blocks may read nothing. Re-splitting of oversized splits, size balancing and halving no longer create splits smaller than min(64 MiB, estimatedSize / parallelism), so small inputs, which may feed transforms generating a lot more data, are still spread over all readers. The initial split is not affected. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 32 +++++++++++---- .../wrappers/BoundedSourceSplitterTest.java | 41 ++++++++++++++----- 2 files changed, 54 insertions(+), 19 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index 324d1358b826..35a467ac64ee 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -78,6 +78,14 @@ public final class BoundedSourceSplitter { /** Lower bound of the balancing target size, relative to {@code estimatedSize/parallelism}. */ static final int MAX_SPLITS_PER_READER_FOR_BALANCING = 4; + /** + * Splits are not re-split below this size by the refinement steps (the initial split is not + * affected), as the fixed cost of opening and reading a split would dominate. Small inputs still + * get one split per reader, as they may feed transforms that generate a lot more data: the floor + * is {@code min(MIN_SPLIT_SIZE_BYTES, estimatedSize/parallelism)}. + */ + static final long MIN_SPLIT_SIZE_BYTES = 64 * MEBIBYTE; + /** How bounded source splits are handed out to readers. */ public enum Assignment { /** Splits are assigned to readers up front, round-robin. */ @@ -140,6 +148,8 @@ public static List> split( throws Exception { int readers = Math.max(1, parallelism); long maxSplitSizeBytes = maxSplitSizeBytes(source, options); + long minSplitSizeBytes = + Math.max(1L, Math.min(MIN_SPLIT_SIZE_BYTES, estimatedSizeBytes / readers)); long desiredSizeBytes = Math.min(Math.max(1L, estimatedSizeBytes / Math.max(1, numSplits)), maxSplitSizeBytes); @@ -148,12 +158,14 @@ public static List> split( // Some sources (e.g. BigQuery) return splits way larger than requested that can still be // split further. - splits = resplitLargerThan(splits, options, desiredSizeBytes); + splits = resplitLargerThan(splits, options, Math.max(desiredSizeBytes, minSplitSizeBytes)); if (allSizesKnown(splits)) { long minTargetBytes = - Math.max(1L, estimatedSizeBytes / ((long) readers * MAX_SPLITS_PER_READER_FOR_BALANCING)); + Math.max( + minSplitSizeBytes, + estimatedSizeBytes / ((long) readers * MAX_SPLITS_PER_READER_FOR_BALANCING)); splits = balanceSizes(splits, options, minTargetBytes, maxSplitSizeBytes, MAX_RESPLIT_ROUNDS); - splits = balanceSplitCount(splits, options, readers); + splits = balanceSplitCount(splits, options, readers, minSplitSizeBytes); splits = orderForRoundRobin(splits, readers); } @@ -227,7 +239,7 @@ private static List> resplitLargerThan( * not ordered. */ private static List> balanceSplitCount( - List> splits, PipelineOptions options, int readers) { + List> splits, PipelineOptions options, int readers, long minSplitSizeBytes) { PriorityQueue> candidates = new PriorityQueue<>(bySizeDescending()); candidates.addAll(splits); List> unsplittable = new ArrayList<>(); @@ -237,7 +249,7 @@ private static List> balanceSplitCount( && !isBalancedEnough(candidates.size() + unsplittable.size(), readers); attempt++) { SizedSource largest = candidates.remove(); - List> halves = halve(largest, options); + List> halves = halve(largest, options, minSplitSizeBytes); if (halves.size() > 1) { candidates.addAll(halves); } else { @@ -259,9 +271,13 @@ private static boolean isBalancedEnough(int count, int readers) { return busiest * readers <= (1 + SPLIT_COUNT_IMBALANCE_TOLERANCE) * count; } - /** Splits {@code split} in two, or returns it alone if that is not possible. */ - private static List> halve(SizedSource split, PipelineOptions options) { - if (split.sizeBytes <= 1) { + /** + * Splits {@code split} in two, or returns it alone if that is not possible or would create splits + * smaller than {@code minSplitSizeBytes}. + */ + private static List> halve( + SizedSource split, PipelineOptions options, long minSplitSizeBytes) { + if (split.sizeBytes <= 1 || split.sizeBytes / 2 < minSplitSizeBytes) { return Collections.singletonList(split); } List> parts = resplit(split, options, split.sizeBytes / 2); diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java index a72b800eeeb6..dd187c8fc0ec 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -36,6 +36,9 @@ /** Tests for {@link BoundedSourceSplitter}. */ public class BoundedSourceSplitterTest { + private static final long MIB = BoundedSourceSplitter.MEBIBYTE; + private static final long GIB = 1024 * MIB; + private final FlinkPipelineOptions options = FlinkPipelineOptions.defaults(); @Test @@ -90,16 +93,12 @@ public void testResplitsLargeSplits() throws Exception { @Test public void testRoundsSplitCountToMultipleOfParallelism() throws Exception { - List files = new ArrayList<>(); - for (int i = 0; i < 5; i++) { - files.add(range(i * 1000, (i + 1) * 1000)); - } - BoundedSource source = new CoarseSource(files); - - List> splits = BoundedSourceSplitter.split(source, options, 4, 5000); + // 5 splits for 4 readers: halving the largest splits ends at 8 splits. + List> splits = + BoundedSourceSplitter.split(equalFiles(5, GIB), options, 4, 5 * GIB); - assertEquals(0, splits.size() % 4); - assertEquals(5000, totalSize(splits)); + assertEquals(8, splits.size()); + assertEquals(5 * GIB, totalSize(splits)); } @Test @@ -116,10 +115,30 @@ public void testHalvesSplitsWithFewSplitsPerReader() throws Exception { // 11 splits for 10 readers: the busiest reader would get 2 splits for an average of 1.1. // Halving stops at 19 splits: 2 splits for an average of 1.9. List> splits = - BoundedSourceSplitter.split(equalFiles(11, 1000), options, 10, 11_000); + BoundedSourceSplitter.split(equalFiles(11, GIB), options, 10, 11 * GIB); assertEquals(19, splits.size()); - assertEquals(11_000, totalSize(splits)); + assertEquals(11 * GIB, totalSize(splits)); + } + + @Test + public void testDoesNotHalveBelowMinSplitSize() throws Exception { + // Halving 100 MiB splits would create splits below the 64 MiB minimum. + List> splits = + BoundedSourceSplitter.split(equalFiles(11, 100 * MIB), options, 10, 1100 * MIB); + + assertEquals(11, splits.size()); + } + + @Test + public void testSmallInputStillGetsOneSplitPerReader() throws Exception { + // A 10 MiB input is below the minimum split size, but may feed transforms generating a lot + // more data, so it is still spread over all readers. + BoundedSource source = new CoarseSource(Collections.singletonList(range(0, 10 * MIB))); + + List> splits = BoundedSourceSplitter.split(source, options, 10, 10 * MIB); + + assertEquals(10, splits.size()); } private static BoundedSource equalFiles(int count, long sizeBytes) { From 29dd1354ff70b09216733785eddc7587a05eedf6 Mon Sep 17 00:00:00 2001 From: jto Date: Mon, 5 Oct 2026 15:46:30 +0200 Subject: [PATCH 08/11] [Flink] Aim for two bounded source splits per reader With static round-robin assignment, a single split per reader passes any size difference straight to the slowest reader, while two splits per reader let the ordering pair large splits with small ones. Re-split splits larger than estimatedSize / (2 x parallelism), and lower the minimum split size to min(64 MiB, estimatedSize / (2 x parallelism)) so small inputs also get two splits per reader. The initial split request is unchanged. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 32 +++++++++++++------ .../wrappers/BoundedSourceSplitterTest.java | 29 ++++++++++------- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index 35a467ac64ee..cb30056787ff 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -41,9 +41,10 @@ * refined in several bounded passes: * *
      - *
    1. Split the source using {@code estimatedSize / parallelism} as the desired split size, - * capped by {@link FlinkPipelineOptions#getFileInputSplitMaxSizeMB()} for file sources. - *
    2. Re-split the splits that are much larger than the desired split size, as some sources + *
    3. Split the source using {@code estimatedSize / numSplits} as the desired split size, capped + * by {@link FlinkPipelineOptions#getFileInputSplitMaxSizeMB()} for file sources. + *
    4. Re-split the splits that are much larger than {@code estimatedSize / (2 x parallelism)}, so + * that each reader gets at least {@link #MIN_SPLITS_PER_READER} splits, and as some sources * return splits way larger than requested. *
    5. While the split sizes vary a lot (high coefficient of variation), re-split the larger * splits. @@ -81,11 +82,19 @@ public final class BoundedSourceSplitter { /** * Splits are not re-split below this size by the refinement steps (the initial split is not * affected), as the fixed cost of opening and reading a split would dominate. Small inputs still - * get one split per reader, as they may feed transforms that generate a lot more data: the floor - * is {@code min(MIN_SPLIT_SIZE_BYTES, estimatedSize/parallelism)}. + * get {@link #MIN_SPLITS_PER_READER} splits per reader, as they may feed transforms that generate + * a lot more data: the floor is {@code min(MIN_SPLIT_SIZE_BYTES, estimatedSize / + * (MIN_SPLITS_PER_READER x parallelism))}. */ static final long MIN_SPLIT_SIZE_BYTES = 64 * MEBIBYTE; + /** + * Number of splits each reader should get. With static round-robin assignment, two splits per + * reader let the ordering pair large splits with small ones, while a single split per reader + * passes any size difference straight to the slowest reader. + */ + static final int MIN_SPLITS_PER_READER = 2; + /** How bounded source splits are handed out to readers. */ public enum Assignment { /** Splits are assigned to readers up front, round-robin. */ @@ -148,17 +157,20 @@ public static List> split( throws Exception { int readers = Math.max(1, parallelism); long maxSplitSizeBytes = maxSplitSizeBytes(source, options); - long minSplitSizeBytes = - Math.max(1L, Math.min(MIN_SPLIT_SIZE_BYTES, estimatedSizeBytes / readers)); + long perReaderTargetBytes = + Math.max(1L, estimatedSizeBytes / ((long) readers * MIN_SPLITS_PER_READER)); + long minSplitSizeBytes = Math.min(MIN_SPLIT_SIZE_BYTES, perReaderTargetBytes); long desiredSizeBytes = Math.min(Math.max(1L, estimatedSizeBytes / Math.max(1, numSplits)), maxSplitSizeBytes); List> splits = sized(source.split(desiredSizeBytes, options), options); int initialSplits = splits.size(); - // Some sources (e.g. BigQuery) return splits way larger than requested that can still be - // split further. - splits = resplitLargerThan(splits, options, Math.max(desiredSizeBytes, minSplitSizeBytes)); + // Aim for MIN_SPLITS_PER_READER splits per reader. This also handles sources (e.g. BigQuery) + // that return splits way larger than requested that can still be split further. + long targetSizeBytes = + Math.max(Math.min(desiredSizeBytes, perReaderTargetBytes), minSplitSizeBytes); + splits = resplitLargerThan(splits, options, targetSizeBytes); if (allSizesKnown(splits)) { long minTargetBytes = Math.max( diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java index dd187c8fc0ec..23996af22f73 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -65,10 +65,11 @@ public void testResplitsWhenTooFewSplits() throws Exception { List> splits = BoundedSourceSplitter.split(source, options, 8, 8000); - assertEquals(8, splits.size()); + // Two splits per reader. + assertEquals(16, splits.size()); assertEquals(8000, totalSize(splits)); for (BoundedSource split : splits) { - assertEquals(1000, split.getEstimatedSizeBytes(options)); + assertEquals(500, split.getEstimatedSizeBytes(options)); } } @@ -92,12 +93,14 @@ public void testResplitsLargeSplits() throws Exception { } @Test - public void testRoundsSplitCountToMultipleOfParallelism() throws Exception { - // 5 splits for 4 readers: halving the largest splits ends at 8 splits. + public void testHalvesSplitsUntilSplitCountIsBalanced() throws Exception { + // Each 1 GiB file is re-split to the 640 MiB target (2 splits per reader): 10 splits, so the + // busiest of the 4 readers gets 3 splits for an average of 2.5. Halving one 640 MiB split gives + // 11 splits: 3 for an average of 2.75. List> splits = BoundedSourceSplitter.split(equalFiles(5, GIB), options, 4, 5 * GIB); - assertEquals(8, splits.size()); + assertEquals(11, splits.size()); assertEquals(5 * GIB, totalSize(splits)); } @@ -112,33 +115,35 @@ public void testAcceptsSplitCountWithManySplitsPerReader() throws Exception { @Test public void testHalvesSplitsWithFewSplitsPerReader() throws Exception { - // 11 splits for 10 readers: the busiest reader would get 2 splits for an average of 1.1. - // Halving stops at 19 splits: 2 splits for an average of 1.9. + // Each 1 GiB file is re-split to the ~563 MiB target (2 splits per reader): 22 splits, so the + // busiest of the 10 readers gets 3 splits for an average of 2.2. Halving stops at 28 splits: 3 + // splits for an average of 2.8. List> splits = BoundedSourceSplitter.split(equalFiles(11, GIB), options, 10, 11 * GIB); - assertEquals(19, splits.size()); + assertEquals(28, splits.size()); assertEquals(11 * GIB, totalSize(splits)); } @Test public void testDoesNotHalveBelowMinSplitSize() throws Exception { - // Halving 100 MiB splits would create splits below the 64 MiB minimum. + // The 11 x 100 MiB splits are below the 110 MiB target for 5 readers, but the busiest reader + // gets 3 splits for an average of 2.2. Halving would create splits below the 64 MiB minimum. List> splits = - BoundedSourceSplitter.split(equalFiles(11, 100 * MIB), options, 10, 1100 * MIB); + BoundedSourceSplitter.split(equalFiles(11, 100 * MIB), options, 5, 1100 * MIB); assertEquals(11, splits.size()); } @Test - public void testSmallInputStillGetsOneSplitPerReader() throws Exception { + public void testSmallInputStillGetsTwoSplitsPerReader() throws Exception { // A 10 MiB input is below the minimum split size, but may feed transforms generating a lot // more data, so it is still spread over all readers. BoundedSource source = new CoarseSource(Collections.singletonList(range(0, 10 * MIB))); List> splits = BoundedSourceSplitter.split(source, options, 10, 10 * MIB); - assertEquals(10, splits.size()); + assertEquals(20, splits.size()); } private static BoundedSource equalFiles(int count, long sizeBytes) { From 75e0d8df4a526581cc8aba99cd3d95b9703c1d2e Mon Sep 17 00:00:00 2001 From: jto Date: Mon, 5 Oct 2026 15:51:55 +0200 Subject: [PATCH 09/11] [Flink] Log the effect of each bounded source splitting step Log each re-split, size balancing round, split count balancing and round-robin ordering with the split count, size min/mean/max and coefficient of variation, and the busiest reader's load relative to the average. Also document that fileInputSplitMaxSizeMB is rarely needed now that bounded sources are re-split by default to even out uneven file sizes. Co-Authored-By: Claude Opus 5.5 --- .../runners/flink/FlinkPipelineOptions.java | 5 +- .../runners/flink/FlinkPipelineOptions.java | 5 +- .../wrappers/BoundedSourceSplitter.java | 119 ++++++++++++++++-- 3 files changed, 115 insertions(+), 14 deletions(-) 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 e2ecf2c879c2..5228dc710cdb 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 @@ -360,7 +360,10 @@ public Long create(PipelineOptions options) { void setFlinkConfDir(String confDir); @Description( - "Set the maximum size of input split when data is read from a filesystem. 0 implies no max size.") + "Set the maximum size in MiB of input splits when data is read from a filesystem. 0 implies " + + "no max size. Rarely needed: bounded sources are re-split by default so that uneven " + + "file sizes are evenly distributed over the readers. Only use it to force splits " + + "smaller than the runner would pick.") @Default.Long(0) Long getFileInputSplitMaxSizeMB(); 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 3b1df33c258b..d7ec2659e965 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 @@ -377,7 +377,10 @@ public Long create(PipelineOptions options) { void setFlinkConfDir(String confDir); @Description( - "Set the maximum size of input split when data is read from a filesystem. 0 implies no max size.") + "Set the maximum size in MiB of input splits when data is read from a filesystem. 0 implies " + + "no max size. Rarely needed: bounded sources are re-split by default so that uneven " + + "file sizes are evenly distributed over the readers. Only use it to force splits " + + "smaller than the runner would pick.") @Default.Long(0) Long getFileInputSplitMaxSizeMB(); diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index cb30056787ff..7aaac51b433c 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -18,9 +18,11 @@ package org.apache.beam.runners.flink.translation.wrappers; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.Comparator; import java.util.List; +import java.util.LongSummaryStatistics; import java.util.PriorityQueue; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -190,7 +192,7 @@ public static List> split( estimatedSizeBytes, desiredSizeBytes, readers, - allSizesKnown(splits) ? coefficientOfVariation(splits) : "unknown"); + allSizesKnown(splits) ? format(coefficientOfVariation(splits)) : "unknown"); List> result = new ArrayList<>(splits.size()); for (SizedSource split : splits) { @@ -218,14 +220,42 @@ private static List> balanceSizes( long minTargetBytes, long maxSplitSizeBytes, int roundsLeft) { - if (roundsLeft == 0 || coefficientOfVariation(splits) <= MAX_SIZE_COEFFICIENT_OF_VARIATION) { + int round = MAX_RESPLIT_ROUNDS - roundsLeft + 1; + double cv = coefficientOfVariation(splits); + if (cv <= MAX_SIZE_COEFFICIENT_OF_VARIATION) { + LOG.info( + "Size balancing round {}: split size coefficient of variation {} is at most {}, done: {}", + round, + format(cv), + MAX_SIZE_COEFFICIENT_OF_VARIATION, + describe(splits)); + return splits; + } + if (roundsLeft == 0) { + LOG.info( + "Size balancing: stopping after {} rounds with split size coefficient of variation {}: " + + "{}", + MAX_RESPLIT_ROUNDS, + format(cv), + describe(splits)); return splits; } long targetBytes = Math.min(Math.max((long) mean(splits), minTargetBytes), maxSplitSizeBytes); + LOG.info( + "Size balancing round {}: split size coefficient of variation {} is above {}, re-splitting " + + "towards {} bytes", + round, + format(cv), + MAX_SIZE_COEFFICIENT_OF_VARIATION, + targetBytes); List> next = resplitLargerThan(splits, options, targetBytes); if (next.size() > splits.size() && allSizesKnown(next)) { return balanceSizes(next, options, minTargetBytes, maxSplitSizeBytes, roundsLeft - 1); } else { + LOG.info( + "Size balancing round {}: no progress (or unknown sizes), keeping: {}", + round, + describe(splits)); return splits; } } @@ -233,16 +263,28 @@ private static List> balanceSizes( /** Re-splits, once, each split larger than {@link #OVERSIZED_SPLIT_FACTOR} x the target. */ private static List> resplitLargerThan( List> splits, PipelineOptions options, long targetBytes) { - return splits.stream() - .flatMap( - split -> { - if (split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes) { - return resplit(split, options, targetBytes).stream(); - } else { - return Stream.of(split); - } - }) - .collect(Collectors.toList()); + long thresholdBytes = (long) (OVERSIZED_SPLIT_FACTOR * targetBytes); + long oversized = splits.stream().filter(split -> split.sizeBytes > thresholdBytes).count(); + List> result = + splits.stream() + .flatMap( + split -> { + if (split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes) { + return resplit(split, options, targetBytes).stream(); + } else { + return Stream.of(split); + } + }) + .collect(Collectors.toList()); + LOG.info( + "Re-split {} of {} splits larger than {} bytes towards {} bytes: {} -> {}", + oversized, + splits.size(), + thresholdBytes, + targetBytes, + describe(splits), + describe(result)); + return result; } /** @@ -270,6 +312,15 @@ private static List> balanceSplitCount( } List> result = new ArrayList<>(candidates); result.addAll(unsplittable); + LOG.info( + "Split count balancing for {} readers: {} -> {} splits ({} could not be halved), busiest " + + "reader gets {} splits for an average of {}", + readers, + splits.size(), + result.size(), + unsplittable.size(), + ((long) result.size() + readers - 1) / readers, + format((double) result.size() / readers)); return result; } @@ -321,9 +372,53 @@ private static List> orderForRoundRobin( } ordered.addAll(row); } + LOG.info( + "Ordered {} splits for round-robin over {} readers: busiest reader load / average went " + + "from {} to {}", + ordered.size(), + readers, + format(busiestReaderLoad(splits, readers)), + format(busiestReaderLoad(ordered, readers))); return ordered; } + /** Bytes of the busiest reader relative to the average, with split i on reader i % readers. */ + private static double busiestReaderLoad(List> splits, int readers) { + long[] loads = new long[readers]; + long total = 0; + int index = 0; + for (SizedSource split : splits) { + loads[index++ % readers] += split.sizeBytes; + total += split.sizeBytes; + } + if (total <= 0) { + return 1.0; + } + return Arrays.stream(loads).max().getAsLong() / ((double) total / readers); + } + + private static String format(double value) { + return String.format("%.3f", value); + } + + /** Short description of the splits for logging. */ + private static String describe(List> splits) { + if (!allSizesKnown(splits)) { + return splits.size() + " splits (sizes unknown)"; + } + LongSummaryStatistics sizes = new LongSummaryStatistics(); + for (SizedSource split : splits) { + sizes.accept(split.sizeBytes); + } + return String.format( + "%d splits, size min/mean/max %d/%.0f/%d bytes, coefficient of variation %.3f", + sizes.getCount(), + sizes.getMin(), + sizes.getAverage(), + sizes.getMax(), + coefficientOfVariation(splits)); + } + private static List> resplit( SizedSource split, PipelineOptions options, long desiredSizeBytes) { try { From b1a4988f62e8130e46d2090983c4a57ecae2dbc4 Mon Sep 17 00:00:00 2001 From: jto Date: Tue, 6 Oct 2026 10:42:50 +0200 Subject: [PATCH 10/11] [Flink] Balance bounded source split count by bytes per reader An uneven split count is not a problem when split sizes differ: the round-robin ordering gives the readers with more splits the smaller ones. Skip halving when the ordered splits already give the busiest reader at most 10% more bytes than the average, instead of judging by split count. On a real BigQuery export read, this avoids thousands of pointless halve attempts reporting a misleading 1.43 count imbalance for a 1.017 byte imbalance. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 38 ++++++++++++++----- .../wrappers/BoundedSourceSplitterTest.java | 33 +++++++++++++--- 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index 7aaac51b433c..843a4c33ff30 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -294,6 +294,18 @@ private static List> resplitLargerThan( */ private static List> balanceSplitCount( List> splits, PipelineOptions options, int readers, long minSplitSizeBytes) { + // What matters is the bytes per reader. When splits differ in size, an uneven split count can + // still be balanced: the round-robin ordering gives the readers with more splits smaller ones. + double load = busiestReaderLoad(snakeOrder(splits, readers), readers); + if (load <= 1 + SPLIT_COUNT_IMBALANCE_TOLERANCE) { + LOG.info( + "Split count balancing for {} readers: skipped, {} splits already give the busiest " + + "reader {} times the average bytes", + readers, + splits.size(), + format(load)); + return splits; + } PriorityQueue> candidates = new PriorityQueue<>(bySizeDescending()); candidates.addAll(splits); List> unsplittable = new ArrayList<>(); @@ -313,9 +325,11 @@ private static List> balanceSplitCount( List> result = new ArrayList<>(candidates); result.addAll(unsplittable); LOG.info( - "Split count balancing for {} readers: {} -> {} splits ({} could not be halved), busiest " - + "reader gets {} splits for an average of {}", + "Split count balancing for {} readers: busiest reader had {} times the average bytes, " + + "{} -> {} splits ({} could not be halved), busiest reader gets {} splits for an " + + "average of {}", readers, + format(load), splits.size(), result.size(), unsplittable.size(), @@ -361,6 +375,19 @@ private static Comparator> bySizeDescending() { */ private static List> orderForRoundRobin( List> splits, int readers) { + List> ordered = snakeOrder(splits, readers); + LOG.info( + "Ordered {} splits for round-robin over {} readers: busiest reader load / average went " + + "from {} to {}", + ordered.size(), + readers, + format(busiestReaderLoad(splits, readers)), + format(busiestReaderLoad(ordered, readers))); + return ordered; + } + + /** Sorts splits by decreasing size in snake order, see {@link #orderForRoundRobin}. */ + private static List> snakeOrder(List> splits, int readers) { List> sorted = new ArrayList<>(splits); sorted.sort(bySizeDescending()); List> ordered = new ArrayList<>(sorted.size()); @@ -372,13 +399,6 @@ private static List> orderForRoundRobin( } ordered.addAll(row); } - LOG.info( - "Ordered {} splits for round-robin over {} readers: busiest reader load / average went " - + "from {} to {}", - ordered.size(), - readers, - format(busiestReaderLoad(splits, readers)), - format(busiestReaderLoad(ordered, readers))); return ordered; } diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java index 23996af22f73..3ad764e1d1ba 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -93,15 +93,16 @@ public void testResplitsLargeSplits() throws Exception { } @Test - public void testHalvesSplitsUntilSplitCountIsBalanced() throws Exception { - // Each 1 GiB file is re-split to the 640 MiB target (2 splits per reader): 10 splits, so the - // busiest of the 4 readers gets 3 splits for an average of 2.5. Halving one 640 MiB split gives - // 11 splits: 3 for an average of 2.75. + public void testResplitsFilesToTwoSplitsPerReader() throws Exception { + // Each 1 GiB file is re-split to the 640 MiB target (2 splits per reader): 10 splits of 640 and + // 384 MiB. The busiest of the 4 readers gets 1408 MiB for an average of 1280 MiB, within 10%, + // so no split is halved. List> splits = BoundedSourceSplitter.split(equalFiles(5, GIB), options, 4, 5 * GIB); - assertEquals(11, splits.size()); + assertEquals(10, splits.size()); assertEquals(5 * GIB, totalSize(splits)); + assertEquals(1408 * MIB, Arrays.stream(perReaderSizes(splits, 4)).max().getAsLong()); } @Test @@ -125,6 +126,28 @@ public void testHalvesSplitsWithFewSplitsPerReader() throws Exception { assertEquals(11 * GIB, totalSize(splits)); } + @Test + public void testSkipsSplitCountBalancingWhenBytesAreBalanced() throws Exception { + // Like 4 files re-split into 10 full pieces and a smaller remainder each. 44 splits for 21 + // readers is uneven by count (3 splits for an average of 2.1), but the round-robin ordering + // gives the readers with 3 splits the small remainders, so bytes are balanced. + List pieces = new ArrayList<>(); + long offset = 0; + for (int i = 0; i < 44; i++) { + long size = (i % 11 == 10 ? 845 : 1790) * MIB; + pieces.add(range(offset, offset + size)); + offset += size; + } + BoundedSource source = new CoarseSource(pieces); + + List> splits = BoundedSourceSplitter.split(source, options, 21, offset); + + assertEquals(44, splits.size()); + long[] perReader = perReaderSizes(splits, 21); + long max = Arrays.stream(perReader).max().getAsLong(); + assertTrue("Unbalanced readers: " + Arrays.toString(perReader), max <= 1.1 * offset / 21); + } + @Test public void testDoesNotHalveBelowMinSplitSize() throws Exception { // The 11 x 100 MiB splits are below the 110 MiB target for 5 readers, but the busiest reader From b9bb0c5fef0c7fa196e325888d4e689e6983e18b Mon Sep 17 00:00:00 2001 From: jto Date: Tue, 6 Oct 2026 10:45:59 +0200 Subject: [PATCH 11/11] [Flink] Log bounded source split sizes in human readable units Format sizes in the splitter logs with binary units and two decimals (e.g. 1.71 MiB, 12.51 GiB) instead of raw byte counts. Co-Authored-By: Claude Opus 5.5 --- .../wrappers/BoundedSourceSplitter.java | 42 +++++++++++++------ .../wrappers/BoundedSourceSplitterTest.java | 10 +++++ 2 files changed, 39 insertions(+), 13 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java index 843a4c33ff30..6dd23d060b57 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -184,13 +184,13 @@ public static List> split( } LOG.info( - "Split bounded source {} in {} splits (initially {}; estimated size {} bytes, desired " - + "split size {} bytes, parallelism {}, split size coefficient of variation {})", + "Split bounded source {} in {} splits (initially {}; estimated size {}, desired split size " + + "{}, parallelism {}, split size coefficient of variation {})", source, splits.size(), initialSplits, - estimatedSizeBytes, - desiredSizeBytes, + formatSize(estimatedSizeBytes), + formatSize(desiredSizeBytes), readers, allSizesKnown(splits) ? format(coefficientOfVariation(splits)) : "unknown"); @@ -243,11 +243,11 @@ private static List> balanceSizes( long targetBytes = Math.min(Math.max((long) mean(splits), minTargetBytes), maxSplitSizeBytes); LOG.info( "Size balancing round {}: split size coefficient of variation {} is above {}, re-splitting " - + "towards {} bytes", + + "towards {}", round, format(cv), MAX_SIZE_COEFFICIENT_OF_VARIATION, - targetBytes); + formatSize(targetBytes)); List> next = resplitLargerThan(splits, options, targetBytes); if (next.size() > splits.size() && allSizesKnown(next)) { return balanceSizes(next, options, minTargetBytes, maxSplitSizeBytes, roundsLeft - 1); @@ -277,11 +277,11 @@ private static List> resplitLargerThan( }) .collect(Collectors.toList()); LOG.info( - "Re-split {} of {} splits larger than {} bytes towards {} bytes: {} -> {}", + "Re-split {} of {} splits larger than {} towards {}: {} -> {}", oversized, splits.size(), - thresholdBytes, - targetBytes, + formatSize(thresholdBytes), + formatSize(targetBytes), describe(splits), describe(result)); return result; @@ -421,6 +421,22 @@ private static String format(double value) { return String.format("%.3f", value); } + private static final String[] SIZE_UNITS = {"B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB"}; + + /** Formats a size in bytes with a binary unit, e.g. {@code 1.71 MiB}. */ + static String formatSize(long bytes) { + if (bytes < 1024) { + return bytes + " B"; + } + double value = bytes; + int unit = 0; + while (value >= 1024 && unit < SIZE_UNITS.length - 1) { + value /= 1024; + unit++; + } + return String.format("%.2f %s", value, SIZE_UNITS[unit]); + } + /** Short description of the splits for logging. */ private static String describe(List> splits) { if (!allSizesKnown(splits)) { @@ -431,11 +447,11 @@ private static String describe(List> splits) { sizes.accept(split.sizeBytes); } return String.format( - "%d splits, size min/mean/max %d/%.0f/%d bytes, coefficient of variation %.3f", + "%d splits, size min/mean/max %s/%s/%s, coefficient of variation %.3f", sizes.getCount(), - sizes.getMin(), - sizes.getAverage(), - sizes.getMax(), + formatSize(sizes.getMin()), + formatSize((long) sizes.getAverage()), + formatSize(sizes.getMax()), coefficientOfVariation(splits)); } diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java index 3ad764e1d1ba..596eeb038c5e 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -203,6 +203,16 @@ public void testKeepsUnsplittableSources() throws Exception { assertEquals(2, splits.size()); } + @Test + public void testFormatsSizesWithBinaryUnits() { + assertEquals("0 B", BoundedSourceSplitter.formatSize(0)); + assertEquals("1023 B", BoundedSourceSplitter.formatSize(1023)); + assertEquals("1.00 KiB", BoundedSourceSplitter.formatSize(1024)); + assertEquals("1.71 MiB", BoundedSourceSplitter.formatSize(1_790_807)); + assertEquals("12.51 GiB", BoundedSourceSplitter.formatSize(13_431_054_034L)); + assertEquals("8.00 EiB", BoundedSourceSplitter.formatSize(Long.MAX_VALUE)); + } + private static RangeSource range(long start, long end) { return new RangeSource(start, end); }