From 9f5c0eb59eee0fddbfd45a8764f98513e8ac2b68 Mon Sep 17 00:00:00 2001 From: jto Date: Thu, 15 Oct 2026 11:30:19 +0200 Subject: [PATCH] [Flink] Re-split bounded sources for even work distribution in batch Beam does not require a source to be fully split in one pass, and many sources do not honour the desired split size: a BigQuery export returns one split per exported file whatever size is requested. The Flink runner used the result of a single split() call, which often gave too few splits for the readers, or splits of very uneven size. Share bounded source splitting between the DataSet and DataStream translations in BoundedSourceSplitter. After the initial split it: - re-splits splits much larger than estimatedSize / (2 x parallelism), so each reader gets at least two splits and oversized splits returned by sources such as BigQuery are cut down, - re-splits the larger splits while split sizes vary a lot, - halves the largest splits while round-robin assignment would give the busiest reader more than 10% above the average bytes, - orders splits by size in snake order, so that assigning split i to reader i % parallelism pairs large splits with small ones. Re-splitting never goes below min(64 MiB, estimatedSize / (2 x parallelism)), so small inputs, which may feed transforms generating a lot more data, are still spread over all readers, and fileInputSplitMaxSizeMB is still honoured. The DataStream enumerators pass their current parallelism separately from numSplits (which may be maxParallelism), so ordering matches the actual assignment. Batch pipelines now use static split assignment by default, in both DataSet (new StaticInputSplitAssigner) and DataStream. Lazy assignment is used when requested via sourceStaticSplitThresholdMb=0, and stays the default for streaming pipelines. The lazy enumerator hands out the largest splits first, in constant time. Each splitting step logs its effect: split count, size min/mean/max and coefficient of variation, and the busiest reader's load relative to the average. Co-Authored-By: Claude Opus 5.5 --- CHANGES.md | 2 + .../runners/flink/FlinkPipelineOptions.java | 24 +- .../runners/flink/FlinkPipelineOptions.java | 24 +- .../wrappers/BoundedSourceSplitter.java | 526 ++++++++++++++++++ .../wrappers/SourceInputFormat.java | 88 ++- .../streaming/io/source/FlinkSource.java | 17 +- .../io/source/FlinkSourceSplitEnumerator.java | 6 +- .../io/source/FlinkSourceSplitUtils.java | 40 +- .../LazyFlinkSourceSplitEnumerator.java | 15 +- .../SizeBasedFlinkSourceSplitEnumerator.java | 17 +- .../wrappers/BoundedSourceSplitterTest.java | 320 +++++++++++ .../wrappers/SourceInputFormatTest.java | 40 ++ .../FlinkSourceSplitEnumeratorTest.java | 5 +- .../flink_java_pipeline_options.html | 6 +- .../flink_python_pipeline_options.html | 6 +- 15 files changed, 1031 insertions(+), 105 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/CHANGES.md b/CHANGES.md index f98790bc2e19..09d21c7fcfba 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -31,10 +31,12 @@ ## New Features / Improvements +* (Java) Flink runner: bounded sources in batch are re-split and balanced, so that work is evenly distributed over the source readers even when sources return too few or unevenly sized splits, e.g. BigQuery exports ([#40607](https://github.com/apache/beam/issues/40607)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes +* (Java) Flink runner: batch pipelines now assign source splits to readers statically by default instead of lazily. Set `--sourceStaticSplitThresholdMb=0` to restore lazy assignment ([#40607](https://github.com/apache/beam/issues/40607)). * X behavior was changed ([#X](https://github.com/apache/beam/issues/X)). ## Deprecations 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..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 @@ -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. @@ -359,22 +360,25 @@ 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(); 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..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 @@ -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. @@ -376,22 +377,25 @@ 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(); 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..fbac43ce378b --- /dev/null +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitter.java @@ -0,0 +1,526 @@ +/* + * 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.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; +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 / numSplits} as the desired split size, capped + * by {@link FlinkPipelineOptions#getFileInputSplitMaxSizeMB()} for file sources. + *
  2. 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. + *
  3. While the split sizes vary a lot (high coefficient of variation), re-split the larger + * splits. + *
  4. Halve the largest splits until round-robin assignment gives each reader about the same + * number of splits. + *
  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; + + /** + * 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 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; + + /** + * 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 {@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. */ + 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 { + 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 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(); + + // 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( + minSplitSizeBytes, + estimatedSizeBytes / ((long) readers * MAX_SPLITS_PER_READER_FOR_BALANCING)); + splits = balanceSizes(splits, options, minTargetBytes, maxSplitSizeBytes, MAX_RESPLIT_ROUNDS); + splits = balanceSplitCount(splits, options, readers, minSplitSizeBytes); + splits = orderForRoundRobin(splits, readers); + } + + LOG.info( + "Split bounded source {} in {} splits (initially {}; estimated size {}, desired split size " + + "{}, parallelism {}, split size coefficient of variation {})", + source, + splits.size(), + initialSplits, + formatSize(estimatedSizeBytes), + formatSize(desiredSizeBytes), + readers, + allSizesKnown(splits) ? format(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 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, + long minTargetBytes, + long maxSplitSizeBytes, + int roundsLeft) { + 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 {}", + round, + format(cv), + MAX_SIZE_COEFFICIENT_OF_VARIATION, + formatSize(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; + } + } + + /** Re-splits, once, each split larger than {@link #OVERSIZED_SPLIT_FACTOR} x the target. */ + private static List> resplitLargerThan( + List> splits, PipelineOptions options, long targetBytes) { + 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 {} towards {}: {} -> {}", + oversized, + splits.size(), + formatSize(thresholdBytes), + formatSize(targetBytes), + describe(splits), + describe(result)); + return result; + } + + /** + * 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> 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<>(); + for (int attempt = 0; + attempt < 2 * readers + && !candidates.isEmpty() + && !isBalancedEnough(candidates.size() + unsplittable.size(), readers); + attempt++) { + SizedSource largest = candidates.remove(); + List> halves = halve(largest, options, minSplitSizeBytes); + if (halves.size() > 1) { + candidates.addAll(halves); + } else { + unsplittable.add(largest); + } + } + List> result = new ArrayList<>(candidates); + result.addAll(unsplittable); + LOG.info( + "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(), + ((long) result.size() + readers - 1) / readers, + format((double) result.size() / readers)); + return result; + } + + /** + * 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 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 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); + 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. + */ + 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()); + 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 != 0) { + Collections.reverse(row); + } + ordered.addAll(row); + } + 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); + } + + 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)) { + 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 %s/%s/%s, coefficient of variation %.3f", + sizes.getCount(), + formatSize(sizes.getMin()), + formatSize((long) sizes.getAverage()), + formatSize(sizes.getMax()), + coefficientOfVariation(splits)); + } + + 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..b061c0382cb7 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 int numTasks = 1; + private 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/FlinkSourceSplitEnumerator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/FlinkSourceSplitEnumerator.java index ed20c7fa5259..03b68562b9af 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 @@ -167,7 +167,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 26a5dd4d9349..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 @@ -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() {} @@ -45,20 +42,12 @@ static ArrayList> splitBoundedSource( BoundedSource boundedSource, PipelineOptions pipelineOptions, int numSplits, + int parallelism, 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, parallelism, estimatedSizeBytes)); } static ArrayList> splitUnboundedSource( @@ -70,25 +59,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 f298e6534b03..1424d96b1466 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/io/source/LazyFlinkSourceSplitEnumerator.java @@ -19,7 +19,9 @@ import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; 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; @@ -49,7 +51,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; @@ -72,7 +74,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; @@ -163,7 +165,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() { @@ -183,7 +189,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.removeFirst(); 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..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 @@ -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; @@ -130,18 +131,21 @@ 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); } 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 +160,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..596eeb038c5e --- /dev/null +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/BoundedSourceSplitterTest.java @@ -0,0 +1,320 @@ +/* + * 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.assertArrayEquals; +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 static final long MIB = BoundedSourceSplitter.MEBIBYTE; + private static final long GIB = 1024 * MIB; + + 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); + + // Two splits per reader. + assertEquals(16, splits.size()); + assertEquals(8000, totalSize(splits)); + for (BoundedSource split : splits) { + assertEquals(500, 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)); + 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 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(10, splits.size()); + assertEquals(5 * GIB, totalSize(splits)); + assertEquals(1408 * MIB, Arrays.stream(perReaderSizes(splits, 4)).max().getAsLong()); + } + + @Test + 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); + + assertEquals(39, splits.size()); + } + + @Test + public void testHalvesSplitsWithFewSplitsPerReader() throws Exception { + // 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(28, splits.size()); + 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 + // 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, 5, 1100 * MIB); + + assertEquals(11, splits.size()); + } + + @Test + 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(20, splits.size()); + } + + 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 + 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(1), new UnsplittableSource(1))); + + List> splits = BoundedSourceSplitter.split(source, options, 4, 2); + + 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); + } + + 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 { + private final long sizeBytes; + + UnsplittableSource(long sizeBytes) { + this.sizeBytes = sizeBytes; + } + + @Override + public List split(long desiredSize, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public long getEstimatedSizeBytes(PipelineOptions options) { + return sizeBytes; + } + } +} 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..9540dc0956f9 100644 --- a/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_java_pipeline_options.html @@ -84,7 +84,7 @@ fileInputSplitMaxSizeMB - 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: 0 @@ -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..4bcda7586179 100644 --- a/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html +++ b/website/www/site/layouts/shortcodes/flink_python_pipeline_options.html @@ -84,7 +84,7 @@ file_input_split_max_size_m_b - 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: 0 @@ -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