Add files using upload-large-folder tool

.venv/lib/python3.11/site-packages/ray/data/_internal/execution/bundle_queue/fifo_bundle_queue.py

@@ -0,0 +1,129 @@
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from .bundle_queue import BundleQueue
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+@dataclass
+class _Node:
+    value: "RefBundle"
+    next: Optional["_Node"] = None
+    prev: Optional["_Node"] = None
+
+
+class FIFOBundleQueue(BundleQueue):
+    """A bundle queue that follows a first-in-first-out policy."""
+
+    def __init__(self):
+        # We manually implement a linked list because we need to remove elements
+        # efficiently, and Python's built-in data structures have O(n) removal time.
+        self._head: Optional[_Node] = None
+        self._tail: Optional[_Node] = None
+        # We use a dictionary to keep track of the nodes corresponding to each bundle.
+        # This allows us to remove a bundle from the queue in O(1) time. We need a list
+        # because a bundle can be added to the queue multiple times. Nodes in each list
+        # are insertion-ordered.
+        self._bundle_to_nodes: Dict["RefBundle", List[_Node]] = defaultdict(deque)
+
+        self._nbytes = 0
+        self._num_bundles = 0
+
+    def __len__(self) -> int:
+        return self._num_bundles
+
+    def __contains__(self, bundle: "RefBundle") -> bool:
+        return bundle in self._bundle_to_nodes
+
+    def add(self, bundle: "RefBundle") -> None:
+        """Add a bundle to the end (right) of the queue."""
+        new_node = _Node(value=bundle, next=None, prev=self._tail)
+        # Case 1: The queue is empty.
+        if self._head is None:
+            assert self._tail is None
+            self._head = new_node
+            self._tail = new_node
+        # Case 2: The queue has at least one element.
+        else:
+            self._tail.next = new_node
+            self._tail = new_node
+
+        self._bundle_to_nodes[bundle].append(new_node)
+
+        self._nbytes += bundle.size_bytes()
+        self._num_bundles += 1
+
+    def pop(self) -> "RefBundle":
+        """Return the first (left) bundle in the queue."""
+        # Case 1: The queue is empty.
+        if not self._head:
+            raise IndexError("You can't pop from an empty queue")
+
+        bundle = self._head.value
+        self.remove(bundle)
+
+        return bundle
+
+    def peek(self) -> Optional["RefBundle"]:
+        """Return the first (left) bundle in the queue without removing it."""
+        if self._head is None:
+            return None
+
+        return self._head.value
+
+    def remove(self, bundle: "RefBundle"):
+        """Remove a bundle from the queue.
+
+        If there are multiple instances of the bundle in the queue, this method only
+        removes the first one.
+        """
+        # Case 1: The queue is empty.
+        if bundle not in self._bundle_to_nodes:
+            raise ValueError(f"The bundle {bundle} is not in the queue.")
+
+        node = self._bundle_to_nodes[bundle].popleft()
+        if not self._bundle_to_nodes[bundle]:
+            del self._bundle_to_nodes[bundle]
+
+        # Case 2: The bundle is the only element in the queue.
+        if self._head is self._tail:
+            self._head = None
+            self._tail = None
+        # Case 3: The bundle is the first element in the queue.
+        elif node is self._head:
+            self._head = node.next
+            self._head.prev = None
+        # Case 4: The bundle is the last element in the queue.
+        elif node is self._tail:
+            self._tail = node.prev
+            self._tail.next = None
+        # Case 5: The bundle is in the middle of the queue.
+        else:
+            node.prev.next = node.next
+            node.next.prev = node.prev
+
+        self._nbytes -= bundle.size_bytes()
+        assert self._nbytes >= 0, (
+            "Expected the total size of objects in the queue to be non-negative, but "
+            f"got {self._nbytes} bytes instead."
+        )
+
+        self._num_bundles -= 1
+
+        return node.value
+
+    def clear(self):
+        self._head = None
+        self._tail = None
+        self._bundle_to_nodes.clear()
+        self._nbytes = 0
+        self._num_bundles = 0
+
+    def estimate_size_bytes(self) -> int:
+        return self._nbytes
+
+    def is_empty(self):
+        return not self._bundle_to_nodes and self._head is None and self._tail is None

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/optimizers.py

