[Flink] Select bounded-source split assignment by estimated size by pkuzmickas · Pull Request #39874 · apache/beam

pkuzmickas · 2026-08-24T10:10:08Z

Summary

Add an opt-in, size-based split-assignment strategy for bounded sources in the Flink DataStream runner.

The default remains lazy pull-based assignment. When users configure a positive threshold, the runner assigns sources estimated below that size per reader statically and keeps larger sources lazy. A negative value forces static assignment.

Fixes #39873.

Why

Moving production workloads from the Flink 1 DataSet runner to the Flink 2 DataStream runner caused large performance regressions for sources that emit inexpensive descriptors for expensive downstream work.

The DataStream runner currently sends all bounded sources to the lazy enumerator. That enumerator gives the next split to whichever reader requests it. This balances sources whose splits contain expensive I/O, but fast-starting readers can claim most descriptors before their peers start. A pointwise downstream edge then preserves the skew during the expensive work.

Static assignment fixes descriptor sources but can slow direct file readers, where dynamic work sharing compensates for different file sizes. The option therefore remains disabled by default and lets each pipeline choose a threshold.

Performance

These anonymized production measurements used matching input partitions and isolated outputs.

Workload	Source shape	Comparison	Result
A	Cheap descriptors followed by expensive work	Flink 2 static vs. Flink 1	~45% faster
B	Cheap descriptors followed by expensive work	Flink 2 static vs. Flink 1	~28% faster
C	Variable-size files read in the source	Threshold-selected lazy vs. forced static	~92 min vs. ~115 min (~20% faster)

The measured descriptor and file-reader examples were approximately 4 GiB and 8 GiB per reader, respectively. This supports a threshold near 6 GiB for those pipelines, but estimated bytes are only a proxy for split cost; Beam does not select a global threshold.

Configuration

lazySourceSplitAssignmentMinSizeMbPerReader (Python: lazy_source_split_assignment_min_size_mb_per_reader) controls bounded-source assignment:

0 (default): always use the existing lazy assignment.
Positive: use static round-robin assignment below the threshold and lazy assignment at or above it.
Negative: always use static assignment.

Implementation

Estimate size and split the source once on the enumerator's asynchronous worker.
Keep the static and lazy enumerators separate, and use a small size-based selector only when a positive threshold is configured.
Queue reader requests while size-based selection is still running, then replay them to the selected enumerator.
Preserve the existing Beam source bundle-size calculation; the new option changes assignment only.
Persist the selected mode and a neutral list of pending splits in version 1 checkpoint state.
Restore without estimating again; repartition pending static splits by split index when parallelism changes.
Upgrade version 0 map checkpoints explicitly: bounded sources resume lazily, and unbounded sources resume statically.
Document the Java and Python options.

Compatibility

Existing pipelines keep lazy bounded-source assignment because the option defaults to 0. Existing checkpoints retain their previous assignment behavior. Estimation failures continue to fail the job, matching the current runner.

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in the description.
Update CHANGES.md with noteworthy changes.
Confirm the Apache ICLA before marking this draft ready.

See the Contributor Guide for more tips on making the review process smoother.

pkuzmickas · 2026-08-24T12:57:16Z

Run Java PreCommit

github-actions · 2026-08-24T14:18:38Z

Assigning reviewers:

R: @shunping for label website.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

pkuzmickas · 2026-08-31T10:22:10Z

Hey @Abacn thank you for the super helpful review of the other PR. I'm working on testing your proposed improvements.

Could you also please review this when you get a chance? 🙏

Abacn

Thanks! Since it's tested working the approach sounds good to me. Had a few initial comments mainly on code base simplifications.

Abacn · 2026-08-31T19:37:09Z

+    }
+    long estimatedBytesPerReader = estimatedSizeBytes / sourceParallelism;
+    long thresholdBytes = FlinkSourceSplitUtils.mebibytesToBytes(thresholdMb);
+    FlinkSourceSplitAssignmentMode selectedMode =


