chore: import upstream snapshot with attribution

python/ray/data/_internal/logical/interfaces/__init__.py

@@ -0,0 +1,32 @@
+from .logical_operator import (
+    LogicalOperator,
+    LogicalOperatorPreservesSchema,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorSupportsPredicatePushdown,
+    LogicalOperatorSupportsProjectionPushdown,
+    LogicalOperatorUnifiesInputSchemas,
+    PredicatePassThroughBehavior,
+)
+from .logical_plan import LogicalPlan
+from .operator import Operator
+from .optimizer import Optimizer, Rule
+from .physical_plan import PhysicalPlan
+from .plan import Plan
+from .source_operator import SourceOperator
+
+__all__ = [
+    "LogicalOperator",
+    "LogicalPlan",
+    "Operator",
+    "Optimizer",
+    "PhysicalPlan",
+    "Plan",
+    "Rule",
+    "SourceOperator",
+    "LogicalOperatorPreservesSchema",
+    "LogicalOperatorSupportsProjectionPushdown",
+    "LogicalOperatorSupportsPredicatePushdown",
+    "LogicalOperatorSupportsPredicatePassThrough",
+    "LogicalOperatorUnifiesInputSchemas",
+    "PredicatePassThroughBehavior",
+]

python/ray/data/_internal/logical/interfaces/logical_operator.py

@@ -0,0 +1,247 @@
+import copy
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field, fields, replace
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional
+
+from .operator import Operator
+from ray.data.block import BlockMetadata
+from ray.data.expressions import Expr
+
+if TYPE_CHECKING:
+    from ray.data.block import Schema
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class LogicalOperator(Operator, ABC):
+    """Abstract class for logical operators.
+
+    A logical operator describes transformation, and later is converted into
+    physical operator.
+    """
+
+    _name: Optional[str] = field(init=False, default=None, repr=False)
+    _input_dependencies: List["LogicalOperator"] = field(
+        init=False, default_factory=list, repr=False
+    )
+
+    @property
+    def name(self) -> str:
+        return self._name or self.__class__.__name__
+
+    @property
+    def num_outputs(self) -> Optional[int]:
+        """Expected number of output blocks, if known."""
+        return None
+
+    def estimated_num_outputs(self) -> Optional[int]:
+        """Returns the estimated number of blocks that
+        would be outputted by this logical operator.
+
+        This method does not execute the plan, so it does not take into consideration
+        block splitting. This method only considers high-level block constraints like
+        `Dataset.repartition(num_blocks=X)`. A more accurate estimation can be given by
+        `PhysicalOperator.num_outputs_total()` during execution.
+        """
+        if self.num_outputs is not None:
+            return self.num_outputs
+        elif len(self.input_dependencies) == 1:
+            return self.input_dependencies[0].estimated_num_outputs()
+        return None
+
+    # Override the following 3 methods to correct type hints.
+
+    @property
+    def input_dependencies(self) -> List["LogicalOperator"]:
+        value = self._input_dependencies
+        for x in value:
+            assert isinstance(x, LogicalOperator), x
+        return value
+
+    @input_dependencies.setter
+    def input_dependencies(self, value: List["LogicalOperator"]) -> None:
+        for x in value:
+            assert isinstance(x, LogicalOperator), x
+        object.__setattr__(self, "_input_dependencies", value)
+
+    def post_order_iter(self) -> Iterator["LogicalOperator"]:
+        return super().post_order_iter()  # type: ignore
+
+    def _apply_transform(
+        self, transform: Callable[["LogicalOperator"], "LogicalOperator"]
+    ) -> "LogicalOperator":
+        input_dependencies = self.input_dependencies
+        transformed_inputs = [
+            input_op._apply_transform(transform) for input_op in input_dependencies
+        ]
+        if all(
+            transformed_input is input_op
+            for transformed_input, input_op in zip(
+                transformed_inputs, input_dependencies
+            )
+        ):
+            target = self
+        else:
+            target = self._with_new_input_dependencies(transformed_inputs)
+        return transform(target)
+
+    def _with_new_input_dependencies(
+        self, input_dependencies: List["LogicalOperator"]
+    ) -> "LogicalOperator":
+        if "input_dependencies" in {field.name for field in fields(self)}:
+            return replace(self, input_dependencies=input_dependencies)
+
+        target = copy.copy(self)
+        object.__setattr__(target, "_input_dependencies", input_dependencies)
+        return target
+
+    def _get_args(self) -> Dict[str, Any]:
+        """This Dict must be serializable"""
+        args: Dict[str, Any] = {}
+        for dataclass_field in fields(self):
+            key = dataclass_field.name
+            value = getattr(self, key)
+            # Keep underscore-prefixed keys to preserve legacy export schema.
+            args[key if key.startswith("_") else f"_{key}"] = value
+        args["_name"] = self.name
+        # Preserve legacy export shape even though output deps are no longer tracked.
+        args["_output_dependencies"] = []
+        # Do not include input dependencies, since we only want to export this
+        # operator-specific args. Adding input_dependencies isn't wrong, but can
+        # lead to slow recursive calls with `sanitize_for_struct`, since logical
+        # operators are dataclasses.
+        args["_input_dependencies"] = []
+        return args
+
+    def infer_schema(self) -> Optional["Schema"]:
+        """Returns the inferred schema of the output blocks."""
+        return None
+
+    def infer_metadata(self) -> "BlockMetadata":
+        """A ``BlockMetadata`` that represents the aggregate metadata of the outputs.
+
+        This method is used by methods like :meth:`~ray.data.Dataset.schema` to
+        efficiently return metadata.
+        """
+        return BlockMetadata(None, None, None, None)
+
+    def is_lineage_serializable(self) -> bool:
+        """Returns whether the lineage of this operator can be serialized.
+
+        An operator is lineage serializable if you can serialize it on one machine and
+        deserialize it on another without losing information. Operators that store
+        object references (e.g., ``InputData``) aren't lineage serializable because the
+        objects aren't available on the deserialized machine.
+        """
+        return True
+
+
+class LogicalOperatorSupportsProjectionPushdown(LogicalOperator):
+    """Mixin for reading operators supporting projection pushdown"""
+
+    def supports_projection_pushdown(self) -> bool:
+        return False
+
+    def get_projection_map(self) -> Optional[Dict[str, str]]:
+        return None
+
+    def apply_projection(
+        self,
+        projection_map: Optional[Dict[str, str]],
+    ) -> LogicalOperator:
+        return self
+
+
+class LogicalOperatorPreservesSchema(LogicalOperator):
+    """Mixin for operators whose output column layout is identical to their
+    single input's. Provides a default ``infer_schema()`` that delegates to
+    the input. Use for ops like ``Filter``, ``Sort``, ``Limit``, etc., that
+    only re-order or filter rows.
+
+    List this mixin last in the bases of subclasses so the concrete operator
+    base (e.g., ``AbstractMap``, ``AbstractAllToAll``) drives ``__init__`` /
+    ``super()`` chains.
+    """
+
+    def infer_schema(self) -> Optional["Schema"]:
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        return self.input_dependencies[0].infer_schema()
+
+
+class LogicalOperatorUnifiesInputSchemas(LogicalOperator):
+    """Mixin for n-ary operators whose output schema is the unification of
+    all inputs' schemas (e.g., ``Union``, ``Mix``). Provides a default
+    ``infer_schema()`` that returns the result of
+    ``unify_schemas_with_validation`` over each input's schema, or
+    ``None`` if any input's schema is unresolvable.
+
+    List this mixin last in the bases of subclasses so the concrete operator
+    base (e.g., ``NAry``) drives ``__init__`` / ``super()`` chains.
+    """
+
+    def infer_schema(self) -> Optional["Schema"]:
+        import pyarrow as pa
+
+        from ray.data._internal.util import unify_schemas_with_validation
+
+        input_schemas = [op.infer_schema() for op in self.input_dependencies]
+        if not all(isinstance(s, pa.Schema) for s in input_schemas):
+            return None
+        try:
+            return unify_schemas_with_validation(input_schemas)
+        except (pa.ArrowTypeError, pa.ArrowInvalid):
+            return None
+
+
+class LogicalOperatorSupportsPredicatePushdown(LogicalOperator):
+    """Mixin for reading operators supporting predicate pushdown"""
+
+    def supports_predicate_pushdown(self) -> bool:
+        return False
+
+    def get_current_predicate(self) -> Optional[Expr]:
+        return None
+
+    def apply_predicate(
+        self,
+        predicate_expr: Expr,
+    ) -> LogicalOperator:
+        return self
+
+
+class PredicatePassThroughBehavior(Enum):
+    """Defines how predicates can be passed through an operator."""
+
+    # Predicate can be pushed through as-is (e.g., Sort, Repartition, RandomShuffle, Limit)
+    PASSTHROUGH = "passthrough"
+
+    # Predicate can be pushed through but needs column rebinding (e.g., Project)
+    PASSTHROUGH_WITH_SUBSTITUTION = "passthrough_with_substitution"
+
+    # Predicate can be pushed into each branch (e.g., Union)
+    PUSH_INTO_BRANCHES = "push_into_branches"
+
+    # Predicate can be conditionally pushed based on columns (e.g., Join)
+    CONDITIONAL = "conditional"
+
+
+class LogicalOperatorSupportsPredicatePassThrough(ABC):
+    """Mixin for operators that allow predicates to be pushed through them.
+
+    This is distinct from LogicalOperatorSupportsPredicatePushdown, which is for
+    operators that can *accept* predicates (like Read). This trait is for operators
+    that allow predicates to *pass through* them.
+    """
+
+    @abstractmethod
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        """Returns the predicate passthrough behavior for this operator."""
+        pass
+
+    def get_column_substitutions(self) -> Optional[Dict[str, str]]:
+        """Returns column renames needed when pushing through (for PASSTHROUGH_WITH_SUBSTITUTION).
+
+        Returns:
+            Dict mapping from old_name -> new_name, or None if no rebinding needed
+        """
+        return None

python/ray/data/_internal/logical/interfaces/logical_plan.py

@@ -0,0 +1,56 @@
+from typing import TYPE_CHECKING, List, Optional
+
+from .logical_operator import LogicalOperator
+from .plan import Plan
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+
+class LogicalPlan(Plan):
+    """The plan with a DAG of logical operators."""
+
+    def __init__(self, dag: LogicalOperator, context: "DataContext"):
+        super().__init__(context)
+        self._dag = dag
+
+    @property
+    def dag(self) -> LogicalOperator:
+        """Get the DAG of logical operators."""
+        return self._dag
+
+    def sources(self) -> List[LogicalOperator]:
+        """List of operators that are sources for this plan's DAG."""
+        # If an operator has no input dependencies, it's a source.
+        if not any(self._dag.input_dependencies):
+            return [self._dag]
+
+        sources = []
+        for op in self._dag.input_dependencies:
+            sources.extend(LogicalPlan(op, self.context).sources())
+        return sources
+
+    def has_lazy_input(self) -> bool:
+        """Return whether this plan has lazy input blocks."""
+        from ray.data._internal.logical.operators import Read
+
+        return all(isinstance(op, Read) for op in self.sources())
+
+    def require_preserve_order(self) -> bool:
+        """Whether this plan requires to preserve order."""
+        from ray.data._internal.logical.operators import Zip
+
+        return any(isinstance(op, Zip) for op in self.dag.post_order_iter())
+
+    def input_files(self) -> Optional[List[str]]:
+        """Get the input files of the dataset, if available."""
+        input_files = self.dag.infer_metadata().input_files
+        if input_files is None:
+            return None
+        return list(set(input_files))
+
+    def initial_num_blocks(self) -> Optional[int]:
+        """Get the estimated number of blocks from the logical plan
+        after applying execution plan optimizations, but prior to
+        fully executing the dataset."""
+        return self.dag.estimated_num_outputs()

python/ray/data/_internal/logical/interfaces/operator.py

@@ -0,0 +1,56 @@
+from abc import ABC, abstractmethod
+from typing import Callable, Iterator, List
+
+
+class Operator(ABC):
+    """Abstract class for operators.
+
+    Operators live on the driver side of the Dataset only.
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Name for this operator."""
+        ...
+
+    @property
+    def dag_str(self) -> str:
+        """String representation of the whole DAG."""
+        if self.input_dependencies:
+            out_str = ", ".join([x.dag_str for x in self.input_dependencies])
+            out_str += " -> "
+        else:
+            out_str = ""
+        out_str += f"{self.__class__.__name__}[{self.name}]"
+        return out_str
+
+    @property
+    @abstractmethod
+    def input_dependencies(self) -> List["Operator"]:
+        """List of operators that provide inputs for this operator."""
+        ...
+
+    def post_order_iter(self) -> Iterator["Operator"]:
+        """Depth-first traversal of this operator and its input dependencies."""
+        for op in self.input_dependencies:
+            yield from op.post_order_iter()
+        yield self
+
+    @abstractmethod
+    def _apply_transform(
+        self, transform: Callable[["Operator"], "Operator"]
+    ) -> "Operator":
+        """Recursively applies transformation (in post-order) to the operators DAG
+
+        NOTE: This operation should be opting in to avoid in-place modifications,
+              instead creating new operations whenever any operator needs to be
+              updated.
+        """
+        ...
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}[{self.name}]"
+
+    def __str__(self) -> str:
+        return repr(self)

python/ray/data/_internal/logical/interfaces/optimizer.py

@@ -0,0 +1,52 @@
+from typing import List, Type
+
+from .plan import Plan
+
+
+class Rule:
+    """Abstract class for optimization rule."""
+
+    def apply(self, plan: Plan) -> Plan:
+        """Apply the optimization rule to the execution plan."""
+        raise NotImplementedError
+
+    @classmethod
+    def dependencies(cls) -> List[Type["Rule"]]:
+        """List of rules that must be applied before this rule."""
+        return []
+
+    @classmethod
+    def dependents(cls) -> List[Type["Rule"]]:
+        """List of rules that must be applied after this rule."""
+        return []
+
+
+class Optimizer:
+    """Abstract class for optimizers.
+
+    An optimizers transforms a DAG of operators with a list of predefined rules.
+    """
+
+    @property
+    def rules(self) -> List[Rule]:
+        """List of predefined rules for this optimizer."""
+        raise NotImplementedError
+
+    def optimize(self, plan: Plan) -> Plan:
+        """Optimize operators with a list of rules."""
+        # Apply rules until the plan is not changed
+        previous_plan = plan
+        while True:
+            for rule in self.rules:
+                plan = rule.apply(plan)
+            # TODO: Eventually we should implement proper equality.
+            # Using str to check equality seems brittle
+            if plan.dag.dag_str == previous_plan.dag.dag_str:
+                break
+            previous_plan = plan
+        return self._post_optimize(plan)
+
+    def _post_optimize(self, plan: Plan) -> Plan:
+        """Post optimize is used for rules or other post-processing
+        that needs to be executed only once after the `optimize` loop."""
+        return plan

python/ray/data/_internal/logical/interfaces/physical_plan.py

@@ -0,0 +1,34 @@
+from typing import TYPE_CHECKING, Dict
+
+from .logical_operator import LogicalOperator
+from .plan import Plan
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import PhysicalOperator
+    from ray.data.context import DataContext
+
+
+class PhysicalPlan(Plan):
+    """The plan with a DAG of physical operators."""
+
+    def __init__(
+        self,
+        dag: "PhysicalOperator",
+        op_map: Dict["PhysicalOperator", LogicalOperator],
+        context: "DataContext",
+    ):
+        super().__init__(context)
+        self._dag = dag
+        self._op_map = op_map
+
+    @property
+    def dag(self) -> "PhysicalOperator":
+        """Get the DAG of physical operators."""
+        return self._dag
+
+    @property
+    def op_map(self) -> Dict["PhysicalOperator", LogicalOperator]:
+        """
+        Get a mapping from physical operators to their corresponding logical operator.
+        """
+        return self._op_map

python/ray/data/_internal/logical/interfaces/plan.py

@@ -0,0 +1,29 @@
+from typing import TYPE_CHECKING
+
+from .operator import Operator
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+
+class Plan:
+    """Abstract class for logical/physical execution plans.
+
+    This plan should hold an operator representing the plan DAG and any auxiliary data
+    that's useful for plan optimization or execution.
+    """
+
+    def __init__(self, context: "DataContext"):
+        self._context = context
+
+    @property
+    def dag(self) -> Operator:
+        raise NotImplementedError
+
+    @property
+    def context(self) -> "DataContext":
+        return self._context
+
+    @context.setter
+    def context(self, context: "DataContext") -> None:
+        self._context = context

python/ray/data/_internal/logical/interfaces/source_operator.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+
+    from ray.data.dataset import RefBundle
+
+
+class SourceOperator(ABC):
+    """Mixin for Logical operators that can be logical source nodes.
+    Subclasses: Read, InputData, FromAbstract.
+    """
+
+    @abstractmethod
+    def output_data(self) -> Optional[List["RefBundle"]]:
+        """The output data of this operator if already known, or ``None``."""
+        pass

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

@@ -0,0 +1,93 @@
+"""Expose logical operator classes in ray.data._internal.logical.operators."""
+
+from ray.data._internal.logical.operators.all_to_all_operator import (
+    AbstractAllToAll,
+    Aggregate,
+    RandomizeBlocks,
+    RandomShuffle,
+    Repartition,
+    Sort,
+)
+from ray.data._internal.logical.operators.count_operator import Count
+from ray.data._internal.logical.operators.from_operators import (
+    AbstractFrom,
+    FromArrow,
+    FromBlocks,
+    FromItems,
+    FromNumpy,
+    FromPandas,
+)
+from ray.data._internal.logical.operators.input_data_operator import InputData
+from ray.data._internal.logical.operators.join_operator import Join, JoinSide, JoinType
+from ray.data._internal.logical.operators.map_operator import (
+    CSE_TEMP_COLUMN_PREFIX,
+    AbstractMap,
+    AbstractUDFMap,
+    Filter,
+    FlatMap,
+    MapBatches,
+    MapRows,
+    Project,
+    StreamingRepartition,
+)
+from ray.data._internal.logical.operators.n_ary_operator import (
+    Mix,
+    MixStoppingCondition,
+    NAry,
+    Union,
+    Zip,
+)
+from ray.data._internal.logical.operators.one_to_one_operator import (
+    AbstractOneToOne,
+    Download,
+    Limit,
+)
+from ray.data._internal.logical.operators.read_operator import (
+    ListFiles,
+    Read,
+    ReadFiles,
+)
+from ray.data._internal.logical.operators.streaming_split_operator import StreamingSplit
+from ray.data._internal.logical.operators.write_operator import Write
+
+__all__ = [
+    "AbstractAllToAll",
+    "AbstractFrom",
+    "AbstractMap",
+    "AbstractOneToOne",
+    "AbstractUDFMap",
+    "Aggregate",
+    "CSE_TEMP_COLUMN_PREFIX",
+    "Count",
+    "Download",
+    "Filter",
+    "FlatMap",
+    "FromArrow",
+    "FromBlocks",
+    "FromItems",
+    "FromNumpy",
+    "FromPandas",
+    "InputData",
+    "Join",
+    "JoinSide",
+    "JoinType",
+    "Limit",
+    "ListFiles",
+    "MapBatches",
+    "MapRows",
+    "Mix",
+    "MixStoppingCondition",
+    "NAry",
+    "Project",
+    "RandomShuffle",
+    "RandomizeBlocks",
+    "Read",
+    "ReadFiles",
+    "Repartition",
+    "Sort",
+    "StreamingRepartition",
+    "StreamingSplit",
+    "Union",
+    "Write",
+    "Zip",
+]

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

