jto commented on code in PR #40608: URL: https://github.com/apache/beam/pull/40608#discussion_r4481954549
########## 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. + * + * <p>Beam does not require a source to be fully split in one pass, so the initial splits are + * refined in several bounded passes: + * + * <ol> + * <li>Split the source using {@code estimatedSize / numSplits} as the desired split size, capped + * by {@link FlinkPipelineOptions#getFileInputSplitMaxSizeMB()} for file sources. + * <li>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. + * <li>While the split sizes vary a lot (high coefficient of variation), re-split the larger + * splits. + * <li>Halve the largest splits until round-robin assignment gives each reader about the same + * number of splits. + * <li>Order the splits so that round-robin assignment gives each reader a similar amount of data. + * </ol> + * + * <p>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 <T> List<BoundedSource<T>> split( + BoundedSource<T> 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 <T> List<BoundedSource<T>> split( + BoundedSource<T> 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); Review Comment: Thanks, but I don't think this can happen in practice. `minSplitSizeBytes` is only used as a floor for deciding which splits to re-split or halve further; it never makes a split larger. Every step that sets a split size is already capped by `maxSplitSizeBytes`: - the initial `source.split(desiredSizeBytes, ...)` request uses `desiredSizeBytes = min(..., maxSplitSizeBytes)`, and file sources split at the requested size; - `resplitLargerThan` only splits further; - `balanceSizes` caps its target at `maxSplitSizeBytes`; - `halve` only makes splits smaller. I tried a test with a 1 GiB file and `fileInputSplitMaxSizeMB=16`: it already produces 64 splits of at most 16 MiB without capping the minimum. The only way to get splits over the max is a file source whose min bundle size exceeds it, and capping the minimum would not change that either. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
