[Data] Support strict=False mode for StreamingRepartition

ci/lint/pydoclint-baseline.txt

@@ -1138,10 +1138,6 @@ python/ray/data/_internal/logical/operators/join_operator.py
     DOC101: Method `Join.__init__`: Docstring contains fewer arguments than in function signature.
     DOC103: Method `Join.__init__`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [aggregator_ray_remote_args: Optional[Dict[str, Any]], join_type: str, left_columns_suffix: Optional[str], left_input_op: LogicalOperator, left_key_columns: Tuple[str], num_partitions: int, partition_size_hint: Optional[int], right_columns_suffix: Optional[str], right_input_op: LogicalOperator, right_key_columns: Tuple[str]].
 --------------------
-python/ray/data/_internal/logical/operators/map_operator.py
-    DOC101: Method `StreamingRepartition.__init__`: Docstring contains fewer arguments than in function signature.
-    DOC103: Method `StreamingRepartition.__init__`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [input_op: LogicalOperator].
---------------------
 python/ray/data/_internal/logical/operators/n_ary_operator.py
     DOC001: Method `__init__` Potential formatting errors in docstring. Error message: No specification for "Args": ""
     DOC001: Function/method `__init__`: Potential formatting errors in docstring. Error message: No specification for "Args": "" (Note: DOC001 could trigger other unrelated violations under this function/method too. Please fix the docstring formatting first.)

python/ray/data/_internal/logical/operators/map_operator.py