@@ -0,0 +1,94 @@
+from typing import List, Optional, Type
+
+from ray.data._internal.logical.interfaces import (
+    LogicalPlan,
+    Optimizer,
+    PhysicalPlan,
+    Rule,
+)
+from ray.data._internal.logical.rules.inherit_batch_format import InheritBatchFormatRule
+from ray.data._internal.logical.rules.inherit_target_max_block_size import (
+    InheritTargetMaxBlockSizeRule,
+)
+from ray.data._internal.logical.rules.operator_fusion import OperatorFusionRule
+from ray.data._internal.logical.rules.randomize_blocks import ReorderRandomizeBlocksRule
+from ray.data._internal.logical.rules.set_read_parallelism import SetReadParallelismRule
+from ray.data._internal.logical.rules.zero_copy_map_fusion import (
+    EliminateBuildOutputBlocks,
+)
+from ray.data._internal.planner.planner import Planner
+from ray.util.annotations import DeveloperAPI
+
+_LOGICAL_RULES = [
+    ReorderRandomizeBlocksRule,
+    InheritBatchFormatRule,
+]
+
+_PHYSICAL_RULES = [
+    InheritTargetMaxBlockSizeRule,
+    SetReadParallelismRule,
+    OperatorFusionRule,
+    EliminateBuildOutputBlocks,
+]
+
+
+@DeveloperAPI
+def register_logical_rule(cls: Type[Rule], insert_index: Optional[int] = None):
+    if cls in _LOGICAL_RULES:
+        return
+
+    if insert_index is None:
+        _LOGICAL_RULES.append(cls)
+    else:
+        _LOGICAL_RULES.insert(insert_index, cls)
+
+
+@DeveloperAPI
+def get_logical_rules() -> List[Type[Rule]]:
+    return list(_LOGICAL_RULES)
+
+
+@DeveloperAPI
+def register_physical_rule(cls: Type[Rule], insert_index: Optional[int] = None):
+    if cls in _PHYSICAL_RULES:
+        return
+
+    if insert_index is None:
+        _PHYSICAL_RULES.append(cls)
+    else:
+        _PHYSICAL_RULES.insert(insert_index, cls)
+
+
+@DeveloperAPI
+def get_physical_rules() -> List[Type[Rule]]:
+    return list(_PHYSICAL_RULES)
+
+
+class LogicalOptimizer(Optimizer):
+    """The optimizer for logical operators."""
+
+    @property
+    def rules(self) -> List[Rule]:
+        return [rule_cls() for rule_cls in _LOGICAL_RULES]
+
+
+class PhysicalOptimizer(Optimizer):
+    """The optimizer for physical operators."""
+
+    @property
+    def rules(self) -> List[Rule]:
+        return [rule_cls() for rule_cls in _PHYSICAL_RULES]
+
+
+def get_execution_plan(logical_plan: LogicalPlan) -> PhysicalPlan:
+    """Get the physical execution plan for the provided logical plan.
+
+    This process has 3 steps:
+    (1) logical optimization: optimize logical operators.
+    (2) planning: convert logical to physical operators.
+    (3) physical optimization: optimize physical operators.
+    """
+    optimized_logical_plan = LogicalOptimizer().optimize(logical_plan)
+    logical_plan._dag = optimized_logical_plan.dag
+    physical_plan = Planner().plan(optimized_logical_plan)
+    return PhysicalOptimizer().optimize(physical_plan)

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/rules/inherit_batch_format.py

@@ -0,0 +1,42 @@
+from collections import deque
+from typing import Iterable
+
+from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan, Rule
+from ray.data._internal.logical.operators.all_to_all_operator import AbstractAllToAll
+from ray.data._internal.logical.operators.map_operator import MapBatches
+
+
+class InheritBatchFormatRule(Rule):
+    """For AbstractAllToAll based operator, apply this rule
+    to inherit batch_format from upstream operator by traversing
+    the entire DAG."""
+
+    def apply(self, plan: LogicalPlan) -> LogicalPlan:
+        optimized_dag: LogicalOperator = self._apply(plan.dag)
+        new_plan = LogicalPlan(dag=optimized_dag, context=plan.context)
+        return new_plan
+
+    def _apply(self, op: LogicalOperator):
+        # Post-order traversal.
+        nodes: Iterable[LogicalOperator] = deque()
+        for node in op.post_order_iter():
+            nodes.appendleft(node)
+
+        while len(nodes) > 0:
+            current_op = nodes.pop()
+
+            if isinstance(current_op, AbstractAllToAll):
+                # traversal up the DAG until we find MapBatches with batch_format
+                # or we reach to source op and do nothing
+                upstream_op = current_op.input_dependencies[0]
+                while upstream_op.input_dependencies:
+                    if (
+                        isinstance(upstream_op, MapBatches)
+                        and upstream_op._batch_format
+                    ):
+                        current_op._batch_format = upstream_op._batch_format
+                        break
+                    upstream_op = upstream_op.input_dependencies[0]
+
+        # just return the default op
+        return op

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/rules/inherit_target_max_block_size.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.logical.interfaces import PhysicalPlan, Rule
+
+
+class InheritTargetMaxBlockSizeRule(Rule):
+    """For each op that has overridden the default target max block size,
+    propagate to upstream ops until we reach an op that has also overridden the
+    target max block size."""
+
+    def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
+        self._propagate_target_max_block_size_to_upstream_ops(plan.dag)
+        return plan
+
+    def _propagate_target_max_block_size_to_upstream_ops(
+        self, dag: PhysicalOperator, target_max_block_size: Optional[int] = None
+    ):
+        if dag.target_max_block_size is not None:
+            # Set the target block size to inherit for
+            # upstream ops.
+            target_max_block_size = dag.target_max_block_size
+        elif target_max_block_size is not None:
+            # Inherit from downstream op.
+            dag.set_target_max_block_size(target_max_block_size)
+
+        for upstream_op in dag.input_dependencies:
+            self._propagate_target_max_block_size_to_upstream_ops(
+                upstream_op, target_max_block_size
+            )

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/rules/limit_pushdown.py

@@ -0,0 +1,133 @@
+import copy
+from collections import deque
+from typing import Iterable, List
+
+from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan, Rule
+from ray.data._internal.logical.operators.one_to_one_operator import (
+    AbstractOneToOne,
+    Limit,
+)
+from ray.data._internal.logical.operators.read_operator import Read
+
+
+class LimitPushdownRule(Rule):
+    """Rule for pushing down the limit operator.
+
+    When a limit operator is present, we apply the limit on the
+    most upstream operator that supports it. Notably, we move the
+    Limit operator downstream from Read op, any other non-OneToOne operator,
+    or any operator which could potentially change the number of output rows.
+
+    In addition, we also fuse consecutive Limit operators into a single
+    Limit operator, i.e. `Limit[n] -> Limit[m]` becomes `Limit[min(n, m)]`.
+    """
+
+    def apply(self, plan: LogicalPlan) -> LogicalPlan:
+        optimized_dag = self._apply_limit_pushdown(plan.dag)
+        optimized_dag = self._apply_limit_fusion(optimized_dag)
+        return LogicalPlan(dag=optimized_dag, context=plan.context)
+
+    def _apply_limit_pushdown(self, op: LogicalOperator) -> LogicalOperator:
+        """Given a DAG of LogicalOperators, traverse the DAG and push down
+        Limit operators, i.e. move Limit operators as far upstream as possible.
+
+        Returns a new LogicalOperator with the Limit operators pushed down."""
+        # Post-order traversal.
+        nodes: Iterable[LogicalOperator] = deque()
+        for node in op.post_order_iter():
+            nodes.appendleft(node)
+
+        while len(nodes) > 0:
+            current_op = nodes.pop()
+
+            # If we encounter a Limit op, move it upstream until it reaches:
+            # - Read operator
+            # - A non-AbstractOneToOne operator (e.g. AbstractAllToAll)
+            # - An AbstractOneToOne operator that could change the number of output rows
+
+            # TODO(scottjlee): in our current abstraction, we have Read extend
+            # AbstractMap (with no input dependency), which extends AbstractOneToOne.
+            # So we have to explicitly separate the Read op in its own check.
+            # We should remove this case once we refactor Read op to no longer
+            # be an AbstractOneToOne op.
+            if isinstance(current_op, Limit):
+                limit_op_copy = copy.copy(current_op)
+
+                # Traverse up the DAG until we reach the first operator that meets
+                # one of the conditions above, which will serve as the new input
+                # into the Limit operator.
+                new_input_into_limit = current_op.input_dependency
+                ops_between_new_input_and_limit: List[LogicalOperator] = []
+                while (
+                    isinstance(new_input_into_limit, AbstractOneToOne)
+                    and not isinstance(new_input_into_limit, Read)
+                    and not getattr(new_input_into_limit, "can_modify_num_rows", False)
+                ):
+                    new_input_into_limit_copy = copy.copy(new_input_into_limit)
+                    ops_between_new_input_and_limit.append(new_input_into_limit_copy)
+                    new_input_into_limit = new_input_into_limit.input_dependency
+
+                # Link the Limit operator and its newly designated input op from above.
+                limit_op_copy._input_dependencies = [new_input_into_limit]
+                new_input_into_limit._output_dependencies = [limit_op_copy]
+
+                # Build the chain of operator dependencies between the new
+                # input and the Limit operator, using copies of traversed operators.
+                ops_between_new_input_and_limit.append(limit_op_copy)
+                for idx in range(len(ops_between_new_input_and_limit) - 1):
+                    curr_op, up_op = (
+                        ops_between_new_input_and_limit[idx],
+                        ops_between_new_input_and_limit[idx + 1],
+                    )
+                    curr_op._input_dependencies = [up_op]
+                    up_op._output_dependencies = [curr_op]
+                    # Add the copied operator to the list of nodes to be traversed.
+                    nodes.append(curr_op)
+
+                # Link the Limit operator to its new input operator.
+                for limit_output_op in current_op.output_dependencies:
+                    limit_output_op._input_dependencies = [
+                        ops_between_new_input_and_limit[0]
+                    ]
+                last_op = ops_between_new_input_and_limit[0]
+                last_op._output_dependencies = current_op.output_dependencies
+
+        return current_op
+
+    def _apply_limit_fusion(self, op: LogicalOperator) -> LogicalOperator:
+        """Given a DAG of LogicalOperators, traverse the DAG and fuse all
+        back-to-back Limit operators, i.e.
+        Limit[n] -> Limit[m] becomes Limit[min(n, m)].
+
+        Returns a new LogicalOperator with the Limit operators fusion applied."""
+
+        # Post-order traversal.
+        nodes: Iterable[LogicalOperator] = deque()
+        for node in op.post_order_iter():
+            nodes.appendleft(node)
+
+        while len(nodes) > 0:
+            current_op = nodes.pop()
+
+            # If we encounter two back-to-back Limit operators, fuse them.
+            if isinstance(current_op, Limit):
+                upstream_op = current_op.input_dependency
+                if isinstance(upstream_op, Limit):
+                    new_limit = min(current_op._limit, upstream_op._limit)
+                    fused_limit_op = Limit(upstream_op.input_dependency, new_limit)
+
+                    # Link the fused Limit operator to its input and output ops, i.e.:
+                    # `upstream_input -> limit_upstream -> limit_downstream -> downstream_output`  # noqa: E501
+                    # becomes `upstream_input -> fused_limit -> downstream_output`
+                    fused_limit_op._input_dependencies = upstream_op.input_dependencies
+                    fused_limit_op._output_dependencies = current_op.output_dependencies
+
+                    # Replace occurrences of the upstream Limit operator in
+                    # output_dependencies with the newly fused Limit operator.
+                    upstream_input = upstream_op.input_dependency
+                    upstream_input._output_dependencies = [fused_limit_op]
+
+                    for current_output in current_op.output_dependencies:
+                        current_output._input_dependencies = [fused_limit_op]
+                    nodes.append(fused_limit_op)
+        return current_op

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/rules/operator_fusion.py

