chore: import upstream snapshot with attribution

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}")