This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 355e8bd48b5 Expose desiredNumUnboundedSourceSplits for Java and Python
portable pipeline (#40345)
355e8bd48b5 is described below
commit 355e8bd48b52fec351fd9067630030789dbdf09f
Author: Yi Hu <[email protected]>
AuthorDate: Thu Oct 1 12:49:05 2026 -0400
Expose desiredNumUnboundedSourceSplits for Java and Python portable
pipeline (#40345)
- Port desiredNumUnboundedSourceSplits currently honored by Dataflow
non-portable
pipeline submission to Java portable (UnboundedSource translated to SDF
wrapper)
and Python SDK
---
.../src/main/java/org/apache/beam/sdk/io/Read.java | 12 +++++++-
.../apache/beam/sdk/options/StreamingOptions.java | 10 ++++++
.../test/java/org/apache/beam/sdk/io/ReadTest.java | 36 ++++++++++++++++++++++
sdks/python/apache_beam/io/unbounded_source.py | 13 ++++++--
.../python/apache_beam/io/unbounded_source_test.py | 21 +++++++++++++
.../python/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 43920abb737..b4b575c881d 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.CheckpointMark.NoopCheckpointMark;
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 class Read {
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<OutputT, CheckpointT> 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 53f5a1870e4..250c871a735 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 47dad6b8565..9ac46a8bc4a 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.CheckpointMark;
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 class ReadTest implements Serializable {
}
}
+ @Test
+ public void testUnboundedSdfSplitRestrictionDesiredNumSplits() throws
Exception {
+ List<Integer> recordedDesiredNumSplits = new ArrayList<>();
+ CustomUnboundedSource source =
+ new CustomUnboundedSource() {
+ @Override
+ public List<? extends UnboundedSource<String, NoOpCheckpointMark>>
split(
+ int desiredNumSplits, PipelineOptions options) {
+ recordedDesiredNumSplits.add(desiredNumSplits);
+ return Collections.singletonList(this);
+ }
+ };
+
+ Read.UnboundedSourceAsSDFWrapperFn<String, NoOpCheckpointMark> wrapperFn =
+ new Read.UnboundedSourceAsSDFWrapperFn<>(null);
+ Read.UnboundedSourceAsSDFWrapperFn.UnboundedSourceRestriction<String,
NoOpCheckpointMark>
+ 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 <T extends Serializable & Consumer<Instant>>
Read.Unbounded<Long> 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 ac5f17ce45a..99e987522f5 100644
--- a/sdks/python/apache_beam/io/unbounded_source.py
+++ b/sdks/python/apache_beam/io/unbounded_source.py
@@ -104,6 +104,8 @@ from apache_beam.coders.coders import TupleCoder
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 @@ class
_UnboundedSourceRestrictionProvider(core.RestrictionProvider):
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 6c6651cd0db..d4646e94de5 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
_UnboundedSourceRestrictionCoder
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 @@ class RestrictionProviderTest(unittest.TestCase):
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 75492344a60..743ab967b04 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -812,6 +812,11 @@ class StreamingOptions(PipelineOptions):
'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):
