Expose desiredNumUnboundedSourceSplits for Java and Python portable pipeline

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<OutputT, CheckpointT> split :
-            restriction.getSource().split(DEFAULT_DESIRED_NUM_SPLITS, pipelineOptions)) {
+            restriction.getSource().split(desiredNumSplits, pipelineOptions)) {
           receiver.output(
               UnboundedSourceRestriction.create(split, null, restriction.getWatermark()));
         }

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;

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<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) {

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.',

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 = []
 

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):