@@ -0,0 +1,265 @@
+from dataclasses import InitVar, dataclass, field, replace
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorPreservesSchema,
+    LogicalOperatorSupportsPredicatePassThrough,
+    PredicatePassThroughBehavior,
+)
+from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
+from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey, SortTaskSpec
+from ray.data._internal.random_config import RandomSeedConfig
+from ray.data.aggregate import AggregateFn
+from ray.data.block import BlockMetadata
+
+if TYPE_CHECKING:
+    from ray.data.block import Schema
+
+__all__ = [
+    "AbstractAllToAll",
+    "Aggregate",
+    "RandomShuffle",
+    "RandomizeBlocks",
+    "Repartition",
+    "Sort",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class AbstractAllToAll(LogicalOperator):
+    """Abstract class for logical operators should be converted to physical
+    AllToAllOperator.
+    """
+
+    def __init__(
+        self,
+        input_dependencies: List[LogicalOperator],
+        sub_progress_bar_names: Optional[List[str]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+        *,
+        name: Optional[str] = None,
+    ):
+        """Initialize an ``AbstractAllToAll`` logical operator.
+
+        Args:
+            input_dependencies: The operators preceding this operator in the plan DAG.
+                The outputs of these operators will be the inputs to this operator.
+            sub_progress_bar_names: Optional sub-stage progress bar names for this
+                operator.
+            ray_remote_args: Args to provide to :func:`ray.remote`.
+            name: Name for this operator. This is the name that will appear when
+                inspecting the logical plan of a Dataset.
+        """
+        object.__setattr__(self, "_input_dependencies", list(input_dependencies))
+        if name is not None:
+            object.__setattr__(self, "_name", name)
+        object.__setattr__(self, "ray_remote_args", ray_remote_args or {})
+        object.__setattr__(self, "sub_progress_bar_names", sub_progress_bar_names)
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class RandomizeBlocks(
+    AbstractAllToAll,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorPreservesSchema,
+):
+    """Logical operator for randomize_block_order."""
+
+    seed_config: Optional[RandomSeedConfig] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    sub_progress_bar_names: Optional[List[str]] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if self.seed_config is None:
+            object.__setattr__(self, "seed_config", RandomSeedConfig())
+        object.__setattr__(self, "_name", "RandomizeBlockOrder")
+
+    def infer_metadata(self) -> "BlockMetadata":
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        assert isinstance(self.input_dependencies[0], LogicalOperator)
+        return self.input_dependencies[0].infer_metadata()
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        # Randomizing block order doesn't affect filtering correctness
+        return PredicatePassThroughBehavior.PASSTHROUGH
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class RandomShuffle(
+    AbstractAllToAll,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorPreservesSchema,
+):
+    """Logical operator for random_shuffle."""
+
+    name: InitVar[str] = "RandomShuffle"
+    seed_config: Optional[RandomSeedConfig] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    sub_progress_bar_names: Optional[List[str]] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+
+    def __post_init__(self, name: str):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if self.seed_config is None:
+            object.__setattr__(self, "seed_config", RandomSeedConfig())
+        if self.sub_progress_bar_names is None:
+            object.__setattr__(
+                self,
+                "sub_progress_bar_names",
+                [
+                    ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME,
+                    ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME,
+                ],
+            )
+        object.__setattr__(self, "_name", name)
+
+    def _with_new_input_dependencies(
+        self, input_dependencies: List[LogicalOperator]
+    ) -> LogicalOperator:
+        return replace(self, input_dependencies=input_dependencies, name=self._name)
+
+    def infer_metadata(self) -> "BlockMetadata":
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        assert isinstance(self.input_dependencies[0], LogicalOperator)
+        return self.input_dependencies[0].infer_metadata()
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        # Random shuffle doesn't affect filtering correctness
+        return PredicatePassThroughBehavior.PASSTHROUGH
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Repartition(
+    AbstractAllToAll,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorPreservesSchema,
+):
+    """Logical operator for repartition."""
+
+    shuffle: bool = False
+    keys: Optional[List[str]] = None
+    sort: bool = False
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    sub_progress_bar_names: Optional[List[str]] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+    num_outputs: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if self.shuffle:
+            sub_progress_bar_names = [
+                ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME,
+                ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME,
+            ]
+        else:
+            sub_progress_bar_names = [
+                ShuffleTaskSpec.SPLIT_REPARTITION_SUB_PROGRESS_BAR_NAME,
+            ]
+        object.__setattr__(self, "sub_progress_bar_names", sub_progress_bar_names)
+
+    def _with_new_input_dependencies(
+        self, input_dependencies: List[LogicalOperator]
+    ) -> LogicalOperator:
+        return replace(
+            self,
+            input_dependencies=input_dependencies,
+            num_outputs=self.num_outputs,
+        )
+
+    def infer_metadata(self) -> "BlockMetadata":
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        assert isinstance(self.input_dependencies[0], LogicalOperator)
+        return self.input_dependencies[0].infer_metadata()
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        # Repartition doesn't affect filtering correctness
+        return PredicatePassThroughBehavior.PASSTHROUGH
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Sort(
+    AbstractAllToAll,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorPreservesSchema,
+):
+    """Logical operator for sort."""
+
+    sort_key: SortKey
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    sub_progress_bar_names: Optional[List[str]] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        object.__setattr__(
+            self,
+            "sub_progress_bar_names",
+            [
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME,
+                ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME,
+                ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME,
+            ],
+        )
+
+    def infer_metadata(self) -> "BlockMetadata":
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        assert isinstance(self.input_dependencies[0], LogicalOperator)
+        return self.input_dependencies[0].infer_metadata()
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        # Sort doesn't affect filtering correctness
+        return PredicatePassThroughBehavior.PASSTHROUGH
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Aggregate(AbstractAllToAll):
+    """Logical operator for aggregate."""
+
+    key: Optional[str | List[str]]
+    aggs: List[AggregateFn]
+    num_partitions: Optional[int] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    sub_progress_bar_names: Optional[List[str]] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        object.__setattr__(
+            self,
+            "sub_progress_bar_names",
+            [
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME,
+                ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME,
+                ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME,
+            ],
+        )
+
+    def infer_schema(self) -> Optional["Schema"]:
+        # Output = key field(s) from input schema + one field per aggregator.
+        # Returns None if any aggregator can't declare its output field
+        # (callers fall back to limit(1)).
+        import pyarrow as pa
+
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        input_schema = self.input_dependencies[0].infer_schema()
+        if not isinstance(input_schema, pa.Schema):
+            return None
+
+        fields: List[pa.Field] = []
+        if self.key is not None:
+            keys = self.key if isinstance(self.key, list) else [self.key]
+            for key in keys:
+                try:
+                    fields.append(input_schema.field(key))
+                except (KeyError, TypeError, ValueError):
+                    return None
+        for agg in self.aggs:
+            f = agg.output_field(input_schema)
+            if f is None:
+                return None
+            fields.append(f)
+        return pa.schema(fields)

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

@@ -0,0 +1,35 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional
+
+from ray.data._internal.logical.interfaces import LogicalOperator
+
+if TYPE_CHECKING:
+    from ray.data.block import Schema
+
+__all__ = [
+    "Count",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Count(LogicalOperator):
+    """Logical operator that represents counting the number of rows in inputs.
+
+    Physical operators that implement this logical operator should produce one or more
+    rows with a single column named `Count.COLUMN_NAME`. When you sum the values in
+    this column, you should get the total number of rows in the dataset.
+    """
+
+    COLUMN_NAME = "__num_rows"
+
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+
+    def infer_schema(self) -> Optional["Schema"]:
+        # Fixed output: one row per partial count with a single ``__num_rows``
+        # int64 column.
+        import pyarrow as pa
+
+        return pa.schema([pa.field(self.COLUMN_NAME, pa.int64(), nullable=False)])

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

@@ -0,0 +1,133 @@
+import abc
+import functools
+from dataclasses import InitVar, dataclass, field
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
+from ray.data._internal.logical.interfaces import LogicalOperator, SourceOperator
+from ray.data._internal.util import unify_ref_bundles_schema
+from ray.data.block import (
+    Block,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+)
+from ray.types import ObjectRef
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+    ArrowTable = Union["pa.Table", bytes]
+
+__all__ = [
+    "AbstractFrom",
+    "FromArrow",
+    "FromBlocks",
+    "FromItems",
+    "FromNumpy",
+    "FromPandas",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class AbstractFrom(LogicalOperator, SourceOperator, metaclass=abc.ABCMeta):
+    """Abstract logical operator for `from_*`."""
+
+    input_blocks: InitVar[List[ObjectRef[Block]]]
+    input_metadata: InitVar[List[BlockMetadataWithSchema]]
+    input_data: List[RefBundle] = field(init=False)
+    _input_dependencies: list[LogicalOperator] = field(
+        init=False, repr=False, default_factory=list
+    )
+
+    def __post_init__(
+        self,
+        input_blocks: List[ObjectRef[Block]],
+        input_metadata: List[BlockMetadataWithSchema],
+    ):
+        assert len(input_blocks) == len(input_metadata), (
+            len(input_blocks),
+            len(input_metadata),
+        )
+
+        # `owns_blocks` is False because this op may be shared by multiple Datasets.
+        object.__setattr__(
+            self,
+            "input_data",
+            [
+                RefBundle(
+                    [BlockEntry(input_blocks[i], input_metadata[i])],
+                    owns_blocks=False,
+                    schema=input_metadata[i].schema,
+                )
+                for i in range(len(input_blocks))
+            ],
+        )
+
+    def output_data(self) -> Optional[List[RefBundle]]:
+        return self.input_data
+
+    @property
+    def num_outputs(self) -> Optional[int]:
+        return len(self.input_data)
+
+    @functools.cached_property
+    def _cached_output_metadata(self) -> BlockMetadata:
+        return BlockMetadata(
+            num_rows=self._num_rows(),
+            size_bytes=self._size_bytes(),
+            input_files=None,
+            exec_stats=None,
+        )
+
+    def _num_rows(self):
+        if all(bundle.num_rows() is not None for bundle in self.input_data):
+            return sum(bundle.num_rows() for bundle in self.input_data)
+        else:
+            return None
+
+    def _size_bytes(self):
+        metadata = [m for bundle in self.input_data for m in bundle.metadata]
+        if all(m.size_bytes is not None for m in metadata):
+            return sum(m.size_bytes for m in metadata)
+        else:
+            return None
+
+    def infer_metadata(self) -> BlockMetadata:
+        return self._cached_output_metadata
+
+    def infer_schema(self):
+        return unify_ref_bundles_schema(self.input_data)
+
+    def is_lineage_serializable(self) -> bool:
+        # This operator isn't serializable because it contains ObjectRefs.
+        return False
+
+
+class FromItems(AbstractFrom):
+    """Logical operator for `from_items`."""
+
+    pass
+
+
+class FromBlocks(AbstractFrom):
+    """Logical operator for `from_blocks`."""
+
+    pass
+
+
+class FromNumpy(AbstractFrom):
+    """Logical operator for `from_numpy`."""
+
+    pass
+
+
+class FromArrow(AbstractFrom):
+    """Logical operator for `from_arrow`."""
+
+    pass
+
+
+class FromPandas(AbstractFrom):
+    """Logical operator for `from_pandas`."""
+
+    pass

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

@@ -0,0 +1,64 @@
+import functools
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from ray.data._internal.execution.interfaces import RefBundle
+from ray.data._internal.logical.interfaces import LogicalOperator, SourceOperator
+from ray.data._internal.util import unify_schemas_with_validation
+from ray.data.block import BlockMetadata
+
+__all__ = [
+    "InputData",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class InputData(LogicalOperator, SourceOperator):
+    """Logical operator for input data.
+
+    This may hold cached blocks from a previous Dataset execution.
+    """
+
+    input_data: List[RefBundle]
+    _input_dependencies: list[LogicalOperator] = field(
+        init=False, repr=False, default_factory=list
+    )
+
+    def output_data(self) -> Optional[List[RefBundle]]:
+        return self.input_data
+
+    @property
+    def num_outputs(self) -> Optional[int]:
+        return len(self.input_data)
+
+    def infer_metadata(self) -> BlockMetadata:
+        return self._cached_output_metadata
+
+    @functools.cached_property
+    def _cached_output_metadata(self) -> BlockMetadata:
+        return BlockMetadata(
+            num_rows=self._num_rows(),
+            size_bytes=self._size_bytes(),
+            input_files=None,
+            exec_stats=None,
+        )
+
+    def _num_rows(self):
+        if all(bundle.num_rows() is not None for bundle in self.input_data):
+            return sum(bundle.num_rows() for bundle in self.input_data)
+        else:
+            return None
+
+    def _size_bytes(self):
+        metadata = [m for bundle in self.input_data for m in bundle.metadata]
+        if all(m.size_bytes is not None for m in metadata):
+            return sum(m.size_bytes for m in metadata)
+        else:
+            return None
+
+    def infer_schema(self):
+        return unify_schemas_with_validation([data.schema for data in self.input_data])
+
+    def is_lineage_serializable(self) -> bool:
+        # This operator isn't serializable because it contains ObjectRefs.
+        return False

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

@@ -0,0 +1,263 @@
+from dataclasses import InitVar, dataclass, field, replace
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorSupportsPredicatePassThrough,
+    PredicatePassThroughBehavior,
+)
+from ray.data._internal.logical.operators.n_ary_operator import NAry
+
+if TYPE_CHECKING:
+    from ray.data.dataset import Schema
+    from ray.data.expressions import Expr
+
+__all__ = [
+    "Join",
+    "JoinSide",
+    "JoinType",
+]
+
+
+class JoinType(Enum):
+    INNER = "inner"
+    LEFT_OUTER = "left_outer"
+    RIGHT_OUTER = "right_outer"
+    FULL_OUTER = "full_outer"
+    LEFT_SEMI = "left_semi"
+    RIGHT_SEMI = "right_semi"
+    LEFT_ANTI = "left_anti"
+    RIGHT_ANTI = "right_anti"
+
+
+class JoinSide(Enum):
+    """Represents which side of a join to push a predicate to.
+
+    The enum values correspond to branch indices (0 for left, 1 for right).
+    """
+
+    LEFT = 0
+    RIGHT = 1
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Join(NAry, LogicalOperatorSupportsPredicatePassThrough):
+    """Logical operator for join."""
+
+    left_input_op: InitVar[LogicalOperator]
+    right_input_op: InitVar[LogicalOperator]
+    join_type: Union[JoinType, str]
+    left_key_columns: Tuple[str]
+    right_key_columns: Tuple[str]
+    num_partitions: int
+    left_columns_suffix: Optional[str] = None
+    right_columns_suffix: Optional[str] = None
+    partition_size_hint: Optional[int] = None
+    aggregator_ray_remote_args: Optional[Dict[str, Any]] = None
+    _input_dependencies: list[LogicalOperator] = field(init=False, repr=False)
+
+    def __post_init__(
+        self,
+        left_input_op: LogicalOperator,
+        right_input_op: LogicalOperator,
+    ):
+        try:
+            join_type_enum = JoinType(self.join_type)
+        except ValueError:
+            raise ValueError(
+                f"Invalid join type: '{self.join_type}'. "
+                f"Supported join types are: {', '.join(jt.value for jt in JoinType)}."
+            )
+
+        object.__setattr__(self, "join_type", join_type_enum)
+        object.__setattr__(
+            self,
+            "_input_dependencies",
+            [left_input_op, right_input_op],
+        )
+        # Mirror the reduce-task remote args onto `ray_remote_args` so fusing a downstream map
+        # into this join's reduce task respects the reduce task's actual args.
+        object.__setattr__(
+            self, "ray_remote_args", dict(self.aggregator_ray_remote_args or {})
+        )
+
+    @property
+    def num_outputs(self) -> Optional[int]:
+        return self.num_partitions
+
+    def _with_new_input_dependencies(
+        self, input_dependencies: List[LogicalOperator]
+    ) -> LogicalOperator:
+        return replace(
+            self,
+            left_input_op=input_dependencies[0],
+            right_input_op=input_dependencies[1],
+            num_partitions=self.num_outputs,
+        )
+
+    @staticmethod
+    def _validate_schemas(
+        left_op_schema: "Schema",
+        right_op_schema: "Schema",
+        left_key_column_names: Tuple[str],
+        right_key_column_names: Tuple[str],
+    ):
+        def _col_names_as_str(keys: Sequence[str]):
+            keys_joined = ", ".join(map(lambda k: f"'{k}'", keys))
+            return f"[{keys_joined}]"
+
+        if len(left_key_column_names) < 1:
+            raise ValueError(
+                f"At least 1 column name to join on has to be provided (got "
+                f"{_col_names_as_str(left_key_column_names)})"
+            )
+
+        if len(left_key_column_names) != len(right_key_column_names):
+            raise ValueError(
+                f"Number of columns provided for left and right datasets has to match "
+                f"(got {_col_names_as_str(left_key_column_names)} and "
+                f"{_col_names_as_str(right_key_column_names)})"
+            )
+
+        def _get_key_column_types(schema: "Schema", keys: Tuple[str]):
+            return (
+                [
+                    _type
+                    for name, _type in zip(schema.names, schema.types)
+                    if name in keys
+                ]
+                if schema
+                else None
+            )
+
+        right_op_key_cols = _get_key_column_types(
+            right_op_schema, left_key_column_names
+        )
+        left_op_key_cols = _get_key_column_types(left_op_schema, right_key_column_names)
+
+        if left_op_key_cols != right_op_key_cols:
+            raise ValueError(
+                f"Key columns are expected to be present and have the same types "
+                "in both left and right operands of the join operation: "
+                f"left has {left_op_schema}, but right has {right_op_schema}"
+            )
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        return PredicatePassThroughBehavior.CONDITIONAL
+
+    def which_side_to_push_predicate(
+        self, predicate_expr: "Expr"
+    ) -> Optional[JoinSide]:
+        """Determine which side of the join to push a predicate to.
+
+        Returns the side to push to, or None if pushdown is not safe.
+
+        Predicate pushdown is safe for:
+        - INNER: Can push to either side
+        - LEFT_OUTER/SEMI/ANTI: Can push to left side (preserved/output side)
+        - RIGHT_OUTER/SEMI/ANTI: Can push to right side (preserved/output side)
+        - FULL_OUTER: Cannot push (both sides can generate nulls)
+
+        The predicate must reference columns from exactly one side of the join,
+        OR reference only join key columns that all exist on one side.
+        """
+        # Get predicate columns and schemas
+        predicate_columns = self._get_referenced_columns(predicate_expr)
+        left_schema = self.input_dependencies[0].infer_schema()
+        right_schema = self.input_dependencies[1].infer_schema()
+
+        if not left_schema or not right_schema:
+            return None
+
+        # Get column sets for each side
+        left_columns = set(left_schema.names)
+        right_columns = set(right_schema.names)
+        left_join_keys = set(self.left_key_columns)
+        right_join_keys = set(self.right_key_columns)
+
+        # Get pushdown rules for this join type
+        can_push_left, can_push_right = self._get_pushdown_rules()
+
+        # Check if predicate can be evaluated on left side
+        # Condition: ALL predicate columns must exist on left (either as regular columns or join keys)
+        can_evaluate_on_left = predicate_columns.issubset(
+            left_columns
+        ) or predicate_columns.issubset(left_join_keys)
+        if can_evaluate_on_left and can_push_left:
+            return JoinSide.LEFT
+
+        # Check if predicate can be evaluated on right side
+        can_evaluate_on_right = predicate_columns.issubset(
+            right_columns
+        ) or predicate_columns.issubset(right_join_keys)
+        if can_evaluate_on_right and can_push_right:
+            return JoinSide.RIGHT
+
+        # Cannot push down
+        return None
+
+    def _get_pushdown_rules(self) -> Tuple[bool, bool]:
+        """Get pushdown rules for the current join type.
+
+        Returns:
+            Tuple of (can_push_left, can_push_right) indicating which sides
+            can accept predicate pushdown for this join type.
+        """
+        pushdown_rules = {
+            JoinType.INNER: (True, True),
+            JoinType.LEFT_OUTER: (True, False),
+            JoinType.RIGHT_OUTER: (False, True),
+            JoinType.LEFT_SEMI: (True, False),
+            JoinType.RIGHT_SEMI: (False, True),
+            JoinType.LEFT_ANTI: (True, False),
+            JoinType.RIGHT_ANTI: (False, True),
+            JoinType.FULL_OUTER: (False, False),
+        }
+        return pushdown_rules.get(self.join_type, (False, False))
+
+    def _get_referenced_columns(self, expr: "Expr") -> set[str]:
+        """Extract all column names referenced in an expression."""
+        from ray.data._internal.planner.plan_expression.expression_visitors import (
+            _ColumnReferenceCollector,
+        )
+
+        visitor = _ColumnReferenceCollector()
+        visitor.visit(expr)
+        return set(visitor.get_column_refs())
+
+    def infer_schema(self) -> Optional["Schema"]:
+        """Infer the output schema by running the shared ``join_tables``
+        utility on empty tables built from the input schemas. The same
+        utility runs at execution time, so plan-time and runtime schemas
+        agree by construction.
+        """
+        import pyarrow as pa
+
+        from ray.data._internal.execution.operators.join import join_tables
+
+        left_schema = self.input_dependencies[0].infer_schema()
+        right_schema = self.input_dependencies[1].infer_schema()
+        if not isinstance(left_schema, pa.Schema) or not isinstance(
+            right_schema, pa.Schema
+        ):
+            return None
+
+        join_type_enum = (
+            self.join_type
+            if isinstance(self.join_type, JoinType)
+            else JoinType(self.join_type)
+        )
+        try:
+            joined = join_tables(
+                left_schema.empty_table(),
+                right_schema.empty_table(),
+                join_type=join_type_enum,
+                left_key_col_names=tuple(self.left_key_columns),
+                right_key_col_names=tuple(self.right_key_columns),
+                left_columns_suffix=self.left_columns_suffix,
+                right_columns_suffix=self.right_columns_suffix,
+            )
+        except (pa.ArrowTypeError, pa.ArrowInvalid, pa.ArrowKeyError, ValueError):
+            return None
+        return joined.schema

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

@@ -0,0 +1,552 @@
+import functools
+import inspect
+import logging
+from dataclasses import dataclass, field
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Union,
+)
+
+from ray.data._internal.compute import ComputeStrategy, TaskPoolStrategy
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorPreservesSchema,
+    LogicalOperatorSupportsPredicatePassThrough,
+    PredicatePassThroughBehavior,
+)
+from ray.data._internal.logical.operators.one_to_one_operator import AbstractOneToOne
+from ray.data.block import UserDefinedFunction
+from ray.data.expressions import (
+    Expr,
+    StarExpr,
+    expand_star_exprs,
+    exprlist_to_fields,
+)
+from ray.data.preprocessor import Preprocessor
+
+if TYPE_CHECKING:
+    from ray.data.block import Schema
+
+__all__ = [
+    "AbstractMap",
+    "AbstractUDFMap",
+    "Filter",
+    "FlatMap",
+    "MapBatches",
+    "MapRows",
+    "Project",
+    "StreamingRepartition",
+    "CSE_TEMP_COLUMN_PREFIX",
+]
+
+
+logger = logging.getLogger(__name__)
+
+CSE_TEMP_COLUMN_PREFIX = "__ray_data_cse_"
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class AbstractMap(AbstractOneToOne):
+    """Abstract class for logical operators that should be converted to physical
+    MapOperator.
+    """
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        input_dependencies: Optional[List[LogicalOperator]] = None,
+        *,
+        can_modify_num_rows: bool,
+        min_rows_per_bundled_input: Optional[int] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        compute: Optional[ComputeStrategy] = None,
+    ):
+        """Initialize an ``AbstractMap`` logical operator that will later
+        be converted into a physical ``MapOperator``.
+
+        Args:
+            name: Name for this operator. This is the name that will appear when
+                inspecting the logical plan of a Dataset.
+            input_dependencies: The operators preceding this operator in the plan
+                DAG. The outputs of these operators will be the inputs to this
+                operator.
+            can_modify_num_rows: Whether the operator can change the row count. False if
+                # of input rows = # of output rows. True otherwise.
+            min_rows_per_bundled_input: Minimum number of rows a single bundle of
+                blocks passed on to the task must possess.
+            ray_remote_args: Args to provide to :func:`ray.remote`.
+            ray_remote_args_fn: A function that returns a dictionary of remote
+                args passed to each map worker. The purpose of this argument is
+                to generate dynamic arguments for each actor/task, and it will
+                be called each time prior to initializing the worker. Args
+                returned from this dict always override the args in
+                ``ray_remote_args``. Note: this is an advanced, experimental
+                feature.
+            compute: The compute strategy, either ``TaskPoolStrategy`` (default)
+                to use Ray tasks, or ``ActorPoolStrategy`` to use an
+                autoscaling actor pool.
+        """
+        super().__init__(
+            input_dependencies=input_dependencies,
+            can_modify_num_rows=can_modify_num_rows,
+            name=name,
+        )
+        object.__setattr__(
+            self, "min_rows_per_bundled_input", min_rows_per_bundled_input
+        )
+        object.__setattr__(self, "ray_remote_args", ray_remote_args or {})
+        object.__setattr__(self, "ray_remote_args_fn", ray_remote_args_fn)
+        object.__setattr__(self, "compute", compute or TaskPoolStrategy())
+        object.__setattr__(self, "per_block_limit", None)
+
+    def set_per_block_limit(self, per_block_limit: int):
+        object.__setattr__(self, "per_block_limit", per_block_limit)
+
+    def _get_args(self) -> Dict[str, Any]:
+        args = super()._get_args()
+        for key in [
+            "can_modify_num_rows",
+            "min_rows_per_bundled_input",
+            "ray_remote_args",
+            "ray_remote_args_fn",
+            "compute",
+            "per_block_limit",
+        ]:
+            args[f"_{key}"] = getattr(self, key)
+        return args
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class AbstractUDFMap(AbstractMap):
+    """Abstract class for logical operators performing a UDF that should be converted
+    to physical MapOperator.
+    """
+
+    fn: UserDefinedFunction
+    fn_args: Optional[Iterable[Any]] = None
+    fn_kwargs: Optional[Dict[str, Any]] = None
+    fn_constructor_args: Optional[Iterable[Any]] = None
+    fn_constructor_kwargs: Optional[Dict[str, Any]] = None
+    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
+
+    def __init__(
+        self,
+        name: str,
+        input_dependencies: List[LogicalOperator],
+        fn: UserDefinedFunction,
+        *,
+        can_modify_num_rows: bool,
+        fn_args: Optional[Iterable[Any]] = None,
+        fn_kwargs: Optional[Dict[str, Any]] = None,
+        fn_constructor_args: Optional[Iterable[Any]] = None,
+        fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
+        min_rows_per_bundled_input: Optional[int] = None,
+        compute: Optional[ComputeStrategy] = None,
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+    ):
+        """Initialize AbstractUDFMap.
+
+        Args:
+            name: Name for this operator. This is the name that will appear when
+                inspecting the logical plan of a Dataset.
+            input_dependencies: The operators preceding this operator in the plan DAG.
+                The outputs of these operators will be the inputs to this operator.
+            fn: User-defined function to be called.
+            can_modify_num_rows: Whether the UDF can change the row count. False if
+                # of input rows = # of output rows. True otherwise.
+            fn_args: Arguments to `fn`.
+            fn_kwargs: Keyword arguments to `fn`.
+            fn_constructor_args: Arguments to provide to the initializor of `fn` if
+                `fn` is a callable class.
+            fn_constructor_kwargs: Keyword Arguments to provide to the initializor of
+                `fn` if `fn` is a callable class.
+            min_rows_per_bundled_input: The target number of rows to pass to
+                ``MapOperator._add_bundled_input()``.
+            compute: The compute strategy, either ``TaskPoolStrategy`` (default) to use
+                Ray tasks, or ``ActorPoolStrategy`` to use an autoscaling actor pool.
+            ray_remote_args_fn: A function that returns a dictionary of remote args
+                passed to each map worker. The purpose of this argument is to generate
+                dynamic arguments for each actor/task, and will be called each time
+                prior to initializing the worker. Args returned from this dict will
+                always override the args in ``ray_remote_args``. Note: this is an
+                advanced, experimental feature.
+            ray_remote_args: Args to provide to :func:`ray.remote`.
+        """
+        name = self._get_operator_name(name, fn)
+        super().__init__(
+            name,
+            input_dependencies,
+            can_modify_num_rows=can_modify_num_rows,
+            min_rows_per_bundled_input=min_rows_per_bundled_input,
+            ray_remote_args=ray_remote_args,
+            compute=compute,
+        )
+        object.__setattr__(self, "fn", fn)
+        object.__setattr__(self, "fn_args", fn_args)
+        object.__setattr__(self, "fn_kwargs", fn_kwargs)
+        object.__setattr__(self, "fn_constructor_args", fn_constructor_args)
+        object.__setattr__(self, "fn_constructor_kwargs", fn_constructor_kwargs)
+        object.__setattr__(self, "ray_remote_args_fn", ray_remote_args_fn)
+
+    def _get_operator_name(self, op_name: str, fn: UserDefinedFunction):
+        """Gets the Operator name including the map `fn` UDF name."""
+        # If the input `fn` is a Preprocessor, the
+        # name is simply the name of the Preprocessor class.
+        if inspect.ismethod(fn) and isinstance(fn.__self__, Preprocessor):
+            return fn.__self__.__class__.__name__
+
+        # Otherwise, it takes the form of `<MapOperator class>(<UDF name>)`,
+        # e.g. `MapBatches(my_udf)`.
+        try:
+            if inspect.isclass(fn):
+                # callable class
+                return f"{op_name}({fn.__name__})"
+            elif inspect.ismethod(fn):
+                # class method
+                return f"{op_name}({fn.__self__.__class__.__name__}.{fn.__name__})"
+            elif inspect.isfunction(fn):
+                # normal function or lambda function.
+                return f"{op_name}({fn.__name__})"
+            elif isinstance(fn, functools.partial):
+                # functools.partial
+                return f"{op_name}({fn.func.__name__})"
+            else:
+                # callable object.
+                return f"{op_name}({fn.__class__.__name__})"
+        except AttributeError as e:
+            logger.error("Failed to get name of UDF %s: %s", fn, e)
+            return "<unknown>"
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class MapBatches(AbstractUDFMap):
+    """Logical operator for map_batches."""
+
+    fn: UserDefinedFunction
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+    can_modify_num_rows: bool = False
+    batch_size: Union[Optional[int], Literal["auto"]] = None
+    batch_format: Optional[str] = "default"
+    zero_copy_batch: bool = True
+    fn_args: Optional[Iterable[Any]] = None
+    fn_kwargs: Optional[Dict[str, Any]] = None
+    fn_constructor_args: Optional[Iterable[Any]] = None
+    fn_constructor_kwargs: Optional[Dict[str, Any]] = None
+    min_rows_per_bundled_input: Optional[int] = None
+    compute: Optional[ComputeStrategy] = None
+    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    per_block_limit: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if self.compute is None:
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        object.__setattr__(
+            self,
+            "_name",
+            self._get_operator_name(self.__class__.__name__, self.fn),
+        )
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class MapRows(AbstractUDFMap):
+    """Logical operator for map."""
+
+    fn: UserDefinedFunction
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+    fn_args: Optional[Iterable[Any]] = None
+    fn_kwargs: Optional[Dict[str, Any]] = None
+    fn_constructor_args: Optional[Iterable[Any]] = None
+    fn_constructor_kwargs: Optional[Dict[str, Any]] = None
+    compute: Optional[ComputeStrategy] = None
+    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    can_modify_num_rows: bool = field(init=False, default=False)
+    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
+    per_block_limit: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if self.compute is None:
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        object.__setattr__(self, "_name", self._get_operator_name("Map", self.fn))
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Filter(AbstractUDFMap, LogicalOperatorPreservesSchema):
+    """Logical operator for filter."""
+
+    predicate_expr: Optional[Expr] = None
+    fn: Optional[UserDefinedFunction] = None
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+    fn_args: Optional[Iterable[Any]] = None
+    fn_kwargs: Optional[Dict[str, Any]] = None
+    fn_constructor_args: Optional[Iterable[Any]] = None
+    fn_constructor_kwargs: Optional[Dict[str, Any]] = None
+    compute: Optional[ComputeStrategy] = None
+    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    can_modify_num_rows: bool = field(init=False, default=True)
+    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
+    per_block_limit: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        provided_params = sum([self.fn is not None, self.predicate_expr is not None])
+        if provided_params != 1:
+            raise ValueError(
+                "Exactly one of 'fn', or 'predicate_expr' must be provided "
+                f"(received fn={self.fn}, predicate_expr={self.predicate_expr})"
+            )
+        if self.compute is None:
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        object.__setattr__(
+            self,
+            "_name",
+            self._get_operator_name(self.__class__.__name__, self.fn),
+        )
+
+    def is_expression_based(self) -> bool:
+        return self.predicate_expr is not None
+
+    def _get_operator_name(self, op_name: str, fn: UserDefinedFunction):
+        if self.is_expression_based():
+            # Get a concise inline string representation of the expression
+            from ray.data._internal.planner.plan_expression.expression_visitors import (
+                _InlineExprReprVisitor,
+            )
+
+            expr_str = _InlineExprReprVisitor().visit(self.predicate_expr)
+
+            # Truncate only the final result if too long
+            max_length = 60
+            if len(expr_str) > max_length:
+                expr_str = expr_str[: max_length - 3] + "..."
+
+            return f"{op_name}({expr_str})"
+        return super()._get_operator_name(op_name, fn)
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Project(AbstractMap, LogicalOperatorSupportsPredicatePassThrough):
+    """Logical operator for all Projection Operations."""
+
+    exprs: list["Expr"]
+    _common_sub_exprs: list["Expr"] = field(
+        default_factory=list,
+        repr=False,
+        kw_only=True,
+    )
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+    compute: Optional[ComputeStrategy] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
+    can_modify_num_rows: bool = field(init=False, default=False)
+    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
+    batch_size: Optional[int] = field(init=False, default=None)
+    batch_format: str = field(init=False, default="pyarrow")
+    zero_copy_batch: bool = field(init=False, default=True)
+    per_block_limit: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        # Eagerly expand ``StarExpr`` when the input schema is known. By the time
+        # optimizer rules see this op, the projection list contains only
+        # explicit ``col()`` and computed expressions, no ``StarExpr``.
+        # When the input schema is opaque (e.g., upstream UDF map), the
+        # ``StarExpr`` is preserved and runtime ``eval_projection``
+        # expands it on a per-block basis.
+        import pyarrow as pa
+
+        input_schema = self.input_dependencies[0].infer_schema()
+        if isinstance(input_schema, pa.Schema):
+            object.__setattr__(
+                self, "exprs", expand_star_exprs(self.exprs, input_schema)
+            )
+        if self.compute is None:
+            object.__setattr__(
+                self,
+                "compute",
+                self._detect_and_get_compute_strategy(self.get_all_exprs()),
+            )
+        for expr in self.exprs:
+            if expr.name is None and not isinstance(expr, StarExpr):
+                raise TypeError(
+                    "All Project expressions must be named (use .alias(name) or col(name)), "
+                    "or be a star() expression."
+                )
+
+    def _detect_and_get_compute_strategy(self, exprs: list["Expr"]) -> ComputeStrategy:
+        """Detect if expressions contain callable class UDFs and return appropriate compute strategy.
+
+        If any expression contains a callable class UDF, returns ActorPoolStrategy.
+        Otherwise returns TaskPoolStrategy.
+        """
+        from ray.data._internal.planner.plan_expression.expression_visitors import (
+            _CallableClassUDFCollector,
+        )
+
+        # Check all expressions for callable class UDFs
+        for expr in exprs:
+            collector = _CallableClassUDFCollector()
+            collector.visit(expr)
+            if collector.get_callable_class_udfs():
+                # Found at least one callable class UDF - use actor semantics
+                from ray.data._internal.compute import ActorPoolStrategy
+
+                return ActorPoolStrategy(min_size=1, max_size=None)
+
+        # No callable class UDFs found - use task-based execution
+        from ray.data._internal.compute import TaskPoolStrategy
+
+        return TaskPoolStrategy()
+
+    def has_star_expr(self) -> bool:
+        return self.get_star_expr() is not None
+
+    def is_idempotent(self) -> bool:
+        """Return whether every output expression of this projection is idempotent."""
+        return all(expr.is_idempotent() for expr in self.exprs)
+
+    def get_star_expr(self) -> Optional[StarExpr]:
+        """Check if this projection contains a star() expression."""
+        for expr in self.exprs:
+            if isinstance(expr, StarExpr):
+                return expr
+
+        return None
+
+    def get_common_sub_exprs(self) -> list["Expr"]:
+        return self._common_sub_exprs
+
+    def get_all_exprs(self) -> list["Expr"]:
+        """Both projection expressions and common expressions"""
+        return [*self._common_sub_exprs, *self.exprs]
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        return PredicatePassThroughBehavior.PASSTHROUGH_WITH_SUBSTITUTION
+
+    def get_column_substitutions(self) -> Optional[Dict[str, str]]:
+        """Returns the column renames from this projection.
+
+        Maps source_column_name -> output_column_name. This is what we need
+        to rebind predicates when pushing through.
+        """
+        # Reuse the existing logic from projection pushdown
+        from ray.data._internal.logical.rules.projection_pushdown import (
+            _extract_input_columns_renaming_mapping,
+        )
+
+        rename_map = _extract_input_columns_renaming_mapping(self.exprs)
+        return rename_map if rename_map else None
+
+    def infer_schema(self) -> Optional["Schema"]:
+        import pyarrow as pa
+
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        input_schema = self.input_dependencies[0].infer_schema()
+        # Only Arrow schemas are supported for static expression resolution.
+        # (``PandasBlockSchema`` chains fall back to ``limit(1)`` execution.)
+        if not isinstance(input_schema, pa.Schema):
+            return None
+        working_schema = input_schema
+        for common_expr in self.get_common_sub_exprs():
+            field = common_expr.to_field(working_schema)
+            if field is None:
+                return None
+            working_schema = working_schema.append(field)
+        fields = exprlist_to_fields(self.exprs, working_schema)
+        if fields is None:
+            return None
+        if self.get_common_sub_exprs():
+            temp_names = {expr.name for expr in self.get_common_sub_exprs()}
+            fields = [field for field in fields if field.name not in temp_names]
+        return pa.schema(fields)
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class FlatMap(AbstractUDFMap):
+    """Logical operator for flat_map."""
+
+    fn: UserDefinedFunction
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+    fn_args: Optional[Iterable[Any]] = None
+    fn_kwargs: Optional[Dict[str, Any]] = None
+    fn_constructor_args: Optional[Iterable[Any]] = None
+    fn_constructor_kwargs: Optional[Dict[str, Any]] = None
+    compute: Optional[ComputeStrategy] = None
+    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    can_modify_num_rows: bool = field(init=False, default=True)
+    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
+    per_block_limit: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if self.compute is None:
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        object.__setattr__(
+            self,
+            "_name",
+            self._get_operator_name(self.__class__.__name__, self.fn),
+        )
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class StreamingRepartition(
+    AbstractMap,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorPreservesSchema,
+):
+    """Logical operator for streaming repartition operation.
+
+    Args:
+        input_dependencies: The operators preceding this operator in the plan DAG.
+        target_num_rows_per_block: The target number of rows per block granularity for
+            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.
+    """
+
+    target_num_rows_per_block: int
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+    strict: bool = False
+    can_modify_num_rows: bool = field(init=False, default=False)
+    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
+    compute: Optional[ComputeStrategy] = None
+    per_block_limit: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if self.target_num_rows_per_block <= 0:
+            raise ValueError(
+                "target_num_rows_per_block must be positive for streaming repartition, "
+                f"got {self.target_num_rows_per_block}"
+            )
+        if self.compute is None:
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        object.__setattr__(
+            self,
+            "_name",
+            f"StreamingRepartition[num_rows_per_block={self.target_num_rows_per_block},strict={self.strict}]",
+        )
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        # StreamingRepartition only re-bundles rows into different block sizes.
+        # It doesn't modify schema or filter rows, so filters can safely pass through.
+        return PredicatePassThroughBehavior.PASSTHROUGH

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

@@ -0,0 +1,211 @@
+import enum
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, Optional
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorUnifiesInputSchemas,
+    PredicatePassThroughBehavior,
+)
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.block import Schema
+
+__all__ = [
+    "Mix",
+    "MixStoppingCondition",
+    "NAry",
+    "Union",
+    "Zip",
+]
+
+
+@PublicAPI(stability="alpha")
+class MixStoppingCondition(enum.Enum):
+    """Controls when a mix pipeline terminates.
+
+    STOP_ON_SHORTEST: Pipeline ends when the shortest dataset is exhausted.
+        Other datasets are truncated.
+    STOP_ON_LONGEST_DROP: Pipeline ends when the longest dataset is exhausted.
+        Shorter datasets drop out once exhausted; later batches are drawn
+        entirely from longer datasets.
+    """
+
+    STOP_ON_SHORTEST = "stop_on_shortest"
+    STOP_ON_LONGEST_DROP = "stop_on_longest_drop"
+
+
+def estimate_num_mix_outputs(
+    per_input_counts: List[Optional[int]],
+    weights: List[float],
+    stopping_condition: MixStoppingCondition,
+) -> Optional[int]:
+    """Estimate total output count for a mix operation.
+
+    Used by both the logical and physical Mix operators to estimate
+    num_outputs_total / num_output_rows_total.
+    """
+    if any(c is None for c in per_input_counts):
+        return None
+    if stopping_condition == MixStoppingCondition.STOP_ON_LONGEST_DROP:
+        return sum(per_input_counts)
+    elif stopping_condition == MixStoppingCondition.STOP_ON_SHORTEST:
+        # Limited by whichever input runs out first relative to its weight.
+        total_weight = sum(weights)
+        return min(
+            int(count / (w / total_weight))
+            for count, w in zip(per_input_counts, weights)
+        )
+    else:
+        raise ValueError(f"Unknown stopping condition: {stopping_condition}")
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class NAry(LogicalOperator):
+    """Base class for n-ary operators, which take multiple input operators."""
+
+    def __init__(
+        self,
+        input_dependencies: List[LogicalOperator],
+    ):
+        """Initialize the n-ary operator.
+
+        Args:
+            input_dependencies: The input operators.
+        """
+        object.__setattr__(self, "_input_dependencies", list(input_dependencies))
+
+    def _with_new_input_dependencies(
+        self, input_dependencies: List[LogicalOperator]
+    ) -> LogicalOperator:
+        return self.__class__(input_dependencies)
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class Zip(NAry):
+    """Logical operator for zip."""
+
+    _input_dependencies: List[LogicalOperator] = field(init=False, repr=False)
+
+    def __init__(
+        self,
+        input_dependencies: List[LogicalOperator],
+    ):
+        for input_op in input_dependencies:
+            assert isinstance(input_op, LogicalOperator), input_op
+        object.__setattr__(self, "_input_dependencies", list(input_dependencies))
+
+    def estimated_num_outputs(self):
+        total_num_outputs = 0
+        for input in self.input_dependencies:
+            num_outputs = input.estimated_num_outputs()
+            if num_outputs is None:
+                return None
+            total_num_outputs = max(total_num_outputs, num_outputs)
+        return total_num_outputs
+
+    def infer_schema(self) -> Optional["Schema"]:
+        # Reuse the runtime ``BlockAccessor.zip`` so plan-time and
+        # execution-time schemas agree by construction (same column
+        # suffixing rules, etc.).
+        import pyarrow as pa
+
+        from ray.data.block import BlockAccessor
+
+        input_schemas = [op.infer_schema() for op in self.input_dependencies]
+        if not input_schemas or not all(
+            isinstance(s, pa.Schema) for s in input_schemas
+        ):
+            return None
+        try:
+            combined = input_schemas[0].empty_table()
+            for s in input_schemas[1:]:
+                combined = BlockAccessor.for_block(combined).zip(s.empty_table())
+        except (pa.ArrowTypeError, pa.ArrowInvalid):
+            return None
+        return combined.schema
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class Mix(NAry, LogicalOperatorUnifiesInputSchemas):
+    """Logical operator for weighted dataset mixing."""
+
+    _name: str = field(init=False, repr=False)
+    _input_dependencies: List[LogicalOperator] = field(init=False, repr=False)
+    weights: List[float] = field(init=False, repr=False)
+    stopping_condition: MixStoppingCondition = field(init=False, repr=False)
+
+    def __init__(
+        self,
+        input_dependencies: List[LogicalOperator],
+        *,
+        weights: List[float],
+        stopping_condition: MixStoppingCondition = MixStoppingCondition.STOP_ON_SHORTEST,
+    ):
+        if len(input_dependencies) != len(weights):
+            raise ValueError(
+                f"Number of input operators ({len(input_dependencies)}) must match "
+                f"number of weights ({len(weights)})."
+            )
+        if any(weight <= 0 for weight in weights):
+            raise ValueError(f"Weights must be positive. Got weights: {weights}")
+
+        for input_op in input_dependencies:
+            assert isinstance(input_op, LogicalOperator), input_op
+        object.__setattr__(self, "_name", self.__class__.__name__)
+        object.__setattr__(self, "_input_dependencies", list(input_dependencies))
+        object.__setattr__(self, "weights", weights)
+        object.__setattr__(self, "stopping_condition", stopping_condition)
+
+    def estimated_num_outputs(self) -> Optional[int]:
+        if self.stopping_condition == MixStoppingCondition.STOP_ON_SHORTEST:
+            return None
+
+        return estimate_num_mix_outputs(
+            [op.estimated_num_outputs() for op in self.input_dependencies],
+            self.weights,
+            self.stopping_condition,
+        )
+
+    def _with_new_input_dependencies(
+        self, input_dependencies: List[LogicalOperator]
+    ) -> LogicalOperator:
+        return self.__class__(
+            input_dependencies,
+            weights=self.weights,
+            stopping_condition=self.stopping_condition,
+        )
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class Union(
+    NAry,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorUnifiesInputSchemas,
+):
+    """Logical operator for union."""
+
+    _input_dependencies: List[LogicalOperator] = field(init=False, repr=False)
+
+    def __init__(
+        self,
+        input_dependencies: List[LogicalOperator],
+    ):
+        for input_op in input_dependencies:
+            assert isinstance(input_op, LogicalOperator), input_op
+        object.__setattr__(self, "_input_dependencies", list(input_dependencies))
+
+    def estimated_num_outputs(self):
+        total_num_outputs = 0
+        for input in self.input_dependencies:
+            num_outputs = input.estimated_num_outputs()
+            if num_outputs is None:
+                return None
+            total_num_outputs += num_outputs
+        return total_num_outputs
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        # Union allows pushing filter into each branch
+        return PredicatePassThroughBehavior.PUSH_INTO_BRANCHES

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

