chore: import upstream snapshot with attribution

python/ray/data/_internal/planner/__init__.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ray.data._internal.planner.planner import Planner
+
+
+def create_planner() -> "Planner":
+    # Import here to avoid circular import.
+    from ray.data._internal.planner.planner import Planner
+
+    return Planner()
+
+
+__all__ = ["create_planner"]

python/ray/data/_internal/planner/aggregate.py

@@ -0,0 +1,111 @@
+from typing import List, Optional, Union
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.execution.util import merge_label_selector
+from ray.data._internal.planner.exchange.aggregate_task_spec import (
+    SortAggregateTaskSpec,
+)
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey, SortTaskSpec
+from ray.data._internal.util import unify_ref_bundles_schema
+from ray.data.aggregate import AggregateFn
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_aggregate_fn(
+    key: Optional[Union[str, List[str]]],
+    aggs: List[AggregateFn],
+    data_context: DataContext,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to aggregate blocks by the specified key column or key
+    function.
+    """
+    assert data_context.shuffle_strategy in [
+        ShuffleStrategy.SORT_SHUFFLE_PULL_BASED,
+        ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED,
+    ]
+
+    if len(aggs) == 0:
+        raise ValueError("Aggregate requires at least one aggregation")
+
+    def fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        blocks = []
+        metadata = []
+        for ref_bundle in refs:
+            blocks.extend(ref_bundle.block_refs)
+            metadata.extend(ref_bundle.metadata)
+        if len(blocks) == 0:
+            return (blocks, {})
+
+        unified_schema = unify_ref_bundles_schema(refs)
+        for agg_fn in aggs:
+            agg_fn._validate(unified_schema)
+
+        num_mappers = len(blocks)
+
+        sort_key = SortKey(key)
+
+        label_selector = data_context.execution_options.label_selector
+        if key is None:
+            num_outputs = 1
+            boundaries = []
+        else:
+            # Use same number of output partitions.
+            num_outputs = num_mappers
+            sample_bar = ctx.sub_progress_bar_dict[
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME
+            ]
+            # Sample boundaries for aggregate key.
+            boundaries = SortTaskSpec.sample_boundaries(
+                blocks,
+                sort_key,
+                num_outputs,
+                sample_bar,
+                label_selector=label_selector,
+            )
+
+        agg_spec = SortAggregateTaskSpec(
+            boundaries=boundaries,
+            key=sort_key,
+            aggs=aggs,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            scheduler = PushBasedShuffleTaskScheduler(agg_spec)
+        elif data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PULL_BASED:
+            scheduler = PullBasedShuffleTaskScheduler(agg_spec)
+        else:
+            raise ValueError(
+                f"Invalid shuffle strategy '{data_context.shuffle_strategy}'"
+            )
+
+        map_ray_remote_args = merge_label_selector({}, label_selector)
+        reduce_ray_remote_args = merge_label_selector({}, label_selector)
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            map_ray_remote_args=map_ray_remote_args,
+            reduce_ray_remote_args=reduce_ray_remote_args,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    return fn

python/ray/data/_internal/planner/checkpoint/__init__.py

@@ -0,0 +1,13 @@
+from .plan_read_files_op import plan_read_files_op_with_checkpoint_filter
+from .plan_read_op import (
+    create_checkpoint_filter_op,
+    plan_read_op_with_checkpoint_filter,
+)
+from .plan_write_op import plan_write_op_with_checkpoint_writer
+
+__all__ = [
+    "plan_read_files_op_with_checkpoint_filter",
+    "create_checkpoint_filter_op",
+    "plan_read_op_with_checkpoint_filter",
+    "plan_write_op_with_checkpoint_writer",
+]

python/ray/data/_internal/planner/checkpoint/plan_read_files_op.py

@@ -0,0 +1,47 @@
+"""Checkpoint-aware planner for the V2 ``ReadFiles`` logical operator.
+
+Mirrors :func:`plan_read_op_with_checkpoint_filter` (the V1 ``Read``
+variant) so V2 uses the same wrapping ActorPool ``CheckpointFilter``
+``MapOperator`` downstream of the read: same
+``_CheckpointFilterFn`` / ``_get_checkpoint_map_transformer``, same
+memory reservation formula, and ``supports_fusion=False`` so the
+filter stays a distinct op.
+
+Registered via ``Planner._get_plan_fns_for_checkpointing`` so it only
+runs when ``DataContext.checkpoint_config`` is set *and* the logical
+plan is a ``Write`` or ``StreamingSplit`` with a ``ReadFiles`` at the
+leaf. V2's plain ``plan_read_files_op`` stays checkpoint-unaware; this
+file is the only place V2 reads pick up a checkpoint filter.
+"""
+from typing import List, Optional
+
+import pyarrow
+import pyarrow.fs
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.logical.operators import ReadFiles
+from ray.data._internal.planner.checkpoint.plan_read_op import (
+    create_checkpoint_filter_op,
+)
+from ray.data._internal.planner.plan_read_files_op import plan_read_files_op
+from ray.data.context import DataContext
+
+
+def plan_read_files_op_with_checkpoint_filter(
+    data_file_dir: Optional[str],
+    data_file_filesystem: Optional["pyarrow.fs.FileSystem"],
+    op: ReadFiles,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Wrap a V2 ``ReadFiles`` physical op with a ``CheckpointFilter``.
+
+    Defers all wrapping behavior to V1's :func:`create_checkpoint_filter_op`
+    so the not-found short-circuit, ``IdColumnCheckpointManager.load_checkpoint``
+    invocation, actor-pool sizing, and ``supports_fusion=False`` placement stay
+    in one place across the V1 and V2 read paths.
+    """
+    physical_read_op = plan_read_files_op(op, physical_children, data_context)
+    return create_checkpoint_filter_op(
+        physical_read_op, data_context, data_file_dir, data_file_filesystem
+    )

python/ray/data/_internal/planner/checkpoint/plan_read_op.py

@@ -0,0 +1,156 @@
+from typing import Iterable, List, Optional
+
+import numpy
+import pyarrow
+import pyarrow.fs as fs
+
+from ray.data._internal.compute import ActorPoolStrategy
+from ray.data._internal.execution.interfaces import PhysicalOperator, TaskContext
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.logical.operators import Read
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.planner.plan_read_op import plan_read_op
+from ray.data.block import Block
+from ray.data.checkpoint.checkpoint_filter import (
+    IdColumnCheckpointManager,
+    NumpyArrayBasedCheckpointFilter,
+)
+from ray.data.context import DataContext
+from ray.data.datasource.path_util import _unwrap_protocol
+from ray.types import ObjectRef
+
+CHECKPOINT_MEMORY_SAFETY_FACTOR = 1.5
+
+
+def create_checkpoint_filter_op(
+    physical_input_op: PhysicalOperator,
+    data_context: DataContext,
+    data_file_dir: Optional[str],
+    data_file_filesystem: Optional["pyarrow.fs.FileSystem"],
+) -> PhysicalOperator:
+    """Wrap ``physical_input_op`` with an actor-pool checkpoint filter operator.
+
+    Args:
+        physical_input_op: The upstream physical operator whose output should
+            be filtered.
+        data_context: The data context carrying the checkpoint config.
+        data_file_dir: Directory where data files are written. Used to clean
+            up orphaned data files from pending (incomplete) checkpoints.
+        data_file_filesystem: Filesystem for data files. Defaults to the
+            checkpoint filesystem if not provided.
+
+    Returns:
+        A ``CheckpointFilter`` ``MapOperator`` downstream of
+        ``physical_input_op``, or ``physical_input_op`` itself if there is no
+        checkpoint data to restore from.
+    """
+    checkpoint_config = data_context.checkpoint_config
+
+    # Return the input op directly if:
+    # 1. the checkpoint directory does not exist.
+    # 2. no valid files under checkpoint_path (for example, it is an empty directory).
+    info = checkpoint_config.filesystem.get_file_info(
+        _unwrap_protocol(checkpoint_config.checkpoint_path)
+    )
+    if info.type == fs.FileType.NotFound:
+        return physical_input_op
+
+    checkpoint_manager = IdColumnCheckpointManager(
+        checkpoint_config=checkpoint_config,
+        data_context=data_context,
+    )
+    # load checkpointed IDs as a numpy ndarray and store it to object store.
+    checkpointed_ids_ref, checkpointed_ids_size = checkpoint_manager.load_checkpoint(
+        data_file_dir, data_file_filesystem
+    )
+    if not checkpointed_ids_ref:
+        return physical_input_op
+
+    map_transformer = _get_checkpoint_map_transformer(
+        data_context, checkpointed_ids_ref
+    )
+
+    checkpoint_op = MapOperator.create(
+        map_transformer=map_transformer,
+        input_op=physical_input_op,
+        data_context=data_context,
+        name="CheckpointFilter",
+        compute_strategy=ActorPoolStrategy(
+            min_size=checkpoint_config.checkpoint_actor_pool_min_size,
+            max_size=checkpoint_config.checkpoint_actor_pool_max_size,
+        ),
+        ray_remote_args={
+            "memory": max(
+                checkpoint_config.checkpoint_actor_memory_bytes,
+                int(checkpointed_ids_size * CHECKPOINT_MEMORY_SAFETY_FACTOR),
+            )
+        },
+        supports_fusion=False,
+    )
+    return checkpoint_op
+
+
+def plan_read_op_with_checkpoint_filter(
+    data_file_dir: Optional[str],
+    data_file_filesystem: Optional["pyarrow.fs.FileSystem"],
+    op: Read,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Plan the read op to physical operators.
+    1. If checkpoint is not enabled, or the checkpoint_path is an empty directory,
+    return the original read physical operator.
+    2. If the checkpoint is valid, translate the logical read operator into two
+    physical operators read->map, where the map operator receives blocks from the
+    read operator and outputs the filtered Blocks.
+    The implementation of the map operator is `ActorPoolMapOperator`. At runtime
+    the number of checkpoint-actors is dynamically scaled. The number of actors
+    is in the range [checkpoint_actor_pool_min_size, checkpoint_actor_pool_max_size].
+    """
+    physical_read_op = plan_read_op(op, physical_children, data_context)
+    return create_checkpoint_filter_op(
+        physical_read_op, data_context, data_file_dir, data_file_filesystem
+    )
+
+
+class _CheckpointFilterFn:
+    def __init__(
+        self,
+        checkpoint_config,
+        checkpointed_ids_ref: ObjectRef[numpy.ndarray],
+    ):
+        self._config = checkpoint_config
+        self._ref = checkpointed_ids_ref
+        self._filter = None
+
+    def init_checkpoint_filter(self):
+        """Called once per actor worker to materialize the filter."""
+        self._filter = NumpyArrayBasedCheckpointFilter(self._config, self._ref)
+
+    def __call__(self, blocks: Iterable[Block], ctx: TaskContext) -> Iterable[Block]:
+        assert self._filter is not None, "checkpoint filter was not initialized!"
+        for block in blocks:
+            filtered_block = self._filter.filter_rows_for_block(block)
+            if filtered_block.num_rows > 0:
+                yield filtered_block
+
+
+def _get_checkpoint_map_transformer(
+    data_context: DataContext, checkpointed_ids_ref: ObjectRef[numpy.ndarray]
+) -> MapTransformer:
+    fn = _CheckpointFilterFn(data_context.checkpoint_config, checkpointed_ids_ref)
+
+    transformer_fn = BlockMapTransformFn(
+        block_fn=fn,
+        output_block_size_option=OutputBlockSizeOption.of(
+            target_max_block_size=data_context.target_max_block_size,
+        ),
+    )
+    return MapTransformer(
+        transform_fns=[transformer_fn],
+        init_fn=fn.init_checkpoint_filter,
+    )

python/ray/data/_internal/planner/checkpoint/plan_write_op.py

@@ -0,0 +1,350 @@
+import warnings
+from typing import Iterable, List, Tuple
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+)
+from ray.data._internal.logical.operators import Write
+from ray.data._internal.planner.plan_write_op import (
+    PENDING_CHECKPOINTS_KWARG_NAME,
+    WRITE_UUID_KWARG_NAME,
+    _plan_write_op_internal,
+    generate_collect_write_stats_fn,
+)
+from ray.data.block import Block, BlockAccessor
+from ray.data.checkpoint.checkpoint_writer import (
+    CheckpointWriter,
+    PendingCheckpoint,
+)
+from ray.data.checkpoint.interfaces import (
+    InvalidCheckpointingOperators,
+)
+from ray.data.context import DataContext
+from ray.data.datasource.datasink import Datasink
+from ray.data.datasource.file_datasink import _FileDatasink
+from ray.data.datasource.filename_provider import _split_base_and_ext
+
+
+def _validate_id_column_exists(id_column: str, block: Block) -> None:
+    """Validate that the ID column exists in the block.
+
+    Args:
+        id_column: The name of the ID column to validate.
+        block: The block to check.
+
+    Raises:
+        ValueError: If the ID column is not present in the block.
+    """
+    block_accessor = BlockAccessor.for_block(block)
+    if id_column not in block_accessor.column_names():
+        raise ValueError(
+            f"ID column {id_column} is "
+            f"absent in the block to be written. Do not drop or rename "
+            f"this column."
+        )
+
+
+def _combine_blocks(
+    blocks: Iterable[Block],
+) -> Tuple[List[Block], Block]:
+    """Combine multiple blocks into a single block.
+
+    This is used by checkpoint transforms to match the behavior of _FileDatasink.write(),
+    which combines all input blocks into one output file.
+
+    Args:
+        blocks: Iterable of blocks to combine.
+
+    Returns:
+        A tuple of (block_list, combined_block) where:
+        - block_list: The original blocks as a list (for later iteration)
+        - combined_block: A single block combining all input blocks
+    """
+    block_list = list(blocks)
+    builder = DelegatingBlockBuilder()
+    for block in block_list:
+        builder.add_block(block)
+    combined_block = builder.build()
+    return block_list, combined_block
+
+
+def plan_write_op_with_checkpoint_writer(
+    op: Write, physical_children: List[PhysicalOperator], data_context: DataContext
+) -> PhysicalOperator:
+    """Plan a write operation with checkpoint support.
+
+    For file-based datasinks (_FileDatasink):
+        Uses 2-phase commit for atomicity:
+        1. Pre-write: computes expected paths, write pending checkpoints
+        2. Write: writes data files
+        3. Post-write: commits checkpoints (renames pending -> committed)
+
+    Writing the pending checkpoint BEFORE the data file is critical: the
+    pending checkpoint is the source of truth for recovery. If failure occurs
+    after data write but before commit, recovery finds the pending checkpoint,
+    deletes the matching data files, and retries cleanly. Writing checkpoints
+    after data files would be non-atomic — if failure occurs between data
+    write and checkpoint write, there's no record of which data files are
+    uncommitted.
+
+    For non-file datasinks (SQLDatasink, etc.):
+        Falls back to post-write checkpointing:
+        1. Write: Write data to destination
+        2. Post-write: Write checkpoints
+
+    Non-file sinks (SQL, MongoDB, etc.) cannot predict a "file path" - data goes
+    to database rows or documents. So we fall back to writing data first, then
+    checkpointing. If failure occurs after data write but before checkpoint
+    write, the same data may be written again on retry without removing the
+    old data (at-least-once semantics for non-idempotent operations).
+    """
+    assert data_context.checkpoint_config is not None
+
+    datasink = op.datasink_or_legacy_datasource
+    if not isinstance(datasink, Datasink):
+        raise InvalidCheckpointingOperators(
+            f"To enable checkpointing, Write operation must use a "
+            f"Datasink and not a legacy Datasource, but got: "
+            f"{type(datasink)}"
+        )
+
+    checkpoint_writer = CheckpointWriter.create(data_context.checkpoint_config)
+    collect_stats_fn = generate_collect_write_stats_fn()
+
+    if isinstance(datasink, _FileDatasink):
+        # File-based datasink: use 2-phase commit for atomicity
+        # Pre-write transform: compute expected paths and write pending checkpoints
+        prepare_checkpoint_fn = _generate_prepare_checkpoint_transform(
+            data_context, datasink, checkpoint_writer
+        )
+
+        # Post-write transform: commit checkpoints
+        commit_checkpoint_fn = _generate_commit_checkpoint_transform(checkpoint_writer)
+
+        pre_transformations = [
+            prepare_checkpoint_fn,
+        ]
+        post_transformations = [
+            commit_checkpoint_fn,
+            collect_stats_fn,
+        ]
+    else:
+        # Non-file datasink (SQL, Mongo, etc.): fall back to non-atomic checkpoint
+        # No 2-phase commit - write checkpoint after data write
+        # This might cause duplicate writes if the write operation is retried.
+        warnings.warn(
+            f"Checkpointing with non-file datasink ({type(datasink).__name__}) "
+            f"uses post-write checkpointing, which provides at-least-once "
+            f"semantics. If a failure occurs after data is written but before "
+            f"the checkpoint is saved, duplicate data may be written on retry. "
+            f"This will be addressed in a future version."
+        )
+        write_checkpoint_fn = _generate_non_atomic_write_checkpoint_transform(
+            data_context, checkpoint_writer
+        )
+        post_transformations = [
+            write_checkpoint_fn,
+            collect_stats_fn,
+        ]
+        pre_transformations = []
+
+    physical_op = _plan_write_op_internal(
+        op,
+        physical_children,
+        data_context,
+        post_transformations=post_transformations,
+        pre_transformations=pre_transformations,
+    )
+
+    return physical_op
+
+
+def _generate_base_filename(
+    datasink: _FileDatasink,
+    ctx: TaskContext,
+) -> str:
+    """Compute the base filename (without extension) for this task's data files.
+
+    This is called BEFORE writing to determine the filename prefix for data files
+    that will be written by this task. Datasinks may write multiple files (with
+    partitioning, max_rows_per_file, etc.), all sharing this base filename.
+
+    Args:
+        datasink: The file datasink being used.
+        ctx: The task context.
+
+    Returns:
+        The base filename without extension (e.g., "write_uuid_000000_000000").
+        Used both as a checkpoint ID for deterministic naming and as a prefix
+        for matching data files during recovery.
+    """
+    write_uuid = ctx.kwargs.get(WRITE_UUID_KWARG_NAME)
+    assert write_uuid is not None, "WRITE_UUID_KWARG_NAME is required"
+
+    filename = datasink.filename_provider.get_filename_for_task(
+        write_uuid, ctx.task_idx
+    )
+
+    # All file datasinks can potentially generate multiple files (e.g., with
+    # partitioning, max_rows_per_file, etc.). Use prefix matching to handle
+    # cases like "{filename}-{i}.parquet".
+    base, _ = _split_base_and_ext(filename)
+    return base
+
+
+def _generate_prepare_checkpoint_transform(
+    data_context: DataContext,
+    datasink: _FileDatasink,
+    checkpoint_writer: CheckpointWriter,
+) -> BlockMapTransformFn:
+    """Generate transform for preparing checkpoints BEFORE data write.
+
+    This transform runs BEFORE the data write to enable rollback on failure.
+    By recording the expected file path in a pending checkpoint first, we can
+    clean up orphaned data files if the task fails after writing data but
+    before committing.
+
+    Steps:
+    1. Combines all blocks (matching _FileDatasink behavior)
+    2. Computes expected data file path prefix from FilenameProvider
+    3. Writes pending checkpoint with expected path prefix as filename
+    4. Stores pending checkpoint info in ctx.kwargs for later commit
+    """
+
+    def prepare_checkpoint(
+        blocks: Iterable[Block], ctx: TaskContext
+    ) -> Iterable[Block]:
+        # Combine all blocks to match _FileDatasink.write() behavior
+        # which combines all input blocks into one output file
+        block_list, combined_block = _combine_blocks(blocks)
+        ba = BlockAccessor.for_block(combined_block)
+
+        if ba.num_rows() > 0:
+            # Validate ID column exists
+            id_column = data_context.checkpoint_config.id_column
+            _validate_id_column_exists(id_column, combined_block)
+
+            # Compute base filename using FilenameProvider
+            # Note: This only depends on write_uuid and task_idx, NOT block content
+            # base_filename is the filename without extension, used as checkpoint_id
+            # for deterministic naming (same on retry, enabling idempotent writes)
+            base_filename = _generate_base_filename(datasink, ctx)
+
+            # Extract ID column data for checkpoint
+            # Project to the single column first, then convert to Arrow to
+            # avoid materializing the entire block as an Arrow table.
+            id_column_data = BlockAccessor.for_block(
+                ba.select(columns=[id_column])
+            ).to_arrow()[id_column]
+
+            # Write pending checkpoint with the base filename as checkpoint_id.
+            # The checkpoint filename will be {base_filename}.pending.parquet.
+            # During recovery, the pending checkpoint basename (without
+            # .pending.parquet) is used as a prefix to match data files.
+            pending = checkpoint_writer.write_pending_checkpoint(
+                id_column_data,
+                checkpoint_id=base_filename,
+            )
+
+            # Store pending checkpoint for commit phase
+            if pending is not None:
+                if PENDING_CHECKPOINTS_KWARG_NAME not in ctx.kwargs:
+                    ctx.kwargs[PENDING_CHECKPOINTS_KWARG_NAME] = []
+                ctx.kwargs[PENDING_CHECKPOINTS_KWARG_NAME].append(pending)
+
+        # Return original blocks for the write transform
+        return iter(block_list)
+
+    return BlockMapTransformFn(
+        prepare_checkpoint,
+        is_udf=False,
+        disable_block_shaping=True,
+    )
+
+
+def _generate_commit_checkpoint_transform(
+    checkpoint_writer: CheckpointWriter,
+) -> BlockMapTransformFn:
+    """Generate transform for committing checkpoints AFTER data write.
+
+    This transform runs AFTER the data write succeeds, completing the 2-phase
+    commit. The commit operation (renaming pending -> committed) is the atomic
+    point: once committed, the data is considered durably written. If failure
+    occurs before this point, recovery will find the pending checkpoint and
+    can safely delete the orphaned data files using the stored path.
+
+    Steps:
+    1. Retrieves pending checkpoints from ctx.kwargs
+    2. Commits each pending checkpoint (rename pending -> committed)
+    """
+
+    def commit_checkpoints(
+        blocks: Iterable[Block], ctx: TaskContext
+    ) -> Iterable[Block]:
+        # Get pending checkpoints written in pre-write phase
+        pending_checkpoints: List[PendingCheckpoint] = ctx.kwargs.get(
+            PENDING_CHECKPOINTS_KWARG_NAME, []
+        )
+
+        # Commit each pending checkpoint
+        for pending in pending_checkpoints:
+            checkpoint_writer.commit_checkpoint(pending)
+
+        return blocks
+
+    return BlockMapTransformFn(
+        commit_checkpoints,
+        is_udf=False,
+        disable_block_shaping=True,
+    )
+
+
+def _generate_non_atomic_write_checkpoint_transform(
+    data_context: DataContext,
+    checkpoint_writer: CheckpointWriter,
+) -> BlockMapTransformFn:
+    """Generate transform for writing checkpoints AFTER data write (non-file datasinks).
+
+    This is a fallback for non-file datasinks (SQL, Mongo, etc.) that don't
+    support deletions. Unlike file-based sinks where we can delete orphaned
+    data files during recovery, these sinks have no way to undo a write once
+    data has been inserted into rows or documents.
+
+    The checkpoint is written directly after the data write completes. This
+    provides at-least-once semantics: if failure occurs after data write but
+    before checkpoint write, the same data will be written again on retry
+    without removing the old data.
+
+    For idempotent operations (upserts with unique keys), this is safe. For
+    non-idempotent operations (inserts), duplicates may result.
+
+    TODO: For datasinks that support deletions (e.g., SQL DELETE by ID), we
+    could store written IDs in pending checkpoints and delete them on recovery,
+    avoiding duplicates even for non-idempotent operations.
+    """
+
+    def write_checkpoint(blocks: Iterable[Block], ctx: TaskContext) -> Iterable[Block]:
+        # Combine all blocks
+        block_list, combined_block = _combine_blocks(blocks)
+        ba = BlockAccessor.for_block(combined_block)
+
+        if ba.num_rows() > 0:
+            # Validate ID column exists
+            id_column = data_context.checkpoint_config.id_column
+            _validate_id_column_exists(id_column, combined_block)
+
+            # Write checkpoint directly (no 2-phase commit)
+            # No data_file_path since non-file datasinks don't have file paths
+            checkpoint_writer.write_block_checkpoint(ba)
+
+        return iter(block_list)
+
+    return BlockMapTransformFn(
+        write_checkpoint,
+        is_udf=False,
+        # NOTE: No need for block-shaping
+        disable_block_shaping=True,
+    )

python/ray/data/_internal/planner/download_partition_actor.py

@@ -0,0 +1,312 @@
+import asyncio
+import logging
+import math
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Iterator, List, Optional, cast
+
+import pyarrow as pa
+import pyarrow.fs as pafs
+from typing_extensions import override
+
+from ray.data._internal.planner._obstore_download import (
+    _FILE_SIZE_COLUMN_PREFIX,
+    RAY_DATA_OBSTORE_RANGE_THRESHOLD,
+    StoreRegistry,
+    _extract_credentials_from_filesystem,
+    _is_obstore_supported_url,
+    _split_obstore_uri,
+)
+from ray.data._internal.util import RetryingPyFileSystem, _arrow_batcher
+from ray.data.block import BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.path_util import _resolve_paths_and_filesystem
+
+logger = logging.getLogger(__name__)
+
+URI_DOWNLOAD_MAX_WORKERS = 16
+URI_HEAD_MAX_CONCURRENCY = 128
+
+
+class PartitionActor:
+    """Actor that partitions download operations based on estimated file sizes.
+
+    For multiple URI columns, estimates the combined size across all columns.
+    Uses threaded HEAD/metadata sampling only (no obstore-specific columns).
+    """
+
+    INIT_SAMPLE_BATCH_SIZE = 25
+
+    def __init__(
+        self,
+        uri_column_names: List[str],
+        data_context: DataContext,
+        filesystem: Optional[pafs.FileSystem] = None,
+    ):
+        self._uri_column_names = uri_column_names
+        self._data_context = data_context
+        self._filesystem = filesystem
+        self._batch_size_estimate = None
+
+    def __call__(self, block: pa.Table) -> Iterator[pa.Table]:
+        block = self._ensure_arrow_table(block)
+        self._validate_uri_columns(block)
+        yield from self._partition_and_yield(block)
+
+    def _ensure_arrow_table(self, block: pa.Table) -> pa.Table:
+        if not isinstance(block, pa.Table):
+            return BlockAccessor.for_block(block).to_arrow()
+        return block
+
+    def _validate_uri_columns(self, block: pa.Table) -> None:
+        for uri_column_name in self._uri_column_names:
+            if uri_column_name not in block.column_names:
+                raise ValueError(
+                    "Ray Data tried to download URIs from a column named "
+                    f"{uri_column_name!r}, but a column with that name doesn't "
+                    "exist. Is the specified download column correct?"
+                )
+
+    def _partition_and_yield(self, block: pa.Table) -> Iterator[pa.Table]:
+        if self._batch_size_estimate is None and block.num_rows > 0:
+            self._batch_size_estimate = self._estimate_nrows_per_partition(block)
+        yield from _arrow_batcher(block, self._batch_size_estimate or 1)
+
+    def _sampled_file_sizes_for_partition_estimate(
+        self, block: pa.Table, uri_column_name: str
+    ) -> List[Optional[int]]:
+        uris = block.column(uri_column_name).to_pylist()
+        sample_uris = uris[: self.INIT_SAMPLE_BATCH_SIZE]
+        # ``_sample_sizes`` returns concrete ``int``s; widen for this API.
+        return cast(List[Optional[int]], self._sample_sizes(sample_uris))
+
+    def _estimate_nrows_per_partition(self, block: pa.Table) -> int:
+        sampled_file_sizes_by_column = {}
+        for uri_column_name in self._uri_column_names:
+            sampled_file_sizes = self._sampled_file_sizes_for_partition_estimate(
+                block, uri_column_name
+            )
+            sampled_file_sizes_by_column[uri_column_name] = sampled_file_sizes
+
+        sampled_file_sizes_by_column = {
+            uri_column_name: [
+                file_size if file_size is not None else 0
+                for file_size in sampled_file_sizes
+            ]
+            for uri_column_name, sampled_file_sizes in sampled_file_sizes_by_column.items()
+        }
+
+        # This is some fancy Python code to compute the file size of each row.
+        row_sizes = [
+            sum(file_sizes_in_row)
+            for file_sizes_in_row in zip(*sampled_file_sizes_by_column.values())
+        ]
+
+        target_nbytes_per_partition = self._data_context.target_max_block_size
+        avg_nbytes_per_row = sum(row_sizes) / len(row_sizes)
+        if avg_nbytes_per_row == 0:
+            logger.warning(
+                "Estimated average row size is 0. Falling back to using the number of "
+                "rows in the block as the partition size."
+            )
+            return len(block)
+
+        if target_nbytes_per_partition is None:
+            # Target max block size is None--keep the whole block as one partition.
+            return len(block)
+
+        nrows_per_partition = math.floor(
+            target_nbytes_per_partition / avg_nbytes_per_row
+        )
+        if nrows_per_partition == 0:
+            # A single file exceeds target_max_block_size. Fall back to one row
+            # per partition so _arrow_batcher doesn't crash on a zero step size.
+            logger.warning(
+                f"Estimated average file size ({avg_nbytes_per_row:.0f} bytes) "
+                f"exceeds target_max_block_size ({target_nbytes_per_partition} bytes). "
+                "Falling back to one row per partition; output blocks may be larger "
+                "than the configured target."
+            )
+            return 1
+        return nrows_per_partition
+
+    def _sample_sizes(self, uris: List[str]) -> List[int]:
+        """Fetch file sizes in parallel using ThreadPoolExecutor."""
+
+        def get_file_size(uri_path, fs):
+            try:
+                return fs.get_file_info(uri_path).size
+            except Exception:
+                return None
+
+        if not uris:
+            return []
+
+        # Get the filesystem from the URIs (assumes all URIs use same filesystem for sampling)
+        # This is for sampling the file sizes which doesn't require a full resolution of the paths.
+        try:
+            paths, fs = _resolve_paths_and_filesystem(uris, filesystem=self._filesystem)
+            fs = RetryingPyFileSystem.wrap(
+                fs, retryable_errors=self._data_context.retried_io_errors
+            )
+        except Exception as e:
+            logger.warning(f"Failed to resolve URIs for size sampling: {e}")
+            return [0] * len(uris)
+
+        # _resolve_paths_and_filesystem silently drops URIs that fail.
+        # Fall back to zeros (triggers HEAD in the download path) rather than risk
+        # a length mismatch with the input block.
+        if len(paths) != len(uris):
+            logger.debug(
+                "Path resolution dropped %d of %d URIs; returning size=0 "
+                "for all so the download path can issue HEAD requests.",
+                len(uris) - len(paths),
+                len(uris),
+            )
+            return [0] * len(uris)
+
+        # Use ThreadPoolExecutor for concurrent size fetching
+        file_sizes: List[Optional[int]] = [None] * len(paths)
+        with ThreadPoolExecutor(max_workers=URI_DOWNLOAD_MAX_WORKERS) as executor:
+            # Submit all size fetch tasks
+            future_to_file_index = {
+                executor.submit(get_file_size, uri_path, fs): file_index
+                for file_index, uri_path in enumerate(paths)
+            }
+
+            # Collect results as they complete (order doesn't matter)
+            for future in as_completed(future_to_file_index):
+                file_index = future_to_file_index[future]
+                try:
+                    size = future.result()
+                    file_sizes[file_index] = size if size is not None else 0
+                except Exception as e:
+                    logger.warning(f"Error fetching file size for download: {e}")
+                    file_sizes[file_index] = 0
+
+        assert all(
+            size is not None for size in file_sizes
+        ), "File size sampling did not complete for all paths"
+        return [size for size in file_sizes if size is not None]
+
+
+class AsyncPartitionActor(PartitionActor):
+    """Partition actor for the obstore download path.
+
+    Uses obstore's async HEAD API for all size-fetching operations, replacing
+    the base class's 16-thread PyArrow ThreadPoolExecutor with fully async
+    requests at up to URI_HEAD_MAX_CONCURRENCY (128) concurrency.
+
+    When range splitting is enabled, attaches per-URI size columns so
+    downstream obstore downloads can skip redundant HEAD requests.
+    """
+
+    def __init__(
+        self,
+        uri_column_names: List[str],
+        data_context: DataContext,
+        filesystem: Optional[pafs.FileSystem] = None,
+    ):
+        super().__init__(uri_column_names, data_context, filesystem)
+        fs_kwargs = _extract_credentials_from_filesystem(filesystem)
+        if fs_kwargs is None:
+            # Fail closed. ``plan_download_op`` routes filesystems we can't
+            # extract credentials from to ``PartitionActor`` (PyArrow path),
+            # so reaching ``AsyncPartitionActor`` with an unextractable FS is
+            # a bug. Silently seeding an empty kwargs dict would hand the
+            # user's filesystem over to obstore's ambient credential chain
+            # (IMDS / env), which is exactly the silent-drop behavior the
+            # routing was designed to prevent.
+            raise RuntimeError(
+                "AsyncPartitionActor was constructed with a filesystem whose "
+                f"credentials cannot be statically extracted ({type(filesystem).__name__}). "
+                "This indicates a dispatch bug: use PartitionActor for such "
+                "filesystems so the user's credentials are not silently dropped."
+            )
+        self._registry = StoreRegistry(retry_config={"max_retries": 10}, **fs_kwargs)
+
+    def __call__(self, block: pa.Table) -> Iterator[pa.Table]:
+        block = self._ensure_arrow_table(block)
+        self._validate_uri_columns(block)
+
+        if block.num_rows > 0 and RAY_DATA_OBSTORE_RANGE_THRESHOLD > 0:
+            first_uri = block.column(self._uri_column_names[0])[0].as_py()
+            if _is_obstore_supported_url(first_uri):
+                block = self._attach_file_sizes(block)
+
+        yield from self._partition_and_yield(block)
+
+    def _sampled_file_sizes_for_partition_estimate(
+        self, block: pa.Table, uri_column_name: str
+    ) -> List[Optional[int]]:
+        size_col = f"{_FILE_SIZE_COLUMN_PREFIX}{uri_column_name}"
+        if size_col in block.column_names:
+            all_sizes = block.column(size_col).to_pylist()
+            return all_sizes[: self.INIT_SAMPLE_BATCH_SIZE]
+        return super()._sampled_file_sizes_for_partition_estimate(
+            block, uri_column_name
+        )
+
+    @override
+    def _sample_sizes(self, uris: List[str]) -> List[int]:
+        """Fetch file sizes concurrently using obstore's async HEAD API.
+
+        Overrides the base class to use obstore instead of PyArrow's threaded
+        get_file_info.  This affects all callers: both the initial partition-
+        size estimation (25-file sample) and _attach_file_sizes (all files).
+
+        For URI schemes not supported by obstore (e.g. hdfs://), falls back
+        to the base class's PyArrow-threaded implementation.
+        """
+        import obstore as obs
+
+        if not uris:
+            return []
+
+        if not _is_obstore_supported_url(uris[0]):
+            return super()._sample_sizes(uris)
+
+        sem = asyncio.Semaphore(URI_HEAD_MAX_CONCURRENCY)
+
+        async def _head_one(uri: str) -> int:
+            try:
+                store_url, path = _split_obstore_uri(uri)
+                store = self._registry.get(store_url)
+                async with sem:
+                    meta = await obs.head_async(store, path)
+                return meta["size"]
+            except Exception:
+                return 0
+
+        async def _head_all() -> List[int]:
+            sizes = await asyncio.gather(*[_head_one(u) for u in uris])
+            failed = [uri for uri, size in zip(uris, sizes) if size == 0]
+            if failed:
+                logger.debug(
+                    "obstore HEAD failed for %d URIs: %s",
+                    len(failed),
+                    failed,
+                )
+            return sizes
+
+        return asyncio.run(_head_all())
+
+    def _attach_file_sizes(self, block: pa.Table) -> pa.Table:
+        """Fetch file sizes for all URIs and attach as hidden columns.
+
+        Only called when obstore is available, range splitting is enabled
+        (RAY_DATA_OBSTORE_RANGE_THRESHOLD > 0), and the URI scheme is
+        supported by obstore.  The hidden columns are consumed by
+        ``_download_uris_with_obstore`` and dropped before output.
+
+        The hidden columns are consumed by ``_download_uris_with_obstore`` and
+        dropped before output. For cloud URIs this uses cheap metadata lookups.
+        For HTTP URIs where sizes are unavailable, stores 0 so the downstream
+        download path falls back to HEAD via obstore.
+        """
+        for uri_column_name in self._uri_column_names:
+            size_col = f"{_FILE_SIZE_COLUMN_PREFIX}{uri_column_name}"
+            uris = block.column(uri_column_name).to_pylist()
+            # Fetches all file sizes (not just a sample).
+            sizes = self._sample_sizes(uris)
+            block = block.append_column(size_col, pa.array(sizes, type=pa.int64()))
+        return block

python/ray/data/_internal/planner/exchange/aggregate_task_spec.py

@@ -0,0 +1,123 @@
+from typing import List, Tuple, Union
+
+from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+from ray.data._internal.table_block import TableBlockAccessor
+from ray.data.aggregate import AggregateFn, AggregateFnV2, Count
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadataWithSchema,
+    KeyType,
+)
+
+
+class SortAggregateTaskSpec(ExchangeTaskSpec):
+    """
+    The implementation for sort-based aggregate tasks.
+
+    Aggregate is done in 2 steps: partial aggregate of individual blocks, and
+    final aggregate of sorted blocks.
+
+    Partial aggregate (`map`): each block is sorted locally, then partitioned into
+    smaller blocks according to the boundaries. Each partitioned block is aggregated
+    separately, then passed to a final aggregate task.
+
+    Final aggregate (`reduce`): each task would receive a block from every worker that
+    consists of items in a certain range. It then merges the sorted blocks and
+    aggregates on-the-fly.
+    """
+
+    def __init__(
+        self,
+        boundaries: List[KeyType],
+        key: SortKey,
+        aggs: List[AggregateFn],
+    ):
+        super().__init__(
+            map_args=[boundaries, key, aggs],
+            reduce_args=[key, aggs],
+        )
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        boundaries: List[KeyType],
+        sort_key: SortKey,
+        aggs: List[AggregateFn],
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        stats = BlockExecStats.builder()
+
+        block = SortAggregateTaskSpec._prune_unused_columns(block, sort_key, aggs)
+        if sort_key.get_columns():
+            partitions = BlockAccessor.for_block(block).sort_and_partition(
+                boundaries,
+                sort_key,
+            )
+        else:
+            partitions = [block]
+        parts = [
+            BlockAccessor.for_block(p)._aggregate(sort_key, aggs) for p in partitions
+        ]
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta_with_schema = BlockMetadataWithSchema.from_block(
+            block, block_exec_stats=stats.build()
+        )
+        return parts + [meta_with_schema]
+
+    @staticmethod
+    def reduce(
+        key: SortKey,
+        aggs: List[AggregateFn],
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        normalized_blocks = TableBlockAccessor.normalize_block_types(
+            mapper_outputs,
+            target_block_type=None,
+        )
+        blocks, meta_with_schema = BlockAccessor.for_block(
+            normalized_blocks[0]
+        )._combine_aggregated_blocks(
+            list(normalized_blocks), key, aggs, finalize=not partial_reduce
+        )
+        return blocks, meta_with_schema
+
+    @staticmethod
+    def _prune_unused_columns(
+        block: Block,
+        sort_key: SortKey,
+        aggs: Tuple[AggregateFn],
+    ) -> Block:
+        """Prune unused columns from block before aggregate."""
+        prune_columns = True
+        columns = set()
+        key = sort_key.get_columns()
+
+        if isinstance(key, str):
+            columns.add(key)
+        elif isinstance(key, list):
+            columns.update(key)
+        elif callable(key):
+            prune_columns = False
+
+        for agg in aggs:
+            if isinstance(agg, AggregateFnV2) and agg.get_target_column():
+                columns.add(agg.get_target_column())
+            elif not isinstance(agg, Count):
+                # Don't prune columns if any aggregate key is not string.
+                prune_columns = False
+
+        block_accessor = BlockAccessor.for_block(block)
+        if (
+            prune_columns
+            and isinstance(block_accessor, TableBlockAccessor)
+            and block_accessor.num_rows() > 0
+        ):
+            return block_accessor.select(list(columns))
+        else:
+            return block

python/ray/data/_internal/planner/exchange/interfaces.py

@@ -0,0 +1,135 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import ray._private.worker
+from ray.data._internal.execution.interfaces import RefBundle
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import convert_bytes_to_human_readable_str
+from ray.data.block import Block
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+
+    from ray.data.block import BlockMetadataWithSchema
+
+logger = logging.getLogger(__name__)
+
+
+class ExchangeTaskSpec:
+    """
+    An interface to specify the exchange map and reduce tasks.
+
+    Subclasses should implement the `map` and `reduce` static methods.
+    `map` method is to transform one input block into multiple output blocks.
+    `reduce` is to combine multiple map output blocks. Both methods are
+    single-task operations. See `ExchangeScheduler` for how to distribute
+    the operations across multiple tasks.
+
+    Any custom arguments for `map` and `reduce` methods should be specified by
+    setting `map_args` and `reduce_args`.
+
+    The concept here is similar to the exchange operator described in
+    "Volcano - An Extensible and Parallel Query Evaluation System"
+    (https://dl.acm.org/doi/10.1109/69.273032).
+    """
+
+    MAP_SUB_PROGRESS_BAR_NAME = "Shuffle Map"
+    REDUCE_SUB_PROGRESS_BAR_NAME = "Shuffle Reduce"
+
+    def __init__(self, map_args: List[Any] = None, reduce_args: List[Any] = None):
+        self._map_args = map_args or []
+        self._reduce_args = reduce_args or []
+        assert isinstance(self._map_args, list)
+        assert isinstance(self._reduce_args, list)
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        """
+        Map function to be run on each input block.
+
+        Returns list of [BlockMetadata, Block1, Block2, ..., BlockN].
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def reduce(
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        """
+        Reduce function to be run for each output block.
+
+        Args:
+            *mapper_outputs: List of map output blocks to reduce.
+            partial_reduce: Whether should partially or fully reduce.
+
+        Returns:
+            The reduced block and its metadata.
+        """
+        raise NotImplementedError
+
+
+class ExchangeTaskScheduler:
+    """
+    An interface to schedule exchange tasks (`exchange_spec`) for multi-nodes
+    execution.
+    """
+
+    def __init__(self, exchange_spec: ExchangeTaskSpec):
+        """Initialize the scheduler.
+
+        Args:
+            exchange_spec: The implementation of exchange tasks to execute.
+        """
+        self._exchange_spec = exchange_spec
+        # If driver memory exceeds this threshold, warn the user. For now, this
+        # only applies to shuffle ops because most other ops are unlikely to use as
+        # much driver memory.
+        self.warn_on_driver_memory_usage_bytes: Optional[
+            int
+        ] = DataContext.get_current().warn_on_driver_memory_usage_bytes
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+        warn_on_driver_memory_usage: Optional[int] = None,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        """
+        Execute the exchange tasks on input `refs`.
+        """
+        raise NotImplementedError
+
+    def warn_on_high_local_memory_store_usage(self) -> None:
+        ray_core_worker = ray._private.worker.global_worker.core_worker
+        local_memory_store_bytes_used = (
+            ray_core_worker.get_local_memory_store_bytes_used()
+        )
+        self.warn_on_driver_memory_usage(
+            local_memory_store_bytes_used,
+            "More than "
+            f"{convert_bytes_to_human_readable_str(local_memory_store_bytes_used)} "
+            "of driver memory used to store Ray Data block data and metadata. "
+            "This job may exit if driver memory is insufficient.\n\n"
+            "This can happen when many tiny blocks are created. "
+            "Check the block size using Dataset.stats() and see "
+            "https://docs.ray.io/en/latest/data/performance-tips.html"
+            " for mitigation.",
+        )
+
+    def warn_on_driver_memory_usage(
+        self, memory_usage_bytes: int, log_str: str
+    ) -> None:
+        if self.warn_on_driver_memory_usage_bytes is None:
+            return
+
+        if memory_usage_bytes > self.warn_on_driver_memory_usage_bytes:
+            logger.warning(log_str)
+            # Double the threshold to avoid verbose warnings.
+            self.warn_on_driver_memory_usage_bytes = memory_usage_bytes * 2

python/ray/data/_internal/planner/exchange/pull_based_shuffle_task_scheduler.py

@@ -0,0 +1,155 @@
+import logging
+from typing import Any, Dict, List, Optional, Tuple
+
+from ray._private.ray_constants import CALLER_MEMORY_USAGE_PER_OBJECT_REF
+from ray.data._internal.execution.interfaces import BlockEntry, RefBundle, TaskContext
+from ray.data._internal.planner.exchange.interfaces import (
+    ExchangeTaskScheduler,
+    ExchangeTaskSpec,
+)
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import (
+    convert_bytes_to_human_readable_str,
+    unzip,
+)
+from ray.data.block import BlockMetadataWithSchema, to_stats
+
+logger = logging.getLogger(__name__)
+
+
+class PullBasedShuffleTaskScheduler(ExchangeTaskScheduler):
+    """
+    The pull-based map-reduce shuffle scheduler.
+
+    Map tasks are first scheduled to generate map output blocks. After all map output
+    are generated, then reduce tasks are scheduled to combine map output blocks
+    together.
+
+    The concept here is similar to
+    "MapReduce: Simplified Data Processing on Large Clusters"
+    (https://dl.acm.org/doi/10.1145/1327452.1327492).
+    """
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        task_ctx: TaskContext,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+        _debug_limit_execution_to_num_blocks: Optional[int] = None,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+
+        # TODO: eagerly delete the input and map output block references in order to
+        # eagerly release the blocks' memory.
+        input_blocks_list = []
+        for ref_bundle in refs:
+            input_blocks_list.extend(ref_bundle.block_refs)
+        input_num_blocks = len(input_blocks_list)
+        input_owned = all(b.owns_blocks for b in refs)
+
+        caller_memory_usage = (
+            input_num_blocks * output_num_blocks * CALLER_MEMORY_USAGE_PER_OBJECT_REF
+        )
+        self.warn_on_driver_memory_usage(
+            caller_memory_usage,
+            "Execution is estimated to use at least "
+            f"{convert_bytes_to_human_readable_str(caller_memory_usage)} "
+            "of driver memory. Ensure that the driver machine has at least "
+            "this much memory to ensure job completion.\n\n"
+            "To reduce the "
+            "amount of driver memory needed, enable push-based shuffle using "
+            "RAY_DATA_PUSH_BASED_SHUFFLE=1 "
+            "(https://docs.ray.io/en/latest/data/performance-tips.html"
+            ").",
+        )
+
+        if map_ray_remote_args is None:
+            map_ray_remote_args = {}
+        if reduce_ray_remote_args is None:
+            reduce_ray_remote_args = {}
+        if "scheduling_strategy" not in reduce_ray_remote_args:
+            reduce_ray_remote_args = reduce_ray_remote_args.copy()
+            reduce_ray_remote_args["scheduling_strategy"] = "SPREAD"
+
+        shuffle_map = cached_remote_fn(self._exchange_spec.map)
+        shuffle_reduce = cached_remote_fn(self._exchange_spec.reduce)
+
+        sub_progress_bar_dict = task_ctx.sub_progress_bar_dict
+        bar_name = ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        map_bar = sub_progress_bar_dict[bar_name]
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            input_blocks_list = input_blocks_list[:_debug_limit_execution_to_num_blocks]
+            logger.debug(f"Limiting execution to {len(input_blocks_list)} map tasks")
+
+        shuffle_map_out = [
+            shuffle_map.options(
+                **map_ray_remote_args,
+                num_returns=1 + output_num_blocks,
+            ).remote(i, block, output_num_blocks, *self._exchange_spec._map_args)
+            for i, block in enumerate(input_blocks_list)
+        ]
+
+        # The first item returned is the BlockMetadata.
+        shuffle_map_metadata_schema = []
+        for i, refs in enumerate(shuffle_map_out):
+            shuffle_map_metadata_schema.append(refs[-1])
+            shuffle_map_out[i] = refs[:-1]
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            while len(shuffle_map_out) < output_num_blocks:
+                # Repeat the first map task's results.
+                shuffle_map_out.append(shuffle_map_out[0][:])
+
+        shuffle_map_metadata_schema = map_bar.fetch_until_complete(
+            shuffle_map_metadata_schema
+        )
+
+        self.warn_on_high_local_memory_store_usage()
+
+        bar_name = ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        reduce_bar = sub_progress_bar_dict[bar_name]
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            output_num_blocks = _debug_limit_execution_to_num_blocks
+            logger.debug(f"Limiting execution to {output_num_blocks} reduce tasks")
+        shuffle_reduce_out = [
+            shuffle_reduce.options(**reduce_ray_remote_args, num_returns=2).remote(
+                *self._exchange_spec._reduce_args,
+                *[shuffle_map_out[i][j] for i in range(input_num_blocks)],
+            )
+            for j in range(output_num_blocks)
+        ]
+
+        # Release map task outputs from the Ray object store.
+        del shuffle_map_out
+
+        new_blocks, new_metadata_schema = [], []
+        if shuffle_reduce_out:
+            new_blocks, new_metadata_schema = unzip(shuffle_reduce_out)
+        new_metadata_schema: List[
+            "BlockMetadataWithSchema"
+        ] = reduce_bar.fetch_until_complete(list(new_metadata_schema))
+
+        self.warn_on_high_local_memory_store_usage()
+
+        output = []
+        for block, meta_with_schema in zip(new_blocks, new_metadata_schema):
+            output.append(
+                RefBundle(
+                    [BlockEntry(block, meta_with_schema.metadata)],
+                    owns_blocks=input_owned,
+                    schema=meta_with_schema.schema,
+                )
+            )
+
+        stats = {
+            "map": to_stats(shuffle_map_metadata_schema),
+            "reduce": to_stats(new_metadata_schema),
+        }
+
+        return (output, stats)

python/ray/data/_internal/planner/exchange/push_based_shuffle_task_scheduler.py

@@ -0,0 +1,847 @@
+import logging
+import math
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, TypeVar, Union
+
+import ray
+from ray._private.ray_constants import CALLER_MEMORY_USAGE_PER_OBJECT_REF
+from ray.data._internal.execution.interfaces import BlockEntry, RefBundle, TaskContext
+from ray.data._internal.planner.exchange.interfaces import (
+    ExchangeTaskScheduler,
+    ExchangeTaskSpec,
+)
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import (
+    convert_bytes_to_human_readable_str,
+    unzip,
+)
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+    _take_first_non_empty_schema,
+    to_stats,
+)
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+if TYPE_CHECKING:
+    from ray.data._internal.progress.base_progress import BaseProgressBar
+    from ray.data.block import BlockMetadataWithSchema
+
+logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+class _MergeTaskSchedule:
+    def __init__(self, output_num_blocks: int, num_merge_tasks_per_round: int):
+        self.output_num_blocks = output_num_blocks
+        self.num_merge_tasks_per_round = num_merge_tasks_per_round
+        self.num_reducers_per_merger = output_num_blocks // num_merge_tasks_per_round
+        self._num_mergers_with_extra_reducer = (
+            output_num_blocks % num_merge_tasks_per_round
+        )
+
+        if self.num_reducers_per_merger == 0:
+            self.num_merge_tasks_per_round = self._num_mergers_with_extra_reducer
+            self.num_reducers_per_merger = 1
+            self._num_mergers_with_extra_reducer = 0
+
+    def __repr__(self):
+        return (
+            f"    num merge tasks per round: {self.num_merge_tasks_per_round}\n"
+            f"    num reduce tasks per merge task: {self.num_reducers_per_merger}\n"
+            "    num merge tasks with extra reduce task: "
+            f"{self._num_mergers_with_extra_reducer}"
+        )
+
+    def get_num_reducers_per_merge_idx(self, merge_idx: int) -> int:
+        """
+        Each intermediate merge task will produce outputs for a partition of P
+        final reduce tasks. This helper function returns P based on the merge
+        task index.
+        """
+        assert merge_idx < self.num_merge_tasks_per_round
+        num_reducers_for_cur_merger = self.num_reducers_per_merger
+        if merge_idx < self._num_mergers_with_extra_reducer:
+            num_reducers_for_cur_merger += 1
+        return num_reducers_for_cur_merger
+
+    def get_merge_idx_for_reducer_idx(self, reducer_idx: int) -> int:
+        if (
+            reducer_idx
+            < (self.num_reducers_per_merger + 1) * self._num_mergers_with_extra_reducer
+        ):
+            merge_idx = reducer_idx // (self.num_reducers_per_merger + 1)
+        else:
+            reducer_idx -= (
+                self.num_reducers_per_merger + 1
+            ) * self._num_mergers_with_extra_reducer
+            merge_idx = (
+                self._num_mergers_with_extra_reducer
+                + reducer_idx // self.num_reducers_per_merger
+            )
+        assert merge_idx < self.num_merge_tasks_per_round
+        return merge_idx
+
+    def round_robin_reduce_idx_iterator(self):
+        """
+        When there are multiple nodes, merge tasks are spread throughout the
+        cluster to improve load-balancing. Each merge task produces outputs for
+        a contiguous partition of reduce tasks. This method creates an iterator
+        that returns reduce task indices round-robin across the merge tasks.
+        This can be used to submit reduce tasks in a way that spreads the load
+        evenly across the cluster.
+        """
+        idx = 0
+        round_idx = 0
+        while idx < self.output_num_blocks:
+            for merge_idx in range(self.num_merge_tasks_per_round):
+                if merge_idx < self._num_mergers_with_extra_reducer:
+                    reduce_idx = merge_idx * (self.num_reducers_per_merger + 1)
+                    num_reducers_for_cur_merger = self.num_reducers_per_merger + 1
+                else:
+                    reduce_idx = self._num_mergers_with_extra_reducer * (
+                        self.num_reducers_per_merger + 1
+                    )
+                    merge_idx -= self._num_mergers_with_extra_reducer
+                    reduce_idx += merge_idx * self.num_reducers_per_merger
+                    num_reducers_for_cur_merger = self.num_reducers_per_merger
+
+                if round_idx >= num_reducers_for_cur_merger:
+                    continue
+
+                reduce_idx += round_idx
+                yield reduce_idx
+                idx += 1
+            round_idx += 1
+
+
+class _PushBasedShuffleStage:
+    def __init__(
+        self,
+        output_num_blocks: int,
+        num_rounds: int,
+        num_map_tasks_per_round: int,
+        merge_task_placement: List[str],
+    ):
+        # The number of rounds of map-merge tasks. Reducer tasks are given the
+        # outputs of the merge tasks as inputs. Reducer tasks receive one input
+        # per round.
+        self.num_rounds = num_rounds
+        # The number of map tasks per round of map-merge tasks. The map task
+        # produces one output per merge task in the same round.
+        self.num_map_tasks_per_round = num_map_tasks_per_round
+
+        node_strategies = {
+            node_id: {
+                "scheduling_strategy": NodeAffinitySchedulingStrategy(
+                    node_id, soft=True
+                )
+            }
+            for node_id in set(merge_task_placement)
+        }
+        self._merge_task_options = [
+            node_strategies[node_id] for node_id in merge_task_placement
+        ]
+
+        self.merge_schedule = _MergeTaskSchedule(
+            output_num_blocks, len(merge_task_placement)
+        )
+
+    def get_estimated_num_refs(self) -> int:
+        # Number of intermediate blocks = Number of rounds x (map tasks per
+        # round * merge tasks per round).
+        num_intermediate_refs = self.num_rounds * (
+            self.num_map_tasks_per_round * self.merge_schedule.num_merge_tasks_per_round
+        )
+        # Number of input blocks + intermediate blocks + output blocks.
+        num_refs_total = (
+            (self.num_rounds * self.num_map_tasks_per_round)
+            + num_intermediate_refs
+            + self.merge_schedule.output_num_blocks
+        )
+        return num_refs_total
+
+    def get_merge_task_options(self, merge_idx):
+        return self._merge_task_options[merge_idx]
+
+    def __repr__(self):
+        return (
+            "num map tasks per round (num args per merge task): "
+            f"{self.num_map_tasks_per_round}\n"
+            f"num rounds (num args per reduce task): {self.num_rounds}\n"
+            "merge task placement: \n"
+            f"{self.merge_schedule}"
+        )
+
+
+class _PipelinedStageExecutor:
+    def __init__(
+        self,
+        stage_iter,
+        num_tasks_per_round: int,
+        max_concurrent_rounds: int = 1,
+        progress_bar: Optional["BaseProgressBar"] = None,
+    ):
+        self._stage_iter = stage_iter
+        self._num_tasks_per_round = num_tasks_per_round
+        self._max_concurrent_rounds = max_concurrent_rounds
+        self._progress_bar = progress_bar
+
+        self._rounds: List[List[ObjectRef]] = []
+        self._task_idx = 0
+
+        self._submit_round()
+
+        self._num_block_bytes_stored_at_driver = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> List["BlockMetadataWithSchema"]:
+        """
+        Submit one round of tasks. If we already have the max concurrent rounds
+        in flight, first wait for the oldest round of tasks to finish.
+        """
+        prev_metadata_and_schema = []
+        if all(len(r) == 0 for r in self._rounds):
+            raise StopIteration
+
+        if len(self._rounds) >= self._max_concurrent_rounds:
+            prev_metadata_schema_refs = self._rounds.pop(0)
+            if prev_metadata_schema_refs:
+                if self._progress_bar is not None:
+                    prev_metadata_and_schema = self._progress_bar.fetch_until_complete(
+                        prev_metadata_schema_refs
+                    )
+                    # TODO(swang): Eagerly free the previous round's args.
+                    # See https://github.com/ray-project/ray/issues/42145.
+                else:
+                    prev_metadata_and_schema = ray.get(prev_metadata_schema_refs)
+
+        self._submit_round()
+
+        return prev_metadata_and_schema
+
+    def _submit_round(self):
+        assert len(self._rounds) < self._max_concurrent_rounds
+        task_round = []
+        for _ in range(self._num_tasks_per_round):
+            try:
+                task_round.append(next(self._stage_iter))
+            except StopIteration:
+                break
+        self._rounds.append(task_round)
+
+
+class _MapStageIterator:
+    def __init__(self, input_blocks_list, shuffle_map, map_args):
+        self._input_blocks_list = input_blocks_list
+        self._shuffle_map = shuffle_map
+        self._map_args = map_args
+
+        self._mapper_idx = 0
+        self._map_results = []
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if not self._input_blocks_list:
+            raise StopIteration
+
+        block = self._input_blocks_list.pop(0)
+        # NOTE(swang): Results are shuffled between map and merge tasks, so
+        # there is no advantage to colocating specific map and merge tasks.
+        # Therefore, we do not specify a node affinity policy for map tasks
+        # in case the caller or Ray has a better scheduling strategy, e.g.,
+        # based on data locality.
+        map_result = self._shuffle_map.remote(
+            self._mapper_idx,
+            block,
+            *self._map_args,
+        )
+        metadata_schema_ref = map_result.pop(-1)
+        self._map_results.append(map_result)
+        self._mapper_idx += 1
+        return metadata_schema_ref
+
+    def pop_map_results(self) -> List[List[ObjectRef]]:
+        map_results = self._map_results
+        self._map_results = []
+        return map_results
+
+
+class _MergeStageIterator:
+    def __init__(
+        self,
+        map_stage_iter: _MapStageIterator,
+        shuffle_merge,
+        stage: _PushBasedShuffleStage,
+        reduce_args,
+    ):
+        self._map_stage_iter = map_stage_iter
+        self._shuffle_merge = shuffle_merge
+        self._stage = stage
+        self._reduce_args = reduce_args
+
+        self._merge_idx = 0
+        self._map_result_buffer = None
+        # Final outputs from the map-merge stage.
+        # This is a map from merge task index to a nested list of merge results
+        # (ObjectRefs). Each merge task index corresponds to a partition of P
+        # final reduce tasks.
+        self._all_merge_results = [
+            [] for _ in range(self._stage.merge_schedule.num_merge_tasks_per_round)
+        ]
+
+    def __next__(self):
+        if not self._map_result_buffer or not self._map_result_buffer[0]:
+            assert self._merge_idx == 0
+            self._map_result_buffer = self._map_stage_iter.pop_map_results()
+
+        if not self._map_result_buffer:
+            raise StopIteration
+
+        # Shuffle the map results for the merge tasks.
+        merge_args = [map_result.pop(0) for map_result in self._map_result_buffer]
+        num_merge_returns = self._stage.merge_schedule.get_num_reducers_per_merge_idx(
+            self._merge_idx
+        )
+        merge_result = self._shuffle_merge.options(
+            num_returns=1 + num_merge_returns,
+            **self._stage.get_merge_task_options(self._merge_idx),
+        ).remote(
+            *merge_args,
+            reduce_args=self._reduce_args,
+        )
+        metadata_schema_ref = merge_result.pop(-1)
+        self._all_merge_results[self._merge_idx].append(merge_result)
+        del merge_result
+
+        self._merge_idx += 1
+        self._merge_idx %= self._stage.merge_schedule.num_merge_tasks_per_round
+        return metadata_schema_ref
+
+    def pop_merge_results(self) -> List[List[ObjectRef]]:
+        """Return a nested list of merge task results. The list at index i
+        stores the outputs of the i-th merge task submitted during each
+        map-merge round. Each merge task returns a list of outputs because it
+        may produce outputs for multiple downstream reduce tasks.
+        """
+        all_merge_results = self._all_merge_results
+        self._all_merge_results = []
+        return all_merge_results
+
+
+class _ReduceStageIterator:
+    def __init__(
+        self,
+        stage: _PushBasedShuffleStage,
+        shuffle_reduce,
+        all_merge_results: List[List[List[ObjectRef]]],
+        ray_remote_args,
+        reduce_args: List[Any],
+        _debug_limit_execution_to_num_blocks: Optional[int],
+    ):
+        self._shuffle_reduce = shuffle_reduce
+        self._stage = stage
+        self._reduce_arg_blocks: List[Tuple[int, List[ObjectRef]]] = []
+        self._ray_remote_args = ray_remote_args
+        self._reduce_args = reduce_args
+
+        for reduce_idx in self._stage.merge_schedule.round_robin_reduce_idx_iterator():
+            merge_idx = self._stage.merge_schedule.get_merge_idx_for_reducer_idx(
+                reduce_idx
+            )
+            reduce_arg_blocks = [
+                merge_results.pop(0) for merge_results in all_merge_results[merge_idx]
+            ]
+            self._reduce_arg_blocks.append((reduce_idx, reduce_arg_blocks))
+
+        assert len(self._reduce_arg_blocks) == stage.merge_schedule.output_num_blocks
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            self._reduce_arg_blocks = self._reduce_arg_blocks[
+                :_debug_limit_execution_to_num_blocks
+            ]
+            logger.debug(
+                f"Limiting execution to {len(self._reduce_arg_blocks)} reduce tasks"
+            )
+
+        for merge_idx, merge_results in enumerate(all_merge_results):
+            assert all(len(merge_result) == 0 for merge_result in merge_results), (
+                "Reduce stage did not process outputs from merge tasks at index: "
+                f"{merge_idx}"
+            )
+
+        self._reduce_results: List[Tuple[int, ObjectRef]] = []
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if not self._reduce_arg_blocks:
+            raise StopIteration
+
+        reduce_idx, reduce_arg_blocks = self._reduce_arg_blocks.pop(0)
+        merge_idx = self._stage.merge_schedule.get_merge_idx_for_reducer_idx(reduce_idx)
+        # Submit one partition of reduce tasks, one for each of the P
+        # outputs produced by the corresponding merge task.
+        # We also add the merge task arguments so that the reduce task
+        # is colocated with its inputs.
+        block, meta_with_schema = self._shuffle_reduce.options(
+            **self._ray_remote_args,
+            **self._stage.get_merge_task_options(merge_idx),
+            num_returns=2,
+        ).remote(*self._reduce_args, *reduce_arg_blocks, partial_reduce=False)
+        self._reduce_results.append((reduce_idx, block))
+        return meta_with_schema
+
+    def pop_reduce_results(self):
+        reduce_results = self._reduce_results
+        self._reduce_results = []
+        return reduce_results
+
+
+class PushBasedShuffleTaskScheduler(ExchangeTaskScheduler):
+    """
+    Push-based shuffle merges intermediate map outputs on the reducer nodes
+    while other map tasks are executing. The merged outputs are merged again
+    during a final reduce stage. This works as follows:
+
+    1. Submit rounds of concurrent map and merge tasks until all map inputs
+    have been processed. In each round, we execute:
+
+       M map tasks
+         Each produces N outputs. Each output contains P blocks.
+       N merge tasks
+         Takes 1 output from each of M map tasks.
+         Each produces P outputs.
+       Where M and N are chosen to maximize parallelism across CPUs. Note that
+       this assumes that all CPUs in the cluster will be dedicated to the
+       shuffle job.
+
+       Map and merge tasks are pipelined so that we always merge the previous
+       round of map outputs while executing the next round of map tasks.
+
+    2. In the final reduce stage:
+       R reduce tasks
+         Takes 1 output from one of the merge tasks from every round.
+
+    Notes:
+        N * P = R = total number of output blocks
+        M / N = merge factor - the ratio of map : merge tasks is to improve
+          pipelined parallelism. For example, if map takes twice as long to
+          execute as merge, then we should set this to 2. If pipeline bubbles
+          appear and the merge tasks are much longer than the map tasks, then
+          the merge factor should be decreased, and vice versa.
+        See paper at https://arxiv.org/abs/2203.05072 for more details.
+    """
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        task_ctx: TaskContext,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+        merge_factor: float = 2,
+        _debug_limit_execution_to_num_blocks: int = None,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        logger.debug("Using experimental push-based shuffle.")
+        # TODO: Preemptively clear the blocks list since we will incrementally delete
+        # the last remaining references as we submit the dependent map tasks during the
+        # map-merge stage.
+
+        # TODO(swang): For jobs whose reduce work is heavier than the map work,
+        # we should support fractional merge factors.
+        # TODO(swang): For large jobs, we should try to choose the merge factor
+        # automatically, e.g., by running one test round of map and merge tasks
+        # and comparing their run times.
+        # TODO(swang): Add option to automatically reduce write amplification
+        # during map-merge stage, by limiting how many partitions can be
+        # processed concurrently.
+        input_blocks_list = []
+        for ref_bundle in refs:
+            input_blocks_list.extend(ref_bundle.block_refs)
+        input_owned = all(b.owns_blocks for b in refs)
+
+        if map_ray_remote_args is None:
+            map_ray_remote_args = {}
+        if reduce_ray_remote_args is None:
+            reduce_ray_remote_args = {}
+        # The placement strategy for reduce tasks is overwritten to colocate
+        # them with their inputs from the merge stage, so remove any
+        # pre-specified scheduling strategy here.
+        reduce_ray_remote_args = reduce_ray_remote_args.copy()
+        reduce_ray_remote_args.pop("scheduling_strategy", None)
+
+        # Compute all constants used for task scheduling.
+        num_cpus_per_node_map = _get_num_cpus_per_node_map()
+        stage = self._compute_shuffle_schedule(
+            num_cpus_per_node_map,
+            len(input_blocks_list),
+            merge_factor,
+            output_num_blocks,
+        )
+
+        caller_memory_usage = (
+            stage.get_estimated_num_refs() * CALLER_MEMORY_USAGE_PER_OBJECT_REF
+        )
+        self.warn_on_driver_memory_usage(
+            caller_memory_usage,
+            "Execution is estimated to use at least "
+            f"{convert_bytes_to_human_readable_str(caller_memory_usage)}"
+            " of driver memory. Ensure that the driver machine has at least "
+            "this much memory to ensure job completion.",
+        )
+
+        # TODO(swang): Use INFO level. Currently there is no easy way to set
+        # the logging level to DEBUG from a driver script, so just print
+        # verbosely for now.
+        # See https://github.com/ray-project/ray/issues/42002.
+        logger.debug(f"Push-based shuffle schedule:\n{stage}")
+
+        map_fn = self._map_partition
+        merge_fn = self._merge
+
+        def map_partition(*args, **kwargs):
+            return map_fn(self._exchange_spec.map, *args, **kwargs)
+
+        def merge(*args, **kwargs):
+            return merge_fn(self._exchange_spec.reduce, *args, **kwargs)
+
+        shuffle_map = cached_remote_fn(map_partition)
+        shuffle_map = shuffle_map.options(
+            **map_ray_remote_args,
+            num_returns=1 + stage.merge_schedule.num_merge_tasks_per_round,
+        )
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            input_blocks_list = input_blocks_list[:_debug_limit_execution_to_num_blocks]
+            logger.debug(f"Limiting execution to {len(input_blocks_list)} map tasks")
+        map_stage_iter = _MapStageIterator(
+            input_blocks_list,
+            shuffle_map,
+            [output_num_blocks, stage.merge_schedule, *self._exchange_spec._map_args],
+        )
+
+        sub_progress_bar_dict = task_ctx.sub_progress_bar_dict
+        bar_name = ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        map_bar = sub_progress_bar_dict[bar_name]
+        map_stage_executor = _PipelinedStageExecutor(
+            map_stage_iter, stage.num_map_tasks_per_round, progress_bar=map_bar
+        )
+
+        shuffle_merge = cached_remote_fn(merge)
+        merge_stage_iter = _MergeStageIterator(
+            map_stage_iter, shuffle_merge, stage, self._exchange_spec._reduce_args
+        )
+        merge_stage_executor = _PipelinedStageExecutor(
+            merge_stage_iter,
+            stage.merge_schedule.num_merge_tasks_per_round,
+            max_concurrent_rounds=2,
+        )
+        # Execute the map-merge stage. This submits tasks in rounds of M map
+        # tasks and N merge tasks each. Task execution between map and merge is
+        # pipelined, so that while executing merge for one round of inputs, we
+        # also execute the map tasks for the following round.
+        map_done = False
+        merge_done = False
+        map_stage_metadata_schema = []
+        merge_stage_metadata_schema = []
+        while not (map_done and merge_done):
+            try:
+                map_stage_metadata_schema += next(map_stage_executor)
+            except StopIteration:
+                map_done = True
+                break
+
+            try:
+                merge_stage_metadata_schema += next(merge_stage_executor)
+            except StopIteration:
+                merge_done = True
+                break
+
+            self.warn_on_high_local_memory_store_usage()
+
+        all_merge_results = merge_stage_iter.pop_merge_results()
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            for merge_idx in range(len(all_merge_results)):
+                while len(all_merge_results[merge_idx]) < stage.num_rounds:
+                    # Repeat the first merge task's results.
+                    all_merge_results[merge_idx].append(
+                        all_merge_results[merge_idx][0][:]
+                    )
+
+        # Execute and wait for the reduce stage.
+        bar_name = ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        reduce_bar = sub_progress_bar_dict[bar_name]
+
+        shuffle_reduce = cached_remote_fn(self._exchange_spec.reduce)
+        reduce_stage_iter = _ReduceStageIterator(
+            stage,
+            shuffle_reduce,
+            all_merge_results,
+            reduce_ray_remote_args,
+            self._exchange_spec._reduce_args,
+            _debug_limit_execution_to_num_blocks,
+        )
+
+        max_reduce_tasks_in_flight = output_num_blocks
+        ctx = DataContext.get_current()
+        if ctx.pipeline_push_based_shuffle_reduce_tasks:
+            # If pipelining is enabled, we should still try to utilize all
+            # cores.
+            max_reduce_tasks_in_flight = min(
+                max_reduce_tasks_in_flight, sum(num_cpus_per_node_map.values())
+            )
+
+        reduce_stage_executor = _PipelinedStageExecutor(
+            reduce_stage_iter,
+            max_reduce_tasks_in_flight,
+            max_concurrent_rounds=2,
+            progress_bar=reduce_bar,
+        )
+        reduce_stage_metadata_schema = []
+        while True:
+            try:
+                reduce_stage_metadata_schema += next(reduce_stage_executor)
+            except StopIteration:
+                break
+
+            self.warn_on_high_local_memory_store_usage()
+
+        new_blocks = reduce_stage_iter.pop_reduce_results()
+        sorted_blocks = [
+            (block[0], block[1], reduce_stage_metadata_schema[i])
+            for i, block in enumerate(new_blocks)
+        ]
+        sorted_blocks.sort(key=lambda x: x[0])
+
+        new_blocks, reduce_stage_metadata_schema = [], []
+        if sorted_blocks:
+            res: Tuple[
+                List[Any], List[ObjectRef[Block]], List[BlockMetadataWithSchema]
+            ] = unzip(sorted_blocks)
+            _, new_blocks, reduce_stage_metadata_schema = res
+        del sorted_blocks
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            output_num_blocks = min(
+                _debug_limit_execution_to_num_blocks, output_num_blocks
+            )
+
+        assert (
+            len(new_blocks) == output_num_blocks
+        ), f"Expected {output_num_blocks} outputs, produced {len(new_blocks)}"
+
+        output = []
+        for block, meta_with_schema in zip(new_blocks, reduce_stage_metadata_schema):
+            output.append(
+                RefBundle(
+                    [BlockEntry(block, meta_with_schema.metadata)],
+                    owns_blocks=input_owned,
+                    schema=meta_with_schema.schema,
+                )
+            )
+
+        stats = {
+            "map": to_stats(map_stage_metadata_schema),
+            "merge": to_stats(merge_stage_metadata_schema),
+            "reduce": to_stats(reduce_stage_metadata_schema),
+        }
+
+        return (output, stats)
+
+    @staticmethod
+    def _map_partition(
+        map_fn,
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        schedule: _MergeTaskSchedule,
+        *map_args: List[Any],
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        mapper_outputs = map_fn(idx, block, output_num_blocks, *map_args)
+
+        # A merge task may produce results for multiple downstream reducer
+        # tasks. Therefore, each map task should give each merge task a
+        # partition of its outputs, where the length of the partition is equal
+        # to the number of reducers downstream to the merge task.
+        partition = []
+        merge_idx = 0
+        while merge_idx < schedule.num_merge_tasks_per_round and mapper_outputs:
+            output = mapper_outputs.pop(0)
+            partition.append(output)
+
+            if len(partition) == schedule.get_num_reducers_per_merge_idx(merge_idx):
+                yield partition
+
+                partition = []
+                merge_idx += 1
+
+        assert not partition
+        assert len(mapper_outputs) == 1, (
+            mapper_outputs,
+            "The last output should be a BlockMetadataWithSchema",
+        )
+        assert isinstance(mapper_outputs[0], BlockMetadataWithSchema)
+        yield mapper_outputs[0]
+
+        assert merge_idx == schedule.num_merge_tasks_per_round, (
+            merge_idx,
+            schedule.num_merge_tasks_per_round,
+        )
+
+    @staticmethod
+    def _merge(
+        reduce_fn,
+        *all_mapper_outputs: List[List[Block]],
+        reduce_args: Optional[List[Any]] = None,
+    ) -> List[Union["BlockMetadataWithSchema", Block]]:
+        """
+        Returns list of [BlockMetadata, O1, O2, O3, ...output_num_blocks].
+        """
+        assert (
+            len({len(mapper_outputs) for mapper_outputs in all_mapper_outputs}) == 1
+        ), "Received different number of map inputs"
+        stats = BlockExecStats.builder()
+        if not reduce_args:
+            reduce_args = []
+
+        num_rows = 0
+        size_bytes = 0
+        schemas = []
+        for i, mapper_outputs in enumerate(zip(*all_mapper_outputs)):
+            block_meta_with_schema: Tuple[Block, "BlockMetadataWithSchema"] = reduce_fn(
+                *reduce_args, *mapper_outputs, partial_reduce=True
+            )
+            block, meta_with_schema = block_meta_with_schema
+            yield block
+
+            block = BlockAccessor.for_block(block)
+            num_rows += block.num_rows()
+            size_bytes += block.size_bytes()
+            del block
+            schemas.append(meta_with_schema.schema)
+
+        schema = _take_first_non_empty_schema(iter(schemas))
+
+        meta = BlockMetadata(
+            num_rows=num_rows,
+            size_bytes=size_bytes,
+            input_files=None,
+            exec_stats=stats.build(),
+        )
+        meta_with_schema = BlockMetadataWithSchema.from_metadata(meta, schema=schema)
+        yield meta_with_schema
+
+    @staticmethod
+    def _compute_shuffle_schedule(
+        num_cpus_per_node_map: Dict[str, int],
+        num_input_blocks: int,
+        merge_factor: float,
+        num_output_blocks: int,
+    ) -> _PushBasedShuffleStage:
+        num_cpus_total = sum(v for v in num_cpus_per_node_map.values())
+        logger.debug(
+            f"Found {num_cpus_total} CPUs available CPUs for push-based shuffle."
+        )
+        num_tasks_per_map_merge_group = merge_factor + 1
+        num_total_merge_tasks = math.ceil(num_input_blocks / merge_factor)
+
+        num_merge_tasks_per_round = 0
+        merge_task_placement = []
+        leftover_cpus = 0
+        # Compute the total number of merge tasks and their node placement.
+        # Each merge task should be grouped with `merge_factor` map tasks for
+        # pipelining. These groups should then be spread across nodes according
+        # to CPU availability for load-balancing.
+        num_input_blocks_remaining = num_input_blocks
+        for node, num_cpus in num_cpus_per_node_map.items():
+            # First find how many merge tasks we should run on this node.
+            # We take the min of the number of CPUs on this node and the number
+            # of input blocks that we haven't scheduled yet, in case there are
+            # fewer input blocks than CPU slots on this node.
+            num_cpu_slots = min(num_cpus, num_input_blocks_remaining)
+            num_merge_tasks_on_cur_node = round(
+                num_cpu_slots / num_tasks_per_map_merge_group
+            )
+            # For small datasets, the number of tasks to run may be less than
+            # the total CPU slots available.
+            num_merge_tasks_on_cur_node = min(
+                num_merge_tasks_on_cur_node, num_total_merge_tasks
+            )
+            for i in range(num_merge_tasks_on_cur_node):
+                merge_task_placement.append(node)
+                # We schedule `merge_factor` many map tasks for every merge
+                # task. Subtract from the number of input blocks remaining to
+                # account for cases where the number of map tasks is smaller
+                # than the available CPU slots.
+                num_input_blocks_remaining -= merge_factor
+                num_cpus -= num_tasks_per_map_merge_group
+            num_merge_tasks_per_round += num_merge_tasks_on_cur_node
+
+            # Handle the case where a single node cannot fit a group of map and
+            # merge tasks, but we can spread the group across multiple distinct
+            # nodes.
+            leftover_cpus += num_cpus
+            if (
+                leftover_cpus >= num_tasks_per_map_merge_group
+                and num_merge_tasks_per_round < num_total_merge_tasks
+            ):
+                merge_task_placement.append(node)
+                num_merge_tasks_per_round += 1
+                leftover_cpus -= num_tasks_per_map_merge_group
+                num_input_blocks_remaining -= merge_factor
+
+            num_input_blocks_remaining = max(0, num_input_blocks_remaining)
+
+        if num_merge_tasks_per_round == 0:
+            # For small datasets, make sure we have at least one merge task.
+            for node, num_cpus in num_cpus_per_node_map.items():
+                if num_cpus >= 1:
+                    merge_task_placement.append(node)
+                    num_merge_tasks_per_round = 1
+                    break
+
+        assert num_merge_tasks_per_round == len(merge_task_placement)
+        assert num_merge_tasks_per_round > 0, num_merge_tasks_per_round
+        # Use the remaining CPUs to execute map tasks.
+        num_map_tasks_per_round = num_cpus_total - num_merge_tasks_per_round
+        num_map_tasks_per_round = min(num_map_tasks_per_round, num_input_blocks)
+        # Make sure there is at least one map task in each round.
+        num_map_tasks_per_round = max(num_map_tasks_per_round, 1)
+
+        num_rounds = math.ceil(num_input_blocks / num_map_tasks_per_round)
+        return _PushBasedShuffleStage(
+            num_output_blocks,
+            num_rounds,
+            num_map_tasks_per_round,
+            merge_task_placement,
+        )
+
+
+def _get_num_cpus_per_node_map() -> Dict[str, int]:
+    total_resources_by_node = ray.state.total_resources_per_node()
+    # Map from per-node resource name to number of CPUs available on that
+    # node.
+    num_cpus_per_node_map = {}
+    for node_id, resources in total_resources_by_node.items():
+        num_cpus = int(resources.get("CPU", 0))
+        if num_cpus == 0:
+            continue
+        num_cpus_per_node_map[node_id] = num_cpus
+    return num_cpus_per_node_map

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

@@ -0,0 +1,152 @@
+import logging
+import math
+from typing import Callable, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+)
+from ray.data.context import MAX_SAFE_BLOCK_SIZE_FACTOR
+
+logger = logging.getLogger(__name__)
+
+
+class ShuffleTaskSpec(ExchangeTaskSpec):
+    """
+    The implementation for shuffle tasks.
+
+    This is used by random_shuffle() and repartition().
+    """
+
+    SPLIT_REPARTITION_SUB_PROGRESS_BAR_NAME = "Split Repartition"
+
+    def __init__(
+        self,
+        target_shuffle_max_block_size: int,
+        random_shuffle: bool = False,
+        random_seed: Optional[int] = None,
+        upstream_map_fn: Optional[Callable[[Iterable[Block]], Iterable[Block]]] = None,
+    ):
+        super().__init__(
+            map_args=[
+                target_shuffle_max_block_size,
+                upstream_map_fn,
+                random_shuffle,
+                random_seed,
+            ],
+            reduce_args=[random_shuffle, random_seed],
+        )
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        target_shuffle_max_block_size: int,
+        upstream_map_fn: Optional[Callable[[Iterable[Block]], Iterable[Block]]],
+        random_shuffle: bool,
+        random_seed: Optional[int],
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        stats = BlockExecStats.builder()
+        if upstream_map_fn:
+            # Create a local TaskContext for the upstream map function.
+            # May be used by expressions that depend on task-level state.
+            local_ctx = TaskContext(task_idx=idx, op_name="shuffle_map")
+            with TaskContext.current(local_ctx):
+                # TODO: Support dynamic block splitting in
+                # all-to-all ops, to avoid having to re-fuse
+                # upstream blocks together.
+                upstream_map_iter = upstream_map_fn([block])
+                mapped_block = next(upstream_map_iter)
+                builder = BlockAccessor.for_block(mapped_block).builder()
+                builder.add_block(mapped_block)
+                for mapped_block in upstream_map_iter:
+                    builder.add_block(mapped_block)
+                # Drop the upstream inputs to reduce memory usage.
+                del mapped_block
+                block = builder.build()
+
+        block = BlockAccessor.for_block(block)
+        if (
+            block.size_bytes()
+            > MAX_SAFE_BLOCK_SIZE_FACTOR * target_shuffle_max_block_size
+        ):
+            logger.warning(
+                "Input block to map task has size "
+                f"{block.size_bytes() // (1024 * 1024)}MiB, which exceeds "
+                "DataContext.get_current().target_shuffle_max_block_size="
+                f"{target_shuffle_max_block_size // (1024 * 1024)}MiB. "
+                "This can lead to out-of-memory errors and can happen "
+                "when map tasks are fused to the shuffle operation. "
+                "To prevent fusion, call Dataset.materialize() on the "
+                "dataset before shuffling."
+            )
+
+        # Randomize the distribution of records to blocks.
+        if random_shuffle:
+            seed_i = random_seed + idx if random_seed is not None else None
+            block = block.random_shuffle(seed_i)
+            block = BlockAccessor.for_block(block)
+
+        # Build a list of slices to return. It's okay to put the results in a
+        # list instead of yielding them as a generator because slicing the
+        # ArrowBlock is zero-copy.
+        slice_sz = max(1, math.ceil(block.num_rows() / output_num_blocks))
+        slices = []
+        for i in range(output_num_blocks):
+            slices.append(block.slice(i * slice_sz, (i + 1) * slice_sz))
+
+        # Randomize the distribution order of the blocks (this prevents empty
+        # outputs when input blocks are very small).
+        if random_shuffle:
+            random = np.random.RandomState(seed_i)
+            random.shuffle(slices)
+
+        num_rows = sum(BlockAccessor.for_block(s).num_rows() for s in slices)
+        assert num_rows == block.num_rows(), (num_rows, block.num_rows())
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta = block.get_metadata(block_exec_stats=stats.build())
+        schema = block.schema()
+        meta_with_schema = BlockMetadataWithSchema.from_metadata(meta, schema=schema)
+        return slices + [meta_with_schema]
+
+    @staticmethod
+    def reduce(
+        random_shuffle: bool,
+        random_seed: Optional[int],
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        # TODO: Support fusion with other downstream operators.
+        stats = BlockExecStats.builder()
+        builder = DelegatingBlockBuilder()
+        for block in mapper_outputs:
+            builder.add_block(block)
+        new_block = builder.build()
+        accessor = BlockAccessor.for_block(new_block)
+        if random_shuffle:
+            new_block = accessor.random_shuffle(
+                random_seed if random_seed is not None else None
+            )
+            accessor = BlockAccessor.for_block(new_block)
+        new_metadata = BlockMetadata(
+            num_rows=accessor.num_rows(),
+            size_bytes=accessor.size_bytes(),
+            input_files=None,
+            exec_stats=stats.build(),
+        )
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta_with_schema = BlockMetadataWithSchema.from_metadata(
+            new_metadata, schema=accessor.schema()
+        )
+        return new_block, meta_with_schema

python/ray/data/_internal/planner/exchange/sort_task_spec.py

@@ -0,0 +1,240 @@
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, TypeVar, Union
+
+import numpy as np
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
+from ray.data._internal.progress.progress_bar import ProgressBar
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.table_block import TableBlockAccessor
+from ray.data._internal.util import NULL_SENTINEL
+from ray.data.block import Block, BlockAccessor, BlockExecStats
+from ray.types import ObjectRef
+
+T = TypeVar("T")
+
+if TYPE_CHECKING:
+    import pyarrow
+
+    from ray.data.block import BlockMetadataWithSchema
+
+
+class SortKey:
+    """SortKey class to convert between different sort args formats."""
+
+    def __init__(
+        self,
+        key: Optional[Union[str, List[str]]] = None,
+        descending: Union[bool, List[bool]] = False,
+        boundaries: Optional[List[T]] = None,
+    ):
+        if key is None:
+            key = []
+        if isinstance(key, str):
+            key = [key]
+        if not (isinstance(key, list) and all(isinstance(k, str) for k in key)):
+            raise ValueError(
+                f"Key must be a string or a list of strings, but got {key}."
+            )
+        if isinstance(descending, bool):
+            descending = [descending for _ in key]
+        elif isinstance(descending, list):
+            if len(descending) != len(key):
+                raise ValueError(
+                    "Length of `descending` does not match the length of the key."
+                )
+        self._columns = key
+        self._descending = descending
+        if boundaries:
+            for item in boundaries:
+                if not isinstance(item, (int, float)):
+                    raise ValueError(
+                        "The type of items in boundaries must be int or float."
+                    )
+            boundaries = list(set(boundaries))
+            boundaries.sort()
+        self._boundaries = boundaries
+
+    def get_columns(self) -> List[str]:
+        return self._columns
+
+    def get_descending(self) -> List[bool]:
+        return self._descending
+
+    def to_arrow_sort_args(self) -> List[Tuple[str, str]]:
+        return [
+            (key, "descending" if desc else "ascending")
+            for key, desc in zip(self._columns, self._descending)
+        ]
+
+    def to_pandas_sort_args(self) -> Tuple[List[str], List[bool]]:
+        return self._columns, [not desc for desc in self._descending]
+
+    def validate_schema(self, schema: Optional[Union[type, "pyarrow.lib.Schema"]]):
+        """Check the key function is valid on the given schema."""
+        if schema is None:
+            # Dataset is empty/cleared, validation not possible.
+            return
+
+        if self._columns and len(schema.names) > 0:
+            schema_names_set = set(schema.names)
+            for column in self._columns:
+                if column not in schema_names_set:
+                    raise ValueError(
+                        f"You specified the column '{column}', but there's no such "
+                        "column in the dataset. The dataset has columns: "
+                        f"{schema.names}"
+                    )
+
+    @property
+    def boundaries(self):
+        return self._boundaries
+
+
+class SortTaskSpec(ExchangeTaskSpec):
+    """
+    The implementation for distributed sort tasks.
+
+    The algorithm is similar to [External Merge Sort]
+    (https://en.wikipedia.org/wiki/External_sorting).
+    Sorting is done in 3 steps: sampling, sorting individual blocks, and
+    merging sorted blocks.
+
+    Sampling (`sample_boundaries`): we get a number of sample items from each block,
+    sort them, and use them to compute boundaries that would partition all items into
+    approximately equal ranges.
+
+    Sorting (`map`): each block is sorted locally, then partitioned into smaller
+    blocks according to the boundaries. Each partitioned block is passed to a merge
+    task.
+
+    Merging (`reduce`): a merge task would receive a block from every worker that
+    consists of items in a certain range. It then merges the sorted blocks into one
+    sorted block and becomes part of the new, sorted block.
+    """
+
+    SORT_SAMPLE_SUB_PROGRESS_BAR_NAME = "Sort Sample"
+
+    def __init__(
+        self,
+        boundaries: List[T],
+        sort_key: SortKey,
+    ):
+        super().__init__(
+            map_args=[boundaries, sort_key],
+            reduce_args=[sort_key],
+        )
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        boundaries: List[T],
+        sort_key: SortKey,
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        stats = BlockExecStats.builder()
+        accessor = BlockAccessor.for_block(block)
+        out = accessor.sort_and_partition(boundaries, sort_key)
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta_with_schema = BlockMetadataWithSchema.from_block(
+            block, block_exec_stats=stats.build()
+        )
+        return out + [meta_with_schema]
+
+    @staticmethod
+    def reduce(
+        sort_key: SortKey,
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        normalized_blocks = TableBlockAccessor.normalize_block_types(
+            mapper_outputs,
+            target_block_type=None,
+        )
+        blocks, meta_with_schema = BlockAccessor.for_block(
+            normalized_blocks[0]
+        ).merge_sorted_blocks(normalized_blocks, sort_key)
+        return blocks, meta_with_schema
+
+    @staticmethod
+    def sample_boundaries(
+        blocks: List[ObjectRef[Block]],
+        sort_key: SortKey,
+        num_reducers: int,
+        sample_bar: Optional[ProgressBar] = None,
+        label_selector: Optional[Dict[str, str]] = None,
+    ) -> List[T]:
+        """
+        Return (num_reducers - 1) items in ascending order from the blocks that
+        partition the domain into ranges with approximately equally many elements.
+        Each boundary item is a tuple of a form (col1_value, col2_value, ...).
+        """
+        columns = sort_key.get_columns()
+        n_samples = int(num_reducers * 10 / len(blocks))
+
+        sample_block = cached_remote_fn(_sample_block)
+        if label_selector:
+            sample_block = sample_block.options(label_selector=label_selector)
+
+        sample_results = [
+            sample_block.remote(block, n_samples, sort_key) for block in blocks
+        ]
+        if sample_bar is None:
+            sample_bar = ProgressBar(
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME,
+                len(blocks) * n_samples,
+                unit="rows",
+            )
+        # TODO(zhilong): Update sort sample bar before finished.
+        samples = sample_bar.fetch_until_complete(sample_results)
+        del sample_results
+        samples: List[Block] = [s for s in samples if len(s) > 0]
+        # The dataset is empty
+        if len(samples) == 0:
+            return [None] * (num_reducers - 1)
+
+        # Convert samples to a sorted list[tuple[...]] where each tuple represents a
+        # sample.
+        # TODO: Once we deprecate pandas blocks, we can avoid this conversion and
+        # directly sort the samples.
+        builder = DelegatingBlockBuilder()
+        for sample in samples:
+            builder.add_block(sample)
+        samples_table = builder.build()
+        samples_dict = BlockAccessor.for_block(samples_table).to_numpy(columns=columns)
+        # This zip does the transposition from list of column values to list of tuples.
+        samples_list = list(zip(*samples_dict.values()))
+
+        def is_na(x):
+            # Check if x is None or NaN. Type casting to np.array first to avoid
+            # isnan failing on strings and other types.
+            if x is None:
+                return True
+            x = np.asarray(x)
+            if np.issubdtype(x.dtype, np.number):
+                return np.isnan(x)
+            return False
+
+        # To allow multi-directional sort, we utilize Python's stable sort: we
+        # sort several times with different directions. We do this in reverse, so
+        # that the last key we sort by is the primary sort key passed by the user.
+        for i, desc in list(enumerate(sort_key.get_descending()))[::-1]:
+            # Sort the list, but Nones should be NULL_SENTINEL to ensure safe sorting.
+            samples_list.sort(
+                key=lambda sample: NULL_SENTINEL if is_na(sample[i]) else sample[i],
+                reverse=desc,
+            )
+
+        # Each boundary corresponds to a quantile of the data.
+        quantile_indices = [
+            int(q * (len(samples_list) - 1))
+            for q in np.linspace(0, 1, num_reducers + 1)
+        ]
+        # Exclude the first and last quantiles because they're 0 and 1.
+        return [samples_list[i] for i in quantile_indices[1:-1]]
+
+
+def _sample_block(block: Block, n_samples: int, sort_key: SortKey) -> Block:
+    return BlockAccessor.for_block(block).sample(n_samples, sort_key)

python/ray/data/_internal/planner/exchange/split_repartition_task_scheduler.py

@@ -0,0 +1,171 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+import ray
+from ray.data._internal.execution.interfaces import BlockEntry, RefBundle, TaskContext
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.planner.exchange.interfaces import (
+    ExchangeTaskScheduler,
+)
+from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.split import _split_at_indices
+from ray.data._internal.util import unzip
+from ray.data.block import (
+    Block,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+)
+from ray.types import ObjectRef
+
+
+class SplitRepartitionTaskScheduler(ExchangeTaskScheduler):
+    """
+    The split (non-shuffle) repartition scheduler.
+
+    First, we calculate global splits needed to produce `output_num_blocks` blocks.
+    After the split blocks are generated accordingly, reduce tasks are scheduled
+    to combine split blocks together.
+    """
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        ctx: TaskContext,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+    ) -> AllToAllTransformFnResult:
+        input_num_rows = 0
+        input_owned_by_consumer = True
+        for ref_bundle in refs:
+            block_num_rows = ref_bundle.num_rows()
+            if block_num_rows is None:
+                raise ValueError(
+                    "Cannot split partition on blocks with unknown number of rows."
+                )
+            input_num_rows += block_num_rows
+            if not ref_bundle.owns_blocks:
+                input_owned_by_consumer = False
+
+        # Compute the (output_num_blocks) indices needed for an equal split of the
+        # input blocks. When output_num_blocks=1, the total number of
+        # input rows is used as the end index during the split calculation,
+        # so that we can combine all input blocks into a single output block.
+        indices = []
+        if output_num_blocks == 1:
+            indices = [input_num_rows]
+        else:
+            cur_idx = 0
+            for _ in range(output_num_blocks - 1):
+                cur_idx += input_num_rows / output_num_blocks
+                indices.append(int(cur_idx))
+        assert len(indices) <= output_num_blocks, (indices, output_num_blocks)
+
+        if map_ray_remote_args is None:
+            map_ray_remote_args = {}
+        if reduce_ray_remote_args is None:
+            reduce_ray_remote_args = {}
+        if "scheduling_strategy" not in reduce_ray_remote_args:
+            reduce_ray_remote_args = reduce_ray_remote_args.copy()
+            reduce_ray_remote_args["scheduling_strategy"] = "SPREAD"
+
+        blocks_with_metadata: List[Tuple[ObjectRef[Block], BlockMetadata]] = []
+        for ref_bundle in refs:
+            blocks_with_metadata.extend(
+                (entry.ref, entry.metadata) for entry in ref_bundle.blocks
+            )
+        split_return = _split_at_indices(
+            blocks_with_metadata,
+            indices,
+            input_owned_by_consumer,
+            label_selector=map_ray_remote_args.get("label_selector"),
+        )
+        split_block_refs, split_metadata = [], []
+        for b, m in zip(*split_return):
+            split_block_refs.append(b)
+            split_metadata.extend(m)
+
+        sub_progress_bar_dict = ctx.sub_progress_bar_dict
+        bar_name = ShuffleTaskSpec.SPLIT_REPARTITION_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        reduce_bar = sub_progress_bar_dict[bar_name]
+
+        reduce_task = cached_remote_fn(self._exchange_spec.reduce)
+        reduce_return = [
+            reduce_task.options(**reduce_ray_remote_args, num_returns=2).remote(
+                *self._exchange_spec._reduce_args,
+                *split_block_refs[j],
+            )
+            for j in range(output_num_blocks)
+            # Only process splits which contain blocks.
+            if len(split_block_refs[j]) > 0
+        ]
+
+        reduce_block_refs, reduce_metadata_schema = [], []
+        if reduce_return:
+            reduce_block_refs, reduce_metadata_schema = unzip(reduce_return)
+        reduce_metadata_schema: List[
+            "BlockMetadataWithSchema"
+        ] = reduce_bar.fetch_until_complete(list(reduce_metadata_schema))
+        reduce_block_refs = list(reduce_block_refs)
+
+        # Handle empty blocks.
+        if len(reduce_block_refs) < output_num_blocks:
+            import pyarrow as pa
+
+            from ray.data._internal.arrow_block import ArrowBlockBuilder
+            from ray.data._internal.pandas_block import (
+                PandasBlockBuilder,
+                PandasBlockSchema,
+            )
+
+            num_empty_blocks = output_num_blocks - len(reduce_block_refs)
+            if len(reduce_metadata_schema) > 0:
+                first_block_schema = reduce_metadata_schema[0].schema
+                if isinstance(first_block_schema, pa.Schema):
+                    builder = ArrowBlockBuilder()
+                elif isinstance(first_block_schema, PandasBlockSchema):
+                    builder = PandasBlockBuilder()
+                else:
+                    raise ValueError(
+                        "Cannot split partition on blocks with unknown block schema:"
+                        f" {first_block_schema}."
+                    )
+            else:
+                # If the result is empty, default to Arrow format for the empty blocks.
+                builder = ArrowBlockBuilder()
+
+            empty_block = builder.build()
+            empty_meta_with_schema = BlockMetadataWithSchema.from_block(
+                empty_block
+            )  # No stats for empty block.
+            empty_block_refs, empty_metadata = zip(
+                *[
+                    (ray.put(empty_block), empty_meta_with_schema)
+                    for _ in range(num_empty_blocks)
+                ]
+            )
+            reduce_block_refs.extend(empty_block_refs)
+            reduce_metadata_schema.extend(empty_metadata)
+
+        output = []
+        assert len(reduce_block_refs) == len(reduce_metadata_schema), (
+            len(reduce_block_refs),
+            len(reduce_metadata_schema),
+        )
+        for block, meta_with_schema in zip(reduce_block_refs, reduce_metadata_schema):
+            output.append(
+                RefBundle(
+                    [BlockEntry(block, meta_with_schema.metadata)],
+                    owns_blocks=input_owned_by_consumer,
+                    schema=meta_with_schema.schema,
+                )
+            )
+        stats = {
+            "split": split_metadata,
+            "reduce": reduce_metadata_schema,
+        }
+
+        return (output, stats)

python/ray/data/_internal/planner/plan_all_to_all_op.py

@@ -0,0 +1,280 @@
+import logging
+from typing import List
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.base_physical_operator import (
+    AllToAllOperator,
+)
+from ray.data._internal.execution.operators.hash_shuffle_v2 import (
+    _SHUFFLE_MAP_RUNTIME_ENV,
+    _concat_reduce,
+    _make_hash_partition_fn,
+    _sort_reduce,
+)
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (  # noqa: E501
+    ShuffleMapOp,
+)
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import (  # noqa: E501
+    ShuffleReduceOp,
+)
+from ray.data._internal.logical.operators import (
+    AbstractAllToAll,
+    Aggregate,
+    RandomizeBlocks,
+    RandomShuffle,
+    Repartition,
+    Sort,
+)
+from ray.data._internal.planner.aggregate import generate_aggregate_fn
+from ray.data._internal.planner.random_shuffle import generate_random_shuffle_fn
+from ray.data._internal.planner.randomize_blocks import generate_randomize_blocks_fn
+from ray.data._internal.planner.repartition import generate_repartition_fn
+from ray.data._internal.planner.sort import generate_sort_fn
+from ray.data.context import DataContext, ShuffleStrategy
+
+logger = logging.getLogger(__name__)
+
+
+def _plan_gpu_shuffle_repartition(
+    data_context: DataContext,
+    logical_op: Repartition,
+    input_physical_op: PhysicalOperator,
+) -> PhysicalOperator:
+    from ray.data._internal.gpu_shuffle.hash_shuffle import GPUShuffleOperator
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+    normalized_key_columns = SortKey(logical_op.keys).get_columns()
+
+    schema = logical_op.infer_schema()
+    columns = list(schema.names) if schema is not None else None
+
+    return GPUShuffleOperator(
+        input_physical_op,
+        data_context,
+        key_columns=tuple(normalized_key_columns),
+        columns=columns,
+        num_partitions=logical_op.num_outputs,
+        should_sort=logical_op.sort,
+    )
+
+
+def _plan_hash_shuffle_repartition(
+    data_context: DataContext,
+    logical_op: Repartition,
+    input_physical_op: PhysicalOperator,
+) -> PhysicalOperator:
+    """Build the two-op (ShuffleMapOp → ShuffleReduceOp) DAG for V2 hash shuffle.
+
+    Returns the reduce op; the executor crawls upstream via its
+    input_dependencies to find the map op.
+    """
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+    normalized_key_columns = SortKey(logical_op.keys).get_columns()
+    key_list = list(normalized_key_columns)
+
+    input_logical_op = input_physical_op._logical_operators[0]
+    estimated_input_blocks = input_logical_op.estimated_num_outputs()
+    target_num_partitions = (
+        logical_op.num_outputs
+        or estimated_input_blocks
+        or data_context.default_hash_shuffle_parallelism
+    )
+
+    partition_fn = _make_hash_partition_fn(key_list, target_num_partitions)
+    reduce_fn = _sort_reduce(key_list) if logical_op.sort else _concat_reduce
+
+    map_op = ShuffleMapOp(
+        input_physical_op,
+        data_context,
+        num_partitions=target_num_partitions,
+        partition_fn=partition_fn,
+        map_runtime_env=_SHUFFLE_MAP_RUNTIME_ENV,
+        name=(
+            f"HashShuffleMap(keys={tuple(key_list)}, "
+            f"partitions={target_num_partitions})"
+        ),
+    )
+    reduce_op = ShuffleReduceOp(
+        map_op,
+        data_context,
+        num_partitions=target_num_partitions,
+        reduce_fn=reduce_fn,
+        disallow_block_splitting=True,
+        name=(
+            f"HashShuffleReduce(keys={tuple(key_list)}, "
+            f"partitions={target_num_partitions})"
+        ),
+    )
+    return reduce_op
+
+
+def _plan_hash_shuffle_aggregate(
+    data_context: DataContext,
+    logical_op: Aggregate,
+    input_physical_op: PhysicalOperator,
+) -> PhysicalOperator:
+    from ray.data._internal.execution.operators.hash_aggregate import (
+        HashAggregateOperator,
+    )
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+    normalized_key_columns = SortKey(logical_op.key).get_columns()
+
+    return HashAggregateOperator(
+        data_context,
+        input_physical_op,
+        key_columns=tuple(normalized_key_columns),  # noqa: type
+        aggregation_fns=tuple(logical_op.aggs),  # noqa: type
+        # NOTE: In case number of partitions is not specified, we fall back to
+        #       default min parallelism configured
+        num_partitions=logical_op.num_partitions,
+        # TODO wire in aggregator args overrides
+    )
+
+
+def _plan_gpu_shuffle_aggregate(
+    data_context: DataContext,
+    logical_op: Aggregate,
+    input_physical_op: PhysicalOperator,
+) -> PhysicalOperator:
+    from ray.data._internal.gpu_shuffle.hash_aggregate import (
+        GPUAggregateFn,
+        GPUHashAggregateOperator,
+        build_gpu_aggregation_plan,
+    )
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+    normalized_key_columns = SortKey(logical_op.key).get_columns()
+    key_columns = tuple(normalized_key_columns)
+    aggregation_fns = tuple(logical_op.aggs)
+    input_schema = logical_op.input_dependencies[0].infer_schema()
+
+    aggregation_plan = build_gpu_aggregation_plan(
+        key_columns, aggregation_fns, input_schema=input_schema
+    )
+    if isinstance(aggregation_plan, str):
+        # Fall back to CPU hash aggregate if GPU aggregation plan is not supported.
+        fallback_reason = aggregation_plan
+        if any(isinstance(agg, GPUAggregateFn) for agg in aggregation_fns):
+            raise ValueError(
+                "GPU aggregation plan is not supported for a GPUAggregateFn "
+                f"aggregate list with key={logical_op.key}, aggs={logical_op.aggs}: "
+                f"{fallback_reason}."
+            )
+        logger.warning(
+            "GPU aggregation plan is not supported for key=%s, aggs=%s: %s; "
+            "falling back to CPU hash aggregate.",
+            logical_op.key,
+            logical_op.aggs,
+            fallback_reason,
+        )
+        return _plan_hash_shuffle_aggregate(data_context, logical_op, input_physical_op)
+
+    return GPUHashAggregateOperator(
+        data_context,
+        input_physical_op,
+        key_columns=key_columns,  # noqa: type
+        aggregation_plan=aggregation_plan,
+        num_partitions=logical_op.num_partitions,
+    )
+
+
+def plan_all_to_all_op(
+    op: AbstractAllToAll,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Get the corresponding physical operators DAG for AbstractAllToAll operators.
+
+    Note this method only converts the given `op`, but not its input dependencies.
+    See Planner.plan() for more details.
+    """
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    if isinstance(op, RandomizeBlocks):
+        fn = generate_randomize_blocks_fn(op, data_context)
+        # Randomize block order does not actually compute anything, so we
+        # want to inherit the upstream op's target max block size.
+
+    elif isinstance(op, RandomShuffle):
+        debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+            "debug_limit_shuffle_execution_to_num_blocks", None
+        )
+        fn = generate_random_shuffle_fn(
+            data_context,
+            op.seed_config,
+            op.num_outputs,
+            op.ray_remote_args,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+
+    elif isinstance(op, Repartition):
+        if op.keys:
+            if data_context.shuffle_strategy == ShuffleStrategy.GPU_SHUFFLE:
+                return _plan_gpu_shuffle_repartition(
+                    data_context, op, input_physical_dag
+                )
+            elif data_context.shuffle_strategy == ShuffleStrategy.HASH_SHUFFLE:
+                return _plan_hash_shuffle_repartition(
+                    data_context, op, input_physical_dag
+                )
+            else:
+                raise ValueError(
+                    "Key-based repartitioning only supported for "
+                    f"`DataContext.shuffle_strategy=HASH_SHUFFLE` or "
+                    f"`DataContext.shuffle_strategy=GPU_SHUFFLE` "
+                    f"(got {data_context.shuffle_strategy})"
+                )
+
+        elif op.shuffle:
+            debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+                "debug_limit_shuffle_execution_to_num_blocks", None
+            )
+        else:
+            debug_limit_shuffle_execution_to_num_blocks = None
+
+        fn = generate_repartition_fn(
+            op.num_outputs,
+            op.shuffle,
+            data_context,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+
+    elif isinstance(op, Sort):
+        debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+            "debug_limit_shuffle_execution_to_num_blocks", None
+        )
+        fn = generate_sort_fn(
+            op.sort_key,
+            data_context,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+
+    elif isinstance(op, Aggregate):
+        if data_context.shuffle_strategy == ShuffleStrategy.GPU_SHUFFLE:
+            return _plan_gpu_shuffle_aggregate(data_context, op, input_physical_dag)
+        elif data_context.shuffle_strategy == ShuffleStrategy.HASH_SHUFFLE:
+            return _plan_hash_shuffle_aggregate(data_context, op, input_physical_dag)
+
+        debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+            "debug_limit_shuffle_execution_to_num_blocks", None
+        )
+        fn = generate_aggregate_fn(
+            op.key,
+            op.aggs,
+            data_context,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+    else:
+        raise ValueError(f"Found unknown logical operator during planning: {op}")
+
+    return AllToAllOperator(
+        fn,
+        input_physical_dag,
+        data_context,
+        num_outputs=op.num_outputs,
+        sub_progress_bar_names=op.sub_progress_bar_names,
+        name=op.name,
+    )

python/ray/data/_internal/planner/plan_download_op.py

@@ -0,0 +1,311 @@
+import logging
+from typing import Iterator, List, Optional
+
+import pyarrow as pa
+
+from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.actor_pool_map_operator import (
+    ActorPoolMapOperator,
+)
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.logical.operators import Download
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.planner._obstore_download import (
+    OBSTORE_AVAILABLE,
+    _log_fallback_warning,
+    _plan_obstore_routing,
+    download_bytes_async,
+)
+from ray.data._internal.planner.download_partition_actor import (
+    URI_DOWNLOAD_MAX_WORKERS,
+    AsyncPartitionActor,
+    PartitionActor,
+)
+from ray.data._internal.util import (
+    RetryingPyFileSystem,
+    _iter_arrow_table_for_target_max_block_size,
+    make_async_gen,
+)
+from ray.data.block import BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.path_util import (
+    _resolve_paths_and_filesystem,
+    _validate_and_wrap_filesystem,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def plan_download_op(
+    op: Download,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    """Plan the download operation with partitioning and downloading stages."""
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    upstream_op_is_download = False
+    if len(input_physical_dag._logical_operators) == 1 and isinstance(
+        input_physical_dag._logical_operators[0], Download
+    ):
+        upstream_op_is_download = True
+
+    uri_column_names = op.uri_column_names
+    uri_column_names_str = ", ".join(uri_column_names)
+    output_bytes_column_names = op.output_bytes_column_names
+    ray_remote_args = op.ray_remote_args
+    filesystem = op.filesystem
+
+    # Import _get_udf from the main planner file
+    from ray.data._internal.planner.plan_udf_map_op import (
+        _generate_transform_fn_for_map_batches,
+        _get_udf,
+    )
+
+    # If we have multiple download operators in a row, we should only include the partition actor
+    # at the start of the chain. This is primarily done to prevent partition actors from bottlenecking
+    # the chain becuase the interleaved operators would be a single actor. As a result, the
+    # URIDownloader physical operator is responsible for outputting appropriately sized blocks.
+    # Decide obstore vs threaded upfront. For fsspec-S3 filesystems backed by
+    # a session we can't statically introspect (Okta / STS / profile-based),
+    # _plan_obstore_routing emits a warning and returns use_obstore=False so
+    # we fall back to the threaded PyArrow path — which uses the user's
+    # filesystem directly and resolves credentials correctly.
+    use_obstore_path = False
+    if OBSTORE_AVAILABLE:
+        use_obstore_path, _ = _plan_obstore_routing(filesystem)
+
+    partition_map_operator = None
+    if not upstream_op_is_download:
+        partition_cls = AsyncPartitionActor if use_obstore_path else PartitionActor
+        # PartitionActor / AsyncPartitionActor are callable classes, so we need
+        # ActorPoolStrategy.
+        partition_compute = ActorPoolStrategy(
+            size=1, enable_true_multi_threading=True
+        )  # Use single actor for partitioning
+
+        fn, init_fn = _get_udf(
+            partition_cls,
+            (),
+            {},
+            (uri_column_names, data_context, filesystem),
+            {},
+            compute=partition_compute,
+        )
+        block_fn = _generate_transform_fn_for_map_batches(fn)
+
+        partition_transform_fns = [
+            BlockMapTransformFn(
+                block_fn,
+                # NOTE: Disable block-shaping to produce blocks as is
+                disable_block_shaping=True,
+            ),
+        ]
+        partition_map_transformer = MapTransformer(
+            partition_transform_fns,
+            init_fn=init_fn,
+        )
+
+        partition_map_operator = ActorPoolMapOperator(
+            partition_map_transformer,
+            input_physical_dag,
+            data_context,
+            name=f"Partition({uri_column_names_str})",
+            # NOTE: Partition actor doesn't use the user-provided `ray_remote_args`
+            #       since those only apply to the actual download tasks. Partitioning is
+            #       a lightweight internal operation that doesn't need custom resource
+            #       requirements.
+            ray_remote_args=None,
+            compute_strategy=partition_compute,  # Use actor-based compute for callable class
+            # NOTE: We set `_generator_backpressure_num_objects` to -1 to unblock
+            #       backpressure since partitioning is extremely fast. Without this, the
+            #       partition actor gets bottlenecked by the Ray Data scheduler, which
+            #       can prevent Ray Data from launching enough download tasks.
+            ray_actor_task_remote_args={"_generator_backpressure_num_objects": -1},
+        )
+
+    if use_obstore_path:
+        download_fn = download_bytes_async
+        logger.debug("Using obstore async download path.")
+    else:
+        download_fn = download_bytes_threaded
+        # The "obstore not installed" warning is only relevant when obstore is
+        # missing entirely. When obstore is available but the filesystem can't
+        # be routed through it, _plan_obstore_routing already logged the reason
+        # (a WARNING for fsspec-S3-unextractable, DEBUG otherwise).
+        if not OBSTORE_AVAILABLE:
+            _log_fallback_warning()
+
+    fn, init_fn = _get_udf(
+        download_fn,
+        (uri_column_names, output_bytes_column_names, data_context, filesystem),
+        {},
+        None,
+        None,
+        None,
+    )
+
+    download_transform_fn = _generate_transform_fn_for_map_batches(fn)
+    transform_fns = [
+        BlockMapTransformFn(
+            download_transform_fn,
+            output_block_size_option=OutputBlockSizeOption.of(
+                target_max_block_size=data_context.target_max_block_size
+            ),
+        ),
+    ]
+
+    download_compute = TaskPoolStrategy()
+    download_map_transformer = MapTransformer(
+        transform_fns,
+        init_fn=init_fn,
+    )
+
+    download_map_operator = MapOperator.create(
+        download_map_transformer,
+        partition_map_operator if partition_map_operator else input_physical_dag,
+        data_context,
+        name=f"Download({uri_column_names_str})",
+        compute_strategy=download_compute,
+        ray_remote_args=ray_remote_args,
+    )
+
+    return download_map_operator
+
+
+def download_bytes_threaded(
+    block: pa.Table,
+    uri_column_names: List[str],
+    output_bytes_column_names: List[str],
+    data_context: DataContext,
+    filesystem: Optional["pa.fs.FileSystem"] = None,
+) -> Iterator[pa.Table]:
+    """Optimized version that uses make_async_gen for concurrent downloads.
+
+    Supports downloading from multiple URI columns in a single operation.
+
+    Args:
+        block: Input PyArrow table containing URI columns.
+        uri_column_names: Names of columns containing URIs to download.
+        output_bytes_column_names: Names for the output columns containing downloaded bytes.
+        data_context: Ray Data context for configuration.
+        filesystem: PyArrow filesystem to use for reading remote files.
+            If None, the filesystem is auto-detected from the path scheme.
+
+    Yields:
+        pa.Table: PyArrow table with the downloaded bytes added as new columns.
+    """
+    if not isinstance(block, pa.Table):
+        block = BlockAccessor.for_block(block).to_arrow()
+
+    output_block = block
+
+    # Download each URI column and add it to the output block
+    for uri_column_name, output_bytes_column_name in zip(
+        uri_column_names, output_bytes_column_names
+    ):
+        # Extract URIs from PyArrow table
+        uris = output_block.column(uri_column_name).to_pylist()
+
+        if len(uris) == 0:
+            continue
+
+        # Resolve the filesystem once before spawning workers; otherwise each
+        # worker infers its own S3FileSystem and fires a duplicate IMDS
+        # credential fetch. Normalize fsspec inputs so RetryingPyFileSystem.wrap
+        # can forward open_input_stream.
+        resolved_fs = _validate_and_wrap_filesystem(filesystem)
+        if resolved_fs is None:
+            for probe_uri in uris:
+                if probe_uri is None:
+                    continue
+                try:
+                    paths, candidate_fs = _resolve_paths_and_filesystem(probe_uri, None)
+                except Exception as e:
+                    logger.debug(f"Could not infer filesystem from '{probe_uri}': {e}")
+                    continue
+                # Skip results that drop the URI (([], ...)) or yield no FS.
+                if paths and candidate_fs is not None:
+                    resolved_fs = candidate_fs
+                    break
+
+        if resolved_fs is None:
+            # No URI resolved a filesystem; workers would only repeat the same
+            # failed inference. Yield None for every row and skip the pool.
+            logger.warning(
+                "Could not resolve a filesystem from any URI in column "
+                f"{uri_column_name!r} ({len(uris)} URIs). Yielding None for "
+                "all rows."
+            )
+            output_block = output_block.add_column(
+                len(output_block.column_names),
+                output_bytes_column_name,
+                pa.array([None] * len(uris), type=pa.binary()),
+            )
+            continue
+
+        wrapped_fs = RetryingPyFileSystem.wrap(
+            resolved_fs, retryable_errors=data_context.retried_io_errors
+        )
+
+        def load_uri_bytes(
+            uri_iterator,
+            wrapped_fs=wrapped_fs,
+            resolved_fs=resolved_fs,
+            uri_column_name=uri_column_name,
+        ):
+            """Download bytes for each URI using the pre-resolved filesystem."""
+            for uri in uri_iterator:
+                read_bytes = None
+                try:
+                    if uri is None:
+                        continue
+                    # Normalize the path only; FS is supplied so no network I/O.
+                    resolved_paths, _ = _resolve_paths_and_filesystem(
+                        uri, filesystem=resolved_fs
+                    )
+                    resolved_path = resolved_paths[0] if resolved_paths else None
+                    if resolved_path is None:
+                        continue
+                    with wrapped_fs.open_input_stream(resolved_path) as f:
+                        read_bytes = f.read()
+                except OSError as e:
+                    logger.debug(
+                        f"OSError reading uri '{uri}' for column '{uri_column_name}': {e}"
+                    )
+                except Exception as e:
+                    # Catch unexpected errors like pyarrow.lib.ArrowInvalid caused by an invalid uri like
+                    # `foo://bar` to avoid failing because of one invalid uri.
+                    logger.warning(
+                        f"Unexpected error reading uri '{uri}' for column '{uri_column_name}': {e}"
+                    )
+                finally:
+                    yield read_bytes
+
+        # Use make_async_gen to resolve and download URI bytes concurrently
+        # preserve_ordering=True ensures results are returned in the same order as input URIs
+        uri_bytes = list(
+            make_async_gen(
+                base_iterator=iter(uris),
+                fn=load_uri_bytes,
+                preserve_ordering=True,
+                num_workers=URI_DOWNLOAD_MAX_WORKERS,
+            )
+        )
+
+        # Add the new column to the PyArrow table
+        output_block = output_block.add_column(
+            len(output_block.column_names),
+            output_bytes_column_name,
+            pa.array(uri_bytes),
+        )
+
+    yield from _iter_arrow_table_for_target_max_block_size(
+        output_block, data_context.target_max_block_size
+    )

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

@@ -0,0 +1,940 @@
+from __future__ import annotations
+
+import ast
+import logging
+import operator
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.dataset as ds
+
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.logical.rules.projection_pushdown import (
+    _extract_input_columns_renaming_mapping,
+)
+from ray.data.block import Block, BlockAccessor, BlockColumn, BlockType
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    DownloadExpr,
+    Expr,
+    LiteralExpr,
+    MonotonicallyIncreasingIdExpr,
+    Operation,
+    RandomExpr,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+    UUIDExpr,
+    _ExprVisitor,
+    col,
+    is_rename_expr,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _pa_is_in(left: Any, right: Any) -> Any:
+    if not isinstance(right, (pa.Array, pa.ChunkedArray)):
+        right = pa.array(right.as_py() if isinstance(right, pa.Scalar) else right)
+    return pc.is_in(left, right)
+
+
+_PANDAS_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
+    Operation.ADD: operator.add,
+    Operation.SUB: operator.sub,
+    Operation.MUL: operator.mul,
+    Operation.DIV: operator.truediv,
+    Operation.MOD: operator.mod,
+    Operation.FLOORDIV: operator.floordiv,
+    Operation.GT: operator.gt,
+    Operation.LT: operator.lt,
+    Operation.GE: operator.ge,
+    Operation.LE: operator.le,
+    Operation.EQ: operator.eq,
+    Operation.NE: operator.ne,
+    Operation.AND: operator.and_,
+    Operation.OR: operator.or_,
+    Operation.NOT: operator.invert,
+    Operation.IS_NULL: pd.isna,
+    Operation.IS_NOT_NULL: pd.notna,
+    Operation.IN: lambda left, right: left.isin(right),
+    Operation.NOT_IN: lambda left, right: ~left.isin(right),
+}
+
+
+def _is_pa_string_type(t: pa.DataType) -> bool:
+    return pa.types.is_string(t) or pa.types.is_large_string(t)
+
+
+def _is_pa_string_like(x: Union[pa.Array, pa.ChunkedArray]) -> bool:
+    t = x.type
+    if pa.types.is_dictionary(t):
+        t = t.value_type
+    return _is_pa_string_type(t)
+
+
+def _pa_decode_dict_string_array(x: Union[pa.Array, pa.ChunkedArray]) -> Any:
+    """Convert Arrow dictionary-encoded string arrays to regular string arrays.
+
+    Dictionary encoding stores strings as indices into a dictionary of unique values.
+    This function converts them back to regular string arrays for string operations.
+
+    Example:
+        # Input: pa.array(['a', 'b']).dictionary_encode()
+        #   -- dictionary: ["a", "b"]
+        #   -- indices: [0, 1]
+        # Output: regular string array ["a", "b"]
+    Args:
+        x: The input array to convert.
+    Returns:
+        The converted string array.
+    """
+    if pa.types.is_dictionary(x.type) and _is_pa_string_type(x.type.value_type):
+        return pc.cast(x, pa.string())
+    return x
+
+
+def _to_pa_string_input(x: Any) -> Any:
+    if isinstance(x, str):
+        return pa.scalar(x)
+    if isinstance(x, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(x):
+        return _pa_decode_dict_string_array(x)
+    actual_type = (
+        str(x.type) if isinstance(x, (pa.Array, pa.ChunkedArray)) else type(x).__name__
+    )
+    raise TypeError(
+        "Expected string or string-like pyarrow Array/ChunkedArray for string "
+        f"concatenation, got {actual_type}."
+    )
+
+
+def _pa_add_or_concat(left: Any, right: Any) -> Any:
+    if isinstance(left, pa.Scalar):
+        left = left.as_py()
+    if isinstance(right, pa.Scalar):
+        right = right.as_py()
+    # If either side is string-like, perform string concatenation.
+    if (
+        isinstance(left, str)
+        or isinstance(right, str)
+        or (isinstance(left, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(left))
+        or (
+            isinstance(right, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(right)
+        )
+    ):
+        left_input = _to_pa_string_input(left)
+        right_input = _to_pa_string_input(right)
+        return pc.binary_join_element_wise(left_input, right_input, "")
+    return pc.add(left, right)
+
+
+_ARROW_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
+    Operation.ADD: _pa_add_or_concat,
+    Operation.SUB: pc.subtract,
+    Operation.MUL: pc.multiply,
+    Operation.DIV: pc.divide,
+    Operation.MOD: lambda left, right: (
+        # Modulo op is essentially:
+        #   r = N - floor(N/M) * M
+        pc.subtract(left, pc.multiply(pc.floor(pc.divide(left, right)), right))
+    ),
+    Operation.FLOORDIV: lambda left, right: pc.floor(pc.divide(left, right)),
+    Operation.GT: pc.greater,
+    Operation.LT: pc.less,
+    Operation.GE: pc.greater_equal,
+    Operation.LE: pc.less_equal,
+    Operation.EQ: pc.equal,
+    Operation.NE: pc.not_equal,
+    Operation.AND: pc.and_kleene,
+    Operation.OR: pc.or_kleene,
+    Operation.NOT: pc.invert,
+    Operation.IS_NULL: pc.is_null,
+    Operation.IS_NOT_NULL: pc.is_valid,
+    Operation.IN: _pa_is_in,
+    Operation.NOT_IN: lambda left, right: pc.invert(_pa_is_in(left, right)),
+}
+
+
+# NOTE: (srinathk) There are 3 distinct stages of handling passed in exprs:
+# 1. Parsing it (as text)
+# 2. Resolving unbound names (to schema)
+# 3. Converting resolved expressions to PA ones
+# Need to break up the abstraction provided by ExpressionEvaluator.
+
+ScalarType = TypeVar("ScalarType")
+
+
+class ExpressionEvaluator:
+    @staticmethod
+    def get_filters(expression: str) -> ds.Expression:
+        """Parse and evaluate the expression to generate a filter condition.
+
+        Args:
+            expression: A string representing the filter expression to parse.
+
+        Returns:
+            A PyArrow compute expression for filtering data.
+
+        """
+        try:
+            tree = ast.parse(expression, mode="eval")
+            return _ConvertToArrowExpressionVisitor().visit(tree.body)
+        except SyntaxError as e:
+            raise ValueError(f"Invalid syntax in the expression: {expression}") from e
+        except Exception as e:
+            logger.exception(f"Error processing expression: {e}")
+            raise
+
+    @staticmethod
+    def parse_native_expression(expression: str) -> "Expr":
+        """Parse and evaluate the expression to generate a Ray Data expression.
+
+        Args:
+            expression: A string representing the filter expression to parse.
+
+        Returns:
+            A Ray Data Expr object for filtering data.
+
+        """
+        try:
+            tree = ast.parse(expression, mode="eval")
+            return _ConvertToNativeExpressionVisitor().visit(tree.body)
+        except SyntaxError as e:
+            raise ValueError(f"Invalid syntax in the expression: {expression}") from e
+        except Exception as e:
+            logger.exception(f"Error processing expression: {e}")
+            raise
+
+
+class _ConvertToArrowExpressionVisitor(ast.NodeVisitor):
+    # TODO: Deprecate this visitor after we remove string support in filter API.
+    def visit_Compare(self, node: ast.Compare) -> ds.Expression:
+        """Handle comparison operations (e.g., a == b, a < b, a in b).
+
+        Args:
+            node: The AST node representing a comparison operation.
+
+        Returns:
+            An expression representing the comparison.
+        """
+        # Handle left operand
+        # TODO Validate columns
+        if isinstance(node.left, ast.Attribute):
+            # Visit and handle attributes
+            left_expr = self.visit(node.left)
+        elif isinstance(node.left, ast.Name):
+            # Treat as a simple field
+            left_expr = self.visit(node.left)
+        elif isinstance(node.left, ast.Constant):
+            # Constant values are used directly
+            left_expr = node.left.value
+        else:
+            raise ValueError(f"Unsupported left operand type: {type(node.left)}")
+
+        comparators = [self.visit(comp) for comp in node.comparators]
+
+        op = node.ops[0]
+        if isinstance(op, ast.In):
+            return pc.is_in(left_expr, comparators[0])
+        elif isinstance(op, ast.NotIn):
+            return ~pc.is_in(left_expr, comparators[0])
+        elif isinstance(op, ast.Eq):
+            return left_expr == comparators[0]
+        elif isinstance(op, ast.NotEq):
+            return left_expr != comparators[0]
+        elif isinstance(op, ast.Lt):
+            return left_expr < comparators[0]
+        elif isinstance(op, ast.LtE):
+            return left_expr <= comparators[0]
+        elif isinstance(op, ast.Gt):
+            return left_expr > comparators[0]
+        elif isinstance(op, ast.GtE):
+            return left_expr >= comparators[0]
+        else:
+            raise ValueError(f"Unsupported operator type: {op}")
+
+    def visit_BoolOp(self, node: ast.BoolOp) -> ds.Expression:
+        """Handle logical operations (e.g., a and b, a or b).
+
+        Args:
+            node: The AST node representing a boolean operation.
+
+        Returns:
+            An expression representing the logical operation.
+        """
+        conditions = [self.visit(value) for value in node.values]
+        combined_expr = conditions[0]
+
+        for condition in conditions[1:]:
+            if isinstance(node.op, ast.And):
+                # Combine conditions with logical AND
+                combined_expr &= condition
+            elif isinstance(node.op, ast.Or):
+                # Combine conditions with logical OR
+                combined_expr |= condition
+            else:
+                raise ValueError(
+                    f"Unsupported logical operator: {type(node.op).__name__}"
+                )
+
+        return combined_expr
+
+    def visit_Name(self, node: ast.Name) -> ds.Expression:
+        """Handle variable (name) nodes and return them as pa.dataset.Expression.
+
+        Even if the name contains periods, it's treated as a single string.
+
+        Args:
+            node: The AST node representing a variable.
+
+        Returns:
+            The variable wrapped as a pa.dataset.Expression.
+        """
+        # Directly use the field name as a string (even if it contains periods)
+        field_name = node.id
+        return pc.field(field_name)
+
+    def visit_Attribute(self, node: ast.Attribute) -> object:
+        """Handle attribute access (e.g., np.nan).
+
+        Args:
+            node: The AST node representing an attribute access.
+
+        Returns:
+            object: The attribute value.
+
+        Raises:
+            ValueError: If the attribute is unsupported.
+        """
+        # Recursively visit the left side (base object or previous attribute)
+        if isinstance(node.value, ast.Attribute):
+            # If the value is an attribute, recursively resolve it
+            left_expr = self.visit(node.value)
+            return pc.field(f"{left_expr}.{node.attr}")
+
+        elif isinstance(node.value, ast.Name):
+            # If the value is a name (e.g., "foo"), we can directly return the field
+            left_name = node.value.id  # The base name, e.g., "foo"
+            return pc.field(f"{left_name}.{node.attr}")
+
+        raise ValueError(f"Unsupported attribute: {node.attr}")
+
+    def visit_List(self, node: ast.List) -> ds.Expression:
+        """Handle list literals.
+
+        Args:
+            node: The AST node representing a list.
+
+        Returns:
+            The list of elements wrapped as a pa.dataset.Expression.
+        """
+        elements = [self.visit(elt) for elt in node.elts]
+        return pa.array(elements)
+
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> ds.Expression:
+        """Handle case where comparator is UnaryOP (e.g., a == -1).
+
+        AST for this expression will be Compare(left=Name(id='a'), ops=[Eq()],
+        comparators=[UnaryOp(op=USub(), operand=Constant(value=1))])
+
+        Args:
+            node: The constant value.
+
+        Returns:
+            A PyArrow scalar expression representing the unary operation result.
+        """
+        op = node.op
+        if isinstance(op, ast.USub):
+            return pc.scalar(-node.operand.value)
+        else:
+            raise ValueError(f"Unsupported unary operator: {op}")
+
+    # TODO (srinathk) Note that visit_Constant does not return pa.dataset.Expression
+    # because to support function in() which takes in a List, the elements in the List
+    # needs to values instead of pa.dataset.Expression per pyarrow.dataset.Expression
+    # specification. May be down the road, we can update it as Arrow relaxes this
+    # constraint.
+    def visit_Constant(self, node: ast.Constant) -> object:
+        """Handle constant values (e.g., numbers, strings).
+
+        Args:
+            node: The AST node representing a constant value.
+
+        Returns:
+            object: The constant value itself (e.g., number, string, or boolean).
+        """
+        return node.value  # Return the constant value directly.
+
+    def visit_Call(self, node: ast.Call) -> ds.Expression:
+        """Handle function calls (e.g., is_nan(a), is_valid(b)).
+
+        Args:
+            node: The AST node representing a function call.
+
+        Returns:
+            The corresponding expression based on the function called.
+
+        Raises:
+            ValueError: If the function is unsupported or has incorrect arguments.
+        """
+        func_name = node.func.id
+        function_map = {
+            "is_nan": lambda arg: arg.is_nan(),
+            "is_null": lambda arg, nan_is_null=False: arg.is_null(
+                nan_is_null=nan_is_null
+            ),
+            "is_valid": lambda arg: arg.is_valid(),
+            "is_in": lambda arg1, arg2: pc.is_in(arg1, arg2),
+        }
+
+        if func_name in function_map:
+            # Visit all arguments of the function call
+            args = [self.visit(arg) for arg in node.args]
+            # Handle the "is_null" function with one or two arguments
+            if func_name == "is_null":
+                if len(args) == 1:
+                    return function_map[func_name](args[0])
+                elif len(args) == 2:
+                    return function_map[func_name](args[0], args[1])
+                else:
+                    raise ValueError("is_null function requires one or two arguments.")
+            # Handle the "is_in" function with exactly two arguments
+            elif func_name == "is_in" and len(args) != 2:
+                raise ValueError("is_in function requires two arguments.")
+            # Ensure the function has one argument (for functions like is_valid)
+            elif func_name != "is_in" and len(args) != 1:
+                raise ValueError(f"{func_name} function requires exactly one argument.")
+            # Call the corresponding function with the arguments
+            return function_map[func_name](*args)
+        else:
+            raise ValueError(f"Unsupported function: {func_name}")
+
+
+class _ConvertToNativeExpressionVisitor(ast.NodeVisitor):
+    """AST visitor that converts string expressions to Ray Data expressions."""
+
+    def visit_Compare(self, node: ast.Compare) -> "Expr":
+        """Handle comparison operations (e.g., a == b, a < b, a in b)."""
+        from ray.data.expressions import BinaryExpr, Operation
+
+        if len(node.ops) != 1 or len(node.comparators) != 1:
+            raise ValueError("Only simple binary comparisons are supported")
+
+        left = self.visit(node.left)
+        right = self.visit(node.comparators[0])
+        op = node.ops[0]
+
+        # Map AST comparison operators to Ray Data operations
+        op_map = {
+            ast.Eq: Operation.EQ,
+            ast.NotEq: Operation.NE,
+            ast.Lt: Operation.LT,
+            ast.LtE: Operation.LE,
+            ast.Gt: Operation.GT,
+            ast.GtE: Operation.GE,
+            ast.In: Operation.IN,
+            ast.NotIn: Operation.NOT_IN,
+        }
+
+        if type(op) not in op_map:
+            raise ValueError(f"Unsupported comparison operator: {type(op).__name__}")
+
+        return BinaryExpr(op_map[type(op)], left, right)
+
+    def visit_BoolOp(self, node: ast.BoolOp) -> "Expr":
+        """Handle logical operations (e.g., a and b, a or b)."""
+        from ray.data.expressions import BinaryExpr, Operation
+
+        conditions = [self.visit(value) for value in node.values]
+        combined_expr = conditions[0]
+
+        for condition in conditions[1:]:
+            if isinstance(node.op, ast.And):
+                combined_expr = BinaryExpr(Operation.AND, combined_expr, condition)
+            elif isinstance(node.op, ast.Or):
+                combined_expr = BinaryExpr(Operation.OR, combined_expr, condition)
+            else:
+                raise ValueError(
+                    f"Unsupported logical operator: {type(node.op).__name__}"
+                )
+
+        return combined_expr
+
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> "Expr":
+        """Handle unary operations (e.g., not a, -5)."""
+        from ray.data.expressions import Operation, UnaryExpr, lit
+
+        if isinstance(node.op, ast.Not):
+            operand = self.visit(node.operand)
+            return UnaryExpr(Operation.NOT, operand)
+        elif isinstance(node.op, ast.USub):
+            operand = self.visit(node.operand)
+            return operand * lit(-1)
+        else:
+            raise ValueError(f"Unsupported unary operator: {type(node.op).__name__}")
+
+    def visit_Name(self, node: ast.Name) -> "Expr":
+        """Handle variable names (column references)."""
+        from ray.data.expressions import col
+
+        return col(node.id)
+
+    def visit_Constant(self, node: ast.Constant) -> "Expr":
+        """Handle constant values (numbers, strings, booleans)."""
+        from ray.data.expressions import lit
+
+        return lit(node.value)
+
+    def visit_List(self, node: ast.List) -> "Expr":
+        """Handle list literals."""
+        from ray.data.expressions import LiteralExpr, lit
+
+        # Visit all elements first
+        visited_elements = [self.visit(elt) for elt in node.elts]
+
+        # Try to extract constant values for literal list
+        elements = []
+        for elem in visited_elements:
+            if isinstance(elem, LiteralExpr):
+                elements.append(elem.value)
+            else:
+                # For compatibility with Arrow visitor, we need to support non-literals
+                # but Ray Data expressions may have limitations here
+                raise ValueError(
+                    "List contains non-constant expressions. Ray Data expressions "
+                    "currently only support lists of constant values."
+                )
+
+        return lit(elements)
+
+    def visit_Attribute(self, node: ast.Attribute) -> "Expr":
+        """Handle attribute access (e.g., for nested column names)."""
+        from ray.data.expressions import col
+
+        # For nested column names like "user.age", combine them with dots
+        if isinstance(node.value, ast.Name):
+            return col(f"{node.value.id}.{node.attr}")
+        elif isinstance(node.value, ast.Attribute):
+            # Recursively handle nested attributes
+            left_expr = self.visit(node.value)
+            if isinstance(left_expr, ColumnExpr):
+                return col(f"{left_expr._name}.{node.attr}")
+
+        raise ValueError(
+            f"Unsupported attribute access: {node.attr}. Node details: {ast.dump(node)}"
+        )
+
+    def visit_Call(self, node: ast.Call) -> "Expr":
+        """Handle function calls for operations like is_null, is_not_null, is_nan, random."""
+        from ray.data.expressions import (
+            BinaryExpr,
+            Operation,
+            UnaryExpr,
+        )
+
+        func_name = node.func.id if isinstance(node.func, ast.Name) else str(node.func)
+
+        if func_name == "is_null":
+            if len(node.args) != 1:
+                raise ValueError("is_null() expects exactly one argument")
+            operand = self.visit(node.args[0])
+            return UnaryExpr(Operation.IS_NULL, operand)
+        # Adding this conditional to keep it consistent with the current implementation,
+        # of carrying Pyarrow's semantic of `is_valid`
+        elif func_name == "is_valid" or func_name == "is_not_null":
+            if len(node.args) != 1:
+                raise ValueError(f"{func_name}() expects exactly one argument")
+            operand = self.visit(node.args[0])
+            return UnaryExpr(Operation.IS_NOT_NULL, operand)
+        elif func_name == "is_nan":
+            if len(node.args) != 1:
+                raise ValueError("is_nan() expects exactly one argument")
+            operand = self.visit(node.args[0])
+            # Use x != x pattern for NaN detection (NaN != NaN is True)
+            return BinaryExpr(Operation.NE, operand, operand)
+        elif func_name == "is_in":
+            if len(node.args) != 2:
+                raise ValueError("is_in() expects exactly two arguments")
+            left = self.visit(node.args[0])
+            right = self.visit(node.args[1])
+            return BinaryExpr(Operation.IN, left, right)
+        elif func_name == "random":
+            raise ValueError(
+                "random() is not supported in string expressions. "
+                "String expressions are deprecated. Please use the expression API instead: "
+                "from ray.data.expressions import random; ds.filter(expr=(random(seed=42)>0.5))"
+            )
+        elif func_name == "uuid":
+            raise ValueError(
+                "uuid() is not supported in string expressions. "
+                "String expressions are deprecated. Please use the expression API instead: "
+                "ds.filter(expr=uuid().str.starts_with('a'))"
+            )
+        else:
+            raise ValueError(f"Unsupported function: {func_name}")
+
+
+class NativeExpressionEvaluator(_ExprVisitor[Union[BlockColumn, ScalarType]]):
+    """Visitor-based expression evaluator that uses Block and BlockColumns
+
+    This evaluator implements the visitor pattern to traverse expression trees
+    and evaluate them against Block data structures. It maintains operation
+    mappings in shared state and returns consistent BlockColumn types.
+    """
+
+    def __init__(self, block: Block):
+        """Initialize the evaluator with a block and operation mappings.
+
+        Args:
+            block: The Block to evaluate expressions against.
+        """
+        self.block = block
+        self.block_accessor = BlockAccessor.for_block(block)
+
+        # Use BlockAccessor to determine operation mappings
+        block_type = self.block_accessor.block_type()
+        if block_type == BlockType.PANDAS:
+            self.ops = _PANDAS_EXPR_OPS_MAP
+        elif block_type == BlockType.ARROW:
+            self.ops = _ARROW_EXPR_OPS_MAP
+        else:
+            raise TypeError(f"Unsupported block type: {block_type}")
+
+    def visit_column(self, expr: ColumnExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a column expression and return the column data.
+
+        Args:
+            expr: The column expression.
+
+        Returns:
+            The column data as a BlockColumn.
+        """
+        return self.block[expr.name]
+
+    def visit_literal(self, expr: LiteralExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a literal expression and return the literal value.
+
+        Args:
+            expr: The literal expression.
+
+        Returns:
+            The literal value.
+        """
+        # Given that expressions support pandas blocks, we need to return the value as is.
+        # Pandas has multiple dtype_backends, so there's no guarantee on the return type.
+        return expr.value
+
+    def visit_binary(self, expr: BinaryExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a binary expression and return the result of the operation.
+
+        Args:
+            expr: The binary expression.
+
+        Returns:
+            The result of the binary operation as a BlockColumn.
+        """
+        left_result = self.visit(expr.left)
+        right_result = self.visit(expr.right)
+
+        return self.ops[expr.op](left_result, right_result)
+
+    def visit_unary(self, expr: UnaryExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a unary expression and return the result of the operation.
+
+        Args:
+            expr: The unary expression.
+
+        Returns:
+            The result of the unary operation as a BlockColumn.
+        """
+        operand_result = self.visit(expr.operand)
+        return self.ops[expr.op](operand_result)
+
+    def visit_udf(self, expr: UDFExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a UDF expression and return the result of the function call.
+
+        Args:
+            expr: The UDF expression.
+
+        Returns:
+            The result of the UDF call as a BlockColumn.
+        """
+        args = [self.visit(arg) for arg in expr.args]
+        kwargs = {k: self.visit(v) for k, v in expr.kwargs.items()}
+
+        result = expr.fn(*args, **kwargs)
+
+        if not isinstance(result, (pd.Series, np.ndarray, pa.Array, pa.ChunkedArray)):
+            function_name = expr.fn.__name__
+            raise TypeError(
+                f"UDF '{function_name}' returned invalid type {type(result).__name__}. "
+                f"Expected type (pandas.Series, numpy.ndarray, pyarrow.Array, "
+                f"pyarrow.ChunkedArray)"
+            )
+
+        return result
+
+    def visit_alias(self, expr: AliasExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit an alias expression and return the renamed result.
+
+        Args:
+            expr: The alias expression.
+
+        Returns:
+            A Block with the data from the inner expression.
+        """
+        # Evaluate the inner expression
+        return self.visit(expr.expr)
+
+    def visit_star(self, expr: StarExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a star expression.
+
+        Args:
+            expr: The star expression.
+
+        Returns:
+            TypeError: StarExpr cannot be evaluated as a regular expression.
+        """
+        # star() should not be evaluated directly - it's handled at Project level
+        raise TypeError(
+            "StarExpr cannot be evaluated as a regular expression. "
+            "It should only be used in Project operations."
+        )
+
+    def visit_download(self, expr: DownloadExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a download expression.
+
+        Args:
+            expr: The download expression.
+
+        Returns:
+            TypeError: DownloadExpr evaluation not yet implemented.
+        """
+        raise TypeError(
+            "DownloadExpr evaluation is not yet implemented in NativeExpressionEvaluator."
+        )
+
+    def visit_monotonically_increasing_id(
+        self, expr: MonotonicallyIncreasingIdExpr
+    ) -> Union[BlockColumn, ScalarType]:
+        """Visit a monotonically_increasing_id expression.
+
+        Args:
+            expr: The monotonically_increasing_id expression.
+
+        Returns:
+            The result of the monotonically_increasing_id expression as a BlockColumn.
+        """
+        ctx = TaskContext.get_current()
+        assert (
+            ctx is not None
+        ), "TaskContext is required for monotonically_increasing_id()"
+
+        # Key the counter by expression instance ID so that multiple expressions
+        # in the same projection will have isolated row count state.
+        # This is required because a single task may process multiple blocks if
+        # the upstream data source does not compress the data into a single block.
+        counter_key = f"_mono_id_{expr._instance_id}_counter"
+
+        start_idx = ctx.kwargs.get(counter_key, 0)
+        num_rows = self.block_accessor.num_rows()
+        end_idx = start_idx + num_rows
+        ctx.kwargs[counter_key] = end_idx
+
+        # int64 (signed): upper 30 bits = task ID, lower 33 bits = row number.
+        # Note end_idx is an exclusive upper bound, as the max row ID is end_idx - 1.
+        ROW_BITS = 33
+        TASK_BITS = 30
+        if end_idx > (1 << ROW_BITS):
+            raise ValueError(
+                f"Cannot generate monotonically increasing IDs: row count for this task exceeds the maximum allowed value of {(1 << ROW_BITS) - 1}"
+            )
+        if ctx.task_idx >= (1 << TASK_BITS):
+            raise ValueError(
+                f"Cannot generate monotonically increasing IDs: number of tasks exceeds the maximum allowed value of {(1 << TASK_BITS) - 1}"
+            )
+
+        partition_mask = ctx.task_idx << ROW_BITS
+        ids = partition_mask + np.arange(start_idx, end_idx, dtype=np.int64)
+
+        block_type = self.block_accessor.block_type()
+        if block_type == BlockType.PANDAS:
+            return pd.Series(ids)
+        elif block_type == BlockType.ARROW:
+            return pa.array(ids)
+        else:
+            raise TypeError(f"Unsupported block type: {block_type}")
+
+    def visit_random(self, expr: RandomExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a random expression and return the result of the operation.
+
+        Args:
+            expr: The random expression.
+
+        Returns:
+            The result of the random operation as a BlockColumn.
+        """
+        from ray.data._internal.planner.plan_expression.synthetic_impl import (
+            eval_random,
+        )
+
+        return eval_random(
+            self.block_accessor.num_rows(),
+            self.block_accessor.block_type(),
+            seed=expr.seed,
+            reseed_after_execution=expr.reseed_after_execution,
+            instance_id=expr._instance_id,
+        )
+
+    def visit_uuid(self, expr: UUIDExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a uuid expression and return the result of the operation.
+
+        Args:
+            expr: The uuid expression.
+
+        Returns:
+            The result of the uuid operation as a BlockColumn.
+        """
+        from ray.data._internal.planner.plan_expression.synthetic_impl import eval_uuid
+
+        return eval_uuid(
+            self.block_accessor.num_rows(), self.block_accessor.block_type()
+        )
+
+
+def eval_expr(expr: Expr, block: Block) -> Union[BlockColumn, ScalarType]:
+    """Evaluate an expression against a block using the visitor pattern.
+
+    Args:
+        expr: The expression to evaluate.
+        block: The Block to evaluate against.
+
+    Returns:
+        The evaluated result as a BlockColumn or a scalar value.
+    """
+    evaluator = NativeExpressionEvaluator(block)
+    return evaluator.visit(expr)
+
+
+def _eval_projection_without_cse(projection_exprs: List[Expr], block: Block) -> Block:
+    """
+    Evaluate a projection (list of expressions) against a block.
+
+    Handles projection semantics including:
+    - Empty projections
+    - Star() expressions for preserving existing columns
+    - Rename detection
+    - Column ordering
+
+    Args:
+        projection_exprs: List of expressions to evaluate (may include StarExpr)
+        block: The block to project
+
+    Returns:
+        A new block with the projected schema
+    """
+    block_accessor = BlockAccessor.for_block(block)
+
+    # Skip projection only for schema-less empty blocks.
+    if block_accessor.num_rows() == 0 and len(block_accessor.column_names()) == 0:
+        return block
+
+    # Handle simple cases early.
+    if len(projection_exprs) == 0:
+        return block_accessor.select([])
+
+    input_column_names = list(block_accessor.column_names())
+    # Collect input column rename map from the projection list
+    input_column_rename_map = _extract_input_columns_renaming_mapping(projection_exprs)
+
+    # Expand star expr (if any). ``Project.__post_init__`` eagerly expands
+    # ``StarExpr`` to explicit ``col()`` refs whenever the
+    # input schema is known, so this runtime branch is hit only on the
+    # UDF-fallback path (Project on top of an opaque-schema input).
+    if isinstance(projection_exprs[0], StarExpr):
+        # Bucket the trailing exprs: rename ``AliasExpr``s of an input
+        # column get placed into the original column's position (so the
+        # output preserves on-disk column order); anything else (e.g.
+        # ``with_column`` computed expressions) is appended afterwards.
+        rename_exprs_by_source: Dict[str, Expr] = {}
+        extra_exprs: List[Expr] = []
+        for expr in projection_exprs[1:]:
+            # e.g. ``col(source)._rename(new_name)`` — bucket by ``source`` for column order.
+            # ``rename_exprs_by_source``: input column name -> that rename ``AliasExpr``.
+            if is_rename_expr(expr) and expr.expr.name in input_column_rename_map:
+                rename_exprs_by_source[expr.expr.name] = expr
+            else:
+                extra_exprs.append(expr)
+
+        ordered_exprs: List[Expr] = []
+        for c in input_column_names:
+            if c in rename_exprs_by_source:
+                ordered_exprs.append(rename_exprs_by_source.pop(c))
+            elif c not in input_column_rename_map:
+                ordered_exprs.append(col(c))
+
+        # Any rename whose source column isn't in the block falls through to
+        # ``extra_exprs`` so evaluation raises a "column not found" error
+        # instead of silently dropping the expression.
+        extra_exprs = list(rename_exprs_by_source.values()) + extra_exprs
+
+        projection_exprs = ordered_exprs + extra_exprs
+
+    names, output_cols = zip(*[(e.name, eval_expr(e, block)) for e in projection_exprs])
+
+    # This clumsy workaround is necessary to be able to fill in Pyarrow tables
+    # that has to be "seeded" from existing table with N rows, and couldn't be
+    # started from a truly empty table.
+    #
+    # TODO fix
+    new_block = BlockAccessor.for_block(block).fill_column("__stub__", None)
+    new_block = BlockAccessor.for_block(new_block).drop(input_column_names)
+
+    for name, output_col in zip(names, output_cols):
+        new_block = BlockAccessor.for_block(new_block).fill_column(name, output_col)
+
+    return BlockAccessor.for_block(new_block).drop(["__stub__"])
+
+
+def _drop_cse_temp_columns(block: Block, temp_columns: List[str]) -> Block:
+    block_accessor = BlockAccessor.for_block(block)
+    drop_columns = [
+        name for name in temp_columns if name in block_accessor.column_names()
+    ]
+    if not drop_columns:
+        return block
+    return block_accessor.drop(drop_columns)
+
+
+def eval_projection(
+    projection_exprs: List[Expr],
+    block: Block,
+    *,
+    common_sub_exprs: Optional[List[Expr]] = None,
+) -> Block:
+    """
+    Evaluate a projection (list of expressions) against a block.
+
+    If CSE common expressions are provided, they are evaluated first into
+    temporary columns on a working block. Visible projection expressions are
+    then evaluated against that working block.
+    """
+    if not common_sub_exprs:
+        return _eval_projection_without_cse(projection_exprs, block)
+
+    working_block = block
+    for common_expr in common_sub_exprs:
+        assert common_expr.name is not None
+        working_block = BlockAccessor.for_block(working_block).fill_column(
+            common_expr.name,
+            eval_expr(common_expr, working_block),
+        )
+
+    output_block = _eval_projection_without_cse(projection_exprs, working_block)
+    temp_columns = [expr.name for expr in common_sub_exprs]
+    return _drop_cse_temp_columns(output_block, temp_columns)

python/ray/data/_internal/planner/plan_expression/expression_visitors.py

@@ -0,0 +1,843 @@
+from collections import Counter
+from dataclasses import dataclass, replace
+from typing import Dict, Hashable, List, TypeVar
+
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    DownloadExpr,
+    Expr,
+    LiteralExpr,
+    MonotonicallyIncreasingIdExpr,
+    Operation,
+    RandomExpr,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+    UUIDExpr,
+    _CallableClassUDF,
+    _ExprVisitor,
+)
+from ray.data.util.expression_utils import (
+    _alias_fingerprint_key,
+    _binary_fingerprint_key,
+    _column_fingerprint_key,
+    _download_fingerprint_key,
+    _literal_fingerprint_key,
+    _monotonically_increasing_id_fingerprint_key,
+    _random_fingerprint_key,
+    _star_fingerprint_key,
+    _udf_fingerprint_key,
+    _unary_fingerprint_key,
+    _uuid_fingerprint_key,
+)
+
+T = TypeVar("T")
+
+# Mapping of operations to their string symbols for inline representation
+_INLINE_OP_SYMBOLS = {
+    Operation.ADD: "+",
+    Operation.SUB: "-",
+    Operation.MUL: "*",
+    Operation.DIV: "/",
+    Operation.MOD: "%",
+    Operation.FLOORDIV: "//",
+    Operation.GT: ">",
+    Operation.LT: "<",
+    Operation.GE: ">=",
+    Operation.LE: "<=",
+    Operation.EQ: "==",
+    Operation.NE: "!=",
+    Operation.AND: "&",
+    Operation.OR: "|",
+    Operation.IN: "in",
+    Operation.NOT_IN: "not in",
+}
+
+
+class _ExprVisitorBase(_ExprVisitor[None]):
+    """Base visitor that provides automatic recursive traversal.
+
+    This class extends _ExprVisitor and provides default implementations
+    for composite nodes that automatically traverse child expressions.
+    """
+
+    def visit_binary(self, expr: "BinaryExpr") -> None:
+        """Default implementation: recursively visit both operands."""
+        super().visit(expr.left)
+        super().visit(expr.right)
+
+    def visit_unary(self, expr: "UnaryExpr") -> None:
+        """Default implementation: recursively visit the operand."""
+        super().visit(expr.operand)
+
+    def visit_alias(self, expr: "AliasExpr") -> None:
+        """Default implementation: recursively visit the inner expression."""
+        super().visit(expr.expr)
+
+    def visit_udf(self, expr: "UDFExpr") -> None:
+        """Default implementation: recursively visit all arguments."""
+        for arg in expr.args:
+            super().visit(arg)
+        for value in expr.kwargs.values():
+            super().visit(value)
+
+    def visit_literal(self, expr: LiteralExpr) -> None:
+        """Visit a literal expression (no columns to collect)."""
+        pass
+
+    def visit_star(self, expr: StarExpr) -> None:
+        """Visit a star expression (no columns to collect)."""
+        pass
+
+    def visit_download(self, expr: "Expr") -> None:
+        """Visit a download expression (no columns to collect)."""
+        pass
+
+    def visit_monotonically_increasing_id(
+        self, expr: "MonotonicallyIncreasingIdExpr"
+    ) -> None:
+        """Visit a monotonically_increasing_id expression (no columns to collect)."""
+        pass
+
+    def visit_random(self, expr: "RandomExpr") -> None:
+        """Visit a synthetic expression (no columns to collect)."""
+        pass
+
+    def visit_uuid(self, expr: "UUIDExpr") -> None:
+        """Visit a uuid expression (no columns to collect)."""
+        pass
+
+
+class _ColumnReferenceCollector(_ExprVisitorBase):
+    """Visitor that collects all column references from expression trees.
+
+    Backed by a ``Counter`` so callers can take either:
+    - ``get_column_refs()`` -> ordered, de-duplicated column names, or
+    - ``get_counts()``      -> per-name reference multiplicity, counting repeats
+      *within* a single expression (``x + x`` -> ``{"x": 2}``).
+
+    ``Counter`` preserves first-insertion order, so ``get_column_refs()`` returns the
+    same ordered, de-duplicated list as a plain insertion-ordered ``dict`` would.
+    """
+
+    def __init__(self):
+        """Initialize with an empty reference counter."""
+        self._col_refs: Counter = Counter()
+
+    def get_column_refs(self) -> List[str]:
+        return list(self._col_refs.keys())
+
+    def get_counts(self) -> Counter:
+        return self._col_refs
+
+    def visit_column(self, expr: ColumnExpr) -> None:
+        """Visit a column expression and count its name.
+
+        Args:
+            expr: The column expression.
+
+        Returns:
+            None (only counts columns as a side effect).
+        """
+        self._col_refs[expr.name] += 1
+
+    def visit_alias(self, expr: AliasExpr) -> None:
+        """Visit an alias expression and collect from its inner expression.
+
+        Args:
+            expr: The alias expression.
+
+        Returns:
+            None (only collects columns as a side effect).
+        """
+        self.visit(expr.expr)
+
+
+class _IdempotencyVisitor(_ExprVisitor[bool]):
+    """Reports whether an expression is safe to duplicate, reorder, or move.
+
+    Returns ``True`` only when every node in the tree is idempotent. The three
+    non-idempotent leaf types (``RandomExpr``, ``UUIDExpr``,
+    ``MonotonicallyIncreasingIdExpr``) return ``False`` and propagate upward: a
+    composite is idempotent iff all of its children are.
+
+    Optimizer rules consult this (via :func:`is_idempotent`) before any rewrite that
+    would change an expression's evaluation count, row set, or position.
+    """
+
+    # --- non-idempotent leaves ---
+    def visit_random(self, expr: RandomExpr) -> bool:
+        # Conservatively non-idempotent even when seeded: CSE matches structurally and
+        # ignores ``_instance_id``, while the runtime RNG counter keys on it, so a
+        # seeded RandomExpr cannot be safely de-duplicated in general.
+        return False
+
+    def visit_uuid(self, expr: UUIDExpr) -> bool:
+        return False
+
+    def visit_monotonically_increasing_id(
+        self, expr: MonotonicallyIncreasingIdExpr
+    ) -> bool:
+        return False
+
+    # --- idempotent leaves ---
+    def visit_column(self, expr: ColumnExpr) -> bool:
+        return True
+
+    def visit_literal(self, expr: LiteralExpr) -> bool:
+        return True
+
+    def visit_star(self, expr: StarExpr) -> bool:
+        return True
+
+    def visit_download(self, expr: DownloadExpr) -> bool:
+        # ``DownloadExpr`` is a leaf with no Expr children. It is idempotent (same URI
+        # yields the same bytes); CSE avoids re-fetching it for *cost* reasons, which
+        # is a separate concern from this correctness contract.
+        return True
+
+    # --- composites: idempotent iff all children are ---
+    #
+    # Children are visited via ``child.is_idempotent()`` (not ``self.visit(child)``)
+    # so each node's result is read from / written to its per-instance cache. This
+    # keeps an all-nodes query (e.g. CSE visiting every occurrence) linear overall
+    # instead of re-walking each subtree.
+    def visit_alias(self, expr: AliasExpr) -> bool:
+        return expr.expr.is_idempotent()
+
+    def visit_unary(self, expr: UnaryExpr) -> bool:
+        return expr.operand.is_idempotent()
+
+    def visit_binary(self, expr: BinaryExpr) -> bool:
+        return expr.left.is_idempotent() and expr.right.is_idempotent()
+
+    def visit_udf(self, expr: UDFExpr) -> bool:
+        # FUTURE EXTENSION POINT: today UDFs are assumed idempotent and we only recurse
+        # into their argument expressions. When per-UDF non-determinism is supported,
+        # gate this on the UDF's declared determinism as well.
+        return all(arg.is_idempotent() for arg in expr.args) and all(
+            value.is_idempotent() for value in expr.kwargs.values()
+        )
+
+
+# Stateless singleton: ``Expr.is_idempotent`` reuses this rather than allocating a
+# visitor per node during the initial (uncached) computation.
+_IDEMPOTENCY_VISITOR = _IdempotencyVisitor()
+
+
+class _CallableClassUDFCollector(_ExprVisitorBase):
+    """Visitor that collects all callable class UDFs from expression trees.
+
+    This visitor traverses expression trees and collects _CallableClassUDF instances
+    that wrap callable classes (as opposed to regular functions).
+    """
+
+    def __init__(self):
+        """Initialize with an empty list of _CallableClassUDF instances."""
+        self._expr_udfs: List[_CallableClassUDF] = []
+
+    def get_callable_class_udfs(self) -> List[_CallableClassUDF]:
+        """Get the list of collected _CallableClassUDF instances.
+
+        Returns:
+            List of _CallableClassUDF instances that wrap callable classes.
+        """
+        return self._expr_udfs
+
+    def visit_column(self, expr: ColumnExpr) -> None:
+        """Visit a column expression (no UDFs to collect)."""
+        pass
+
+    def visit_udf(self, expr: UDFExpr) -> None:
+        """Visit a UDF expression and collect it if it's a callable class.
+
+        Args:
+            expr: The UDF expression.
+
+        Returns:
+            None (only collects UDFs as a side effect).
+        """
+        # Check if fn is an _CallableClassUDF (indicates callable class)
+        if isinstance(expr.fn, _CallableClassUDF):
+            self._expr_udfs.append(expr.fn)
+
+        # Continue visiting child expressions
+        super().visit_udf(expr)
+
+
+class _ColumnSubstitutionVisitor(_ExprVisitor[Expr]):
+    """Visitor rebinding column references in ``Expression``s.
+
+    This visitor traverses given ``Expression`` trees and substitutes column references
+    according to a provided substitution map.
+    """
+
+    def __init__(self, column_ref_substitutions: Dict[str, Expr]):
+        """Initialize with a column substitution map.
+
+        Args:
+            column_ref_substitutions: Mapping from column names to replacement expressions.
+        """
+        self._col_ref_substitutions = column_ref_substitutions
+
+    def visit_column(self, expr: ColumnExpr) -> Expr:
+        """Visit a column expression and substitute it.
+
+        Args:
+            expr: The column expression.
+
+        Returns:
+            The substituted expression or the original if no substitution exists.
+        """
+        substitution = self._col_ref_substitutions.get(expr.name)
+
+        return substitution if substitution is not None else expr
+
+    def visit_literal(self, expr: LiteralExpr) -> Expr:
+        """Visit a literal expression (no rewriting needed).
+
+        Args:
+            expr: The literal expression.
+
+        Returns:
+            The original literal expression.
+        """
+        return expr
+
+    def visit_binary(self, expr: BinaryExpr) -> Expr:
+        """Visit a binary expression and rewrite its operands.
+
+        Args:
+            expr: The binary expression.
+
+        Returns:
+            A new binary expression with rewritten operands.
+        """
+        return BinaryExpr(
+            expr.op,
+            self.visit(expr.left),
+            self.visit(expr.right),
+        )
+
+    def visit_unary(self, expr: UnaryExpr) -> Expr:
+        """Visit a unary expression and rewrite its operand.
+
+        Args:
+            expr: The unary expression.
+
+        Returns:
+            A new unary expression with rewritten operand.
+        """
+        return UnaryExpr(expr.op, self.visit(expr.operand))
+
+    def visit_udf(self, expr: UDFExpr) -> Expr:
+        """Visit a UDF expression and rewrite its arguments.
+
+        Args:
+            expr: The UDF expression.
+
+        Returns:
+            A new UDF expression with rewritten arguments.
+        """
+        new_args = [self.visit(arg) for arg in expr.args]
+        new_kwargs = {key: self.visit(value) for key, value in expr.kwargs.items()}
+        return replace(expr, args=new_args, kwargs=new_kwargs)
+
+    def visit_alias(self, expr: AliasExpr) -> Expr:
+        """Visit an alias expression and rewrite its inner expression.
+
+        Args:
+            expr: The alias expression.
+
+        Returns:
+            A new alias expression with rewritten inner expression and preserved name.
+        """
+        # We unalias returned expression to avoid nested aliasing
+        visited = self.visit(expr.expr)._unalias()
+        # NOTE: We're carrying over all of the other aspects of the alias
+        #       only replacing inner expre
+        return replace(
+            expr,
+            expr=visited,
+            # Alias expression will remain a renaming one (ie replacing source column)
+            # so long as it's referencing another column (and not otherwise)
+            #
+            # TODO replace w/ standalone rename expr
+            _is_rename=expr._is_rename and _is_col_expr(visited),
+        )
+
+    def visit_download(self, expr: "Expr") -> Expr:
+        """Visit a download expression (no rewriting needed).
+
+        Args:
+            expr: The download expression.
+
+        Returns:
+            The original download expression.
+        """
+        return expr
+
+    def visit_star(self, expr: StarExpr) -> Expr:
+        """Visit a star expression (no rewriting needed).
+
+        Args:
+            expr: The star expression.
+
+        Returns:
+            The original star expression.
+        """
+        return expr
+
+    def visit_monotonically_increasing_id(
+        self, expr: MonotonicallyIncreasingIdExpr
+    ) -> Expr:
+        """Visit a monotonically_increasing_id expression (no rewriting needed).
+
+        Args:
+            expr: The monotonically_increasing_id expression.
+
+        Returns:
+            The original expression.
+        """
+        return expr
+
+    def visit_random(self, expr: "RandomExpr") -> Expr:
+        """Visit a random expression (no rewriting needed).
+
+        Args:
+            expr: The random expression.
+
+        Returns:
+            The original random expression.
+        """
+        return expr
+
+    def visit_uuid(self, expr: "UUIDExpr") -> Expr:
+        """Visit a uuid expression (no rewriting needed).
+
+        Args:
+            expr: The uuid expression.
+
+        Returns:
+            The original uuid expression.
+        """
+        return expr
+
+
+def _is_col_expr(expr: Expr) -> bool:
+    return isinstance(expr, ColumnExpr) or (
+        isinstance(expr, AliasExpr) and isinstance(expr.expr, ColumnExpr)
+    )
+
+
+class _TreeReprVisitor(_ExprVisitor[str]):
+    """Visitor that generates a readable tree representation of expressions. Returns in pre-order traversal."""
+
+    def __init__(self, prefix: str = "", is_last: bool = True):
+        """
+        Initialize the tree representation visitor.
+
+        Args:
+            prefix: The prefix string for indentation (accumulated from parent nodes)
+            is_last: Whether this node is the last child of its parent
+        """
+        self.prefix = prefix
+        self.is_last = is_last
+        self._max_length = 50  # Maximum length of the node label
+
+    def _make_tree_lines(
+        self,
+        node_label: str,
+        children: List[tuple[str, "Expr"]] = None,
+        expr: "Expr" = None,
+    ) -> str:
+        """
+        Format a node and its children with tree box-drawing characters.
+
+        Args:
+            node_label: The label for this node (e.g., "ADD")
+            children: List of (label, child_expr) tuples to render as children
+            expr: The expression node (used to extract datatype)
+
+        Returns:
+            Multi-line string representation of the tree
+        """
+        lines = [node_label]
+
+        if children:
+            for i, (label, child_expr) in enumerate(children):
+                is_last_child = i == len(children) - 1
+
+                # Build prefix for the child based on whether current node is last
+                child_prefix = self.prefix + ("    " if self.is_last else "│   ")
+
+                # Choose connector: └── for last child, ├── for others
+                connector = "└── " if is_last_child else "├── "
+
+                # Recursively visit the child with updated prefix
+                child_visitor = _TreeReprVisitor(child_prefix, is_last_child)
+                child_lines = child_visitor.visit(child_expr).split("\n")
+
+                # Add the first line with label and connector
+                if label:
+                    lines.append(f"{child_prefix}{connector}{label}: {child_lines[0]}")
+                else:
+                    lines.append(f"{child_prefix}{connector}{child_lines[0]}")
+
+                # Add remaining lines from child with proper indentation
+                for line in child_lines[1:]:
+                    lines.append(line)
+
+        return "\n".join(lines)
+
+    def visit_column(self, expr: "ColumnExpr") -> str:
+        return self._make_tree_lines(f"COL({expr.name!r})", expr=expr)
+
+    def visit_literal(self, expr: "LiteralExpr") -> str:
+        # Truncate long values for readability
+        value_repr = repr(expr.value)
+        if len(value_repr) > self._max_length:
+            value_repr = value_repr[: self._max_length - 3] + "..."
+        return self._make_tree_lines(f"LIT({value_repr})", expr=expr)
+
+    def visit_binary(self, expr: "BinaryExpr") -> str:
+        return self._make_tree_lines(
+            f"{expr.op.name}",
+            children=[
+                ("left", expr.left),
+                ("right", expr.right),
+            ],
+            expr=expr,
+        )
+
+    def visit_unary(self, expr: "UnaryExpr") -> str:
+        return self._make_tree_lines(
+            f"{expr.op.name}",
+            children=[("operand", expr.operand)],
+            expr=expr,
+        )
+
+    def visit_alias(self, expr: "AliasExpr") -> str:
+        rename_marker = " [rename]" if expr._is_rename else ""
+        return self._make_tree_lines(
+            f"ALIAS({expr.name!r}){rename_marker}",
+            children=[("", expr.expr)],
+            expr=expr,
+        )
+
+    def visit_udf(self, expr: "UDFExpr") -> str:
+        # Get function name for better readability
+        fn_name = getattr(expr.fn, "__name__", str(expr.fn))
+
+        children = []
+        # Add positional arguments
+        for i, arg in enumerate(expr.args):
+            children.append((f"arg[{i}]", arg))
+
+        # Add keyword arguments
+        for key, value in expr.kwargs.items():
+            children.append((f"kwarg[{key!r}]", value))
+
+        return self._make_tree_lines(
+            f"UDF({fn_name})",
+            children=children if children else None,
+            expr=expr,
+        )
+
+    def visit_download(self, expr: "DownloadExpr") -> str:
+        return self._make_tree_lines(f"DOWNLOAD({expr.uri_column_name!r})", expr=expr)
+
+    def visit_star(self, expr: "StarExpr") -> str:
+        return self._make_tree_lines("COL(*)", expr=expr)
+
+    def visit_monotonically_increasing_id(
+        self, expr: "MonotonicallyIncreasingIdExpr"
+    ) -> str:
+        return self._make_tree_lines("MONOTONICALLY_INCREASING_ID()", expr=expr)
+
+    def visit_random(self, expr: "RandomExpr") -> str:
+        if expr.seed is None:
+            label = "RANDOM()"
+        else:
+            label = f"RANDOM(seed={expr.seed}, reseed_after_execution={expr.reseed_after_execution})"
+        return self._make_tree_lines(label, expr=expr)
+
+    def visit_uuid(self, expr: "UUIDExpr") -> str:
+        return self._make_tree_lines("UUID()", expr=expr)
+
+
+class _InlineExprReprVisitor(_ExprVisitor[str]):
+    """Visitor that generates concise inline string representations of expressions.
+
+    This visitor creates single-line string representations suitable for displaying
+    in operator names, log messages, etc. It aims to be human-readable while keeping
+    the representation compact.
+    """
+
+    def __init__(self, max_literal_length: int = 20):
+        """Initialize the inline representation visitor.
+
+        Args:
+            max_literal_length: Maximum length for literal value representations
+        """
+        self._max_literal_length = max_literal_length
+
+    def visit_column(self, expr: "ColumnExpr") -> str:
+        """Visit a column expression and return its inline representation."""
+        return f"col({expr.name!r})"
+
+    def visit_literal(self, expr: "LiteralExpr") -> str:
+        """Visit a literal expression and return its inline representation."""
+        value_repr = repr(expr.value)
+        if len(value_repr) > self._max_literal_length:
+            value_repr = value_repr[: self._max_literal_length - 3] + "..."
+        return value_repr
+
+    def visit_binary(self, expr: "BinaryExpr") -> str:
+        """Visit a binary expression and return its inline representation."""
+        left_str = self.visit(expr.left)
+        right_str = self.visit(expr.right)
+
+        # Add parentheses around child binary expressions to avoid ambiguity
+        if isinstance(expr.left, BinaryExpr):
+            left_str = f"({left_str})"
+        if isinstance(expr.right, BinaryExpr):
+            right_str = f"({right_str})"
+
+        op_str = _INLINE_OP_SYMBOLS.get(expr.op, expr.op.name.lower())
+        return f"{left_str} {op_str} {right_str}"
+
+    def visit_unary(self, expr: "UnaryExpr") -> str:
+        """Visit a unary expression and return its inline representation."""
+        operand_str = self.visit(expr.operand)
+
+        # Add parentheses around binary expression operands to avoid ambiguity
+        if isinstance(expr.operand, BinaryExpr):
+            operand_str = f"({operand_str})"
+
+        # Map operations to symbols/functions
+        if expr.op == Operation.NOT:
+            return f"~{operand_str}"
+        elif expr.op == Operation.IS_NULL:
+            return f"{operand_str}.is_null()"
+        elif expr.op == Operation.IS_NOT_NULL:
+            return f"{operand_str}.is_not_null()"
+        else:
+            return f"{expr.op.name.lower()}({operand_str})"
+
+    def visit_alias(self, expr: "AliasExpr") -> str:
+        """Visit an alias expression and return its inline representation."""
+        inner_str = self.visit(expr.expr)
+        return f"{inner_str}.alias({expr.name!r})"
+
+    def visit_udf(self, expr: "UDFExpr") -> str:
+        """Visit a UDF expression and return its inline representation."""
+        # Get function name for better readability
+        # For callable objects (instances with __call__), use the class name
+        fn_name = getattr(expr.fn, "__name__", expr.fn.__class__.__name__)
+
+        # Build argument list
+        args_str = []
+        for arg in expr.args:
+            args_str.append(self.visit(arg))
+        for key, value in expr.kwargs.items():
+            args_str.append(f"{key}={self.visit(value)}")
+
+        args_repr = ", ".join(args_str) if args_str else ""
+        return f"{fn_name}({args_repr})"
+
+    def visit_download(self, expr: "DownloadExpr") -> str:
+        """Visit a download expression and return its inline representation."""
+        return f"download({expr.uri_column_name!r})"
+
+    def visit_star(self, expr: "StarExpr") -> str:
+        """Visit a star expression and return its inline representation."""
+        return "col(*)"
+
+    def visit_monotonically_increasing_id(
+        self, expr: "MonotonicallyIncreasingIdExpr"
+    ) -> str:
+        """Visit a monotonically_increasing_id expression and return its inline representation."""
+        return "monotonically_increasing_id()"
+
+    def visit_random(self, expr: "RandomExpr") -> str:
+        """Visit a random expression and return its inline representation."""
+        return "random()"
+
+    def visit_uuid(self, expr: "UUIDExpr") -> str:
+        """Visit a uuid expression and return its inline representation."""
+        return "uuid()"
+
+
+class _StructuralFingerprintVisitor(_ExprVisitor[Hashable]):
+    """Visitor that computes a hashable structural fingerprint for an expression.
+
+    Two expressions that are structurally equivalent produce equal fingerprints,
+    so the fingerprint can be used as a cheap bucketing key before falling back to
+    full ``structurally_equals`` comparison (e.g. for common sub-expression
+    elimination).
+    """
+
+    def visit_column(self, expr: ColumnExpr) -> Hashable:
+        return _column_fingerprint_key(expr)
+
+    def visit_literal(self, expr: LiteralExpr) -> Hashable:
+        return _literal_fingerprint_key(expr)
+
+    def visit_binary(self, expr: BinaryExpr) -> Hashable:
+        return _binary_fingerprint_key(
+            expr,
+            self.visit(expr.left),
+            self.visit(expr.right),
+        )
+
+    def visit_unary(self, expr: UnaryExpr) -> Hashable:
+        return _unary_fingerprint_key(expr, self.visit(expr.operand))
+
+    def visit_udf(self, expr: UDFExpr) -> Hashable:
+        return _udf_fingerprint_key(
+            expr,
+            tuple(self.visit(arg) for arg in expr.args),
+            tuple(
+                (k, self.visit(v))
+                for k, v in sorted(expr.kwargs.items(), key=lambda item: item[0])
+            ),
+        )
+
+    def visit_alias(self, expr: AliasExpr) -> Hashable:
+        return _alias_fingerprint_key(expr, self.visit(expr.expr))
+
+    def visit_download(self, expr: DownloadExpr) -> Hashable:
+        return _download_fingerprint_key(expr)
+
+    def visit_star(self, expr: StarExpr) -> Hashable:
+        return _star_fingerprint_key()
+
+    def visit_monotonically_increasing_id(
+        self, expr: MonotonicallyIncreasingIdExpr
+    ) -> Hashable:
+        return _monotonically_increasing_id_fingerprint_key(expr)
+
+    def visit_random(self, expr: RandomExpr) -> Hashable:
+        return _random_fingerprint_key(expr)
+
+    def visit_uuid(self, expr: UUIDExpr) -> Hashable:
+        return _uuid_fingerprint_key(expr)
+
+
+@dataclass(frozen=True)
+class _ExpressionOccurrence:
+    expr: Expr
+    key: Hashable
+    depth: int
+
+
+class _StructuralFingerprintOccurrenceCollector(_ExprVisitor[Hashable]):
+    """Collect expression occurrences while computing structural keys bottom-up."""
+
+    def __init__(self):
+        self._occurrences: List[_ExpressionOccurrence] = []
+        self._depth = 0
+
+    def get_occurrences(self) -> List[_ExpressionOccurrence]:
+        return self._occurrences
+
+    def _visit_child(self, expr: Expr) -> Hashable:
+        self._depth += 1
+        try:
+            return self.visit(expr)
+        finally:
+            self._depth -= 1
+
+    def _record(self, expr: Expr, key: Hashable) -> Hashable:
+        self._occurrences.append(
+            _ExpressionOccurrence(
+                expr=expr,
+                key=key,
+                depth=self._depth,
+            )
+        )
+        return key
+
+    def visit_column(self, expr: ColumnExpr) -> Hashable:
+        return self._record(expr, _column_fingerprint_key(expr))
+
+    def visit_literal(self, expr: LiteralExpr) -> Hashable:
+        return self._record(expr, _literal_fingerprint_key(expr))
+
+    def visit_binary(self, expr: BinaryExpr) -> Hashable:
+        return self._record(
+            expr,
+            _binary_fingerprint_key(
+                expr,
+                self._visit_child(expr.left),
+                self._visit_child(expr.right),
+            ),
+        )
+
+    def visit_unary(self, expr: UnaryExpr) -> Hashable:
+        return self._record(
+            expr,
+            _unary_fingerprint_key(expr, self._visit_child(expr.operand)),
+        )
+
+    def visit_udf(self, expr: UDFExpr) -> Hashable:
+        return self._record(
+            expr,
+            _udf_fingerprint_key(
+                expr,
+                tuple(self._visit_child(arg) for arg in expr.args),
+                tuple(
+                    (k, self._visit_child(v))
+                    for k, v in sorted(expr.kwargs.items(), key=lambda item: item[0])
+                ),
+            ),
+        )
+
+    def visit_alias(self, expr: AliasExpr) -> Hashable:
+        return self._record(
+            expr,
+            _alias_fingerprint_key(expr, self._visit_child(expr.expr)),
+        )
+
+    def visit_download(self, expr: DownloadExpr) -> Hashable:
+        return self._record(expr, _download_fingerprint_key(expr))
+
+    def visit_star(self, expr: StarExpr) -> Hashable:
+        return self._record(expr, _star_fingerprint_key())
+
+    def visit_monotonically_increasing_id(
+        self, expr: MonotonicallyIncreasingIdExpr
+    ) -> Hashable:
+        return self._record(expr, _monotonically_increasing_id_fingerprint_key(expr))
+
+    def visit_random(self, expr: RandomExpr) -> Hashable:
+        return self._record(expr, _random_fingerprint_key(expr))
+
+    def visit_uuid(self, expr: UUIDExpr) -> Hashable:
+        return self._record(expr, _uuid_fingerprint_key(expr))
+
+
+def get_column_references(expr: Expr) -> List[str]:
+    """Extract all column references from an expression.
+
+    This is a convenience function that creates a _ColumnReferenceCollector,
+    visits the expression tree, and returns the list of referenced column names.
+
+    Args:
+        expr: The expression to extract column references from.
+
+    Returns:
+        List of column names referenced in the expression, in order of appearance.
+
+    Example:
+        >>> from ray.data.expressions import col
+        >>> expr = (col("a") > 5) & (col("b") == "test")
+        >>> get_column_references(expr)
+        ['a', 'b']
+    """
+    collector = _ColumnReferenceCollector()
+    collector.visit(expr)
+    return collector.get_column_refs()

python/ray/data/_internal/planner/plan_expression/synthetic_impl.py

@@ -0,0 +1,124 @@
+import uuid
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data.block import BlockColumn, BlockType
+
+
+def eval_random(
+    num_rows: int,
+    block_type: BlockType,
+    *,
+    seed: int | None = None,
+    reseed_after_execution: bool = True,
+    instance_id: str | None = None,
+) -> BlockColumn:
+    """Implementation of the random expression.
+
+    Args:
+        num_rows: The number of rows to generate random values for.
+        block_type: The type of block to generate random values for.
+        seed: The seed to use for the random number generator.
+        reseed_after_execution: Whether to reseed the random number generator after each execution.
+        instance_id: Unique identifier for the random expression instance, used to isolate
+            block count state when a single task processes multiple blocks.
+
+    Returns:
+        A BlockColumn containing the random values.
+
+    Raises:
+        TypeError: If the block type is not supported.
+    """
+
+    if seed is not None:
+        # Numpy allows using a seed sequence (list of integers) to initialize
+        # a random number generator. This allows us to maintain reproduciblity while
+        # ensuring randomness in parallel execution.
+        # See https://numpy.org/doc/2.2/reference/random/parallel.html#sequence-of-integer-seeds
+        # Below we uses four components to create a seed sequence (fastest changing component first):
+        # 1. A per-block counter within the task (to differentiate blocks in the same task)
+        # 2. An index based on the remote task in Ray Data
+        # 3. An incrementing index of Ray Dataset execution (e.g., multiple training epochs)
+        # 4. A base seed fixed by the user
+
+        ctx = TaskContext.get_current()
+
+        if ctx is None:
+            import warnings
+
+            warnings.warn(
+                "TaskContext is not available for random() expression with seed. "
+                "Falling back to task_idx=0 for all tasks, which reduces the parallelism "
+                "benefits of random number generation. If you see this warning, please "
+                "report it as it may indicate an execution context issue.",
+                stacklevel=2,
+            )
+            task_idx = 0
+            block_idx = 0
+        else:
+            task_idx = ctx.task_idx
+
+            # Key the counter by expression instance ID so that multiple expressions
+            # in the same projection will have isolated block count state.
+            # This is required because a single task may process multiple blocks if
+            # the upstream data source does not compress the data into a single block.
+            if instance_id is not None:
+                counter_key = f"_random_{instance_id}_counter"
+                block_idx = ctx.kwargs.get(counter_key, 0)
+                ctx.kwargs[counter_key] = block_idx + 1
+            else:
+                block_idx = 0
+
+        if reseed_after_execution:
+            from ray.data.context import DataContext
+
+            data_context = (
+                DataContext.get_current()
+            )  # get or create DataContext, never None
+            execution_idx = data_context._execution_idx
+        else:
+            execution_idx = 0
+
+        # Numpy recommends fastest changing component to be the first element
+        block_seed = [block_idx, task_idx, execution_idx, seed]
+    else:
+        block_seed = None
+
+    rng = np.random.default_rng(block_seed)
+    random_values = rng.random(num_rows)
+
+    # Convert to appropriate format based on block type
+    if block_type == BlockType.PANDAS:
+        return pd.Series(random_values, dtype=np.float64)
+    elif block_type == BlockType.ARROW:
+        return pa.array(random_values, type=pa.float64())
+
+    raise TypeError(f"Unsupported block type: {block_type}")
+
+
+def eval_uuid(
+    num_rows: int,
+    block_type: BlockType,
+) -> BlockColumn:
+    """Implementation of the uuid expression.
+
+    Args:
+        num_rows: The number of rows to generate uuid values for.
+        block_type: The type of block to generate uuid values for.
+
+    Returns:
+        A BlockColumn containing the uuid values.
+
+    Raises:
+        TypeError: If the block type is not supported.
+    """
+    arr = [str(uuid.uuid4()) for _ in range(num_rows)]
+    if block_type == BlockType.PANDAS:
+        return pd.Series(arr, dtype=pd.StringDtype())
+    elif block_type == BlockType.ARROW:
+        return pa.array(arr, type=pa.string())
+
+    raise TypeError(f"Unsupported block type: {block_type}")

python/ray/data/_internal/planner/plan_list_files_op.py

@@ -0,0 +1,173 @@
+"""Physical planner for the V2 ``ListFiles`` source operator.
+
+Emits ``FileManifest`` blocks by (a) sharding user-supplied paths into
+parallel listing tasks, (b) invoking the configured ``FileIndexer``, and
+optionally (c) globally shuffling + size-balanced bucketing before the
+downstream ``ReadFiles`` physical op consumes them.
+
+Checkpoint filtering is not attached here — it's wrapped around the
+downstream ``ReadFiles`` physical op by
+:func:`plan_read_files_op_with_checkpoint_filter`, matching V1's
+dispatch pattern.
+"""
+from __future__ import annotations
+
+import logging
+from functools import partial
+from typing import List
+
+import numpy as np
+import pyarrow as pa
+
+import ray
+from ray.data._internal.datasource_v2.listing.file_manifest import (
+    PATH_COLUMN_NAME,
+)
+from ray.data._internal.datasource_v2.listing.listing_utils import (
+    list_files_for_each_block,
+    partition_files,
+    shuffle_files,
+)
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.input_data_buffer import (
+    InputDataBuffer,
+)
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+    MapTransformFn,
+)
+from ray.data._internal.logical.operators import ListFiles
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+# Cap on the number of parallel listing tasks. In practice most reads
+# pass a single directory (one task); this matters when users hand in
+# thousands of explicit paths.
+DEFAULT_MAX_NUM_LIST_FILES_TASKS = 200
+
+
+def plan_list_files_op(
+    op: ListFiles,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    assert len(physical_children) == 0
+
+    # NOTE: Avoid capturing ``op`` inside closures — only its field values.
+    file_extensions = op.file_extensions
+    partition_filter = op.partition_filter
+    filesystem = op.filesystem
+    indexer = op.file_indexer
+    partitioner = op.file_partitioner
+
+    shuffle_config = op.shuffle_config_factory()
+
+    transform_fns: List[MapTransformFn] = [
+        BlockMapTransformFn(
+            partial(
+                list_files_for_each_block,
+                indexer=indexer,
+                filesystem=filesystem,
+                file_extensions=file_extensions,
+                partition_filter=partition_filter,
+                preserve_order=data_context.execution_options.preserve_order,
+            ),
+            # Disable block-shaping: produce manifest blocks as-is.
+            disable_block_shaping=True,
+        ),
+    ]
+
+    if shuffle_config is not None:
+        transform_fns.append(
+            BlockMapTransformFn(
+                partial(
+                    shuffle_files,
+                    shuffle_config=shuffle_config,
+                    execution_idx=data_context._execution_idx,
+                ),
+                disable_block_shaping=True,
+            )
+        )
+
+    if partitioner is not None:
+        transform_fns.append(
+            BlockMapTransformFn(
+                partial(partition_files, partitioner=partitioner),
+                disable_block_shaping=True,
+            )
+        )
+
+    map_transformer = MapTransformer(transform_fns)
+
+    map_op = MapOperator.create(
+        map_transformer,
+        _create_input_data_buffer(
+            op,
+            data_context,
+            # Shuffle needs every manifest on a single task to compute one
+            # global RNG over the full listing.
+            should_parallelize=shuffle_config is None,
+        ),
+        data_context,
+        name="ListFiles",
+        # Listing is extremely fast; default backpressure would starve the
+        # downstream reader of inputs.
+        ray_remote_args={"_generator_backpressure_num_objects": -1},
+        # Don't fuse into the downstream ``ReadFiles`` — listing and reading
+        # have different resource profiles.
+        supports_fusion=False,
+    )
+    map_op.throttling_disabled = lambda: True
+    return map_op
+
+
+def _create_input_data_buffer(
+    op: ListFiles,
+    data_context: DataContext,
+    *,
+    should_parallelize: bool,
+) -> InputDataBuffer:
+    """Wrap ``op.paths`` into listing-input RefBundles.
+
+    Each bundle's block is a 1-column arrow table ``{"__path": [paths...]}``
+    that :func:`list_files_for_each_block` expands into manifest blocks.
+    """
+    if should_parallelize and op.paths:
+        path_splits = np.array_split(
+            list(op.paths),
+            min(DEFAULT_MAX_NUM_LIST_FILES_TASKS, len(op.paths)),
+        )
+    else:
+        path_splits = [list(op.paths)]
+
+    input_data: List[RefBundle] = []
+    for path_split in path_splits:
+        paths = list(path_split)
+        if not paths:
+            continue
+        block = pa.Table.from_pydict({PATH_COLUMN_NAME: paths})
+        metadata = BlockAccessor.for_block(block).get_metadata(
+            input_files=None, block_exec_stats=None
+        )
+        block_ref: ray.ObjectRef[Block] = ray.put(block)
+        ref_bundle = RefBundle(
+            (
+                # pyrefly: ignore[bad-argument-type]
+                BlockEntry(block_ref, metadata),
+            ),
+            # ``owns_blocks=False``: these are the root of the DAG and
+            # must not be freed eagerly, or the DAG can't be reconstructed.
+            owns_blocks=False,
+            schema=BlockAccessor.for_block(block).schema(),
+        )
+        input_data.append(ref_bundle)
+
+    return InputDataBuffer(data_context, input_data=input_data)

python/ray/data/_internal/planner/plan_read_files_op.py

@@ -0,0 +1,95 @@
+"""Physical planner for the V2 ``ReadFiles`` logical operator.
+
+``ReadFiles`` consumes ``FileManifest`` blocks from an upstream
+``ListFiles`` physical op. This planner wires one map transform —
+``do_read`` — that calls ``scanner.create_reader().read(manifest)`` for
+each incoming bucket.
+
+V2 reads never rename columns at the read stage; column renaming is
+always handled by a ``Project`` operator above ``ReadFiles``.
+
+Listing, shuffling, and size-balanced bucketing previously lived here;
+they've moved to :func:`plan_list_files_op` where they belong.
+
+Checkpoint wrapping (when ``data_context.checkpoint_config`` is set) is
+handled by the companion
+:func:`ray.data._internal.planner.checkpoint.plan_read_files_op.plan_read_files_op_with_checkpoint_filter`,
+registered via the planner's ``_get_plan_fns_for_checkpointing`` hook —
+same dispatch shape V1 uses for ``plan_read_op_with_checkpoint_filter``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Iterable, List
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.scanners.file_scanner import FileScanner
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.logical.operators import ReadFiles
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data.block import Block
+from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+def plan_read_files_op(
+    op: ReadFiles,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    """Convert a ``ReadFiles`` logical op into a reader ``MapOperator``.
+
+    Expects exactly one physical child: the upstream ``ListFiles`` op,
+    which produces balanced manifest blocks via its transform chain.
+    """
+    assert len(physical_children) == 1
+    upstream = physical_children[0]
+
+    # NOTE: Avoid capturing the whole ``op`` in closures — only field values.
+    scanner = op.scanner
+    block_udf = op.block_udf
+
+    def do_read(blocks: Iterable[Block], _: TaskContext) -> Iterable[Block]:
+        reader = scanner.create_reader()
+        # File-level predicate pruning (partition predicates pushed down
+        # onto the scanner) runs per incoming manifest block. Only
+        # ``FileScanner`` subclasses expose ``prune_manifest``; the base
+        # implementation is an identity no-op, and ``ArrowFileScanner``
+        # overrides it to evaluate ``partition_predicate``.
+        for block in blocks:
+            manifest = FileManifest(block)
+            if isinstance(scanner, FileScanner):
+                manifest = scanner.prune_manifest(manifest)
+            if len(manifest) == 0:
+                continue
+            for table in reader.read(manifest):
+                if block_udf is not None:
+                    table = block_udf(table)
+                yield table
+
+    return MapOperator.create(
+        MapTransformer(
+            [
+                BlockMapTransformFn(
+                    do_read,
+                    is_udf=False,
+                    output_block_size_option=OutputBlockSizeOption.of(
+                        target_max_block_size=data_context.target_max_block_size,
+                    ),
+                ),
+            ]
+        ),
+        upstream,
+        data_context,
+        name=op.name,
+        compute_strategy=op.compute,
+        ray_remote_args=op.ray_remote_args,
+    )

python/ray/data/_internal/planner/plan_read_op.py

@@ -0,0 +1,135 @@
+import logging
+import warnings
+from typing import Iterable, List
+
+import ray
+from ray import ObjectRef
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.logical.operators import Read
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.util import _warn_on_high_parallelism
+from ray.data.block import Block, BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource.datasource import ReadTask
+from ray.experimental.locations import get_local_object_locations
+from ray.util.debug import log_once
+
+TASK_SIZE_WARN_THRESHOLD_BYTES = 1024 * 1024  # 1 MiB
+
+logger = logging.getLogger(__name__)
+
+
+def _derive_metadata(read_task: ReadTask, read_task_ref: ObjectRef) -> BlockMetadata:
+    # NOTE: Use the `get_local_object_locations` API to get the size of the
+    # serialized ReadTask, instead of pickling.
+    # Because the ReadTask may capture ObjectRef objects, which cannot
+    # be serialized out-of-band.
+    locations = get_local_object_locations([read_task_ref])
+    task_size = locations[read_task_ref]["object_size"]
+    if task_size > TASK_SIZE_WARN_THRESHOLD_BYTES and log_once(
+        f"large_read_task_{read_task.read_fn.__name__}"
+    ):
+        warnings.warn(
+            "The serialized size of your read function named "
+            f"'{read_task.read_fn.__name__}' is {memory_string(task_size)}. This size "
+            "is relatively large. As a result, Ray might excessively "
+            "spill objects during execution. To fix this issue, avoid accessing "
+            f"`self` or other large objects in '{read_task.read_fn.__name__}'."
+        )
+
+    return BlockMetadata(
+        num_rows=1,
+        size_bytes=task_size,
+        exec_stats=None,
+        input_files=None,
+    )
+
+
+def plan_read_op(
+    op: Read,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Get the corresponding DAG of physical operators for Read.
+
+    Note this method only converts the given `op`, but not its input dependencies.
+    See Planner.plan() for more details.
+    """
+    assert len(physical_children) == 0
+
+    def get_input_data(target_max_block_size) -> List[RefBundle]:
+        parallelism = op.get_detected_parallelism()
+        assert (
+            parallelism is not None
+        ), "Read parallelism must be set by the optimizer before execution"
+
+        # Get the original read tasks
+        read_tasks = op.datasource_or_legacy_reader.get_read_tasks(
+            parallelism,
+            per_task_row_limit=op.per_block_limit,
+            data_context=data_context,
+        )
+
+        _warn_on_high_parallelism(parallelism, len(read_tasks))
+
+        ret = []
+        for read_task in read_tasks:
+            read_task_ref = ray.put(read_task)
+            ref_bundle = RefBundle(
+                (
+                    BlockEntry(
+                        # TODO: figure out a better way to pass read
+                        # tasks other than ray.put().
+                        read_task_ref,
+                        _derive_metadata(read_task, read_task_ref),
+                    ),
+                ),
+                # `owns_blocks` is False, because these refs are the root of the
+                # DAG. We shouldn't eagerly free them. Otherwise, the DAG cannot
+                # be reconstructed.
+                owns_blocks=False,
+                schema=None,
+            )
+            ret.append(ref_bundle)
+        return ret
+
+    inputs = InputDataBuffer(data_context, input_data_factory=get_input_data)
+
+    def do_read(blocks: Iterable[ReadTask], _: TaskContext) -> Iterable[Block]:
+        for read_task in blocks:
+            yield from read_task()
+
+    # Create a MapTransformer for a read operator
+    map_transformer = MapTransformer(
+        [
+            BlockMapTransformFn(
+                do_read,
+                is_udf=False,
+                output_block_size_option=OutputBlockSizeOption.of(
+                    target_max_block_size=data_context.target_max_block_size,
+                ),
+            ),
+        ]
+    )
+
+    return MapOperator.create(
+        map_transformer,
+        inputs,
+        data_context,
+        name=op.name,
+        compute_strategy=op.compute,
+        ray_remote_args=op.ray_remote_args,
+        isolate_workers=data_context.isolate_read_workers,
+    )

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -0,0 +1,983 @@
+import asyncio
+import collections
+import inspect
+import logging
+import queue
+from dataclasses import dataclass
+from threading import Thread
+from types import GeneratorType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+)
+
+if TYPE_CHECKING:
+    from ray.data.expressions import _CallableClassSpec
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+
+import ray
+from ray._common.utils import env_integer, get_or_create_event_loop
+from ray.data._internal.compute import ActorPoolStrategy, ComputeStrategy, get_compute
+from ray.data._internal.execution.bundle_queue import ExactMultipleSize, RebundleQueue
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BatchMapTransformFn,
+    BlockMapTransformFn,
+    MapTransformCallable,
+    MapTransformer,
+    Row,
+    RowMapTransformFn,
+)
+from ray.data._internal.execution.util import make_callable_class_single_threaded
+from ray.data._internal.logical.operators import (
+    AbstractUDFMap,
+    Filter,
+    FlatMap,
+    MapBatches,
+    MapRows,
+    Project,
+    StreamingRepartition,
+)
+from ray.data._internal.numpy_support import _is_valid_column_values
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.util import _truncated_repr
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    CallableClass,
+    DataBatch,
+    UserDefinedFunction,
+    _is_cudf_dataframe,
+)
+from ray.data.context import DataContext
+from ray.data.exceptions import UserCodeException
+from ray.util.rpdb import _is_ray_debugger_post_mortem_enabled
+
+logger = logging.getLogger(__name__)
+
+
+# Controls default max-concurrency setting for async row-based UDFs
+DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY = env_integer(
+    "RAY_DATA_DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY", 16
+)
+
+# Controls default max-concurrency setting for async batch-based UDFs
+DEFAULT_ASYNC_BATCH_UDF_MAX_CONCURRENCY = env_integer(
+    "RAY_DATA_DEFAULT_ASYNC_BATCH_UDF_MAX_CONCURRENCY", 4
+)
+
+
+@dataclass
+class UDFSpec:
+    """Specification for a callable class UDF to be instantiated in an actor.
+
+    Attributes:
+        spec: The callable class specification (contains class and constructor args)
+        instantiation_class: The class to instantiate (may be wrapped, e.g., for concurrency)
+    """
+
+    spec: "_CallableClassSpec"
+    instantiation_class: type
+
+
+class _MapActorContext:
+    def __init__(
+        self,
+        is_async: bool = False,
+        udf_instances: Optional[Dict[int, Any]] = None,
+    ):
+        """Initialize the map actor context.
+
+        Args:
+            is_async: Whether any UDF is async
+            udf_instances: Dict mapping UDF class ID to instantiated instance
+        """
+        self.is_async = is_async
+        self.udf_map_asyncio_loop = None
+        self.udf_map_asyncio_thread = None
+        self.udf_instances = udf_instances or {}
+
+        if is_async:
+            self._init_async()
+
+    def _init_async(self):
+        # Only used for callable class with async generator `__call__` method.
+        loop = get_or_create_event_loop()
+
+        def run_loop():
+            asyncio.set_event_loop(loop)
+            loop.run_forever()
+
+        thread = Thread(target=run_loop, daemon=True)
+        thread.start()
+        self.udf_map_asyncio_loop = loop
+        self.udf_map_asyncio_thread = thread
+
+
+def plan_project_op(
+    op: Project,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    # Extract expressions before defining the closure to prevent cloudpickle from
+    # serializing the entire op object (which may contain references to non-serializable
+    # datasources with weak references, e.g., PyIceberg tables)
+    projection_exprs = op.exprs
+    common_sub_exprs = op.get_common_sub_exprs()
+
+    compute = get_compute(op.compute)
+
+    # Create init_fn to initialize all callable class UDFs at actor startup
+    from ray.data.util.expression_utils import (
+        _create_callable_class_udf_init_fn,
+    )
+
+    init_fn = _create_callable_class_udf_init_fn(op.get_all_exprs())
+
+    def _project_block(block: Block) -> Block:
+        try:
+            from ray.data._internal.planner.plan_expression.expression_evaluator import (
+                eval_projection,
+            )
+
+            return eval_projection(
+                projection_exprs,
+                block,
+                common_sub_exprs=common_sub_exprs,
+            )
+        except Exception as e:
+            _try_wrap_udf_exception(e)
+
+    map_transformer = MapTransformer(
+        [
+            BlockMapTransformFn(
+                _generate_transform_fn_for_map_block(_project_block),
+                disable_block_shaping=(len(op.exprs) == 0),
+            )
+        ],
+        init_fn=init_fn,
+    )
+    return MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        ray_remote_args=op.ray_remote_args,
+        ray_remote_args_fn=op.ray_remote_args_fn,
+    )
+
+
+def plan_streaming_repartition_op(
+    op: StreamingRepartition,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+    compute = get_compute(op.compute)
+    transform_fn = BlockMapTransformFn(
+        lambda blocks, ctx: blocks,
+        output_block_size_option=OutputBlockSizeOption.of(
+            target_num_rows_per_block=op.target_num_rows_per_block,  # To split n*target_max_block_size row into n blocks
+        ),
+    )
+    map_transformer = MapTransformer([transform_fn])
+
+    if op.strict:
+        ref_bundler = RebundleQueue(ExactMultipleSize(op.target_num_rows_per_block))
+    else:
+        ref_bundler = None
+
+    operator = MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        ref_bundler=ref_bundler,
+        ray_remote_args=op.ray_remote_args,
+        ray_remote_args_fn=op.ray_remote_args_fn,
+    )
+
+    return operator
+
+
+def plan_filter_op(
+    op: Filter,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    output_block_size_option = OutputBlockSizeOption.of(
+        target_max_block_size=data_context.target_max_block_size,
+    )
+
+    predicate_expr = op.predicate_expr
+    compute = get_compute(op.compute)
+    if predicate_expr is not None:
+
+        def filter_block_fn(
+            blocks: Iterable[Block], ctx: TaskContext
+        ) -> Iterable[Block]:
+            for block in blocks:
+                block_accessor = BlockAccessor.for_block(block)
+                filtered_block = block_accessor.filter(predicate_expr)
+                yield filtered_block
+
+        init_fn = None
+        transform_fn = BlockMapTransformFn(
+            filter_block_fn,
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+    else:
+        udf_is_callable_class = isinstance(op.fn, CallableClass)
+        filter_fn, init_fn = _get_udf(
+            op.fn,
+            op.fn_args,
+            op.fn_kwargs,
+            op.fn_constructor_args if udf_is_callable_class else None,
+            op.fn_constructor_kwargs if udf_is_callable_class else None,
+            compute=compute,
+        )
+
+        transform_fn = RowMapTransformFn(
+            _generate_transform_fn_for_filter(filter_fn),
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+
+    map_transformer = MapTransformer([transform_fn], init_fn=init_fn)
+
+    return MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        ray_remote_args=op.ray_remote_args,
+        ray_remote_args_fn=op.ray_remote_args_fn,
+    )
+
+
+def plan_udf_map_op(
+    op: AbstractUDFMap,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    """Get the corresponding physical operators DAG for AbstractUDFMap operators.
+
+    Note this method only converts the given `op`, but not its input dependencies.
+    See Planner.plan() for more details.
+    """
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    output_block_size_option = OutputBlockSizeOption.of(
+        target_max_block_size=data_context.target_max_block_size,
+    )
+
+    compute = get_compute(op.compute)
+    udf_is_callable_class = isinstance(op.fn, CallableClass)
+    fn, init_fn = _get_udf(
+        op.fn,
+        op.fn_args,
+        op.fn_kwargs,
+        op.fn_constructor_args if udf_is_callable_class else None,
+        op.fn_constructor_kwargs if udf_is_callable_class else None,
+        compute=compute,
+    )
+
+    if isinstance(op, MapBatches):
+        transform_fn = BatchMapTransformFn(
+            _generate_transform_fn_for_map_batches(fn),
+            batch_size=op.batch_size,
+            batch_format=op.batch_format,
+            zero_copy_batch=op.zero_copy_batch,
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+
+    else:
+        if isinstance(op, MapRows):
+            udf_fn = _generate_transform_fn_for_map_rows(fn)
+        elif isinstance(op, FlatMap):
+            udf_fn = _generate_transform_fn_for_flat_map(fn)
+        else:
+            raise ValueError(f"Found unknown logical operator during planning: {op}")
+
+        transform_fn = RowMapTransformFn(
+            udf_fn,
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+
+    map_transformer = MapTransformer([transform_fn], init_fn=init_fn)
+
+    return MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        min_rows_per_bundle=op.min_rows_per_bundled_input,
+        ray_remote_args_fn=op.ray_remote_args_fn,
+        ray_remote_args=op.ray_remote_args,
+        per_block_limit=op.per_block_limit,
+    )
+
+
+def _get_udf(
+    op_fn: Callable,
+    op_fn_args: Tuple[Any, ...],
+    op_fn_kwargs: Dict[str, Any],
+    op_fn_constructor_args: Optional[Tuple[Any, ...]],
+    op_fn_constructor_kwargs: Optional[Dict[str, Any]],
+    compute: Optional[ComputeStrategy],
+):
+    # Note, it's important to define these standalone variables.
+    # So the parsed functions won't need to capture the entire operator, which may not
+    # be serializable.
+    udf = op_fn
+    fn_args = op_fn_args or ()
+    fn_kwargs = op_fn_kwargs or {}
+
+    if isinstance(udf, CallableClass):
+        from ray.data.expressions import _CallableClassSpec
+
+        fn_constructor_args = op_fn_constructor_args or ()
+        fn_constructor_kwargs = op_fn_constructor_kwargs or {}
+
+        is_async_udf = _is_async_udf(udf.__call__)
+
+        # Capture original class BEFORE wrapping for use as dict key
+        original_udf_class = udf
+
+        if (
+            not is_async_udf
+            and isinstance(compute, ActorPoolStrategy)
+            and not compute.enable_true_multi_threading
+        ):
+            # NOTE: By default Actor-based UDFs are restricted to run within a
+            # single-thread (when enable_true_multi_threading=False).
+            #
+            # Historically, this has been done to allow block-fetching, batching, etc to
+            # be overlapped with the actual UDF invocation, while avoiding the
+            # pitfalls of concurrent GPU access (like OOMs, etc) when specifying
+            # max_concurrency > 1.
+            udf = make_callable_class_single_threaded(udf)
+
+        # Create the callable class spec for this UDF
+        callable_class_spec = _CallableClassSpec(
+            cls=original_udf_class,
+            args=fn_constructor_args,
+            kwargs=fn_constructor_kwargs,
+        )
+
+        # Use the shared init function creator (handles both map_batches and expressions)
+        init_fn = create_actor_context_init_fn(
+            udf_specs=[UDFSpec(spec=callable_class_spec, instantiation_class=udf)]
+        )
+
+        # Capture the spec for lookup on the actor
+        captured_spec = callable_class_spec
+
+        if inspect.iscoroutinefunction(udf.__call__):
+            # Async coroutine UDF: wrapper must be async to work with async transform machinery
+            async def _wrapped_udf_map_fn(item: Any) -> Any:
+                assert ray.data._map_actor_context is not None
+                assert ray.data._map_actor_context.is_async
+
+                try:
+                    # Use spec's key for lookup
+                    udf_key = captured_spec.make_key()
+                    udf_instance = ray.data._map_actor_context.udf_instances[udf_key]
+                    # Direct await - already in async context
+                    return await udf_instance(
+                        item,
+                        *fn_args,
+                        **fn_kwargs,
+                    )
+                except Exception as e:
+                    _try_wrap_udf_exception(e)
+
+        elif inspect.isasyncgenfunction(udf.__call__):
+
+            async def _wrapped_udf_map_fn(item: Any) -> Any:
+                assert ray.data._map_actor_context is not None
+                assert ray.data._map_actor_context.is_async
+
+                try:
+                    # Use spec's key for lookup
+                    udf_key = captured_spec.make_key()
+                    udf_instance = ray.data._map_actor_context.udf_instances[udf_key]
+                    gen = udf_instance(
+                        item,
+                        *fn_args,
+                        **fn_kwargs,
+                    )
+
+                    async for res in gen:
+                        yield res
+                except Exception as e:
+                    _try_wrap_udf_exception(e, item)
+
+        else:
+            assert isinstance(
+                udf.__call__, Callable
+            ), f"Expected Callable, got {udf.__call__} ({type(udf.__call__)})"
+
+            def _wrapped_udf_map_fn(item: Any) -> Any:
+                assert ray.data._map_actor_context is not None
+                assert not ray.data._map_actor_context.is_async
+                try:
+                    # Use spec's key for lookup
+                    udf_key = captured_spec.make_key()
+                    udf_instance = ray.data._map_actor_context.udf_instances[udf_key]
+                    return udf_instance(
+                        item,
+                        *fn_args,
+                        **fn_kwargs,
+                    )
+                except Exception as e:
+                    _try_wrap_udf_exception(e)
+
+    else:
+
+        def _wrapped_udf_map_fn(item: Any) -> Any:
+            try:
+                return udf(item, *fn_args, **fn_kwargs)
+            except Exception as e:
+                _try_wrap_udf_exception(e)
+
+        def init_fn():
+            pass
+
+    return _wrapped_udf_map_fn, init_fn
+
+
+def _try_wrap_udf_exception(e: Exception, item: Any = None):
+    """If the Ray Debugger is enabled, keep the full stack trace unmodified
+    so that the debugger can stop at the initial unhandled exception.
+    Otherwise, clear the stack trace to omit noisy internal code path."""
+    ctx = ray.data.DataContext.get_current()
+    if _is_ray_debugger_post_mortem_enabled() or ctx.raise_original_map_exception:
+        raise e
+    else:
+        raise UserCodeException("UDF failed to process a data block.") from e
+
+
+# Following are util functions for converting UDFs to `MapTransformCallable`s.
+
+
+def _validate_batch_output(batch: Block) -> None:
+    allowed = isinstance(
+        batch,
+        (
+            list,
+            pa.Table,
+            np.ndarray,
+            collections.abc.Mapping,
+            pd.core.frame.DataFrame,
+            dict,
+        ),
+    ) or _is_cudf_dataframe(batch)
+    if not allowed:
+        raise ValueError(
+            "The `fn` you passed to `map_batches` returned a value of type "
+            f"{type(batch)}. This isn't allowed -- `map_batches` expects "
+            "`fn` to return a `pandas.DataFrame`, `pyarrow.Table`, "
+            "`cudf.DataFrame`, `numpy.ndarray`, `list`, or "
+            "`dict[str, numpy.ndarray]`."
+        )
+
+    if isinstance(batch, list):
+        raise ValueError(
+            f"Error validating {_truncated_repr(batch)}: "
+            "Returning a list of objects from `map_batches` is not "
+            "allowed in Ray 2.5. To return Python objects, "
+            "wrap them in a named dict field, e.g., "
+            "return `{'results': objects}` instead of just `objects`."
+        )
+
+    # Handle cudf.DataFrame before the Mapping check, since cudf.DataFrame
+    # implements the Mapping protocol. Mirrors the order in batch_to_block.
+    if _is_cudf_dataframe(batch):
+        return
+
+    if isinstance(batch, collections.abc.Mapping):
+        for key, value in list(batch.items()):
+            if not _is_valid_column_values(value):
+                raise ValueError(
+                    f"Error validating {_truncated_repr(batch)}: "
+                    "The `fn` you passed to `map_batches` returned a "
+                    f"`dict`. `map_batches` expects all `dict` values "
+                    f"to be `list` or `np.ndarray` type, but the value "
+                    f"corresponding to key {key!r} is of type "
+                    f"{type(value)}. To fix this issue, convert "
+                    f"the {type(value)} to a `np.ndarray`."
+                )
+
+
+class _TransformingBatchIterator(Iterator[DataBatch]):
+    """Iterator that applies a UDF to batches.
+
+    Unlike a generator, local variables in __next__ go out of scope when the method
+    returns, avoiding holding references to yielded values.
+
+    Uses a deque with popleft() to actually release references when items are consumed,
+    rather than keeping them in an iterator.
+    """
+
+    def __init__(self, batches: Iterable[DataBatch], fn: UserDefinedFunction):
+        self._input_iter = iter(batches)
+        self._fn = fn
+        self._cur_output_iter: Optional[Iterator[DataBatch]] = None
+
+    def __iter__(self) -> "_TransformingBatchIterator":
+        return self
+
+    def __next__(self) -> DataBatch:
+        while True:
+            # Check if there's pending output iter we'd continue fetching
+            # from
+            if self._cur_output_iter is not None:
+                try:
+                    out_batch = next(self._cur_output_iter)
+                except StopIteration:
+                    pass
+                else:
+                    _validate_batch_output(out_batch)
+                    return out_batch
+
+            # Fetch the next batch from upstream
+            input_batch = next(self._input_iter)
+
+            if (
+                not isinstance(input_batch, collections.abc.Mapping)
+                and not _is_cudf_dataframe(input_batch)
+                and BlockAccessor.for_block(input_batch).num_rows() == 0
+            ):
+                # For empty input blocks, we directly output them without
+                # calling the UDF.
+                # TODO(hchen): This workaround is because some all-to-all
+                # operators output empty blocks with no schema.
+                self._cur_output_iter = _ReleasingIterator(
+                    collections.deque([input_batch])
+                )
+            else:
+                try:
+                    res = self._fn(input_batch)
+
+                    if not isinstance(res, GeneratorType):
+                        # NOTE: It's critical that we're utilizing *releasing* iterator
+                        #       to avoid capturing intermediate objects along the whole
+                        #       iterator chain
+                        self._cur_output_iter = _ReleasingIterator(
+                            collections.deque([res])
+                        )
+                    else:
+                        # In cases when UDF returns a generator we iterate over it
+                        # as is (given that we can't release intermediate state from
+                        # UDF anyway)
+                        self._cur_output_iter = res
+                except ValueError as e:
+                    read_only_msgs = [
+                        "assignment destination is read-only",
+                        "buffer source array is read-only",
+                    ]
+                    err_msg = str(e)
+                    if any(msg in err_msg for msg in read_only_msgs):
+                        raise ValueError(
+                            f"Batch mapper function {self._fn.__name__} tried to mutate a "
+                            "zero-copy read-only batch. To be able to mutate the "
+                            "batch, pass zero_copy_batch=False to map_batches(); "
+                            "this will create a writable copy of the batch before "
+                            "giving it to fn. To elide this copy, modify your mapper "
+                            "function so it doesn't try to mutate its input."
+                        ) from e
+                    else:
+                        raise e from None
+
+
+def _generate_transform_fn_for_map_batches(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[DataBatch, DataBatch]:
+
+    if _is_async_udf(fn):
+        transform_fn = _generate_transform_fn_for_async_map(
+            fn,
+            _validate_batch_output,
+            max_concurrency=DEFAULT_ASYNC_BATCH_UDF_MAX_CONCURRENCY,
+        )
+
+    else:
+
+        def transform_fn(
+            batches: Iterable[DataBatch], _: TaskContext
+        ) -> Iterable[DataBatch]:
+            return _TransformingBatchIterator(batches, fn)
+
+    return transform_fn
+
+
+def _is_async_udf(fn: UserDefinedFunction) -> bool:
+    return inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn)
+
+
+def create_actor_context_init_fn(
+    udf_specs: List[UDFSpec],
+):
+    """Create an init function for registering callable class UDFs in actor context.
+
+    This is the shared core logic between map_batches (single UDF) and expressions (multiple UDFs).
+
+    Args:
+        udf_specs: List of UDF specifications
+
+    Returns:
+        An init function that sets up all UDFs in the actor context
+    """
+
+    def init_fn():
+        import ray
+
+        if ray.data._map_actor_context is None:
+            # Check if any UDF is async
+            has_async_udf = any(
+                _is_async_udf(spec.instantiation_class.__call__) for spec in udf_specs
+            )
+
+            # Create instances for all callable class UDFs
+            udf_instances = {}
+            for spec in udf_specs:
+                # Use the spec's key for deduplication and lookup
+                udf_key = spec.spec.make_key()
+                if udf_key not in udf_instances:
+                    # Instantiate using the wrapped/processed class
+                    udf_instances[udf_key] = spec.instantiation_class(
+                        *spec.spec.args, **spec.spec.kwargs
+                    )
+
+            # Single unified context for all UDFs
+            ray.data._map_actor_context = _MapActorContext(
+                is_async=has_async_udf,
+                udf_instances=udf_instances,
+            )
+
+    return init_fn
+
+
+def _validate_row_output(item):
+    if not isinstance(item, collections.abc.Mapping):
+        raise ValueError(
+            f"Error validating {_truncated_repr(item)}: "
+            "Standalone Python objects are not "
+            "allowed in Ray >= 2.5. To return Python objects from map(), "
+            "wrap them in a dict, e.g., "
+            "return `{'item': item}` instead of just `item`."
+        )
+
+
+def _generate_transform_fn_for_map_rows(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Row, Row]:
+
+    if _is_async_udf(fn):
+        transform_fn = _generate_transform_fn_for_async_map(
+            fn,
+            _validate_row_output,
+            # NOTE: UDF concurrency is limited
+            max_concurrency=DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY,
+        )
+
+    else:
+
+        def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+            for row in rows:
+                out_row = fn(row)
+                _validate_row_output(out_row)
+                yield out_row
+
+    return transform_fn
+
+
+def _generate_transform_fn_for_flat_map(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Row, Iterable[Row]]:
+    if _is_async_udf(fn):
+        # UDF is a callable class with async generator `__call__` method.
+        transform_fn = _generate_transform_fn_for_async_map(
+            fn,
+            _validate_row_output,
+            max_concurrency=DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY,
+            is_flat_map=True,
+        )
+
+    else:
+
+        def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+            for row in rows:
+                for out_row in fn(row):
+                    _validate_row_output(out_row)
+                    yield out_row
+
+    return transform_fn
+
+
+def _generate_transform_fn_for_filter(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Row, Row]:
+    def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+        for row in rows:
+            if fn(row):
+                yield row
+
+    return transform_fn
+
+
+def _generate_transform_fn_for_map_block(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Block, Block]:
+    def transform_fn(blocks: Iterable[Block], _: TaskContext) -> Iterable[Block]:
+        for block in blocks:
+            out_block = fn(block)
+            yield out_block
+
+    return transform_fn
+
+
+_SENTINEL = object()
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def _generate_transform_fn_for_async_map(
+    fn: UserDefinedFunction,
+    validate_fn: Callable,
+    *,
+    max_concurrency: int,
+    is_flat_map: bool = False,
+) -> MapTransformCallable:
+    assert max_concurrency > 0, "Max concurrency must be positive"
+
+    if inspect.isasyncgenfunction(fn):
+
+        async def _apply_udf(item: T) -> List[U]:
+            gen = fn(item)
+            # NOTE: Async generator is unrolled inside the task to maintain
+            #       requested concurrency level (`max_concurrent_batches`)
+            return [out async for out in gen]
+
+    elif inspect.iscoroutinefunction(fn):
+
+        async def _apply_udf(item: T) -> List[U]:
+            res = await fn(item)
+            return res if is_flat_map else [res]
+
+    else:
+        raise ValueError(f"Expected a coroutine function, got {fn}")
+
+    # Goals of the algorithm applying async UDF application to the provided iterator
+    # are following:
+    #
+    #   - No more than `max_concurrency` async tasks are running
+    #     at any given moment
+    #   - Slow consumption from the output queue should result in
+    #     the processing to get back-pressured (so that output queue
+    #     doesn't grow unbounded)
+    #   - Order of the items (rows/batches) produced by this method
+    #     *must be* deterministic (though is not guaranteed to be specified
+    #     if max_concurrency > 1)
+    #
+    # To achieve that, algorithm applying async UDF to elements of the provided sequence
+    # is structured like following:
+    #
+    #   - Task scheduling and subsequent results re-ordering are performed as
+    #     different stages (inside `_schedule` and `_report` methods respectively)
+    #
+    #   - Scheduling stage aim to schedule and run no more than `max_concurrency` tasks
+    #     at any given moment
+    #
+    #   - Once task completes it's added into task completion queue for its results to be
+    #     subsequently reported with deterministic ordering). Task completion queue is
+    #     capped at `maxsize=max_concurrency` elements to make sure scheduling stage is
+    #     throttled (and task completion queue isn't growing unbounded) in case when
+    #     reporting stage isn't able to keep up.
+    #
+    #   - Reporting stage dequeues completed tasks from completion queue, reorders
+    #     them (to *always* produce deterministic ordering) and adds its results into
+    #     output queue.
+    #
+    #   - Output queue is capped at `maxsize=max_concurrency` elements to make sure that
+    #     reporting stage is throttled (and output queue doesn't grow unbounded) in case
+    #     when consumer (Ray task itself) isn't able to keep up
+    #
+    async def _execute_transform(it: Iterator[T], output_queue: queue.Queue) -> None:
+        loop = asyncio.get_running_loop()
+
+        # NOTE: Individual tasks could complete in arbitrary order.
+        #       To make sure that the ordering produced by this transformation
+        #       is deterministic we utilize subsequent reordering stage to
+        #       to keep the output ordering the same as that one of the input
+        #       iterator.
+        completed_tasks_queue = asyncio.Queue(maxsize=max_concurrency)
+        # NOTE: This method is nested to support Python 3.9 where we only can
+        #       init `asyncio.Queue` inside the async function
+        async def _reorder() -> None:
+            completed_task_map: Dict[int, asyncio.Task] = dict()
+            next_idx = 0
+            completed_scheduling = False
+
+            try:
+                while not completed_scheduling:
+                    task, idx = await completed_tasks_queue.get()
+
+                    if isinstance(task, Exception):
+                        raise task
+                    elif task is _SENTINEL:
+                        completed_scheduling = True
+                    else:
+                        completed_task_map[idx] = task
+
+                    while next_idx in completed_task_map:
+                        next_task = completed_task_map.pop(next_idx)
+
+                        # NOTE: Once output queue fills up, this will block
+                        #       therefore serving as back-pressure for scheduling tasks
+                        #       preventing it from scheduling new tasks.
+                        # NOTE: This will block the whole event-loop not just this task
+                        output_queue.put(await next_task)
+
+                        next_idx += 1
+
+                assert (
+                    len(completed_task_map) == 0
+                ), f"{next_idx=}, {completed_task_map.keys()=}"
+                sentinel = _SENTINEL
+
+            except BaseException as e:
+                sentinel = e
+            finally:
+                output_queue.put(sentinel)
+
+        # NOTE: Reordering is an async process. Keep a strong reference to
+        # the created task: ``loop.create_task`` only registers a weak
+        # reference with the event loop, so without a strong reference the
+        # task could be garbage collected mid-execution and the reordering
+        # would silently stop.
+        reorder_task = loop.create_task(_reorder())
+
+        cur_task_map: Dict[asyncio.Task, int] = dict()
+        consumed = False
+
+        sentinel = _SENTINEL
+        enumerated_it = enumerate(it)
+
+        try:
+            while True:
+                while len(cur_task_map) < max_concurrency and not consumed:
+                    try:
+                        idx, item = next(enumerated_it)
+                        # Launch async task while keeping track of its
+                        # index in the enumerated sequence
+                        task = loop.create_task(_apply_udf(item))
+                        cur_task_map[task] = idx
+                    except StopIteration:
+                        consumed = True
+                        break
+
+                # Check if any running tasks remaining
+                if not cur_task_map:
+                    break
+
+                done, pending = await asyncio.wait(
+                    cur_task_map.keys(), return_when=asyncio.FIRST_COMPLETED
+                )
+
+                for task in done:
+                    # Report completed tasks along w/ its corresponding
+                    # index in the input sequence
+                    #
+                    # NOTE: Once completed tasks queue fills up, this will block
+                    #       therefore serving as back-pressure for scheduling tasks
+                    #       preventing it from scheduling new tasks
+                    await completed_tasks_queue.put((task, cur_task_map[task]))
+
+                    cur_task_map.pop(task)
+
+        except BaseException as e:
+            for cur_task in cur_task_map:
+                if not cur_task.done():
+                    cur_task.cancel()
+
+            sentinel = e
+        finally:
+            assert len(cur_task_map) == 0, f"{cur_task_map}"
+            await completed_tasks_queue.put((sentinel, None))
+            # Wait for the reorder task to finish draining ``completed_tasks_queue``
+            # and pushing remaining results to the output queue. This both keeps a
+            # strong reference to the task alive until completion (preventing GC)
+            # and surfaces any unexpected exception raised inside ``_reorder``.
+            await reorder_task
+
+    def _transform(batch_iter: Iterable[T], task_context: TaskContext) -> Iterable[U]:
+        output_queue = queue.Queue(maxsize=max_concurrency)
+
+        loop = ray.data._map_actor_context.udf_map_asyncio_loop
+
+        asyncio.run_coroutine_threadsafe(
+            _execute_transform(iter(batch_iter), output_queue), loop
+        )
+
+        while True:
+            items = output_queue.get()
+            if items is _SENTINEL:
+                break
+            elif isinstance(items, Exception):
+                raise items
+            else:
+                # NOTE: Sequences from individual UDFs are combined into a single
+                #       sequence here, as compared to letting individual UDFs to
+                #       add into the output queue to guarantee *deterministic* ordering
+                #       (necessary for Ray Data to be able to guarantee task retries
+                #       producing the same results)
+                for item in items:
+                    validate_fn(item)
+                    yield item
+
+    return _transform
+
+
+class _ReleasingIterator(Iterator[T]):
+    def __init__(self, d: collections.deque):
+        self._d = d
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if not self._d:
+            raise StopIteration
+
+        return self._d.popleft()

python/ray/data/_internal/planner/plan_write_op.py

@@ -0,0 +1,151 @@
+import itertools
+import uuid
+from typing import TYPE_CHECKING, Callable, Iterator, List, Optional, Union
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.datasink import Datasink
+from ray.data.datasource.datasource import Datasource
+
+if TYPE_CHECKING:
+    from ray.data._internal.logical.operators import Write
+
+WRITE_UUID_KWARG_NAME = "write_uuid"
+# Key for storing pending checkpoint paths for commit phase
+PENDING_CHECKPOINTS_KWARG_NAME = "_pending_checkpoints"
+
+
+def generate_write_fn(
+    datasink_or_legacy_datasource: Union[Datasink, Datasource], **write_args
+) -> Callable[[Iterator[Block], TaskContext], Iterator[Block]]:
+    def fn(blocks: Iterator[Block], ctx: TaskContext) -> Iterator[Block]:
+        """Writes the blocks to the given datasink or legacy datasource.
+
+        Outputs the original blocks to be written."""
+        # Create a copy of the iterator, so we can return the original blocks.
+        it1, it2 = itertools.tee(blocks, 2)
+        if isinstance(datasink_or_legacy_datasource, Datasink):
+            ctx.kwargs["_datasink_write_return"] = datasink_or_legacy_datasource.write(
+                it1, ctx
+            )
+        else:
+            datasink_or_legacy_datasource.write(it1, ctx, **write_args)
+
+        return it2
+
+    return fn
+
+
+def generate_collect_write_stats_fn() -> BlockMapTransformFn:
+    # If the write op succeeds, the resulting Dataset is a list of
+    # one Block which contain stats/metrics about the write.
+    # Otherwise, an error will be raised. The Datasource can handle
+    # execution outcomes with `on_write_complete()`` and `on_write_failed()``.
+    def fn(blocks: Iterator[Block], ctx: TaskContext) -> Iterator[Block]:
+        """Handles stats collection for block writes."""
+        block_accessors = [BlockAccessor.for_block(block) for block in blocks]
+        total_num_rows = sum(ba.num_rows() for ba in block_accessors)
+        total_size_bytes = sum(ba.size_bytes() for ba in block_accessors)
+
+        # NOTE: Write tasks can return anything, so we need to wrap it in a valid block
+        # type.
+        import pandas as pd
+
+        block = pd.DataFrame(
+            {
+                "num_rows": [total_num_rows],
+                "size_bytes": [total_size_bytes],
+                "write_return": [ctx.kwargs.get("_datasink_write_return", None)],
+            }
+        )
+        return iter([block])
+
+    return BlockMapTransformFn(
+        fn,
+        is_udf=False,
+        disable_block_shaping=True,
+    )
+
+
+def plan_write_op(
+    op: "Write",
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    collect_stats_fn = generate_collect_write_stats_fn()
+
+    return _plan_write_op_internal(
+        op,
+        physical_children,
+        data_context,
+        post_transformations=[collect_stats_fn],
+    )
+
+
+def _plan_write_op_internal(
+    op: "Write",
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+    post_transformations: List[BlockMapTransformFn],
+    pre_transformations: Optional[List[BlockMapTransformFn]] = None,
+) -> PhysicalOperator:
+    """Plan a write operation with optional pre and post write transformations.
+
+    Args:
+        op: The write operator.
+        physical_children: The physical children operators.
+        data_context: The data context.
+        post_transformations: Transformations to run AFTER the write.
+        pre_transformations: Transformations to run BEFORE the write.
+            Useful for 2-phase commit where pending checkpoint is written first.
+
+    Returns:
+        The physical operator for the write operation.
+    """
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    datasink = op.datasink_or_legacy_datasource
+    write_fn = generate_write_fn(datasink, **op.write_args)
+
+    # Build transform chain: pre_write -> write -> post_write
+    pre_transforms = pre_transformations or []
+    write_transform = BlockMapTransformFn(
+        write_fn,
+        is_udf=False,
+        # NOTE: No need for block-shaping
+        disable_block_shaping=True,
+    )
+    transform_fns = pre_transforms + [write_transform] + post_transformations
+
+    map_transformer = MapTransformer(transform_fns)
+
+    # Set up on_start callback for datasinks.
+    # This allows on_write_start to receive the schema from the first input bundle,
+    # enabling schema-dependent initialization (e.g., Iceberg schema evolution).
+    on_start = None
+    if isinstance(datasink, Datasink):
+        on_start = datasink.on_write_start
+
+    map_op = MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name="Write",
+        # Add a UUID to write tasks to prevent filename collisions. This a UUID for the
+        # overall write operation, not the individual write tasks.
+        map_task_kwargs={WRITE_UUID_KWARG_NAME: uuid.uuid4().hex},
+        ray_remote_args=op.ray_remote_args,
+        min_rows_per_bundle=op.min_rows_per_bundled_input,
+        compute_strategy=op.compute,
+        on_start=on_start,
+    )
+
+    return map_op

python/ray/data/_internal/planner/planner.py

@@ -0,0 +1,448 @@
+import warnings
+from functools import partial
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, TypeVar
+
+if TYPE_CHECKING:
+    import pyarrow.fs
+
+from ray.data._internal.execution.execution_callback import ExecutionCallback
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.aggregate_num_rows import (
+    AggregateNumRows,
+)
+from ray.data._internal.execution.operators.hash_shuffle_v2 import (
+    _SHUFFLE_MAP_RUNTIME_ENV,
+    _make_hash_partition_fn,
+)
+from ray.data._internal.execution.operators.input_data_buffer import (
+    InputDataBuffer,
+)
+from ray.data._internal.execution.operators.join import (
+    JoinOperator,
+    _make_join_reduce_fn,
+)
+from ray.data._internal.execution.operators.limit_operator import LimitOperator
+from ray.data._internal.execution.operators.mix_operator import MixOperator
+from ray.data._internal.execution.operators.output_splitter import OutputSplitter
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (
+    ShuffleMapOp,
+)
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import (
+    ShuffleReduceOp,
+)
+from ray.data._internal.execution.operators.union_operator import UnionOperator
+from ray.data._internal.execution.operators.zip_operator import ZipOperator
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalPlan,
+    PhysicalPlan,
+)
+from ray.data._internal.logical.operators import (
+    AbstractAllToAll,
+    AbstractFrom,
+    AbstractUDFMap,
+    Count,
+    Download,
+    Filter,
+    InputData,
+    Join,
+    JoinType,
+    Limit,
+    ListFiles,
+    Mix,
+    Project,
+    Read,
+    ReadFiles,
+    StreamingRepartition,
+    StreamingSplit,
+    Union,
+    Write,
+    Zip,
+)
+from ray.data._internal.planner.checkpoint import (
+    plan_read_files_op_with_checkpoint_filter,
+    plan_read_op_with_checkpoint_filter,
+    plan_write_op_with_checkpoint_writer,
+)
+from ray.data._internal.planner.plan_all_to_all_op import plan_all_to_all_op
+from ray.data._internal.planner.plan_download_op import plan_download_op
+from ray.data._internal.planner.plan_list_files_op import plan_list_files_op
+from ray.data._internal.planner.plan_read_files_op import plan_read_files_op
+from ray.data._internal.planner.plan_read_op import plan_read_op
+from ray.data._internal.planner.plan_udf_map_op import (
+    plan_filter_op,
+    plan_project_op,
+    plan_streaming_repartition_op,
+    plan_udf_map_op,
+)
+from ray.data._internal.planner.plan_write_op import plan_write_op
+from ray.data._internal.usage import create_usage_callback
+from ray.data.checkpoint.load_checkpoint_callback import LoadCheckpointCallback
+from ray.data.context import DataContext
+from ray.data.datasource.file_datasink import _FileDatasink
+
+LogicalOperatorType = TypeVar("LogicalOperatorType", bound=LogicalOperator)
+PlanLogicalOpFn = Callable[
+    [LogicalOperatorType, List[PhysicalOperator], DataContext], PhysicalOperator
+]
+
+
+def plan_input_data_op(
+    logical_op: InputData,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Get the corresponding DAG of physical operators for InputData."""
+    assert len(physical_children) == 0
+
+    return InputDataBuffer(
+        data_context,
+        input_data=logical_op.input_data,
+    )
+
+
+def plan_from_op(
+    op: AbstractFrom,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    assert len(physical_children) == 0
+    return InputDataBuffer(data_context, op.input_data)
+
+
+def plan_zip_op(_, physical_children, data_context):
+    assert len(physical_children) >= 2
+    return ZipOperator(data_context, *physical_children)
+
+
+def plan_mix_op(logical_op, physical_children, data_context):
+    assert len(physical_children) >= 1
+    return MixOperator(
+        data_context,
+        *physical_children,
+        weights=logical_op.weights,
+        stopping_condition=logical_op.stopping_condition,
+    )
+
+
+def plan_union_op(_, physical_children, data_context):
+    assert len(physical_children) >= 2
+    return UnionOperator(data_context, *physical_children)
+
+
+def plan_limit_op(logical_op, physical_children, data_context):
+    assert len(physical_children) == 1
+    return LimitOperator(logical_op.limit, physical_children[0], data_context)
+
+
+def plan_count_op(logical_op, physical_children, data_context):
+    assert len(physical_children) == 1
+    return AggregateNumRows(
+        [physical_children[0]], data_context, column_name=Count.COLUMN_NAME
+    )
+
+
+def _plan_join_shuffle_v2(
+    logical_op: Join,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    left_keys = list(logical_op.left_key_columns)
+    right_keys = list(logical_op.right_key_columns)
+    num_partitions = logical_op.num_partitions
+    join_type = JoinType(logical_op.join_type)
+
+    left_map = ShuffleMapOp(
+        physical_children[0],
+        data_context,
+        num_partitions=num_partitions,
+        partition_fn=_make_hash_partition_fn(left_keys, num_partitions),
+        map_runtime_env=_SHUFFLE_MAP_RUNTIME_ENV,
+        name=f"JoinShuffleMapLeft(keys={tuple(left_keys)}, parts={num_partitions})",
+    )
+    right_map = ShuffleMapOp(
+        physical_children[1],
+        data_context,
+        num_partitions=num_partitions,
+        partition_fn=_make_hash_partition_fn(right_keys, num_partitions),
+        map_runtime_env=_SHUFFLE_MAP_RUNTIME_ENV,
+        name=f"JoinShuffleMapRight(keys={tuple(right_keys)}, parts={num_partitions})",
+    )
+
+    reduce_fn = _make_join_reduce_fn(
+        join_type=join_type,
+        left_key_col_names=tuple(left_keys),
+        right_key_col_names=tuple(right_keys),
+        left_columns_suffix=logical_op.left_columns_suffix,
+        right_columns_suffix=logical_op.right_columns_suffix,
+        left_schema=logical_op.input_dependencies[0].infer_schema(),
+        right_schema=logical_op.input_dependencies[1].infer_schema(),
+    )
+    return ShuffleReduceOp(
+        [left_map, right_map],
+        data_context,
+        num_partitions=num_partitions,
+        reduce_fn=reduce_fn,
+        disallow_block_splitting=False,
+        reduce_ray_remote_args=logical_op.aggregator_ray_remote_args,
+        name=f"JoinShuffleReduce(num_partitions={num_partitions})",
+    )
+
+
+def plan_join_op(
+    logical_op: Join,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    assert len(physical_children) == 2
+    if data_context.use_hash_shuffle_v2:
+        return _plan_join_shuffle_v2(logical_op, physical_children, data_context)
+    return JoinOperator(
+        data_context=data_context,
+        left_input_op=physical_children[0],
+        right_input_op=physical_children[1],
+        join_type=logical_op.join_type,
+        left_key_columns=logical_op.left_key_columns,
+        right_key_columns=logical_op.right_key_columns,
+        left_columns_suffix=logical_op.left_columns_suffix,
+        right_columns_suffix=logical_op.right_columns_suffix,
+        num_partitions=logical_op.num_outputs,
+        partition_size_hint=logical_op.partition_size_hint,
+        aggregator_ray_remote_args_override=logical_op.aggregator_ray_remote_args,
+    )
+
+
+def plan_streaming_split_op(
+    logical_op: StreamingSplit,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+):
+    assert len(physical_children) == 1
+    return OutputSplitter(
+        physical_children[0],
+        n=logical_op.num_splits,
+        equal=logical_op.equal,
+        data_context=data_context,
+        locality_hints=logical_op.locality_hints,
+    )
+
+
+class Planner:
+    """The planner to convert optimized logical to physical operators.
+
+    Note that planner is only doing operators conversion. Physical optimization work is
+    done by physical optimizer.
+    """
+
+    _DEFAULT_PLAN_FNS = {
+        Read: plan_read_op,
+        ReadFiles: plan_read_files_op,
+        ListFiles: plan_list_files_op,
+        InputData: plan_input_data_op,
+        Write: plan_write_op,
+        AbstractFrom: plan_from_op,
+        Filter: plan_filter_op,
+        AbstractUDFMap: plan_udf_map_op,
+        AbstractAllToAll: plan_all_to_all_op,
+        Mix: plan_mix_op,
+        Union: plan_union_op,
+        Zip: plan_zip_op,
+        Limit: plan_limit_op,
+        Count: plan_count_op,
+        Project: plan_project_op,
+        StreamingRepartition: plan_streaming_repartition_op,
+        Join: plan_join_op,
+        StreamingSplit: plan_streaming_split_op,
+        Download: plan_download_op,
+    }
+    # Operators that support checkpoint filtering. Subclasses can override.
+    _CHECKPOINT_FILTER_OPS = (Read, ReadFiles)
+
+    def __init__(self):
+        self._supports_checkpointing = False
+        self._plan_fns_for_checkpointing = {}
+
+    def plan(
+        self, logical_plan: LogicalPlan
+    ) -> Tuple[PhysicalPlan, List["ExecutionCallback"]]:
+        """Convert logical to physical operators recursively in post-order."""
+        checkpoint_config = logical_plan.context.checkpoint_config
+
+        callbacks = [cls() for cls in logical_plan.context.execution_callback_classes]
+        callbacks.append(create_usage_callback(logical_plan))
+
+        if checkpoint_config is not None and self._check_supports_checkpointing(
+            logical_plan
+        ):
+            self._supports_checkpointing = True
+            data_file_dir, data_file_fs = self._get_data_file_info(logical_plan)
+
+            checkpoint_callback = self._create_checkpoint_callback(
+                checkpoint_config,
+            )
+
+            callbacks.append(checkpoint_callback)
+
+            # Dynamically set the plan functions for checkpointing because they
+            # need to a reference to the checkpoint ref.
+            self._plan_fns_for_checkpointing = self._get_plan_fns_for_checkpointing(
+                data_file_dir, data_file_fs
+            )
+
+        elif checkpoint_config is not None:
+            assert not self._check_supports_checkpointing(logical_plan)
+            warnings.warn(
+                "You've enabled checkpointing, but the logical plan doesn't support "
+                "checkpointing. Checkpointing will be disabled."
+            )
+        physical_dag, op_map = self._plan_recursively(
+            logical_plan.dag, logical_plan.context
+        )
+        physical_plan = PhysicalPlan(physical_dag, op_map, logical_plan.context)
+        return physical_plan, callbacks
+
+    def get_plan_fn(self, logical_op: LogicalOperator) -> PlanLogicalOpFn:
+        if self._supports_checkpointing:
+            assert self._plan_fns_for_checkpointing
+            plan_fn = find_plan_fn(logical_op, self._plan_fns_for_checkpointing)
+            if plan_fn is not None:
+                return plan_fn
+
+        plan_fn = find_plan_fn(logical_op, self._DEFAULT_PLAN_FNS)
+        if plan_fn is not None:
+            return plan_fn
+
+        raise ValueError(
+            f"Found unknown logical operator during planning: {logical_op}"
+        )
+
+    def _plan_recursively(
+        self, logical_op: LogicalOperator, data_context: DataContext
+    ) -> Tuple[PhysicalOperator, Dict[LogicalOperator, PhysicalOperator]]:
+        """Plan a logical operator and its input dependencies recursively.
+
+        Args:
+            logical_op: The logical operator to plan.
+            data_context: The data context.
+
+        Returns:
+            A tuple of the physical operator corresponding to the logical operator, and
+            a mapping from physical to logical operators.
+        """
+        op_map: Dict[PhysicalOperator, LogicalOperator] = {}
+
+        # Plan the input dependencies first.
+        physical_children = []
+        for child in logical_op.input_dependencies:
+            physical_child, child_op_map = self._plan_recursively(child, data_context)
+            physical_children.append(physical_child)
+            op_map.update(child_op_map)
+
+        plan_fn = self.get_plan_fn(logical_op)
+        # We will call `set_logical_operators()` in the following for-loop,
+        # no need to do it here.
+        physical_op = plan_fn(logical_op, physical_children, data_context)
+
+        # Traverse up the DAG, and set the mapping from physical to logical operators.
+        # At this point, all physical operators without logical operators set
+        # must have been created by the current logical operator.
+        queue = [physical_op]
+        while queue:
+            curr_physical_op = queue.pop()
+            if curr_physical_op._logical_operators:
+                continue
+
+            curr_physical_op.set_logical_operators(logical_op)
+            # Add this operator to the op_map so optimizer can find it
+            op_map[curr_physical_op] = logical_op
+            queue.extend(curr_physical_op.input_dependencies)
+
+        # Also add the final operator (in case the loop didn't catch it)
+        op_map[physical_op] = logical_op
+        return physical_op, op_map
+
+    def _create_checkpoint_callback(
+        self,
+        checkpoint_config,
+    ) -> LoadCheckpointCallback:
+        """Factory method to create the LoadCheckpointCallback.
+
+        Subclasses can override this to use a different callback implementation.
+        """
+        return LoadCheckpointCallback(
+            checkpoint_config,
+        )
+
+    @staticmethod
+    def _get_data_file_info(logical_plan: LogicalPlan):
+        """Extract the data file directory and filesystem from the Write op's datasink.
+
+        Returns (path, filesystem) for file-based datasinks, or (None, None)
+        for non-file datasinks.
+        """
+        last_op = logical_plan.dag
+        if isinstance(last_op, Write):
+            datasink = last_op.datasink_or_legacy_datasource
+            if isinstance(datasink, _FileDatasink):
+                return datasink.unresolved_path, datasink.filesystem
+        return None, None
+
+    def _get_plan_fns_for_checkpointing(
+        self,
+        data_file_dir: Optional[str] = None,
+        data_file_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+    ) -> Dict[Type[LogicalOperator], PlanLogicalOpFn]:
+        plan_fns = {
+            Read: partial(
+                plan_read_op_with_checkpoint_filter, data_file_dir, data_file_filesystem
+            ),
+            ReadFiles: partial(
+                plan_read_files_op_with_checkpoint_filter,
+                data_file_dir,
+                data_file_filesystem,
+            ),
+            Write: plan_write_op_with_checkpoint_writer,
+        }
+        return plan_fns
+
+    def _check_supports_checkpointing(self, logical_plan: LogicalPlan) -> bool:
+        """Check if the logical plan supports checkpointing.
+
+        Subclasses can override _CHECKPOINT_FILTER_OPS to support more operators.
+        """
+        if not isinstance(logical_plan.dag, (Write, StreamingSplit)):
+            return False
+
+        def _all_paths_contain_checkpoint_filter(op: LogicalOperator) -> bool:
+            if isinstance(op, self._CHECKPOINT_FILTER_OPS):
+                return True
+            return all(
+                _all_paths_contain_checkpoint_filter(input_dep)
+                for input_dep in op.input_dependencies
+            )
+
+        return _all_paths_contain_checkpoint_filter(logical_plan.dag)
+
+
+def find_plan_fn(
+    logical_op: LogicalOperator, plan_fns: Dict[Type[LogicalOperator], PlanLogicalOpFn]
+) -> Optional[PlanLogicalOpFn]:
+    """Find the plan function for a logical operator.
+
+    This function goes through the plan functions in order and returns the first one
+    that is an instance of the logical operator type.
+
+    Args:
+        logical_op: The logical operator to find the plan function for.
+        plan_fns: The dictionary of plan functions.
+
+    Returns:
+        The plan function for the logical operator, or None if no plan function is
+        found.
+    """
+    # TODO: This implementation doesn't account for type hierarchies conflicts or
+    # multiple inheritance.
+    for op_type, plan_fn in plan_fns.items():
+        if isinstance(logical_op, op_type):
+            return plan_fn
+    return None

python/ray/data/_internal/planner/random_shuffle.py

@@ -0,0 +1,101 @@
+from typing import Any, Dict, List, Optional
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.execution.operators.map_transformer import MapTransformer
+from ray.data._internal.execution.util import merge_label_selector
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
+from ray.data._internal.random_config import (
+    RandomSeedConfig,
+    get_single_integer_random_seed,
+)
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_random_shuffle_fn(
+    data_context: DataContext,
+    seed_config: RandomSeedConfig,
+    num_outputs: Optional[int] = None,
+    ray_remote_args: Optional[Dict[str, Any]] = None,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to randomly shuffle each records of blocks."""
+
+    # If no seed has been specified, pin timestamp based one
+    # so that task could be safely retried (w/o changing their output)
+    seed = get_single_integer_random_seed(seed_config, data_context)
+
+    def fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        num_input_blocks = sum(len(r.blocks) for r in refs)
+
+        # If map_transformer is specified (e.g. from fusing
+        # MapOperator->AllToAllOperator), we pass a map function which
+        # is applied to each block before shuffling.
+        map_transformer: Optional[MapTransformer] = ctx.upstream_map_transformer
+        upstream_map_fn = None
+        nonlocal ray_remote_args
+        if map_transformer:
+            # NOTE: We override target max-block sizing of the previous
+            #       transformation to avoid unnecessary block shaping (if any)
+            map_transformer.override_target_max_block_size(None)
+
+            def upstream_map_fn(blocks):
+                DataContext._set_current(data_context)
+                return map_transformer.apply_transform(blocks, ctx)
+
+            # If there is a fused upstream operator,
+            # also use the ray_remote_args from the fused upstream operator.
+            ray_remote_args = ctx.upstream_map_ray_remote_args
+
+        shuffle_spec = ShuffleTaskSpec(
+            target_shuffle_max_block_size=(
+                ctx.target_max_block_size_override or data_context.target_max_block_size
+            ),
+            random_shuffle=True,
+            random_seed=seed,
+            upstream_map_fn=upstream_map_fn,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            if num_outputs is not None:
+                raise NotImplementedError(
+                    "Push-based shuffle doesn't support setting num_blocks yet."
+                )
+            scheduler = PushBasedShuffleTaskScheduler(shuffle_spec)
+        else:
+            scheduler = PullBasedShuffleTaskScheduler(shuffle_spec)
+
+        label_selector = data_context.execution_options.label_selector
+        map_ray_remote_args = merge_label_selector(
+            ray_remote_args or {}, label_selector
+        )
+        reduce_ray_remote_args = merge_label_selector(
+            ray_remote_args or {}, label_selector
+        )
+        return scheduler.execute(
+            refs,
+            num_outputs or num_input_blocks,
+            task_ctx=ctx,
+            map_ray_remote_args=map_ray_remote_args,
+            reduce_ray_remote_args=reduce_ray_remote_args,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    return fn

python/ray/data/_internal/planner/randomize_blocks.py

@@ -0,0 +1,60 @@
+from typing import List
+
+import numpy as np
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    BlockEntry,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.logical.operators import RandomizeBlocks
+from ray.data._internal.random_config import get_single_integer_random_seed
+from ray.data.context import DataContext
+
+
+def generate_randomize_blocks_fn(
+    op: RandomizeBlocks,
+    data_context: DataContext,
+) -> AllToAllTransformFn:
+    """Generate function to randomize order of blocks."""
+
+    seed = get_single_integer_random_seed(op.seed_config, data_context)
+
+    def fn(
+        refs: List[RefBundle],
+        context: TaskContext,
+    ) -> AllToAllTransformFnResult:
+
+        nonlocal op
+        blocks_with_metadata = []
+        index_to_schema = [None] * len(refs)
+        for i, ref_bundle in enumerate(refs):
+            index_to_schema[i] = ref_bundle.schema
+            blocks_with_metadata.extend(
+                (entry.ref, entry.metadata, i) for entry in ref_bundle.blocks
+            )
+
+        if len(blocks_with_metadata) == 0:
+            return refs, {op.name: []}
+        else:
+            rng = np.random.default_rng(seed)
+            input_owned = all(b.owns_blocks for b in refs)
+            rng.shuffle(blocks_with_metadata)
+            output = []
+            stats_list = []
+            for block, meta, i in blocks_with_metadata:
+                stats_list.append(meta.to_stats())
+                output.append(
+                    RefBundle(
+                        [BlockEntry(block, meta)],
+                        owns_blocks=input_owned,
+                        schema=index_to_schema[i],
+                    )
+                )
+            return output, {op.name: stats_list}
+
+    return fn

python/ray/data/_internal/planner/repartition.py

@@ -0,0 +1,104 @@
+from typing import List, Optional, Tuple
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.execution.operators.map_transformer import MapTransformer
+from ray.data._internal.execution.util import merge_label_selector
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
+from ray.data._internal.planner.exchange.split_repartition_task_scheduler import (
+    SplitRepartitionTaskScheduler,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_repartition_fn(
+    num_outputs: int,
+    shuffle: bool,
+    data_context: DataContext,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to partition each records of blocks."""
+
+    def shuffle_repartition_fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        # If map_transformer is specified (e.g. from fusing
+        # MapOperator->AllToAllOperator), we pass a map function which
+        # is applied to each block before shuffling.
+        map_transformer: Optional["MapTransformer"] = ctx.upstream_map_transformer
+        upstream_map_fn = None
+        if map_transformer:
+            # NOTE: We override target max-block sizing of the previous
+            #       transformation to avoid unnecessary block shaping (if any)
+            map_transformer.override_target_max_block_size(None)
+
+            def upstream_map_fn(blocks):
+                DataContext._set_current(data_context)
+                return map_transformer.apply_transform(blocks, ctx)
+
+        shuffle_spec = ShuffleTaskSpec(
+            target_shuffle_max_block_size=(
+                ctx.target_max_block_size_override or data_context.target_max_block_size
+            ),
+            random_shuffle=False,
+            upstream_map_fn=upstream_map_fn,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            scheduler = PushBasedShuffleTaskScheduler(shuffle_spec)
+        else:
+            scheduler = PullBasedShuffleTaskScheduler(shuffle_spec)
+
+        label_selector = data_context.execution_options.label_selector
+        map_ray_remote_args = merge_label_selector({}, label_selector)
+        reduce_ray_remote_args = merge_label_selector({}, label_selector)
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            map_ray_remote_args=map_ray_remote_args,
+            reduce_ray_remote_args=reduce_ray_remote_args,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    def split_repartition_fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        shuffle_spec = ShuffleTaskSpec(
+            target_shuffle_max_block_size=(
+                ctx.target_max_block_size_override or data_context.target_max_block_size
+            ),
+            random_shuffle=False,
+        )
+        scheduler = SplitRepartitionTaskScheduler(shuffle_spec)
+        label_selector = data_context.execution_options.label_selector
+        map_ray_remote_args = merge_label_selector({}, label_selector)
+        reduce_ray_remote_args = merge_label_selector({}, label_selector)
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            map_ray_remote_args=map_ray_remote_args,
+            reduce_ray_remote_args=reduce_ray_remote_args,
+        )
+
+    if shuffle:
+        return shuffle_repartition_fn
+    return split_repartition_fn

python/ray/data/_internal/planner/sort.py

@@ -0,0 +1,90 @@
+from functools import partial
+from typing import List, Optional, Tuple
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.util import merge_label_selector
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey, SortTaskSpec
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import unify_ref_bundles_schema
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_sort_fn(
+    sort_key: SortKey,
+    data_context: DataContext,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to sort blocks by the specified key column or key function."""
+
+    def fn(
+        sort_key: SortKey,
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        blocks = []
+        for ref_bundle in refs:
+            blocks.extend(ref_bundle.block_refs)
+        if len(blocks) == 0:
+            return (blocks, {})
+
+        sort_key.validate_schema(unify_ref_bundles_schema(refs))
+
+        num_mappers = len(blocks)
+        # Use same number of output partitions.
+        num_outputs = num_mappers
+
+        label_selector = data_context.execution_options.label_selector
+        # Sample boundaries for sort key.
+        if not sort_key.boundaries:
+            sample_bar = ctx.sub_progress_bar_dict[
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME
+            ]
+            boundaries = SortTaskSpec.sample_boundaries(
+                blocks,
+                sort_key,
+                num_outputs,
+                sample_bar,
+                label_selector=label_selector,
+            )
+        else:
+            # For user-specified boundaries (which only partition by the primary
+            # sort key), reverse `boundaries` so that the partitions are produced
+            # in descending order, as desired.
+            boundaries = [(b,) for b in sort_key.boundaries]
+            if sort_key.get_descending()[0]:
+                boundaries = boundaries[::-1]
+            num_outputs = len(boundaries) + 1
+        sort_spec = SortTaskSpec(boundaries=boundaries, sort_key=sort_key)
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            scheduler = PushBasedShuffleTaskScheduler(sort_spec)
+        else:
+            scheduler = PullBasedShuffleTaskScheduler(sort_spec)
+
+        map_ray_remote_args = merge_label_selector({}, label_selector)
+        reduce_ray_remote_args = merge_label_selector({}, label_selector)
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            map_ray_remote_args=map_ray_remote_args,
+            reduce_ray_remote_args=reduce_ray_remote_args,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    # NOTE: use partial function to pass parameters to avoid error like
+    # "UnboundLocalError: local variable ... referenced before assignment",
+    # because `key` and `descending` variables are reassigned in `fn()`.
+    return partial(fn, sort_key)