Repository navigation
Conversation
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 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #40608 +/- ##
============================================
+ Coverage 56.16% 65.09% +8.92%
- Complexity 2288 14031 +11743
============================================
Files 1124 1730 +606
Lines 178336 108698 -69638
Branches 1489 11338 +9849
============================================
- Hits 100170 70757 -29413
+ Misses 75645 33720 -41925
- Partials 2521 4221 +1700
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Assigning reviewers: R: @Abacn for label website. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
This PR introduces BoundedSourceSplitter to re-split and balance BoundedSource splits across readers in the Flink runner, and adds configurable static, lazy, and size-based split assignment strategies across DataSet (SourceInputFormat) and DataStream (FlinkSource). The overall design, backward-compatible checkpoint state serialization, and test suite are well-structured, though there are a few important edge cases around unknown source size estimates, fileInputSplitMaxSizeMB interaction with the 64 MiB split floor, and split ordering in SourceInputFormat lazy mode.
🔍 General Feedback
- Unknown top-level size estimates (
<= 0orLong.MAX_VALUE):BoundedSourceSplitter.split()computes target and minimum split sizes directly fromestimatedSizeByteswithout checking whether the estimate is unknown (<= 0orLong.MAX_VALUE), which can collapse the 64 MiB floor to 1 byte or inflate split targets to Exabytes. - Split ordering for lazy assignment in DataSet: Unlike
LazyFlinkSourceSplitEnumerator(which was updated to pop from the front of the list), Flink'sDefaultInputSplitAssignerinSourceInputFormatpops from the end of the array, handing out the smallest splits first and leaving the largest splits for the end.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
| int readers = Math.max(1, parallelism); | ||
| long maxSplitSizeBytes = maxSplitSizeBytes(source, options); | ||
| long perReaderTargetBytes = | ||
| Math.max(1L, estimatedSizeBytes / ((long) readers * MIN_SPLITS_PER_READER)); |
There was a problem hiding this comment.
🔴 Critical: When estimatedSizeBytes is unknown or zero (<= 0 or Long.MAX_VALUE—which both SourceInputFormat.useLazyAssignment and SizeBasedFlinkSourceSplitEnumerator.selectAssignmentMode explicitly check for), using estimatedSizeBytes directly here causes pathological behavior:
- When
estimatedSizeBytes <= 0:perReaderTargetBytes,minSplitSizeBytes,desiredSizeBytes, andtargetSizeBytesall evaluate to1L(1 byte), bypassingMIN_SPLIT_SIZE_BYTES(64 MiB). If the top-level source reports0or-1before splitting (or throws ingetEstimatedSizeBytes) while its initial child splits report known sizes,resplitLargerThan(splits, options, 1L)andbalanceSplitCountwill attempt to re-split the child sources into 1-byte splits, potentially creating millions of splits and exhausting JobManager memory. - When
estimatedSizeBytes == Long.MAX_VALUE:perReaderTargetBytes,desiredSizeBytes,targetSizeBytes, andminTargetBytesevaluate to multi-Petabyte/Exabyte values (~0.5–8 EiB), preventingresplitLargerThanandbalanceSizesfrom ever re-splitting even when the initial splits have known sizes. Moreover, if any child split itself reportssizeBytes == Long.MAX_VALUE,resplitLargerThanat line 175 runs beforeif (allSizesKnown(splits))at line 176 and treatsLong.MAX_VALUE > thresholdBytesas oversized.
Consider checking whether estimatedSizeBytes is known (estimatedSizeBytes > 0 && estimatedSizeBytes < Long.MAX_VALUE), falling back to the sum of splits sizes when allSizesKnown(splits) is true, and skipping resplitLargerThan when !allSizesKnown(splits).
| 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); |
There was a problem hiding this comment.
🟠 High: minSplitSizeBytes is not capped by maxSplitSizeBytes. The updated documentation for fileInputSplitMaxSizeMB states: "Only use it to force splits smaller than the runner would pick." However, if a user sets fileInputSplitMaxSizeMB below 64 MiB (e.g. 16 MiB) on a large FileBasedSource whose initial split() returns per-file splits, minSplitSizeBytes stays at 64 MiB (MIN_SPLIT_SIZE_BYTES), which forces targetSizeBytes (line 174) and minTargetBytes (line 179) up to 64 MiB and prevents resplitLargerThan and halve from splitting down to maxSplitSizeBytes.
| long minSplitSizeBytes = Math.min(MIN_SPLIT_SIZE_BYTES, perReaderTargetBytes); | |
| long minSplitSizeBytes = Math.min(Math.min(MIN_SPLIT_SIZE_BYTES, perReaderTargetBytes), maxSplitSizeBytes); |
| sourceInputSplits.length, | ||
| initialSource); | ||
| if (lazyAssignment) { | ||
| return new DefaultInputSplitAssigner(sourceInputSplits); |
There was a problem hiding this comment.
🟡 Medium: BoundedSourceSplitter.split() orders splits with the largest splits first and the smallest splits at the end. In LazyFlinkSourceSplitEnumerator, assignNextSplit was updated to call pendingSplits.removeFirst() so the largest splits are assigned first. However, Flink's DefaultInputSplitAssigner pops splits from the end of its internal list (splits.remove(splits.size() - 1)), which means in DataSet lazy mode the smallest splits are handed out first and the largest splits are saved for last, worsening tail latency.
| splits.stream() | ||
| .flatMap( | ||
| split -> { | ||
| if (split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes) { |
There was a problem hiding this comment.
🟢 Low: Line 267 checks split.sizeBytes > thresholdBytes (comparing against the long value), whereas line 272 checks split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes (comparing in double, which can lose precision for large long values or differ from the oversized count logged on line 278). Use thresholdBytes here as well.
| if (split.sizeBytes > OVERSIZED_SPLIT_FACTOR * targetBytes) { | |
| if (split.sizeBytes > thresholdBytes) { |
| value /= 1024; | ||
| unit++; | ||
| } | ||
| return String.format("%.2f %s", value, SIZE_UNITS[unit]); |
There was a problem hiding this comment.
🟢 Low: String.format without an explicit Locale uses the default JVM locale. On locales that use a comma decimal separator (e.g. de_DE or fr_FR), formatSize will produce "1,00 KiB" instead of "1.00 KiB", causing BoundedSourceSplitterTest.testFormatsSizesWithBinaryUnits to fail. Pass Locale.ROOT (here and in format / describe).
| return String.format("%.2f %s", value, SIZE_UNITS[unit]); | |
| return String.format( java.util.Locale.ROOT, "%.2f %s", value, SIZE_UNITS[unit]); |
1 participant
Addresses #40607.
Improves how the Flink runner splits bounded sources in batch pipelines, so that work is evenly distributed over the source readers. The same logic is shared by the DataSet (
SourceInputFormat) and DataStream (FlinkSourceenumerators) translations, in a newBoundedSourceSplitter.Motivation
The runner used to call
BoundedSource.split(estimatedSize / parallelism)once and use whatever came back. Beam does not require sources to split in one pass, and many don't honour the desired size: for example a BigQuery export returns one split per exported file, whatever size is requested. The result was often too few splits for the readers, or splits of very uneven size. Combined with lazy split assignment, a few readers could end up with much more work than the others.What changes
After the initial
split()call, the splitter refines the splits in a few bounded passes:estimatedSize / (2 x parallelism)are split again, so each reader gets at least two splits, and sources returning splits far larger than requested (e.g. BigQuery exports) are cut down.ito readeri % parallelismpairs large splits with small ones.Steps 2 to 4 need per-split size estimates and are skipped when they are not available. Re-splitting never goes below
min(64 MiB, estimatedSize / (2 x parallelism)): small splits mostly add fixed per-split overhead, but small inputs, which may feed transforms that generate a lot more data, are still spread over all readers.fileInputSplitMaxSizeMBis still honoured as a cap.The DataStream enumerators now pass the actual reader count (
currentParallelism()) to the splitter, separately fromnumSplits(which may bemaxParallelism), so ordering and balancing match how splits are actually assigned.Split assignment: static by default in batch
sourceStaticSplitThresholdMbno longer has a default. When unset, batch pipelines use static assignment and streaming pipelines keep lazy assignment. Explicit values behave as before:0= lazy, negative = static, positive = size-based.StaticInputSplitAssigner(splitigoes to taski % parallelism, like the DataStream enumerator), and fall back to Flink'sDefaultInputSplitAssignerwhen lazy assignment is configured.This is a behaviour change for batch pipelines that relied on the previous lazy default. They can restore it with
--sourceStaticSplitThresholdMb=0.Observability
Each step logs its effect at INFO, in human-readable units: split count, size min/mean/max and coefficient of variation, and the busiest reader's load relative to the average. For example:
Results
On a production-sized batch job (45 TiB of Avro, parallelism 3500, DataStream translation), comparing the source stage before and after this change, per subtask:
On a BigQuery export read (714 files for 3750 readers), the export files are now re-split into 7854 splits, so all readers get work, with the busiest reader at 1.017x the average bytes. Before, about 80% of the readers had nothing to read.
Trade-offs
getEstimatedSizeBytes()on every split, andsplit()again on some of them. For file-based sources this is in-memory work once the file pattern has been expanded, but it may be slower for sources whose size estimates are remote calls.Testing
BoundedSourceSplitterTest: re-splitting of oversized splits, size balancing, split count balancing by bytes per reader, the minimum split size, small inputs, round-robin ordering for the actual reader count, unsplittable sources, assignment mode resolution and size formatting.SourceInputFormatTest: static assignment by default and lazy assignment when requested.FlinkSourceSplitEnumeratorTestupdated for the new default.Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README for a list of workflows and how to trigger them.