@@ -0,0 +1,185 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorPreservesSchema,
+    LogicalOperatorSupportsPredicatePassThrough,
+    PredicatePassThroughBehavior,
+)
+from ray.data.block import BlockMetadata
+
+if TYPE_CHECKING:
+    import pyarrow
+
+    from ray.data.block import Schema
+
+__all__ = [
+    "AbstractOneToOne",
+    "Download",
+    "Limit",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False, init=False)
+class AbstractOneToOne(LogicalOperator):
+    """Abstract class for one-to-one logical operators, which
+    have one input and one output dependency.
+    """
+
+    def __init__(
+        self,
+        input_dependencies: Optional[List[LogicalOperator]],
+        can_modify_num_rows: bool,
+        *,
+        name: Optional[str] = None,
+    ):
+        """Initialize an AbstractOneToOne operator.
+
+        Args:
+            input_dependencies: The operators preceding this operator in the plan DAG.
+                The outputs of these operators will be the inputs to this operator.
+            can_modify_num_rows: Whether the UDF can change the row count. False if
+                # of input rows = # of output rows. True otherwise.
+            name: Name for this operator. This is the name that will appear when
+                inspecting the logical plan of a Dataset.
+        """
+        object.__setattr__(self, "_input_dependencies", list(input_dependencies or []))
+        if name is not None:
+            object.__setattr__(self, "_name", name)
+        object.__setattr__(self, "can_modify_num_rows", can_modify_num_rows)
+
+    def infer_metadata(self) -> BlockMetadata:
+        """Best-effort output metadata derived from the single input dependency.
+
+        One-to-one operators that don't modify the row count (e.g. ``Project``)
+        preserve the row count and don't grow the data, so the input's row count
+        and byte size are valid output estimates -- an upper bound for the common
+        column-selection case, where projecting away columns only shrinks the
+        data. Operators that can modify the row count (e.g. ``Filter``,
+        ``FlatMap``) can't reuse these estimates, so they fall back to ``None``.
+
+        Propagating ``size_bytes`` is what keeps size-dependent planning correct
+        when a one-to-one op is pushed below a join/shuffle (e.g. by projection
+        pushdown): otherwise ``_try_estimate_output_bytes`` sees ``size_bytes=None``
+        and the hash-shuffle aggregator silently falls back to a fixed default
+        memory reservation instead of one derived from the dataset size.
+        """
+        if len(self.input_dependencies) != 1:
+            return BlockMetadata(
+                num_rows=None, size_bytes=None, input_files=None, exec_stats=None
+            )
+        input_meta = self.input_dependencies[0].infer_metadata()
+        if self.can_modify_num_rows:
+            return BlockMetadata(
+                num_rows=None,
+                size_bytes=None,
+                input_files=input_meta.input_files,
+                exec_stats=None,
+            )
+        return BlockMetadata(
+            num_rows=input_meta.num_rows,
+            size_bytes=input_meta.size_bytes,
+            input_files=input_meta.input_files,
+            exec_stats=None,
+        )
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Limit(
+    AbstractOneToOne,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorPreservesSchema,
+):
+    """Logical operator for limit."""
+
+    limit: int
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+    can_modify_num_rows: bool = field(init=False, default=True)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        object.__setattr__(self, "_name", f"limit={self.limit}")
+
+    def infer_metadata(self) -> BlockMetadata:
+        return BlockMetadata(
+            num_rows=self._num_rows(),
+            size_bytes=None,
+            input_files=self._input_files(),
+            exec_stats=None,
+        )
+
+    def _num_rows(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        assert isinstance(self.input_dependencies[0], LogicalOperator)
+        input_rows = self.input_dependencies[0].infer_metadata().num_rows
+        if input_rows is not None:
+            return min(input_rows, self.limit)
+        else:
+            return None
+
+    def _input_files(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        assert isinstance(self.input_dependencies[0], LogicalOperator)
+        return self.input_dependencies[0].infer_metadata().input_files
+
+    def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
+        # Pushing filter through limit is safe: Filter(Limit(data, n), pred)
+        # becomes Limit(Filter(data, pred), n), which filters earlier
+        return PredicatePassThroughBehavior.PASSTHROUGH
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Download(AbstractOneToOne):
+    """Logical operator for download operation.
+
+    Supports downloading from multiple URI columns in a single operation.
+    """
+
+    uri_column_names: List[str]
+    output_bytes_column_names: List[str]
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    filesystem: Optional["pyarrow.fs.FileSystem"] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+    can_modify_num_rows: bool = field(init=False, default=False)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if len(self.uri_column_names) != len(self.output_bytes_column_names):
+            raise ValueError(
+                f"Number of URI columns ({len(self.uri_column_names)}) must match "
+                f"number of output columns ({len(self.output_bytes_column_names)})"
+            )
+
+    def infer_metadata(self) -> BlockMetadata:
+        # Download preserves the row count but appends a binary blob column per
+        # requested output, so the output is *larger* than the input -- often by
+        # orders of magnitude. Propagating the input's ``size_bytes`` (as the
+        # row-preserving default does) would be a misleading under-estimate that
+        # could starve downstream size-dependent planning (e.g. hash-shuffle
+        # aggregator memory reservation). Keep the accurate row count and input
+        # files, but report ``size_bytes=None`` since we can't estimate the
+        # downloaded bytes ahead of execution.
+        meta = super().infer_metadata()
+        return BlockMetadata(
+            num_rows=meta.num_rows,
+            size_bytes=None,
+            input_files=meta.input_files,
+            exec_stats=None,
+        )
+
+    def infer_schema(self) -> Optional["Schema"]:
+        # Output = input schema with one binary column appended per requested
+        # output name. The runtime (``download_bytes_threaded``) always appends
+        # via ``add_column`` without removing any pre-existing column of the
+        # same name, so name collisions produce duplicate columns here too.
+        import pyarrow as pa
+
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        input_schema = self.input_dependencies[0].infer_schema()
+        if not isinstance(input_schema, pa.Schema):
+            return None
+        fields = list(input_schema)
+        for name in self.output_bytes_column_names:
+            fields.append(pa.field(name, pa.binary(), nullable=True))
+        return pa.schema(fields)

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

@@ -0,0 +1,476 @@
+import functools
+import math
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Union
+
+from ray.data._internal.compute import ComputeStrategy
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorSupportsPredicatePushdown,
+    LogicalOperatorSupportsProjectionPushdown,
+    SourceOperator,
+)
+from ray.data._internal.logical.operators.map_operator import AbstractMap
+from ray.data.block import (
+    Block,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+)
+from ray.data.context import DataContext
+from ray.data.datasource.datasource import Datasource, Reader
+from ray.data.expressions import Expr
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+    from pyarrow.fs import FileSystem
+
+    from ray.data._internal.datasource_v2.listing.file_indexer import FileIndexer
+    from ray.data._internal.datasource_v2.partitioners.file_partitioner import (
+        FilePartitioner,
+    )
+    from ray.data._internal.datasource_v2.scanners.scanner import Scanner
+    from ray.data.datasource.file_based_datasource import FileShuffleConfig
+    from ray.data.datasource.partitioning import PathPartitionFilter
+
+__all__ = [
+    "ListFiles",
+    "Read",
+    "ReadFiles",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Read(
+    AbstractMap,
+    SourceOperator,
+    LogicalOperatorSupportsProjectionPushdown,
+    LogicalOperatorSupportsPredicatePushdown,
+):
+    """Logical operator for read."""
+
+    datasource: Datasource
+    datasource_or_legacy_reader: Union[Datasource, Reader]
+    parallelism: int
+    num_outputs: Optional[int] = None
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    compute: Optional[ComputeStrategy] = None
+    detected_parallelism: Optional[int] = None
+    can_modify_num_rows: bool = field(init=False, default=True)
+    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
+    ray_remote_args_fn: None = field(init=False, default=None)
+    per_block_limit: Optional[int] = None
+    _input_dependencies: list = field(init=False, repr=False, default_factory=list)
+
+    def __post_init__(self):
+        if self.compute is None:
+            from ray.data._internal.compute import TaskPoolStrategy
+
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        if self.ray_remote_args is None:
+            object.__setattr__(self, "ray_remote_args", {})
+        object.__setattr__(self, "_name", f"Read{self.datasource.get_name()}")
+        object.__setattr__(self, "_input_dependencies", [])
+
+    def output_data(self):
+        return None
+
+    def set_detected_parallelism(self, parallelism: int) -> "Read":
+        """
+        Set the true parallelism that should be used during execution. This
+        should be specified by the user or detected by the optimizer.
+        """
+        object.__setattr__(self, "detected_parallelism", parallelism)
+        return self
+
+    def get_detected_parallelism(self) -> int:
+        """
+        Get the true parallelism that should be used during execution.
+        """
+        return self.detected_parallelism
+
+    def estimated_num_outputs(self) -> Optional[int]:
+        return self.num_outputs or self._estimate_num_outputs()
+
+    def infer_metadata(self) -> BlockMetadata:
+        """A ``BlockMetadata`` that represents the aggregate metadata of the outputs.
+
+        This method gets metadata from the read tasks. It doesn't trigger any actual
+        execution.
+        """
+        return self._cached_output_metadata.metadata
+
+    def infer_schema(self):
+        return self._cached_output_metadata.schema
+
+    def _estimate_num_outputs(self) -> Optional[int]:
+        metadata = self._cached_output_metadata.metadata
+
+        # Handle edge-case of empty dataset
+        if metadata.size_bytes == 0:
+            return 0
+
+        target_max_block_size = DataContext.get_current().target_max_block_size
+
+        # In either case of
+        #   - Total byte-size estimate not available
+        #   - Target max-block-size not being configured
+        #
+        # We fallback to estimating number of outputs to be equivalent to the
+        # number of input files being read (if any)
+        if metadata.size_bytes is None or target_max_block_size is None:
+            # NOTE: If there's no input files specified, return the count (could be 0)
+            return (
+                len(metadata.input_files) if metadata.input_files is not None else None
+            )
+
+        # Otherwise, estimate total number of blocks from estimated total
+        # byte size
+        return math.ceil(metadata.size_bytes / target_max_block_size)
+
+    @functools.cached_property
+    def _cached_output_metadata(self) -> "BlockMetadataWithSchema":
+        # Legacy datasources might not implement `get_read_tasks`.
+        if self.datasource.should_create_reader:
+            empty_meta = BlockMetadata(None, None, None, None)
+            return BlockMetadataWithSchema.from_metadata(empty_meta, schema=None)
+
+        # HACK: Try to get a single read task to get the metadata.
+        read_tasks = self.datasource.get_read_tasks(1)
+        if len(read_tasks) == 0:
+            # If there are no read tasks, the dataset is probably empty.
+            empty_meta = BlockMetadata(
+                num_rows=0,
+                size_bytes=0,
+                input_files=None,
+                exec_stats=None,
+            )
+            return BlockMetadataWithSchema.from_metadata(empty_meta, schema=None)
+
+        # `get_read_tasks` isn't guaranteed to return exactly one read task.
+        metadata = [read_task.metadata for read_task in read_tasks]
+
+        if all(meta.num_rows is not None for meta in metadata):
+            num_rows = sum(meta.num_rows for meta in metadata)
+            original_num_rows = num_rows
+            # Apply per-block limit if set
+            if self.per_block_limit is not None:
+                num_rows = min(num_rows, self.per_block_limit)
+        else:
+            num_rows = None
+            original_num_rows = None
+
+        if all(meta.size_bytes is not None for meta in metadata):
+            size_bytes = sum(meta.size_bytes for meta in metadata)
+            # Pro-rate the byte size if we applied a row limit
+            if (
+                self.per_block_limit is not None
+                and original_num_rows is not None
+                and original_num_rows > 0
+            ):
+                size_bytes = int(size_bytes * (num_rows / original_num_rows))
+        else:
+            size_bytes = None
+
+        input_files = []
+        for meta in metadata:
+            if meta.input_files is not None:
+                input_files.extend(meta.input_files)
+
+        meta = BlockMetadata(
+            num_rows=num_rows,
+            size_bytes=size_bytes,
+            input_files=input_files,
+            exec_stats=None,
+        )
+        schemas = [
+            read_task.schema for read_task in read_tasks if read_task.schema is not None
+        ]
+        from ray.data._internal.util import unify_schemas_with_validation
+
+        schema = None
+        if schemas:
+            schema = unify_schemas_with_validation(schemas)
+        return BlockMetadataWithSchema.from_metadata(meta, schema=schema)
+
+    def supports_projection_pushdown(self) -> bool:
+        return self.datasource.supports_projection_pushdown()
+
+    def get_projection_map(self) -> Optional[Dict[str, str]]:
+        return self.datasource.get_projection_map()
+
+    def apply_projection(
+        self,
+        projection_map: Optional[Dict[str, str]],
+    ) -> "Read":
+        projected_datasource = self.datasource.apply_projection(projection_map)
+        return replace(
+            self,
+            datasource=projected_datasource,
+            datasource_or_legacy_reader=projected_datasource,
+            num_outputs=self.num_outputs,
+        )
+
+    def supports_predicate_pushdown(self) -> bool:
+        return self.datasource.supports_predicate_pushdown()
+
+    def get_current_predicate(self) -> Optional[Expr]:
+        return self.datasource.get_current_predicate()
+
+    def apply_predicate(self, predicate_expr: Expr) -> "Read":
+        predicated_datasource = self.datasource.apply_predicate(predicate_expr)
+
+        # A datasource returns its own instance to signal "no pushdown applied"
+        # (e.g. ``ParquetDatasource`` does this when a mixed-column conjunct
+        # leaves a residual). Preserve identity here so ``PredicatePushdown``'s
+        # ``result_op is input_op`` no-op check keeps the ``Filter`` above.
+        if self.datasource is predicated_datasource:
+            return self
+
+        return replace(
+            self,
+            datasource=predicated_datasource,
+            datasource_or_legacy_reader=predicated_datasource,
+            num_outputs=self.num_outputs,
+        )
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class ReadFiles(
+    AbstractMap,
+    LogicalOperatorSupportsProjectionPushdown,
+    LogicalOperatorSupportsPredicatePushdown,
+):
+    """Logical operator for DataSourceV2 reads.
+
+    Consumes ``FileManifest`` blocks produced by a :class:`ListFiles`
+    source operator upstream. Owns the :class:`Scanner` (with any pushed
+    column/predicate/limit state) and the post-pushdown schema. Listing,
+    shuffling, and size-balanced bucketing happen in the upstream op;
+    this op's physical planner just reads each manifest bucket via
+    ``scanner.create_reader().read(manifest)``.
+
+    V2 reads never rename columns at the read stage — column renaming
+    is always handled by a ``Project`` operator above ``ReadFiles``.
+    This simplifies projection and predicate pushdown by eliminating
+    the "predicate above uses new names, predicate below uses old
+    names" rebinding dance.
+    """
+
+    datasource_name: str
+    scanner: "Scanner"
+    schema: "pa.Schema"
+    parallelism: int
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    compute: Optional[ComputeStrategy] = None
+    # Optional post-read block transform. Used by ``read_parquet``'s
+    # ``_block_udf`` and ``tensor_column_schema`` (the latter is folded
+    # into a ``_block_udf`` by ``_resolve_parquet_args`` before it gets
+    # here). Applied in ``plan_read_files_op.do_read`` after each
+    # table is read.
+    block_udf: Optional[Callable[[Block], Block]] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+    can_modify_num_rows: bool = field(init=False, default=True)
+    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
+    ray_remote_args_fn: None = field(init=False, default=None)
+    # Declared so the inherited ``AbstractMap._get_args`` can resolve it; V2
+    # limit pushdown is applied via ``scanner.push_limit`` (see
+    # ``LimitPushdownRule._apply_per_block_limit_if_supported``), not this field.
+    per_block_limit: Optional[int] = field(init=False, default=None)
+    num_outputs: Optional[int] = None
+    _name: str = field(init=False, repr=False)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        assert isinstance(
+            self.input_dependencies[0], LogicalOperator
+        ), self.input_dependencies[0]
+        if self.compute is None:
+            from ray.data._internal.compute import TaskPoolStrategy
+
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        if self.ray_remote_args is None:
+            object.__setattr__(self, "ray_remote_args", {})
+        object.__setattr__(self, "_name", f"ReadFiles{self.datasource_name}")
+
+    def infer_schema(self) -> "pa.Schema":
+        # Scanner schema reflects any applied projection pushdown
+        # (``scanner.prune_columns`` / empty projection from
+        # ``select_columns([])``); the stored ``self.schema`` is the
+        # unprojected one and only used for construction.
+        schema = self.scanner.read_schema()
+        # When a ``block_udf`` is attached (e.g. ``read_parquet`` was
+        # called with ``tensor_column_schema`` or ``_block_udf``), probe
+        # its effect on the schema so downstream consumers see the
+        # post-transform column types. Mirrors V1 ``ParquetDatasource``'s
+        # dummy-table trick. Falls back to the scanner schema if the
+        # probe fails — the UDF may require a non-empty input.
+        if self.block_udf is not None:
+            try:
+                transformed = self.block_udf(schema.empty_table()).schema
+                schema = transformed.with_metadata(schema.metadata)
+            except Exception:
+                pass
+        return schema
+
+    def infer_metadata(self) -> BlockMetadata:
+        """Return empty metadata; downstream callers fall back to materialization.
+
+        Prior ``ReadFiles`` versions reached into a driver-side file cache to
+        compute size hints. With listing owned by an upstream
+        ``ListFiles`` op, metadata-for-sizing is computed from the
+        materialized manifest at execution time — the logical op doesn't
+        try to pre-estimate.
+        """
+        return BlockMetadata(None, None, None, None)
+
+    def supports_projection_pushdown(self) -> bool:
+        from ray.data._internal.datasource_v2.logical_optimizers import (
+            SupportsColumnPruning,
+        )
+
+        return isinstance(self.scanner, SupportsColumnPruning)
+
+    def get_projection_map(self) -> Optional[Dict[str, str]]:
+        if not self.supports_projection_pushdown():
+            return None
+        columns = self.scanner.pruned_column_names()
+        if columns is None:
+            return None
+        # The read stage never renames at the read layer; the projection
+        # map is always an identity (original name -> original name).
+        # Renaming is always carried by an ``AliasExpr`` in a ``Project``
+        # operator above the read.
+        return {name: name for name in columns}
+
+    def apply_projection(
+        self,
+        projection_map: Optional[Dict[str, str]],
+    ) -> "ReadFiles":
+        if projection_map is None:
+            return self
+        from ray.data._internal.datasource_v2.logical_optimizers import (
+            SupportsColumnPruning,
+        )
+
+        assert isinstance(self.scanner, SupportsColumnPruning)
+
+        # V2 reads only prune columns at the read stage. Any rename info
+        # in ``projection_map`` is dropped here; the optimizer rule keeps
+        # a ``Project`` op on top of ``ReadFiles`` to carry rename
+        # ``AliasExpr`` instances. Only the keys (column names to keep)
+        # are used.
+        new_scanner = self.scanner.prune_columns(list(projection_map.keys()))
+        return replace(self, scanner=new_scanner)
+
+    def supports_predicate_pushdown(self) -> bool:
+        from ray.data._internal.datasource_v2.logical_optimizers import (
+            SupportsFilterPushdown,
+        )
+
+        return isinstance(self.scanner, SupportsFilterPushdown)
+
+    def get_current_predicate(self) -> Optional[Expr]:
+        return getattr(self.scanner, "predicate", None)
+
+    def apply_predicate(self, predicate_expr: Expr) -> LogicalOperator:
+        from ray.data._internal.datasource.parquet_datasource import (
+            _split_predicate_by_columns,
+        )
+        from ray.data._internal.datasource_v2.logical_optimizers import (
+            SupportsFilterPushdown,
+            SupportsPartitionPruning,
+        )
+        from ray.data._internal.logical.operators.map_operator import Filter
+
+        assert isinstance(self.scanner, SupportsFilterPushdown)
+
+        partition_cols: Set[str] = (
+            self.scanner.partition_columns
+            if isinstance(self.scanner, SupportsPartitionPruning)
+            else set()
+        )
+
+        if not partition_cols:
+            new_scanner, _residual = self.scanner.push_filters(predicate_expr)
+            return replace(self, scanner=new_scanner)
+
+        split = _split_predicate_by_columns(predicate_expr, partition_cols)
+
+        if split.data_predicate is None and split.partition_predicate is None:
+            # Entire predicate is residual (e.g. a single mixed-column
+            # ``OR``); nothing safe to push. Returning ``self`` tells
+            # ``PredicatePushdown`` to keep the ``Filter`` above us.
+            return self
+
+        new_scanner = self.scanner
+        if split.partition_predicate is not None:
+            new_scanner = new_scanner.prune_partitions(split.partition_predicate)
+        if split.data_predicate is not None:
+            new_scanner, _residual = new_scanner.push_filters(split.data_predicate)
+
+        new_op = replace(self, scanner=new_scanner)
+
+        if split.residual_predicate is None:
+            return new_op
+
+        # Residual conjuncts can't be pushed through either ``push_filters``
+        # (pyarrow only binds data columns) or ``prune_partitions`` (path
+        # parser only binds partition columns), so re-emit them as a
+        # ``Filter`` above the new ``ReadFiles``. Without this, we'd keep
+        # the splittable parts and silently drop the residual — letting
+        # rows through that the original predicate would have rejected.
+        return Filter(
+            predicate_expr=split.residual_predicate, input_dependencies=[new_op]
+        )
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class ListFiles(LogicalOperator, SourceOperator):
+    """Logical source op that lists files and yields ``FileManifest`` blocks.
+
+    Extracted from the prior monolithic ``ReadFiles`` so listing, shuffling,
+    and size-balanced bucketing live in one place (see
+    :func:`ray.data._internal.planner.plan_list_files_op.plan_list_files_op`).
+    Downstream, ``ReadFiles`` consumes the manifest blocks produced here.
+    """
+
+    paths: List[str]
+    file_indexer: "FileIndexer"
+    filesystem: "FileSystem"
+    # Original user-supplied paths. Lineage-tracking pins this to the
+    # caller's intent rather than the resolved absolute paths.
+    source_paths: List[str]
+    file_partitioner: Optional["FilePartitioner"] = None
+    file_extensions: Optional[List[str]] = None
+    partition_filter: Optional["PathPartitionFilter"] = None
+    # A factory (not a stored config) so the shuffle seed is re-sampled
+    # per execution when the config asks for it.
+    shuffle_config_factory: Callable[[], Optional["FileShuffleConfig"]] = field(
+        default=lambda: None
+    )
+    _name: str = field(init=False, repr=False)
+    _input_dependencies: List[LogicalOperator] = field(
+        init=False, repr=False, default_factory=list
+    )
+
+    def __post_init__(self):
+        object.__setattr__(self, "_name", self.__class__.__name__)
+
+    def output_data(self) -> Optional[list]:
+        return None
+
+    def infer_schema(self) -> "pa.Schema":
+        # ``FileManifest`` columns are fixed: __path, __file_size.
+        import pyarrow as pa
+
+        from ray.data._internal.datasource_v2.listing.file_manifest import (
+            FILE_SIZE_COLUMN_NAME,
+            PATH_COLUMN_NAME,
+        )
+
+        return pa.schema(
+            [
+                pa.field(PATH_COLUMN_NAME, pa.string()),
+                pa.field(FILE_SIZE_COLUMN_NAME, pa.int64()),
+            ]
+        )

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

@@ -0,0 +1,27 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, Optional
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorPreservesSchema,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import NodeIdStr
+
+__all__ = [
+    "StreamingSplit",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class StreamingSplit(LogicalOperatorPreservesSchema):
+    """Logical operator that represents splitting the input data to `n` splits."""
+
+    num_splits: int
+    equal: bool
+    locality_hints: Optional[List["NodeIdStr"]] = None
+    input_dependencies: List[LogicalOperator] = field(repr=False, kw_only=True)
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)

python/ray/data/_internal/logical/operators/tests/test_read_files_logical.py

@@ -0,0 +1,216 @@
+"""Unit tests for :class:`ReadFiles`.
+
+Verifies pushdown scaffolding (projection/predicate capability dispatch,
+immutable scanner substitution) and schema inference without triggering
+physical execution. Each test wires a minimal ``ListFiles`` upstream
+op so ``ReadFiles`` (which now has one input dependency) can be
+constructed.
+"""
+
+import os
+from pathlib import Path
+from typing import Optional, Union
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.parquet as pq
+import pytest
+
+from ray.data._internal.datasource_v2.listing.file_indexer import (
+    NonSamplingFileIndexer,
+)
+from ray.data._internal.datasource_v2.listing.listing_utils import (
+    sample_files,
+)
+from ray.data._internal.datasource_v2.parquet_datasource_v2 import (
+    ParquetDatasourceV2,
+)
+from ray.data._internal.datasource_v2.scanners.parquet_scanner import (
+    ParquetScanner,
+)
+from ray.data._internal.logical.operators import Filter, ListFiles, ReadFiles
+from ray.data.datasource.partitioning import Partitioning, PartitionStyle
+from ray.data.expressions import Expr, col
+
+
+def _mk_parquet(path: Path, table: pa.Table) -> None:
+    pq.write_table(table, str(path))
+
+
+def _mk_read_files(tmp_path: Path) -> ReadFiles:
+    f = tmp_path / "data.parquet"
+    _mk_parquet(f, pa.table({"a": [1, 2, 3], "b": ["x", "y", "z"]}))
+
+    datasource = ParquetDatasourceV2([str(f)])
+    indexer = NonSamplingFileIndexer(ignore_missing_paths=False)
+    sample = sample_files(indexer, datasource.paths, datasource.filesystem)
+    schema = datasource.infer_schema(sample)
+    scanner = datasource.create_scanner(schema=schema)
+
+    list_files_op = ListFiles(
+        paths=list(datasource.paths),
+        file_indexer=indexer,
+        filesystem=datasource.filesystem,
+        source_paths=list(datasource.paths),
+        file_extensions=datasource.file_extensions,
+    )
+
+    return ReadFiles(
+        datasource_name=datasource.name,
+        scanner=scanner,
+        schema=schema,
+        parallelism=-1,
+        input_dependencies=[list_files_op],
+    )
+
+
+def _mk_partitioned_read_files(tmp_path: Path) -> ReadFiles:
+    """Hive-partitioned dataset with partition column ``country``."""
+    for country, value in (("US", 1), ("CA", 2)):
+        d = tmp_path / f"country={country}"
+        os.makedirs(d, exist_ok=True)
+        _mk_parquet(d / "data.parquet", pa.table({"a": [value], "b": [str(value)]}))
+
+    partitioning = Partitioning(
+        PartitionStyle.HIVE, base_dir=str(tmp_path), field_names=["country"]
+    )
+    datasource = ParquetDatasourceV2([str(tmp_path)], partitioning=partitioning)
+    indexer = NonSamplingFileIndexer(ignore_missing_paths=False)
+    sample = sample_files(indexer, datasource.paths, datasource.filesystem)
+    schema = datasource.infer_schema(sample)
+    scanner = datasource.create_scanner(schema=schema, partitioning=partitioning)
+
+    list_files_op = ListFiles(
+        paths=list(datasource.paths),
+        file_indexer=indexer,
+        filesystem=datasource.filesystem,
+        source_paths=list(datasource.paths),
+        file_extensions=datasource.file_extensions,
+    )
+
+    return ReadFiles(
+        datasource_name=datasource.name,
+        scanner=scanner,
+        schema=schema,
+        parallelism=-1,
+        input_dependencies=[list_files_op],
+    )
+
+
+def test_construction_stores_schema_and_infer_schema_returns_it(tmp_path):
+    op = _mk_read_files(tmp_path)
+    assert op.infer_schema().names == ["a", "b"]
+
+
+def test_input_dependency_is_list_files(tmp_path):
+    op = _mk_read_files(tmp_path)
+    assert isinstance(op.input_dependencies[0], ListFiles)
+
+
+def test_supports_projection_pushdown_true_for_parquet_scanner(tmp_path):
+    op = _mk_read_files(tmp_path)
+    assert op.supports_projection_pushdown() is True
+
+
+def test_apply_projection_returns_new_op_with_pruned_scanner(tmp_path):
+    op = _mk_read_files(tmp_path)
+    new_op = op.apply_projection({"a": "a"})
+
+    assert new_op is not op
+    assert isinstance(new_op.scanner, ParquetScanner)
+    assert new_op.scanner.columns == ("a",)
+    # Original scanner untouched
+    assert isinstance(op.scanner, ParquetScanner)
+    assert op.scanner.columns is None
+
+
+def test_apply_projection_none_is_noop(tmp_path):
+    op = _mk_read_files(tmp_path)
+    assert op.apply_projection(None) is op
+
+
+def test_supports_predicate_pushdown(tmp_path):
+    assert _mk_read_files(tmp_path).supports_predicate_pushdown() is True
+
+
+def _assert_pred_equals(
+    actual: Optional[Union[Expr, pc.Expression]], expected: Optional[Expr]
+) -> None:
+    if expected is None:
+        assert actual is None
+        return
+    assert actual is not None
+    # ``ArrowFileScanner.push_filters`` stores the data predicate as a
+    # ``pc.Expression`` (via ``.to_pyarrow()``), while ``partition_predicate``
+    # stays a Ray ``Expr``. Dispatch on which we got.
+    if isinstance(actual, Expr):
+        assert actual.structurally_equals(expected)
+    else:
+        assert actual.equals(expected.to_pyarrow())
+
+
+@pytest.mark.parametrize(
+    "partitioned,predicate,expected_data,expected_partition",
+    [
+        (False, col("a") > 1, col("a") > 1, None),
+        (True, col("country") == "US", None, col("country") == "US"),
+        (
+            True,
+            (col("a") > 0) & (col("country") == "US"),
+            col("a") > 0,
+            col("country") == "US",
+        ),
+    ],
+    ids=["data_only", "partition_only", "mixed_and"],
+)
+def test_apply_predicate_splits_data_and_partition(
+    tmp_path, partitioned, predicate, expected_data, expected_partition
+):
+    op = (_mk_partitioned_read_files if partitioned else _mk_read_files)(tmp_path)
+
+    new_op = op.apply_predicate(predicate)
+
+    assert isinstance(new_op, ReadFiles) and new_op is not op
+    new_scanner = new_op.scanner
+    assert isinstance(new_scanner, ParquetScanner)
+    _assert_pred_equals(new_scanner.predicate, expected_data)
+    _assert_pred_equals(new_scanner.partition_predicate, expected_partition)
+    # Original scanner untouched.
+    orig_scanner = op.scanner
+    assert isinstance(orig_scanner, ParquetScanner)
+    assert orig_scanner.predicate is None
+    assert orig_scanner.partition_predicate is None
+
+
+def test_apply_predicate_mixed_or_keeps_filter_above(tmp_path):
+    op = _mk_partitioned_read_files(tmp_path)
+
+    # Mixed-column ``OR`` can't be safely split — neither bucket is
+    # populated, so ``apply_predicate`` returns ``self`` and the rule
+    # leaves the ``Filter`` above ``ReadFiles`` untouched.
+    result = op.apply_predicate((col("a") > 0) | (col("country") == "US"))
+    assert result is op
+
+
+def test_apply_predicate_mixed_and_with_unsplittable_residual(tmp_path):
+    op = _mk_partitioned_read_files(tmp_path)
+
+    # Top-level ``AND`` of one pure-data, one pure-partition, and one
+    # mixed-OR conjunct: the first two push, the OR stays as a residual
+    # ``Filter`` so we don't silently drop it.
+    pure_data = col("a") > 0
+    pure_partition = col("country") == "US"
+    mixed_or = (col("a") < 100) | (col("country") == "CA")
+
+    result = op.apply_predicate(pure_data & pure_partition & mixed_or)
+
+    assert isinstance(result, Filter)
+    new_read = result.input_dependencies[0]
+    assert isinstance(new_read, ReadFiles)
+    new_scanner = new_read.scanner
+    assert isinstance(new_scanner, ParquetScanner)
+    _assert_pred_equals(new_scanner.predicate, pure_data)
+    _assert_pred_equals(new_scanner.partition_predicate, pure_partition)
+    # The residual carried by the new Filter is exactly the mixed-OR
+    # conjunct that couldn't be pushed.
+    _assert_pred_equals(result.predicate_expr, mixed_or)

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

@@ -0,0 +1,46 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Union
+
+from ray.data._internal.compute import ComputeStrategy
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorPreservesSchema,
+)
+from ray.data._internal.logical.operators.map_operator import AbstractMap
+from ray.data.datasource.datasink import Datasink
+from ray.data.datasource.datasource import Datasource
+
+__all__ = [
+    "Write",
+]
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class Write(AbstractMap, LogicalOperatorPreservesSchema):
+    """Logical operator for write."""
+
+    datasink_or_legacy_datasource: Union[Datasink, Datasource]
+    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
+    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
+    compute: Optional[ComputeStrategy] = None
+    write_args: Dict[str, Any] = field(default_factory=dict)
+    can_modify_num_rows: bool = field(init=False, default=True)
+    min_rows_per_bundled_input: Optional[int] = field(init=False)
+    ray_remote_args_fn: None = field(init=False, default=None)
+    per_block_limit: Optional[int] = None
+
+    def __post_init__(self):
+        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
+        if isinstance(self.datasink_or_legacy_datasource, Datasink):
+            min_rows_per_bundled_input = (
+                self.datasink_or_legacy_datasource.min_rows_per_write
+            )
+        else:
+            min_rows_per_bundled_input = None
+        if self.compute is None:
+            from ray.data._internal.compute import TaskPoolStrategy
+
+            object.__setattr__(self, "compute", TaskPoolStrategy())
+        object.__setattr__(
+            self, "min_rows_per_bundled_input", min_rows_per_bundled_input
+        )

python/ray/data/_internal/logical/optimizers.py

@@ -0,0 +1,99 @@
+from typing import TYPE_CHECKING, List, Tuple
+
+from ray.data._internal.planner import create_planner
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.execution_callback import ExecutionCallback
+
+from .ruleset import Ruleset
+from ray.data._internal.logical.interfaces import (
+    LogicalPlan,
+    Optimizer,
+    PhysicalPlan,
+    Rule,
+)
+from ray.data._internal.logical.rules import (
+    CombineShuffles,
+    CommonSubExprElimination,
+    ConfigureMapTaskMemoryUsingOutputSize,
+    FuseOperators,
+    InheritTargetMaxBlockSizeRule,
+    LimitPushdownRule,
+    PredicatePushdown,
+    ProjectionPushdown,
+    SetReadParallelismRule,
+)
+from ray.util.annotations import DeveloperAPI
+
+_LOGICAL_RULESET = Ruleset(
+    [
+        LimitPushdownRule,
+        ProjectionPushdown,
+        PredicatePushdown,
+        CombineShuffles,
+    ]
+)
+
+_PHYSICAL_RULESET = Ruleset(
+    [
+        InheritTargetMaxBlockSizeRule,
+        SetReadParallelismRule,
+        FuseOperators,
+        ConfigureMapTaskMemoryUsingOutputSize,
+    ]
+)
+
+
+@DeveloperAPI
+def get_logical_ruleset() -> Ruleset:
+    return _LOGICAL_RULESET
+
+
+@DeveloperAPI
+def get_physical_ruleset() -> Ruleset:
+    return _PHYSICAL_RULESET
+
+
+class LogicalOptimizer(Optimizer):
+    """The optimizer for logical operators."""
+
+    @property
+    def rules(self) -> List[Rule]:
+        return [rule_cls() for rule_cls in get_logical_ruleset()]
+
+    def _post_optimize(self, plan: LogicalPlan) -> LogicalPlan:
+        # CommonSubExprElimination is only supposed to run once
+        # isolated from the optimizer rule loop as it applies to
+        # a single Projection operator not a chain of operators.
+        return CommonSubExprElimination().apply(plan)
+
+
+class PhysicalOptimizer(Optimizer):
+    """The optimizer for physical operators."""
+
+    @property
+    def rules(self) -> List[Rule]:
+        return [rule_cls() for rule_cls in get_physical_ruleset()]
+
+
+def get_execution_plan(
+    logical_plan: LogicalPlan,
+) -> Tuple[PhysicalPlan, List["ExecutionCallback"]]:
+    """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.
+    """
+    # 1. Logical -> Logical (Optimized)
+    optimized_logical_plan = LogicalOptimizer().optimize(logical_plan)
+
+    # 2. Rewire Logical -> Logical (Optimized)
+    logical_plan._dag = optimized_logical_plan.dag
+
+    # 3. Logical (Optimized) -> Physical
+    physical_plan, callbacks = create_planner().plan(optimized_logical_plan)
+
+    # 4. Physical (Optimized) -> Physical
+    return PhysicalOptimizer().optimize(physical_plan), callbacks

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

@@ -0,0 +1,32 @@
+"""Expose rule classes in ray.data._internal.logical.rules."""
+
+from .combine_shuffles import CombineShuffles
+from .common_subexpr_elimination import CommonSubExprElimination
+from .configure_map_task_memory import (
+    ConfigureMapTaskMemoryRule,
+    ConfigureMapTaskMemoryUsingOutputSize,
+)
+from .inherit_target_max_block_size import InheritTargetMaxBlockSizeRule
+from .limit_pushdown import LimitPushdownRule
+from .operator_fusion import FuseOperators, are_remote_args_compatible
+from .predicate_pushdown import PredicatePushdown
+from .projection_pushdown import ProjectionPushdown
+from .set_read_parallelism import (
+    SetReadParallelismRule,
+    compute_additional_split_factor,
+)
+
+__all__ = [
+    "CombineShuffles",
+    "CommonSubExprElimination",
+    "ConfigureMapTaskMemoryRule",
+    "ConfigureMapTaskMemoryUsingOutputSize",
+    "FuseOperators",
+    "InheritTargetMaxBlockSizeRule",
+    "LimitPushdownRule",
+    "PredicatePushdown",
+    "ProjectionPushdown",
+    "SetReadParallelismRule",
+    "are_remote_args_compatible",
+    "compute_additional_split_factor",
+]

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

@@ -0,0 +1,86 @@
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalPlan,
+    Plan,
+    Rule,
+)
+from ray.data._internal.logical.operators import (
+    Aggregate,
+    Repartition,
+    Sort,
+    StreamingRepartition,
+)
+
+__all__ = [
+    "CombineShuffles",
+]
+
+
+class CombineShuffles(Rule):
+    """This logical rule combines chained shuffles together. For example,
+    ``Repartition`` and ``StreamingRepartition`` ops fusing them into a single one.
+    """
+
+    def apply(self, plan: Plan) -> Plan:
+        assert isinstance(plan, LogicalPlan)
+
+        original_dag = plan.dag
+        transformed_dag = original_dag._apply_transform(self._combine)
+
+        if transformed_dag is original_dag:
+            return plan
+
+        # TODO replace w/ Plan.copy
+        return LogicalPlan(
+            dag=transformed_dag,
+            context=plan.context,
+        )
+
+    @classmethod
+    def _combine(self, op: LogicalOperator) -> LogicalOperator:
+        # Repartitions should have exactly 1 input
+        if len(op.input_dependencies) != 1:
+            return op
+
+        input_op = op.input_dependencies[0]
+
+        if isinstance(input_op, Repartition) and isinstance(op, Repartition):
+            shuffle = input_op.shuffle or op.shuffle
+            return Repartition(
+                num_outputs=op.num_outputs,
+                input_dependencies=[input_op.input_dependencies[0]],
+                shuffle=shuffle,
+                keys=op.keys,
+                sort=op.sort,
+            )
+        elif isinstance(input_op, StreamingRepartition) and isinstance(
+            op, StreamingRepartition
+        ):
+            strict = input_op.strict or op.strict
+            return StreamingRepartition(
+                target_num_rows_per_block=op.target_num_rows_per_block,
+                input_dependencies=[input_op.input_dependencies[0]],
+                strict=strict,
+            )
+        elif isinstance(input_op, Repartition) and isinstance(op, Aggregate):
+            return Aggregate(
+                key=op.key,
+                aggs=op.aggs,
+                input_dependencies=[input_op.input_dependencies[0]],
+                num_partitions=op.num_partitions,
+            )
+        elif isinstance(input_op, StreamingRepartition) and isinstance(op, Repartition):
+            return Repartition(
+                num_outputs=op.num_outputs,
+                input_dependencies=[input_op.input_dependencies[0]],
+                shuffle=op.shuffle,
+                keys=op.keys,
+                sort=op.sort,
+            )
+        elif isinstance(input_op, Sort) and isinstance(op, Sort):
+            return Sort(
+                sort_key=op.sort_key,
+                input_dependencies=[input_op.input_dependencies[0]],
+            )
+
+        return op

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

