chore: import upstream snapshot with attribution

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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