Reader may return estimatedBytes as 0 if it does not support estimation. This is common in both Beam builtin and user IO Connectors. To make the option truely opt-in I understand the condition should read (thresholdBytes <=0 || estimatedBytesPerReader >= thresholdBytes)

Great catch! Thanks.

Abacn · 2026-08-31T19:42:16Z

+          + "and lazy assignment for sources at or above it. Any negative value always uses "
+          + "static assignment.")
+  @Default.Long(0)
+  Long getLazySourceSplitAssignmentMinSizeMbPerReader();


A few comments

We can leave it opt-in for now, however ideally the default option should be good enough for most common use cases, provide a reasonably performant outcome. I understand now it's expected to set this to some positive values to benefit from it.

The naming is pretty long. In addition, I don't see "lazySplit" as a term anywhere in documentation. Not sure how it ended up with this naming.

Consider "getSourceStaticSplitThresholdMb" and just note (1) positive sizes below this threshold uses static split (2) it's per reader in the description string (3) note in description this is for DataStream mode only.

Makes sense, using the name you suggested and modified the description 👍

We can leave it opt-in for now, however ideally the default option should be good enough for most common use cases, provide a reasonably performant outcome

We tested several thresholds against different production workflows, but did not find one that performed well enough across them to make a good general default unfortunately at this time :(

Abacn · 2026-08-31T19:49:02Z

+  private static final int SOURCE_PARALLELISM = 2;
+  private static final int REQUESTED_SPLITS = 4;
+
+  @Test


Thanks for adding thorough test coverage! However, adding 600+ lines of test code introduces maintenance overhead and duplication. We can cut ~300+ lines while keeping 100% of the effective test coverage by simplifying:

Consolidate threshold decision tests: testSmallBoundedSource..., testLargeBoundedSource..., testUnknownEstimates..., testDefaultConfiguration..., testEmptyBoundedSource..., and testConfigurationCanForce... all test selectAssignmentMode. These 6 methods can be collapsed into a single parameterized or table-driven test using assertAssignmentMode.

Deduplicate restore tests: testStaticRestoreReturnsPendingSplitsToOriginalOwnersAtSameParallelism is redundant with testRestoreRepartitionsStaticSplitsForNewParallelismWithoutEstimating (rescale already covers the same-parallelism case).

Remove out-of-scope test & mock: testStaticAssignmentRespectsFileInputSplitMaxSize and TestFileBasedSource test fileInputSplitMaxSizeMB, which is an existing option not introduced by this PR.

Deduplicate early reader registration tests: testSignalsNoMoreSplitsToEarlyReaderWithoutAssignment and testSizeBasedSelectionAssignsStaticSplitsToEarlyReaders test the same registration race condition.

Trim custom mock classes: Removing TestFileBasedSource and simplifying TestEstimatedSizeBoundedSource (e.g. dropping failEstimation/createFailing) to eliminate mock boilerplate.

Abacn · 2026-08-31T19:52:57Z

 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
+package org.apache.beam.runners.flink.translation.wrappers.streaming.io.source;


THis thin enum can be put as an internal class of FlinkSourceEnumeratorState

Abacn · 2026-08-31T19:54:44Z

+import java.util.List;
+
+/** Checkpoint state shared by the lazy and static source split assignment strategies. */
+public final class FlinkSourceEnumeratorState<T> implements Serializable {


Since there is a dedicate FlinkSourceEnumeratorStateSerializer, why still need implements java.io.Serializable? If it's needed it suggests the Serializer isn't effective in some serialization paths

FlinkSourceEnumeratorStateSerializer handles checkpoint-format versioning and upgrades the previous map-based state. For the new version, the checkpoint data (the assignment mode and list of pending splits) still uses the existing Java object serialization, so FlinkSourceEnumeratorState needs to implement Serializable, as the previous state did. I avoided adding field-by-field encoding to the new serializer for simplicity. Please let me know if you had another approach in mind.

Abacn · 2026-08-31T19:56:43Z

+import org.apache.beam.sdk.options.PipelineOptions;
+
+/** Shared Beam source sizing and splitting helpers. */
+final class FlinkSourceSplitUtils {


It's only used by FlinkSourceSplitEnumerator. Consider just put helper methods there.

In general XxxUtils classes are discouraged for specific helper methods (admittedly it exists everywhere in code base)

These helpers are also used by LazyFlinkSourceSplitEnumerator and SizeBasedFlinkSourceSplitEnumerator. Moving them into FlinkSourceSplitEnumerator would make the lazy and size based implementations depend on the static enumerator for shared source splitting logic. I kept them in a neutral helper for that reason, but I’m happy to use a different approach if you prefer.

pkuzmickas

Made most of the changes you asked for with a few comments. Thanks!! 🙏

pkuzmickas · 2026-09-01T12:56:59Z

+import java.util.List;
+
+/** Checkpoint state shared by the lazy and static source split assignment strategies. */
+public final class FlinkSourceEnumeratorState<T> implements Serializable {


FlinkSourceEnumeratorStateSerializer handles checkpoint-format versioning and upgrades the previous map-based state. For the new version, the checkpoint data (the assignment mode and list of pending splits) still uses the existing Java object serialization, so FlinkSourceEnumeratorState needs to implement Serializable, as the previous state did. I avoided adding field-by-field encoding to the new serializer for simplicity. Please let me know if you had another approach in mind.

pkuzmickas · 2026-09-01T13:08:21Z

+import org.apache.beam.sdk.options.PipelineOptions;
+
+/** Shared Beam source sizing and splitting helpers. */
+final class FlinkSourceSplitUtils {


These helpers are also used by LazyFlinkSourceSplitEnumerator and SizeBasedFlinkSourceSplitEnumerator. Moving them into FlinkSourceSplitEnumerator would make the lazy and size based implementations depend on the static enumerator for shared source splitting logic. I kept them in a neutral helper for that reason, but I’m happy to use a different approach if you prefer.

pkuzmickas · 2026-09-01T13:12:50Z

+    }
+    long estimatedBytesPerReader = estimatedSizeBytes / sourceParallelism;
+    long thresholdBytes = FlinkSourceSplitUtils.mebibytesToBytes(thresholdMb);
+    FlinkSourceSplitAssignmentMode selectedMode =


Great catch! Thanks.

pkuzmickas · 2026-09-01T13:18:31Z

+          + "and lazy assignment for sources at or above it. Any negative value always uses "
+          + "static assignment.")
+  @Default.Long(0)
+  Long getLazySourceSplitAssignmentMinSizeMbPerReader();


Makes sense, using the name you suggested and modified the description 👍

We can leave it opt-in for now, however ideally the default option should be good enough for most common use cases, provide a reasonably performant outcome

We tested several thresholds against different production workflows, but did not find one that performed well enough across them to make a good general default unfortunately at this time :(

github-actions Bot added website runners flink labels Aug 24, 2026

[Flink] Select bounded-source split assignment by size

f110427

pkuzmickas force-pushed the pkuzmickas/flink-source-assignment-by-size branch from f362fb0 to f110427 Compare August 24, 2026 10:47

pkuzmickas marked this pull request as ready for review August 24, 2026 13:15

github-actions Bot added the Next Action: Reviewers label Aug 24, 2026

Abacn reviewed Aug 31, 2026

View reviewed changes

[Flink] Refine bounded-source split assignment

2678c10

pkuzmickas commented Sep 1, 2026

View reviewed changes

pkuzmickas requested a review from Abacn September 1, 2026 13:27

Merge apache/beam master into Flink source assignment branch

0d2ec58

Uh oh!

Conversation

pkuzmickas commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why

Performance

Configuration

Implementation

Compatibility

Uh oh!

pkuzmickas commented Aug 24, 2026

Uh oh!

github-actions Bot commented Aug 24, 2026

Uh oh!

pkuzmickas commented Aug 31, 2026

Uh oh!

Abacn left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pkuzmickas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pkuzmickas commented Aug 24, 2026 •

edited

Loading