@@ -0,0 +1,459 @@
+import uuid
+from collections import defaultdict
+from dataclasses import dataclass, replace
+from typing import Dict, Hashable, List, Optional
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalPlan,
+    Rule,
+)
+from ray.data._internal.logical.operators import CSE_TEMP_COLUMN_PREFIX, Project
+from ray.data._internal.planner.plan_expression.expression_visitors import (
+    _ExpressionOccurrence,
+    _StructuralFingerprintOccurrenceCollector,
+    _StructuralFingerprintVisitor,
+)
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    DownloadExpr,
+    Expr,
+    LiteralExpr,
+    MonotonicallyIncreasingIdExpr,
+    RandomExpr,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+    UUIDExpr,
+)
+
+__all__ = ["CommonSubExprElimination"]
+
+
+@dataclass
+class _Candidate:
+    """A structurally duplicated expression that CSE may materialize once."""
+
+    expr: Expr
+    key: Hashable
+    occurrences: List[_ExpressionOccurrence]
+    temp_name: Optional[str] = None
+
+    @property
+    def max_depth(self) -> int:
+        """Return the deepest AST depth at which this candidate appears."""
+        return max(o.depth for o in self.occurrences)
+
+
+_IGNORED_CSE_ROOT_TYPES = (ColumnExpr, LiteralExpr, AliasExpr, StarExpr)
+
+
+def _is_ignored_cse_root(expr: Expr) -> bool:
+    return isinstance(expr, _IGNORED_CSE_ROOT_TYPES)
+
+
+def _collect_occurrences(exprs: List[Expr]) -> List[_ExpressionOccurrence]:
+    """Collect structural occurrences from projection expressions.
+
+    Args:
+        exprs: Visible expression list from a ``Project`` operator.
+
+    Returns:
+        A flat list of occurrences, one per visited AST node, with each node's
+        structural key and depth computed by
+        ``_StructuralFingerprintOccurrenceCollector``.
+    """
+    collector = _StructuralFingerprintOccurrenceCollector()
+    for expr in exprs:
+        collector.visit(expr)
+
+    return collector.get_occurrences()
+
+
+def _add_to_structural_group(
+    groups: List[List[_ExpressionOccurrence]], occurrence: _ExpressionOccurrence
+) -> None:
+    """Add one occurrence to the first structurally equivalent group.
+
+    Args:
+        groups: Existing structural groups for one fingerprint key. This list is
+            mutated in place.
+        occurrence: Occurrence with the same fingerprint key to place.
+
+    Returns:
+        ``None``. The occurrence is appended to the first group whose
+        representative passes ``structurally_equals``; otherwise, a new
+        singleton group is appended.
+    """
+    for group in groups:
+        if occurrence.expr.structurally_equals(group[0].expr):
+            group.append(occurrence)
+            return
+    groups.append([occurrence])
+
+
+def _find_candidates(exprs: List[Expr]) -> List[_Candidate]:
+    """Find duplicated, CSE-eligible sub-expressions in a projection.
+
+    Args:
+        exprs: Visible expression list from a ``Project`` operator.
+
+    Returns:
+        Candidates whose non-ignored root expressions occur more than once.
+        Ignored roots such as columns, literals, aliases, and stars are skipped
+        before structural grouping so wide projections do not pay unnecessary
+        exact-comparison cost for leaves that will never be materialized.
+    """
+    occurrences = _collect_occurrences(exprs)
+
+    buckets: Dict[Hashable, List[List[_ExpressionOccurrence]]] = defaultdict(list)
+    for occurrence in occurrences:
+        # Skip ignored roots (columns, literals, aliases, stars) and non-idempotent
+        # expressions (random/uuid/monotonically_increasing_id). Materializing a
+        # non-idempotent sub-expression once would incorrectly collapse independent
+        # evaluations (e.g. two ``uuid()`` calls) into a single shared column.
+        if (
+            not _is_ignored_cse_root(occurrence.expr)
+            and occurrence.expr.is_idempotent()
+        ):
+            _add_to_structural_group(buckets[occurrence.key], occurrence)
+
+    candidates: List[_Candidate] = []
+    for key, structural_groups in buckets.items():
+        for group in structural_groups:
+            if len(group) <= 1:
+                continue
+
+            candidates.append(
+                _Candidate(
+                    expr=group[0].expr,
+                    key=key,
+                    occurrences=group,
+                )
+            )
+
+    return candidates
+
+
+def _assign_temp_names(candidates: List[_Candidate]) -> None:
+    """Assign hidden temporary column names to selected CSE candidates.
+
+    Args:
+        candidates: Candidate list returned by ``_find_candidates``. Each
+            candidate is mutated in place.
+    """
+    project_token = uuid.uuid4().hex[:6]
+
+    for i, candidate in enumerate(candidates):
+        temp_name = f"{CSE_TEMP_COLUMN_PREFIX}{project_token}_{i}"
+        candidate.temp_name = temp_name
+
+
+@dataclass(frozen=True)
+class _Replacement:
+    expr: Expr
+    key: Hashable
+    temp_name: str
+
+
+class _ReplacementIndex:
+    """Lookup structure for exact replacement matches during expression rewrite."""
+
+    def __init__(self, replacements: List[_Replacement]):
+        """Build a fingerprint bucket index from replacement records.
+
+        Args:
+            replacements: Zero or more replacements that are already safe to
+                reference.
+        """
+        self._fingerprint = _StructuralFingerprintVisitor()
+        self._replacements_by_key: Dict[Hashable, List[_Replacement]] = defaultdict(
+            list
+        )
+        for replacement in replacements:
+            self._replacements_by_key[replacement.key].append(replacement)
+
+    def find(self, expr: Expr) -> Optional[_Replacement]:
+        """Return the replacement for an exactly matching expression, if any.
+
+        Args:
+            expr: Expression node being considered for replacement.
+
+        Returns:
+            The matching replacement when both the structural fingerprint and
+            ``structurally_equals`` match; otherwise, ``None``.
+        """
+        key = self._fingerprint.visit(expr)
+        for replacement in self._replacements_by_key.get(key, []):
+            if expr.structurally_equals(replacement.expr):
+                return replacement
+        return None
+
+
+def _candidate_to_replacement(candidate: _Candidate) -> _Replacement:
+    """Convert a temp-named candidate into a rewriter replacement.
+
+    Args:
+        candidate: Candidate after temp-name assignment.
+
+    Returns:
+        An immutable replacement carrying the representative expression,
+        fingerprint key, and temp column name.
+    """
+    assert candidate.temp_name is not None
+    return _Replacement(
+        expr=candidate.expr,
+        key=candidate.key,
+        temp_name=candidate.temp_name,
+    )
+
+
+class _CSEExpressionRewriter:
+    """Rewrite expression trees to read already-materialized CSE temp columns."""
+
+    def __init__(self, replacements: List[_Replacement]):
+        """Create a rewriter over the replacements available in this phase.
+
+        Args:
+            replacements: Replacement list that may be referenced while
+                rewriting.
+        """
+        self._replacement_index = _ReplacementIndex(replacements)
+
+    def rewrite_visible_expr(self, expr: Expr) -> Expr:
+        """Rewrite a user-visible projection expression.
+
+        Args:
+            expr: One expression from ``Project.exprs`` after CSE candidates
+                have been selected.
+
+        Returns:
+            A semantically equivalent expression where the root itself may be
+            replaced by a temp column when the entire visible expression is
+            common.
+        """
+        return self._rewrite(expr, allow_root_replacement=True)
+
+    def rewrite_materialization_expr(self, expr: Expr) -> Expr:
+        """Rewrite the right-hand side for a hidden materialization expression.
+
+        Args:
+            expr: Representative expression for a candidate that will be emitted
+                as ``rhs.alias(temp_name)``.
+
+        Returns:
+            An expression that may replace child sub-expressions with previously
+            materialized temp columns, but never replaces the root candidate
+            itself.
+        """
+        return self._rewrite(expr, allow_root_replacement=False)
+
+    def _maybe_replace(self, expr: Expr) -> Optional[ColumnExpr]:
+        """Return a temp-column reference for a matching expression node.
+
+        Args:
+            expr: Expression node encountered by the recursive rewriter.
+
+        Returns:
+            ``ColumnExpr(temp_name)`` when the node matches a replacement
+            exactly; otherwise, ``None`` so traversal can continue.
+        """
+        replacement = self._replacement_index.find(expr)
+        if replacement is None:
+            return None
+        return ColumnExpr(replacement.temp_name)
+
+    def _rewrite(self, expr: Expr, *, allow_root_replacement: bool) -> Expr:
+        """Recursively rewrite one expression node.
+
+        Args:
+            expr: Expression tree to rewrite.
+            allow_root_replacement: Whether the current root may become a temp
+                column reference.
+
+        Returns:
+            The original immutable leaf/synthetic expression, a ``ColumnExpr``
+            replacement, or a rebuilt expression whose children have been
+            rewritten.
+
+        Raises:
+            TypeError: If ``expr`` has an expression node type that CSE does not
+                know how to rewrite.
+        """
+        if allow_root_replacement:
+            replacement = self._maybe_replace(expr)
+            if replacement is not None:
+                return replacement
+
+        if isinstance(expr, (ColumnExpr, LiteralExpr, StarExpr)):
+            return expr
+
+        if isinstance(expr, BinaryExpr):
+            return BinaryExpr(
+                expr.op,
+                self._rewrite(expr.left, allow_root_replacement=True),
+                self._rewrite(expr.right, allow_root_replacement=True),
+            )
+
+        if isinstance(expr, UnaryExpr):
+            return UnaryExpr(
+                expr.op,
+                self._rewrite(expr.operand, allow_root_replacement=True),
+            )
+
+        if isinstance(expr, AliasExpr):
+            visited = self._rewrite(
+                expr.expr,
+                allow_root_replacement=True,
+            )._unalias()
+            return replace(
+                expr,
+                expr=visited,
+                _is_rename=expr._is_rename and isinstance(visited, ColumnExpr),
+            )
+
+        if isinstance(expr, UDFExpr):
+            return replace(
+                expr,
+                args=[
+                    self._rewrite(arg, allow_root_replacement=True) for arg in expr.args
+                ],
+                kwargs={
+                    key: self._rewrite(value, allow_root_replacement=True)
+                    for key, value in expr.kwargs.items()
+                },
+            )
+
+        if isinstance(
+            expr,
+            (
+                DownloadExpr,
+                MonotonicallyIncreasingIdExpr,
+                RandomExpr,
+                UUIDExpr,
+            ),
+        ):
+            return expr
+
+        raise TypeError(f"Unsupported expression type for CSE rewrite: {type(expr)}")
+
+
+def _build_cse_materialization_plan(candidates: List[_Candidate]) -> List[Expr]:
+    """Build hidden expressions that compute common sub-expressions first.
+
+    Args:
+        candidates: Temp-named candidates selected for a ``Project``.
+
+    Returns:
+        Hidden alias expressions ordered deepest-first, so nested common
+        expressions are materialized before parent expressions that can
+        reference their temporary columns.
+    """
+    ordered_candidates = sorted(
+        candidates,
+        key=lambda candidate: candidate.max_depth,
+        reverse=True,
+    )
+
+    materializations: List[Expr] = []
+    replacements: List[_Replacement] = []
+
+    for candidate in ordered_candidates:
+        assert candidate.temp_name is not None
+        rewriter = _CSEExpressionRewriter(replacements)
+        rhs = rewriter.rewrite_materialization_expr(candidate.expr)
+        materializations.append(rhs.alias(candidate.temp_name))
+        replacements.append(_candidate_to_replacement(candidate))
+
+    return materializations
+
+
+def _rewrite_visible_exprs(
+    exprs: List[Expr],
+    candidates: List[_Candidate],
+) -> List[Expr]:
+    """Rewrite all user-visible projection expressions to use CSE temps.
+
+    Args:
+        exprs: Original visible ``Project.exprs``.
+        candidates: Temp-named candidates selected for that project.
+
+    Returns:
+        A new visible expression list where any matching common sub-expression,
+        including an entire root expression, reads from the corresponding hidden
+        temp column.
+    """
+    replacements = [_candidate_to_replacement(candidate) for candidate in candidates]
+    rewriter = _CSEExpressionRewriter(replacements)
+    return [rewriter.rewrite_visible_expr(expr) for expr in exprs]
+
+
+class CommonSubExprElimination(Rule):
+    """Logical optimizer rule that materializes duplicated projection expressions.
+
+    This rule rewrites eligible ``Project`` operators so hidden common
+    sub-expression aliases are evaluated before the visible projection
+    expressions that read them.
+    """
+
+    # Non-idempotent expressions (random/uuid/monotonically_increasing_id) are
+    # excluded from CSE candidacy via ``is_idempotent`` in ``_find_candidates``; the
+    # same contract guards ProjectionPushdown, PredicatePushdown, and LimitPushdown.
+    # TODO: Extend the contract to per-UDF non-determinism (see ``_IdempotencyVisitor``)
+    # and to other rules (e.g. constant folding) that can change expression
+    # evaluation count, timing, or placement.
+
+    def apply(self, plan: LogicalPlan) -> LogicalPlan:
+        """Apply CSE to every eligible project in a logical plan.
+
+        Args:
+            plan: Logical plan that may contain repeated projection
+                sub-expressions.
+
+        Returns:
+            The original plan when no project changes, or a new logical plan
+            with transformed project nodes and the original context preserved.
+        """
+        dag = plan.dag
+        new_dag = dag._apply_transform(self._try_optimize_project)
+        return LogicalPlan(new_dag, plan.context) if dag is not new_dag else plan
+
+    @classmethod
+    def _try_optimize_project(cls, op: LogicalOperator) -> LogicalOperator:
+        """Rewrite one logical operator if it is an eligible ``Project``.
+
+        Args:
+            op: Logical operator visited during DAG traversal.
+
+        Returns:
+            The original operator for non-projects, already-CSE-optimized
+            projects, or projects with no candidates. Otherwise, returns a new
+            ``Project`` containing hidden materialization expressions and
+            rewritten visible expressions while preserving the original project
+            execution settings.
+        """
+        if not isinstance(op, Project):
+            return op
+
+        if op.get_common_sub_exprs():
+            return op
+
+        candidates = _find_candidates(op.exprs)
+        if not candidates:
+            return op
+
+        _assign_temp_names(candidates)
+        common_exprs = _build_cse_materialization_plan(candidates)
+        rewritten_exprs = _rewrite_visible_exprs(op.exprs, candidates)
+
+        return Project(
+            exprs=rewritten_exprs,
+            input_dependencies=op.input_dependencies,
+            compute=op.compute,
+            ray_remote_args=op.ray_remote_args,
+            ray_remote_args_fn=op.ray_remote_args_fn,
+            per_block_limit=op.per_block_limit,
+            _common_sub_exprs=common_exprs,
+        )

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

