Repository: ray-project/ray
File: python/ray/data/_internal/logical/rules/operator_fusion.py (Python)

Code before:
        if isinstance(down_logical_op, StreamingRepartition):
            if not (
                isinstance(up_logical_op, MapBatches)
                and down_logical_op.target_num_rows_per_block is not None
                and down_logical_op.target_num_rows_per_block > 0
            ):
                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
            )
        # StreamingRepartition can only fuse in non-strict mode.
        # In strict mode, it does not support further fusion.
        if isinstance(up_logical_op, StreamingRepartition):
            return not up_logical_op._strict

        # Otherwise, ops are compatible for fusion.
        return True

    def _get_fused_streaming_repartition_operator(
        self, down_op: PhysicalOperator, up_op: PhysicalOperator
    ) -> PhysicalOperator:
        assert self._can_fuse(down_op, up_op), (
            "Current rule supports fusing MapBatches->StreamingRepartition, but received: "
            f"{type(up_op).__name__} -> {type(down_op).__name__}"
        )

        name = up_op.name + "->" + down_op.name

        down_logical_op = self._op_map.pop(down_op)
        up_logical_op = self._op_map.pop(up_op)
        assert isinstance(up_logical_op, MapBatches)
        assert isinstance(down_logical_op, StreamingRepartition)

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

Review comment (suggestion):
We actually don't want to fuse SR > Map, b/c that will reduce parallelism for Map (i believe we'd have the test for that) 

Code after:
        if isinstance(down_logical_op, StreamingRepartition):
            if not (
                isinstance(up_logical_op, MapBatches)
                and down_logical_op.target_num_rows_per_block is not None
                and down_logical_op.target_num_rows_per_block > 0
            ):
                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.
        if isinstance(up_logical_op, StreamingRepartition):
            return False

        # Otherwise, ops are compatible for fusion.
        return True

    def _get_fused_streaming_repartition_operator(
        self, down_op: PhysicalOperator, up_op: PhysicalOperator
    ) -> PhysicalOperator:
        assert self._can_fuse(down_op, up_op), (
            "Current rule supports fusing MapBatches->StreamingRepartition, but received: "
            f"{type(up_op).__name__} -> {type(down_op).__name__}"
        )

        name = up_op.name + "->" + down_op.name

        down_logical_op = self._op_map.pop(down_op)
        up_logical_op = self._op_map.pop(up_op)
        assert isinstance(up_logical_op, MapBatches)
        assert isinstance(down_logical_op, StreamingRepartition)

        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, (