@@ -425,19 +425,33 @@ def __init__(
 
 class StreamingRepartition(AbstractMap):
     """Logical operator for streaming repartition operation.
+
     Args:
+        input_op: The operator preceding this operator in the plan DAG.
         target_num_rows_per_block: The target number of rows per block granularity for
-           streaming repartition.
+            streaming repartition.
+        strict: If True, guarantees that all output blocks, except for the last one,
+            will have exactly target_num_rows_per_block rows. If False, uses best-effort
+            bundling and may produce at most one block smaller than target_num_rows_per_block
+            per input block without forcing exact sizes through block splitting.
+            Defaults to False.
     """
 
     def __init__(
         self,
         input_op: LogicalOperator,
         target_num_rows_per_block: int,
+        strict: bool = False,
     ):
+        if target_num_rows_per_block <= 0:
+            raise ValueError(
+                "target_num_rows_per_block must be positive for streaming repartition, "
+                f"got {target_num_rows_per_block}"
+            )
         super().__init__(
-            f"StreamingRepartition[num_rows_per_block={target_num_rows_per_block}]",
+            f"StreamingRepartition[num_rows_per_block={target_num_rows_per_block},strict={strict}]",
             input_op,
             can_modify_num_rows=False,
         )
         self.target_num_rows_per_block = target_num_rows_per_block
+        self._strict = strict

python/ray/data/_internal/logical/rules/combine_shuffles.py

@@ -56,9 +56,11 @@ def _combine(self, op: LogicalOperator) -> LogicalOperator:
         elif isinstance(input_op, StreamingRepartition) and isinstance(
             op, StreamingRepartition
         ):
+            strict = input_op._strict or op._strict
             return StreamingRepartition(
                 input_op.input_dependencies[0],
                 target_num_rows_per_block=op.target_num_rows_per_block,
+                strict=strict,
             )
         elif isinstance(input_op, Repartition) and isinstance(op, Aggregate):
             return Aggregate(

python/ray/data/_internal/logical/rules/operator_fusion.py

@@ -278,22 +278,31 @@ def _can_fuse(self, down_op: PhysicalOperator, up_op: PhysicalOperator) -> bool:
 
         # only allow fusion of MapBatches -> StreamingRepartition
         if isinstance(down_logical_op, StreamingRepartition):
-            return (
+            if not (
                 isinstance(up_logical_op, MapBatches)
-                and up_logical_op.batch_size is not None
                 and down_logical_op.target_num_rows_per_block is not None
                 and down_logical_op.target_num_rows_per_block > 0
-                # When the batch_size is a multiple of target_num_rows_per_block, fusing would still produce exactly identical sequence of blocks.
-                # See `_fuse_streaming_repartition_operators_in_dag` docstring for details.
-                # TODO: when the StreamingRepartition supports none_strict_mode, we can fuse
-                # `MapBatches -> StreamingRepartition` no matter what the `batch_size` and `target_num_rows` are.
-                # https://anyscale1.atlassian.net/browse/DATA-1731
+            ):
+                return False
+
+            # Non-strict mode: can always fuse, no matter what batch_size is.
+            # This allows fusion without cross-task buffering by using default bundler.
+            if not down_logical_op._strict:
+                return True
+
+            # Strict mode: only fuse when batch_size is a multiple of target_num_rows_per_block.
+            # When batch_size % target == 0, each batch can be perfectly sliced into chunks
+            # without cross-task buffering. See `_fuse_streaming_repartition_operators_in_dag`
+            # docstring for details.
+            return (
+                up_logical_op.batch_size is not None
                 and up_logical_op.batch_size % down_logical_op.target_num_rows_per_block
                 == 0
             )
-        # Other operators cannot fuse with StreamingRepartition.
+        # StreamingRepartition can only fuse in non-strict mode.
+        # In strict mode, it does not support further fusion.
         if isinstance(up_logical_op, StreamingRepartition):
-            return False
+            return not up_logical_op._strict
 
         # Otherwise, ops are compatible for fusion.
         return True
@@ -312,9 +321,30 @@ def _get_fused_streaming_repartition_operator(
         up_logical_op = self._op_map.pop(up_op)
         assert isinstance(up_logical_op, MapBatches)
         assert isinstance(down_logical_op, StreamingRepartition)
-        assert up_logical_op.batch_size % down_logical_op.target_num_rows_per_block == 0
+
         batch_size = up_logical_op.batch_size
 
+        # Choose ref_bundler and fusion behavior based on strict mode
+        if down_logical_op._strict:
+            # Strict mode: use StreamingRepartitionRefBundler for stitching.
+            # Only works when batch_size % target == 0 (verified in _can_fuse).
+            assert batch_size % down_logical_op.target_num_rows_per_block == 0, (
+                f"Strict mode fusion requires batch_size ({batch_size}) to be "
+                f"a multiple of target_num_rows_per_block "
+                f"({down_logical_op.target_num_rows_per_block})"
+            )
+            ref_bundler = StreamingRepartitionRefBundler(batch_size)
+            # No further fusion because StreamingRepartitionRefBundler is stateful
+            # and maintains internal buffering state across bundles.
+            supports_fusion = False
+        else:
+            # Non-strict mode: use default pass-through bundler.
+            # Works with any batch_size without cross-task buffering.
+            ref_bundler = None
+            # Can fuse further because the default bundler is stateless
+            # and processes each bundle independently.
+            supports_fusion = True
+
         compute = self._fuse_compute_strategy(
             up_logical_op.compute, down_logical_op.compute
         )
@@ -331,19 +361,23 @@ def _get_fused_streaming_repartition_operator(
         input_op = input_deps[0]
 
         assert up_op.data_context is down_op.data_context
+
+        # In non-strict mode, use min_rows_per_bundle to ensure creating batches with batch_size.
+        # In strict mode, ref_bundler handles bundling, so do not set min_rows_per_bundle.
+        min_rows = None if down_logical_op._strict else batch_size
+
         op = MapOperator.create(
             up_op.get_map_transformer().fuse(down_op.get_map_transformer()),
             input_op,
             up_op.data_context,
             name=name,
             compute_strategy=compute,
-            ref_bundler=StreamingRepartitionRefBundler(batch_size),
+            min_rows_per_bundle=min_rows,
+            ref_bundler=ref_bundler,
             map_task_kwargs=map_task_kwargs,
             ray_remote_args=ray_remote_args,
             ray_remote_args_fn=ray_remote_args_fn,
-            # For now, we don't want to over-fuse StreamingRepartition with other map operators,
-            # so the result operator does not support further fusion.
-            supports_fusion=False,
+            supports_fusion=supports_fusion,
         )
         op.set_logical_operators(*up_op._logical_operators, *down_op._logical_operators)
         for map_task_kwargs_fn in itertools.chain(

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -195,14 +195,18 @@ def plan_streaming_repartition_op(
     )
     map_transformer = MapTransformer([transform_fn])
 
-    # Disable fusion for streaming repartition with the downstream op.
+    if op._strict:
+        ref_bundler = StreamingRepartitionRefBundler(op.target_num_rows_per_block)
+    else:
+        ref_bundler = None
+
     operator = MapOperator.create(
         map_transformer,
         input_physical_dag,
         data_context,
         name=op.name,
         compute_strategy=compute,
-        ref_bundler=StreamingRepartitionRefBundler(op.target_num_rows_per_block),
+        ref_bundler=ref_bundler,
         ray_remote_args=op.ray_remote_args,
         ray_remote_args_fn=op.ray_remote_args_fn,
     )

python/ray/data/dataset.py

@@ -1632,6 +1632,7 @@ def repartition(
         num_blocks: Optional[int] = None,
         target_num_rows_per_block: Optional[int] = None,
         *,
+        strict: bool = False,
         shuffle: bool = False,
         keys: Optional[List[str]] = None,
         sort: bool = False,
@@ -1681,6 +1682,13 @@ def repartition(
                 optimal execution, based on the `target_num_rows_per_block`. This is
                 the current behavior because of the implementation and may change in
                 the future.
+            strict: If ``True``, ``repartition`` guarantees that all output blocks,
+                except for the last one, will have exactly ``target_num_rows_per_block`` rows.
+                If ``False``, ``repartition`` uses best-effort bundling and may produce at most
+                one block smaller than ``target_num_rows_per_block`` per input block without
+                forcing exact sizes through block splitting.
+                This parameter is only used with ``target_num_rows_per_block``.
+                Defaults to ``False``.
             shuffle: Whether to perform a distributed shuffle during the
                 repartition. When shuffle is enabled, each output block
                 contains a subset of data rows from each input block, which
@@ -1717,6 +1725,13 @@ def repartition(
                 warnings.warn(
                     "`shuffle` is ignored when `target_num_rows_per_block` is set."
                 )
+        else:
+            if strict:
+                # strict is used in row-based repartition only
+                warnings.warn(
+                    "`strict` is ignored when `target_num_rows_per_block` is not set. "
+                    "Use `target_num_rows_per_block` instead of `num_blocks` to enable `strict` mode."
+                )
 
         if (num_blocks is None) and (target_num_rows_per_block is None):
             raise ValueError(
@@ -1738,6 +1753,7 @@ def repartition(
             op = StreamingRepartition(
                 self._logical_plan.dag,
                 target_num_rows_per_block=target_num_rows_per_block,
+                strict=strict,
             )
         else:
             op = Repartition(

python/ray/data/tests/test_operator_fusion.py

@@ -771,12 +771,12 @@ def test_streaming_repartition_map_batches_fusion_order_and_params(
 
     if order == "map_then_sr":
         ds = ds.map_batches(lambda x: x, batch_size=batch_size)
-        ds = ds.repartition(target_num_rows_per_block=target_num_rows)
-        expected_fused_name = f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_num_rows}]"
+        ds = ds.repartition(target_num_rows_per_block=target_num_rows, strict=True)
+        expected_fused_name = f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_num_rows},strict=True]"
     else:  # sr_then_map
-        ds = ds.repartition(target_num_rows_per_block=target_num_rows)
+        ds = ds.repartition(target_num_rows_per_block=target_num_rows, strict=True)
         ds = ds.map_batches(lambda x: x, batch_size=batch_size)
-        expected_fused_name = f"StreamingRepartition[num_rows_per_block={target_num_rows}]->MapBatches(<lambda>)"
+        expected_fused_name = f"StreamingRepartition[num_rows_per_block={target_num_rows},strict=True]->MapBatches(<lambda>)"
 
     assert len(ds.take_all()) == n
 
@@ -816,7 +816,7 @@ def test_streaming_repartition_no_further_fuse(
     stats1 = ds1.stats()
 
     assert (
-        f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_rows}]"
+        f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_rows},strict=False]"
         in stats1
     ), stats1
     assert "MapBatches(<lambda>)->MapBatches(<lambda>)" in stats1
@@ -844,6 +844,48 @@ def test_filter_operator_no_upstream_fusion(ray_start_regular_shared_2_cpus, cap
     assert "TaskPoolMapOperator[MapBatches(<lambda>)->Filter(<lambda>)]" in captured
 
 
+def test_streaming_repartition_multiple_fusion_non_strict(
+    ray_start_regular_shared_2_cpus,
+):
+    """Test that non-strict mode allows multiple operators to fuse with StreamingRepartition.
+
+    Case 1: Map > Map > SR (non-strict)
+    Case 2: Map > SR (non-strict) > Map
+    """
+    n = 100
+    target_rows = 20
+
+    # Case 1: Map > Map > SR (non-strict)
+    ds1 = ray.data.range(n, override_num_blocks=2)
+    ds1 = ds1.map_batches(lambda x: x, batch_size=None)
+    ds1 = ds1.map_batches(lambda x: x, batch_size=None)
+    ds1 = ds1.repartition(target_num_rows_per_block=target_rows, strict=False)
+
+    assert len(ds1.take_all()) == n
+    stats1 = ds1.stats()
+
+    # Verify all three operators are fused together
+    assert (
+        f"MapBatches(<lambda>)->MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_rows},strict=False]"
+        in stats1
+    ), f"Expected full fusion in stats: {stats1}"
+
+    # Case 2: Map > SR (non-strict) > Map
+    ds2 = ray.data.range(n, override_num_blocks=2)
+    ds2 = ds2.map_batches(lambda x: x, batch_size=None)
+    ds2 = ds2.repartition(target_num_rows_per_block=target_rows, strict=False)
+    ds2 = ds2.map_batches(lambda x: x, batch_size=None)
+
+    assert len(ds2.take_all()) == n
+    stats2 = ds2.stats()
+
+    # Verify all three operators are fused together
+    assert (
+        f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_rows},strict=False]->MapBatches(<lambda>)"
+        in stats2
+    ), f"Expected full fusion in stats: {stats2}"
+
+
 def test_combine_repartition_aggregate(
     ray_start_regular_shared_2_cpus, configure_shuffle_method, capsys
 ):