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

Code before:

        # Do not fuse if either op specifies a `_ray_remote_args_fn`,
        # since it is not known whether the generated args will be compatible.
        if getattr(up_logical_op, "_ray_remote_args_fn", None) or getattr(
            down_logical_op, "_ray_remote_args_fn", None
        ):
            return False

        if not self._can_merge_target_max_block_size(
            up_op.target_max_block_size_override,
            down_op.target_max_block_size_override,
        ):
            return False

        # only allow fusion of MapBatches -> StreamingRepartition
        if isinstance(down_logical_op, StreamingRepartition):
            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
            ):
                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 % 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__}"
        )

Review comment (performance):
I don't think this logic is correct -- if _batch_size is None we'd still allow to fuse StreamingRepartition

Code after:
            getattr(up_logical_op, "ray_remote_args", {}),
            getattr(down_logical_op, "ray_remote_args", {}),
        ):
            return False

        # Do not fuse if either op specifies a `ray_remote_args_fn`,
        # since it is not known whether the generated args will be compatible.
        if getattr(up_logical_op, "ray_remote_args_fn", None) or getattr(
            down_logical_op, "ray_remote_args_fn", None
        ):
            return False

        if not self._can_merge_target_max_block_size(
            up_op.target_max_block_size_override,
            down_op.target_max_block_size_override,
        ):
            return False

        # only allow fusion of MapBatches -> StreamingRepartition
        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: