chore: import upstream snapshot with attribution

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

@@ -0,0 +1,14 @@
+def get_task_pool_map_operator_cls():
+    from ray.data._internal.execution.operators.task_pool_map_operator import (
+        TaskPoolMapOperator,
+    )
+
+    return TaskPoolMapOperator
+
+
+def get_actor_pool_map_operator_cls():
+    from ray.data._internal.execution.operators.actor_pool_map_operator import (
+        ActorPoolMapOperator,
+    )
+
+    return ActorPoolMapOperator

python/ray/data/_internal/execution/operators/aggregate_num_rows.py

@@ -0,0 +1,70 @@
+import ray
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.block import BlockAccessor
+from ray.data.context import DataContext
+
+
+class AggregateNumRows(PhysicalOperator):
+    """Count number of rows in input bundles.
+
+    This operator aggregates the number of rows in input bundles using the bundles'
+    block metadata. It outputs a single row with the specified column name.
+    """
+
+    def __init__(
+        self,
+        input_dependencies,
+        data_context: DataContext,
+        column_name: str,
+    ):
+        super().__init__(
+            "AggregateNumRows",
+            input_dependencies,
+            data_context,
+        )
+
+        self._column_name = column_name
+
+        self._num_rows = 0
+        self._has_outputted = False
+        self._estimated_num_output_bundles = 1
+        self._estimated_output_num_rows = 1
+
+    def has_next(self) -> bool:
+        return self._inputs_complete and not self._has_outputted
+
+    def _get_next_inner(self) -> RefBundle:
+        assert self._inputs_complete
+
+        builder = DelegatingBlockBuilder()
+        builder.add({self._column_name: self._num_rows})
+        block = builder.build()
+        block_ref = ray.put(block)
+
+        metadata = BlockAccessor.for_block(block).get_metadata()
+        schema = BlockAccessor.for_block(block).schema()
+        bundle = RefBundle(
+            [BlockEntry(block_ref, metadata)], owns_blocks=True, schema=schema
+        )
+
+        self._block_ref_counter.on_block_produced(
+            block_ref, metadata.size_bytes or 0, self.id
+        )
+        self._has_outputted = True
+        return bundle
+
+    def get_stats(self) -> StatsDict:
+        return {}
+
+    def _add_input_inner(self, refs, input_index) -> None:
+        assert refs.num_rows() is not None
+        self._num_rows += refs.num_rows()
+
+    def throttling_disabled(self) -> bool:
+        return True

python/ray/data/_internal/execution/operators/base_physical_operator.py

@@ -0,0 +1,268 @@
+import abc
+import typing
+from typing import List, Optional
+
+from typing_extensions import override
+
+from ray.data._internal.execution.bundle_queue import BaseBundleQueue, FIFOBundleQueue
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+if typing.TYPE_CHECKING:
+    from ray.data._internal.progress.base_progress import BaseProgressBar
+
+
+class InternalQueueOperatorMixin(PhysicalOperator, abc.ABC):
+    @property
+    @abc.abstractmethod
+    def _input_queues(self) -> List["BaseBundleQueue"]:
+        """Return all the internal input buffer queues for this operator."""
+        ...
+
+    @property
+    @abc.abstractmethod
+    def _output_queues(self) -> List["BaseBundleQueue"]:
+        """Return all the internal output buffer queues for this operator."""
+        ...
+
+    def internal_input_queue_num_blocks(self) -> int:
+        """Returns Operator's internal input queue size (in blocks)"""
+        return sum(input_buffer.num_blocks() for input_buffer in self._input_queues)
+
+    def internal_input_queue_num_bytes(self) -> int:
+        """Returns Operator's internal input queue size (in bytes)"""
+        return sum(
+            input_buffer.estimate_size_bytes() for input_buffer in self._input_queues
+        )
+
+    def internal_output_queue_num_blocks(self) -> int:
+        """Returns Operator's internal output queue size (in blocks)"""
+        return sum(output_buffer.num_blocks() for output_buffer in self._output_queues)
+
+    def internal_output_queue_num_bytes(self) -> int:
+        """Returns Operator's internal output queue size (in bytes)"""
+        return sum(
+            output_buffer.estimate_size_bytes() for output_buffer in self._output_queues
+        )
+
+    def clear_internal_input_queue(self) -> None:
+        """Clear internal input queue(s)."""
+        for input_buffer in self._input_queues:
+            input_buffer.clear()
+
+    def clear_internal_output_queue(self) -> None:
+        """Clear internal output queue(s)."""
+        for output_buffer in self._output_queues:
+            output_buffer.clear()
+
+    def mark_execution_finished(self) -> None:
+        """Mark execution as finished and clear internal queues.
+
+        This default implementation calls the parent's mark_execution_finished()
+        and then clears internal input and output queues.
+        """
+        super().mark_execution_finished()
+        self.clear_internal_input_queue()
+        self.clear_internal_output_queue()
+
+
+class OneToOneOperator(PhysicalOperator):
+    """An operator that has one input and one output dependency.
+
+    This operator serves as the base for map, filter, limit, etc.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        target_max_block_size_override: Optional[int] = None,
+    ):
+        """Create a OneToOneOperator.
+
+        Args:
+            name: The name of this operator.
+            input_op: Operator generating input data for this op.
+            data_context: The :class:`DataContext` to use for this operator.
+            target_max_block_size_override: The target maximum number of bytes to
+                include in an output block.
+        """
+        super().__init__(name, [input_op], data_context, target_max_block_size_override)
+
+    @property
+    def input_dependency(self) -> PhysicalOperator:
+        return self.input_dependencies[0]
+
+
+class AllToAllOperator(
+    InternalQueueOperatorMixin, SubProgressBarMixin, PhysicalOperator
+):
+    """A blocking operator that executes once its inputs are complete.
+
+    This operator implements distributed sort / shuffle operations, etc.
+    """
+
+    def __init__(
+        self,
+        bulk_fn: AllToAllTransformFn,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        target_max_block_size_override: Optional[int] = None,
+        num_outputs: Optional[int] = None,
+        sub_progress_bar_names: Optional[List[str]] = None,
+        name: str = "AllToAll",
+    ):
+        """Create an AllToAllOperator.
+        Args:
+            bulk_fn: The blocking transformation function to run. The inputs are the
+                list of input ref bundles, and the outputs are the output ref bundles
+                and a stats dict.
+            input_op: Operator generating input data for this op.
+            data_context: The DataContext instance containing configuration settings.
+            target_max_block_size_override: The target maximum number of bytes to
+                include in an output block.
+            num_outputs: The number of expected output bundles for progress bar.
+            sub_progress_bar_names: The names of internal sub progress bars.
+            name: The name of this operator.
+        """
+        self._bulk_fn = bulk_fn
+        self._next_task_index = 0
+        self._num_outputs = num_outputs
+        self._output_rows = 0
+        self._sub_progress_bar_names = sub_progress_bar_names
+        self._sub_progress_bar_dict = None
+        self._input_buffer: FIFOBundleQueue = FIFOBundleQueue()
+        self._output_buffer: FIFOBundleQueue = FIFOBundleQueue()
+        self._stats: StatsDict = {}
+        super().__init__(name, [input_op], data_context, target_max_block_size_override)
+
+    @property
+    @override
+    def _input_queues(self) -> List["BaseBundleQueue"]:
+        return [self._input_buffer]
+
+    @property
+    @override
+    def _output_queues(self) -> List["BaseBundleQueue"]:
+        return [self._output_buffer]
+
+    def num_outputs_total(self) -> Optional[int]:
+        return (
+            self._num_outputs
+            if self._num_outputs
+            else self.input_dependencies[0].num_outputs_total()
+        )
+
+    def num_output_rows_total(self) -> Optional[int]:
+        return (
+            self._output_rows
+            if self._output_rows
+            else self.input_dependencies[0].num_output_rows_total()
+        )
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.has_completed()
+        assert input_index == 0, input_index
+        self._input_buffer.add(refs)
+        self._metrics.on_input_queued(refs, input_index=0)
+
+    def all_inputs_done(self) -> None:
+        ctx = TaskContext(
+            task_idx=self._next_task_index,
+            op_name=self.name,
+            sub_progress_bar_dict=self._sub_progress_bar_dict,
+            target_max_block_size_override=self.target_max_block_size_override,
+        )
+        # NOTE: We don't account object store memory use from intermediate `bulk_fn`
+        # outputs (e.g., map outputs for map-reduce).
+
+        input_bundles = self._input_buffer.to_list()
+        output_buffer, self._stats = self._bulk_fn(input_bundles, ctx)
+        self._output_buffer = FIFOBundleQueue(output_buffer)
+
+        for bundle in output_buffer:
+            for entry in bundle.blocks:
+                self._block_ref_counter.on_block_produced(
+                    entry.ref, entry.metadata.size_bytes or 0, self.id
+                )
+
+        while self._input_buffer.has_next():
+            refs = self._input_buffer.get_next()
+            self._metrics.on_input_dequeued(refs, input_index=0)
+
+        for ref in self._output_buffer:
+            self._metrics.on_output_queued(ref)
+
+        self._next_task_index += 1
+
+        super().all_inputs_done()
+
+    def has_next(self) -> bool:
+        return len(self._output_buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        bundle = self._output_buffer.get_next()
+        self._metrics.on_output_dequeued(bundle)
+        self._output_rows += bundle.num_rows()
+        return bundle
+
+    def get_stats(self) -> StatsDict:
+        return self._stats
+
+    def get_transformation_fn(self) -> AllToAllTransformFn:
+        return self._bulk_fn
+
+    def progress_str(self) -> str:
+        return f"{self.num_output_rows_total() or 0} rows output"
+
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        return self._sub_progress_bar_names
+
+    def set_sub_progress_bar(self, name: str, pg: "BaseProgressBar"):
+        if self._sub_progress_bar_dict is None:
+            self._sub_progress_bar_dict = {}
+        self._sub_progress_bar_dict[name] = pg
+
+    def supports_fusion(self):
+        return True
+
+    def throttling_disabled(self) -> bool:
+        # Disable resource allocation and throttling for the operator
+        return True
+
+
+class NAryOperator(PhysicalOperator):
+    """An operator that has multiple input dependencies and one output.
+
+    This operator serves as the base for union, zip, etc.
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        *input_ops: PhysicalOperator,
+        name: Optional[str] = None,
+    ):
+        """Create a NAryOperator.
+
+        Args:
+            data_context: The DataContext instance containing configuration settings.
+            *input_ops: Operators generating input data for this op.
+            name: Optional override for the operator display name.
+        """
+        if name is None:
+            input_names = ", ".join([op._name for op in input_ops])
+            name = f"{self.__class__.__name__}({input_names})"
+        super().__init__(
+            name,
+            list(input_ops),
+            data_context,
+        )

python/ray/data/_internal/execution/operators/hash_aggregate.py

@@ -0,0 +1,219 @@
+import logging
+import math
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.hash_shuffle import (
+    BlockTransformer,
+    HashShufflingOperatorBase,
+    ShuffleAggregation,
+)
+from ray.data._internal.util import GiB, MiB
+from ray.data.aggregate import AggregateFn
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+
+logger = logging.getLogger(__name__)
+
+
+class ReducingAggregation(ShuffleAggregation):
+    """Stateless aggregation that reduces blocks using aggregation functions.
+
+    This implementation performs incremental reduction during compaction,
+    combining multiple partially-aggregated blocks into one. The final
+    aggregation is performed during finalization.
+    """
+
+    def __init__(
+        self,
+        key_columns: Tuple[str, ...],
+        aggregation_fns: Tuple[AggregateFn, ...],
+    ):
+        self._sort_key: "SortKey" = self._get_sort_key(key_columns)
+        self._aggregation_fns: Tuple[AggregateFn, ...] = aggregation_fns
+
+    @classmethod
+    def is_compacting(cls):
+        return True
+
+    def compact(self, partition_shards: List[Block]) -> Block:
+        assert len(partition_shards) > 0, "Provided sequence must be non-empty"
+
+        return self._combine(partition_shards, finalize=False)
+
+    def finalize(self, partition_shards_map: Dict[int, List[Block]]) -> Iterator[Block]:
+        assert (
+            len(partition_shards_map) == 1
+        ), f"Single input-sequence is expected (got {len(partition_shards_map)})"
+
+        blocks = partition_shards_map[0]
+
+        if not blocks:
+            return
+
+        yield self._combine(blocks, finalize=True)
+
+    def _combine(self, blocks: List[Block], *, finalize: bool) -> Block:
+        """Internal method to combine blocks with optional finalization."""
+        assert len(blocks) > 0
+
+        block_accessor = BlockAccessor.for_block(blocks[0])
+        combined_block, _ = block_accessor._combine_aggregated_blocks(
+            blocks,
+            sort_key=self._sort_key,
+            aggs=self._aggregation_fns,
+            finalize=finalize,
+        )
+
+        return combined_block
+
+    @staticmethod
+    def _get_sort_key(key_columns: Tuple[str, ...]) -> "SortKey":
+        from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+        return SortKey(key=list(key_columns), descending=False)
+
+
+class HashAggregateOperator(HashShufflingOperatorBase):
+
+    _DEFAULT_MIN_NUM_SHARDS_COMPACTION_THRESHOLD = 100
+    _DEFAULT_MAX_NUM_SHARDS_COMPACTION_THRESHOLD = 2000
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        input_op: PhysicalOperator,
+        key_columns: Tuple[str],
+        aggregation_fns: Tuple[AggregateFn],
+        *,
+        num_partitions: Optional[int] = None,
+        aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None,
+    ):
+        # Use new stateless ReducingAggregation factory
+        def _create_reducing_aggregation() -> ReducingAggregation:
+            return ReducingAggregation(
+                key_columns=key_columns,
+                aggregation_fns=aggregation_fns,
+            )
+
+        super().__init__(
+            name_factory=(
+                lambda num_partitions: f"HashAggregate(key_columns={key_columns}, "
+                f"num_partitions={num_partitions})"
+            ),
+            input_ops=[input_op],
+            data_context=data_context,
+            key_columns=[key_columns],
+            num_input_seqs=1,
+            num_partitions=(
+                # NOTE: In case of global aggregations (ie with no key columns specified),
+                #       we override number of partitions to 1, since the whole dataset
+                #       will be reduced to just a single row
+                num_partitions
+                if len(key_columns) > 0
+                else 1
+            ),
+            partition_aggregation_factory=_create_reducing_aggregation,
+            input_block_transformer=_create_aggregating_transformer(
+                key_columns, aggregation_fns
+            ),
+            aggregator_ray_remote_args_override=aggregator_ray_remote_args_override,
+            shuffle_progress_bar_name="Shuffle",
+            finalize_progress_bar_name="Aggregation",
+        )
+
+    def _get_operator_num_cpus_override(self) -> float:
+        return self.data_context.hash_aggregate_operator_actor_num_cpus_override
+
+    @classmethod
+    def _estimate_aggregator_memory_allocation(
+        cls,
+        *,
+        num_aggregators: int,
+        num_partitions: int,
+        estimated_dataset_bytes: int,
+    ) -> int:
+        partition_byte_size_estimate = math.ceil(
+            estimated_dataset_bytes / num_partitions
+        )
+
+        # Estimate of object store memory required to accommodate all partitions
+        # handled by a single aggregator
+        aggregator_shuffle_object_store_memory_required: int = math.ceil(
+            estimated_dataset_bytes / num_aggregators
+        )
+        # Estimate of memory required to accommodate single partition as an output
+        # (inside Object Store)
+        output_object_store_memory_required: int = partition_byte_size_estimate
+
+        aggregator_total_memory_required: int = (
+            # Inputs (object store)
+            aggregator_shuffle_object_store_memory_required
+            +
+            # Output (object store)
+            output_object_store_memory_required
+        )
+
+        logger.info(
+            f"Estimated memory requirement for aggregating aggregator "
+            f"(partitions={num_partitions}, "
+            f"aggregators={num_aggregators}, "
+            f"dataset (estimate)={estimated_dataset_bytes / GiB:.1f}GiB): "
+            f"shuffle={aggregator_shuffle_object_store_memory_required / MiB:.1f}MiB, "
+            f"output={output_object_store_memory_required / MiB:.1f}MiB, "
+            f"total={aggregator_total_memory_required / MiB:.1f}MiB, "
+        )
+
+        return aggregator_total_memory_required
+
+    @classmethod
+    def _get_min_max_partition_shards_compaction_thresholds(
+        cls,
+    ) -> Optional[Tuple[int, int]]:
+        return (
+            cls._DEFAULT_MIN_NUM_SHARDS_COMPACTION_THRESHOLD,
+            cls._DEFAULT_MAX_NUM_SHARDS_COMPACTION_THRESHOLD,
+        )
+
+
+def _create_aggregating_transformer(
+    key_columns: Tuple[str], aggregation_fns: Tuple[AggregateFn]
+) -> BlockTransformer:
+    """Method creates input block transformer performing partial aggregation of
+    the block applied prior to block being shuffled (to reduce amount of bytes shuffled)"""
+
+    sort_key = ReducingAggregation._get_sort_key(key_columns)
+
+    def _aggregate(block: Block) -> Block:
+        from ray.data._internal.planner.exchange.aggregate_task_spec import (
+            SortAggregateTaskSpec,
+        )
+
+        # TODO unify blocks schemas, to avoid validating every block
+        # Validate block's schema compatible with aggregations
+        for agg_fn in aggregation_fns:
+            agg_fn._validate(BlockAccessor.for_block(block).schema())
+
+        # Project block to only carry columns used in aggregation
+        pruned_block = SortAggregateTaskSpec._prune_unused_columns(
+            block,
+            sort_key,
+            aggregation_fns,
+        )
+
+        # NOTE: If columns to aggregate on have been provided,
+        #       sort the block on these before aggregation
+        if sort_key.get_columns():
+            target_block = BlockAccessor.for_block(pruned_block).sort(sort_key)
+        else:
+            target_block = pruned_block
+
+        return BlockAccessor.for_block(target_block)._aggregate(
+            sort_key, aggregation_fns
+        )
+
+    return _aggregate

python/ray/data/_internal/execution/operators/hash_shuffle_v2.py

@@ -0,0 +1,54 @@
+from typing import Dict, Iterable, List
+
+import pyarrow as pa
+
+from ray.data._internal.arrow_ops.transform_pyarrow import hash_partition
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_tasks import (
+    PartitionFn,
+    ReduceFn,
+)
+
+# Isolate shuffle map workers into a dedicated worker pool so that
+# ReadParquet/Project tasks don't run on the same workers.  Without this,
+# shared memory pages from object store accesses (mmap'd during
+# combine_chunks) accumulate across task types and inflate worker RSS.
+_SHUFFLE_MAP_RUNTIME_ENV = {"env_vars": {"RAY_DATA_SHUFFLE_MAP_WORKER": "1"}}
+
+
+def _make_hash_partition_fn(key_columns: List[str], num_partitions: int) -> PartitionFn:
+    """Return a partition function that hash-partitions by key_columns."""
+
+    def _partition(block: pa.Table) -> Dict[int, pa.Table]:
+        return hash_partition(
+            block, hash_cols=key_columns, num_partitions=num_partitions
+        )
+
+    return _partition
+
+
+def _concat_reduce(
+    partition_id: int, tables_by_input: List[List[pa.Table]]
+) -> Iterable[pa.Table]:
+    """Concatenate all shards of a (single-input) partition into one block."""
+    tables = tables_by_input[0]
+    if not tables:
+        return
+    yield pa.concat_tables(tables) if len(tables) > 1 else tables[0]
+
+
+def _sort_reduce(key_columns: List[str]) -> ReduceFn:
+    """Return a reduce function that concatenates then sorts by key_columns.
+
+    Requires blocking mode because sorting needs all shards before emitting.
+    """
+
+    def _reduce(
+        partition_id: int, tables_by_input: List[List[pa.Table]]
+    ) -> Iterable[pa.Table]:
+        tables = tables_by_input[0]
+        if not tables:
+            return
+        combined = pa.concat_tables(tables) if len(tables) > 1 else tables[0]
+        yield combined.sort_by([(k, "ascending") for k in key_columns])
+
+    return _reduce

python/ray/data/_internal/execution/operators/input_data_buffer.py

@@ -0,0 +1,103 @@
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.block_ref_counter import BlockRefCounter
+
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+
+class InputDataBuffer(PhysicalOperator):
+    """Defines the input data for the operator DAG.
+
+    For example, this may hold cached blocks from a previous Dataset execution, or
+    the arguments for read tasks.
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        input_data: Optional[List[RefBundle]] = None,
+        input_data_factory: Optional[Callable[[int], List[RefBundle]]] = None,
+    ):
+        """Create an InputDataBuffer.
+
+        Args:
+            data_context: :class:`~ray.data.context.DataContext`
+                object to use injestion.
+            input_data: The list of bundles to output from this operator.
+            input_data_factory: The factory to get input data, if input_data is None.
+        """
+        super().__init__("Input", [], data_context)
+        if input_data is not None:
+            assert input_data_factory is None
+            # Copy the input data to avoid mutating the original list.
+            self._input_data = input_data[:]
+            self._is_input_initialized = True
+            self._initialize_metadata()
+        else:
+            # Initialize input lazily when execution is started.
+            assert input_data_factory is not None
+            self._input_data_factory = input_data_factory
+            self._is_input_initialized = False
+        self._input_data_index = 0
+        self.mark_execution_finished()
+
+    def start(
+        self,
+        options: ExecutionOptions,
+        block_ref_counter: "BlockRefCounter",
+    ) -> None:
+        if not self._is_input_initialized:
+            self._input_data = self._input_data_factory(
+                self.target_max_block_size_override
+                or self.data_context.target_max_block_size
+            )
+            self._is_input_initialized = True
+            self._initialize_metadata()
+        # InputDataBuffer does not take inputs from other operators,
+        # so we record input metrics here
+        for bundle in self._input_data:
+            self._metrics.on_input_received(bundle)
+        super().start(options, block_ref_counter)
+
+    def has_next(self) -> bool:
+        return self._input_data_index < len(self._input_data)
+
+    def _get_next_inner(self) -> RefBundle:
+        # We can't pop the input data. If we do, Ray might garbage collect the block
+        # references, and Ray won't be able to reconstruct downstream objects.
+        bundle = self._input_data[self._input_data_index]
+        self._input_data_index += 1
+        return bundle
+
+    def get_stats(self) -> StatsDict:
+        return {}
+
+    def _add_input_inner(self, refs, input_index) -> None:
+        raise ValueError("Inputs are not allowed for this operator.")
+
+    def _initialize_metadata(self):
+        assert self._input_data is not None and self._is_input_initialized
+        self._estimated_num_output_bundles = len(self._input_data)
+
+        block_metadata = []
+        total_rows = 0
+        for bundle in self._input_data:
+            block_metadata.extend(bundle.metadata)
+            bundle_num_rows = bundle.num_rows()
+            if total_rows is not None and bundle_num_rows is not None:
+                total_rows += bundle_num_rows
+            else:
+                # total row is unknown
+                total_rows = None
+        if total_rows:
+            self._estimated_num_output_rows = total_rows
+        self._stats = {
+            "input": block_metadata,
+        }

python/ray/data/_internal/execution/operators/join.py

@@ -0,0 +1,562 @@
+import logging
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Set, Tuple, Type
+
+from ray.data._internal.arrow_block import ArrowBlockAccessor
+from ray.data._internal.arrow_ops.transform_pyarrow import (
+    MIN_PYARROW_VERSION_RUN_END_ENCODED_TYPES,
+    MIN_PYARROW_VERSION_VIEW_TYPES,
+)
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.hash_shuffle import (
+    HashShufflingOperatorBase,
+    ShuffleAggregation,
+    _combine,
+)
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_tasks import (
+    ReduceFn,
+)
+from ray.data._internal.logical.operators import JoinType
+from ray.data._internal.util import GiB, MiB
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.data._internal.utils.transform_pyarrow import _is_pa_extension_type
+from ray.data.block import Block
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+
+@dataclass(frozen=True)
+class _DatasetPreprocessingResult:
+    """Result of join preprocessing containing split tables.
+
+    Separates tables into supported (directly joinable) and unsupported
+    (requires indexing) column projections.
+    """
+
+    supported_projection: "pa.Table"
+    unsupported_projection: "pa.Table"
+
+
+_JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP = {
+    JoinType.INNER: "inner",
+    JoinType.LEFT_OUTER: "left outer",
+    JoinType.RIGHT_OUTER: "right outer",
+    JoinType.FULL_OUTER: "full outer",
+    JoinType.LEFT_SEMI: "left semi",
+    JoinType.RIGHT_SEMI: "right semi",
+    JoinType.LEFT_ANTI: "left anti",
+    JoinType.RIGHT_ANTI: "right anti",
+}
+
+
+logger = logging.getLogger(__name__)
+
+
+class JoiningAggregation(ShuffleAggregation):
+    """Stateless aggregation for distributed joining of 2 sequences.
+
+    This implementation performs hash-based distributed joining by:
+        - Accumulating identical keys from both sequences into the same partition
+        - Performing join on individual partitions independently
+
+    For actual joining, Pyarrow native joining functionality is utilised.
+    """
+
+    def __init__(
+        self,
+        *,
+        join_type: JoinType,
+        left_key_col_names: Tuple[str, ...],
+        right_key_col_names: Tuple[str, ...],
+        left_columns_suffix: Optional[str] = None,
+        right_columns_suffix: Optional[str] = None,
+        data_context: DataContext,
+    ):
+        assert (
+            len(left_key_col_names) > 0
+        ), "At least 1 column to join on has to be provided"
+        assert len(right_key_col_names) == len(
+            left_key_col_names
+        ), "Number of columns for both left and right join operands has to match"
+
+        assert join_type in _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP, (
+            f"Join type is not currently supported (got: {join_type}; "  # noqa: C416
+            f"supported: {[jt for jt in JoinType]})"  # noqa: C416
+        )
+
+        self._left_key_col_names: Tuple[str, ...] = left_key_col_names
+        self._right_key_col_names: Tuple[str, ...] = right_key_col_names
+        self._join_type: JoinType = join_type
+
+        self._left_columns_suffix: Optional[str] = left_columns_suffix
+        self._right_columns_suffix: Optional[str] = right_columns_suffix
+
+    def finalize(self, partition_shards_map: Dict[int, List[Block]]) -> Iterator[Block]:
+        """Performs join on blocks from left (seq 0) and right (seq 1) sequences."""
+
+        assert (
+            len(partition_shards_map) == 2
+        ), f"Two input-sequences are expected (got {len(partition_shards_map)})"
+
+        left_partition_shards = partition_shards_map[0]
+        right_partition_shards = partition_shards_map[1]
+
+        left_table = _combine(left_partition_shards)
+        right_table = _combine(right_partition_shards)
+
+        yield join_tables(
+            left_table,
+            right_table,
+            join_type=self._join_type,
+            left_key_col_names=self._left_key_col_names,
+            right_key_col_names=self._right_key_col_names,
+            left_columns_suffix=self._left_columns_suffix,
+            right_columns_suffix=self._right_columns_suffix,
+        )
+
+
+def _make_join_reduce_fn(
+    *,
+    join_type: JoinType,
+    left_key_col_names: Tuple[str, ...],
+    right_key_col_names: Tuple[str, ...],
+    left_columns_suffix: Optional[str] = None,
+    right_columns_suffix: Optional[str] = None,
+    left_schema: Optional[Any] = None,
+    right_schema: Optional[Any] = None,
+) -> ReduceFn:
+    """Build a V2-shuffle reduce fn that joins two co-partitioned inputs."""
+    import pyarrow as pa
+
+    def _side_table(tables: List[Block], schema: Optional[Any]) -> Optional["pa.Table"]:
+        if tables:
+            return _combine(tables)
+        if isinstance(schema, pa.Schema):
+            return schema.empty_table()
+        return None
+
+    def _reduce(
+        partition_id: int, tables_by_input: List[List[Block]]
+    ) -> Iterator[Block]:
+        assert (
+            len(tables_by_input) == 2
+        ), f"Join reduce expects two inputs (got {len(tables_by_input)})"
+        left_table = _side_table(tables_by_input[0], left_schema)
+        right_table = _side_table(tables_by_input[1], right_schema)
+        if left_table is None or right_table is None:
+            # TODO(you-cheng): A whole input side is empty AND its schema can't be inferred
+            # (0 blocks + un-inferable schema, e.g. a map_batches side), so
+            # _side_table returns None and we skip the partition. This silently
+            # drops the preserved side's rows for preserving joins, left_outer/
+            # full_outer and left_anti/right_anti.
+            return
+        yield join_tables(
+            left_table,
+            right_table,
+            join_type=join_type,
+            left_key_col_names=left_key_col_names,
+            right_key_col_names=right_key_col_names,
+            left_columns_suffix=left_columns_suffix,
+            right_columns_suffix=right_columns_suffix,
+        )
+
+    return _reduce
+
+
+def join_tables(
+    left_table: "pa.Table",
+    right_table: "pa.Table",
+    *,
+    join_type: JoinType,
+    left_key_col_names: Tuple[str, ...],
+    right_key_col_names: Tuple[str, ...],
+    left_columns_suffix: Optional[str] = None,
+    right_columns_suffix: Optional[str] = None,
+) -> "pa.Table":
+    """Apply preprocess -> ``pa.Table.join`` -> postprocess to two input tables.
+
+    Shared between the physical executor (``JoiningAggregation.finalize``)
+    and plan-time schema inference (``Join.infer_schema``), which calls
+    this with empty tables built from the input schemas. Plan-time and
+    runtime schemas therefore agree by construction.
+    """
+    left_on = list(left_key_col_names)
+    right_on = list(right_key_col_names)
+
+    # Eagerly validate suffix conflicts so callers get a clear error instead
+    # of the opaque PyArrow schema-merge error ('Field X exists 2 times').
+    # Skip for semi/anti joins: only one side's columns appear in the result,
+    # so overlapping non-key names between left and right are harmless.
+    if join_type not in (
+        JoinType.LEFT_SEMI,
+        JoinType.LEFT_ANTI,
+        JoinType.RIGHT_SEMI,
+        JoinType.RIGHT_ANTI,
+    ):
+        left_cols = set(left_table.schema.names)
+        # PyArrow drops right key columns from output (coalescing them into
+        # the left keys), so only right non-key columns can collide with
+        # left columns. Subtracting only right_on (not left_on) correctly
+        # handles asymmetric key names (left_on != right_on).
+        right_output_cols = set(right_table.schema.names) - set(right_on)
+        collisions = left_cols & right_output_cols
+        if left_columns_suffix is None and right_columns_suffix is None and collisions:
+            raise ValueError(
+                "Left and right columns suffixes cannot be both None "
+                f"(overlapping columns: {sorted(collisions)})"
+            )
+
+    # Preprocess: split unsupported columns and add index columns if needed
+    preprocess_result_l, preprocess_result_r = _preprocess(
+        left_table, right_table, left_on, right_on, join_type
+    )
+
+    # Perform the join on supported columns
+    arrow_join_type = _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP[join_type]
+
+    supported = preprocess_result_l.supported_projection.join(
+        preprocess_result_r.supported_projection,
+        join_type=arrow_join_type,
+        keys=left_on,
+        right_keys=right_on,
+        left_suffix=left_columns_suffix,
+        right_suffix=right_columns_suffix,
+    )
+
+    # Add back unsupported columns
+    return _postprocess(
+        supported,
+        preprocess_result_l.unsupported_projection,
+        preprocess_result_r.unsupported_projection,
+    )
+
+
+def _preprocess(
+    left_table: "pa.Table",
+    right_table: "pa.Table",
+    left_on: List[str],
+    right_on: List[str],
+    join_type: JoinType,
+) -> Tuple[_DatasetPreprocessingResult, _DatasetPreprocessingResult]:
+    """Split inputs into supported/unsupported columns and add indices."""
+    supported_l, unsupported_l = _split_unsupported_columns(left_table)
+    supported_r, unsupported_r = _split_unsupported_columns(right_table)
+
+    # Handle joins on unsupported columns
+    conflicting_columns: Set[str] = set(unsupported_l.column_names) & set(left_on)
+    if conflicting_columns:
+        raise ValueError(
+            f"Cannot join on columns with unjoinable types. "
+            f"Left join key columns {conflicting_columns} have unjoinable types "
+            f"(map, union, list, struct, etc.) which cannot be used for join operations."
+        )
+
+    conflicting_columns: Set[str] = set(unsupported_r.column_names) & set(right_on)
+    if conflicting_columns:
+        raise ValueError(
+            f"Cannot join on columns with unjoinable types. "
+            f"Right join key columns {conflicting_columns} have unjoinable types "
+            f"(map, union, list, struct, etc.) which cannot be used for join operations."
+        )
+
+    # Index if we have unsupported columns
+    should_index_l = _should_index_side("left", supported_l, unsupported_l, join_type)
+    should_index_r = _should_index_side("right", supported_r, unsupported_r, join_type)
+
+    # Add index columns for back-referencing if we have unsupported columns
+    if should_index_l:
+        supported_l = _append_index_column(
+            table=supported_l, col_name=_index_name("left")
+        )
+    if should_index_r:
+        supported_r = _append_index_column(
+            table=supported_r, col_name=_index_name("right")
+        )
+
+    left = _DatasetPreprocessingResult(
+        supported_projection=supported_l,
+        unsupported_projection=unsupported_l,
+    )
+    right = _DatasetPreprocessingResult(
+        supported_projection=supported_r,
+        unsupported_projection=unsupported_r,
+    )
+    return left, right
+
+
+def _postprocess(
+    supported: "pa.Table",
+    unsupported_l: "pa.Table",
+    unsupported_r: "pa.Table",
+) -> "pa.Table":
+    """Re-attach unsupported columns to the joined table via the index column."""
+    should_index_l = _index_name("left") in supported.schema.names
+    should_index_r = _index_name("right") in supported.schema.names
+
+    # Add back unsupported columns (join type logic is in should_index_* variables)
+    if should_index_l:
+        supported = _add_back_unsupported_columns(
+            joined_table=supported,
+            unsupported_table=unsupported_l,
+            index_col_name=_index_name("left"),
+        )
+
+    if should_index_r:
+        supported = _add_back_unsupported_columns(
+            joined_table=supported,
+            unsupported_table=unsupported_r,
+            index_col_name=_index_name("right"),
+        )
+
+    return supported
+
+
+def _index_name(suffix: str) -> str:
+    return f"__rd_index_level_{suffix}__"
+
+
+def _should_index_side(
+    side: str,
+    supported_table: "pa.Table",
+    unsupported_table: "pa.Table",
+    join_type: JoinType,
+) -> bool:
+    """
+    Determine whether to create an index column for a given side of the join.
+
+    A column is "supported" if it is "joinable", and "unsupported" otherwise.
+    A supported_table is a table with only "supported" columns. Index columns are
+    needed when we have both supported and unsupported columns in a table, and
+    that table's columns will appear in the final result.
+
+    Args:
+        side: "left" or "right" to indicate which side of the join
+        supported_table: Table containing ONLY joinable columns
+        unsupported_table: Table containing ONLY unjoinable columns
+        join_type: The join type, used to decide whether this side appears in
+            the result (semi/anti joins drop one side).
+
+    Returns:
+        True if an index column should be created for this side
+    """
+    # Must have both supported and unsupported columns to need indexing.
+    # We cannot rely on row_count because it can return a non-zero row count
+    # for an empty-schema.
+    if len(supported_table.schema) == 0 or len(unsupported_table.schema) == 0:
+        return False
+
+    # For semi/anti joins, only index the side that appears in the result
+    if side == "left":
+        # Left side appears in result for all joins except right_semi/right_anti
+        return join_type not in [JoinType.RIGHT_SEMI, JoinType.RIGHT_ANTI]
+    else:  # side == "right"
+        # Right side appears in result for all joins except left_semi/left_anti
+        return join_type not in [JoinType.LEFT_SEMI, JoinType.LEFT_ANTI]
+
+
+def _split_unsupported_columns(
+    table: "pa.Table",
+) -> Tuple["pa.Table", "pa.Table"]:
+    """
+    Split a PyArrow table into two tables based on column joinability.
+
+    Separates columns into supported types and unsupported types that cannot be
+    directly joined on but should be preserved in results.
+
+    Args:
+        table: Input PyArrow table to split
+
+    Returns:
+        Tuple of (supported_table, unsupported_table) where:
+        - supported_table contains columns with primitive/joinable types
+        - unsupported_table contains columns with complex/unjoinable types
+    """
+    supported, unsupported = [], []
+    for idx in range(len(table.columns)):
+        col: "pa.ChunkedArray" = table.column(idx)
+        col_type: "pa.DataType" = col.type
+
+        if _is_pa_extension_type(col_type) or JoinOperator._is_pa_join_not_supported(
+            col_type
+        ):
+            unsupported.append(idx)
+        else:
+            supported.append(idx)
+
+    return table.select(supported), table.select(unsupported)
+
+
+def _add_back_unsupported_columns(
+    joined_table: "pa.Table",
+    unsupported_table: "pa.Table",
+    index_col_name: str,
+) -> "pa.Table":
+    # Extract the index column array and drop the column from the joined table
+    i = joined_table.schema.get_field_index(index_col_name)
+    indices = joined_table.column(i)
+    joined_table = joined_table.remove_column(i)
+
+    # Project the unsupported columns using the indices and combine with joined table
+    projected = ArrowBlockAccessor(unsupported_table).take(indices)
+    return ArrowBlockAccessor(joined_table).hstack(projected)
+
+
+def _append_index_column(table: "pa.Table", col_name: str) -> "pa.Table":
+    import numpy as np
+    import pyarrow as pa
+
+    index_col = pa.array(np.arange(table.num_rows))
+    return table.append_column(col_name, index_col)
+
+
+class JoinOperator(HashShufflingOperatorBase):
+    def __init__(
+        self,
+        data_context: DataContext,
+        left_input_op: PhysicalOperator,
+        right_input_op: PhysicalOperator,
+        left_key_columns: Tuple[str],
+        right_key_columns: Tuple[str],
+        join_type: JoinType,
+        *,
+        num_partitions: Optional[int] = None,
+        left_columns_suffix: Optional[str] = None,
+        right_columns_suffix: Optional[str] = None,
+        partition_size_hint: Optional[int] = None,
+        aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None,
+        shuffle_aggregation_type: Optional[Type[ShuffleAggregation]] = None,
+    ):
+        # Use new stateless JoiningAggregation factory
+        def _create_joining_aggregation() -> JoiningAggregation:
+            if shuffle_aggregation_type is not None:
+                if not issubclass(shuffle_aggregation_type, ShuffleAggregation):
+                    raise TypeError(
+                        f"shuffle_aggregation_type must be a subclass of {ShuffleAggregation}, "
+                        f"got {shuffle_aggregation_type}"
+                    )
+
+            aggregation_class = shuffle_aggregation_type or JoiningAggregation
+
+            return aggregation_class(
+                join_type=join_type,
+                left_key_col_names=left_key_columns,
+                right_key_col_names=right_key_columns,
+                left_columns_suffix=left_columns_suffix,
+                right_columns_suffix=right_columns_suffix,
+                data_context=data_context,
+            )
+
+        super().__init__(
+            name_factory=(
+                lambda num_partitions: f"Join(num_partitions={num_partitions})"
+            ),
+            input_ops=[left_input_op, right_input_op],
+            data_context=data_context,
+            key_columns=[left_key_columns, right_key_columns],
+            num_input_seqs=2,
+            num_partitions=num_partitions,
+            partition_size_hint=partition_size_hint,
+            partition_aggregation_factory=_create_joining_aggregation,
+            aggregator_ray_remote_args_override=aggregator_ray_remote_args_override,
+            shuffle_progress_bar_name="Shuffle",
+            finalize_progress_bar_name="Join",
+        )
+
+    @staticmethod
+    def _is_pa_join_not_supported(dtype: "pa.DataType") -> bool:
+        """
+        The latest pyarrow versions do not support joins where the
+        tables contain the following types below (lists,
+        structs, maps, unions, extension types, etc.)
+
+        Args:
+            dtype: The input type of column.
+
+        Returns:
+            True if the type cannot be present (non join-key) during joins.
+            False if the type can be present.
+        """
+        import pyarrow as pa
+
+        pyarrow_version = get_pyarrow_version()
+        is_v12 = pyarrow_version >= MIN_PYARROW_VERSION_RUN_END_ENCODED_TYPES
+        is_v16 = pyarrow_version >= MIN_PYARROW_VERSION_VIEW_TYPES
+
+        return (
+            pa.types.is_map(dtype)
+            or pa.types.is_union(dtype)
+            or pa.types.is_list(dtype)
+            or pa.types.is_struct(dtype)
+            or pa.types.is_null(dtype)
+            or pa.types.is_large_list(dtype)
+            or pa.types.is_fixed_size_list(dtype)
+            or (is_v12 and pa.types.is_run_end_encoded(dtype))
+            or (
+                is_v16
+                and (
+                    pa.types.is_binary_view(dtype)
+                    or pa.types.is_string_view(dtype)
+                    or pa.types.is_list_view(dtype)
+                )
+            )
+        )
+
+    def _get_operator_num_cpus_override(self) -> float:
+        return self.data_context.join_operator_actor_num_cpus_override
+
+    @classmethod
+    def _estimate_aggregator_memory_allocation(
+        cls,
+        *,
+        num_aggregators: int,
+        num_partitions: int,
+        estimated_dataset_bytes: int,
+    ) -> int:
+        partition_byte_size_estimate = math.ceil(
+            estimated_dataset_bytes / num_partitions
+        )
+
+        # Estimate of object store memory required to accommodate all partitions
+        # handled by a single aggregator
+        aggregator_shuffle_object_store_memory_required: int = math.ceil(
+            estimated_dataset_bytes / num_aggregators
+        )
+        # Estimate of memory required to perform actual (in-memory) join
+        # operation (inclusive of 50% overhead allocated for Pyarrow join
+        # implementation)
+        #
+        # NOTE:
+        #   - 2x due to budgeted 100% overhead of Arrow's in-memory join
+        join_memory_required: int = math.ceil(partition_byte_size_estimate * 2)
+        # Estimate of memory required to accommodate single partition as an output
+        # (inside Object Store)
+        #
+        # NOTE: x2 due to 2 sequences involved in joins
+        output_object_store_memory_required: int = partition_byte_size_estimate
+
+        aggregator_total_memory_required: int = (
+            # Inputs (object store)
+            aggregator_shuffle_object_store_memory_required
+            +
+            # Join (heap)
+            join_memory_required
+            +
+            # Output (object store)
+            output_object_store_memory_required
+        )
+
+        logger.info(
+            f"Estimated memory requirement for joining aggregator "
+            f"(partitions={num_partitions}, "
+            f"aggregators={num_aggregators}, "
+            f"dataset (estimate)={estimated_dataset_bytes / GiB:.1f}GiB): "
+            f"shuffle={aggregator_shuffle_object_store_memory_required / MiB:.1f}MiB, "
+            f"joining={join_memory_required / MiB:.1f}MiB, "
+            f"output={output_object_store_memory_required / MiB:.1f}MiB, "
+            f"total={aggregator_total_memory_required / MiB:.1f}MiB, "
+        )
+
+        return aggregator_total_memory_required

python/ray/data/_internal/execution/operators/limit_operator.py

@@ -0,0 +1,147 @@
+from collections import deque
+from dataclasses import replace
+from typing import Deque, List, Optional, Tuple
+
+import ray
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    OneToOneOperator,
+)
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data.block import Block, BlockAccessor, BlockMetadata, BlockStats
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+
+class LimitOperator(OneToOneOperator):
+    """Physical operator for limit."""
+
+    def __init__(
+        self,
+        limit: int,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+    ):
+        self._limit = limit
+        self._consumed_rows = 0
+        self._buffer: Deque[RefBundle] = deque()
+        self._name = f"limit={limit}"
+        self._output_blocks_stats: List[BlockStats] = []
+        self._cur_output_bundles = 0
+        super().__init__(self._name, input_op, data_context)
+        if self._limit <= 0:
+            self.mark_execution_finished()
+
+    def _limit_reached(self) -> bool:
+        return self._consumed_rows >= self._limit
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.has_completed()
+        assert input_index == 0, input_index
+        if self._limit_reached():
+            return
+        out_blocks: List[ObjectRef[Block]] = []
+        out_metadata: List[BlockMetadata] = []
+        for entry in refs.blocks:
+            block = entry.ref
+            metadata = entry.metadata
+            num_rows = metadata.num_rows
+            assert num_rows is not None
+            if self._consumed_rows + num_rows <= self._limit:
+                out_blocks.append(block)
+                out_metadata.append(metadata)
+                self._output_blocks_stats.append(metadata.to_stats())
+                self._consumed_rows += num_rows
+            else:
+                # Slice the last block.
+                def slice_fn(block, metadata, num_rows) -> Tuple[Block, BlockMetadata]:
+                    block = BlockAccessor.for_block(block).slice(
+                        0, num_rows, copy=False
+                    )
+                    metadata = replace(
+                        metadata,
+                        num_rows=num_rows,
+                        size_bytes=BlockAccessor.for_block(block).size_bytes(),
+                    )
+                    return block, metadata
+
+                slice_task = cached_remote_fn(slice_fn, num_cpus=0, num_returns=2)
+                label_selector = self.data_context.execution_options.label_selector
+                if label_selector:
+                    slice_task = slice_task.options(label_selector=label_selector)
+                block, metadata_ref = slice_task.remote(
+                    block,
+                    metadata,
+                    self._limit - self._consumed_rows,
+                )
+                out_blocks.append(block)
+                metadata = ray.get(metadata_ref)
+                # Slicing creates a new block; register it for memory tracking.
+                self._block_ref_counter.on_block_produced(
+                    block, metadata.size_bytes or 0, self.id
+                )
+                out_metadata.append(metadata)
+                self._output_blocks_stats.append(metadata.to_stats())
+                self._consumed_rows = self._limit
+                break
+        self._cur_output_bundles += 1
+        out_refs = RefBundle(
+            [BlockEntry(b, m) for b, m in zip(out_blocks, out_metadata)],
+            owns_blocks=refs.owns_blocks,
+            schema=refs.schema,
+        )
+        self._buffer.append(out_refs)
+        self._metrics.on_output_queued(out_refs)
+        if self._limit_reached():
+            self.mark_execution_finished()
+
+        # We cannot estimate if we have only consumed empty blocks,
+        # or if the input dependency's total number of output bundles is unknown.
+        num_inputs = self.input_dependencies[0].num_outputs_total()
+        if self._consumed_rows > 0 and num_inputs is not None:
+            # Estimate number of output bundles
+            # Check the case where _limit > # of input rows
+            estimated_total_output_rows = min(
+                self._limit, self._consumed_rows / self._cur_output_bundles * num_inputs
+            )
+            # _consumed_rows / _limit is roughly equal to
+            # _cur_output_bundles / total output blocks
+            self._estimated_num_output_bundles = round(
+                estimated_total_output_rows
+                / self._consumed_rows
+                * self._cur_output_bundles
+            )
+
+    def has_next(self) -> bool:
+        return len(self._buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        output = self._buffer.popleft()
+        self._metrics.on_output_dequeued(output)
+        return output
+
+    def get_stats(self) -> StatsDict:
+        return {self._name: self._output_blocks_stats}
+
+    def num_outputs_total(self) -> Optional[int]:
+        # Before execution is completed, we don't know how many output
+        # bundles we will have. We estimate based off the consumption so far.
+        if self.has_execution_finished():
+            return self._cur_output_bundles
+        return self._estimated_num_output_bundles
+
+    def num_output_rows_total(self) -> Optional[int]:
+        # The total number of rows is simply the limit or the number
+        # of input rows, whichever is smaller
+        input_num_rows = self.input_dependencies[0].num_output_rows_total()
+        if input_num_rows is None:
+            return None
+        return min(self._limit, input_num_rows)
+
+    def throttling_disabled(self) -> bool:
+        return True

python/ray/data/_internal/execution/operators/map_transformer.py

@@ -0,0 +1,605 @@
+import itertools
+import time
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+from ray._common.utils import env_integer
+from ray.data._internal.block_batching.block_batching import batch_blocks
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption
+from ray.data.block import (
+    BatchFormat,
+    Block,
+    BlockAccessor,
+    CustomOpStats,
+    DataBatch,
+)
+
+_DEFAULT_BATCH_SIZE_BYTES: int = env_integer(
+    "RAY_DATA_DEFAULT_BATCH_SIZE_BYTES", 16 * 1024 * 1024  # 16 MiB
+)
+
+# Allowed input/output data types for a MapTransformFn.
+Row = Dict[str, Any]
+MapTransformFnData = Union[Block, Row, DataBatch]
+
+
+class CustomOpStatsReporter:
+    """Per-task reporter that carries transforms' :class:`CustomOpStats`.
+
+    ``_map_task`` creates one per task and threads it into the transform chain.
+    Each producing transform calls ``op_stats_reporter.report(stats)`` once,
+    before yielding output blocks, to append its :class:`CustomOpStats` to the
+    reporter. Fused transforms each contribute one entry, so the reporter holds a
+    list. ``_map_task`` reads :meth:`get_stats` after each output block and stamps
+    the list onto the block metadata as part of ``TaskExecWorkerStats``
+    """
+
+    def __init__(self) -> None:
+        self._stats: List[CustomOpStats] = []
+
+    def report(self, stats: CustomOpStats) -> None:
+        """Append a producing transform's per-task CustomOpStats."""
+        self._stats.append(stats)
+
+    def get_stats(self) -> List[CustomOpStats]:
+        """Return all reported CustomOpStats (empty if none were reported)."""
+        return self._stats
+
+    def clear(self) -> None:
+        """Drop any reported stats (called before each task attempt)."""
+        self._stats = []
+
+
+# Narrow callback handed to producing transforms to report per-task
+# :class:`CustomOpStats`.
+CustomOpStatsReportFn = Callable[[CustomOpStats], None]
+
+
+def _noop_report_custom_op_stats(stats: CustomOpStats) -> None:
+    """Stateless default report callback for callers that don't collect stats."""
+
+
+IN = TypeVar("IN")
+OUT = TypeVar("OUT")
+# A transform callable accepts either ``(data, ctx)`` or, when it reports
+# per-task CustomOpStats, ``(data, ctx, report_custom_op_stats)``.
+MapTransformCallable = Union[
+    Callable[[Iterable[IN], TaskContext], Iterable[OUT]],
+    Callable[[Iterable[IN], TaskContext, CustomOpStatsReportFn], Iterable[OUT]],
+]
+
+
+class MapTransformFnDataType(Enum):
+    """An enum that represents the input/output data type of a MapTransformFn."""
+
+    Block = 0
+    Row = 1
+    Batch = 2
+
+
+class MapTransformFn(ABC):
+    """Represents a single transform function in a MapTransformer."""
+
+    def __init__(
+        self,
+        fn: Callable,
+        input_type: MapTransformFnDataType,
+        *,
+        is_udf: bool = False,
+        output_block_size_option: Optional[OutputBlockSizeOption] = None,
+        reports_custom_op_stats: bool = False,
+    ):
+        """Initialize a :class:`MapTransformFn`.
+
+        Args:
+            fn: The wrapped transform callable. Invoked with ``(data, ctx)``, or
+                ``(data, ctx, report_custom_op_stats)`` when
+                ``reports_custom_op_stats=True``.
+            input_type: Expected type of the input data.
+            is_udf: Whether this transformation is UDF or not.
+            output_block_size_option: (Optional) Output block size configuration.
+            reports_custom_op_stats: If ``True``, the wrapped callable accepts a
+                third ``report_custom_op_stats`` callback argument and may report
+                per-task :class:`CustomOpStats` to the driver. Defaults to
+                ``False``, in which case the callable is invoked with
+                ``(data, ctx)`` only.
+        """
+        self._fn = fn
+        self._input_type = input_type
+        self._output_block_size_option = output_block_size_option
+        self._is_udf = is_udf
+        self._reports_custom_op_stats = reports_custom_op_stats
+
+    @abstractmethod
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        pass
+
+    def _apply_transform(
+        self,
+        ctx: TaskContext,
+        inputs: Iterable[MapTransformFnData],
+        report_custom_op_stats: CustomOpStatsReportFn = _noop_report_custom_op_stats,
+    ) -> Iterable[MapTransformFnData]:
+        """Call the wrapped fn, passing ``report_custom_op_stats`` only if it opted in.
+
+        Keeps the common ``(data, ctx)`` signature for the vast majority of
+        transforms; only those constructed with ``reports_custom_op_stats=True``
+        receive the report callback.
+        """
+        if self._reports_custom_op_stats:
+            return self._fn(inputs, ctx, report_custom_op_stats)
+        return self._fn(inputs, ctx)
+
+    def _pre_process(self, blocks: Iterable[Block]) -> Iterable[MapTransformFnData]:
+        return blocks
+
+    def _shape_blocks(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        """Shape results into blocks using a buffer."""
+        return _BlockShapingIterator(
+            results, self._input_type, self._output_block_size_option
+        )
+
+    def __call__(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+        report_custom_op_stats: CustomOpStatsReportFn = _noop_report_custom_op_stats,
+    ) -> Iterable[Block]:
+        batches = self._pre_process(blocks)
+        results = self._apply_transform(ctx, batches, report_custom_op_stats)
+        return self._post_process(results)
+
+    @property
+    def output_block_size_option(self):
+        return self._output_block_size_option
+
+    def override_target_max_block_size(self, target_max_block_size: Optional[int]):
+        if self._output_block_size_option is not None and (
+            self._output_block_size_option.disable_block_shaping
+            or self._output_block_size_option.target_num_rows_per_block is not None
+        ):
+            raise ValueError(
+                "Cannot override target_max_block_size if block shaping is disabled or target_num_rows_per_block is set"
+            )
+        self._output_block_size_option = OutputBlockSizeOption.of(
+            target_max_block_size=target_max_block_size
+        )
+
+    @property
+    def target_max_block_size(self):
+        if self._output_block_size_option is None:
+            return None
+        else:
+            return self._output_block_size_option.target_max_block_size
+
+    @property
+    def target_num_rows_per_block(self):
+        if self._output_block_size_option is None:
+            return None
+        else:
+            return self._output_block_size_option.target_num_rows_per_block
+
+
+class MapTransformer:
+    """Encapsulates the data transformation logic of a physical MapOperator.
+
+    A MapTransformer may consist of one or more steps, each of which is represented
+    as a MapTransformFn. The first MapTransformFn must take blocks as input, and
+    the last MapTransformFn must output blocks. The intermediate data types can
+    be blocks, rows, or batches.
+    """
+
+    class _UDFTimingIterator(Iterator[MapTransformFnData]):
+        """Iterator that times UDF execution"""
+
+        def __init__(
+            self, input: Iterable[MapTransformFnData], transformer: "MapTransformer"
+        ):
+            self._input = input
+            self._transformer = transformer
+
+        def __iter__(self) -> "MapTransformer._UDFTimingIterator":
+            return self
+
+        def __next__(self) -> MapTransformFnData:
+            start = time.perf_counter()
+            try:
+                return next(self._input)
+            finally:
+                self._transformer._report_udf_time(time.perf_counter() - start)
+
+    def __init__(
+        self,
+        transform_fns: List[MapTransformFn],
+        *,
+        init_fn: Optional[Callable[[], None]] = None,
+        output_block_size_option_override: Optional[OutputBlockSizeOption] = None,
+    ):
+        """Initialize a :class:`MapTransformer`.
+
+        Args:
+            transform_fns: A list of `MapTransformFn`s that will be executed sequentially
+                to transform data.
+            init_fn: A function that will be called before transforming data.
+                Used for the actor-based map operator.
+            output_block_size_option_override: (Optional) Output block size configuration.
+        """
+
+        self._transform_fns: List[MapTransformFn] = []
+        self._init_fn = init_fn if init_fn is not None else lambda: None
+        self._output_block_size_option_override = output_block_size_option_override
+        self._udf_time_s = 0
+
+        # Add transformations
+        self.add_transform_fns(transform_fns)
+
+    def add_transform_fns(self, transform_fns: List[MapTransformFn]) -> None:
+        """Set the transform functions."""
+        assert len(transform_fns) > 0
+        self._transform_fns = self._combine_transformations(
+            self._transform_fns, transform_fns
+        )
+
+    def get_transform_fns(self) -> List[MapTransformFn]:
+        """Get the transform functions."""
+        return self._transform_fns
+
+    def override_target_max_block_size(self, target_max_block_size: Optional[int]):
+        self._output_block_size_option_override = OutputBlockSizeOption.of(
+            target_max_block_size=target_max_block_size
+        )
+
+    @property
+    def target_max_block_size_override(self) -> Optional[int]:
+        if self._output_block_size_option_override is None:
+            return None
+        else:
+            return self._output_block_size_option_override.target_max_block_size
+
+    def init(self) -> None:
+        """Initialize the transformer.
+
+        Should be called before applying the transform.
+        """
+        self._init_fn()
+
+    def apply_transform(
+        self,
+        input_blocks: Iterable[Block],
+        ctx: TaskContext,
+        report_custom_op_stats: CustomOpStatsReportFn = _noop_report_custom_op_stats,
+    ) -> Iterable[Block]:
+        """Apply the transform functions to the input blocks.
+
+        Args:
+            input_blocks: The blocks to transform.
+            ctx: The task context for this transform.
+            report_custom_op_stats: Callback a producing transform calls to report
+                its :class:`CustomOpStats`. ``_map_task`` passes its reporter's
+                ``report``; defaults to a stateless no-op for callers (e.g. tests)
+                that don't collect custom stats.
+
+        Returns:
+            An iterable of the transformed output blocks.
+        """
+        # NOTE: We only need to configure last transforming function to do
+        #       appropriate block sizing
+        last_transform = self._transform_fns[-1]
+
+        if self.target_max_block_size_override is not None:
+            last_transform.override_target_max_block_size(
+                self.target_max_block_size_override
+            )
+
+        iter = input_blocks
+        # Apply the transform functions sequentially to the input iterable.
+        for transform_fn in self._transform_fns:
+            iter = transform_fn(iter, ctx, report_custom_op_stats)
+            if transform_fn._is_udf:
+                iter = self._UDFTimingIterator(iter, self)
+
+        return iter
+
+    def fuse(self, other: "MapTransformer") -> "MapTransformer":
+        """Fuse two `MapTransformer`s together."""
+        assert (
+            self.target_max_block_size_override == other.target_max_block_size_override
+            or (
+                self.target_max_block_size_override is None
+                or other.target_max_block_size_override is None
+            )
+        )
+        # Define them as standalone variables to avoid fused_init_fn capturing the
+        # entire `MapTransformer` object.
+        self_init_fn = self._init_fn
+        other_init_fn = other._init_fn
+
+        def fused_init_fn():
+            self_init_fn()
+            other_init_fn()
+
+        combined_transform_fns = self._combine_transformations(
+            self._transform_fns,
+            other._transform_fns,
+        )
+
+        transformer = MapTransformer(
+            combined_transform_fns,
+            init_fn=fused_init_fn,
+            output_block_size_option_override=OutputBlockSizeOption.of(
+                target_max_block_size=(
+                    self.target_max_block_size_override
+                    or other.target_max_block_size_override
+                ),
+            ),
+        )
+
+        return transformer
+
+    @classmethod
+    def _combine_transformations(
+        cls, ones: List[MapTransformFn], others: List[MapTransformFn]
+    ) -> list[Any]:
+        return ones + others
+
+    def udf_time_s(self, reset: bool) -> float:
+        cur_time_s = self._udf_time_s
+        if reset:
+            self._udf_time_s = 0
+        return cur_time_s
+
+    def _report_udf_time(self, udf_time: float) -> None:
+        self._udf_time_s += udf_time
+
+
+class RowMapTransformFn(MapTransformFn):
+    """A rows-to-rows MapTransformFn."""
+
+    def __init__(
+        self,
+        row_fn: MapTransformCallable[Row, Row],
+        *,
+        is_udf: bool = False,
+        output_block_size_option: OutputBlockSizeOption,
+        reports_custom_op_stats: bool = False,
+    ):
+        super().__init__(
+            row_fn,
+            input_type=MapTransformFnDataType.Row,
+            is_udf=is_udf,
+            output_block_size_option=output_block_size_option,
+            reports_custom_op_stats=reports_custom_op_stats,
+        )
+
+    def _pre_process(self, blocks: Iterable[Block]) -> Iterable[MapTransformFnData]:
+        return _RowBasedIterator(blocks)
+
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        return self._shape_blocks(results)
+
+    def __repr__(self) -> str:
+        return f"RowMapTransformFn({self._fn})"
+
+
+def _peek_first_nonempty_block(
+    blocks: Iterable[Block],
+) -> Tuple[Optional[BlockAccessor], Iterable[Block]]:
+    """Advance the iterator past leading empty blocks to find the first non-empty block,
+    returning the corresponding accessor and a reconstructed iterator of all blocks.
+    We must reconstruct the iterator because we consume blocks as we advance through the iterator."""
+    blocks_iter = iter(blocks)
+    consumed = []
+    for block in blocks_iter:
+        consumed.append(block)
+        accessor = BlockAccessor.for_block(block)
+        if accessor.num_rows() > 0 and accessor.size_bytes() > 0:
+            return accessor, itertools.chain(consumed, blocks_iter)
+    return None, iter(consumed)
+
+
+def _compute_auto_batch_size(
+    blocks: Iterable[Block],
+    target_batch_size_bytes: int = _DEFAULT_BATCH_SIZE_BYTES,
+) -> Tuple[Optional[int], Iterable[Block]]:
+    """Peek at the first non-empty block to estimate the batch size to use for the
+    'auto' batch_size option."""
+    sample, blocks = _peek_first_nonempty_block(blocks)
+    if sample is None:
+        return None, blocks
+    bytes_per_row = sample.size_bytes() / sample.num_rows()
+    computed_batch_size = max(1, int(target_batch_size_bytes / bytes_per_row))
+    return computed_batch_size, blocks
+
+
+class BatchMapTransformFn(MapTransformFn):
+    """A batch-to-batch MapTransformFn."""
+
+    def __init__(
+        self,
+        batch_fn: MapTransformCallable[DataBatch, DataBatch],
+        *,
+        is_udf: bool = False,
+        batch_size: Union[Optional[int], Literal["auto"]] = None,
+        batch_format: Optional[BatchFormat] = None,
+        zero_copy_batch: bool = True,
+        output_block_size_option: Optional[OutputBlockSizeOption] = None,
+        target_batch_size_bytes: int = _DEFAULT_BATCH_SIZE_BYTES,
+        reports_custom_op_stats: bool = False,
+    ):
+        super().__init__(
+            batch_fn,
+            input_type=MapTransformFnDataType.Batch,
+            is_udf=is_udf,
+            output_block_size_option=output_block_size_option,
+            reports_custom_op_stats=reports_custom_op_stats,
+        )
+
+        self._batch_size = batch_size
+        self._batch_format = batch_format
+        self._zero_copy_batch = zero_copy_batch
+        self._target_batch_size_bytes = target_batch_size_bytes
+
+    def _pre_process(self, blocks: Iterable[Block]) -> Iterable[MapTransformFnData]:
+        # TODO make batch-udf zero-copy by default
+        if self._batch_size == "auto":
+            batch_size, blocks = _compute_auto_batch_size(
+                blocks, target_batch_size_bytes=self._target_batch_size_bytes
+            )
+        else:
+            batch_size = self._batch_size
+        ensure_copy = not self._zero_copy_batch and batch_size is not None
+        return batch_blocks(
+            blocks=iter(blocks),
+            stats=None,
+            batch_size=batch_size,
+            batch_format=self._batch_format,
+            ensure_copy=ensure_copy,
+        )
+
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        return self._shape_blocks(results)
+
+    def __repr__(self) -> str:
+        return f"BatchMapTransformFn({self._fn=}, {self._batch_format=}, {self._batch_size=}, {self._zero_copy_batch=})"
+
+
+class BlockMapTransformFn(MapTransformFn):
+    """A block-to-block MapTransformFn."""
+
+    def __init__(
+        self,
+        block_fn: MapTransformCallable[Block, Block],
+        *,
+        is_udf: bool = False,
+        disable_block_shaping: bool = False,
+        output_block_size_option: Optional[OutputBlockSizeOption] = None,
+        reports_custom_op_stats: bool = False,
+    ):
+        """
+        Initializes the object with a transformation function, accompanying options, and
+        configuration for handling blocks during processing.
+
+        Args:
+            block_fn: Callable function to apply a transformation to a block.
+            is_udf: Specifies if the transformation function is a user-defined
+                function (defaults to ``False``).
+            disable_block_shaping: Disables block-shaping, making transformer to
+                produce blocks as is.
+            output_block_size_option: (Optional) Configure output block sizing.
+            reports_custom_op_stats: If ``True``, ``block_fn`` accepts a third
+                ``report_custom_op_stats`` callback argument and may report
+                per-task :class:`CustomOpStats` to the driver.
+        """
+
+        super().__init__(
+            block_fn,
+            input_type=MapTransformFnDataType.Block,
+            is_udf=is_udf,
+            output_block_size_option=output_block_size_option,
+            reports_custom_op_stats=reports_custom_op_stats,
+        )
+
+        self._disable_block_shaping = disable_block_shaping
+
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        # Short-circuit for block transformations for which no
+        # block-shaping is required
+        if self._disable_block_shaping:
+            return results
+
+        return self._shape_blocks(results)
+
+    def __repr__(self) -> str:
+        return f"BlockMapTransformFn({self._fn=}, {self._output_block_size_option=})"
+
+
+class _BlockShapingIterator(Iterator[Block]):
+    """Iterator that shapes results into blocks using a buffer.
+
+    Unlike a generator, local variables in __next__ go out of scope when the method
+    returns, avoiding holding references to yielded values.
+    """
+
+    def __init__(
+        self,
+        results: Iterable[MapTransformFnData],
+        input_type: MapTransformFnDataType,
+        output_block_size_option: Optional[OutputBlockSizeOption],
+    ):
+        self._results_iter = iter(results)
+        self._buffer = BlockOutputBuffer(output_block_size_option)
+        self._finalized = False
+
+        if input_type == MapTransformFnDataType.Block:
+            self._append_buffer = self._buffer.add_block
+        elif input_type == MapTransformFnDataType.Batch:
+            self._append_buffer = self._buffer.add_batch
+        else:
+            assert input_type == MapTransformFnDataType.Row
+            self._append_buffer = self._buffer.add
+
+    def __iter__(self) -> "_BlockShapingIterator":
+        return self
+
+    def __next__(self) -> Block:
+        while True:
+            # First, yield any ready blocks from buffer
+            if self._buffer.has_next():
+                return self._buffer.next()
+
+            # If finalized, no more data
+            elif self._finalized:
+                raise StopIteration
+
+            try:
+                # Fetch more results
+                result = next(self._results_iter)
+                self._append_buffer(result)
+            except StopIteration:
+                self._buffer.finalize()
+                self._finalized = True
+
+
+class _RowBasedIterator(Iterator[Row]):
+    """Iterator that extracts rows from blocks.
+
+    Unlike a generator, local variables in __next__ go out of scope when the method
+    returns, avoiding holding references to yielded values.
+    """
+
+    def __init__(self, blocks: Iterable[Block]):
+        self._blocks_iter = iter(blocks)
+        self._cur_row_iter: Optional[Iterator[Row]] = None
+
+    def __iter__(self) -> "_RowBasedIterator":
+        return self
+
+    def __next__(self) -> Row:
+        while True:
+            # Try to get next row from current block
+            if self._cur_row_iter is not None:
+                try:
+                    return next(self._cur_row_iter)
+                except StopIteration:
+                    pass
+
+            # Get iterator from the next block
+            block = next(self._blocks_iter)
+
+            self._cur_row_iter = BlockAccessor.for_block(block).iter_rows(
+                public_row_format=True
+            )

python/ray/data/_internal/execution/operators/mix_operator.py

@@ -0,0 +1,232 @@
+from typing import List, Optional
+
+from typing_extensions import override
+
+from ray.data._internal.execution.bundle_queue import BaseBundleQueue, FIFOBundleQueue
+from ray.data._internal.execution.interfaces import (
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+    NAryOperator,
+)
+from ray.data._internal.logical.operators.n_ary_operator import (
+    MixStoppingCondition,
+    estimate_num_mix_outputs,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+
+class MixOperator(InternalQueueOperatorMixin, NAryOperator):
+    """An operator that interleaves blocks from multiple input operators
+    into a single output stream, respecting target row ratios specified
+    by weights.
+
+    Tracks cumulative row counts per input and always pulls from whichever
+    input has fallen furthest behind its target ratio. This ensures the
+    output row ratio converges to the target weights regardless of input
+    block size variance.
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        *input_ops: PhysicalOperator,
+        weights: List[float],
+        stopping_condition: MixStoppingCondition = MixStoppingCondition.STOP_ON_SHORTEST,
+    ):
+        assert len(input_ops) >= 1
+        assert len(weights) == len(input_ops)
+        if any(w <= 0 for w in weights):
+            raise ValueError("Weights must be positive.")
+
+        total_weight = sum(weights)
+        self._weights = [w / total_weight for w in weights]
+
+        self._stopping_condition = stopping_condition
+
+        self._input_buffers: List[BaseBundleQueue] = [
+            FIFOBundleQueue() for _ in range(len(input_ops))
+        ]
+        self._output_buffer: BaseBundleQueue = FIFOBundleQueue()
+
+        # Cumulative rows output per input, used for deficit calculation.
+        self._rows_seen: List[int] = [0] * len(input_ops)
+        self._input_done_flags: List[bool] = [False] * len(input_ops)
+        self._stopped: bool = False
+
+        self._stats: StatsDict = {"Mix": []}
+
+        input_names = ", ".join([op._name for op in input_ops])
+        weights_str = [round(w, 2) for w in self._weights]
+        name = f"Mix({input_names}, weights={weights_str})"
+        super().__init__(data_context, *input_ops, name=name)
+
+    # ------------------------------------------------------------------
+    # InternalQueueOperatorMixin interface
+    # ------------------------------------------------------------------
+
+    @property
+    @override
+    def _input_queues(self) -> List[BaseBundleQueue]:
+        return self._input_buffers
+
+    @property
+    @override
+    def _output_queues(self) -> List[BaseBundleQueue]:
+        return [self._output_buffer]
+
+    # ------------------------------------------------------------------
+    # PhysicalOperator interface
+    # ------------------------------------------------------------------
+
+    @override
+    def mark_execution_finished(self) -> None:
+        # Override InternalQueueOperatorMixin's version to preserve the
+        # output buffer for draining. Only clear input queues.
+        PhysicalOperator.mark_execution_finished(self)
+        self.clear_internal_input_queue()
+
+    @override
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.has_completed()
+        assert 0 <= input_index < len(self._input_dependencies), input_index
+        if self._stopped:
+            return
+        self._input_buffers[input_index].add(refs)
+        self._metrics.on_input_queued(refs, input_index=input_index)
+        self._try_output()
+
+    @override
+    def input_done(self, input_index: int) -> None:
+        self._input_done_flags[input_index] = True
+        self._try_output()
+
+    @override
+    def all_inputs_done(self) -> None:
+        super().all_inputs_done()
+        self._try_output()
+
+    @override
+    def has_next(self) -> bool:
+        return len(self._output_buffer) > 0
+
+    @override
+    def _get_next_inner(self) -> RefBundle:
+        refs = self._output_buffer.get_next()
+        self._metrics.on_output_dequeued(refs)
+        return refs
+
+    @override
+    def num_outputs_total(self) -> Optional[int]:
+        if self._stopping_condition == MixStoppingCondition.STOP_ON_SHORTEST:
+            # Can't accurately estimate output block count because weights
+            # control row ratios, not block ratios. With non-uniform block
+            # sizes, the block count doesn't follow the weight distribution.
+            return None
+
+        return estimate_num_mix_outputs(
+            [op.num_outputs_total() for op in self.input_dependencies],
+            self._weights,
+            self._stopping_condition,
+        )
+
+    @override
+    def num_output_rows_total(self) -> Optional[int]:
+        return estimate_num_mix_outputs(
+            [op.num_output_rows_total() for op in self.input_dependencies],
+            self._weights,
+            self._stopping_condition,
+        )
+
+    @override
+    def get_stats(self) -> StatsDict:
+        return self._stats
+
+    @override
+    def throttling_disabled(self) -> bool:
+        # TODO: Disable throttling along with Union once NAry operator resource accounting is fixed.
+        return False
+
+    # ------------------------------------------------------------------
+    # Output selection
+    # ------------------------------------------------------------------
+
+    def _is_input_exhausted(self, index: int) -> bool:
+        """An input is exhausted when it's done and its buffer is empty."""
+        return (
+            self._input_done_flags[index] and not self._input_buffers[index].has_next()
+        )
+
+    def _select_most_behind_input(self) -> int:
+        """Select which input to pull from next.
+
+        Returns the index of the non-exhausted input that has fallen furthest
+        behind its target row ratio. Ties are broken by weight (prefer higher),
+        then by index. Returns -1 if all inputs are exhausted.
+        """
+        total = sum(self._rows_seen)
+        best_index = -1
+        most_behind = float("-inf")
+        best_weight = -1.0
+
+        for i in range(len(self._input_buffers)):
+            if self._is_input_exhausted(i):
+                continue
+            # How far behind this input is: positive means underrepresented.
+            gap = self._weights[i] * total - self._rows_seen[i]
+            # Tie-break by weight.
+            if gap > most_behind or (
+                gap == most_behind and self._weights[i] > best_weight
+            ):
+                most_behind = gap
+                best_weight = self._weights[i]
+                best_index = i
+
+        return best_index
+
+    def _try_output(self) -> None:
+        """Move blocks from input buffers to the output buffer.
+
+        On each iteration, selects the input furthest behind its target ratio.
+        If that input has blocks, one is moved to the output. If not, we wait
+        rather than pulling from a different input — this keeps the output
+        deterministic regardless of block arrival timing.
+        """
+        if self._stopped:
+            return
+
+        while True:
+            if self._stopping_condition == MixStoppingCondition.STOP_ON_SHORTEST:
+                if any(
+                    self._is_input_exhausted(i) for i in range(len(self._input_buffers))
+                ):
+                    self._stopped = True
+                    self.mark_execution_finished()
+                    return
+            elif self._stopping_condition != MixStoppingCondition.STOP_ON_LONGEST_DROP:
+                raise ValueError(
+                    f"Unknown stopping condition: {self._stopping_condition}"
+                )
+
+            best_index = self._select_most_behind_input()
+            if best_index == -1:
+                return
+
+            if not self._input_buffers[best_index].has_next():
+                # Selected input has no blocks yet — wait rather than
+                # pulling from a lower-deficit input.
+                return
+
+            # Move one block from the selected input to the output buffer.
+            bundle = self._input_buffers[best_index].get_next()
+            self._metrics.on_input_dequeued(bundle, input_index=best_index)
+
+            num_rows = bundle.num_rows()
+            assert num_rows is not None
+            self._rows_seen[best_index] += num_rows
+
+            self._output_buffer.add(bundle)
+            self._metrics.on_output_queued(bundle)

python/ray/data/_internal/execution/operators/output_splitter.py

@@ -0,0 +1,451 @@
+import logging
+import math
+import time
+from dataclasses import replace
+from typing import TYPE_CHECKING, Any, Collection, Dict, List, Optional, Tuple
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.block_ref_counter import BlockRefCounter
+
+from typing_extensions import override
+
+from ray._common.utils import env_float
+from ray.data._internal.execution.bundle_queue import (
+    BaseBundleQueue,
+    FIFOBundleQueue,
+    HashLinkedQueue,
+)
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    ExecutionOptions,
+    NodeIdStr,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+)
+from ray.data._internal.execution.util import locality_string
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_OUTPUT_SPLITTER_MAX_BUFFERING_FACTOR = env_float(
+    "RAY_DATA_DEFAULT_OUTPUT_SPLITTER_MAX_BUFFERING_FACTOR", 2
+)
+
+
+class OutputSplitter(InternalQueueOperatorMixin, PhysicalOperator):
+    """An operator that splits the given data into `n` output splits.
+
+    The output bundles of this operator will have a `bundle.output_split_idx` attr
+    set to an integer from [0..n-1]. This operator tries to divide the rows evenly
+    across output splits. If the `equal` option is set, the operator will furthermore
+    guarantee an exact split of rows across outputs, truncating the Dataset.
+
+    Implementation wise, this operator keeps an internal buffer of bundles. The buffer
+    has a minimum size calculated to enable a good locality hit rate, as well as ensure
+    we can satisfy the `equal` requirement.
+
+    OutputSplitter does not provide any ordering guarantees.
+    """
+
+    def __init__(
+        self,
+        input_op: PhysicalOperator,
+        n: int,
+        equal: bool,
+        data_context: DataContext,
+        locality_hints: Optional[List[NodeIdStr]] = None,
+    ):
+        super().__init__(
+            f"split({n}, equal={equal})",
+            [input_op],
+            data_context,
+            num_output_splits=n,
+        )
+        self._equal = equal
+        # Buffer of bundles not yet assigned to output splits.
+        self._buffer: HashLinkedQueue = HashLinkedQueue()
+        # The outputted bundles with output_split attribute set.
+        self._output_queue: FIFOBundleQueue = FIFOBundleQueue()
+        # The number of rows output to each output split so far.
+        self._num_output: List[int] = [0 for _ in range(n)]
+        # The time of the overhead for the output splitter (operator level)
+        self._output_splitter_overhead_time = 0
+
+        if locality_hints is not None:
+            if n != len(locality_hints):
+                raise ValueError(
+                    "Locality hints list must have length `n`: "
+                    f"len({locality_hints}) != {n}"
+                )
+        self._locality_hints = locality_hints
+
+        # To optimize locality, we might defer dispatching of the bundles to allow
+        # for better node affinity by allowing next receiver to wait for a block
+        # with preferred locality (minimizing data movement).
+        #
+        # However, to guarantee liveness we cap buffering to not exceed
+        #
+        #   DEFAULT_OUTPUT_SPLITTER_MAX_BUFFERING_FACTOR * N
+        #
+        # Where N is the number of outputs the sequence is being split into
+        if locality_hints:
+            self._max_buffer_size = DEFAULT_OUTPUT_SPLITTER_MAX_BUFFERING_FACTOR * n
+        else:
+            self._max_buffer_size = 0
+
+        self._locality_hits = 0
+        self._locality_misses = 0
+
+        logger.debug(
+            f"OutputSplitter created: {n=}, {equal=}, {locality_hints=}, "
+            f"{self._max_buffer_size=}"
+        )
+
+    @property
+    @override
+    def _input_queues(self) -> List["BaseBundleQueue"]:
+        return [self._buffer]
+
+    @property
+    @override
+    def _output_queues(self) -> List["BaseBundleQueue"]:
+        return [self._output_queue]
+
+    def num_outputs_total(self) -> Optional[int]:
+        # OutputSplitter does not change the number of blocks,
+        # so we can return the number of blocks from the input op.
+        return self.input_dependencies[0].num_outputs_total()
+
+    def num_output_rows_total(self) -> Optional[int]:
+        # The total number of rows is the same as the number of input rows.
+        return self.input_dependencies[0].num_output_rows_total()
+
+    def start(
+        self,
+        options: ExecutionOptions,
+        block_ref_counter: "BlockRefCounter",
+    ) -> None:
+        if options.preserve_order:
+            # If preserve_order is set, we need to ignore locality hints to ensure determinism.
+            self._locality_hints = None
+            self._max_buffer_size = 0
+
+        super().start(options, block_ref_counter)
+
+    def throttling_disabled(self) -> bool:
+        """Disables resource-based throttling.
+
+        It doesn't make sense to throttle the inputs to this operator, since all that
+        would do is lower the buffer size and prevent us from emitting outputs /
+        reduce the locality hit rate.
+        """
+        return True
+
+    def has_next(self) -> bool:
+        return self._output_queue.has_next()
+
+    def _get_next_inner(self) -> RefBundle:
+        output = self._output_queue.get_next()
+        self._metrics.on_output_dequeued(output)
+        return output
+
+    def get_stats(self) -> StatsDict:
+        return {"split": []}  # TODO(ekl) add split metrics?
+
+    def _extra_metrics(self) -> Dict[str, Any]:
+        stats = {}
+        for i, num in enumerate(self._num_output):
+            stats[f"num_output_{i}"] = num
+        stats["output_splitter_overhead_time"] = self._output_splitter_overhead_time
+        return stats
+
+    def _add_input_inner(self, bundle, input_index) -> None:
+        if bundle.num_rows() is None:
+            raise ValueError("OutputSplitter requires bundles with known row count")
+        self._buffer.add(bundle)
+        self._metrics.on_input_queued(bundle, input_index=0)
+        # Try dispatch buffered bundles
+        self._try_dispatch_bundles()
+
+    def all_inputs_done(self) -> None:
+        super().all_inputs_done()
+
+        # First, attempt to dispatch bundles based on the locality preferences
+        # (if configured)
+        if self._locality_hints:
+            # NOTE: If equal distribution is not requested, we will force
+            #       the dispatching
+            self._try_dispatch_bundles(force=not self._equal)
+
+            if not self._equal:
+                assert not self._buffer, "All bundles should have been dispatched"
+                return
+
+        if not self._buffer:
+            return
+
+        # Otherwise:
+        # Need to finalize distribution of buffered data to output splits.
+        buffer_size = self._buffer.num_rows()
+        max_n = max(self._num_output)
+
+        # First calculate the min rows to add per output to equalize them.
+        allocation = [max_n - n for n in self._num_output]
+        remainder = buffer_size - sum(allocation)
+        # Invariant: buffer should always be large enough to equalize.
+        assert remainder >= 0, (remainder, buffer_size, allocation)
+
+        # Equally distribute remaining rows in buffer to outputs.
+        x = remainder // len(allocation)
+        allocation = [a + x for a in allocation]
+
+        # Execute the split.
+        for i, count in enumerate(allocation):
+            bundles = self._split_from_buffer(count)
+            for b in bundles:
+                b = replace(b, output_split_idx=i)
+                self._output_queue.add(b)
+                self._metrics.on_output_queued(b)
+        # Drain truncated remainder through the metrics layer.
+        # A bare self._buffer.clear() would bypass on_input_dequeued,
+        # orphaning RefBundle references in _metrics._internal_inqueues
+        # that pin ObjectRefs in the object store.
+        self.clear_internal_input_queue()
+
+    def progress_str(self) -> str:
+        if self._locality_hints:
+            return locality_string(self._locality_hits, self._locality_misses)
+        else:
+            return "[locality disabled]"
+
+    def _try_dispatch_bundles(self, force: bool = False) -> None:
+        start_time = time.perf_counter()
+
+        # Currently, there are 2 modes of operation when dispatching
+        # accumulated bundles:
+        #
+        #   1. Best-effort: we do a single pass over the whole buffer
+        #      and try to dispatch all bundles either
+        #
+        #      a) Based on their locality (if feasible)
+        #      b) Longest-waiting if buffer exceeds max-size threshold
+        #
+        #   2. Mandatory: when whole buffer has to be dispatched (for ex,
+        #      upon completion of the dataset execution)
+        #
+        for _ in range(len(self._buffer)):
+            # Get target output index of the next receiver
+            target_output_index = self._select_next_output_index()
+            # Look up preferred bundle
+            preferred_bundle = self._find_preferred_bundle(target_output_index)
+
+            if preferred_bundle:
+                target_bundle = preferred_bundle
+            elif len(self._buffer) >= self._max_buffer_size or force:
+                # If we're not able to find a preferred bundle and buffer size is above
+                # the cap, we pop the longest awaiting and pass to the next receiver
+                target_bundle = self._buffer.peek_next()
+                assert target_bundle is not None
+            else:
+                # Provided that we weren't able to either locate preferred bundle
+                # or dequeue the head one, we bail out from iteration
+                break
+
+            # In case, when we can't safely dispatch (to avoid violating distribution
+            # requirements), short-circuit
+            if not self._can_safely_dispatch(
+                target_output_index, target_bundle.num_rows()
+            ):
+                break
+
+            # Pop preferred bundle from the buffer
+            self._buffer.remove(target_bundle)
+            self._metrics.on_input_dequeued(target_bundle, input_index=0)
+
+            target_bundle = replace(target_bundle, output_split_idx=target_output_index)
+
+            self._num_output[target_output_index] += target_bundle.num_rows()
+            self._output_queue.add(target_bundle)
+            self._metrics.on_output_queued(target_bundle)
+
+            if self._locality_hints:
+                if preferred_bundle:
+                    self._locality_hits += 1
+                else:
+                    self._locality_misses += 1
+
+        self._output_splitter_overhead_time += time.perf_counter() - start_time
+
+    def _select_next_output_index(self) -> int:
+        # Greedily dispatch to the consumer with the least data so far.
+        i, _ = min(enumerate(self._num_output), key=lambda t: t[1])
+        return i
+
+    def _find_preferred_bundle(self, target_output_index: int) -> Optional[RefBundle]:
+        if self._locality_hints:
+            preferred_loc = self._locality_hints[target_output_index]
+
+            # TODO make this more efficient (adding inverse hash-map)
+            for bundle in self._buffer:
+                if preferred_loc in self._get_locations(bundle):
+                    return bundle
+
+        return None
+
+    def _can_safely_dispatch(self, target_index: int, target_num_rows: int) -> bool:
+        if not self._equal:
+            # If not in equals mode, dispatch away with no buffer requirements.
+            return True
+
+        # Simulate dispatching a bundle to the target receiver
+        output_distribution = self._num_output.copy()
+        output_distribution[target_index] += target_num_rows
+        buffer_requirement = self._calculate_buffer_requirement(output_distribution)
+        # Subtract target bundle size from the projected buffer
+        buffer_size = self._buffer.num_rows() - target_num_rows
+        # Check if we have enough rows LEFT after dispatching to equalize.
+        return buffer_size >= buffer_requirement
+
+    def _calculate_buffer_requirement(self, output_distribution: List[int]) -> int:
+        # Calculate the new number of rows that we'd need to equalize the row
+        # distribution after the bundle dispatch.
+        max_n = max(output_distribution)
+        return sum([max_n - n for n in output_distribution])
+
+    def _split_from_buffer(self, nrow: int) -> List[RefBundle]:
+        output = []
+        acc = 0
+        label_selector = self.data_context.execution_options.label_selector
+        while acc < nrow:
+            b = self._buffer.get_next()
+            self._metrics.on_input_dequeued(b, input_index=0)
+            if acc + b.num_rows() <= nrow:
+                output.append(b)
+                acc += b.num_rows()
+            else:
+                input_refs = {entry.ref for entry in b.blocks}
+                left, right = _split(b, nrow - acc, label_selector)
+                # Only register genuinely new blocks created by _split_block.
+                for part in (left, right):
+                    for entry in part.blocks:
+                        if entry.ref not in input_refs:
+                            self._block_ref_counter.on_block_produced(
+                                entry.ref,
+                                entry.metadata.size_bytes or 0,
+                                self.id,
+                            )
+                output.append(left)
+                acc += left.num_rows()
+                self._buffer.add(right)
+                self._metrics.on_input_queued(right, input_index=0)
+                assert acc == nrow, (acc, nrow)
+
+        assert sum(b.num_rows() for b in output) == nrow, (acc, nrow)
+        return output
+
+    @staticmethod
+    def _get_locations(bundle: RefBundle) -> Collection[NodeIdStr]:
+        """Fetches list of node ids holding the objects of the given bundle.
+
+        This method may be overridden for testing.
+
+        Args:
+            bundle: The ``RefBundle`` whose object locations to look up.
+
+        Returns:
+            A list of node ids where the objects in the bundle are located
+        """
+        preferred_locations = bundle.get_preferred_object_locations()
+
+        return preferred_locations.keys()
+
+
+def _split(
+    bundle: RefBundle,
+    left_size: int,
+    label_selector: Optional[Dict[str, str]] = None,
+) -> Tuple[RefBundle, RefBundle]:
+    left_blocks, left_meta = [], []
+    right_blocks, right_meta = [], []
+    acc = 0
+    for entry in bundle.blocks:
+        b = entry.ref
+        m = entry.metadata
+        if acc >= left_size:
+            right_blocks.append(b)
+            right_meta.append(m)
+        elif acc + m.num_rows <= left_size:
+            left_blocks.append(b)
+            left_meta.append(m)
+            acc += m.num_rows
+        else:
+            # Trouble case: split it up.
+            lm, rm = _split_meta(m, left_size - acc)
+            lb, rb = _split_block(b, left_size - acc, label_selector)
+            left_meta.append(lm)
+            right_meta.append(rm)
+            left_blocks.append(lb)
+            right_blocks.append(rb)
+            acc += lm.num_rows
+            assert acc == left_size
+    left = RefBundle(
+        [BlockEntry(b, m) for b, m in zip(left_blocks, left_meta)],
+        owns_blocks=bundle.owns_blocks,
+        schema=bundle.schema,
+    )
+    right = RefBundle(
+        [BlockEntry(b, m) for b, m in zip(right_blocks, right_meta)],
+        owns_blocks=bundle.owns_blocks,
+        schema=bundle.schema,
+    )
+    assert left.num_rows() == left_size
+    assert left.num_rows() + right.num_rows() == bundle.num_rows()
+    return left, right
+
+
+def _split_meta(
+    m: BlockMetadata, left_size: int
+) -> Tuple[BlockMetadata, BlockMetadata]:
+    left_bytes = int(math.floor(m.size_bytes * (left_size / m.num_rows)))
+    left = BlockMetadata(
+        num_rows=left_size,
+        size_bytes=left_bytes,
+        input_files=m.input_files,
+        exec_stats=None,
+    )
+    right = BlockMetadata(
+        num_rows=m.num_rows - left_size,
+        size_bytes=m.size_bytes - left_bytes,
+        input_files=m.input_files,
+        exec_stats=None,
+    )
+    return left, right
+
+
+def _split_block(
+    b: ObjectRef[Block],
+    left_size: int,
+    label_selector: Optional[Dict[str, str]] = None,
+) -> Tuple[ObjectRef[Block], ObjectRef[Block]]:
+    split_single_block = cached_remote_fn(_split_single_block)
+    options: Dict[str, Any] = {"num_cpus": 0, "num_returns": 2}
+    if label_selector:
+        options["label_selector"] = label_selector
+    left, right = split_single_block.options(**options).remote(b, left_size)
+    return left, right
+
+
+def _split_single_block(b: Block, left_size: int) -> Tuple[Block, Block]:
+    acc = BlockAccessor.for_block(b)
+    left = acc.slice(0, left_size)
+    right = acc.slice(left_size, acc.num_rows())
+    assert BlockAccessor.for_block(left).num_rows() == left_size
+    assert BlockAccessor.for_block(right).num_rows() == (acc.num_rows() - left_size)
+    return left, right

python/ray/data/_internal/execution/operators/shuffle_operators/shuffle_map_operator.py

@@ -0,0 +1,454 @@
+import dataclasses
+import functools
+import logging
+import typing
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, Tuple
+
+import ray
+from ray.data._internal.execution.bundle_queue import (
+    BaseBundleQueue,
+    FIFOBundleQueue,
+)
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    ExecutionResources,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    MetadataOpTask,
+    ObjectStoreUsage,
+    OpTask,
+    estimate_total_num_of_blocks,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+)
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_tasks import (
+    SHUFFLE_PEAK_MEMORY_MULTIPLIER,
+    PartitionFn,
+    _shuffle_map_task,
+)
+from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
+from ray.data.block import Block, BlockMetadata, BlockStats
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+if typing.TYPE_CHECKING:
+    from ray.data._internal.progress.base_progress import BaseProgressBar
+
+logger = logging.getLogger(__name__)
+
+
+_PARTITION_ID_SENTINEL = "__partition__"
+
+
+def make_partition_sentinel(partition_id: int) -> Tuple[str, ...]:
+    return (f"{_PARTITION_ID_SENTINEL}{partition_id}",)
+
+
+def extract_partition_id(bundle: RefBundle) -> int:
+    """Recover the partition_id stamped onto an upstream bundle."""
+    for entry in bundle.blocks:
+        files = entry.metadata.input_files
+        if not files:
+            continue
+        for f in files:
+            if f.startswith(_PARTITION_ID_SENTINEL):
+                return int(f[len(_PARTITION_ID_SENTINEL) :])
+    raise ValueError("ShuffleMapOp bundle is missing a partition_id sentinel.")
+
+
+class ShuffleMapOp(InternalQueueOperatorMixin, PhysicalOperator, SubProgressBarMixin):
+    """Map phase of a shuffle: partition inputs and group shards by partition.
+
+    Each map task splits its input into num_partitions shards.  Shards land in a
+    per-partition staging queue as tasks finish.  Once upstream is done and no
+    map tasks remain, each staging queue is drained and merged into a single
+    block per mapper.  That yields one output bundle per partition;
+
+    Args:
+        input_op: Upstream physical operator.
+        data_context: Runtime configuration.
+        num_partitions: Total number of output partitions.
+        partition_fn: Function mapping a pa.Table to Dict[int, pa.Table].
+        pre_map_merge_threshold: Byte threshold per node at which buffered
+            blocks are merged into a single map task.  Set to 0 to disable.
+        map_runtime_env: Optional runtime_env for map tasks; useful to
+            isolate map workers from other ops.
+        map_cpus: CPU request per map task.
+        name: Display name shown in progress bars and logs.
+    """
+
+    _DEFAULT_SHUFFLE_MAP_TASK_NUM_CPUS = 1.0
+    _DEFAULT_PRE_MAP_MERGE_THRESHOLD = 1024 * 1024 * 1024  # 1 GB
+
+    def __init__(
+        self,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        *,
+        num_partitions: int,
+        partition_fn: PartitionFn,
+        pre_map_merge_threshold: int = _DEFAULT_PRE_MAP_MERGE_THRESHOLD,
+        map_runtime_env: Optional[Dict[str, Any]] = None,
+        map_cpus: float = _DEFAULT_SHUFFLE_MAP_TASK_NUM_CPUS,
+        name: str = "ShuffleMap",
+    ):
+        super().__init__(
+            name=name,
+            input_dependencies=[input_op],
+            data_context=data_context,
+        )
+
+        self._num_partitions: int = num_partitions
+        self._partition_fn: PartitionFn = partition_fn
+
+        # -- Map task config -------------------------------------------------
+        self._shuffle_map_task_num_cpus: float = map_cpus
+        self._map_runtime_env: Optional[Dict[str, Any]] = map_runtime_env
+
+        # -- Pre-map merge ---------------------------------------------------
+        self._pre_map_merge_threshold: int = pre_map_merge_threshold
+        self._merge_buffer_refs_by_node: Dict[
+            str, List[ObjectRef[Block]]
+        ] = defaultdict(list)
+        self._merge_buffer_bytes_by_node: Dict[str, int] = defaultdict(int)
+        self._merge_buffer_bundles_by_node: Dict[str, List[RefBundle]] = defaultdict(
+            list
+        )
+
+        # -- Map task tracking -----------------------------------------------
+        self._next_shuffle_map_task_idx: int = 0
+        self._shuffle_map_tasks: Dict[int, MetadataOpTask] = {}
+        self._map_resource_usage = ExecutionResources.zero()
+
+        # -- Per-partition staging queues ------------------------------------
+        self._partition_staging: List[FIFOBundleQueue] = [
+            FIFOBundleQueue() for _ in range(num_partitions)
+        ]
+
+        # -- Per-partition total bytes ---------------------------------------
+        self._partition_bytes: Dict[int, int] = defaultdict(int)
+
+        # -- Output queue  ---------------------------------------------------
+        self._output_queue: FIFOBundleQueue = FIFOBundleQueue()
+        self._partition_bundles_emitted: bool = False
+
+        # -- Stats -----------------------------------------------------------
+        self._total_input_rows: int = 0
+        self._total_input_bytes: int = 0
+        self._map_blocks_stats: List[BlockStats] = []
+
+        # -- Sub-progress bars -----------------------------------------------
+        self._map_bar: Optional["BaseProgressBar"] = None
+
+    @property
+    def _input_queues(self) -> List[BaseBundleQueue]:
+        return []
+
+    @property
+    def _output_queues(self) -> List[BaseBundleQueue]:
+        return [*self._partition_staging, self._output_queue]
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert input_index == 0
+
+        if self._pre_map_merge_threshold > 0:
+            preferred_locs = refs.get_preferred_object_locations()
+            node_id = (
+                max(preferred_locs, key=lambda n: preferred_locs[n])
+                if preferred_locs
+                else "unknown"
+            )
+
+            for block_ref, block_metadata in zip(refs.block_refs, refs.metadata):
+                self._merge_buffer_refs_by_node[node_id].append(block_ref)
+                self._merge_buffer_bytes_by_node[node_id] += (
+                    block_metadata.size_bytes or 0
+                )
+            self._merge_buffer_bundles_by_node[node_id].append(refs)
+
+            if (
+                self._merge_buffer_bytes_by_node[node_id]
+                >= self._pre_map_merge_threshold
+            ):
+                self._flush_merge_buffer(node_id)
+        else:
+            self._submit_shuffle_map_task(
+                list(refs.block_refs),
+                [refs],
+                estimated_bytes=sum((m.size_bytes or 0) for m in refs.metadata),
+            )
+
+    def all_inputs_done(self) -> None:
+        super().all_inputs_done()
+        for node_id in list(self._merge_buffer_refs_by_node.keys()):
+            self._flush_merge_buffer(node_id)
+        self._maybe_emit_partition_bundles()
+
+    def _flush_merge_buffer(self, node_id: str) -> None:
+        block_refs = self._merge_buffer_refs_by_node.pop(node_id, [])
+        bundles = self._merge_buffer_bundles_by_node.pop(node_id, [])
+        estimated_bytes = self._merge_buffer_bytes_by_node.pop(node_id, 0)
+        if not block_refs:
+            for bundle in bundles:
+                bundle.destroy_if_owned()
+            return
+        self._submit_shuffle_map_task(
+            block_refs,
+            bundles,
+            estimated_bytes=estimated_bytes,
+            target_node_id=node_id if node_id != "unknown" else None,
+        )
+
+    def _submit_shuffle_map_task(
+        self,
+        block_refs: List[ObjectRef[Block]],
+        input_bundles: List[RefBundle],
+        estimated_bytes: int = 0,
+        target_node_id: Optional[str] = None,
+    ) -> None:
+        cur_task_idx = self._next_shuffle_map_task_idx
+        self._next_shuffle_map_task_idx += 1
+
+        resources: Dict[str, Any] = {"num_cpus": self._shuffle_map_task_num_cpus}
+        if estimated_bytes > 0:
+            resources["memory"] = estimated_bytes * SHUFFLE_PEAK_MEMORY_MULTIPLIER
+
+        ray_options: Dict[str, Any] = {
+            **resources,
+            "num_returns": self._num_partitions + 1,
+        }
+        if target_node_id is not None:
+            ray_options["scheduling_strategy"] = NodeAffinitySchedulingStrategy(
+                target_node_id, soft=True
+            )
+        if self._map_runtime_env is not None:
+            ray_options["runtime_env"] = self._map_runtime_env
+
+        map_refs = _shuffle_map_task.options(**ray_options).remote(
+            *block_refs,
+            partition_fn=self._partition_fn,
+            num_partitions=self._num_partitions,
+            compression=self.data_context.hash_shuffle_compression,
+        )
+        metadata_ref = map_refs[0]
+        partition_refs = list(map_refs[1:])
+
+        task = MetadataOpTask(
+            task_index=cur_task_idx,
+            object_ref=metadata_ref,
+            task_done_callback=functools.partial(
+                self._handle_map_done, cur_task_idx, partition_refs, input_bundles
+            ),
+            task_resource_bundle=ExecutionResources.from_resource_dict(resources),
+        )
+        self._shuffle_map_tasks[cur_task_idx] = task
+        requested = task.get_requested_resource_bundle()
+        assert requested is not None
+        self._map_resource_usage = self._map_resource_usage.add(requested)
+
+        all_blocks_meta = tuple(
+            BlockEntry(ref=ref, metadata=meta)
+            for bundle in input_bundles
+            for ref, meta in zip(bundle.block_refs, bundle.metadata)
+        )
+        self._metrics.on_task_submitted(
+            cur_task_idx,
+            RefBundle(all_blocks_meta, schema=None, owns_blocks=False),
+            task_id=task.get_task_id(),
+        )
+
+        if self._map_bar is not None:
+            _, _, num_rows = estimate_total_num_of_blocks(
+                cur_task_idx + 1,
+                self.upstream_op_num_outputs(),
+                self._metrics,
+                total_num_tasks=None,
+            )
+            self._map_bar.update(total=num_rows)
+
+    def _handle_map_done(
+        self,
+        task_idx: int,
+        partition_refs: List[ObjectRef[Block]],
+        input_bundles: List[RefBundle],
+    ) -> None:
+        task = self._shuffle_map_tasks.pop(task_idx)
+        requested = task.get_requested_resource_bundle()
+        assert requested is not None
+        self._map_resource_usage = self._map_resource_usage.subtract(requested)
+
+        # `task_done_callback` fires only after the metadata ref is ready,
+        # so this is just local deserialization.
+        input_meta, shard_sizes, output_schema = ray.get(task.get_waitable())
+
+        for partition_id, ref in enumerate(partition_refs):
+            rows, nbytes = shard_sizes.get(partition_id, (0, 0))
+            shard_meta = BlockMetadata(
+                num_rows=rows,
+                size_bytes=nbytes,
+                exec_stats=None,
+                input_files=None,
+            )
+            shard_bundle = RefBundle(
+                (BlockEntry(ref=ref, metadata=shard_meta),),
+                schema=output_schema,
+                owns_blocks=True,
+            )
+            self._partition_staging[partition_id].add(shard_bundle)
+            self._partition_bytes[partition_id] += nbytes
+
+        for bundle in input_bundles:
+            bundle.destroy_if_owned()
+
+        self._total_input_rows += input_meta.num_rows or 0
+        self._total_input_bytes += input_meta.size_bytes or 0
+        self._map_blocks_stats.append(input_meta.to_stats())
+
+        self._metrics.on_task_finished(
+            task_idx,
+            None,
+            task_exec_stats=None,
+            task_exec_driver_stats=None,
+        )
+
+        if self._map_bar is not None:
+            self._map_bar.update(increment=input_meta.num_rows or 0)
+
+        self._maybe_emit_partition_bundles()
+
+    def _maybe_emit_partition_bundles(self) -> None:
+        """Drain each partition's staging queue into one output bundle.
+
+        Every partition is staged (empty partitions carry a schema-only shard),
+        so this emits exactly num_partitions bundles.
+        """
+        if self._partition_bundles_emitted:
+            return
+        if self._shuffle_map_tasks or self._merge_buffer_refs_by_node:
+            return
+        if not self._inputs_complete:
+            return
+
+        self._partition_bundles_emitted = True
+
+        for partition_id in range(self._num_partitions):
+            staging = self._partition_staging[partition_id]
+            if not staging.has_next():
+                continue
+            shards: List[RefBundle] = []
+            while staging.has_next():
+                shards.append(staging.get_next())
+
+            merged = RefBundle.merge_ref_bundles(shards)
+            # Stamp the partition_id sentinel onto the merged bundle's
+            # first block so the downstream reducer can recover the
+            # partition this bundle represents.
+            stamped_blocks = []
+            for i, entry in enumerate(merged.blocks):
+                meta = entry.metadata
+                if i == 0:
+                    meta = dataclasses.replace(
+                        meta, input_files=make_partition_sentinel(partition_id)
+                    )
+                stamped_blocks.append(BlockEntry(ref=entry.ref, metadata=meta))
+            stamped = RefBundle(
+                tuple(stamped_blocks),
+                schema=merged.schema,
+                owns_blocks=merged.owns_blocks,
+            )
+            self._output_queue.add(stamped)
+            self._metrics.on_output_queued(stamped)
+
+    def has_next(self) -> bool:
+        return self._output_queue.has_next()
+
+    def _get_next_inner(self) -> RefBundle:
+        bundle: RefBundle = self._output_queue.get_next()
+        self._metrics.on_output_dequeued(bundle)
+        return bundle
+
+    def get_partition_bytes(self) -> Dict[int, int]:
+        return dict(self._partition_bytes)
+
+    def get_active_tasks(self) -> List[OpTask]:
+        return list(self._shuffle_map_tasks.values())
+
+    def has_execution_finished(self) -> bool:
+        if (
+            self._shuffle_map_tasks
+            or self._merge_buffer_refs_by_node
+            or not self._partition_bundles_emitted
+            or self._output_queue.has_next()
+        ):
+            return False
+        return super().has_execution_finished()
+
+    def has_completed(self) -> bool:
+        return (
+            not self._shuffle_map_tasks
+            and not self._merge_buffer_refs_by_node
+            and self._partition_bundles_emitted
+            and not self._output_queue.has_next()
+            and super().has_completed()
+        )
+
+    def _do_shutdown(self, force: bool = False) -> None:
+        super()._do_shutdown(force)
+        self._shuffle_map_tasks.clear()
+        self._merge_buffer_refs_by_node.clear()
+        for bundles in self._merge_buffer_bundles_by_node.values():
+            for bundle in bundles:
+                bundle.destroy_if_owned()
+        self._merge_buffer_bundles_by_node.clear()
+        self._merge_buffer_bytes_by_node.clear()
+        for queue in self._partition_staging:
+            queue.clear()
+        self._output_queue.clear()
+
+    def get_stats(self) -> Dict[str, List[BlockStats]]:
+        return {self._name: self._map_blocks_stats}
+
+    def num_output_rows_total(self) -> Optional[int]:
+        return self._total_input_rows if self._total_input_rows > 0 else None
+
+    def current_logical_usage(self) -> ExecutionResources:
+        return ExecutionResources(
+            cpu=self._map_resource_usage.cpu,
+            memory=self._map_resource_usage.memory,
+        )
+
+    def estimate_object_store_usage(self, state) -> ObjectStoreUsage:
+        return ObjectStoreUsage(internal=0, outputs=0)
+
+    def incremental_resource_usage(self) -> ExecutionResources:
+        avg_input = self._metrics.average_bytes_inputs_per_task
+        memory = int(avg_input * SHUFFLE_PEAK_MEMORY_MULTIPLIER) if avg_input else 0
+        return ExecutionResources(
+            cpu=self._shuffle_map_task_num_cpus,
+            memory=memory,
+        )
+
+    def min_scheduling_resources(self) -> ExecutionResources:
+        return self.incremental_resource_usage()
+
+    def progress_str(self) -> str:
+        maps_done = self._next_shuffle_map_task_idx - len(self._shuffle_map_tasks)
+        parts = [f"map: {maps_done}/{self._next_shuffle_map_task_idx}"]
+        total_merge_buf = sum(
+            len(refs) for refs in self._merge_buffer_refs_by_node.values()
+        )
+        if total_merge_buf:
+            parts.append(f"merge_buf: {total_merge_buf}")
+        return ", ".join(parts)
+
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        return ["Map"]
+
+    def set_sub_progress_bar(self, name: str, pg: "BaseProgressBar") -> None:
+        if name == "Map":
+            self._map_bar = pg

python/ray/data/_internal/execution/operators/shuffle_operators/shuffle_reduce_operator.py

@@ -0,0 +1,448 @@
+import functools
+import logging
+import typing
+from collections import deque
+from typing import Any, Dict, List, Optional, Union
+
+import pyarrow as pa
+
+import ray
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    ExecutionResources,
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    DataOpTask,
+    OpTask,
+    TaskExecDriverStats,
+    estimate_total_num_of_blocks,
+)
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (
+    ShuffleMapOp,
+    extract_partition_id,
+)
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_tasks import (
+    SHUFFLE_PEAK_MEMORY_MULTIPLIER,
+    ReduceFn,
+    _shuffle_reduce_task,
+)
+from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
+from ray.data.block import BlockAccessor, BlockStats, TaskExecWorkerStats, to_stats
+from ray.data.context import DataContext
+
+if typing.TYPE_CHECKING:
+    from ray.data._internal.execution.operators.map_transformer import MapTransformer
+    from ray.data._internal.progress.base_progress import BaseProgressBar
+
+logger = logging.getLogger(__name__)
+
+
+class ShuffleReduceOp(PhysicalOperator, SubProgressBarMixin):
+    """Reduce phase of a shuffle.
+
+    Supports one or more co-partitioned upstream `ShuffleMapOp`s.  With a single
+    input this is the unary reduce used by repartition/sort.  With multiple
+    inputs (e.g. join) every input must be partitioned into the same
+    `num_partitions`; this op pairs up the per-partition bundles across all
+    inputs and hands the reducer one shard list per input.
+
+    Args:
+        input_op: Upstream `ShuffleMapOp`, or a list of them (one per input).
+            For an N-input reduce the reducer receives shards in this order.
+        data_context: Runtime configuration.
+        num_partitions: Total number of output partitions.  Must match the
+            value used by every paired `ShuffleMapOp`.
+        reduce_fn: Function called once per partition with all shards to
+            combine them into output blocks.  Receives
+            `(partition_id, tables_by_input)` where `tables_by_input` is
+            aligned with `input_op`.
+        disallow_block_splitting: If True, output blocks are emitted as-is
+            without being reshaped to `target_max_block_size`.
+        reduce_ray_remote_args: Remote args for the reducer tasks.
+        name: Display name shown in progress bars and logs.
+        fused_output_map_transformer: Set by ``FuseOperators`` when a
+            ``TaskPoolMapOperator`` directly downstream is fused into this
+            reduce: each reduce task applies it to its output blocks before
+            yielding.
+        fused_output_map_task_kwargs: Per-task kwargs the fused map injects into
+            its ``TaskContext``.
+        fused_output_map_target_max_block_size_override: The fused map op's
+            block-size override.
+    """
+
+    _DEFAULT_SHUFFLE_REDUCE_TASK_NUM_CPUS = 1.0
+
+    def __init__(
+        self,
+        input_op: Union[ShuffleMapOp, List[ShuffleMapOp]],
+        data_context: DataContext,
+        *,
+        num_partitions: int,
+        reduce_fn: ReduceFn,
+        disallow_block_splitting: bool = False,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+        name: str = "ShuffleReduce",
+        fused_output_map_transformer: Optional["MapTransformer"] = None,
+        fused_output_map_task_kwargs: Optional[Dict[str, Any]] = None,
+        fused_output_map_target_max_block_size_override: Optional[int] = None,
+    ):
+        input_ops: List[PhysicalOperator] = (
+            [input_op] if isinstance(input_op, ShuffleMapOp) else list(input_op)
+        )
+        assert input_ops, "ShuffleReduceOp requires at least one upstream ShuffleMapOp"
+        super().__init__(
+            name=name,
+            input_dependencies=input_ops,
+            data_context=data_context,
+        )
+
+        self._num_inputs: int = len(input_ops)
+        self._num_partitions: int = num_partitions
+        self._reduce_fn: ReduceFn = reduce_fn
+        self._disallow_block_splitting: bool = disallow_block_splitting
+
+        # -- Reduce task config & tracking -----------------------------------
+        self._reduce_ray_remote_args: Dict[str, Any] = dict(
+            reduce_ray_remote_args or {}
+        )
+        self._shuffle_reduce_tasks: Dict[int, DataOpTask] = {}
+        self._num_reduce_tasks_submitted: int = 0
+
+        # -- Per-partition pairing across inputs -----------------------------
+        # partition_id -> input_index -> the single bundle that input emitted
+        # for that partition.  A reduce task is submitted once all inputs have
+        # delivered their bundle for a partition.  With a single input a bundle
+        # pairs immediately.
+        self._pending_inputs: Dict[int, Dict[int, RefBundle]] = {}
+
+        # -- Fused downstream map --------------------------------------------
+        self._fused_output_map_transformer = fused_output_map_transformer
+        self._fused_output_map_task_kwargs = fused_output_map_task_kwargs or {}
+        self._fused_output_map_target_max_block_size_override = (
+            fused_output_map_target_max_block_size_override
+        )
+
+        # -- Output queue ----------------------------------------------------
+        self._output_queue: deque = deque()
+
+        # -- Stats -----------------------------------------------------------
+        self._output_blocks_stats: List[BlockStats] = []
+
+        # -- Sub-progress bars -----------------------------------------------
+        self._reduce_bar: Optional["BaseProgressBar"] = None
+
+    def _reduce_task_remote_args(self, memory_estimate: int) -> Dict[str, Any]:
+        remote_args: Dict[str, Any] = {
+            "num_cpus": self._DEFAULT_SHUFFLE_REDUCE_TASK_NUM_CPUS,
+            "scheduling_strategy": "SPREAD",
+        }
+        if memory_estimate > 0:
+            remote_args["memory"] = memory_estimate
+        remote_args.update(self._reduce_ray_remote_args)
+        remote_args["num_returns"] = "streaming"
+        return remote_args
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        """Buffer this input's partition-bundle; submit when all inputs paired.
+
+        Each upstream bundle is a single partition's shards (M blocks from M
+        mappers) from one input.  The partition_id is encoded in the first
+        block's `input_files`.  A reduce task runs only once every input has
+        delivered its bundle for that partition (immediately for the common
+        single-input case), so the reducer sees all inputs' shards together.
+        This is the framework-gated entry point — the executor only calls it
+        when all configured backpressure policies say the op can accept another
+        input.
+        """
+        assert 0 <= input_index < self._num_inputs
+
+        if not refs.block_refs:
+            refs.destroy_if_owned()
+            return
+
+        partition_id = extract_partition_id(refs)
+
+        # Single-input empty-partition fast path: emit one empty block instead
+        # of launching a reduce task.  Skipped for multi-input reduces (an outer
+        # join's empty side can still produce rows) and when a downstream map is
+        # fused in (the map must run even on empty partitions, e.g. a Write).
+        schema = refs.schema
+        if (
+            self._num_inputs == 1
+            and self._fused_output_map_transformer is None
+            and isinstance(schema, pa.Schema)
+            and not any((m.num_rows or 0) for m in refs.metadata)
+        ):
+            self._emit_empty_partition(refs, schema)
+            return
+
+        pending = self._pending_inputs.setdefault(partition_id, {})
+        assert input_index not in pending, (
+            f"input {input_index} already delivered a bundle for partition "
+            f"{partition_id}; each ShuffleMapOp must emit at most one bundle "
+            f"per partition"
+        )
+        pending[input_index] = refs
+        if len(pending) == self._num_inputs:
+            del self._pending_inputs[partition_id]
+            self._submit_reduce_task(
+                partition_id, [pending[i] for i in range(self._num_inputs)]
+            )
+
+    def all_inputs_done(self) -> None:
+        super().all_inputs_done()
+        # Every upstream input is now exhausted.  A partition still missing an
+        # input's bundle will never receive it -- that input ran no map tasks for
+        # this partition's key space (e.g. a block-less input).  Flush such
+        # partitions with an empty placeholder for each missing input so the op
+        # can complete instead of hanging on a never-paired partition.
+        for partition_id in list(self._pending_inputs.keys()):
+            pending = self._pending_inputs.pop(partition_id)
+            bundles = [
+                pending.get(i) or RefBundle((), schema=None, owns_blocks=True)
+                for i in range(self._num_inputs)
+            ]
+            self._submit_reduce_task(partition_id, bundles)
+
+    def _submit_reduce_task(self, partition_id: int, bundles: List[RefBundle]) -> None:
+        """Submit one reduce task for a fully-paired partition."""
+        shard_refs_by_input = []
+        metrics_blocks = []
+        estimated_bytes = 0
+        for bundle in bundles:
+            shard_refs_by_input.append(list(bundle.block_refs))
+            metrics_blocks.extend(bundle.blocks)
+            estimated_bytes += sum((m.size_bytes or 0) for m in bundle.metadata)
+
+        reduce_options = self._reduce_task_remote_args(
+            int(estimated_bytes * SHUFFLE_PEAK_MEMORY_MULTIPLIER)
+            if estimated_bytes > 0
+            else 0
+        )
+
+        target_max_block_size = (
+            None
+            if self._disallow_block_splitting
+            else self.data_context.target_max_block_size
+        )
+
+        map_task_context = None
+        if self._fused_output_map_transformer is not None:
+            map_task_context = TaskContext(
+                task_idx=partition_id,
+                op_name=self.name,
+                target_max_block_size_override=(
+                    self._fused_output_map_target_max_block_size_override
+                ),
+            )
+            map_task_context.kwargs.update(self._fused_output_map_task_kwargs)
+
+        block_gen = _shuffle_reduce_task.options(**reduce_options).remote(
+            shard_refs_by_input,  # pyrefly: ignore[bad-argument-type]
+            partition_id,
+            self._reduce_fn,
+            target_max_block_size,
+            self.data_context.hash_shuffle_reduce_batch_size,
+            self.data_context.hash_shuffle_reduce_get_timeout_s,
+            self._fused_output_map_transformer,
+            map_task_context,
+            self.data_context,
+        )
+        metrics_bundle = RefBundle(
+            tuple(metrics_blocks), schema=None, owns_blocks=False
+        )
+
+        data_task = DataOpTask(
+            task_index=partition_id,
+            streaming_gen=block_gen,
+            block_ref_counter=self._block_ref_counter,
+            producer_id=self.id,
+            output_ready_callback=functools.partial(
+                self._handle_reduce_output_ready, partition_id
+            ),
+            task_done_callback=functools.partial(
+                self._handle_reduce_done, partition_id, bundles
+            ),
+            task_resource_bundle=ExecutionResources.from_resource_dict(reduce_options),
+            operator_name=self.name,
+        )
+
+        assert partition_id not in self._shuffle_reduce_tasks, (
+            f"partition_id {partition_id} already has an in-flight reducer "
+            f"task; ShuffleMapOp must emit at most one bundle per partition"
+        )
+        self._shuffle_reduce_tasks[partition_id] = data_task
+        self._num_reduce_tasks_submitted += 1
+        self._metrics.on_task_submitted(
+            partition_id, metrics_bundle, task_id=data_task.get_task_id()
+        )
+
+    def _emit_empty_partition(self, refs: RefBundle, schema: pa.Schema) -> None:
+        """Emit one empty output block for an empty partition.
+
+        The partition contributed no rows, so there is nothing to reduce; we
+        build the empty block from the schema the map stage propagated onto
+        the bundle and queue it as this partition's single output block.
+        """
+        empty_block = schema.empty_table()
+        block_meta = BlockAccessor.for_block(empty_block).get_metadata()
+        out_bundle = RefBundle(
+            (
+                BlockEntry(
+                    ref=ray.put(empty_block),  # pyrefly: ignore[bad-argument-type]
+                    metadata=block_meta,
+                ),
+            ),
+            schema=schema,
+            owns_blocks=True,
+        )
+        refs.destroy_if_owned()
+
+        # Empty partition creates a new block; register it for memory tracking.
+        self._block_ref_counter.on_block_produced(
+            out_bundle.blocks[0].ref,  # pyrefly: ignore[bad-argument-type]
+            block_meta.size_bytes or 0,
+            self.id,
+        )
+        self._num_reduce_tasks_submitted += 1
+        self._output_queue.append(out_bundle)
+        self._metrics.on_output_queued(out_bundle)
+        _, num_outputs, num_rows = estimate_total_num_of_blocks(
+            self._num_reduce_tasks_submitted,
+            self.upstream_op_num_outputs(),
+            self._metrics,
+            total_num_tasks=self._num_partitions,
+        )
+        self._estimated_num_output_bundles = num_outputs
+        self._estimated_output_num_rows = num_rows
+        if self._reduce_bar is not None:
+            self._reduce_bar.update(increment=0, total=self.num_output_rows_total())
+
+    def has_next(self) -> bool:
+        return len(self._output_queue) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        bundle: RefBundle = self._output_queue.popleft()
+        self._metrics.on_output_dequeued(bundle)
+        self._output_blocks_stats.extend(to_stats(bundle.metadata))
+        return bundle
+
+    def get_active_tasks(self) -> List[OpTask]:
+        return list(self._shuffle_reduce_tasks.values())
+
+    def _handle_reduce_output_ready(self, partition_id: int, bundle: RefBundle) -> None:
+        self._output_queue.append(bundle)
+        self._metrics.on_output_queued(bundle)
+        self._metrics.on_task_output_generated(task_index=partition_id, output=bundle)
+        _, num_outputs, num_rows = estimate_total_num_of_blocks(
+            self._num_reduce_tasks_submitted,
+            self.upstream_op_num_outputs(),
+            self._metrics,
+            total_num_tasks=self._num_partitions,
+        )
+        self._estimated_num_output_bundles = num_outputs
+        self._estimated_output_num_rows = num_rows
+        if self._reduce_bar is not None:
+            self._reduce_bar.update(
+                increment=bundle.num_rows() or 0,
+                total=self.num_output_rows_total(),
+            )
+
+    def _handle_reduce_done(
+        self,
+        partition_id: int,
+        input_bundles: List[RefBundle],
+        exc: Optional[Exception],
+        task_exec_stats: Optional[TaskExecWorkerStats],
+        task_exec_driver_stats: Optional[TaskExecDriverStats],
+    ) -> None:
+        """Callback when a reduce task finishes (with or without exception)."""
+        for input_bundle in input_bundles:
+            input_bundle.destroy_if_owned()
+        if partition_id not in self._shuffle_reduce_tasks:
+            return
+        self._shuffle_reduce_tasks.pop(partition_id)
+        self._metrics.on_task_finished(
+            task_index=partition_id,
+            exception=exc,
+            task_exec_stats=task_exec_stats,
+            task_exec_driver_stats=task_exec_driver_stats,
+        )
+        if exc:
+            logger.error(
+                f"Reduce of partition {partition_id} failed: {exc}", exc_info=exc
+            )
+
+    def has_execution_finished(self) -> bool:
+        if self._shuffle_reduce_tasks or self._output_queue or self._pending_inputs:
+            return False
+        return super().has_execution_finished()
+
+    def has_completed(self) -> bool:
+        return (
+            not self._shuffle_reduce_tasks
+            and not self._output_queue
+            and not self._pending_inputs
+            and super().has_completed()
+        )
+
+    def _do_shutdown(self, force: bool = False) -> None:
+        super()._do_shutdown(force)
+        self._shuffle_reduce_tasks.clear()
+        self._output_queue.clear()
+        for pending in self._pending_inputs.values():
+            for bundle in pending.values():
+                bundle.destroy_if_owned()
+        self._pending_inputs.clear()
+
+    def get_stats(self) -> Dict[str, List[BlockStats]]:
+        return {self._name: self._output_blocks_stats}
+
+    def num_output_rows_total(self) -> Optional[int]:
+        # Multi-input reduces (e.g. join) can grow or shrink the row count, so it
+        # is unknown until the reducers run; a single-input reduce preserves it.
+        if self._num_inputs > 1:
+            return None
+        upstream = self.input_dependencies[0]
+        assert isinstance(upstream, ShuffleMapOp)
+        return upstream.num_output_rows_total()
+
+    def current_logical_usage(self) -> ExecutionResources:
+        usage = ExecutionResources.zero()
+        for task in self._shuffle_reduce_tasks.values():
+            bundle = task.get_requested_resource_bundle()
+            if bundle is None:
+                continue
+            usage = usage.add(ExecutionResources(cpu=bundle.cpu, memory=bundle.memory))
+        return usage
+
+    def incremental_resource_usage(self) -> ExecutionResources:
+        """Per-task resource ask for the framework's budget allocator."""
+        memory = 0
+        for upstream in self.input_dependencies:
+            assert isinstance(upstream, ShuffleMapOp)
+            sizes = [b for b in upstream.get_partition_bytes().values() if b > 0]
+            if sizes:
+                avg_bytes = sum(sizes) / len(sizes)
+                memory += int(avg_bytes * SHUFFLE_PEAK_MEMORY_MULTIPLIER)
+        return ExecutionResources.from_resource_dict(
+            self._reduce_task_remote_args(memory)
+        )
+
+    def min_scheduling_resources(self) -> ExecutionResources:
+        return self.incremental_resource_usage()
+
+    def progress_str(self) -> str:
+        submitted = self._num_reduce_tasks_submitted
+        done = submitted - len(self._shuffle_reduce_tasks)
+        return f"reduce: {done}/{submitted}"
+
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        return ["Reduce"]
+
+    def set_sub_progress_bar(self, name: str, pg: "BaseProgressBar") -> None:
+        if name == "Reduce":
+            self._reduce_bar = pg

python/ray/data/_internal/execution/operators/shuffle_operators/shuffle_tasks.py

@@ -0,0 +1,330 @@
+"""Shared remote tasks + helpers for ShuffleMapOp / ShuffleReduceOp."""
+
+import logging
+import math
+import time
+import typing
+from dataclasses import replace
+from typing import Callable, Dict, Generator, Iterable, List, Optional, Tuple, Union
+
+import pyarrow as pa
+
+import ray
+from ray import ObjectRef
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.util import yield_block_with_stats
+from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption
+from ray.data._internal.table_block import TableBlockAccessor
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+    BlockType,
+    TaskExecWorkerStats,
+)
+from ray.data.context import DataContext
+from ray.exceptions import GetTimeoutError
+
+if typing.TYPE_CHECKING:
+    from ray.data._internal.execution.operators.map_transformer import MapTransformer
+
+logger = logging.getLogger(__name__)
+
+PartitionFn = Callable[[pa.Table], Dict[int, pa.Table]]
+ReduceFn = Callable[[int, List[List[pa.Table]]], Iterable[Block]]
+
+# Peak working-set of a shuffle map/reduce task is ~2x the input bytes
+SHUFFLE_PEAK_MEMORY_MULTIPLIER = 2
+
+
+def _ipc_write_options(compression: Optional[str]) -> pa.ipc.IpcWriteOptions:
+    """Arrow IPC write options for the given shard compression codec.
+
+    Args:
+        compression: A pyarrow codec name such as "lz4" or "zstd", or "none"
+            (or None) to write shards uncompressed. See pyarrow.Codec for the
+            full list of supported codecs:
+            https://arrow.apache.org/docs/python/generated/pyarrow.Codec.html
+
+    Returns:
+        IpcWriteOptions for encoding shards; no compression for "none"/None.
+    """
+    if not compression or compression == "none":
+        return pa.ipc.IpcWriteOptions()
+    return pa.ipc.IpcWriteOptions(compression=pa.Codec(compression))
+
+
+def _partition_blocks_to_shards(
+    blocks: Tuple[Block, ...], partition_fn: PartitionFn
+) -> Dict[int, List[pa.Table]]:
+    """Partition each block independently, grouping shards by partition id."""
+    partition_accumulators: Dict[int, List[pa.Table]] = {}
+    for block in blocks:
+        block = TableBlockAccessor.try_convert_block_type(
+            block, block_type=BlockType.ARROW
+        )
+        if block.num_rows == 0:
+            continue
+        assert isinstance(block, pa.Table), f"Expected pa.Table, got {type(block)}"
+        if any(col.num_chunks > 1 for col in block.columns):
+            block = block.combine_chunks()
+        block_partitions = partition_fn(block)
+        for partition_id, shard in block_partitions.items():
+            if shard.num_rows > 0:
+                partition_accumulators.setdefault(partition_id, []).append(shard)
+        del block, block_partitions
+    return partition_accumulators
+
+
+def _encode_partition_ipc(
+    table: pa.Table,
+    ipc_write_options: pa.ipc.IpcWriteOptions,
+) -> pa.Buffer:
+    """Encode one partition's shard as a single Arrow IPC stream."""
+    if table.num_columns > 0:
+        table = table.combine_chunks()
+
+    sink = pa.BufferOutputStream()
+    with pa.ipc.new_stream(sink, table.schema, options=ipc_write_options) as writer:
+        for batch in table.to_batches():
+            writer.write_batch(batch)
+    return sink.getvalue()
+
+
+@ray.remote  # pyrefly: ignore[no-matching-overload]
+def _shuffle_map_task(
+    *blocks: Block,
+    partition_fn: PartitionFn,
+    num_partitions: int,
+    compression: Optional[str],
+) -> Tuple[
+    Union[Tuple[BlockMetadata, Dict[int, Tuple[int, int]], "pa.Schema"], pa.Buffer],
+    ...,
+]:
+    """Map stage: partition the input blocks and return one shard per partition."""
+    stats = BlockExecStats.builder()
+
+    # Use BlockAccessor so we also work for non-Arrow blocks (pandas, numpy)
+    accessors = [BlockAccessor.for_block(b) for b in blocks]
+    total_rows = sum(a.num_rows() for a in accessors)
+    total_bytes = sum((a.size_bytes() or 0) for a in accessors)
+
+    ipc_write_options = _ipc_write_options(compression)
+    output_schema = TableBlockAccessor.try_convert_block_type(
+        blocks[0], block_type=BlockType.ARROW
+    ).schema
+    empty_shard = _encode_partition_ipc(output_schema.empty_table(), ipc_write_options)
+
+    partition_accumulators = (
+        {} if total_rows == 0 else _partition_blocks_to_shards(blocks, partition_fn)
+    )
+
+    shard_sizes: Dict[int, Tuple[int, int]] = {}
+    partition_bufs: List[pa.Buffer] = []
+    for partition_id in range(num_partitions):
+        tables = partition_accumulators.pop(partition_id, None)
+        if not tables:
+            partition_bufs.append(empty_shard)
+            continue
+        merged = pa.concat_tables(tables) if len(tables) > 1 else tables[0]
+        shard_sizes[partition_id] = (merged.num_rows, merged.nbytes)
+        partition_bufs.append(_encode_partition_ipc(merged, ipc_write_options))
+        del merged
+
+    input_meta = BlockAccessor.for_block(blocks[0]).get_metadata(
+        block_exec_stats=stats.build(block_ser_time_s=0),
+    )
+    input_meta = replace(input_meta, num_rows=total_rows, size_bytes=total_bytes)
+    return (input_meta, shard_sizes, output_schema), *partition_bufs
+
+
+def _read_partition_ipc(buf: pa.Buffer) -> Optional[pa.Table]:
+    """Decompress one partition shard."""
+    if len(buf) == 0:
+        return None
+    reader = pa.ipc.open_stream(buf)
+    schema = reader.schema
+    batches: List[pa.RecordBatch] = []
+    while True:
+        try:
+            batch = reader.read_next_batch()
+        except StopIteration:
+            break
+        if batch.num_rows > 0:
+            batches.append(batch)
+    return pa.Table.from_batches(batches, schema=schema)
+
+
+# Warn once a shard fetch has stalled for this fraction of the fail timeout
+_REDUCE_GET_WARN_AT_FRACTION = 1 / 3
+
+
+def _get_shard_batch(
+    batch: List[ObjectRef],
+    partition_id: int,
+    batch_index: int,
+    num_batches: int,
+    timeout_s: float,
+) -> List[Optional[pa.Buffer]]:
+    """``ray.get`` a batch of shard refs, warning then failing if the fetch stalls.
+
+    Args:
+        batch: Shard ObjectRefs to fetch (a slice of one partition's shards).
+        partition_id: Partition this reducer owns (for logging).
+        batch_index: 0-based index of this batch within the partition.
+        num_batches: Total number of batches for the partition (for logging).
+        timeout_s: ``ray.get`` timeout in seconds.  A non-positive value disables
+            the timeout (single blocking fetch).
+
+    Returns:
+        The dereferenced shard buffers (some entries may be ``None``).
+
+    Raises:
+        GetTimeoutError: If the shards are not available within ``timeout_s``.
+    """
+    if timeout_s <= 0:
+        return ray.get(batch)
+
+    wait_start_s = time.perf_counter()
+    warn_timeout_s = timeout_s * _REDUCE_GET_WARN_AT_FRACTION
+    try:
+        return ray.get(batch, timeout=warn_timeout_s)
+    except GetTimeoutError:
+        logger.warning(
+            f"Shuffle reduce task for partition {partition_id} has waited "
+            f"{time.perf_counter() - wait_start_s:.0f}s for {len(batch)} "
+            f"shard(s) in batch {batch_index + 1}/{num_batches}."
+        )
+
+    try:
+        return ray.get(batch, timeout=timeout_s - warn_timeout_s)
+    except GetTimeoutError:
+        logger.error(
+            f"Shuffle reduce task for partition {partition_id} timed out after "
+            f"{time.perf_counter() - wait_start_s:.0f}s waiting for {len(batch)} "
+            f"shard(s) in batch {batch_index + 1}/{num_batches}."
+        )
+        raise
+
+
+def _gather_input_shards(
+    shard_refs: List[ObjectRef],
+    partition_id: int,
+    batch_size: int,
+    get_timeout_s: float,
+) -> List[pa.Table]:
+    """Fetch + decompress every shard of one input for one partition."""
+    tables: List[pa.Table] = []
+    num_batches = math.ceil(len(shard_refs) / batch_size) if batch_size else 0
+    for batch_index, batch_start in enumerate(range(0, len(shard_refs), batch_size)):
+        batch = shard_refs[batch_start : batch_start + batch_size]
+        for buf in _get_shard_batch(
+            batch, partition_id, batch_index, num_batches, get_timeout_s
+        ):
+            if buf is None:
+                continue
+            table = _read_partition_ipc(buf)
+            if table is None:
+                continue
+            tables.append(table)
+    return tables
+
+
+@ray.remote
+def _shuffle_reduce_task(
+    shard_refs_by_input: List[List[ObjectRef]],
+    partition_id: int,
+    reduce_fn: ReduceFn,
+    target_max_block_size: Optional[int],
+    batch_size: int,
+    get_timeout_s: float,
+    map_transformer: Optional["MapTransformer"],
+    map_task_context: Optional["TaskContext"],
+    data_context: Optional["DataContext"],
+) -> Generator[Union[Block, bytes], None, None]:
+    """Reduce stage: fetch this partition's shards and run reduce_fn over them.
+
+    ``shard_refs_by_input`` carries one shard-ref list per upstream input -- one
+    for single-input shuffles (repartition/sort), several for multi-input ones
+    (e.g. join).  Every input's full shard list is accumulated, then reduce_fn is
+    called once with ``tables_by_input`` aligned with ``shard_refs_by_input``.
+
+    Args:
+        shard_refs_by_input: Per-input lists of ObjectRefs to this partition's
+            IPC shards from every mapper.  May contain None for empty shards.
+        partition_id: Partition this reducer owns.
+        reduce_fn: User-supplied reduce callable.
+        target_max_block_size: Output block size.  None emits blocks as-is.
+        batch_size: Number of shard refs to ray.get() at a time.
+        get_timeout_s: Timeout for batch ray.get().
+        map_transformer: Fused downstream map applied to reduce output (or None).
+        map_task_context: TaskContext for the fused map, built by the reduce op
+            -- carries task_idx, op_name, the block-size override, and per-task
+            kwargs (e.g. a Write's ``write_uuid``); None when nothing is fused.
+        data_context: DataContext to install for the fused map (or None).
+    """
+    start_time_s = time.perf_counter()
+
+    output_buffer: Optional[BlockOutputBuffer] = None
+
+    def _yield_with_stats(block: Block):
+        """Yield a block then its pickled metadata (streaming-gen protocol)."""
+
+        def build_metadata(block_ser_time_s):
+            exec_stats = BlockExecStats.builder()
+            exec_stats.finish()
+            return BlockMetadataWithSchema.from_block(
+                block,
+                block_exec_stats=exec_stats.build(block_ser_time_s=block_ser_time_s),
+                task_exec_stats=TaskExecWorkerStats(
+                    task_wall_time_s=time.perf_counter() - start_time_s,
+                ),
+            )
+
+        yield from yield_block_with_stats(block, build_metadata)
+
+    def _flush(tables_by_input: List[List[pa.Table]]):
+        nonlocal output_buffer
+        if output_buffer is None:
+            output_buffer = BlockOutputBuffer(
+                OutputBlockSizeOption.of(
+                    target_max_block_size=target_max_block_size,
+                )
+            )
+        for block in reduce_fn(partition_id, tables_by_input):
+            output_buffer.add_block(block)
+            # Yield raw blocks: a fused map (and `_yield_with_stats`) is applied
+            # downstream of ``_reduce_output_blocks``.
+            yield from output_buffer.iter_ready_blocks()
+
+    def _reduce_output_blocks():
+        # Gather every input's full shard list, then call reduce_fn exactly once
+        # with all inputs together (no streaming: a multi-input reducer needs
+        # every input's shards, and single-input reducers run blocking too).
+        tables_by_input = [
+            _gather_input_shards(shard_refs, partition_id, batch_size, get_timeout_s)
+            for shard_refs in shard_refs_by_input
+        ]
+        if any(tables_by_input):
+            yield from _flush(tables_by_input)
+
+        # Finalize the buffer to flush any partial block.
+        if output_buffer is not None:
+            output_buffer.finalize()
+            yield from output_buffer.iter_ready_blocks()
+
+    if map_transformer is None:
+        for block in _reduce_output_blocks():
+            yield from _yield_with_stats(block)
+    else:
+        assert map_task_context is not None and data_context is not None
+        with DataContext.current(data_context), TaskContext.current(map_task_context):
+            map_transformer.override_target_max_block_size(
+                map_task_context.target_max_block_size_override
+            )
+            for block in map_transformer.apply_transform(
+                _reduce_output_blocks(), map_task_context
+            ):
+                yield from _yield_with_stats(block)

python/ray/data/_internal/execution/operators/sub_progress.py

@@ -0,0 +1,31 @@
+import typing
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+if typing.TYPE_CHECKING:
+    from ray.data._internal.progress.base_progress import BaseProgressBar
+
+
+class SubProgressBarMixin(ABC):
+    """Abstract class for operators that support sub-progress bars"""
+
+    @abstractmethod
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        """
+        Returns list of sub-progress bar names
+
+        This is used to create the sub-progress bars in the progress manager.
+        Note that sub-progress bars will be created in the order returned by
+        this method.
+        """
+        ...
+
+    @abstractmethod
+    def set_sub_progress_bar(self, name: str, pg: "BaseProgressBar"):
+        """
+        Sets sub-progress bars
+
+        name: name of sub-progress bar
+        pg: a progress bar. Can be sub-progress bars for rich, tqdm, etc.
+        """
+        ...

python/ray/data/_internal/execution/operators/task_pool_map_operator.py

@@ -0,0 +1,281 @@
+import copy
+import warnings
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+from typing_extensions import override
+
+from ray.data._internal.execution.bundle_queue import (
+    BaseBundleQueue,
+    RebundleQueue,
+)
+from ray.data._internal.execution.interfaces import (
+    ExecutionResources,
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.operators.map_operator import (
+    MapOperator,
+    _map_task,
+)
+from ray.data._internal.execution.operators.map_transformer import MapTransformer
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data.context import DataContext
+
+
+class TaskPoolMapOperator(MapOperator):
+    """A MapOperator implementation that executes tasks on a task pool."""
+
+    def __init__(
+        self,
+        map_transformer: MapTransformer,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        name: str = "TaskPoolMap",
+        target_max_block_size_override: Optional[int] = None,
+        min_rows_per_bundle: Optional[int] = None,
+        ref_bundler: Optional[RebundleQueue] = None,
+        max_concurrency: Optional[int] = None,
+        supports_fusion: bool = True,
+        map_task_kwargs: Optional[Dict[str, Any]] = None,
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+        on_start: Optional[Callable[[Optional["pa.Schema"]], None]] = None,
+        isolate_workers: bool = False,
+        default_logical_memory_enabled: bool = False,
+    ):
+        """Create an TaskPoolMapOperator instance.
+
+        Args:
+            map_transformer: The :class:`MapTransformer` to apply to each ref
+                bundle input.
+            input_op: Operator generating input data for this op.
+            data_context: The :class:`DataContext` to use for this operator.
+            name: The name of this operator.
+            target_max_block_size_override: Override for target max-block-size.
+            min_rows_per_bundle: The number of rows to gather per batch passed to the
+                transform_fn, or None to use the block size. Setting the batch size is
+                important for the performance of GPU-accelerated transform functions.
+                The actual rows passed may be less if the dataset is small.
+            ref_bundler: The ref bundler to use for this operator.
+            max_concurrency: The maximum number of Ray tasks to use concurrently,
+                or None to use as many tasks as possible.
+            supports_fusion: Whether this operator supports fusion with other operators.
+            map_task_kwargs: A dictionary of kwargs to pass to the map task. You can
+                access these kwargs through the `TaskContext.kwargs` dictionary.
+            ray_remote_args_fn: A function that returns a dictionary of remote args
+                passed to each map worker. The purpose of this argument is to generate
+                dynamic arguments for each actor/task, and will be called each time
+                prior to initializing the worker. Args returned from this dict will
+                always override the args in ``ray_remote_args``. Note: this is an
+                advanced, experimental feature.
+            ray_remote_args: Customize the :func:`ray.remote` args for this op's tasks.
+            on_start: Optional callback invoked with the schema from the first input
+                bundle before any tasks are submitted.
+            isolate_workers: If ``True``, ensure that other operators' tasks don't get
+                scheduled on the same worker processes as this operator's. This flag
+                is useful to prevent side-effects from affecting other operators, like
+                large PyArrow memory allocations.
+            default_logical_memory_enabled: If ``True``, the operator launches tasks
+                with a default logical ``memory``. The method for choosing the
+                default is an implementation detail.
+        """
+        super().__init__(
+            map_transformer,
+            input_op,
+            data_context,
+            name,
+            target_max_block_size_override,
+            min_rows_per_bundle,
+            ref_bundler,
+            supports_fusion,
+            map_task_kwargs,
+            ray_remote_args_fn,
+            ray_remote_args,
+            on_start,
+            default_logical_memory_enabled,
+        )
+
+        self._isolate_workers = isolate_workers
+
+        if max_concurrency is not None and max_concurrency <= 0:
+            raise ValueError(f"max_concurrency have to be > 0 (got {max_concurrency})")
+
+        self._max_concurrency = max_concurrency
+        self._current_logical_usage = ExecutionResources.zero()
+
+        # NOTE: Unlike static Ray remote args, dynamic arguments extracted from the
+        #       blocks themselves are going to be passed inside `fn.options(...)`
+        #       invocation
+        ray_remote_static_args = {
+            **(self._ray_remote_args or {}),
+            "num_returns": "streaming",
+            "_labels": {self._OPERATOR_ID_LABEL_KEY: self.id},
+        }
+
+        # Ray Core doesn't share workers for tasks with different `runtime_env`s. We use
+        # this property to implicitly ensure that this operator's tasks run on isolated
+        # workers.
+        if self._isolate_workers:
+            ray_remote_static_args = self._add_unique_runtime_env(
+                ray_remote_static_args
+            )
+
+        self._map_task = cached_remote_fn(_map_task, **ray_remote_static_args)
+
+    def _add_unique_runtime_env(
+        self, ray_remote_args: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Return a copy of the remote args with a runtime env that's unique to this
+        operator.
+        """
+        ray_remote_args = copy.deepcopy(ray_remote_args)
+
+        runtime_env = ray_remote_args.get("runtime_env", {})
+        env_vars = ray_remote_args.get("env_vars", {})
+        env_vars["__RAY_DATA_OPERATOR_ID"] = self.id
+        runtime_env["env_vars"] = env_vars
+
+        ray_remote_args["runtime_env"] = runtime_env
+        return ray_remote_args
+
+    @property
+    def isolate_workers(self) -> bool:
+        """Return whether this operator launches tasks on isolated worker processes.
+
+        If ``True``, other operators' tasks won't get scheduled on the same worker
+        processes as this operator's. This flag is useful to prevent side-effects
+        from affecting other operators, like large PyArrow memory allocations.
+        """
+        return self._isolate_workers
+
+    @property
+    @override
+    def _input_queues(self) -> List["BaseBundleQueue"]:
+        return [self._block_ref_bundler]
+
+    @property
+    @override
+    def _output_queues(self) -> List["BaseBundleQueue"]:
+        return [self._output_queue]
+
+    def _try_schedule_task(self, bundle: RefBundle, strict: bool):
+        # Notify first input for deferred initialization (e.g., Iceberg schema evolution).
+        self._notify_first_input(bundle)
+        # Submit the task as a normal Ray task.
+        ctx = TaskContext(
+            task_idx=self._next_data_task_idx,
+            op_name=self.name,
+            target_max_block_size_override=self.target_max_block_size_override,
+        )
+
+        dynamic_ray_remote_args = self._get_dynamic_ray_remote_args(input_bundle=bundle)
+        dynamic_ray_remote_args["name"] = self.name
+        logical_usage = ExecutionResources.from_resource_dict(dynamic_ray_remote_args)
+
+        if (
+            "_generator_backpressure_num_objects" not in dynamic_ray_remote_args
+            and self.data_context._max_num_blocks_in_streaming_gen_buffer is not None
+        ):
+            # The `_generator_backpressure_num_objects` parameter should be
+            # `2 * _max_num_blocks_in_streaming_gen_buffer` because we yield
+            # 2 objects for each block: the block and the block metadata.
+            dynamic_ray_remote_args["_generator_backpressure_num_objects"] = (
+                2 * self.data_context._max_num_blocks_in_streaming_gen_buffer
+            )
+
+        gen = self._map_task.options(**dynamic_ray_remote_args).remote(
+            self._map_transformer_ref,
+            self._data_context_ref,
+            ctx,
+            *bundle.block_refs,
+            slices=bundle.slices,
+            **self.get_map_task_kwargs(),
+        )
+
+        self._current_logical_usage = self._current_logical_usage.add(logical_usage)
+
+        def task_done_callback():
+            self._current_logical_usage = self._current_logical_usage.subtract(
+                logical_usage
+            )
+
+        self._submit_data_task(gen, bundle, task_done_callback=task_done_callback)
+
+    def progress_str(self) -> str:
+        return ""
+
+    def current_logical_usage(self) -> ExecutionResources:
+        return self._current_logical_usage
+
+    def pending_logical_usage(self) -> ExecutionResources:
+        return ExecutionResources()
+
+    def incremental_resource_usage(self) -> ExecutionResources:
+        return self.per_task_resource_allocation()
+
+    def per_task_resource_allocation(self) -> ExecutionResources:
+        return ExecutionResources(
+            cpu=self._ray_remote_args.get("num_cpus", 0),
+            gpu=self._ray_remote_args.get("num_gpus", 0),
+            memory=self._ray_remote_args.get("memory", 0),
+        )
+
+    def min_scheduling_resources(
+        self: "PhysicalOperator",
+    ) -> ExecutionResources:
+        return self.incremental_resource_usage()
+
+    def get_max_concurrency_limit(self) -> Optional[int]:
+        return self._max_concurrency
+
+    def min_max_resource_requirements(
+        self,
+    ) -> Tuple[ExecutionResources, ExecutionResources]:
+        """Returns min/max resource requirements for this operator.
+
+        - Min: resources needed for one task (minimum to make progress)
+        - Max: resources for max_concurrency tasks (if set), else infinite
+        """
+        per_task = self.per_task_resource_allocation()
+        obj_store_per_task = (
+            self._metrics.obj_store_mem_max_pending_output_per_task or 0
+        )
+
+        min_resource_usage = per_task.copy(object_store_memory=obj_store_per_task)
+
+        # Cap resources to 0 if this operator doesn't use them.
+        # This prevents operators from hoarding resource budget they don't need.
+        max_concurrency = (
+            self._max_concurrency if self._max_concurrency is not None else float("inf")
+        )
+        max_resource_usage = ExecutionResources(
+            cpu=0 if per_task.cpu == 0 else per_task.cpu * max_concurrency,
+            gpu=0 if per_task.gpu == 0 else per_task.gpu * max_concurrency,
+            memory=0 if per_task.memory == 0 else per_task.memory * max_concurrency,
+            # Set the max `object_store_memory` requirement to 'inf', because we
+            # don't know how much data each task can output.
+            object_store_memory=float("inf"),
+        )
+
+        return min_resource_usage, max_resource_usage
+
+    def all_inputs_done(self):
+        super().all_inputs_done()
+
+        if (
+            self._max_concurrency is not None
+            and self._metrics.num_inputs_received < self._max_concurrency
+        ):
+            warnings.warn(
+                f"The maximum number of concurrent tasks for '{self.name}' is set to "
+                f"{self._max_concurrency}, but the operator only received "
+                f"{self._metrics.num_inputs_received} input(s). This means that the "
+                f"operator can launch at most {self._metrics.num_inputs_received} "
+                "task(s), which is less than the concurrency limit. You might be able "
+                "to increase the number of concurrent tasks by configuring "
+                "`override_num_blocks` earlier in the pipeline."
+            )

python/ray/data/_internal/execution/operators/union_operator.py

@@ -0,0 +1,161 @@
+from typing import TYPE_CHECKING, List, Optional
+
+from typing_extensions import override
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.block_ref_counter import BlockRefCounter
+
+from ray.data._internal.execution.bundle_queue import BaseBundleQueue, FIFOBundleQueue
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+    NAryOperator,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+
+class UnionOperator(InternalQueueOperatorMixin, NAryOperator):
+    """An operator that combines output blocks from
+    two or more input operators into a single output."""
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        *input_ops: PhysicalOperator,
+    ):
+        """Create a UnionOperator.
+
+        Args:
+            data_context: The :class:`DataContext` to use for this operator.
+            *input_ops: Operators generating input data for this operator to union.
+        """
+
+        # By default, union does not preserve the order of output blocks.
+        # To preserve the order, configure ExecutionOptions accordingly.
+        self._preserve_order = False
+
+        # Intermediary buffers used to store blocks from each input dependency.
+        # Only used when `self._prserve_order` is True.
+        self._input_buffers: List["FIFOBundleQueue"] = [
+            FIFOBundleQueue() for _ in range(len(input_ops))
+        ]
+
+        self._input_done_flags: List[bool] = [False] * len(input_ops)
+
+        self._output_buffer = FIFOBundleQueue()
+        self._stats: StatsDict = {"Union": []}
+        self._current_input_index = 0
+        super().__init__(data_context, *input_ops)
+
+    @property
+    @override
+    def _input_queues(self) -> List["BaseBundleQueue"]:
+        return self._input_buffers
+
+    @property
+    @override
+    def _output_queues(self) -> List["BaseBundleQueue"]:
+        return [self._output_buffer]
+
+    def start(
+        self,
+        options: ExecutionOptions,
+        block_ref_counter: "BlockRefCounter",
+    ):
+        # Whether to preserve deterministic ordering of output blocks.
+        # When True, blocks are emitted in round-robin order across inputs,
+        # ensuring the same input always produces the same output order.
+        self._preserve_order = options.preserve_order
+        super().start(options, block_ref_counter)
+
+    def num_outputs_total(self) -> Optional[int]:
+        num_outputs = 0
+        for input_op in self.input_dependencies:
+            input_num_outputs = input_op.num_outputs_total()
+            if input_num_outputs is None:
+                return None
+            num_outputs += input_num_outputs
+        return num_outputs
+
+    def num_output_rows_total(self) -> Optional[int]:
+        total_rows = 0
+        for input_op in self.input_dependencies:
+            input_num_rows = input_op.num_output_rows_total()
+            if input_num_rows is None:
+                return None
+            total_rows += input_num_rows
+        return total_rows
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.has_completed()
+        assert 0 <= input_index <= len(self._input_dependencies), input_index
+        if self._preserve_order:
+            self._input_buffers[input_index].add(refs)
+            self._metrics.on_input_queued(refs, input_index=input_index)
+            self._try_round_robin()
+        else:
+            self._output_buffer.add(refs)
+            self._metrics.on_output_queued(refs)
+
+    def input_done(self, input_index: int) -> None:
+        self._input_done_flags[input_index] = True
+        if self._preserve_order:
+            self._try_round_robin()
+
+    def all_inputs_done(self) -> None:
+        super().all_inputs_done()
+
+        if not self._preserve_order:
+            return
+        self._try_round_robin()
+        assert all(not buffer.has_next() for buffer in self._input_buffers)
+
+    def has_next(self) -> bool:
+        # Check if the output buffer still contains at least one block.
+        return len(self._output_buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        refs = self._output_buffer.get_next()
+        self._metrics.on_output_dequeued(refs)
+        return refs
+
+    def get_stats(self) -> StatsDict:
+        return self._stats
+
+    def _try_round_robin(self) -> None:
+        """Try to move blocks from input buffers to output in round-robin order.
+
+        Pulls one block from the current input, then advances to the next.
+        If the current input's buffer is empty but not done, we return
+        without advancing to the next input so the scheduling won't be blocked.
+
+        This ensures deterministic ordering of output blocks:
+        - We iterate through inputs in a fixed order (0, 1, 2, ..., 0, 1, ...).
+        - We only advance to the next input after consuming exactly one block
+          from the current input (or if the current input is exhausted).
+        - If an input is not ready (empty but not done), we return
+          rather than skipping it, preserving the round-robin sequence.
+        """
+        num_inputs = len(self._input_buffers)
+
+        while True:
+            buffer = self._input_buffers[self._current_input_index]
+
+            if buffer.has_next():
+                refs = buffer.get_next()
+                self._metrics.on_input_dequeued(
+                    refs, input_index=self._current_input_index
+                )
+                self._output_buffer.add(refs)
+                self._metrics.on_output_queued(refs)
+            elif not self._input_done_flags[self._current_input_index] or all(
+                not buffer.has_next() for buffer in self._input_buffers
+            ):
+                return
+
+            self._current_input_index = (self._current_input_index + 1) % num_inputs

python/ray/data/_internal/execution/operators/zip_operator.py

@@ -0,0 +1,333 @@
+import collections
+import itertools
+from dataclasses import replace
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+from typing_extensions import override
+
+import ray
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.bundle_queue import BaseBundleQueue, FIFOBundleQueue
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+    NAryOperator,
+)
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.split import _split_at_indices
+from ray.data._internal.stats import StatsDict
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    to_stats,
+)
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+
+    from ray.data.block import BlockMetadataWithSchema
+
+
+class ZipOperator(InternalQueueOperatorMixin, NAryOperator):
+    """An operator that zips its inputs together.
+
+    NOTE: the implementation is bulk for now, which materializes all its inputs in
+    object store, before starting execution. Should re-implement it as a streaming
+    operator in the future.
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        *input_ops: PhysicalOperator,
+    ):
+        """Create a ZipOperator.
+
+        Args:
+            data_context: The :class:`DataContext` to use for this operator.
+            *input_ops: Operators generating input data for this operator to zip.
+        """
+        assert len(input_ops) >= 2
+        self._input_buffers: List[FIFOBundleQueue] = [
+            FIFOBundleQueue() for _ in range(len(input_ops))
+        ]
+        self._output_buffer: FIFOBundleQueue = FIFOBundleQueue()
+        self._stats: StatsDict = {}
+        super().__init__(
+            data_context,
+            *input_ops,
+        )
+
+    @property
+    @override
+    def _input_queues(self) -> List["BaseBundleQueue"]:
+        return self._input_buffers
+
+    @property
+    @override
+    def _output_queues(self) -> List["BaseBundleQueue"]:
+        return [self._output_buffer]
+
+    def num_outputs_total(self) -> Optional[int]:
+        num_outputs = None
+        for input_op in self.input_dependencies:
+            input_num_outputs = input_op.num_outputs_total()
+            if input_num_outputs is None:
+                continue
+            if num_outputs is None:
+                num_outputs = input_num_outputs
+            else:
+                num_outputs = max(num_outputs, input_num_outputs)
+        return num_outputs
+
+    def num_output_rows_total(self) -> Optional[int]:
+        num_rows = None
+        for input_op in self.input_dependencies:
+            input_num_rows = input_op.num_output_rows_total()
+            if input_num_rows is None:
+                continue
+            if num_rows is None:
+                num_rows = input_num_rows
+            else:
+                num_rows = max(num_rows, input_num_rows)
+        return num_rows
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.has_completed()
+        assert 0 <= input_index <= len(self._input_dependencies), input_index
+        self._input_buffers[input_index].add(refs)
+        self._metrics.on_input_queued(refs, input_index=input_index)
+
+    def all_inputs_done(self) -> None:
+        assert len(self._output_buffer) == 0, len(self._output_buffer)
+
+        # Start with the first input buffer
+        while self._input_buffers[0].has_next():
+            refs = self._input_buffers[0].get_next()
+            self._output_buffer.add(refs)
+            self._metrics.on_input_dequeued(refs, input_index=0)
+
+        # Process each additional input buffer
+        for idx, input_buffer in enumerate(self._input_buffers[1:], start=1):
+            output_buffer, self._stats = self._zip(self._output_buffer, input_buffer)
+            self._output_buffer = FIFOBundleQueue(bundles=output_buffer)
+
+            # Clear the input buffer AFTER using it in _zip
+            while input_buffer.has_next():
+                refs = input_buffer.get_next()
+                self._metrics.on_input_dequeued(refs, input_index=idx)
+
+        # Zipping creates new blocks; register them for memory tracking.
+        for ref in self._output_buffer:
+            for entry in ref.blocks:
+                self._block_ref_counter.on_block_produced(
+                    entry.ref, entry.metadata.size_bytes or 0, self.id
+                )
+            self._metrics.on_output_queued(ref)
+
+        super().all_inputs_done()
+
+    def has_next(self) -> bool:
+        return len(self._output_buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        refs = self._output_buffer.get_next()
+        self._metrics.on_output_dequeued(refs)
+        return refs
+
+    def get_stats(self) -> StatsDict:
+        return self._stats
+
+    def throttling_disabled(self) -> bool:
+        # TODO revert once zip becomes streaming
+        return True
+
+    def _zip(
+        self,
+        left_input: FIFOBundleQueue,
+        right_input: FIFOBundleQueue,
+    ) -> Tuple[collections.deque[RefBundle], StatsDict]:
+        """Zip the RefBundles from `left_input` and `right_input` together.
+
+        Zip is done in 2 steps: aligning blocks, and zipping blocks from
+        both sides.
+
+        Aligning blocks (optional): check the blocks from `left_input` and
+        `right_input` are aligned or not, i.e. if having different number of blocks, or
+        having different number of rows in some blocks. If not aligned, repartition the
+        smaller input with `_split_at_indices` to align with larger input.
+
+        Zipping blocks: after blocks from both sides are aligned, zip
+        blocks from both sides together in parallel.
+        """
+        left_entries: List[BlockEntry] = []
+        for bundle in left_input:
+            left_entries.extend(bundle.blocks)
+        right_entries: List[BlockEntry] = []
+        for bundle in right_input:
+            right_entries.extend(bundle.blocks)
+
+        left_block_rows, left_block_bytes = self._calculate_blocks_rows_and_bytes(
+            left_entries
+        )
+        right_block_rows, right_block_bytes = self._calculate_blocks_rows_and_bytes(
+            right_entries
+        )
+
+        # Check that both sides have the same number of rows.
+        # TODO(Clark): Support different number of rows via user-directed
+        # dropping/padding.
+        total_left_rows = sum(left_block_rows)
+        total_right_rows = sum(right_block_rows)
+        if total_left_rows != total_right_rows:
+            raise ValueError(
+                "Cannot zip datasets of different number of rows: "
+                f"{total_left_rows}, {total_right_rows}"
+            )
+
+        # Whether the left and right input sides are inverted
+        input_side_inverted = False
+        if sum(right_block_bytes) > sum(left_block_bytes):
+            # Make sure that right side is smaller, so we minimize splitting
+            # work when aligning both sides.
+            # TODO(Clark): Improve this heuristic for minimizing splitting work,
+            # e.g. by generating the splitting plans for each route (via
+            # _generate_per_block_split_indices) and choosing the plan that splits
+            # the least cumulative bytes.
+            left_entries, right_entries = right_entries, left_entries
+            left_block_rows, right_block_rows = right_block_rows, left_block_rows
+            input_side_inverted = True
+
+        # Get the split indices that will align both sides.
+        indices = list(itertools.accumulate(left_block_rows))
+        indices.pop(-1)
+
+        # Split other at the alignment indices, such that for every block from
+        # left side, we have a list of blocks from right side that have the same
+        # cumulative number of rows as that left block.
+        # NOTE: _split_at_indices has a no-op fastpath if the blocks are already
+        # aligned.
+        # Determine the ownership of the blocks being split, accounting for the
+        # potential swap above. We must not free blocks that are shared with
+        # other operators (e.g., when the input RefBundle has owns_blocks=False
+        # because it comes from a materialized dataset).
+        split_side_owned = all(
+            b.owns_blocks for b in (left_input if input_side_inverted else right_input)
+        )
+        label_selector = self.data_context.execution_options.label_selector
+        aligned_right_blocks_with_metadata = _split_at_indices(
+            [(e.ref, e.metadata) for e in right_entries],
+            indices,
+            owned_by_consumer=split_side_owned,
+            block_rows=right_block_rows,
+            label_selector=label_selector,
+        )
+        del right_entries
+
+        left_blocks = [e.ref for e in left_entries]
+        right_blocks_list = aligned_right_blocks_with_metadata[0]
+        del left_entries, aligned_right_blocks_with_metadata
+
+        zip_one_block = cached_remote_fn(_zip_one_block, num_returns=2)
+        if label_selector:
+            zip_one_block = zip_one_block.options(label_selector=label_selector)
+
+        output_blocks = []
+        output_metadata_schema = []
+        for left_block, right_blocks in zip(left_blocks, right_blocks_list):
+            # For each block from left side, zip it together with 1 or more blocks from
+            # right side. We're guaranteed to have that left_block has the same number
+            # of rows as right_blocks has cumulatively.
+            res, meta_with_schema = zip_one_block.remote(
+                left_block, *right_blocks, inverted=input_side_inverted
+            )
+            output_blocks.append(res)
+            output_metadata_schema.append(meta_with_schema)
+
+        # Early release memory.
+        del left_blocks, right_blocks_list
+
+        # TODO(ekl) it might be nice to have a progress bar here.
+        output_metadata_schema: List[BlockMetadataWithSchema] = ray.get(
+            output_metadata_schema
+        )
+
+        output_refs: collections.deque[RefBundle] = collections.deque()
+        input_owned = all(b.owns_blocks for b in left_input)
+        for block, meta_with_schema in zip(output_blocks, output_metadata_schema):
+            output_refs.append(
+                RefBundle(
+                    [BlockEntry(block, meta_with_schema.metadata)],
+                    owns_blocks=input_owned,
+                    schema=meta_with_schema.schema,
+                )
+            )
+        stats = {self._name: to_stats(output_metadata_schema)}
+
+        # Clean up inputs.
+        for ref in left_input:
+            ref.destroy_if_owned()
+        for ref in right_input:
+            ref.destroy_if_owned()
+
+        return output_refs, stats
+
+    def _calculate_blocks_rows_and_bytes(
+        self,
+        entries: List[BlockEntry],
+    ) -> Tuple[List[int], List[int]]:
+        """Calculate the number of rows and size in bytes for a list of blocks with
+        metadata.
+        """
+        get_num_rows_and_bytes = cached_remote_fn(_get_num_rows_and_bytes)
+        label_selector = self.data_context.execution_options.label_selector
+        if label_selector:
+            get_num_rows_and_bytes = get_num_rows_and_bytes.options(
+                label_selector=label_selector
+            )
+        block_rows = []
+        block_bytes = []
+        for entry in entries:
+            metadata = entry.metadata
+            if metadata.num_rows is None or metadata.size_bytes is None:
+                # Need to fetch number of rows or size in bytes, so just fetch both.
+                num_rows, size_bytes = ray.get(get_num_rows_and_bytes.remote(entry.ref))
+                # Cache on the block metadata.
+                metadata = replace(metadata, num_rows=num_rows, size_bytes=size_bytes)
+            block_rows.append(metadata.num_rows)
+            block_bytes.append(metadata.size_bytes)
+        return block_rows, block_bytes
+
+
+def _zip_one_block(
+    block: Block, *other_blocks: Block, inverted: bool = False
+) -> Tuple[Block, "BlockMetadataWithSchema"]:
+    """Zip together `block` with `other_blocks`."""
+    stats = BlockExecStats.builder()
+    # Concatenate other blocks.
+    # TODO(Clark): Extend BlockAccessor.zip() to work with N other blocks,
+    # so we don't need to do this concatenation.
+    builder = DelegatingBlockBuilder()
+    for other_block in other_blocks:
+        builder.add_block(other_block)
+    other_block = builder.build()
+    if inverted:
+        # Swap blocks if ordering was inverted during block alignment splitting.
+        block, other_block = other_block, block
+    # Zip block and other blocks.
+    result = BlockAccessor.for_block(block).zip(other_block)
+    from ray.data.block import BlockMetadataWithSchema
+
+    return result, BlockMetadataWithSchema.from_block(
+        result, block_exec_stats=stats.build()
+    )
+
+
+def _get_num_rows_and_bytes(block: Block) -> Tuple[int, int]:
+    block = BlockAccessor.for_block(block)
+    return block.num_rows(), block.size_bytes()