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

Code before:
            "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, (
                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
        )
        assert compute is not None

        map_task_kwargs = {**up_op._map_task_kwargs, **down_op._map_task_kwargs}

        ray_remote_args = up_logical_op.ray_remote_args
        ray_remote_args_fn = (
            up_logical_op.ray_remote_args_fn or down_logical_op.ray_remote_args_fn
        )
        input_deps = up_op.input_dependencies
        assert len(input_deps) == 1
        input_op = input_deps[0]

        assert up_op.data_context is down_op.data_context


Review comment (suggestion):
We'd not be blocking any subsequent fusion like that 

Let's add a test that we're able to fuse multiple ops like this:

 - Map > Map > SR
 - Map > SR > SR

Code after:
            "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, (
                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)
        else:
            # Non-strict mode: use default pass-through bundler.
            # Works with any batch_size without cross-task buffering.
            ref_bundler = None

        compute = self._fuse_compute_strategy(
            up_logical_op.compute, down_logical_op.compute
        )
        assert compute is not None

        map_task_kwargs = {**up_op._map_task_kwargs, **down_op._map_task_kwargs}

        ray_remote_args = up_logical_op.ray_remote_args
        ray_remote_args_fn = (
            up_logical_op.ray_remote_args_fn or down_logical_op.ray_remote_args_fn
        )
        input_deps = up_op.input_dependencies
        assert len(input_deps) == 1
        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()),