@@ -0,0 +1,491 @@
+import itertools
+from typing import List, Optional, Tuple
+
+from ray.data._internal.compute import (
+    ActorPoolStrategy,
+    ComputeStrategy,
+    TaskPoolStrategy,
+)
+from ray.data._internal.execution.interfaces import (
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.operators.actor_pool_map_operator import (
+    ActorPoolMapOperator,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    AllToAllOperator,
+)
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.task_pool_map_operator import (
+    TaskPoolMapOperator,
+)
+from ray.data._internal.logical.interfaces import PhysicalPlan, Rule
+from ray.data._internal.logical.operators.all_to_all_operator import (
+    AbstractAllToAll,
+    RandomShuffle,
+    Repartition,
+)
+from ray.data._internal.logical.operators.map_operator import (
+    AbstractMap,
+    AbstractUDFMap,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+# Scheduling strategy can be inherited from upstream operator if not specified.
+INHERITABLE_REMOTE_ARGS = ["scheduling_strategy"]
+
+
+class OperatorFusionRule(Rule):
+    """Fuses linear chains of compatible physical operators."""
+
+    def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
+        self._op_map = plan.op_map.copy()
+        # Do DFS fusion on compatible pairwise operators in two passes.
+        # In the first pass, only fuse back-to-back map operators together.
+        fused_dag = self._fuse_map_operators_in_dag(plan.dag)
+
+        # Now that we have fused together all back-to-back map operators,
+        # we fuse together MapOperator -> AllToAllOperator pairs.
+        fused_dag = self._fuse_all_to_all_operators_in_dag(fused_dag)
+
+        # Update output dependencies after fusion.
+        # TODO(hchen): Instead of updating the depdencies manually,
+        # we need a better abstraction for manipulating the DAG.
+        self._remove_output_depes(fused_dag)
+        self._update_output_depes(fused_dag)
+
+        new_plan = PhysicalPlan(fused_dag, self._op_map, plan.context)
+        return new_plan
+
+    def _remove_output_depes(self, op: PhysicalOperator) -> None:
+        for input in op._input_dependencies:
+            input._output_dependencies = []
+            self._remove_output_depes(input)
+
+    def _update_output_depes(self, op: PhysicalOperator) -> None:
+        for input in op._input_dependencies:
+            input._output_dependencies.append(op)
+            self._update_output_depes(input)
+
+    def _fuse_map_operators_in_dag(self, dag: PhysicalOperator) -> MapOperator:
+        """Starting at the given operator, traverses up the DAG of operators
+        and recursively fuses compatible MapOperator -> MapOperator pairs.
+        Returns the current (root) operator after completing upstream operator fusions.
+        """
+        upstream_ops = dag.input_dependencies
+        while (
+            len(upstream_ops) == 1
+            and isinstance(dag, MapOperator)
+            and isinstance(upstream_ops[0], MapOperator)
+            and self._can_fuse(dag, upstream_ops[0])
+        ):
+            # Fuse operator with its upstream op.
+            dag = self._get_fused_map_operator(dag, upstream_ops[0])
+            upstream_ops = dag.input_dependencies
+
+        # Done fusing back-to-back map operators together here,
+        # move up the DAG to find the next map operators to fuse.
+        dag._input_dependencies = [
+            self._fuse_map_operators_in_dag(upstream_op) for upstream_op in upstream_ops
+        ]
+        return dag
+
+    def _fuse_all_to_all_operators_in_dag(
+        self, dag: AllToAllOperator
+    ) -> AllToAllOperator:
+        """Starting at the given operator, traverses up the DAG of operators
+        and recursively fuses compatible MapOperator -> AllToAllOperator pairs.
+
+        Also, sets the target block size of the immediately upstream map op to
+        match the shuffle block size. We use a larger block size for shuffles
+        because tiny blocks are bad for I/O performance.
+
+        Returns the current (root) operator after completing upstream operator fusions.
+        """
+        upstream_ops = dag.input_dependencies
+        while (
+            len(upstream_ops) == 1
+            and isinstance(dag, AllToAllOperator)
+            and isinstance(upstream_ops[0], MapOperator)
+            and self._can_fuse(dag, upstream_ops[0])
+        ):
+            # Fuse operator with its upstream op.
+            dag = self._get_fused_all_to_all_operator(dag, upstream_ops[0])
+            upstream_ops = dag.input_dependencies
+
+        # Done fusing MapOperator -> AllToAllOperator together here,
+        # move up the DAG to find the next pair of operators to fuse.
+        dag._input_dependencies = [
+            self._fuse_all_to_all_operators_in_dag(upstream_op)
+            for upstream_op in upstream_ops
+        ]
+        return dag
+
+    def _can_fuse(self, down_op: PhysicalOperator, up_op: PhysicalOperator) -> bool:
+        """Returns whether the provided downstream operator can be fused with the given
+        upstream operator.
+
+        We currently support fusing two operators if the following are all true:
+            * We are fusing either MapOperator -> MapOperator or
+              MapOperator -> AllToAllOperator.
+            * They either use the same compute configuration, or the upstream operator
+              uses a task pool while the downstream operator uses an actor pool.
+            * If both operators involve callable classes, the callable classes are
+              the same class AND constructor args are the same for both.
+            * They have compatible remote arguments.
+        """
+        if not up_op.supports_fusion() or not down_op.supports_fusion():
+            return False
+
+        # We currently only support fusing for the following cases:
+        # - TaskPoolMapOperator -> TaskPoolMapOperator/ActorPoolMapOperator
+        # - TaskPoolMapOperator -> AllToAllOperator
+        # (only RandomShuffle and Repartition LogicalOperators are currently supported)
+        if not (
+            (
+                isinstance(up_op, TaskPoolMapOperator)
+                and isinstance(down_op, (TaskPoolMapOperator, ActorPoolMapOperator))
+            )
+            or (
+                isinstance(up_op, TaskPoolMapOperator)
+                and isinstance(down_op, AllToAllOperator)
+            )
+        ):
+            return False
+
+        down_logical_op = self._op_map[down_op]
+        up_logical_op = self._op_map[up_op]
+
+        if up_op.get_additional_split_factor() > 1:
+            return False
+
+        # If the downstream operator takes no input, it cannot be fused with
+        # the upstream operator.
+        if not down_logical_op._input_dependencies:
+            return False
+
+        # We currently only support fusing for the following cases:
+        # - AbstractMap -> AbstractMap
+        # - AbstractMap -> RandomShuffle
+        # - AbstractMap -> Repartition (shuffle=True)
+        if not (
+            (
+                isinstance(up_logical_op, AbstractMap)
+                and isinstance(down_logical_op, AbstractMap)
+            )
+            or (
+                isinstance(up_logical_op, AbstractMap)
+                and isinstance(down_logical_op, RandomShuffle)
+            )
+            or (
+                isinstance(up_logical_op, AbstractMap)
+                and isinstance(down_logical_op, Repartition)
+            )
+        ):
+            return False
+
+        # Do not fuse Repartition operator if shuffle is disabled
+        # (i.e. using split shuffle).
+        if isinstance(down_logical_op, Repartition) and not down_logical_op._shuffle:
+            return False
+
+        if isinstance(down_logical_op, AbstractMap) and isinstance(
+            up_logical_op, AbstractMap
+        ):
+            if (
+                self._fuse_compute_strategy(
+                    up_logical_op._compute,
+                    down_logical_op._compute,
+                )
+                is None
+            ):
+                return False
+
+        # Only fuse if the ops' remote arguments are compatible.
+        if not _are_remote_args_compatible(
+            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,
+            down_op.target_max_block_size,
+            up_op.data_context,
+        ):
+            return False
+
+        # Otherwise, ops are compatible for fusion.
+        return True
+
+    def _fuse_compute_strategy(
+        self, up_compute: ComputeStrategy, down_compute: ComputeStrategy
+    ) -> Optional[ComputeStrategy]:
+        """Fuse the compute strategies of the upstream and downstream operators.
+        Returns None if they are not compatible.
+
+        Task->Task and Task->Actor are allowed.
+        Actor->Actor and Actor->Task are not allowed.
+        """
+        if isinstance(up_compute, ActorPoolStrategy):
+            return None
+        assert isinstance(up_compute, TaskPoolStrategy)
+        if isinstance(down_compute, TaskPoolStrategy):
+            # For Task->Task, the sizes must match.
+            if up_compute.size != down_compute.size:
+                return None
+            return down_compute
+        else:
+            assert isinstance(down_compute, ActorPoolStrategy)
+            # For Task->Actor, if Task's size is set, it must match Actor's max_size.
+            if up_compute.size is not None and up_compute.size != down_compute.max_size:
+                return None
+            return down_compute
+
+    def _can_merge_target_max_block_size(
+        self,
+        up_target_max_block_size: Optional[int],
+        down_target_max_block_size: Optional[int],
+        data_context: DataContext,
+    ):
+        # If the upstream op overrode the target max block size, only fuse if
+        # they are equal.
+        if up_target_max_block_size is not None:
+            if down_target_max_block_size is None:
+                down_target_max_block_size = data_context.target_max_block_size
+            if up_target_max_block_size != down_target_max_block_size:
+                return False
+        return True
+
+    def _get_merged_target_max_block_size(
+        self,
+        up_target_max_block_size: Optional[int],
+        down_target_max_block_size: Optional[int],
+    ):
+        if up_target_max_block_size is not None:
+            # If the upstream op overrode the target max block size, we can
+            # only merge if the downstream op matches or uses the default.
+            assert (
+                down_target_max_block_size is None
+                or down_target_max_block_size == up_target_max_block_size
+            )
+            return up_target_max_block_size
+        else:
+            # Upstream op inherits the downstream op's target max block size,
+            # because the downstream op is the one that outputs the final
+            # blocks.
+            return down_target_max_block_size
+
+    def _get_fused_map_operator(
+        self, down_op: MapOperator, up_op: MapOperator
+    ) -> MapOperator:
+        assert self._can_fuse(down_op, up_op), (
+            "Current rule supports fusing MapOperator->MapOperator, but received: "
+            f"{type(up_op).__name__} -> {type(down_op).__name__}"
+        )
+
+        # Fuse operator names.
+        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(down_logical_op, AbstractMap)
+        assert isinstance(up_logical_op, AbstractMap)
+
+        # Merge minimum block sizes.
+        down_min_rows_per_bundled_input = down_logical_op._min_rows_per_bundled_input
+        up_min_rows_per_bundled_input = up_logical_op._min_rows_per_bundled_input
+        if (
+            down_min_rows_per_bundled_input is not None
+            and up_min_rows_per_bundled_input is not None
+        ):
+            min_rows_per_bundled_input = max(
+                down_min_rows_per_bundled_input, up_min_rows_per_bundled_input
+            )
+        elif up_min_rows_per_bundled_input is not None:
+            min_rows_per_bundled_input = up_min_rows_per_bundled_input
+        else:
+            min_rows_per_bundled_input = down_min_rows_per_bundled_input
+
+        target_max_block_size = self._get_merged_target_max_block_size(
+            up_op.target_max_block_size, down_op.target_max_block_size
+        )
+
+        compute = self._fuse_compute_strategy(
+            up_logical_op._compute, down_logical_op._compute
+        )
+        assert compute is not None
+        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
+        )
+        # Make the upstream operator's inputs the new, fused operator's inputs.
+        input_deps = up_op.input_dependencies
+        assert len(input_deps) == 1
+        input_op = input_deps[0]
+
+        # Fused physical map operator.
+        assert up_op.data_context is down_op.data_context
+        op = MapOperator.create(
+            up_op.get_map_transformer().fuse(down_op.get_map_transformer()),
+            input_op,
+            up_op.data_context,
+            target_max_block_size=target_max_block_size,
+            name=name,
+            compute_strategy=compute,
+            min_rows_per_bundle=min_rows_per_bundled_input,
+            ray_remote_args=ray_remote_args,
+            ray_remote_args_fn=ray_remote_args_fn,
+        )
+        op.set_logical_operators(*up_op._logical_operators, *down_op._logical_operators)
+        for map_task_kwargs_fn in itertools.chain(
+            up_op._map_task_kwargs_fns, down_op._map_task_kwargs_fns
+        ):
+            op.add_map_task_kwargs_fn(map_task_kwargs_fn)
+
+        # Build a map logical operator to be used as a reference for further fusion.
+        # TODO(Scott): This is hacky, remove this once we push fusion to be purely based
+        # on a lower-level operator spec.
+        if isinstance(up_logical_op, AbstractUDFMap):
+            input_op = up_logical_op.input_dependency
+        else:
+            # Bottom out at the source logical op (e.g. Read()).
+            input_op = up_logical_op
+        if isinstance(down_logical_op, AbstractUDFMap):
+            logical_op = AbstractUDFMap(
+                name,
+                input_op,
+                down_logical_op._fn,
+                down_logical_op._fn_args,
+                down_logical_op._fn_kwargs,
+                down_logical_op._fn_constructor_args,
+                down_logical_op._fn_constructor_kwargs,
+                min_rows_per_bundled_input,
+                compute,
+                ray_remote_args_fn,
+                ray_remote_args,
+            )
+        else:
+            # The downstream op is AbstractMap instead of AbstractUDFMap.
+            logical_op = AbstractMap(
+                name,
+                input_op,
+                min_rows_per_bundled_input=min_rows_per_bundled_input,
+                ray_remote_args_fn=ray_remote_args_fn,
+                ray_remote_args=ray_remote_args,
+            )
+        self._op_map[op] = logical_op
+        # Return the fused physical operator.
+        return op
+
+    def _get_fused_all_to_all_operator(
+        self, down_op: AllToAllOperator, up_op: MapOperator
+    ) -> AllToAllOperator:
+        assert self._can_fuse(down_op, up_op), (
+            "Current rule supports fusing MapOperator -> AllToAllOperator"
+            f", but received: {type(up_op).__name__} -> {type(down_op).__name__}"
+        )
+
+        # Fuse operator names.
+        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(down_logical_op, AbstractAllToAll)
+        assert isinstance(up_logical_op, AbstractMap)
+
+        # Fuse transformation functions.
+        ray_remote_args = up_logical_op._ray_remote_args
+        down_transform_fn = down_op.get_transformation_fn()
+        up_map_transformer = up_op.get_map_transformer()
+
+        def fused_all_to_all_transform_fn(
+            blocks: List[RefBundle], ctx: TaskContext
+        ) -> Tuple[List[RefBundle], StatsDict]:
+            """To fuse MapOperator->AllToAllOperator, we store the map function
+            in the TaskContext so that it may be used by the downstream
+            AllToAllOperator's transform function."""
+            ctx.upstream_map_transformer = up_map_transformer
+            ctx.upstream_map_ray_remote_args = ray_remote_args
+            return down_transform_fn(blocks, ctx)
+
+        # Make the upstream operator's inputs the new, fused operator's inputs.
+        input_deps = up_op.input_dependencies
+        assert len(input_deps) == 1
+        input_op = input_deps[0]
+
+        target_max_block_size = self._get_merged_target_max_block_size(
+            up_op.target_max_block_size, down_op.target_max_block_size
+        )
+
+        assert up_op.data_context is down_op.data_context
+        op = AllToAllOperator(
+            fused_all_to_all_transform_fn,
+            input_op,
+            up_op.data_context,
+            target_max_block_size=target_max_block_size,
+            num_outputs=down_op._num_outputs,
+            # Transfer over the existing sub-progress bars from
+            # the AllToAllOperator (if any) into the fused operator.
+            sub_progress_bar_names=down_op._sub_progress_bar_names,
+            name=name,
+        )
+        # Bottom out at the source logical op (e.g. Read()).
+        input_op = up_logical_op
+
+        if isinstance(down_logical_op, RandomShuffle):
+            logical_op = RandomShuffle(
+                input_op,
+                name=name,
+                ray_remote_args=ray_remote_args,
+            )
+        elif isinstance(down_logical_op, Repartition):
+            logical_op = Repartition(
+                input_op,
+                num_outputs=down_logical_op._num_outputs,
+                shuffle=down_logical_op._shuffle,
+            )
+        self._op_map[op] = logical_op
+        # Return the fused physical operator.
+        return op
+
+
+def _are_remote_args_compatible(prev_args, next_args):
+    """Check if Ray remote arguments are compatible for merging."""
+    prev_args = _canonicalize(prev_args)
+    next_args = _canonicalize(next_args)
+    remote_args = next_args.copy()
+    for key in INHERITABLE_REMOTE_ARGS:
+        # NOTE: We only carry over inheritable value in case
+        #       of it not being provided in the remote args
+        if key in prev_args and key not in remote_args:
+            remote_args[key] = prev_args[key]
+
+    if prev_args != remote_args:
+        return False
+    return True
+
+
+def _canonicalize(remote_args: dict) -> dict:
+    """Returns canonical form of given remote args."""
+    remote_args = remote_args.copy()
+    if "num_cpus" not in remote_args or remote_args["num_cpus"] is None:
+        remote_args["num_cpus"] = 1
+    if "num_gpus" not in remote_args or remote_args["num_gpus"] is None:
+        remote_args["num_gpus"] = 0
+    resources = remote_args.get("resources", {})
+    for k, v in list(resources.items()):
+        if v is None or v == 0.0:
+            del resources[k]
+    remote_args["resources"] = resources
+    return remote_args

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/rules/set_read_parallelism.py

@@ -0,0 +1,132 @@
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+from ray import available_resources as ray_available_resources
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.logical.interfaces import PhysicalPlan, Rule
+from ray.data._internal.logical.operators.read_operator import Read
+from ray.data._internal.util import _autodetect_parallelism
+from ray.data.context import WARN_PREFIX, DataContext
+from ray.data.datasource.datasource import Datasource, Reader
+
+logger = logging.getLogger(__name__)
+
+
+def compute_additional_split_factor(
+    datasource_or_legacy_reader: Union[Datasource, Reader],
+    parallelism: int,
+    mem_size: int,
+    target_max_block_size: int,
+    cur_additional_split_factor: Optional[int] = None,
+) -> Tuple[int, str, int, Optional[int]]:
+    ctx = DataContext.get_current()
+    detected_parallelism, reason, _ = _autodetect_parallelism(
+        parallelism, target_max_block_size, ctx, datasource_or_legacy_reader, mem_size
+    )
+    num_read_tasks = len(
+        datasource_or_legacy_reader.get_read_tasks(detected_parallelism)
+    )
+    expected_block_size = None
+    if mem_size:
+        expected_block_size = mem_size / num_read_tasks
+        logger.debug(
+            f"Expected in-memory size {mem_size}," f" block size {expected_block_size}"
+        )
+        size_based_splits = round(max(1, expected_block_size / target_max_block_size))
+    else:
+        size_based_splits = 1
+    if cur_additional_split_factor:
+        size_based_splits *= cur_additional_split_factor
+    logger.debug(f"Size based split factor {size_based_splits}")
+    estimated_num_blocks = num_read_tasks * size_based_splits
+    logger.debug(f"Blocks after size splits {estimated_num_blocks}")
+
+    available_cpu_slots = ray_available_resources().get("CPU", 1)
+    if (
+        parallelism != -1
+        and num_read_tasks >= available_cpu_slots * 4
+        and num_read_tasks >= 5000
+    ):
+        logger.warning(
+            f"{WARN_PREFIX} The requested number of read blocks of {parallelism} "
+            "is more than 4x the number of available CPU slots in the cluster of "
+            f"{available_cpu_slots}. This can "
+            "lead to slowdowns during the data reading phase due to excessive "
+            "task creation. Reduce the value to match with the available "
+            "CPU slots in the cluster, or set override_num_blocks to -1 for Ray Data "
+            "to automatically determine the number of read tasks blocks."
+            "You can ignore this message if the cluster is expected to autoscale."
+        )
+
+    # Add more output splitting for each read task if needed.
+    # TODO(swang): For parallelism=-1 (user did not explicitly set
+    # parallelism), and if the following operator produces much larger blocks,
+    # we should scale down the target max block size here instead of using
+    # splitting, which can have higher memory usage.
+    if estimated_num_blocks < detected_parallelism and estimated_num_blocks > 0:
+        k = math.ceil(detected_parallelism / estimated_num_blocks)
+        estimated_num_blocks = estimated_num_blocks * k
+        return detected_parallelism, reason, estimated_num_blocks, k
+
+    return detected_parallelism, reason, estimated_num_blocks, None
+
+
+class SetReadParallelismRule(Rule):
+    """
+    This rule sets the read op's task parallelism based on the target block
+    size, the requested parallelism, the number of read files, and the
+    available resources in the cluster.
+
+    If the parallelism is lower than requested, this rule also sets a split
+    factor to split the output blocks of the read task, so that the following
+    operator will have the desired parallelism.
+    """
+
+    def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
+        ops = [plan.dag]
+
+        while len(ops) > 0:
+            op = ops.pop(0)
+            if isinstance(op, InputDataBuffer):
+                continue
+            logical_op = plan.op_map[op]
+            if isinstance(logical_op, Read):
+                self._apply(op, logical_op)
+            ops += op.input_dependencies
+
+        return plan
+
+    def _apply(self, op: PhysicalOperator, logical_op: Read):
+        (
+            detected_parallelism,
+            reason,
+            estimated_num_blocks,
+            k,
+        ) = compute_additional_split_factor(
+            logical_op._datasource_or_legacy_reader,
+            logical_op._parallelism,
+            logical_op._mem_size,
+            op.actual_target_max_block_size,
+            op._additional_split_factor,
+        )
+
+        if logical_op._parallelism == -1:
+            assert reason != ""
+            logger.debug(
+                f"Using autodetected parallelism={detected_parallelism} "
+                f"for operator {logical_op.name} to satisfy {reason}."
+            )
+        logical_op.set_detected_parallelism(detected_parallelism)
+
+        if k is not None:
+            logger.debug(
+                f"To satisfy the requested parallelism of {detected_parallelism}, "
+                f"each read task output is split into {k} smaller blocks."
+            )
+
+        if k is not None:
+            op.set_additional_split_factor(k)
+
+        logger.debug(f"Estimated num output blocks {estimated_num_blocks}")

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/rules/zero_copy_map_fusion.py

@@ -0,0 +1,88 @@
+from abc import abstractmethod
+from typing import List
+
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BuildOutputBlocksMapTransformFn,
+    MapTransformFn,
+    MapTransformFnDataType,
+)
+from ray.data._internal.logical.interfaces.optimizer import Rule
+from ray.data._internal.logical.interfaces.physical_plan import PhysicalPlan
+
+
+class ZeroCopyMapFusionRule(Rule):
+    """Base abstract class for all zero-copy map fusion rules.
+
+    A zero-copy map fusion rule is a rule that optimizes the transform_fn chain of
+    a fused MapOperator. The optimization is usually done by removing unnecessary
+    data conversions.
+
+    This base abstract class defines the common util functions. And subclasses
+    should implement the `_optimize` method for the concrete optimization
+    strategy.
+    """
+
+    def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
+        self._traverse(plan.dag)
+        return plan
+
+    def _traverse(self, op):
+        """Traverse the DAG and apply the optimization to each MapOperator."""
+        if isinstance(op, MapOperator):
+            map_transformer = op.get_map_transformer()
+            transform_fns = map_transformer.get_transform_fns()
+            new_transform_fns = self._optimize(transform_fns)
+            # Physical operators won't be shared,
+            # so it's safe to modify the transform_fns in place.
+            map_transformer.set_transform_fns(new_transform_fns)
+
+        for input_op in op.input_dependencies:
+            self._traverse(input_op)
+
+    @abstractmethod
+    def _optimize(self, transform_fns: List[MapTransformFn]) -> List[MapTransformFn]:
+        """Optimize the transform_fns chain of a MapOperator.
+
+        Args:
+            transform_fns: The old transform_fns chain.
+        Returns:
+            The optimized transform_fns chain.
+        """
+        ...
+
+
+class EliminateBuildOutputBlocks(ZeroCopyMapFusionRule):
+    """This rule eliminates unnecessary BuildOutputBlocksMapTransformFn,
+    if the previous fn already outputs blocks.
+
+    This happens for the "Read -> Map/Write" fusion.
+    """
+
+    def _optimize(self, transform_fns: List[MapTransformFn]) -> List[MapTransformFn]:
+        # For the following subsquence,
+        # 1. Any MapTransformFn with block output.
+        # 2. BuildOutputBlocksMapTransformFn
+        # 3. Any MapTransformFn with block input.
+        # We drop the BuildOutputBlocksMapTransformFn in the middle.
+        new_transform_fns = []
+
+        for i in range(len(transform_fns)):
+            cur_fn = transform_fns[i]
+            drop = False
+            if (
+                i > 0
+                and i < len(transform_fns) - 1
+                and isinstance(cur_fn, BuildOutputBlocksMapTransformFn)
+            ):
+                prev_fn = transform_fns[i - 1]
+                next_fn = transform_fns[i + 1]
+                if (
+                    prev_fn.output_type == MapTransformFnDataType.Block
+                    and next_fn.input_type == MapTransformFnDataType.Block
+                ):
+                    drop = True
+            if not drop:
+                new_transform_fns.append(cur_fn)
+
+        return new_transform_fns