@@ -0,0 +1,86 @@
+import abc
+import copy
+from typing import Any, Dict, Optional
+
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.logical.interfaces import Rule
+from ray.data._internal.logical.interfaces.physical_plan import PhysicalPlan
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+__all__ = [
+    "ConfigureMapTaskMemoryRule",
+    "ConfigureMapTaskMemoryUsingOutputSize",
+]
+
+
+class ConfigureMapTaskMemoryRule(Rule, abc.ABC):
+    def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
+        for op in plan.dag.post_order_iter():
+            if not isinstance(op, MapOperator):
+                continue
+
+            def ray_remote_args_fn(
+                op: MapOperator = op, original_ray_remote_args_fn=op._ray_remote_args_fn
+            ) -> Dict[str, Any]:
+                assert isinstance(op, MapOperator), op
+
+                static_ray_remote_args = copy.deepcopy(op._ray_remote_args)
+
+                dynamic_ray_remote_args = {}
+                if original_ray_remote_args_fn is not None:
+                    dynamic_ray_remote_args = original_ray_remote_args_fn()
+
+                if (
+                    "memory" not in static_ray_remote_args
+                    and "memory" not in dynamic_ray_remote_args
+                    # If this rule configures memory but the user hasn't specified
+                    # memory in the placement group, then Ray won't be able to
+                    # schedule tasks.
+                    and not any(
+                        isinstance(
+                            scheduling_strategy, PlacementGroupSchedulingStrategy
+                        )
+                        for scheduling_strategy in (
+                            static_ray_remote_args.get("scheduling_strategy"),
+                            dynamic_ray_remote_args.get("scheduling_strategy"),
+                            op.data_context.scheduling_strategy,
+                            op.data_context.scheduling_strategy_large_args,
+                        )
+                    )
+                ):
+                    memory = self.estimate_per_task_memory_requirement(op)
+                    if memory is not None:
+                        dynamic_ray_remote_args["memory"] = memory
+
+                return dynamic_ray_remote_args
+
+            op._ray_remote_args_fn = ray_remote_args_fn
+
+        return plan
+
+    @abc.abstractmethod
+    def estimate_per_task_memory_requirement(self, op: MapOperator) -> Optional[int]:
+        """Estimate the per-task memory requirement for the given map operator.
+
+        This is used to configure the `memory` argument in `ray.remote`.
+        """
+        ...
+
+
+class ConfigureMapTaskMemoryUsingOutputSize(ConfigureMapTaskMemoryRule):
+    def estimate_per_task_memory_requirement(self, op: MapOperator) -> Optional[int]:
+        # Typically, this configuration won't make a difference because
+        # `average_bytes_per_output` is usually ~128 MiB and each core usually has
+        # 4 GiB of memory. However, if `num_cpus` is small (e.g., 0.01) or
+        # `target_max_block_size` is large (e.g., 1GB), then tasks can OOM even
+        # if it just uses enough memory to produce an output block. By setting
+        # `memory` to the average output size, we can mitigate this case.
+        #
+        # We set it to 1 target block size out of assumption that *at least* 1 copy
+        # of data (to process heap) will be made during processing.
+        #
+        # Note that, unless object store memory is manually specified, by default Ray's
+        # "memory" resource is exclusive of the Object Store memory allocated on the
+        # node (i.e., its total allocatable value is Total memory - Object Store
+        # memory).
+        return op.metrics.average_bytes_per_output

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

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.logical.interfaces import PhysicalPlan, Rule
+
+__all__ = [
+    "InheritTargetMaxBlockSizeRule",
+]
+
+
+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_override is not None:
+            # Set the target block size to inherit for
+            # upstream ops.
+            target_max_block_size = dag.target_max_block_size_override
+        elif target_max_block_size is not None:
+            # Inherit from downstream op.
+            dag.override_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
+            )

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

