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..6dd23d060b57 --- /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 == 1) { + 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..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/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 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 bd07d00cd049..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; @@ -156,7 +158,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() { @@ -176,7 +182,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..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