From 43e878dd4343590c26ee8fd5069c853e1ddfd765 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 29 Sep 2026 15:42:44 -0400 Subject: [PATCH] Expose desiredNumUnboundedSourceSplits for Java and Python portable pipeline - Port desiredNumUnboundedSourceSplits currently honored by Dataflow non-portable pipeline submission to Java portable (UnboundedSource translated to SDF wrapper) and Python SDK --- .../java/org/apache/beam/sdk/io/Read.java | 12 ++++++- .../beam/sdk/options/StreamingOptions.java | 10 ++++++ .../java/org/apache/beam/sdk/io/ReadTest.java | 36 +++++++++++++++++++ .../python/apache_beam/io/unbounded_source.py | 13 +++++-- .../apache_beam/io/unbounded_source_test.py | 21 +++++++++++ .../apache_beam/options/pipeline_options.py | 5 +++ 6 files changed, 94 insertions(+), 3 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/Read.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/Read.java index 43920abb7371..b4b575c881dc 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/Read.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/Read.java @@ -43,6 +43,7 @@ import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.transforms.Deduplicate; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement; @@ -550,9 +551,18 @@ public void splitRestriction( receiver.output(restriction); } + int desiredNumSplits = DEFAULT_DESIRED_NUM_SPLITS; + if (pipelineOptions != null) { + int configuredSplits = + pipelineOptions.as(StreamingOptions.class).getDesiredNumUnboundedSourceSplits(); + if (configuredSplits > 0) { + desiredNumSplits = configuredSplits; + } + } + try { for (UnboundedSource split : - restriction.getSource().split(DEFAULT_DESIRED_NUM_SPLITS, pipelineOptions)) { + restriction.getSource().split(desiredNumSplits, pipelineOptions)) { receiver.output( UnboundedSourceRestriction.create(split, null, restriction.getWatermark())); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java index 53f5a1870e46..250c871a7356 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/StreamingOptions.java @@ -45,6 +45,16 @@ public interface StreamingOptions extends ApplicationNameOptions, PipelineOption void setUpdateCompatibilityVersion(@Nullable String updateCompatibilityVersion); + /** + * The desired number of initial splits for UnboundedSources. If this value is <=0, a default + * number of splits will be chosen. + */ + @Description("The desired number of initial splits for UnboundedSources.") + @Default.Integer(0) + int getDesiredNumUnboundedSourceSplits(); + + void setDesiredNumUnboundedSourceSplits(int value); + static boolean updateCompatibilityVersionLessThan(PipelineOptions options, String version) { if (options == null) { return false; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/ReadTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/ReadTest.java index 47dad6b85650..9ac46a8bc4a0 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/ReadTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/ReadTest.java @@ -49,8 +49,11 @@ import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.testing.NeedsRunner; import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestOutputReceiver; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.testing.UsesUnboundedPCollections; import org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo; @@ -222,6 +225,39 @@ public void testWatermarkAdvanceOnClaimFail() { } } + @Test + public void testUnboundedSdfSplitRestrictionDesiredNumSplits() throws Exception { + List recordedDesiredNumSplits = new ArrayList<>(); + CustomUnboundedSource source = + new CustomUnboundedSource() { + @Override + public List> split( + int desiredNumSplits, PipelineOptions options) { + recordedDesiredNumSplits.add(desiredNumSplits); + return Collections.singletonList(this); + } + }; + + Read.UnboundedSourceAsSDFWrapperFn wrapperFn = + new Read.UnboundedSourceAsSDFWrapperFn<>(null); + Read.UnboundedSourceAsSDFWrapperFn.UnboundedSourceRestriction + restriction = wrapperFn.initialRestriction(source); + TestOutputReceiver< + Read.UnboundedSourceAsSDFWrapperFn.UnboundedSourceRestriction< + String, NoOpCheckpointMark>> + receiver = new TestOutputReceiver<>(); + + PipelineOptions defaultOptions = PipelineOptionsFactory.create(); + wrapperFn.splitRestriction(restriction, receiver, defaultOptions); + assertEquals(Collections.singletonList(20), recordedDesiredNumSplits); + + recordedDesiredNumSplits.clear(); + PipelineOptions customOptions = PipelineOptionsFactory.create(); + customOptions.as(StreamingOptions.class).setDesiredNumUnboundedSourceSplits(5); + wrapperFn.splitRestriction(restriction, receiver, customOptions); + assertEquals(Collections.singletonList(5), recordedDesiredNumSplits); + } + private > Read.Unbounded newUnboundedReadInterceptingWatermark( long numElements, T interceptedWatermarkReceiver) { diff --git a/sdks/python/apache_beam/io/unbounded_source.py b/sdks/python/apache_beam/io/unbounded_source.py index ac5f17ce45a5..99e987522f56 100644 --- a/sdks/python/apache_beam/io/unbounded_source.py +++ b/sdks/python/apache_beam/io/unbounded_source.py @@ -104,6 +104,8 @@ def get_checkpoint_mark_coder(self): from apache_beam.coders.coders import _MemoizingPickleCoder from apache_beam.io import iobase from apache_beam.io.watermark_estimators import ManualWatermarkEstimator +from apache_beam.options.pipeline_options import StreamingOptions +from apache_beam.options.pipeline_options_context import get_pipeline_options from apache_beam.runners import sdf_utils from apache_beam.transforms import PTransform from apache_beam.transforms import core @@ -709,11 +711,18 @@ def split(self, element, yield restriction return + desired_num_splits = _DEFAULT_DESIRED_NUM_SPLITS + options = get_pipeline_options() + if options is not None: + configured_splits = options.view_as( + StreamingOptions).desired_num_unbounded_source_splits + if configured_splits > 0: + desired_num_splits = configured_splits + # ``source.split`` is user code and may refuse to split; fall back to a # single restriction on error. try: - split_sources = list( - restriction.source.split(_DEFAULT_DESIRED_NUM_SPLITS, None)) + split_sources = list(restriction.source.split(desired_num_splits, None)) except Exception: # pylint: disable=broad-except _LOGGER.warning( 'Exception while splitting UnboundedSource. Source not split.', diff --git a/sdks/python/apache_beam/io/unbounded_source_test.py b/sdks/python/apache_beam/io/unbounded_source_test.py index 6c6651cd0db3..d4646e94de59 100644 --- a/sdks/python/apache_beam/io/unbounded_source_test.py +++ b/sdks/python/apache_beam/io/unbounded_source_test.py @@ -47,6 +47,8 @@ from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionProvider from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionTracker from apache_beam.io.watermark_estimators import ManualWatermarkEstimator +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.options.pipeline_options_context import scoped_pipeline_options from apache_beam.runners import sdf_utils from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that @@ -361,6 +363,25 @@ def split(self, desired_num_splits, options=None): self.assertTrue( all(split.finalization_checkpoint_mark is None for split in splits)) + def test_initial_split_uses_desired_num_unbounded_source_splits_option(self): + split_log = [] + + class _SplitSource(UnboundedCountingSource): + @override + def split(self, desired_num_splits, options=None): + split_log.append((desired_num_splits, options)) + return [self] + + source = _SplitSource(5) + provider = _UnboundedSourceRestrictionProvider() + restriction = _UnboundedSourceRestriction(source=source) + options = PipelineOptions(['--desired_num_unbounded_source_splits=5']) + + with scoped_pipeline_options(options): + list(provider.split(source, restriction)) + + self.assertEqual(split_log, [(5, None)]) + def test_initial_split_does_not_split_checkpointed_restriction(self): split_log = [] diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 75492344a604..743ab967b04c 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -812,6 +812,11 @@ def _add_argparse_args(cls, parser): 'version of the Beam SDK. ' 'See for example, https://cloud.google.com/dataflow/docs/guides/' 'updating-a-pipeline') + parser.add_argument( + '--desired_num_unbounded_source_splits', + type=int, + default=0, + help='The desired number of initial splits for UnboundedSources.') class CrossLanguageOptions(PipelineOptions):