@@ -0,0 +1,275 @@
+import copy
+import logging
+from dataclasses import is_dataclass, replace
+from typing import List, Type
+
+from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan, Rule
+from ray.data._internal.logical.operators import (
+    AbstractMap,
+    AbstractOneToOne,
+    Download,
+    Limit,
+    Project,
+    Read,
+    ReadFiles,
+    Union,
+)
+
+__all__ = [
+    "LimitPushdownRule",
+]
+
+
+logger = logging.getLogger(__name__)
+
+
+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. We are conservative and only
+    push through operators that we know for certain do not modify row counts:
+    - Project operations (column selection)
+    - MapRows operations (row-wise transformations that preserve row count)
+    - Union operations (limits are prepended to each branch)
+
+    We stop at:
+    - Any operator that can modify the number of output rows (Sort, Shuffle, Aggregate, Read etc.)
+
+    For per-block limiting, we also set per-block limits on Read operators to optimize
+    I/O while keeping the Limit operator for exact row count control.
+
+    In addition, we also fuse consecutive Limit operators into a single
+    Limit operator, i.e. `Limit[n] -> Limit[m]` becomes `Limit[min(n, m)]`.
+    """
+
+    @classmethod
+    def dependencies(cls) -> List[Type["Rule"]]:
+        # Run ProjectionPushdown and PredicatePushdown first. A `Project`
+        # (from `select_columns`, and from `read_parquet(columns=...)` which is
+        # rewired to it) or a `Filter` sits directly above the read. If limit
+        # pushdown runs first it slides the `Limit` in between that operator and
+        # the read, after which projection/predicate pushdown can no longer
+        # reach the read -- the column selection / filter is stranded above the
+        # `Limit` and the reader reads every column / every row. Applying those
+        # pushdowns first lets the selection and predicate be absorbed into the
+        # read while still adjacent, so the reader prunes columns and filters
+        # rows; the `Limit` then pushes down past the already-pruned read.
+        from ray.data._internal.logical.rules.predicate_pushdown import (
+            PredicatePushdown,
+        )
+        from ray.data._internal.logical.rules.projection_pushdown import (
+            ProjectionPushdown,
+        )
+
+        return [ProjectionPushdown, PredicatePushdown]
+
+    def apply(self, plan: LogicalPlan) -> LogicalPlan:
+        # The DAG's root is the most downstream operator.
+        def transform(node: LogicalOperator) -> LogicalOperator:
+            if isinstance(node, Limit):
+                # First, try to fuse with upstream Limit if possible (reuse fusion logic)
+                upstream_op = node.input_dependencies[0]
+                if isinstance(upstream_op, Limit):
+                    # Fuse consecutive Limits: Limit[n] -> Limit[m] becomes Limit[min(n,m)]
+                    new_limit = min(node.limit, upstream_op.limit)
+                    return Limit(
+                        new_limit,
+                        input_dependencies=[upstream_op.input_dependencies[0]],
+                    )
+
+                # If no fusion, apply pushdown logic
+                if isinstance(upstream_op, Union):
+                    return self._push_limit_into_union(node)
+                else:
+                    return self._push_limit_down(node)
+
+            return node
+
+        optimized_dag = plan.dag._apply_transform(transform)
+        return LogicalPlan(dag=optimized_dag, context=plan.context)
+
+    def _apply_limit_pushdown(self, op: LogicalOperator) -> LogicalOperator:
+        """Push down Limit operators in the given operator DAG.
+
+        This implementation uses ``LogicalOperator._apply_transform`` to
+        post-order-traverse the DAG and rewrite each ``Limit`` node via
+        :py:meth:`_push_limit_down`.
+        """
+
+        def transform(node: LogicalOperator) -> LogicalOperator:
+            if isinstance(node, Limit):
+                if isinstance(node.input_dependencies[0], Union):
+                    return self._push_limit_into_union(node)
+                return self._push_limit_down(node)
+            return node
+
+        # ``_apply_transform`` returns the (potentially new) root of the DAG.
+        return op._apply_transform(transform)
+
+    def _push_limit_into_union(self, limit_op: Limit) -> Limit:
+        """Push `limit_op` INTO every branch of its upstream Union
+        and preserve the global limit.
+
+        Existing topology:
+            child₁ , child₂ , …  ->  Union  ->  Limit
+
+        New topology:
+            child₁ -> Limit ->│
+                               │
+            child₂ -> Limit ->┤ Union ──► Limit   (original)
+                               │
+            …    -> Limit  ->│
+
+        Example (skip duplicate limit on a branch that already has it):
+            before:
+                child -> Limit(n) -> Union -> Limit(n)
+            after:
+                child -> Limit(n) -> Union -> Limit(n)   (no extra branch limit inserted)
+        """
+        union_op = limit_op.input_dependencies[0]
+        assert isinstance(union_op, Union)
+
+        def _branch_has_limit(op: LogicalOperator, limit: int) -> bool:
+            current = op
+            while (
+                isinstance(current, AbstractOneToOne)
+                and not current.can_modify_num_rows
+                and current.input_dependencies
+            ):
+                if isinstance(current, Limit):
+                    return current.limit == limit
+                # Safe to use the first dependency: current is one-to-one here.
+                current = current.input_dependencies[0]
+
+            return isinstance(current, Limit) and current.limit == limit
+
+        # Insert a branch-local Limit and push it further upstream.
+        branch_tails: List[LogicalOperator] = []
+        for child in union_op.input_dependencies:
+            # Avoid inserting a duplicate Limit on a branch that already has the same
+            # limit upstream of row-preserving ops.
+            if _branch_has_limit(child, limit_op.limit):
+                branch_tails.append(child)
+                continue
+            raw_limit = Limit(limit_op.limit, input_dependencies=[child])
+
+            if isinstance(raw_limit.input_dependencies[0], Union):
+                # This represents the limit operator appended after the union.
+                pushed_tail = self._push_limit_into_union(raw_limit)
+            else:
+                # This represents the operator that takes place of the original limit position.
+                pushed_tail = self._push_limit_down(raw_limit)
+            branch_tails.append(pushed_tail)
+
+        new_union = Union(branch_tails)
+        return Limit(limit_op.limit, input_dependencies=[new_union])
+
+    def _push_limit_down(self, limit_op: Limit) -> LogicalOperator:
+        """Push a single limit down through compatible operators conservatively.
+
+        Creates entirely new operators instead of mutating existing ones.
+        """
+        # Traverse up the DAG until we reach the first operator that meets
+        # one of the stopping conditions
+        current_op = limit_op.input_dependencies[0]
+        num_rows_preserving_ops: List[LogicalOperator] = []
+        while (
+            isinstance(current_op, AbstractOneToOne)
+            and not current_op.can_modify_num_rows
+        ):
+            if isinstance(current_op, Project) and not current_op.is_idempotent():
+                # Do not push the limit past a projection producing a non-idempotent
+                # column (e.g. monotonically_increasing_id): its value depends on row
+                # position / cardinality, which a reordered limit would change.
+                break
+
+            if isinstance(current_op, AbstractMap):
+                min_rows = current_op.min_rows_per_bundled_input
+                if min_rows is not None and min_rows > limit_op.limit:
+                    # Avoid pushing the limit past batch-based maps that require more
+                    # rows than the limit to produce stable outputs (e.g. schema).
+                    logger.info(
+                        f"Skipping push down of limit {limit_op.limit} through map {current_op} because it requires {min_rows} rows to produce stable outputs"
+                    )
+                    break
+            num_rows_preserving_ops.append(current_op)
+            current_op = current_op.input_dependencies[0]
+
+        # If we couldn't push through any operators, return original
+        if not num_rows_preserving_ops:
+            return limit_op
+        # Apply per-block limit to the deepest operator if it supports it
+        limit_input = self._apply_per_block_limit_if_supported(
+            current_op, limit_op.limit
+        )
+
+        # Build the new operator chain: Chain non-preserving number of rows -> Limit -> Operators preserving number of rows
+        new_limit = Limit(limit_op.limit, input_dependencies=[limit_input])
+        result_op = new_limit
+
+        # Recreate the intermediate operators and apply per-block limits
+        for op_to_recreate in reversed(num_rows_preserving_ops):
+            recreated_op = self._recreate_operator_with_new_input(
+                op_to_recreate, result_op
+            )
+            result_op = recreated_op
+
+        return result_op
+
+    def _apply_per_block_limit_if_supported(
+        self, op: LogicalOperator, limit: int
+    ) -> LogicalOperator:
+        """Apply per-block limit to operators that support it."""
+        if isinstance(op, AbstractMap):
+            if is_dataclass(op):
+                if isinstance(op, Read):
+                    return replace(
+                        op,
+                        per_block_limit=limit,
+                        num_outputs=op.num_outputs,
+                    )
+                if isinstance(op, ReadFiles):
+                    from ray.data._internal.datasource_v2.logical_optimizers import (
+                        SupportsLimitPushdown,
+                    )
+
+                    if isinstance(op.scanner, SupportsLimitPushdown):
+                        return replace(
+                            op,
+                            scanner=op.scanner.push_limit(limit),
+                        )
+                    return op
+                assert len(op.input_dependencies) == 1, len(op.input_dependencies)
+                return replace(
+                    op,
+                    input_dependencies=[op.input_dependencies[0]],
+                    per_block_limit=limit,
+                )
+            new_op = copy.copy(op)
+            new_op.set_per_block_limit(limit)
+            return new_op
+        return op
+
+    def _recreate_operator_with_new_input(
+        self, original_op: LogicalOperator, new_input: LogicalOperator
+    ) -> LogicalOperator:
+        """Create a new operator of the same type as original_op but with new_input as its input."""
+
+        if isinstance(original_op, Limit):
+            return Limit(original_op.limit, input_dependencies=[new_input])
+        if isinstance(original_op, Download):
+            return Download(
+                uri_column_names=original_op.uri_column_names,
+                output_bytes_column_names=original_op.output_bytes_column_names,
+                input_dependencies=[new_input],
+                ray_remote_args=original_op.ray_remote_args,
+                filesystem=original_op.filesystem,
+            )
+        if isinstance(original_op, AbstractMap) and is_dataclass(original_op):
+            return replace(original_op, input_dependencies=[new_input])
+
+        # Use copy and replace input dependencies approach
+        new_op = copy.copy(original_op)
+        new_op.input_dependencies = [new_input]
+        return new_op

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

@@ -0,0 +1,890 @@
+import itertools
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from ray.data._internal.compute import (
+    ActorPoolStrategy,
+    ComputeStrategy,
+    TaskPoolStrategy,
+)
+from ray.data._internal.execution.bundle_queue import ExactMultipleSize, RebundleQueue
+from ray.data._internal.execution.interfaces import (
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+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.limit_operator import LimitOperator
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import (
+    ShuffleReduceOp,
+)
+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 import (
+    AbstractAllToAll,
+    AbstractMap,
+    AbstractUDFMap,
+    MapBatches,
+    RandomShuffle,
+    Repartition,
+    StreamingRepartition,
+    Write,
+)
+from ray.data.datasource.file_datasink import _FileDatasink
+from ray.util.annotations import DeveloperAPI
+
+__all__ = [
+    "FuseOperators",
+    "are_remote_args_compatible",
+]
+
+
+# Scheduling strategy and label selector can be inherited from upstream operator if not specified.
+INHERITABLE_REMOTE_ARGS = ["scheduling_strategy", "label_selector"]
+
+
+logger = logging.getLogger(__name__)
+
+
+class FuseOperators(Rule):
+    """Fuses linear chains of compatible physical operators."""
+
+    def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
+        self._op_map = plan.op_map.copy()
+        # TODO(xgui): Currently we have to fuse streaming_repartition before map fusion
+        # because the result of map fusion loses the batch_size information.
+        # We should fix this by not losing the batch_size information when fusing map operators.
+        fused_dag = self._fuse_streaming_repartition_operators_in_dag(plan.dag)
+        # 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(fused_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)
+
+        # Fuse a downstream task-pool map into the V2 hash-shuffle reduce phase.
+        # Runs after map fusion so a downstream map chain is already collapsed
+        # into one TaskPoolMapOperator.
+        fused_dag = self._fuse_map_into_shuffle_reduce_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_deps(fused_dag)
+        self._update_output_deps(fused_dag)
+
+        new_plan = PhysicalPlan(fused_dag, self._op_map, plan.context)
+        return new_plan
+
+    def _remove_output_deps(self, op: PhysicalOperator) -> None:
+        for input in op.input_dependencies:
+            input._output_dependencies = []
+            self._remove_output_deps(input)
+
+    def _update_output_deps(self, op: PhysicalOperator) -> None:
+        for input in op.input_dependencies:
+            input._output_dependencies.append(op)
+            self._update_output_deps(input)
+
+    def _fuse_map_into_shuffle_reduce_in_dag(
+        self, dag: PhysicalOperator, has_downstream_limit: bool = False
+    ) -> PhysicalOperator:
+        """Starting at the given operator, traverses up the DAG and fuses a
+        task-pool map sitting directly downstream of a V2 hash-shuffle reduce
+        into the reduce (a ``ShuffleReduceOp -> TaskPoolMapOperator`` pair).
+
+        Returns the current (root) operator after completing upstream fusions.
+        """
+        if self._can_fuse_map_into_shuffle_reduce(dag, has_downstream_limit):
+            dag = self._get_fused_map_into_shuffle_reduce_operator(
+                dag, dag.input_dependencies[0]
+            )
+
+        has_downstream_limit = has_downstream_limit or isinstance(dag, LimitOperator)
+        dag._input_dependencies = [
+            self._fuse_map_into_shuffle_reduce_in_dag(upstream_op, has_downstream_limit)
+            for upstream_op in dag.input_dependencies
+        ]
+        return dag
+
+    def _can_fuse_map_into_shuffle_reduce(
+        self, dag: PhysicalOperator, has_downstream_limit: bool
+    ) -> bool:
+        """Whether ``dag`` is a task-pool map that can be fused into the V2
+        hash-shuffle reduce immediately upstream of it.
+        """
+        # `dag` must be a fusable task-pool map.
+        if not (isinstance(dag, TaskPoolMapOperator) and dag.supports_fusion()):
+            return False
+
+        # Don't fuse a map with a `concurrency=` cap: the reduce runs one task
+        # per partition with no concurrency cap, so fusing would silently ignore
+        # the limit.
+        if dag.get_max_concurrency_limit() is not None:
+            return False
+
+        # Don't fuse under a downstream limit. A standalone map is throttled at
+        # task admission , but a fused reduce task runs the map over its whole
+        # partition before the limit can stop it this could materializing far
+        # more than requested.
+        if has_downstream_limit:
+            return False
+
+        # A non-file-datasink write defers `on_write_start` to the map op (e.g.
+        # Iceberg schema evolution), which the fused reduce never runs. File
+        # datasinks run it driver-side in `Dataset.write_datasink`, so they're
+        # safe; non-write maps have no such hook.
+        # TODO: support non-file-datasink writes by running the map's `on_start`
+        # hook in the fused reduce op.
+        logical_op = self._op_map.get(dag)
+        if isinstance(logical_op, Write) and not isinstance(
+            logical_op.datasink_or_legacy_datasource, _FileDatasink
+        ):
+            return False
+
+        # The sole upstream must be a reduce that hasn't already fused with a map.
+        upstream_ops = dag.input_dependencies
+        if len(upstream_ops) != 1 or not isinstance(upstream_ops[0], ShuffleReduceOp):
+            return False
+        reduce_op = upstream_ops[0]
+        if reduce_op._fused_output_map_transformer is not None:
+            return False
+
+        return are_op_remote_args_compatible(self._op_map[reduce_op], self._op_map[dag])
+
+    def _fuse_streaming_repartition_operators_in_dag(
+        self, dag: PhysicalOperator
+    ) -> PhysicalOperator:
+        """Fuse (MapBatches -> StreamingRepartition) pair.
+
+        This will ensure the map_batch's function receive the correct number of rows.
+        We also ensure the output rows is `batch_size`.
+
+        Why don't we fuse `StreamingRepartition -> MapBatches`?
+
+        ----------------------------------------------------------------------------------------------------
+        |                      | Number of `map_batches` tasks                                             |
+        |----------------------|---------------------------------------------------------------------------|
+        | Fused                | num_input_blocks (which is <= num output blocks of StreamingRepartition) |
+        | Not fused            | num output blocks of StreamingRepartition                                 |
+        ----------------------------------------------------------------------------------------------------
+
+        When fused, the number of tasks equals the number of input blocks, which is
+        <= the number of output blocks of StreamingRepartition. If StreamingRepartition
+        is supposed to break down blocks to increase parallelism, that won't happen
+        when fused. So we don't fuse.
+
+        Why do we fuse `MapBatches -> StreamingRepartition` (when `batch_size % target_num_rows == 0`)?
+        ----------------------------------------------------------
+        |                      | Number of `map_batches` tasks  |
+        |----------------------|--------------------------------|
+        | Fused                | total_rows / batch_size        |
+        | Not fused            | total_rows / batch_size        |
+        ----------------------------------------------------------
+
+        Parallelism is unchanged, so we fuse to avoid intermediate materialization.
+        """
+        upstream_ops = dag.input_dependencies
+        while (
+            len(upstream_ops) == 1
+            and isinstance(self._op_map[dag], StreamingRepartition)
+            and isinstance(self._op_map[upstream_ops[0]], MapBatches)
+            and self._can_fuse(dag, upstream_ops[0])
+        ):
+            dag = self._get_fused_streaming_repartition_operator(dag, upstream_ops[0])
+            upstream_ops = dag.input_dependencies
+
+        dag._input_dependencies = [
+            self._fuse_streaming_repartition_operators_in_dag(upstream_op)
+            for upstream_op in upstream_ops
+        ]
+        return dag
+
+    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)
+                and self._can_fuse_map_ops(up_logical_op, down_logical_op)
+            )
+            or (
+                isinstance(up_logical_op, AbstractMap)
+                and isinstance(down_logical_op, RandomShuffle)
+            )
+            # Do not fuse Repartition operator if shuffle is disabled
+            # (i.e. using split shuffle).
+            or (
+                isinstance(up_logical_op, AbstractMap)
+                and isinstance(down_logical_op, Repartition)
+                and down_logical_op.shuffle
+            )
+        ):
+            return False
+
+        # Only fuse if the ops' remote arguments are compatible.
+        if not are_op_remote_args_compatible(up_logical_op, down_logical_op):
+            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.
+            # "auto" batch_size is resolved at task runtime, so divisibility is unknown at
+            # plan time — skip fusion and let the operators run separately.
+            return (
+                isinstance(up_logical_op.batch_size, int)
+                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_map_into_shuffle_reduce_operator(
+        self, down_op: TaskPoolMapOperator, up_op: ShuffleReduceOp
+    ) -> ShuffleReduceOp:
+        name = up_op.name + "->" + down_op.name
+
+        up_logical_op = self._op_map.pop(up_op)
+        self._op_map.pop(down_op)
+
+        fused_op = ShuffleReduceOp(
+            up_op.input_dependencies,
+            up_op.data_context,
+            num_partitions=up_op._num_partitions,
+            reduce_fn=up_op._reduce_fn,
+            disallow_block_splitting=up_op._disallow_block_splitting,
+            reduce_ray_remote_args=up_op._reduce_ray_remote_args,
+            name=name,
+            fused_output_map_transformer=down_op.get_map_transformer(),
+            fused_output_map_task_kwargs=down_op.get_map_task_kwargs(),
+            fused_output_map_target_max_block_size_override=(
+                down_op.target_max_block_size_override
+            ),
+        )
+        fused_op.set_logical_operators(
+            *up_op._logical_operators, *down_op._logical_operators
+        )
+        self._op_map[fused_op] = up_logical_op
+        return fused_op
+
+    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, (
+                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 = RebundleQueue(ExactMultipleSize(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.
+        # "auto" batch_size is resolved at task runtime, so we cannot set a fixed
+        # min_rows_per_bundle at plan time — leave it as None and let bundling use its default.
+        min_rows = (
+            None if (down_logical_op.strict or batch_size == "auto") 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=ref_bundler,
+            min_rows_per_bundle=min_rows,
+            map_task_kwargs=map_task_kwargs,
+            ray_remote_args=ray_remote_args,
+            ray_remote_args_fn=ray_remote_args_fn,
+            supports_fusion=True,
+        )
+        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)
+
+        input_op = up_logical_op.input_dependencies[0]
+        logical_op = AbstractUDFMap(
+            name,
+            [input_op],
+            up_logical_op.fn,
+            can_modify_num_rows=up_logical_op.can_modify_num_rows,
+            fn_args=up_logical_op.fn_args,
+            fn_kwargs=up_logical_op.fn_kwargs,
+            fn_constructor_args=up_logical_op.fn_constructor_args,
+            fn_constructor_kwargs=up_logical_op.fn_constructor_kwargs,
+            min_rows_per_bundled_input=min_rows,
+            compute=compute,
+            ray_remote_args_fn=ray_remote_args_fn,
+            ray_remote_args=ray_remote_args,
+        )
+        self._op_map[op] = logical_op
+        return op
+
+    @classmethod
+    def _fuse_compute_strategy(
+        cls, 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],
+    ) -> bool:
+        if (
+            up_target_max_block_size is not None
+            and down_target_max_block_size is not None
+        ):
+            # NOTE: In case of both ops overriding `target_max_block_size` only
+            #       merge them if settings are equal
+            return down_target_max_block_size == up_target_max_block_size
+
+        return True
+
+    def _get_merged_target_max_block_size(
+        self,
+        up_target_max_block_size: Optional[int],
+        down_target_max_block_size: Optional[int],
+    ) -> Optional[int]:
+        assert self._can_merge_target_max_block_size(
+            up_target_max_block_size, down_target_max_block_size
+        )
+
+        return up_target_max_block_size or 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)
+
+        # Derive min num rows per input bundle
+        min_rows_per_bundled_input = self._derive_bundle_min_num_rows(
+            down_logical_op, up_logical_op
+        )
+
+        target_max_block_size = self._get_merged_target_max_block_size(
+            up_op.target_max_block_size_override, down_op.target_max_block_size_override
+        )
+
+        compute = self._fuse_compute_strategy(
+            up_logical_op.compute, down_logical_op.compute
+        )
+        assert compute is not None
+
+        # Merge map task kwargs
+        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
+        )
+        # 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]
+
+        # Fuse on_start callbacks from both operators.
+        # This preserves deferred initialization (e.g., on_write_start for Write ops).
+        up_on_start = up_op._on_start
+        down_on_start = down_op._on_start
+        if up_on_start is not None and down_on_start is not None:
+
+            def fused_on_start(schema):
+                up_on_start(schema)
+                down_on_start(schema)
+
+            on_start = fused_on_start
+        else:
+            on_start = up_on_start or down_on_start
+
+        # Preserve StreamingRepartitionRefBundler if either operator has one.
+        # This is critical for strict-mode streaming repartition to maintain
+        # exact block size guarantees during further fusion.
+        ref_bundler = None
+        if isinstance(up_op._block_ref_bundler, RebundleQueue) and isinstance(
+            up_op._block_ref_bundler._strategy, ExactMultipleSize
+        ):
+            ref_bundler = up_op._block_ref_bundler
+        elif isinstance(down_op._block_ref_bundler, RebundleQueue) and isinstance(
+            down_op._block_ref_bundler._strategy, ExactMultipleSize
+        ):
+            ref_bundler = down_op._block_ref_bundler
+
+        isolate_workers = (
+            isinstance(up_op, TaskPoolMapOperator) and up_op.isolate_workers
+        ) or (isinstance(down_op, TaskPoolMapOperator) and down_op.isolate_workers)
+
+        # 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_override=target_max_block_size,
+            name=name,
+            compute_strategy=compute,
+            min_rows_per_bundle=min_rows_per_bundled_input
+            if ref_bundler is None
+            else None,
+            ref_bundler=ref_bundler,
+            map_task_kwargs=map_task_kwargs,
+            ray_remote_args=ray_remote_args,
+            ray_remote_args_fn=ray_remote_args_fn,
+            on_start=on_start,
+            isolate_workers=isolate_workers,
+        )
+        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_dependencies[0]
+        else:
+            # Bottom out at the source logical op (e.g. Read()).
+            input_op = up_logical_op
+
+        can_modify_num_rows = (
+            up_logical_op.can_modify_num_rows or down_logical_op.can_modify_num_rows
+        )
+        if isinstance(down_logical_op, AbstractUDFMap):
+            logical_op = AbstractUDFMap(
+                name,
+                [input_op],
+                down_logical_op.fn,
+                fn_args=down_logical_op.fn_args,
+                fn_kwargs=down_logical_op.fn_kwargs,
+                fn_constructor_args=down_logical_op.fn_constructor_args,
+                fn_constructor_kwargs=down_logical_op.fn_constructor_kwargs,
+                min_rows_per_bundled_input=min_rows_per_bundled_input,
+                compute=compute,
+                can_modify_num_rows=can_modify_num_rows,
+                ray_remote_args_fn=ray_remote_args_fn,
+                ray_remote_args=ray_remote_args,
+            )
+        else:
+            # The downstream op is AbstractMap instead of AbstractUDFMap.
+            logical_op = AbstractMap(
+                name,
+                [input_op],
+                can_modify_num_rows=can_modify_num_rows,
+                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
+
+    @classmethod
+    def _derive_bundle_min_num_rows(
+        cls,
+        down_logical_op: AbstractMap,
+        up_logical_op: AbstractMap,
+    ) -> Optional[int]:
+        us_bundle_min_rows_req = up_logical_op.min_rows_per_bundled_input
+        ds_bundle_min_rows_req = down_logical_op.min_rows_per_bundled_input
+
+        # In case neither of the ops specify `min_rows_per_bundled_input`,
+        # return None
+        if us_bundle_min_rows_req is None and ds_bundle_min_rows_req is None:
+            return None
+
+        # Target min bundle size is selected as max of upstream and downstream ones
+        # such that it could satisfy both of their requirements
+        return max(
+            ds_bundle_min_rows_req or 0,
+            us_bundle_min_rows_req or 0,
+        )
+
+    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,
+        ) -> AllToAllTransformFnResult:
+            """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_override, down_op.target_max_block_size_override
+        )
+
+        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_override=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(
+                name=name,
+                input_dependencies=[input_op],
+                ray_remote_args=ray_remote_args,
+            )
+        elif isinstance(down_logical_op, Repartition):
+            logical_op = Repartition(
+                num_outputs=down_logical_op.num_outputs,
+                input_dependencies=[input_op],
+                shuffle=down_logical_op.shuffle,
+            )
+        self._op_map[op] = logical_op
+        # Return the fused physical operator.
+        return op
+
+    @classmethod
+    def _can_fuse_map_ops(
+        cls,
+        upstream_op: AbstractMap,
+        downstream_op: AbstractMap,
+    ) -> bool:
+        if (
+            cls._fuse_compute_strategy(
+                upstream_op.compute,
+                downstream_op.compute,
+            )
+            is None
+        ):
+            return False
+
+        # Do not fuse Map operators in case:
+        #
+        #   - Upstream could (potentially) drastically modify number of rows, while
+        #   - Downstream has `min_rows_per_input_bundle` specified
+        #
+        # Fusing such transformations is not desirable as it could
+        #
+        #   - Drastically reduce parallelism for the upstream up (for ex, if
+        #   fusing ``Read->MapBatches(batch_size=...)`` with large enough batch-size
+        #   could drastically reduce parallelism level of the Read op)
+        #
+        #   - Potentially violate batching semantic by fusing
+        #   ``Filter->MapBatches(batch_size=...)``
+        #
+        if (
+            upstream_op.can_modify_num_rows
+            # For historical consistency, we allow fusing `MapBatches` even if it
+            # can modify the number of rows. Before #60448, `MapBatches` was
+            # incorrectly marked as not modifying row counts, so it was always
+            # fused. We preserve that behavior here to avoid regressions.
+            #
+            # For the full history, see https://github.com/ray-project/ray/pull/60756.
+            and not isinstance(upstream_op, MapBatches)
+        ) and downstream_op.min_rows_per_bundled_input is not None:
+            logger.debug(
+                f"Upstream operator '{upstream_op}' could be modifying # of input "
+                f"rows, while downstream operator '{downstream_op}' expects at least "
+                f"{downstream_op.min_rows_per_bundled_input} rows in a batch. "
+                f"Skipping fusion"
+            )
+
+            return False
+
+        return True
+
+
+def are_op_remote_args_compatible(
+    up_logical_op: Union[AbstractMap, AbstractAllToAll],
+    down_logical_op: Union[AbstractMap, AbstractAllToAll],
+) -> bool:
+    """Check whether two logical ops can be fused based on their Ray remote args.
+
+    Two ops are compatible only if their ``ray_remote_args`` are mergeable and
+    neither op specifies a ``ray_remote_args_fn``, since the args it generates
+    are not known ahead of time.
+    """
+    # Do not fuse if either op specifies a `ray_remote_args_fn`,
+    # since it is not known whether the generated args will be compatible.
+    # Only `AbstractMap` ops carry a `ray_remote_args_fn`.
+    for logical_op in (up_logical_op, down_logical_op):
+        if isinstance(logical_op, AbstractMap) and logical_op.ray_remote_args_fn:
+            return False
+
+    # Only fuse if the ops' remote arguments are compatible.
+    return are_remote_args_compatible(
+        up_logical_op.ray_remote_args,
+        down_logical_op.ray_remote_args,
+    )
+
+
+@DeveloperAPI
+def are_remote_args_compatible(
+    prev_args: Dict[str, Any], next_args: Dict[str, Any]
+) -> bool:
+    """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

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

@@ -0,0 +1,446 @@
+import copy
+from dataclasses import dataclass, is_dataclass, replace
+from typing import List, Optional
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorSupportsPredicatePassThrough,
+    LogicalOperatorSupportsPredicatePushdown,
+    LogicalPlan,
+    PredicatePassThroughBehavior,
+    Rule,
+)
+from ray.data._internal.logical.operators import (
+    AbstractAllToAll,
+    AbstractMap,
+    Filter,
+    Join,
+    Limit,
+    Project,
+    RandomShuffle,
+    Repartition,
+    Union,
+)
+from ray.data._internal.planner.plan_expression.expression_visitors import (
+    _ColumnSubstitutionVisitor,
+)
+from ray.data.expressions import BinaryExpr, Expr, Operation, col
+
+__all__ = [
+    "PredicatePushdown",
+]
+
+
+@dataclass(frozen=True)
+class _ConvertibilitySplit:
+    """Result of splitting a predicate by PyArrow convertibility.
+
+    Attributes:
+        convertible: The conjuncts that can be lowered to PyArrow and pushed
+            into the datasource. ``None`` if nothing is convertible.
+        residual: The conjuncts that cannot be lowered and must remain as a
+            ``Filter``. ``None`` if everything is convertible.
+    """
+
+    convertible: Optional[Expr]
+    residual: Optional[Expr]
+
+
+class PredicatePushdown(Rule):
+    """Pushes down predicates across the graph.
+
+    This rule performs the following optimizations:
+    1. Combines chained Filter operators with compatible expressions
+    2. Pushes filter expressions through eligible operators using trait-based rules
+    3. Pushes filters into data sources that support predicate pushdown
+
+    Eligibility is determined by the LogicalOperatorSupportsPredicatePassThrough trait, which operators
+    implement to declare their pushdown behavior:
+    - PASSTHROUGH: Filter passes through unchanged (Sort, Repartition, Shuffle, Limit)
+    - PASSTHROUGH_WITH_SUBSTITUTION: Filter passes through with column rebinding (Project)
+    - PUSH_INTO_BRANCHES: Filter is pushed into each branch (Union)
+    - CONDITIONAL: Filter may be pushed based on analysis (Join - analyzes which side
+      the predicate references and pushes to that side if safe for the join type)
+    """
+
+    def apply(self, plan: LogicalPlan) -> LogicalPlan:
+        """Apply predicate pushdown optimization to the logical plan."""
+        dag = plan.dag
+        new_dag = dag._apply_transform(self._try_fuse_filters)
+        new_dag = new_dag._apply_transform(self._try_push_down_predicate)
+        return LogicalPlan(new_dag, plan.context) if dag is not new_dag else plan
+
+    @classmethod
+    def _is_valid_filter_operator(cls, op: LogicalOperator) -> bool:
+        return isinstance(op, Filter) and op.is_expression_based()
+
+    @classmethod
+    def _try_fuse_filters(cls, op: LogicalOperator) -> LogicalOperator:
+        """Fuse consecutive Filter operators with compatible expressions."""
+        if not cls._is_valid_filter_operator(op):
+            return op
+
+        input_op = op.input_dependencies[0]
+        if not cls._is_valid_filter_operator(input_op):
+            return op
+
+        # Do not fuse across a non-idempotent predicate (random/uuid/
+        # monotonically_increasing_id): combining moves where a predicate is
+        # evaluated (and thus which rows it sees), changing results.
+        if (
+            not op.predicate_expr.is_idempotent()
+            or not input_op.predicate_expr.is_idempotent()
+        ):
+            return op
+
+        # Combine predicates
+        combined_predicate = op.predicate_expr & input_op.predicate_expr
+
+        # Create new filter on the input of the lower filter
+        return Filter(
+            predicate_expr=combined_predicate,
+            input_dependencies=[input_op.input_dependencies[0]],
+        )
+
+    @classmethod
+    def _can_push_filter_through_projection(
+        cls, filter_op: "Filter", projection_op: Project
+    ) -> bool:
+        """Check if a filter can be pushed through a projection operator.
+
+        Returns False (blocks pushdown) if filter references:
+        - Columns removed by select: select(['a']).filter(col('b'))
+        - Computed columns: with_column('d', 4).filter(col('d'))
+        - Old column names after rename: rename({'b': 'B'}).filter(col('b'))
+
+        Returns True (allows pushdown) for:
+        - Columns present in output: select(['a', 'b']).filter(col('a'))
+        - New column names after rename: rename({'b': 'B'}).filter(col('B'))
+        - Rename chains with name reuse: rename({'a': 'b', 'b': 'c'}).filter(col('b'))
+          (where 'b' is valid output created by a->b)
+        """
+        from ray.data._internal.planner.plan_expression.expression_visitors import (
+            _ColumnReferenceCollector,
+        )
+        from ray.data.expressions import AliasExpr, is_rename_expr
+
+        # Do not push a filter below a projection that produces a non-idempotent
+        # column (random/uuid/monotonically_increasing_id): reordering changes the row
+        # set / position the expression is evaluated over (e.g.
+        # monotonically_increasing_id reassigned over the filtered subset).
+        if not projection_op.is_idempotent():
+            return False
+
+        collector = _ColumnReferenceCollector()
+        collector.visit(filter_op.predicate_expr)
+        predicate_columns = set(collector.get_column_refs() or [])
+
+        output_columns = set()
+        new_names = set()
+        original_columns_being_renamed = set()
+
+        for expr in projection_op.exprs:
+            if expr.name is not None:
+                # Collect output column names
+                output_columns.add(expr.name)
+
+            # Process AliasExpr (computed columns or renames)
+            if isinstance(expr, AliasExpr):
+                new_names.add(expr.name)
+
+                # Check computed column: with_column('d', 4) creates AliasExpr(lit(4), 'd')
+                if expr.name in predicate_columns and not is_rename_expr(expr):
+                    return False  # Computed column
+
+                # Track old names being renamed for later check
+                if is_rename_expr(expr):
+                    original_columns_being_renamed.add(expr.expr.name)
+
+        # Check if filter references columns removed by explicit select.
+        # Valid if: projection includes all columns (star, UDF-fallback path)
+        # OR predicate columns exist in the explicit output set (typed path,
+        # where ``StarExpr`` is expanded into explicit ``col()`` refs in
+        # ``Project.__post_init__`` when the input schema is known).
+        has_required_columns = (
+            projection_op.has_star_expr() or predicate_columns.issubset(output_columns)
+        )
+        if not has_required_columns:
+            return False
+
+        # Find old names that are:
+        # 1. Being renamed away (in original_columns_being_renamed), AND
+        # 2. Referenced in predicate (in predicate_columns), AND
+        # 3. NOT recreated as new names (not in new_names)
+        #
+        # Examples:
+        #   rename({'b': 'B'}).filter(col('b'))
+        #     → {'b'} & {'b'} - {'B'} = {'b'} → BLOCKS (old name 'b' no longer exists)
+        #
+        #   rename({'a': 'b', 'b': 'c'}).filter(col('b'))
+        #     → {'a','b'} & {'b'} - {'b','c'} = {} → ALLOWS (new 'b' created by a->b)
+        #
+        #   rename({'b': 'B'}).filter(col('B'))
+        #     → {'b'} & {'B'} - {'B'} = {} → ALLOWS (using new name 'B')
+        invalid_old_names = (
+            original_columns_being_renamed & predicate_columns
+        ) - new_names
+        if invalid_old_names:
+            return False  # Old name after rename
+
+        return True
+
+    @classmethod
+    def _substitute_predicate_columns(
+        cls, predicate_expr: Expr, column_rename_map: dict[str, str]
+    ) -> Expr:
+        """Rebind column references in a predicate expression.
+
+        When pushing a predicate through a projection with column renames,
+        we need to rewrite column references from new names to old names.
+
+        Args:
+            predicate_expr: The predicate with new column names
+            column_rename_map: Mapping from old_name -> new_name
+
+        Returns:
+            The predicate rewritten to use old column names
+        """
+        # Invert the mapping: new_name -> old_name (as col expression)
+        # This is because the predicate uses new names and we need to map
+        # them back to old names
+        column_mapping = {
+            new_col: col(old_col) for old_col, new_col in column_rename_map.items()
+        }
+
+        visitor = _ColumnSubstitutionVisitor(column_mapping)
+        return visitor.visit(predicate_expr)
+
+    @classmethod
+    def _combine_with_and(
+        cls, left: Optional[Expr], right: Optional[Expr]
+    ) -> Optional[Expr]:
+        """Combine two optional predicates with ``AND``, ignoring ``None``."""
+        if left is not None and right is not None:
+            return left & right
+        return left if left is not None else right
+
+    @classmethod
+    def _split_by_convertibility(cls, predicate: Expr) -> _ConvertibilitySplit:
+        """Split a predicate into PyArrow convertible and residual parts.
+
+        Walks the top level ``AND`` chain and buckets each conjunct by whether
+        it can be lowered to PyArrow. The convertible part can be pushed into
+        the datasource while the residual part stays as a ``Filter`` above it.
+
+        Args:
+            predicate: The predicate expression to split.
+
+        Returns:
+            A ``_ConvertibilitySplit`` whose ``convertible`` and ``residual``
+            fields hold the two parts. Both are optional.
+        """
+        if isinstance(predicate, BinaryExpr) and predicate.op == Operation.AND:
+            left = cls._split_by_convertibility(predicate.left)
+            right = cls._split_by_convertibility(predicate.right)
+            return _ConvertibilitySplit(
+                convertible=cls._combine_with_and(left.convertible, right.convertible),
+                residual=cls._combine_with_and(left.residual, right.residual),
+            )
+
+        if predicate._is_pyarrow_convertible():
+            return _ConvertibilitySplit(convertible=predicate, residual=None)
+        return _ConvertibilitySplit(convertible=None, residual=predicate)
+
+    @classmethod
+    def _try_push_down_predicate(cls, op: LogicalOperator) -> LogicalOperator:
+        """Push Filter down through the operator tree."""
+        if not cls._is_valid_filter_operator(op):
+            return op
+        filter_op: Filter = op
+        input_op = filter_op.input_dependencies[0]
+        predicate_expr = filter_op.predicate_expr
+
+        # Case 1: Check if operator supports predicate pushdown (e.g., Read).
+        # The read stage never renames columns (renaming is always carried
+        # by an ``AliasExpr`` in a ``Project`` operator above the read), so
+        # the predicate above the read is already in the same column
+        # namespace the scanner sees — no rebinding is required here.
+        if (
+            isinstance(input_op, LogicalOperatorSupportsPredicatePushdown)
+            and input_op.supports_predicate_pushdown()
+        ):
+            # Datasources evaluate pushed predicates via PyArrow. A predicate
+            # that can't be lowered to PyArrow (e.g. it contains a UDF) must
+            # stay as a Filter. Split the top level AND chain so the convertible
+            # conjuncts can still be pushed while the residual ones are kept as
+            # a Filter above the read.
+            split = cls._split_by_convertibility(predicate_expr)
+
+            if split.convertible is None:
+                return filter_op
+
+            result_op = input_op.apply_predicate(split.convertible)
+
+            # If the operator is unchanged (e.g., predicate references partition columns
+            # that can't be pushed down), keep the Filter operator
+            if result_op is input_op:
+                return filter_op
+
+            # Convertible conjuncts were pushed into the read. Re-apply any
+            # residual (non-convertible) conjuncts as a Filter above it.
+            if split.residual is None:
+                return result_op
+            return Filter(predicate_expr=split.residual, input_dependencies=[result_op])
+
+        # Datasource pushdown (Case 1) only lowers PyArrow-convertible (hence
+        # idempotent) conjuncts. Beyond that, do not relocate a filter whose predicate
+        # is non-idempotent (random/uuid/monotonically_increasing_id): pushing it
+        # through a pass-through operator, into Union branches, or to a Join side
+        # changes the row set / position the expression is evaluated over.
+        if not predicate_expr.is_idempotent():
+            return filter_op
+
+        # Case 2: Check if operator allows predicates to pass through
+        if isinstance(input_op, LogicalOperatorSupportsPredicatePassThrough):
+            behavior = input_op.predicate_passthrough_behavior()
+
+            if behavior in (
+                PredicatePassThroughBehavior.PASSTHROUGH,
+                PredicatePassThroughBehavior.PASSTHROUGH_WITH_SUBSTITUTION,
+            ):
+                # Both cases push through a single input with optional column rebinding
+                assert len(input_op.input_dependencies) == 1, (
+                    f"{behavior.value} operators must have exactly 1 input, "
+                    f"got {len(input_op.input_dependencies)}"
+                )
+
+                # Apply column substitution if needed
+                if (
+                    behavior
+                    == PredicatePassThroughBehavior.PASSTHROUGH_WITH_SUBSTITUTION
+                ):
+                    # Check if we can safely push the filter through this projection
+                    if isinstance(
+                        input_op, Project
+                    ) and not cls._can_push_filter_through_projection(
+                        filter_op, input_op
+                    ):
+                        return filter_op
+
+                    rename_map = input_op.get_column_substitutions()
+                    if rename_map:
+                        predicate_expr = cls._substitute_predicate_columns(
+                            predicate_expr, rename_map
+                        )
+
+                # Push filter through and recursively try to push further
+                new_filter = Filter(
+                    predicate_expr=predicate_expr,
+                    input_dependencies=[input_op.input_dependencies[0]],
+                )
+                pushed_filter = cls._try_push_down_predicate(new_filter)
+
+                # Return input_op with the pushed filter as its input
+                return cls._clone_op_with_new_inputs(input_op, [pushed_filter])
+
+            elif behavior == PredicatePassThroughBehavior.PUSH_INTO_BRANCHES:
+                # Push into each branch (e.g., Union)
+                # Apply filter to each branch and recursively push down
+                new_inputs = []
+                for branch_op in input_op.input_dependencies:
+                    branch_filter = Filter(
+                        predicate_expr=predicate_expr, input_dependencies=[branch_op]
+                    )
+                    pushed_branch = cls._try_push_down_predicate(branch_filter)
+                    new_inputs.append(pushed_branch)
+
+                # Return operator with filtered branches
+                return cls._clone_op_with_new_inputs(input_op, new_inputs)
+
+            elif behavior == PredicatePassThroughBehavior.CONDITIONAL:
+                # Handle conditional pushdown (e.g., Join)
+                return cls._push_filter_through_conditionally(filter_op, input_op)
+
+        return filter_op
+
+    @classmethod
+    def _push_filter_through_conditionally(
+        cls, filter_op: Filter, conditional_op: LogicalOperator
+    ) -> LogicalOperator:
+        """Handle conditional pushdown for operators like Join.
+
+        For operators with multiple inputs, we can push predicates that reference
+        only one side down to that side, when semantically safe.
+        """
+        # Check if operator supports conditional pushdown by having the required method
+        if not hasattr(conditional_op, "which_side_to_push_predicate"):
+            return filter_op
+
+        push_side = conditional_op.which_side_to_push_predicate(
+            filter_op.predicate_expr
+        )
+
+        if push_side is None:
+            # Cannot push through
+            return filter_op
+
+        # Use the enum value directly as branch index
+        branch_idx = push_side.value
+
+        # Push to the appropriate branch
+        new_inputs = list(conditional_op.input_dependencies)
+        branch_filter = Filter(
+            predicate_expr=filter_op.predicate_expr,
+            input_dependencies=[new_inputs[branch_idx]],
+        )
+        new_inputs[branch_idx] = cls._try_push_down_predicate(branch_filter)
+
+        # Return operator with updated input
+        return cls._clone_op_with_new_inputs(conditional_op, new_inputs)
+
+    @classmethod
+    def _clone_op_with_new_inputs(
+        cls, op: LogicalOperator, new_inputs: List[LogicalOperator]
+    ) -> LogicalOperator:
+        """Clone an operator with new inputs.
+
+        Args:
+            op: The operator to clone
+            new_inputs: List of new input operators (can be single element list)
+
+        Returns:
+            A shallow copy of the operator with updated input dependencies
+        """
+        if isinstance(op, Limit):
+            assert len(new_inputs) == 1, len(new_inputs)
+            return Limit(op.limit, input_dependencies=[new_inputs[0]])
+        if isinstance(op, AbstractMap) and is_dataclass(op):
+            assert len(new_inputs) == 1, len(new_inputs)
+            return replace(op, input_dependencies=[new_inputs[0]])
+        if isinstance(op, AbstractAllToAll) and is_dataclass(op):
+            assert len(new_inputs) == 1, len(new_inputs)
+            kwargs = {"input_dependencies": [new_inputs[0]]}
+            if isinstance(op, Repartition):
+                kwargs["num_outputs"] = op.num_outputs
+            if isinstance(op, RandomShuffle):
+                kwargs["name"] = op.name
+            return replace(op, **kwargs)
+        if isinstance(op, Join) and is_dataclass(op):
+            assert len(new_inputs) == 2, len(new_inputs)
+            return Join(
+                new_inputs[0],
+                new_inputs[1],
+                op.join_type,
+                op.left_key_columns,
+                op.right_key_columns,
+                num_partitions=op.num_outputs,
+                left_columns_suffix=op.left_columns_suffix,
+                right_columns_suffix=op.right_columns_suffix,
+                partition_size_hint=op.partition_size_hint,
+                aggregator_ray_remote_args=op.aggregator_ray_remote_args,
+            )
+        if isinstance(op, Union) and is_dataclass(op):
+            return Union(new_inputs)
+        new_op = copy.copy(op)
+        new_op.input_dependencies = new_inputs
+        return new_op

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

@@ -0,0 +1,523 @@
+from typing import Dict, List, Optional, Set, Tuple
+
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalOperatorSupportsProjectionPushdown,
+    LogicalPlan,
+    Rule,
+)
+from ray.data._internal.logical.operators import Project
+from ray.data._internal.planner.plan_expression.expression_visitors import (
+    _ColumnReferenceCollector,
+    _ColumnSubstitutionVisitor,
+    _is_col_expr,
+)
+from ray.data.expressions import (
+    AliasExpr,
+    Expr,
+    StarExpr,
+    is_rename_expr,
+)
+
+__all__ = [
+    "ProjectionPushdown",
+]
+
+
+def _collect_referenced_columns(exprs: List[Expr]) -> Optional[List[str]]:
+    """
+    Extract all column names referenced by the given expressions.
+
+    Recursively traverses expression trees to find all ColumnExpr nodes
+    and collects their names.
+
+    Example: For expression "col1 + col2", returns {"col1", "col2"}
+    """
+    # ``StarExpr`` is eagerly expanded to explicit ``col()`` refs in
+    # ``Project.__post_init__`` when the input schema is known. So this
+    # branch is hit only on the UDF-fallback path (Project on top of an
+    # opaque-schema input like ``MapBatches``), where we can't enumerate
+    # columns and have to fall back to "all columns" (``None``).
+    if any(isinstance(expr, StarExpr) for expr in exprs):
+        return None
+
+    collector = _ColumnReferenceCollector()
+    for expr in exprs or []:
+        collector.visit(expr)
+
+    return collector.get_column_refs()
+
+
+def _analyze_upstream_project(
+    upstream_project: Project,
+) -> Tuple[Set[str], dict[str, Expr], Set[str]]:
+    """
+    Analyze what the upstream project produces and identifies removed columns.
+
+    Example: Upstream exprs [col("x").alias("y")] → removed_by_renames = {"x"} if "x" not in output
+    """
+    output_column_names = {
+        expr.name for expr in upstream_project.exprs if not isinstance(expr, StarExpr)
+    }
+
+    # Compose column definitions in the form of a mapping of
+    #   - Target column name
+    #   - Target expression
+    output_column_defs = {
+        expr.name: expr for expr in _filter_out_star(upstream_project.exprs)
+    }
+
+    # Identify upstream input columns removed by renaming (ie not propagated into
+    # its output)
+    upstream_column_renaming_map = _extract_input_columns_renaming_mapping(
+        upstream_project.exprs
+    )
+
+    return (
+        output_column_names,
+        output_column_defs,
+        set(upstream_column_renaming_map.keys()),
+    )
+
+
+def _validate_fusion(
+    downstream_project: Project,
+    upstream_has_all: bool,
+    upstream_output_columns: Set[str],
+    removed_by_renames: Set[str],
+) -> Tuple[bool, Set[str]]:
+    """
+    Validate if fusion is possible without rewriting expressions.
+
+    Args:
+        downstream_project: The downstream Project operator
+        upstream_has_all: True if the upstream Project has all columns, False otherwise
+        upstream_output_columns: Set of column names that are available in the upstream Project
+        removed_by_renames: Set of column names that are removed by renames in the upstream Project
+
+    Returns:
+        Tuple of (is_valid, missing_columns)
+        - is_valid: True if all expressions can be fused, False otherwise
+        - missing_columns: Set of column names that are referenced but not available
+
+    Example: Downstream refs "x" but upstream renamed "x" to "y" and dropped "x"
+             → (False, {"x"})
+    """
+    missing_columns = set()
+
+    for expr in downstream_project.exprs:
+        if isinstance(expr, StarExpr):
+            continue
+
+        column_refs = _collect_referenced_columns([expr])
+        column_refs_set = set(column_refs or [])
+
+        columns_from_original = column_refs_set - (
+            column_refs_set & upstream_output_columns
+        )
+
+        # Validate accessibility
+        if not upstream_has_all and columns_from_original:
+            # Example: Upstream selects ["a", "b"], Downstream refs "c" → can't fuse
+            missing_columns.update(columns_from_original)
+
+        if any(col in removed_by_renames for col in columns_from_original):
+            # Example: Upstream renames "x" to "y" (dropping "x"), Downstream refs "x" → can't fuse
+            removed_cols = {
+                col for col in columns_from_original if col in removed_by_renames
+            }
+            missing_columns.update(removed_cols)
+
+    is_valid = len(missing_columns) == 0
+    return is_valid, missing_columns
+
+
+def _would_duplicate_nonidempotent_expr(
+    upstream_project: Project,
+    downstream_project: Project,
+    upstream_column_defs: Dict[str, Expr],
+) -> bool:
+    """Return ``True`` if fusing would materialize a non-idempotent column >1 time.
+
+    Fusion inlines each upstream output column definition into every downstream
+    reference. For a non-idempotent definition (random/uuid/monotonically_increasing_id)
+    this changes its evaluation count, so we block the fusion and let the upstream
+    evaluate it exactly once.
+
+    Materialization count for an upstream column post-fusion is its downstream reference
+    multiplicity, plus one if it survives as a passthrough (the composition/star case).
+    """
+    nonidem_cols = {
+        name
+        for name, def_expr in upstream_column_defs.items()
+        if not def_expr.is_idempotent()
+    }
+    if not nonidem_cols:
+        return False
+
+    # Downstream reference multiplicity (counts ``x + x`` as 2).
+    counter = _ColumnReferenceCollector()
+    # Note: We ignore _common_sub_exprs field as CSE rule is applied post-optimization.
+    # If order is to be changed, we also need to invoke the _ColumnReferenceCollector
+    # on _common_sub_exprs.
+    for e in _filter_out_star(downstream_project.exprs):
+        counter.visit(e)
+    ref_counts = counter.get_counts()
+
+    # In the composition (downstream-star) case the upstream column also survives in
+    # the fused output unless downstream renames it away, adding one materialization.
+    passthrough: Set[str] = set()
+    if downstream_project.has_star_expr():
+        renamed_away = _extract_input_columns_renaming_mapping(downstream_project.exprs)
+        passthrough = {
+            e.name
+            for e in upstream_project.exprs
+            if not isinstance(e, StarExpr) and e.name not in renamed_away
+        }
+
+    for name in nonidem_cols:
+        total = ref_counts.get(name, 0) + (1 if name in passthrough else 0)
+        if total > 1:
+            return True
+    return False
+
+
+def _try_fuse(upstream_project: Project, downstream_project: Project) -> Project:
+    """
+    Attempt to merge two consecutive Project operations into one.
+
+    Example: Upstream: [star(), col("x").alias("y")], Downstream: [star(), (col("y") + 1).alias("z")] → Fused: [star(), (col("x") + 1).alias("z")]
+    """
+    # Check resource compatibility before attempting fusion
+    # This ensures with_column respects resource boundaries like map_batches does
+    from ray.data._internal.logical.rules.operator_fusion import (
+        FuseOperators,
+        are_op_remote_args_compatible,
+    )
+
+    # Check if remote args (num_cpus, num_gpus, etc.) are compatible and that
+    # neither op specifies a `ray_remote_args_fn`.
+    if not are_op_remote_args_compatible(upstream_project, downstream_project):
+        # Resources don't match - cannot fuse
+        return downstream_project
+
+    # Check if compute strategies are compatible
+    fused_compute = FuseOperators._fuse_compute_strategy(
+        upstream_project.compute, downstream_project.compute
+    )
+    if fused_compute is None:
+        # Compute strategies incompatible - cannot fuse
+        return downstream_project
+
+    upstream_has_star: bool = upstream_project.has_star_expr()
+
+    # TODO add validations that
+    #   - exprs only depend on input attrs (ie no dep on output of other exprs)
+
+    # Analyze upstream
+    (
+        upstream_output_cols,
+        upstream_column_defs,
+        upstream_input_cols_removed,
+    ) = _analyze_upstream_project(upstream_project)
+
+    # Validate fusion possibility
+    is_valid, missing_columns = _validate_fusion(
+        downstream_project,
+        upstream_has_star,
+        upstream_output_cols,
+        upstream_input_cols_removed,
+    )
+
+    if not is_valid:
+        # Raise KeyError to match expected error type in tests
+        raise KeyError(
+            f"Column(s) {sorted(missing_columns)} not found. "
+            f"Available columns: {sorted(upstream_output_cols) if not upstream_has_star else 'all columns (has star)'}"
+        )
+
+    if _would_duplicate_nonidempotent_expr(
+        upstream_project, downstream_project, upstream_column_defs
+    ):
+        # Fusing would inline a non-idempotent expression (random/uuid/
+        # monotonically_increasing_id) into multiple references, changing its
+        # evaluation count. Leave the two Projects unfused so the upstream evaluates
+        # it exactly once and downstream reads the materialized column.
+        return downstream_project
+
+    # Following invariants are upheld for each ``Project`` logical op:
+    #
+    #   1. ``Project``s list of expressions are bound to op's input columns **only**
+    #   (ie there could be no inter-dependency b/w expressions themselves)
+    #
+    #   2. `Each of expressions on the `Project``s list constitutes an output
+    #   column definition, where column's name is derived from ``expr.name`` and
+    #   column itself is derived by executing that expression against the op's
+    #   input block.
+    #
+    # Therefore to abide by and satisfy aforementioned invariants, when fusing
+    # 2 ``Project`` operators, following scenarios are considered:
+    #
+    #   1. Composition: downstream including (and potentially renaming) upstream
+    #      output columns (this is the case when downstream holds ``StarExpr``).
+    #
+    #   2. Projection: downstream projecting upstream output columns (by for ex,
+    #      only selecting & transforming some of the upstream output columns).
+    #
+
+    # Upstream output column refs inside downstream expressions need to be bound
+    # to upstream output column definitions to satisfy invariant #1 (common for both
+    # composition/projection cases)
+    v = _ColumnSubstitutionVisitor(upstream_column_defs)
+
+    rebound_downstream_exprs = [
+        v.visit(e) for e in _filter_out_star(downstream_project.exprs)
+    ]
+
+    if not downstream_project.has_star_expr():
+        # Projection case: this is when downstream is a *selection* (ie, not including
+        # the upstream columns with ``StarExpr``). With eager expansion of
+        # ``StarExpr`` in ``Project.__post_init__`` this is the common case
+        # for typed chains (no ``StarExpr`` reaches the optimizer).
+        #
+        # Example:
+        #   Upstream: Project([col("a").alias("b")])
+        #   Downstream: Project([col("b").alias("c")])
+        #
+        #   Result: Project([col("a").alias("c")])
+        new_exprs = rebound_downstream_exprs
+    else:
+        # Composition case: downstream has ``StarExpr`` (entailing that downstream
+        # output will be including all of the upstream output columns). This
+        # is the UDF-fallback path; for typed chains
+        # ``Project.__post_init__`` would have replaced the ``StarExpr``
+        # with explicit ``col()`` refs already.
+        #
+        # Example 1:
+        #   Upstream: [star(), col("a").alias("b")],
+        #   Downstream: [star(), col("b").alias("c")]
+        #
+        #   Result: [star(), col("a").alias("b"), col("a").alias("c")]
+        #
+        # Example 2:
+        #   Input (columns): ["a", "b"]
+        #   Upstream: [star({"b": "z"}), col("a").alias("x")],
+        #   Downstream: [star({"x": "y"}), col("z")]
+        #
+        #   Result: [star(), col("a").alias("y"), col("b").alias("z")]
+
+        # Extract downstream's input column rename map (downstream inputs are
+        # upstream's outputs)
+        downstream_input_column_rename_map = _extract_input_columns_renaming_mapping(
+            downstream_project.exprs
+        )
+        # Collect upstream output column expression "projected" to become
+        # downstream expressions
+        projected_upstream_output_col_exprs = []
+
+        # When fusing 2 projections
+        for e in upstream_project.exprs:
+            # NOTE: We have to filter out upstream output columns that are
+            #       being *renamed* by downstream expression
+            if e.name not in downstream_input_column_rename_map:
+                projected_upstream_output_col_exprs.append(e)
+
+        new_exprs = projected_upstream_output_col_exprs + rebound_downstream_exprs
+
+    return Project(
+        exprs=new_exprs,
+        input_dependencies=[upstream_project.input_dependencies[0]],
+        compute=fused_compute,
+        ray_remote_args=downstream_project.ray_remote_args,
+    )
+
+
+def _filter_out_star(exprs: List[Expr]) -> List[Expr]:
+    return [e for e in exprs if not isinstance(e, StarExpr)]
+
+
+class ProjectionPushdown(Rule):
+    """
+    Optimization rule that pushes projections (column selections) down the query plan.
+
+    This rule performs two optimizations:
+    1. Fuses consecutive Project operations to eliminate redundant projections
+    2. Pushes projections into data sources (e.g., Read operations) to enable
+       column pruning at the storage layer
+    """
+
+    def apply(self, plan: LogicalPlan) -> LogicalPlan:
+        """Apply projection pushdown optimization to the entire plan."""
+        dag = plan.dag
+        # Insert a pruning projection below consuming ops (e.g. ``Aggregate``)
+        # first, so the fuse/push steps can carry the narrowed columns into the
+        # read.
+        new_dag = dag._apply_transform(self._prune_aggregate_input)
+        new_dag = new_dag._apply_transform(self._try_fuse_projects)
+        new_dag = new_dag._apply_transform(self._push_projection_into_read_op)
+        return LogicalPlan(new_dag, plan.context) if dag is not new_dag else plan
+
+    @classmethod
+    def _prune_aggregate_input(cls, op: LogicalOperator) -> LogicalOperator:
+        """Insert a ``Project`` below an ``Aggregate`` that keeps only the
+        columns it consumes (group keys + each aggregation's target column).
+
+        The aggregation drops every other column anyway, so pruning them before
+        the shuffle avoids dragging unused (often wide) columns through it --
+        which otherwise inflates both the aggregator's memory reservation and
+        the bytes shuffled. The inserted projection is fused/pushed toward the
+        read by the steps in ``apply``.
+        """
+        from dataclasses import replace
+
+        from ray.data._internal.logical.operators.all_to_all_operator import Aggregate
+        from ray.data.aggregate import AggregateFnV2
+        from ray.data.expressions import col
+
+        if not isinstance(op, Aggregate):
+            return op
+
+        keys = op.key if isinstance(op.key, list) else ([op.key] if op.key else [])
+        required: List[str] = list(keys)
+        for agg in op.aggs:
+            # A generic ``AggregateFn`` may read arbitrary columns, so only
+            # prune when every aggregation declares the column it reads.
+            if not isinstance(agg, AggregateFnV2):
+                return op
+            target = agg.get_target_column()
+            if target is not None:
+                required.append(target)
+
+        # Order-preserving dedup; empty means nothing safe to prune to (e.g. a
+        # global count reading no columns).
+        required = list(dict.fromkeys(required))
+        if not required:
+            return op
+
+        input_op = op.input_dependencies[0]
+        schema = input_op.infer_schema()
+        if schema is None or not hasattr(schema, "names"):
+            return op  # unknown schema: can't prove pruning helps
+
+        # Insert only when ``required`` is a strict subset of the input columns:
+        # this guarantees there's something to drop and keeps the rule
+        # idempotent (once the input yields exactly ``required`` nothing more is
+        # inserted, so the fixed-point optimizer terminates).
+        if not set(required) < set(schema.names):
+            return op
+
+        prune = Project(exprs=[col(c) for c in required], input_dependencies=[input_op])
+        return replace(op, input_dependencies=[prune])
+
+    @classmethod
+    def _try_fuse_projects(cls, op: LogicalOperator) -> LogicalOperator:
+        """
+        Optimize a single Project operator.
+
+        Steps:
+        1. Iteratively fuse with upstream Project operations
+        2. Push the resulting projection into the data source if possible
+        """
+        if not isinstance(op, Project):
+            return op
+
+        # Step 1: Iteratively fuse with upstream Project operations
+        current_project: Project = op
+
+        upstream_op = current_project.input_dependencies[0]
+        if not isinstance(upstream_op, Project):
+            return op
+
+        fused = _try_fuse(upstream_op, current_project)
+
+        return fused
+
+    @classmethod
+    def _push_projection_into_read_op(cls, op: LogicalOperator) -> LogicalOperator:
+
+        if not isinstance(op, Project):
+            return op
+
+        current_project: Project = op
+
+        # Step 2: Push projection into the data source if supported
+        input_op = current_project.input_dependencies[0]
+        if (
+            isinstance(input_op, LogicalOperatorSupportsProjectionPushdown)
+            and input_op.supports_projection_pushdown()
+        ):
+            # Collect the set of input columns this ``Project`` reads.
+            # ``None`` means "all columns" (a ``StarExpr`` is present, so
+            # we can't enumerate). Renames are NOT pushed into the read:
+            # column renaming always stays as an ``AliasExpr`` in a
+            # ``Project`` on top of the read. This keeps the read stage
+            # in a single column namespace (the scanner's on-disk names)
+            # and avoids the predicate-pushdown rebinding dance.
+            if current_project.has_star_expr():
+                # UDF-fallback path: if ``StarExpr`` survives to this rule
+                # the input schema was unknown at construction time, so
+                # we can't enumerate columns and have to push "all".
+                required_columns = None
+            else:
+                # Otherwise, collect required columns to push projection down
+                # into the reader. (``Project.__post_init__`` expands
+                # ``StarExpr`` to explicit ``col()`` refs when the input
+                # schema is known, so this is the common path.)
+                required_columns = _collect_referenced_columns(current_project.exprs)
+
+            # Build a pure-prune projection map (identity, no renames).
+            projection_map = (
+                None
+                if required_columns is None
+                else {name: name for name in required_columns}
+            )
+            projected_input_op = input_op.apply_projection(projection_map)
+
+            # If the ``Project`` is a pure-prune (only ``col()`` refs,
+            # no renames, no computed expressions), the projection
+            # pushdown into the read op fully subsumes it — discard it.
+            # Otherwise (renames or computed expressions present), keep
+            # the ``Project`` on top so it runs above the (pruned) read.
+            # Physical operator fusion later merges the kept ``Project``
+            # into the same ``MapOperator`` as the read, so the
+            # runtime cost is the same either way.
+            has_renames = any(isinstance(e, AliasExpr) for e in current_project.exprs)
+            all_col_refs = all(
+                _is_col_expr(e) for e in _filter_out_star(current_project.exprs)
+            )
+            is_pure_prune = not has_renames and all_col_refs
+            if is_pure_prune:
+                return projected_input_op
+
+            return Project(
+                exprs=current_project.exprs,
+                input_dependencies=[projected_input_op],
+                compute=current_project.compute,
+                ray_remote_args=current_project.ray_remote_args,
+            )
+
+        return current_project
+
+
+def _extract_input_columns_renaming_mapping(
+    projection_exprs: List[Expr],
+) -> Dict[str, str]:
+    """Fetches renaming mapping of all input columns names being renamed (replaced).
+    Format is source column name -> new column name.
+    """
+
+    return dict(
+        [
+            _get_renaming_mapping(expr)
+            for expr in _filter_out_star(projection_exprs)
+            if is_rename_expr(expr)
+        ]
+    )
+
+
+def _get_renaming_mapping(expr: Expr) -> Tuple[str, str]:
+    assert is_rename_expr(expr)
+
+    alias: AliasExpr = expr
+
+    return alias.expr.name, alias.name

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

@@ -0,0 +1,148 @@
+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 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
+
+__all__ = [
+    "SetReadParallelismRule",
+    "compute_additional_split_factor",
+]
+
+
+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: Optional[int],
+    cur_additional_split_factor: Optional[int] = None,
+) -> Tuple[int, str, int, Optional[int]]:
+    """Returns parallelism to use and the min safe parallelism to avoid OOMs."""
+
+    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}, block size {expected_block_size}"
+        )
+        if target_max_block_size is None:
+            # Unlimited block size -> no extra splits
+            size_based_splits = 1
+        else:
+            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(plan, op, logical_op)
+            ops += op.input_dependencies
+
+        return plan
+
+    def _apply(self, plan: PhysicalPlan, op: PhysicalOperator, logical_op: Read):
+        estimated_in_mem_bytes = logical_op.infer_metadata().size_bytes
+
+        (
+            detected_parallelism,
+            reason,
+            estimated_num_blocks,
+            k,
+        ) = compute_additional_split_factor(
+            logical_op.datasource_or_legacy_reader,
+            logical_op.parallelism,
+            estimated_in_mem_bytes,
+            op.target_max_block_size_override or op.data_context.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}."
+            )
+        plan.op_map[op] = 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}")

python/ray/data/_internal/logical/ruleset.py

@@ -0,0 +1,136 @@
+import collections
+from dataclasses import dataclass, field
+from typing import Dict, Iterator, List, Optional, Set, Tuple, Type
+
+from ray.data._internal.logical.interfaces import Rule
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class Ruleset:
+    """A collection of rules to apply to a plan.
+
+    This is a utility class to ensure that, if rules depend on each other, they're
+    applied in a correct order.
+    """
+
+    @dataclass(frozen=True)
+    class _Node:
+        rule: Type[Rule]
+        dependents: List["Ruleset._Node"] = field(default_factory=list)
+
+    def __init__(self, rules: Optional[List[Type[Rule]]] = None):
+        if rules is None:
+            rules = []
+
+        self._rules = list(rules)
+
+    def add(self, rule: Type[Rule]):
+        if rule in self._rules:
+            raise ValueError(f"Rule {rule} already in ruleset")
+
+        self._rules.append(rule)
+
+        if self._contains_cycle():
+            raise ValueError("Cannot add rule that would create a cycle")
+
+    def remove(self, rule: Type[Rule]):
+        if rule not in self._rules:
+            raise ValueError(f"Rule {rule} not found in ruleset")
+
+        self._rules.remove(rule)
+
+    def __iter__(self) -> Iterator[Type[Rule]]:
+        """Iterate over the rules in this ruleset.
+
+        This method yields rules in dependency order. For example, if B depends on A,
+        then this method yields A before B. Each rule is yielded exactly once, and a
+        rule is only yielded once *all* of its dependencies have been yielded (so a
+        rule that several others must precede is not emitted early or duplicated).
+        Insertion order breaks ties among rules that are ready at the same time.
+        """
+        order, _ = self._topological_order()
+        for node in order:
+            yield node.rule
+
+    def _topological_order(self) -> Tuple[List["Ruleset._Node"], int]:
+        """Order the nodes by dependency using Kahn's algorithm.
+
+        Returns the topologically-ordered nodes and the total node count.
+        A node is enqueued the moment its in-degree (count of not-yet-emitted
+        dependencies) hits zero; since an in-degree only decreases and we
+        enqueue solely on the zero-crossing, each node is emitted exactly once.
+        Insertion order breaks ties among nodes that are ready together.
+
+        Nodes that participate in a cycle never reach in-degree zero, so they
+        are absent from the result -- i.e. ``len(order) < total`` exactly when
+        the graph contains a cycle.
+        """
+        nodes, indegree = self._build_graph()
+        queue = collections.deque(n for n in nodes if indegree[id(n)] == 0)
+        order: List["Ruleset._Node"] = []
+        while queue:
+            node = queue.popleft()
+            order.append(node)
+            for dep in node.dependents:
+                indegree[id(dep)] -= 1
+                if indegree[id(dep)] == 0:
+                    queue.append(dep)
+        return order, len(nodes)
+
+    def _build_graph(
+        self,
+    ) -> Tuple[List["Ruleset._Node"], Dict[int, int]]:
+        """Build the dependency DAG.
+
+        Returns the nodes (one per rule, in insertion order) and their
+        in-degrees -- the number of rules that must be applied before each.
+        The in-degree map is keyed by node identity (``id``) rather than rule
+        type so that distinct nodes never share a counter, and is computed as
+        the edges are added (every incoming edge bumps the target's in-degree)
+        rather than re-derived by a second traversal. A node whose in-degree
+        is zero is a root.
+        """
+        rule_to_node: Dict[Type[Rule], "Ruleset._Node"] = {
+            rule: Ruleset._Node(rule) for rule in self._rules
+        }
+        indegree: Dict[int, int] = {id(node): 0 for node in rule_to_node.values()}
+
+        # De-duplicate edges. The same ordering can be declared from both ends
+        # -- rule A lists B in ``dependencies()`` while B lists A in
+        # ``dependents()`` -- which would otherwise add the edge twice,
+        # double-counting the in-degree and duplicating ``dependents`` entries.
+        seen_edges: Set[Tuple[int, int]] = set()
+
+        def add_edge(before: "Ruleset._Node", after: "Ruleset._Node") -> None:
+            """Record that ``before`` must be applied before ``after``."""
+            edge = (id(before), id(after))
+            if edge in seen_edges:
+                return
+            seen_edges.add(edge)
+            before.dependents.append(after)
+            indegree[id(after)] += 1
+
+        for rule in self._rules:
+            node = rule_to_node[rule]
+
+            # Rules that must be applied *before* this rule: dependency -> node.
+            for dependency in rule.dependencies():
+                if dependency in rule_to_node:
+                    add_edge(rule_to_node[dependency], node)
+
+            # Rules that must be applied *after* this rule: node -> dependent.
+            for dependent in rule.dependents():
+                if dependent in rule_to_node:
+                    add_edge(node, rule_to_node[dependent])
+
+        return list(rule_to_node.values()), indegree
+
+    def _contains_cycle(self) -> bool:
+        # Kahn's traversal drops any node stuck in a cycle (its in-degree never
+        # reaches zero), so a shortfall between the ordered nodes and the total
+        # means a cycle exists. This correctly flags a graph that mixes an
+        # acyclic component with a disjoint cycle -- a plain "does any root
+        # exist?" check would be fooled by the acyclic root.
+        order, total = self._topological_order()
+        return len(order) != total