.venv/lib/python3.11/site-packages/ray/data/_internal/logical/util.py

@@ -0,0 +1,112 @@
+import json
+import re
+import threading
+from typing import Dict
+
+from ray._private.usage.usage_lib import TagKey, record_extra_usage_tag
+from ray.data._internal.logical.interfaces import LogicalOperator
+from ray.data._internal.logical.operators.map_operator import AbstractUDFMap
+from ray.data._internal.logical.operators.read_operator import Read
+from ray.data._internal.logical.operators.write_operator import Write
+
+# The dictionary for the operator name and count.
+_recorded_operators = dict()
+_recorded_operators_lock = threading.Lock()
+
+# The white list of operator names allowed to be recorded.
+_op_name_white_list = [
+    # Read
+    "ReadBigQuery",
+    "ReadRange",
+    "ReadMongo",
+    "ReadParquet",
+    "ReadParquetBulk",
+    "ReadImage",
+    "ReadJSON",
+    "ReadCSV",
+    "ReadText",
+    "ReadNumpy",
+    "ReadTFRecord",
+    "ReadBinary",
+    "ReadTorch",
+    "ReadAvro",
+    "ReadWebDataset",
+    "ReadSQL",
+    "ReadDatabricksUC",
+    "ReadLance",
+    "ReadHuggingFace",
+    "ReadCustom",
+    # From
+    "FromArrow",
+    "FromItems",
+    "FromNumpy",
+    "FromPandas",
+    # Write
+    "WriteBigQuery",
+    "WriteParquet",
+    "WriteJSON",
+    "WriteCSV",
+    "WriteTFRecord",
+    "WriteNumpy",
+    "WriteMongo",
+    "WriteWebDataset",
+    "WriteSQL",
+    "WriteCustom",
+    # Map
+    "Map",
+    "MapBatches",
+    "Filter",
+    "FlatMap",
+    # All-to-all
+    "RandomizeBlockOrder",
+    "RandomShuffle",
+    "Repartition",
+    "Sort",
+    "Aggregate",
+    # N-ary
+    "Zip",
+    "Union",
+]
+
+
+def record_operators_usage(op: LogicalOperator):
+    """Record logical operator usage with Ray telemetry."""
+    ops_dict = dict()
+    _collect_operators_to_dict(op, ops_dict)
+    ops_json_str = ""
+    with _recorded_operators_lock:
+        for op, count in ops_dict.items():
+            _recorded_operators.setdefault(op, 0)
+            _recorded_operators[op] += count
+        ops_json_str = json.dumps(_recorded_operators)
+
+    record_extra_usage_tag(TagKey.DATA_LOGICAL_OPS, ops_json_str)
+
+
+def _collect_operators_to_dict(op: LogicalOperator, ops_dict: Dict[str, int]):
+    """Collect the logical operator name and count into `ops_dict`."""
+    for child in op.input_dependencies:
+        _collect_operators_to_dict(child, ops_dict)
+
+    op_name = op.name
+
+    # Check read and write operator, and anonymize user-defined data source.
+    if isinstance(op, Read):
+        op_name = f"Read{op._datasource.get_name()}"
+        if op_name not in _op_name_white_list:
+            op_name = "ReadCustom"
+    elif isinstance(op, Write):
+        op_name = f"Write{op._datasink_or_legacy_datasource.get_name()}"
+        if op_name not in _op_name_white_list:
+            op_name = "WriteCustom"
+    elif isinstance(op, AbstractUDFMap):
+        # Remove the function name from the map operator name.
+        # E.g., Map(<lambda>) -> Map
+        op_name = re.sub("\\(.*\\)$", "", op_name)
+
+    # Anonymize any operator name if not in white list.
+    if op_name not in _op_name_white_list:
+        op_name = "Unknown"
+
+    ops_dict.setdefault(op_name, 0)
+    ops_dict[op_name] += 1