chore: import upstream snapshot with attribution

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

@@ -0,0 +1,57 @@
+import enum
+import logging
+import os
+from typing import TYPE_CHECKING, Optional
+
+from .ranker import DefaultRanker, Ranker
+
+if TYPE_CHECKING:
+    from ... import DataContext
+    from .resource_manager import OpResourceAllocator, ResourceManager
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION = os.environ.get(
+    "RAY_DATA_USE_OP_RESOURCE_ALLOCATOR_VERSION", "V1"
+)
+
+
+class OpResourceAllocatorVersion(str, enum.Enum):
+    V1 = "V1"  # ReservationOpResourceAllocator
+
+
+def create_resource_allocator(
+    resource_manager: "ResourceManager",
+    data_context: "DataContext",
+) -> Optional["OpResourceAllocator"]:
+    """Creates ``OpResourceAllocator`` instances.``"""
+
+    if not data_context.op_resource_reservation_enabled:
+        # This is a historical kill-switch to disable resource allocator, that
+        # will be soon deprecated and removed.
+        return None
+
+    logger.debug(
+        f"Using op resource allocator version: "
+        f"{DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION!r}"
+    )
+
+    if DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION == OpResourceAllocatorVersion.V1:
+        from .resource_manager import ReservationOpResourceAllocator
+
+        return ReservationOpResourceAllocator(
+            resource_manager,
+            reservation_ratio=data_context.op_resource_reservation_ratio,
+        )
+
+    else:
+        raise ValueError(
+            "Resource allocator version of "
+            f"'{DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION}' is not supported"
+        )
+
+
+def create_ranker() -> Ranker:
+    """Create a ranker instance based on environment and configuration."""
+    return DefaultRanker()

python/ray/data/_internal/execution/autoscaling_requester.py

@@ -0,0 +1,126 @@
+import math
+import threading
+import time
+from typing import Dict, List
+
+import ray
+
+# Resource requests are considered stale after this number of seconds, and
+# will be purged.
+RESOURCE_REQUEST_TIMEOUT = 60
+PURGE_INTERVAL = RESOURCE_REQUEST_TIMEOUT * 2
+
+# When the autoscaling is driven by memory pressure and there are abundant
+# CPUs to support incremental CPUs needed to launch more tasks, we'll translate
+# memory pressure into an artificial request of CPUs. The amount of CPUs we'll
+# request is ARTIFICIAL_CPU_SCALING_FACTOR * ray.cluster_resources()["CPU"].
+ARTIFICIAL_CPU_SCALING_FACTOR = 1.2
+
+
+@ray.remote(num_cpus=0, max_restarts=-1, max_task_retries=-1)
+class AutoscalingRequester:
+    """Actor to make resource requests to autoscaler for the datasets.
+
+    The resource requests are set to timeout after RESOURCE_REQUEST_TIMEOUT seconds.
+    For those live requests, we keep track of the last request made for each execution,
+    which overrides all previous requests it made; then sum the requested amounts
+    across all executions as the final request to the autoscaler.
+    """
+
+    def __init__(self):
+        # execution_id -> (List[Dict], expiration timestamp)
+        self._resource_requests = {}
+        # TTL for requests.
+        self._timeout = RESOURCE_REQUEST_TIMEOUT
+
+        self._self_handle = ray.get_runtime_context().current_actor
+
+        # Start a thread to purge expired requests periodically.
+        def purge_thread_run():
+            while True:
+                time.sleep(PURGE_INTERVAL)
+                # Call purge_expired_requests() as an actor task,
+                # so we don't need to handle multi-threading.
+                ray.get(self._self_handle.purge_expired_requests.remote())
+
+        self._purge_thread = threading.Thread(target=purge_thread_run, daemon=True)
+        self._purge_thread.start()
+
+    def purge_expired_requests(self):
+        self._purge()
+        ray.autoscaler.sdk.request_resources(bundles=self._aggregate_requests())
+
+    def request_resources(self, req: List[Dict], execution_id: str):
+        # Purge expired requests before making request to autoscaler.
+        self._purge()
+        # For the same execution_id, we track the latest resource request and
+        # its expiration timestamp.
+        self._resource_requests[execution_id] = (
+            req,
+            time.time() + self._timeout,
+        )
+        # We aggregate the resource requests across all execution_id's to Ray
+        # autoscaler.
+        ray.autoscaler.sdk.request_resources(bundles=self._aggregate_requests())
+
+    def _purge(self):
+        # Purge requests that are stale.
+        now = time.time()
+        for k, (_, t) in list(self._resource_requests.items()):
+            if t < now:
+                self._resource_requests.pop(k)
+
+    def _aggregate_requests(self) -> List[Dict]:
+        req = []
+        for _, (r, _) in self._resource_requests.items():
+            req.extend(r)
+
+        def get_cpus(req):
+            num_cpus = 0
+            for r in req:
+                if "CPU" in r:
+                    num_cpus += r["CPU"]
+            return num_cpus
+
+        # Round up CPUs to exceed total cluster CPUs so it can actually upscale.
+        # This is to handle the issue where the autoscaling is driven by memory
+        # pressure (rather than CPUs) from streaming executor. In such case, simply
+        # asking for incremental CPUs (e.g. 1 CPU for each ready operator) may not
+        # actually be able to trigger autoscaling if existing CPUs in cluster can
+        # already satisfy the incremental CPUs request.
+        num_cpus = get_cpus(req)
+        if num_cpus > 0:
+            total = ray.cluster_resources()
+            if "CPU" in total and num_cpus <= total["CPU"]:
+                delta = (
+                    math.ceil(ARTIFICIAL_CPU_SCALING_FACTOR * total["CPU"]) - num_cpus
+                )
+                req.extend([{"CPU": 1}] * delta)
+
+        return req
+
+    def _test_set_timeout(self, ttl):
+        """Set the timeout. This is for test only"""
+        self._timeout = ttl
+
+
+# Creating/getting an actor from multiple threads is not safe.
+# https://github.com/ray-project/ray/issues/41324
+_autoscaling_requester_lock: threading.RLock = threading.RLock()
+
+
+def get_or_create_autoscaling_requester_actor():
+    # Pin the autoscaling requester actor to the local node so it fate-shares with the driver.
+    # Note: for Ray Client, the ray.get_runtime_context().get_node_id() should
+    # point to the head node.
+    label_selector = {
+        ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
+    }
+    with _autoscaling_requester_lock:
+        return AutoscalingRequester.options(
+            name="AutoscalingRequester",
+            namespace="AutoscalingRequester",
+            get_if_exists=True,
+            lifetime="detached",
+            label_selector=label_selector,
+        ).remote()

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

@@ -0,0 +1,43 @@
+from typing import TYPE_CHECKING, List
+
+from .backpressure_policy import BackpressurePolicy
+from .concurrency_cap_backpressure_policy import ConcurrencyCapBackpressurePolicy
+from .downstream_capacity_backpressure_policy import (
+    DownstreamCapacityBackpressurePolicy,
+)
+from .resource_budget_backpressure_policy import ResourceBudgetBackpressurePolicy
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+# Default enabled backpressure policies and its config key.
+# Use `DataContext.set_config` to config it.
+ENABLED_BACKPRESSURE_POLICIES = [
+    ConcurrencyCapBackpressurePolicy,
+    ResourceBudgetBackpressurePolicy,
+    DownstreamCapacityBackpressurePolicy,
+]
+ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY = "backpressure_policies.enabled"
+
+
+def get_backpressure_policies(
+    data_context: DataContext,
+    topology: "Topology",
+    resource_manager: "ResourceManager",
+) -> List[BackpressurePolicy]:
+    policies = data_context.get_config(
+        ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY, ENABLED_BACKPRESSURE_POLICIES
+    )
+
+    return [policy(data_context, topology, resource_manager) for policy in policies]
+
+
+__all__ = [
+    "BackpressurePolicy",
+    "ConcurrencyCapBackpressurePolicy",
+    "DownstreamCapacityBackpressurePolicy",
+    "ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY",
+    "get_backpressure_policies",
+]

python/ray/data/_internal/execution/backpressure_policy/backpressure_policy.py

@@ -0,0 +1,70 @@
+from abc import ABC
+from typing import TYPE_CHECKING, Optional
+
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+
+class BackpressurePolicy(ABC):
+    """Interface for back pressure policies."""
+
+    @property
+    def name(self) -> str:
+        """Human-readable name for UX/progress bar display.
+
+        Defaults to the class name. Subclasses can override for a custom name.
+        """
+        return type(self).__name__
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ):
+        """Initialize the backpressure policy.
+
+        Args:
+            data_context: The data context.
+            topology: The execution topology.
+            resource_manager: The resource manager.
+        """
+        self._data_context = data_context
+        self._topology = topology
+        self._resource_manager = resource_manager
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        """Determine if we can add a new input to the operator. If returns False, the
+        operator will be backpressured and will not be able to run new tasks.
+        Used in `streaming_executor_state.py::select_operator_to_run()`.
+
+        Returns: True if we can add a new input to the operator, False otherwise.
+
+        Note, if multiple backpressure policies are enabled, the operator will be
+        backpressured if any of the policies returns False.
+        """
+        return True
+
+    def max_task_output_bytes_to_read(self, op: "PhysicalOperator") -> Optional[int]:
+        """Return the maximum bytes of pending task outputs can be read for
+        the given operator. None means no limit.
+
+        This is used for output backpressure to limit how much data an operator
+        can read from its running tasks.
+
+        Note, if multiple backpressure policies return non-None values for an operator,
+        the minimum of those values will be used as the limit.
+
+        Args:
+            op: The operator to get the limit for.
+
+        Returns:
+            The maximum bytes that can be read, or None if no limit.
+        """
+        return None

python/ray/data/_internal/execution/backpressure_policy/concurrency_cap_backpressure_policy.py

@@ -0,0 +1,253 @@
+import logging
+import math
+import warnings
+from collections import defaultdict
+from typing import TYPE_CHECKING, Dict
+
+from .backpressure_policy import BackpressurePolicy
+from .downstream_capacity_backpressure_policy import (
+    get_available_object_store_budget_fraction,
+)
+from ray._common.utils import env_float
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.task_pool_map_operator import (
+    TaskPoolMapOperator,
+)
+from ray.util.annotations import Deprecated, RayDeprecationWarning
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.operators.map_operator import MapOperator
+
+logger = logging.getLogger(__name__)
+
+
+@Deprecated(
+    message="ConcurrencyCapBackpressurePolicy is deprecated and will be removed "
+    "on or after Ray 2.59.",
+)
+class ConcurrencyCapBackpressurePolicy(BackpressurePolicy):
+    """A backpressure policy that caps the concurrency of each operator.
+    This policy dynamically limits the number of concurrent tasks per operator
+    based on the output queue growth rate.
+
+      - Maintain asymmetric EWMA of total enqueued output bytes as the
+        typical level: `level`.
+      - Maintain asymmetric EWMA of absolute residual vs the *previous* level as a
+        scale proxy: `dev = EWMA(|q - level_prev|)`.
+      - Define deadband: Deadband is the acceptable range of the output queue size
+        around the typical level where the queue size is expected to stay stable.
+        deadband [lower, upper] = [level - K_DEV*dev, level + K_DEV*dev].
+      - If q > upper -> target cap = running - BACKOFF_FACTOR  (back off)
+        If q < lower -> target cap = running + RAMPUP_FACTOR  (ramp up)
+        Else         -> target cap = running      (hold)
+      - Apply user-configured max concurrency cap, admit iff running < target cap.
+
+    NOTE: Only support setting concurrency cap for `TaskPoolMapOperator` for now.
+    TODO(chengsu): Consolidate with actor scaling logic of `ActorPoolMapOperator`.
+    """
+
+    @property
+    def name(self) -> str:
+        return "ConcurrencyCap"
+
+    # Smoothing factor for the asymmetric EWMA (slow fall, faster rise).
+    EWMA_ALPHA = env_float("RAY_DATA_CONCURRENCY_CAP_EWMA_ALPHA", 0.1)
+    EWMA_ALPHA_UP = 1.0 - (1.0 - EWMA_ALPHA) ** 2  # fast rise
+    # Deadband width in units of the EWMA absolute deviation estimate.
+    K_DEV = env_float("RAY_DATA_CONCURRENCY_CAP_K_DEV", 1.0)
+    # Factor to back off when the queue is too large.
+    BACKOFF_FACTOR = env_float("RAY_DATA_CONCURRENCY_CAP_BACKOFF_FACTOR", 1)
+    # Factor to ramp up when the queue is too small.
+    RAMPUP_FACTOR = env_float("RAY_DATA_CONCURRENCY_CAP_RAMPUP_FACTOR", 1)
+    # Threshold for per-Op object store budget (available) vs total
+    # (available / total) ratio to enable dynamic output queue size backpressure.
+    AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD = env_float(
+        "RAY_DATA_CONCURRENCY_CAP_AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD", 0.1
+    )
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Configured per-operator caps (+inf if unset).
+        self._concurrency_caps: Dict["PhysicalOperator", float] = {}
+
+        # EWMA state for level
+        self._q_level_nbytes: Dict["PhysicalOperator", float] = defaultdict(float)
+
+        # EWMA state for dev
+        self._q_level_dev: Dict["PhysicalOperator", float] = defaultdict(float)
+
+        # Per-operator cached threshold (bootstrapped from first sample).
+        self._queue_level_thresholds: Dict["PhysicalOperator", int] = defaultdict(int)
+
+        # Last effective cap for change logs.
+        self._last_effective_caps: Dict["PhysicalOperator", int] = {}
+
+        # Initialize caps from operators (infinite if unset)
+        for op, _ in self._topology.items():
+            if (
+                isinstance(op, TaskPoolMapOperator)
+                and op.get_max_concurrency_limit() is not None
+            ):
+                self._concurrency_caps[op] = op.get_max_concurrency_limit()
+            else:
+                self._concurrency_caps[op] = float("inf")
+
+        # Whether to cap the concurrency of an operator based on its and downstream's queue size.
+        self.enable_dynamic_output_queue_size_backpressure = (
+            self._data_context.enable_dynamic_output_queue_size_backpressure
+        )
+
+        if self.enable_dynamic_output_queue_size_backpressure:
+            warnings.warn(
+                "ConcurrencyCapBackpressurePolicy is deprecated and will be "
+                "removed on or after Ray 2.59.",
+                RayDeprecationWarning,
+                stacklevel=2,
+            )
+
+        dynamic_output_queue_size_backpressure_configs = ""
+        if self.enable_dynamic_output_queue_size_backpressure:
+            dynamic_output_queue_size_backpressure_configs = (
+                f", EWMA_ALPHA={self.EWMA_ALPHA}, K_DEV={self.K_DEV}, "
+                f"BACKOFF_FACTOR={self.BACKOFF_FACTOR}, RAMPUP_FACTOR={self.RAMPUP_FACTOR}, "
+                f"AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD={self.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD}"
+            )
+        logger.debug(
+            f"ConcurrencyCapBackpressurePolicy caps: {self._concurrency_caps}, "
+            f"enabled: {self.enable_dynamic_output_queue_size_backpressure}{dynamic_output_queue_size_backpressure_configs}"
+        )
+
+    def _update_ewma_asymmetric(self, prev_value: float, sample: float) -> float:
+        """
+        Update EWMA with asymmetric behavior: fast rise, slow fall.
+        Args:
+            prev_value: Previous EWMA value
+            sample: New sample value
+
+        Returns:
+            Updated EWMA value
+        """
+        if prev_value <= 0:
+            return sample
+
+        # fast rise if sample > prev_value, slow fall otherwise
+        alpha = self.EWMA_ALPHA_UP if sample > prev_value else self.EWMA_ALPHA
+        return (1 - alpha) * prev_value + alpha * sample
+
+    def _update_level_and_dev(self, op: "PhysicalOperator", q_bytes: int) -> None:
+        """Update EWMA level and dev (residual w.r.t. previous level)."""
+        q = float(q_bytes)
+
+        level_prev = self._q_level_nbytes[op]
+        dev_prev = self._q_level_dev[op]
+
+        # Deviation vs the previous level
+        dev_sample = abs(q - level_prev) if level_prev > 0 else 0.0
+        dev = self._update_ewma_asymmetric(dev_prev, dev_sample)
+
+        # Now update the level itself
+        level = self._update_ewma_asymmetric(level_prev, q)
+
+        self._q_level_nbytes[op] = level
+        self._q_level_dev[op] = dev
+
+        # For visibility, store the integer center of the band
+        self._queue_level_thresholds[op] = max(1, int(level))
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        """Return whether `op` may accept another input now."""
+        num_tasks_running = op.metrics.num_tasks_running
+
+        # Skip dynamic backpressure if:
+        # - Not a MapOperator
+        # - Not eligible for Op for Backpressure
+        # - Dynamic backpressure based on output queue size is disabled
+        # - Downstream is a materializing op which requires full materialization
+        if (
+            not isinstance(op, MapOperator)
+            or not self._resource_manager.is_op_eligible(op)
+            or not self.enable_dynamic_output_queue_size_backpressure
+            or self._resource_manager._is_blocking_materializing_op(op)
+        ):
+            return num_tasks_running < self._concurrency_caps[op]
+
+        # For this Op, if the objectstore budget (available) to total
+        # ratio is above threshold, skip dynamic output queue size backpressure.
+        available_budget_fraction = get_available_object_store_budget_fraction(
+            self._resource_manager, op, consider_downstream_ineligible_ops=True
+        )
+        if (
+            available_budget_fraction is not None
+            and available_budget_fraction > self.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD
+        ):
+            # If the objectstore budget (available) to total
+            # ratio is above threshold, skip dynamic output queue size
+            # backpressure, but still enforce the configured cap.
+            return num_tasks_running < self._concurrency_caps[op]
+
+        # Current total queued bytes (this op + downstream)
+        current_queue_size_bytes = self._resource_manager.get_mem_op_internal(
+            op
+        ) + self._resource_manager.get_mem_op_outputs(
+            op, include_ineligible_downstream=True
+        )
+
+        # Update EWMA state (level & dev) and compute effective cap. Note that
+        # we don't update the EWMA state if the objectstore budget (available) vs total
+        # ratio is above threshold, because the level and dev adjusts quickly.
+        self._update_level_and_dev(op, current_queue_size_bytes)
+        effective_cap = self._effective_cap(
+            op, num_tasks_running, current_queue_size_bytes
+        )
+
+        last = self._last_effective_caps.get(op, None)
+        if last != effective_cap:
+            logger.debug(
+                f"Cap change {op.name}: {last if last is not None else 'None'} -> "
+                f"{effective_cap} (running={num_tasks_running}, queue={current_queue_size_bytes}, "
+                f"thr={self._queue_level_thresholds[op]})"
+            )
+            self._last_effective_caps[op] = effective_cap
+
+        return num_tasks_running < effective_cap
+
+    def _effective_cap(
+        self,
+        op: "PhysicalOperator",
+        num_tasks_running: int,
+        current_queue_size_bytes: int,
+    ) -> int:
+        """A simple controller around EWMA level.
+        Args:
+            op: The operator to compute the effective cap for.
+            num_tasks_running: The number of tasks currently running.
+            current_queue_size_bytes: Current total queued bytes for this operator + downstream.
+        Returns:
+            The effective cap.
+        """
+        cap_cfg = self._concurrency_caps[op]
+
+        level = float(self._q_level_nbytes[op])
+        dev = max(1.0, float(self._q_level_dev[op]))
+        upper = level + self.K_DEV * dev
+        lower = level - self.K_DEV * dev
+
+        if current_queue_size_bytes > upper:
+            # back off
+            target = num_tasks_running - self.BACKOFF_FACTOR
+        elif current_queue_size_bytes < lower:
+            # ramp up
+            target = num_tasks_running + self.RAMPUP_FACTOR
+        else:
+            # hold
+            target = num_tasks_running
+
+        # Clamp to [1, configured_cap]
+        target = max(1, target)
+        if not math.isinf(cap_cfg):
+            target = min(target, int(cap_cfg))
+        return int(target)

python/ray/data/_internal/execution/backpressure_policy/downstream_capacity_backpressure_policy.py

@@ -0,0 +1,229 @@
+import logging
+from typing import TYPE_CHECKING, Optional
+
+from .backpressure_policy import BackpressurePolicy
+from ray._common.utils import env_float
+from ray.data._internal.execution.resource_manager import (
+    ResourceManager,
+)
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+logger = logging.getLogger(__name__)
+
+
+def get_available_object_store_budget_fraction(
+    resource_manager: "ResourceManager",
+    op: "PhysicalOperator",
+    consider_downstream_ineligible_ops: bool,
+) -> Optional[float]:
+    """Get available object store memory budget fraction for the operator.
+
+    Args:
+        resource_manager: The resource manager to use.
+        op: The operator to get the budget fraction for.
+        consider_downstream_ineligible_ops: If True, include downstream ineligible
+            ops in the calculation. If False, only consider this op's usage/budget.
+
+    Returns:
+        The available budget fraction, or None if not available.
+    """
+    op_usage = resource_manager.get_op_usage(
+        op, include_ineligible_downstream=consider_downstream_ineligible_ops
+    )
+    op_budget = resource_manager.get_budget(op)
+    if op_usage is None or op_budget is None:
+        return None
+
+    total_usage = op_usage.object_store_memory
+
+    total_budget = op_budget.object_store_memory
+    total_mem = total_usage + total_budget
+    if total_mem == 0:
+        return None
+
+    return total_budget / total_mem
+
+
+def get_utilized_object_store_budget_fraction(
+    resource_manager: "ResourceManager",
+    op: "PhysicalOperator",
+    consider_downstream_ineligible_ops: bool,
+) -> Optional[float]:
+    """Get utilized object store memory budget fraction for the operator.
+
+    Args:
+        resource_manager: The resource manager to use.
+        op: The operator to get the utilized fraction for.
+        consider_downstream_ineligible_ops: If True, include downstream ineligible
+            ops in the calculation. If False, only consider this op's usage/budget.
+
+    Returns:
+        The utilized budget fraction, or None if not available.
+    """
+    available_fraction = get_available_object_store_budget_fraction(
+        resource_manager,
+        op,
+        consider_downstream_ineligible_ops=consider_downstream_ineligible_ops,
+    )
+    if available_fraction is None:
+        return None
+    return 1 - available_fraction
+
+
+class DownstreamCapacityBackpressurePolicy(BackpressurePolicy):
+    """Backpressure policy based on downstream processing capacity.
+
+    To backpressure a given operator, use queue size build up / downstream capacity ratio.
+    This ratio represents the upper limit of buffering in object store between pipeline stages
+    to optimize for throughput.
+    """
+
+    # Threshold for per-Op object store budget utilization vs total
+    # (utilization / total) ratio to enable downstream capacity backpressure.
+    OBJECT_STORE_BUDGET_UTIL_THRESHOLD = env_float(
+        "RAY_DATA_DOWNSTREAM_CAPACITY_OBJECT_STORE_BUDGET_UTIL_THRESHOLD", 0.5
+    )
+
+    @property
+    def name(self) -> str:
+        return "DownstreamCapacity"
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ):
+        super().__init__(data_context, topology, resource_manager)
+
+        self._backpressure_capacity_ratio = (
+            self._data_context.downstream_capacity_backpressure_ratio
+        )
+        self._prev_should_backpressure: dict["PhysicalOperator", bool] = {}
+
+        if self._backpressure_capacity_ratio is not None:
+            logger.debug(
+                "DownstreamCapacityBackpressurePolicy enabled with backpressure "
+                f"capacity ratio: {self._backpressure_capacity_ratio}"
+            )
+
+    def _get_queue_size_bytes(self, op: "PhysicalOperator") -> int:
+        """Get the output current queue size
+        (this operator + ineligible downstream operators) in bytes for the given operator.
+        """
+        op_outputs_usage = self._topology[op].output_queue_bytes()
+        # Also account the downstream ineligible operators' memory usage.
+        op_outputs_usage += sum(
+            self._resource_manager.get_op_usage(next_op).object_store_memory
+            for next_op in self._resource_manager._get_downstream_ineligible_ops(op)
+        )
+        return op_outputs_usage
+
+    def _get_downstream_capacity_size_bytes(self, op: "PhysicalOperator") -> int:
+        """Get the downstream capacity size for the given operator.
+
+        Downstream capacity size is the sum of the pending task inputs of the
+        downstream eligible operators.
+
+        If an output dependency is ineligible, skip it and recurse down to find
+        eligible output dependencies. If there are no output dependencies,
+        return external consumer bytes.
+        """
+        if not op.output_dependencies:
+            # No output dependencies, return external consumer bytes.
+            return self._resource_manager.get_external_consumer_bytes()
+
+        total_capacity_size_bytes = 0
+        for output_dependency in op.output_dependencies:
+            if self._resource_manager.is_op_eligible(output_dependency):
+                # Output dependency is eligible, add its pending task inputs.
+                total_capacity_size_bytes += (
+                    output_dependency.metrics.obj_store_mem_pending_task_inputs or 0
+                )
+            else:
+                # Output dependency is ineligible, recurse down to find eligible ops.
+                total_capacity_size_bytes += self._get_downstream_capacity_size_bytes(
+                    output_dependency
+                )
+        return total_capacity_size_bytes
+
+    def _should_skip_backpressure(self, op: "PhysicalOperator") -> bool:
+        """Check if backpressure should be skipped for the operator.
+        TODO(srinathk10): Extract this to common logic to skip invoking BackpressurePolicy.
+        """
+        if self._backpressure_capacity_ratio is None:
+            # Downstream capacity backpressure is disabled.
+            return True
+
+        if not self._resource_manager.is_op_eligible(op):
+            # Operator is not eligible for backpressure.
+            return True
+
+        if self._resource_manager._is_blocking_materializing_op(op):
+            # Operator is materializing, so no need to perform backpressure.
+            return True
+
+        return False
+
+    def _get_queue_ratio(self, op: "PhysicalOperator") -> float:
+        """Get queue/capacity ratio for the operator."""
+        queue_size_bytes = self._get_queue_size_bytes(op)
+        downstream_capacity_size_bytes = self._get_downstream_capacity_size_bytes(op)
+        if downstream_capacity_size_bytes == 0:
+            # No downstream capacity to backpressure against, so no backpressure.
+            return 0
+        return queue_size_bytes / downstream_capacity_size_bytes
+
+    def _should_apply_backpressure(self, op: "PhysicalOperator") -> bool:
+        """Check if backpressure should be applied for the operator.
+
+        Returns True if backpressure should be applied, False otherwise.
+        """
+        if self._should_skip_backpressure(op):
+            return False
+
+        utilized_budget_fraction = get_utilized_object_store_budget_fraction(
+            self._resource_manager, op, consider_downstream_ineligible_ops=True
+        )
+        queue_ratio = self._get_queue_ratio(op)
+        if (
+            utilized_budget_fraction is not None
+            and utilized_budget_fraction <= self.OBJECT_STORE_BUDGET_UTIL_THRESHOLD
+        ):
+            # Utilized budget fraction is below threshold, so should skip backpressure.
+            result = False
+        else:
+            # Apply backpressure if queue ratio exceeds the threshold.
+            result = queue_ratio > self._backpressure_capacity_ratio
+
+        prev = self._prev_should_backpressure.get(op)
+        if prev != result:
+            queue_size_bytes = self._get_queue_size_bytes(op)
+            downstream_capacity_bytes = self._get_downstream_capacity_size_bytes(op)
+            logger.debug(
+                f"Backpressure change {op.name}: {prev} -> {result} "
+                f"(queue_ratio={queue_ratio:.2f}, {queue_size_bytes=}, "
+                f"{downstream_capacity_bytes=}, {utilized_budget_fraction=})"
+            )
+            self._prev_should_backpressure[op] = result
+
+        return result
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        """Determine if we can add input to the operator based on
+        downstream capacity.
+        """
+        return not self._should_apply_backpressure(op)
+
+    def max_task_output_bytes_to_read(self, op: "PhysicalOperator") -> Optional[int]:
+        """Return the maximum bytes of pending task outputs can be read for
+        the given operator. None means no limit."""
+        if self._should_apply_backpressure(op):
+            return 0
+        return None

python/ray/data/_internal/execution/backpressure_policy/resource_budget_backpressure_policy.py

@@ -0,0 +1,36 @@
+import logging
+from typing import TYPE_CHECKING, Optional
+
+from .backpressure_policy import BackpressurePolicy
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class ResourceBudgetBackpressurePolicy(BackpressurePolicy):
+    """A backpressure policy based on resource budgets in ResourceManager."""
+
+    @property
+    def name(self) -> str:
+        return "ResourceBudget"
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        if self._resource_manager._op_resource_allocator is not None:
+            return self._resource_manager._op_resource_allocator.can_submit_new_task(op)
+
+        return True
+
+    def max_task_output_bytes_to_read(self, op: "PhysicalOperator") -> Optional[int]:
+        """Determine maximum bytes to read based on the resource budgets.
+
+        Args:
+            op: The operator to get the limit for.
+
+        Returns:
+            The maximum bytes that can be read, or None if no limit.
+        """
+        return self._resource_manager.max_task_output_bytes_to_read(op)

python/ray/data/_internal/execution/block_ref_counter.py

@@ -0,0 +1,85 @@
+import threading
+from collections import defaultdict
+from typing import Callable, Dict, Optional
+
+import ray
+from ray._private.worker import global_worker
+
+
+class BlockRefCounter:
+    """Tracks object-store memory usage per operator via Ray Core callbacks.
+
+    The callback fires when:
+    - All Python ObjectRefs wrapping the block's ObjectID are garbage-collected, AND
+    - All Ray tasks that received the block as an argument have completed.
+    """
+
+    def __init__(
+        self,
+        add_object_out_of_scope_callback: Optional[
+            Callable[["ray.ObjectRef", Callable[[bytes], None]], bool]
+        ] = None,
+    ):
+        if add_object_out_of_scope_callback is None:
+            add_object_out_of_scope_callback = (
+                global_worker.core_worker.add_object_out_of_scope_callback  # pyrefly: ignore[missing-attribute]
+            )
+        self._add_callback_fn = add_object_out_of_scope_callback
+        # IDs of live blocks. Stale callbacks (fired after clear()) check
+        # membership here and no-op, preventing negative _bytes_by_producer.
+        self._registered_ids: set[bytes] = set()
+        # (producer_id -> total live bytes); maintained incrementally for O(1) reads.
+        self._bytes_by_producer: Dict[str, int] = defaultdict(int)
+        self._lock = threading.Lock()
+
+    def on_block_produced(
+        self,
+        block_ref: "ray.ObjectRef",
+        size_bytes: int,
+        producer_id: str,
+    ) -> None:
+        """Register a block and attribute its memory to producer_id.
+
+        Registers a Ray Core out-of-scope callback so that when all references
+        to block_ref are gone the bytes are automatically removed from the
+        producer's usage.
+
+        Idempotent: calling twice with the same block_ref is a no-op.
+        """
+        id_binary = block_ref.binary()
+        with self._lock:
+            if id_binary in self._registered_ids:
+                return
+            self._registered_ids.add(id_binary)
+            self._bytes_by_producer[producer_id] += size_bytes
+
+        def _on_object_freed(id_bytes: bytes) -> None:
+            with self._lock:
+                if id_bytes not in self._registered_ids:
+                    # Already cleared (e.g. by clear()), nothing to do.
+                    return
+                self._registered_ids.discard(id_bytes)
+                self._bytes_by_producer[producer_id] -= size_bytes
+
+        # TODO(srayhome): This raises ValueError for blocks not owned by this
+        # worker (e.g., materialized dataset passed to streaming_split). We may
+        # need to guard this with RefBundle.owns_blocks or skip registration at
+        # the InputDataBuffer level.
+        registered = self._add_callback_fn(block_ref, _on_object_freed)
+        if not registered:
+            _on_object_freed(id_binary)
+
+    def get_object_store_memory_usage(self, producer_id: str) -> int:
+        """Total bytes of live blocks attributed to producer_id."""
+        with self._lock:
+            return self._bytes_by_producer.get(producer_id, 0)
+
+    def clear(self) -> None:
+        """Reset all accounting, e.g. on executor shutdown.
+
+        Any previously registered Ray Core callbacks firing after clear()
+        will be silently ignored because _registered_ids is empty.
+        """
+        with self._lock:
+            self._registered_ids.clear()
+            self._bytes_by_producer.clear()

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

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from .base import (
+    BaseBundleQueue,
+    QueueWithRemoval,
+)
+from .bundler import EstimateSize, ExactMultipleSize, RebundleQueue
+from .fifo import FIFOBundleQueue
+from .hash_link import HashLinkedQueue
+from .reordering import ReorderingBundleQueue
+from .thread_safe import ThreadSafeBundleQueue
+
+
+def create_bundle_queue() -> QueueWithRemoval:
+    return HashLinkedQueue()
+
+
+__all__ = [
+    "BaseBundleQueue",
+    "create_bundle_queue",
+    "HashLinkedQueue",
+    "RebundleQueue",
+    "EstimateSize",
+    "ReorderingBundleQueue",
+    "FIFOBundleQueue",
+    "ExactMultipleSize",
+    "QueueWithRemoval",
+    "ThreadSafeBundleQueue",
+]

python/ray/data/_internal/execution/bundle_queue/base.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import abc
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Optional,
+)
+
+from ray.data._internal.execution.util import memory_string
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+class BundleQueue(abc.ABC):
+    @abc.abstractmethod
+    def estimate_size_bytes(self) -> int:
+        """Returns the estimated size in bytes of all bundles."""
+        ...
+
+    @abc.abstractmethod
+    def num_blocks(self) -> int:
+        """Returns the total # of blocks across all bundles."""
+        ...
+
+    @abc.abstractmethod
+    def num_bundles(self) -> int:
+        """Returns the total # of bundles."""
+        ...
+
+    @abc.abstractmethod
+    def num_rows(self) -> int:
+        """Return the total # of rows across all bundles."""
+        ...
+
+    @abc.abstractmethod
+    def _on_enqueue_bundle(self, bundle: RefBundle):
+        """Hook called before a bundle is added to the queue."""
+        ...
+
+    @abc.abstractmethod
+    def _on_dequeue_bundle(self, bundle: RefBundle):
+        """Hook called after a bundle is removed from the queue."""
+        ...
+
+    @abc.abstractmethod
+    def _add_inner(self, bundle: RefBundle, **kwargs: Any):
+        """Add a bundle to the internal data structure."""
+        ...
+
+    @abc.abstractmethod
+    def _get_next_inner(self) -> RefBundle:
+        """Remove and return the next bundle from the internal data structure."""
+        ...
+
+    def add(self, bundle: RefBundle, **kwargs: Any):
+        """Add a bundle to the tail(end) of the queue. Base classes should override
+        the `_add_inner` method for simple use cases. For more complex metrics tracking,
+        they can override this method.
+
+        Args:
+            bundle: The bundle to add.
+            **kwargs: Additional queue-specific parameters (e.g., `key` for ordered queues).
+                This is used for `finalize`.
+        """
+        self._on_enqueue_bundle(bundle)
+        self._add_inner(bundle, **kwargs)
+
+    def get_next(self) -> RefBundle:
+        """Remove and return the head of the queue. Base classes should override
+        the `_get_next_inner` method for simple use cases. For more complex metrics tracking,
+        they can override this method.
+
+        Raises:
+            IndexError: If the queue is empty.
+
+        Returns:
+            The `RefBundle` at the head of the queue.
+        """
+        bundle = self._get_next_inner()
+        self._on_dequeue_bundle(bundle)
+        return bundle
+
+    @abc.abstractmethod
+    def peek_next(self) -> Optional[RefBundle]:
+        """Return the head of the queue. The only invariant is
+        that the # of blocks, rows, and bytes must remain unchanged
+        before and after this method call.
+
+        If queue.has_next() == False, return `None`.
+        """
+        ...
+
+    @abc.abstractmethod
+    def has_next(self) -> bool:
+        """Check if the queue has a valid bundle."""
+        ...
+
+    @abc.abstractmethod
+    def clear(self):
+        """Remove all bundles from the queue."""
+        ...
+
+    @abc.abstractmethod
+    def finalize(self, **kwargs: Any):
+        """Signal that no additional bundles will be added to the bundler so
+        the bundler can be finalized. The keys of kwargs provided should be the same
+        as the ones passed into the `add()` method. This is important for ordered
+        queues."""
+        ...
+
+    def __len__(self) -> int:
+        """Return the total # bundles."""
+        return self.num_bundles()
+
+    def __repr__(self) -> str:
+        """Return a string representation showing queue metrics."""
+        nbytes = memory_string(self.estimate_size_bytes())
+        return (
+            f"{self.__class__.__name__}("
+            f"num_bundles={len(self)}, "
+            f"num_blocks={self.num_blocks()}, "
+            f"num_rows={self.num_rows()}, "
+            f"nbytes={nbytes})"
+        )
+
+
+class BaseBundleQueue(BundleQueue):
+    """Base class for storing bundles. Here and subclasses should adhere to the mental
+    model that "first", "front", or "head" is the next bundle to be dequeued. Consequently,
+    "last", "back", or "tail" is the last bundle to be dequeued.
+
+    Subclasses may choose to use the _on_dequeue_bundle and _on_enqueue_bundle methods to
+    track num_blocks, nbytes, etc... If not, they should override those methods.
+    """
+
+    def __init__(self):
+        self._nbytes: int = 0
+        self._num_blocks: int = 0
+        self._num_bundles: int = 0
+        self._num_rows: int = 0
+
+    def _on_enqueue_bundle(self, bundle: RefBundle):
+        self._nbytes += bundle.size_bytes()
+        self._num_blocks += len(bundle.block_refs)
+        self._num_bundles += 1
+        self._num_rows += bundle.num_rows() or 0
+
+    def _on_dequeue_bundle(self, bundle: RefBundle):
+        self._nbytes -= bundle.size_bytes()
+        self._num_blocks -= len(bundle.block_refs)
+        self._num_bundles -= 1
+        self._num_rows -= bundle.num_rows() or 0
+
+    def estimate_size_bytes(self) -> int:
+        """Return the estimated size in bytes of all bundles."""
+        return self._nbytes
+
+    def num_blocks(self) -> int:
+        """Return the total # of blocks across all bundles."""
+        return self._num_blocks
+
+    def num_bundles(self) -> int:
+        return self._num_bundles
+
+    def num_rows(self) -> int:
+        """Return the total # of rows across all bundles."""
+        return self._num_rows
+
+    def _reset_metrics(self):
+        self._num_rows = 0
+        self._num_blocks = 0
+        self._num_bundles = 0
+        self._nbytes = 0
+
+    def add(self, bundle: RefBundle, **kwargs: Any):
+        """Add a bundle to the tail(end) of the queue. Base classes should override
+        the `_add_inner` method for simple use cases. For more complex metrics tracking,
+        they can override this method.
+
+        Args:
+            bundle: The bundle to add.
+            **kwargs: Additional queue-specific parameters (e.g., `key` for ordered queues).
+                This is used for `finalize`.
+        """
+        self._on_enqueue_bundle(bundle)
+        self._add_inner(bundle, **kwargs)
+
+    def _add_inner(self, bundle: RefBundle, **kwargs: Any) -> None:
+        raise NotImplementedError
+
+    def get_next(self) -> RefBundle:
+        """Remove and return the head of the queue. Base classes should override
+        the `_get_next_inner` method for simple use cases. For more complex metrics tracking,
+        they can override this method.
+
+        Raises:
+            IndexError: If the queue is empty.
+
+        Returns:
+            The `RefBundle` at the head of the queue.
+        """
+        bundle = self._get_next_inner()
+        self._on_dequeue_bundle(bundle)
+        return bundle
+
+    def _get_next_inner(self) -> RefBundle:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def peek_next(self) -> Optional[RefBundle]:
+        """Return the head of the queue. The only invariant is
+        that the # of blocks, rows, and bytes must remain unchanged
+        before and after this method call.
+
+        If queue.has_next() == False, return `None`.
+        """
+        ...
+
+    @abc.abstractmethod
+    def has_next(self) -> bool:
+        """Check if the queue has a valid bundle."""
+        ...
+
+    @abc.abstractmethod
+    def clear(self):
+        """Remove all bundles from the queue."""
+        ...
+
+    def finalize(self, **kwargs: Any):
+        """Signal that no additional bundles will be added to the bundler so
+        the bundler can be finalized. The keys of kwargs provided should be the same
+        as the ones passed into the `add()` method. This is important for ordered
+        queues."""
+        return None
+
+
+class QueueWithRemoval(BaseBundleQueue):
+    """Base class for storing bundles AND supporting remove(bundle)
+    and contains(bundle) operations."""
+
+    def __contains__(self, bundle: RefBundle) -> bool:
+        """Return whether the key is in the queue."""
+        ...
+
+    def remove(self, bundle: RefBundle) -> RefBundle:
+        """Remove the specified bundle from the queue. If multiple instances exist, remove the first one."""
+        bundle = self._remove_inner(bundle)
+        self._on_dequeue_bundle(bundle)
+        return bundle
+
+    def _remove_inner(self, bundle: RefBundle) -> RefBundle:
+        raise NotImplementedError

python/ray/data/_internal/execution/bundle_queue/bundler.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import abc
+from collections import deque
+from typing import TYPE_CHECKING, Any, Deque, List, Optional, Tuple
+
+from typing_extensions import override
+
+from ray.data._internal.execution.bundle_queue import (
+    BaseBundleQueue,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+class RebundlingStrategy(abc.ABC):
+    """Base class for strategies describing how to rebundle queues."""
+
+    @abc.abstractmethod
+    def can_build_ready_bundle(self, num_pending_rows: int) -> bool:
+        """Signifies whether we can build a ready bundle. A ready bundle is a bundle
+        that will be returned from `get_next()` calls. Pending bundles merge into Ready bundles."""
+        ...
+
+    @abc.abstractmethod
+    def rows_needed_from_last_pending_bundle(
+        self, total_pending_rows: int, last_pending_bundle: RefBundle
+    ) -> int:
+        """Used to determine how to rebundle and slice an existing bundle.
+
+        Args:
+            total_pending_rows: The number of rows in a batch of pending bundles that will be merged to form
+                a ready bundle, including the last_pending_bundle.
+            last_pending_bundle: The last pending bundles in that batch ^. The term *last* means the bundle that caused
+                `can_build_ready_bundle(num_pending_rows)` to be `True` for the first time.
+
+        Returns:
+            The # of rows needed from the last pending bundle. This should be > 0, unless bundle.num_rows() is None.
+        """
+        ...
+
+
+class EstimateSize(RebundlingStrategy):
+    """Rebundles RefBundles to get them close to a particular number of rows."""
+
+    def __init__(self, min_rows_per_bundle: Optional[int]):
+        """Creates a strategy for combining bundles close to a particular row count.
+
+        Args:
+            min_rows_per_bundle: The target number of rows per bundle. Note that we
+                bundle up to this target, but only exceed it if not doing so would
+                result in an empty bundle. If None, this behaves like a normal queue.
+        """
+
+        assert (
+            min_rows_per_bundle is None or min_rows_per_bundle >= 0
+        ), "Min rows per bundle has to be non-negative"
+
+        self._min_rows_per_bundle: Optional[int] = min_rows_per_bundle
+
+    @override
+    def can_build_ready_bundle(self, num_pending_rows: int) -> bool:
+        return num_pending_rows > 0 and (
+            self._min_rows_per_bundle is None
+            or num_pending_rows >= self._min_rows_per_bundle
+        )
+
+    @override
+    def rows_needed_from_last_pending_bundle(
+        self, total_pending_rows: int, last_pending_bundle: RefBundle
+    ) -> int:
+        """Returns all the rows in the pending bundle, since we only care about an estimate"""
+        return last_pending_bundle.num_rows() or 0
+
+
+class ExactMultipleSize(RebundlingStrategy):
+    def __init__(self, target_num_rows_per_block: int):
+        assert (
+            target_num_rows_per_block > 0
+        ), "target_num_rows_per_block must be positive for streaming repartition."
+        self._target_num_rows = target_num_rows_per_block
+
+    @override
+    def can_build_ready_bundle(self, num_pending_rows: int) -> bool:
+        return num_pending_rows >= self._target_num_rows
+
+    @override
+    def rows_needed_from_last_pending_bundle(
+        self, total_pending_rows: int, last_pending_bundle: RefBundle
+    ) -> int:
+        """Returns an exact MULTIPLE of target_num_rows from the last pending bundle."""
+        pending_rows = last_pending_bundle.num_rows() or 0
+        assert total_pending_rows - pending_rows < self._target_num_rows, (
+            f"Total pending rows={total_pending_rows} should be less than target_num_rows={self._target_num_rows}, "
+            "because last_pending_bundle should trigger building ready bundles"
+        )
+        extra_rows = total_pending_rows % self._target_num_rows
+        assert extra_rows < pending_rows
+        return pending_rows - extra_rows
+
+
+"""**For `ExactMultipleSize` strategy ONLY**
+
+    Streaming repartition builds fixed-size outputs from a stream of inputs.
+
+    We construct batches here to produce exactly sized outputs from arbitrary [start, end) slices across input blocks.
+    The task builder submits a map task only after the total number of rows accumulated across pending blocks reaches
+    target num rows (except during the final flush, which may emit a smaller tail block). This allows us to create
+    target-sized batches without materializing entire large blocks on the driver.
+
+    Detailed Implementation:
+    1. When a new bundle arrives, buffer it in the pending list.
+    2. Whenever the total number of rows in the pending bundles reaches the target row count, try to build a ready bundle.
+    3. Determine the slice needed from the final bundle so the ready bundle holds an exact multiple of the target rows,
+       and add the remaining bundle to the pending bundles for the next iteration.
+    4. Submit that ready bundle to a remote map task; the task slices each block according to the slice metadata stored
+       in the RefBundle (the bundle now contains n * target rows for n ≥ 1).
+    5. We configured the `OutputBlockSizeOption.target_num_rows_per_block` to the target number of rows per block in
+       plan_streaming_repartition_op so the output buffer further splits the n * target rows into n blocks of exactly
+       the target size.
+    6. Once upstream input is exhausted, flush any leftover pending bundles and repeat steps 1-5 for the tail.
+    7. The resulting blocks have lengths `[target, …, target, (total_rows % target)]`; ordering isn't guaranteed, but the
+       remainder block should appear near the end.
+
+"""
+
+
+class RebundleQueue(BaseBundleQueue):
+    """Incrementally builds task inputs to produce multiples of target-sized outputs."""
+
+    def __init__(self, strategy: RebundlingStrategy):
+        super().__init__()
+
+        self._strategy = strategy
+        self._pending_bundles: Deque[RefBundle] = deque()
+        self._ready_bundles: Deque[RefBundle] = deque()
+
+        self._curr_consumed_bundles: List[RefBundle] = []
+        # The original bundles that formed a ready bundle
+        self._consumed_bundles_list: Deque[List[RefBundle]] = deque()
+        self._total_pending_rows: int = 0
+
+    def _merge_bundles(self):
+        """Combine *ALL* pending_bundles into a single, ready bundle."""
+        from ray.data._internal.execution.interfaces import RefBundle
+
+        merged_bundle = RefBundle.merge_ref_bundles(self._pending_bundles)
+        # Update the metrics
+        self._ready_bundles.append(merged_bundle)
+        self._on_enqueue_bundle(merged_bundle)
+
+        # Clear the pending queue since all bundles have been processed
+        for bundle in self._pending_bundles:
+            self._on_dequeue_bundle(bundle)
+        self._pending_bundles.clear()
+        self._total_pending_rows = 0
+
+    def _try_build_ready_bundle(self, flush_remaining: bool) -> int:
+        """Attempts to build a ready bundle from a list of pending bundles by:
+
+        - Checking the threshold to build a ready bundle defined by `RebundlingStrategy`
+        - Appropiately keeping track of queue metrics
+
+        Returns `True` if ready bundle built, otherwise `False`
+        """
+
+        ready_bundles_built: int = 0
+        if self._pending_bundles and self._strategy.can_build_ready_bundle(
+            self._total_pending_rows
+        ):
+            last_pending_bundle = self._pending_bundles.pop()
+
+            # We now know `pending_bundle` is the bundle that enabled us to
+            # build a ready bundle. Therefore, we may need to slice the bundle.
+            rows_needed = self._strategy.rows_needed_from_last_pending_bundle(
+                total_pending_rows=self._total_pending_rows,
+                last_pending_bundle=last_pending_bundle,
+            )
+            assert rows_needed > 0, (
+                "A refbundle has zero row-count but triggered building a ready bundle"
+                "This is a bug in the Ray Data code."
+            )
+            remaining_bundle: Optional[RefBundle] = None
+            last_num_rows = last_pending_bundle.num_rows() or 0
+            if rows_needed < last_num_rows:
+                sliced_bundle, remaining_bundle = last_pending_bundle.slice(rows_needed)
+                # The original bundle was enqueued in add(). We need to dequeue it
+                # and enqueue the sliced portion, since _merge_bundles will dequeue
+                # sliced_bundle (which has different metrics than the original).
+                self._on_dequeue_bundle(last_pending_bundle)
+                self._on_enqueue_bundle(sliced_bundle)
+                self._pending_bundles.append(sliced_bundle)
+            else:
+                assert rows_needed == last_num_rows
+                self._pending_bundles.append(last_pending_bundle)
+
+            self._merge_bundles()
+            ready_bundles_built += 1
+
+            if remaining_bundle is not None:
+                # Add back remaining sliced bundle that was not included to build
+                # a ready bundle.
+                self._pending_bundles.appendleft(remaining_bundle)
+                self._total_pending_rows += remaining_bundle.num_rows() or 0
+                self._on_enqueue_bundle(remaining_bundle)
+
+        # If we're flushing and have leftover bundles, convert them to a ready bundle.
+        # Note: add() eagerly calls _try_build_ready_bundle after every insertion, so
+        # pending rows are always below the threshold when finalize() is called. This
+        # means at most one ready bundle is built per call (only the flush path fires).
+        if flush_remaining and self._pending_bundles:
+            self._merge_bundles()
+            ready_bundles_built += 1
+
+        return ready_bundles_built
+
+    @override
+    def add(self, bundle: RefBundle, **kwargs: Any):
+        from ray.data._internal.execution.interfaces import RefBundle
+
+        num_rows = bundle.num_rows() or 0
+        if num_rows == 0:
+            if self._pending_bundles:
+                last = self._pending_bundles.pop()
+                self._on_dequeue_bundle(last)
+                merged = RefBundle.merge_ref_bundles([last, bundle])
+                self._pending_bundles.append(merged)
+                self._on_enqueue_bundle(merged)
+            else:
+                self._pending_bundles.append(bundle)
+                self._on_enqueue_bundle(bundle)
+            return
+        self._total_pending_rows += num_rows
+        self._pending_bundles.append(bundle)
+        self._on_enqueue_bundle(bundle)
+        self._curr_consumed_bundles.append(bundle)
+        ready_bundles_built = self._try_build_ready_bundle(flush_remaining=False)
+        if ready_bundles_built > 0:
+            assert ready_bundles_built == 1
+            self._consumed_bundles_list.append(self._curr_consumed_bundles)
+            self._curr_consumed_bundles = []
+
+    @override
+    def has_next(self) -> bool:
+        return len(self._ready_bundles) > 0
+
+    @override
+    def _get_next_inner(self) -> RefBundle:
+        if not self.has_next():
+            raise ValueError("You can't pop from empty queue")
+        ready_bundle = self._ready_bundles.popleft()
+        # discard the original bundle
+        self._consumed_bundles_list.popleft()
+        return ready_bundle
+
+    def get_next_with_original(self) -> Tuple[RefBundle, List[RefBundle]]:
+        if not self.has_next():
+            raise ValueError("You can't pop from empty queue")
+        ready_bundle = self._ready_bundles.popleft()
+        self._on_dequeue_bundle(ready_bundle)
+        consumed_bundle = self._consumed_bundles_list.popleft()
+        return ready_bundle, consumed_bundle
+
+    @override
+    def peek_next(self) -> Optional[RefBundle]:
+        if not self.has_next():
+            return None
+        return self._ready_bundles[0]
+
+    @override
+    def finalize(self, **kwargs: Any):
+        if len(self._pending_bundles) > 0:
+            ready_bundles_built = self._try_build_ready_bundle(flush_remaining=True)
+            assert ready_bundles_built == 1
+            self._consumed_bundles_list.append(self._curr_consumed_bundles)
+            self._curr_consumed_bundles = []
+
+    @override
+    def clear(self):
+        self._reset_metrics()
+        self._pending_bundles.clear()
+        self._ready_bundles.clear()
+        self._curr_consumed_bundles.clear()
+        self._consumed_bundles_list.clear()
+        self._total_pending_rows = 0

python/ray/data/_internal/execution/bundle_queue/fifo.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from collections import deque
+from typing import TYPE_CHECKING, Any, Deque, Iterator, List, Optional
+
+from typing_extensions import override
+
+from .base import BaseBundleQueue
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+class FIFOBundleQueue(BaseBundleQueue):
+    """A bundle queue that follows fifo-policy. Conceptually
+    [ ] <- [ ] <- [ ] ...
+     ^ where the leftmost is popped first
+
+    NOTE: Not thread-safe
+    """
+
+    def __init__(self, bundles: Optional[List[RefBundle]] = None):
+        super().__init__()
+        self._inner: Deque[RefBundle] = deque([])
+        if bundles is not None:
+            for bundle in bundles:
+                self.add(bundle)
+
+    @override
+    def _add_inner(self, bundle: RefBundle, **kwargs: Any):
+        self._inner.append(bundle)
+
+    @override
+    def _get_next_inner(self) -> RefBundle:
+        if not self.has_next():
+            raise ValueError(
+                f"Popping from empty {self.__class__.__name__} is prohibited"
+            )
+
+        bundle = self._inner.popleft()
+        return bundle
+
+    @override
+    def peek_next(self) -> Optional[RefBundle]:
+        if not self.has_next():
+            return None
+        return self._inner[0]
+
+    @override
+    def has_next(self) -> bool:
+        return len(self) > 0
+
+    @override
+    def finalize(self, **kwargs: Any):
+        pass
+
+    def __iter__(self) -> Iterator[RefBundle]:
+        yield from self._inner
+
+    def to_list(self) -> List[RefBundle]:
+        return list(self._inner)
+
+    @override
+    def clear(self):
+        self._reset_metrics()
+        self._inner.clear()

python/ray/data/_internal/execution/bundle_queue/hash_link.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Deque, Dict, Iterator, Optional
+
+from typing_extensions import override
+
+from .base import QueueWithRemoval
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+@dataclass
+class _Node:
+    value: RefBundle
+    next: Optional[_Node] = None
+    prev: Optional[_Node] = None
+
+
+class HashLinkedQueue(QueueWithRemoval):
+    """A bundle queue that supports these operations quickly:
+    - contains(bundle)
+    - remove(bundle)
+
+    NOTE: Not thread-safe
+    """
+
+    def __init__(self):
+        super().__init__()
+        # We manually implement a linked list because we need to remove elements
+        # efficiently, and Python's built-in data structures have O(n) removal time.
+        self._head: Optional[_Node] = None
+        self._tail: Optional[_Node] = None
+        # We use a dictionary to keep track of the nodes corresponding to each bundle.
+        # This allows us to remove a bundle from the queue in O(1) time. We need a list
+        # because a bundle can be added to the queue multiple times. Nodes in each list
+        # are insertion-ordered.
+        self._bundle_to_nodes: Dict[RefBundle, Deque[_Node]] = defaultdict(deque)
+
+    @override
+    def __contains__(self, bundle: RefBundle) -> bool:
+        return bundle in self._bundle_to_nodes
+
+    @override
+    def _add_inner(self, bundle: RefBundle, **kwargs: Any):
+        new_node = _Node(value=bundle, next=None, prev=self._tail)
+        # Case 1: The queue is empty.
+        if self._head is None:
+            assert self._tail is None
+            self._head = new_node
+            self._tail = new_node
+        # Case 2: The queue has at least one element.
+        else:
+            self._tail.next = new_node
+            self._tail = new_node
+
+        self._bundle_to_nodes[bundle].append(new_node)
+
+    @override
+    def _get_next_inner(self) -> RefBundle:
+        # Case 1: The queue is empty.
+        if not self._head:
+            raise IndexError("You can't pop from an empty queue")
+
+        bundle = self._head.value
+        self._remove_inner(bundle)
+
+        return bundle
+
+    @override
+    def has_next(self) -> bool:
+        return self._num_bundles > 0
+
+    @override
+    def peek_next(self) -> Optional[RefBundle]:
+        if self._head is None:
+            return None
+
+        return self._head.value
+
+    @override
+    def _remove_inner(self, bundle: RefBundle) -> RefBundle:
+        # Case 1: The queue is empty.
+        if bundle not in self._bundle_to_nodes:
+            raise ValueError(f"The bundle {bundle} is not in the queue.")
+
+        node = self._bundle_to_nodes[bundle].popleft()
+        if not self._bundle_to_nodes[bundle]:
+            del self._bundle_to_nodes[bundle]
+
+        node = self._remove_node(node)
+        return node.value
+
+    def _remove_node(self, node: _Node) -> _Node:
+        # Case 2: The bundle is the only element in the queue.
+        if self._head is self._tail:
+            self._head = None
+            self._tail = None
+        # Case 3: The bundle is the first element in the queue.
+        elif node is self._head:
+            self._head = node.next
+            self._head.prev = None
+        # Case 4: The bundle is the last element in the queue.
+        elif node is self._tail:
+            self._tail = node.prev
+            self._tail.next = None
+        # Case 5: The bundle is in the middle of the queue.
+        else:
+            node.prev.next = node.next
+            node.next.prev = node.prev
+
+        return node
+
+    def __iter__(self) -> Iterator[RefBundle]:
+        curr = self._head
+        while curr:
+            yield curr.value
+            curr = curr.next
+
+    def clear(self):
+        self._reset_metrics()
+        self._bundle_to_nodes.clear()
+        self._head = None
+        self._tail = None

python/ray/data/_internal/execution/bundle_queue/reordering.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from typing import TYPE_CHECKING, DefaultDict, Deque, Optional, Set
+
+from typing_extensions import override
+
+from ray.data._internal.execution.bundle_queue import BaseBundleQueue
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+class ReorderingBundleQueue(BaseBundleQueue):
+    """A queue that iterates over the bundles in the order of provided "keys" rather than
+    insertion order (for bundles inserted with the same key, insertion order is used)
+
+    User of this queue has to adhere to following invariants of this queue:
+
+    1. (!) Used keys have to be a *contiguous* range of `[0, N]`
+
+    Failure to follow this requirement might result in this queue getting
+    irreversibly stuck.
+
+    NOTE: Not thread-safe
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self._inner: DefaultDict[int, Deque[RefBundle]] = defaultdict(lambda: deque())
+        self._current_key: int = 0
+        self._finalized_keys: Set[int] = set()
+
+    def _move_to_next_key(self):
+        """Move the output index to the next task.
+
+        This method should only be called when the current task is complete and all
+        outputs have been taken.
+        """
+        assert len(self._inner[self._current_key]) == 0
+        assert self._current_key in self._finalized_keys
+
+        self._current_key += 1
+
+    @override
+    def _add_inner(self, bundle: RefBundle, key: int) -> None:
+        assert key is not None
+        self._inner[key].append(bundle)
+
+    @override
+    def has_next(self) -> bool:
+        while (
+            self._current_key in self._finalized_keys
+            and len(self._inner[self._current_key]) == 0
+        ):
+            self._move_to_next_key()
+
+        return len(self._inner[self._current_key]) > 0
+
+    @override
+    def _get_next_inner(self) -> RefBundle:
+        # It's vital to invoke `has_next` here, to potentially advance the pointer
+        # to the next key
+        if not self.has_next():
+            raise ValueError("Cannot pop from empty queue.")
+
+        return self._inner[self._current_key].popleft()
+
+    @override
+    def peek_next(self) -> Optional[RefBundle]:
+        # It's vital to invoke `has_next` here, to potentially advance the pointer
+        # to the next key
+        if not self.has_next():
+            return None
+
+        return self._inner[self._current_key][0]
+
+    @override
+    def finalize(self, key: int):
+        assert key is not None and key >= self._current_key
+        self._finalized_keys.add(key)
+
+    @override
+    def clear(self):
+        self._reset_metrics()
+        self._inner.clear()
+        self._finalized_keys.clear()
+        self._current_key = 0

python/ray/data/_internal/execution/bundle_queue/thread_safe.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING, Any, Optional
+
+from .base import BundleQueue
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+class ThreadSafeBundleQueue(BundleQueue):
+    """A thread-safe wrapper for a ``BundleQueue``.
+
+    Delegates all operations to the wrapped queue while a lock.
+
+    NOTE: Not safe to use with ``__contains__``/``remove`` from ``QueueWithRemoval``.
+    """
+
+    def __init__(self, inner: BundleQueue):
+        self._inner = inner
+        self._lock = threading.Lock()
+
+    def estimate_size_bytes(self) -> int:
+        with self._lock:
+            return self._inner.estimate_size_bytes()
+
+    def num_blocks(self) -> int:
+        with self._lock:
+            return self._inner.num_blocks()
+
+    def num_bundles(self) -> int:
+        with self._lock:
+            return self._inner.num_bundles()
+
+    def num_rows(self) -> int:
+        with self._lock:
+            return self._inner.num_rows()
+
+    def _on_enqueue_bundle(self, bundle: RefBundle):
+        raise NotImplementedError("Use add() for thread-safe access")
+
+    def _on_dequeue_bundle(self, bundle: RefBundle):
+        raise NotImplementedError("Use get_next() for thread-safe access")
+
+    def _add_inner(self, bundle: RefBundle, **kwargs: Any):
+        raise NotImplementedError("Use add() for thread-safe access")
+
+    def _get_next_inner(self) -> RefBundle:
+        raise NotImplementedError("Use get_next() for thread-safe access")
+
+    def add(self, bundle: RefBundle, **kwargs: Any):
+        with self._lock:
+            self._inner.add(bundle, **kwargs)
+
+    def get_next(self) -> RefBundle:
+        with self._lock:
+            return self._inner.get_next()
+
+    def peek_next(self) -> Optional[RefBundle]:
+        with self._lock:
+            return self._inner.peek_next()
+
+    def has_next(self) -> bool:
+        with self._lock:
+            return self._inner.has_next()
+
+    def clear(self):
+        with self._lock:
+            self._inner.clear()
+
+    def finalize(self, **kwargs: Any):
+        with self._lock:
+            return self._inner.finalize(**kwargs)

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

@@ -0,0 +1,11 @@
+from ray.data._internal.execution.callbacks.insert_issue_detectors import (
+    IssueDetectionExecutionCallback,
+)
+from ray.data._internal.execution.callbacks.resource_allocator_prometheus_callback import (
+    ResourceAllocatorPrometheusCallback,
+)
+
+__all__ = [
+    "IssueDetectionExecutionCallback",
+    "ResourceAllocatorPrometheusCallback",
+]

python/ray/data/_internal/execution/callbacks/execution_idx_update_callback.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+from ray.data._internal.execution.execution_callback import (
+    ExecutionCallback,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+
+
+class ExecutionIdxUpdateCallback(ExecutionCallback):
+    def after_execution_succeeds(self, executor: "StreamingExecutor"):
+        dataset_context = executor._data_context
+        dataset_context._execution_idx += 1

python/ray/data/_internal/execution/callbacks/insert_issue_detectors.py

@@ -0,0 +1,23 @@
+from typing import TYPE_CHECKING
+
+from ray.data._internal.execution.execution_callback import (
+    ExecutionCallback,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+from ray.data._internal.issue_detection.issue_detector_manager import (
+    IssueDetectorManager,
+)
+
+
+class IssueDetectionExecutionCallback(ExecutionCallback):
+    """ExecutionCallback that handles issue detection."""
+
+    def before_execution_starts(self, executor: "StreamingExecutor"):
+        # Initialize issue detector in StreamingExecutor
+        executor._issue_detector_manager = IssueDetectorManager(executor)
+
+    def on_execution_step(self, executor: "StreamingExecutor"):
+        # Invoke all issue detectors
+        executor._issue_detector_manager.invoke_detectors()

python/ray/data/_internal/execution/callbacks/resource_allocator_prometheus_callback.py

@@ -0,0 +1,115 @@
+import math
+from typing import TYPE_CHECKING, Dict
+
+from ray.data._internal.execution.execution_callback import ExecutionCallback
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.resource_manager import ResourceManager
+from ray.util.metrics import Gauge
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+
+
+class ResourceAllocatorPrometheusCallback(ExecutionCallback):
+    """Updates Prometheus metrics related to resource allocation.
+
+    This callback monitors the StreamingExecutor and updates Prometheus
+    Gauges for CPU, GPU, memory, and object store memory budgets for each
+    operator at every execution step.
+    """
+
+    def __init__(self):
+        self._cpu_budget_gauge: Gauge = Gauge(
+            "data_cpu_budget",
+            "Budget (CPU) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._gpu_budget_gauge: Gauge = Gauge(
+            "data_gpu_budget",
+            "Budget (GPU) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._memory_budget_gauge: Gauge = Gauge(
+            "data_memory_budget",
+            "Budget (Memory) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._osm_budget_gauge: Gauge = Gauge(
+            "data_object_store_memory_budget",
+            "Budget (Object Store Memory) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._max_bytes_to_read_gauge: Gauge = Gauge(
+            "data_max_bytes_to_read",
+            description="Maximum bytes to read from streaming generator buffer.",
+            tag_keys=("dataset", "operator"),
+        )
+
+    def on_execution_step(self, executor: "StreamingExecutor") -> None:
+        """Called by the executor after every scheduling loop step."""
+        topology = executor._topology
+        resource_manager = executor._resource_manager
+        dataset_id = executor._dataset_id
+
+        if topology is None or resource_manager is None:
+            return
+
+        for i, op in enumerate(topology):
+            tags = {
+                "dataset": dataset_id,
+                "operator": executor._get_operator_id(op, i),
+            }
+            self._update_budget_metrics(op, tags, resource_manager)
+            self._update_max_bytes_to_read_metric(op, tags, resource_manager)
+
+    def after_execution_succeeds(self, executor: "StreamingExecutor") -> None:
+        """Updates metrics upon successful execution to ensure final states are captured."""
+        self.on_execution_step(executor)
+
+    def after_execution_fails(
+        self, executor: "StreamingExecutor", error: Exception
+    ) -> None:
+        """Updates metrics upon execution failure to ensure final states are captured."""
+        self.on_execution_step(executor)
+
+    def _update_budget_metrics(
+        self,
+        op: PhysicalOperator,
+        tags: Dict[str, str],
+        resource_manager: ResourceManager,
+    ):
+        budget = resource_manager.get_budget(op)
+        if budget is None:
+            cpu_budget = 0
+            gpu_budget = 0
+            memory_budget = 0
+            object_store_memory_budget = 0
+        else:
+            cpu_budget = -1 if math.isinf(budget.cpu) else budget.cpu
+            gpu_budget = -1 if math.isinf(budget.gpu) else budget.gpu
+            memory_budget = -1 if math.isinf(budget.memory) else budget.memory
+            object_store_memory_budget = (
+                -1
+                if math.isinf(budget.object_store_memory)
+                else budget.object_store_memory
+            )
+
+        self._cpu_budget_gauge.set(cpu_budget, tags=tags)
+        self._gpu_budget_gauge.set(gpu_budget, tags=tags)
+        self._memory_budget_gauge.set(memory_budget, tags=tags)
+        self._osm_budget_gauge.set(object_store_memory_budget, tags=tags)
+
+    def _update_max_bytes_to_read_metric(
+        self,
+        op: PhysicalOperator,
+        tags: Dict[str, str],
+        resource_manager: ResourceManager,
+    ):
+        if resource_manager.op_resource_allocator_enabled():
+            resource_allocator = resource_manager.op_resource_allocator
+            output_budget_bytes = resource_allocator.get_output_budget(op)
+            if output_budget_bytes is not None:
+                if math.isinf(output_budget_bytes):
+                    # Convert inf to -1 to represent unlimited bytes to read
+                    output_budget_bytes = -1
+                self._max_bytes_to_read_gauge.set(output_budget_bytes, tags=tags)

python/ray/data/_internal/execution/dataset_state.py

@@ -0,0 +1,22 @@
+import enum
+
+
+class DatasetState(enum.IntEnum):
+    """Enum representing the possible states of a dataset during execution."""
+
+    UNKNOWN = 0
+    RUNNING = 1
+    FINISHED = 2
+    FAILED = 3
+    PENDING = 4
+
+    def __str__(self):
+        return self.name
+
+    @classmethod
+    def from_string(cls, text):
+        """Get enum by name."""
+        try:
+            return cls[text]  # This uses the name to lookup the enum
+        except KeyError:
+            return cls.UNKNOWN

python/ray/data/_internal/execution/execution_callback.py

@@ -0,0 +1,24 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+
+
+class ExecutionCallback:
+    """Callback interface for execution events."""
+
+    def before_execution_starts(self, executor: "StreamingExecutor"):
+        """Called before the Dataset execution starts."""
+        ...
+
+    def on_execution_step(self, executor: "StreamingExecutor"):
+        """Called at each step of the Dataset execution loop."""
+        ...
+
+    def after_execution_succeeds(self, executor: "StreamingExecutor"):
+        """Called after the Dataset execution succeeds."""
+        ...
+
+    def after_execution_fails(self, executor: "StreamingExecutor", error: Exception):
+        """Called after the Dataset execution fails."""
+        ...

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

@@ -0,0 +1,22 @@
+from .common import NodeIdStr
+from .execution_options import ExecutionOptions, ExecutionResources
+from .executor import Executor, OutputIterator
+from .physical_operator import PhysicalOperator, ReportsExtraResourceUsage
+from .ref_bundle import BlockEntry, BlockSlice, RefBundle
+from .task_context import TaskContext
+from .transform_fn import AllToAllTransformFn
+
+__all__ = [
+    "AllToAllTransformFn",
+    "BlockEntry",
+    "BlockSlice",
+    "ExecutionOptions",
+    "ExecutionResources",
+    "Executor",
+    "NodeIdStr",
+    "OutputIterator",
+    "PhysicalOperator",
+    "RefBundle",
+    "ReportsExtraResourceUsage",
+    "TaskContext",
+]

python/ray/data/_internal/execution/interfaces/common.py

@@ -0,0 +1,186 @@
+import bisect
+import json
+from typing import Dict, List, Tuple
+
+from ray.data._internal.util import GiB, KiB, MiB
+from ray.util.metrics import Histogram
+
+# Node id string returned by `ray.get_runtime_context().get_node_id()`.
+NodeIdStr = str
+
+# Used for time-based histograms (e.g. task completion time, block completion time)
+histogram_buckets_s = [
+    0.1,
+    0.25,
+    0.5,
+    1.0,
+    2.5,
+    5.0,
+    7.5,
+    10.0,
+    15.0,
+    20.0,
+    25.0,
+    50.0,
+    75.0,
+    100.0,
+    150.0,
+    500.0,
+    1000.0,
+    2500.0,
+    5000.0,
+]
+
+# Used for size-based histograms (e.g. block size in bytes)
+histogram_buckets_bytes = [
+    KiB,
+    8 * KiB,
+    64 * KiB,
+    128 * KiB,
+    256 * KiB,
+    512 * KiB,
+    MiB,
+    8 * MiB,
+    64 * MiB,
+    128 * MiB,
+    256 * MiB,
+    512 * MiB,
+    GiB,
+    4 * GiB,
+    16 * GiB,
+    64 * GiB,
+    128 * GiB,
+    256 * GiB,
+    512 * GiB,
+    1024 * GiB,
+    4096 * GiB,
+]
+
+# Used for row count-based histograms (e.g. block size in rows)
+histogram_bucket_rows = [
+    1,
+    5,
+    10,
+    25,
+    50,
+    100,
+    250,
+    500,
+    1_000,
+    2_500,
+    5_000,
+    10_000,
+    25_000,
+    50_000,
+    100_000,
+    250_000,
+    500_000,
+    1_000_000,
+    2_500_000,
+    5_000_000,
+    10_000_000,
+]
+
+
+class RuntimeMetricsHistogram:
+    """
+    Class that tracks a histogram of values.
+
+    Contains helper methods to record the values and apply those values to a `ray.util.metrics.Histogram` metric.
+    """
+
+    def __init__(self, boundaries: List[float]):
+        self._boundaries = boundaries
+        # Initialize bucket counts to 0 (+1 additional bucket to represent the +Inf bucket)
+        self._bucket_counts = [0 for _ in range(len(boundaries) + 1)]
+        self._memoized_avg = None
+
+    def observe(self, value: float, num_observations: int = 1):
+        self._bucket_counts[self._find_bucket_index(value)] += num_observations
+        self._memoized_avg = None
+
+    def export_to(
+        self,
+        metric: Histogram,
+        tags: Dict[str, str],
+    ):
+        """
+        This method calculates the difference between the current bucket counts and the previous bucket counts,
+        and applies those observations to the metric.
+
+        This method stores the previous_bucket_counts in the metric as `last_applied_bucket_counts_for_tags`.
+        """
+        if getattr(metric, "last_applied_bucket_counts_for_tags", None) is None:
+            metric.last_applied_bucket_counts_for_tags = {}
+        tags_key = json.dumps(tags, sort_keys=True)
+        previous_bucket_counts = metric.last_applied_bucket_counts_for_tags.get(
+            tags_key
+        )
+
+        for i in range(len(self._bucket_counts)):
+            # Pick a value between the boundaries so the sample falls into the right bucket.
+            # We need to calculate the mid point because choosing the exact boundary value
+            # seems to have unreliable behavior on which bucket it ends up in.
+            boundary_upper_bound = (
+                self._boundaries[i]
+                if i < len(self._bucket_counts) - 1
+                # Since choosing an exact boundary value is unreliable to if it'll
+                # end up in the upper or lower bucket, we add a small buffer to the
+                # last boundary. The amount of the value doesn't matter much
+                # since it's the last bucket and should go to infinity.
+                else self._boundaries[-1] + 100
+            )
+            boundary_lower_bound = self._boundaries[i - 1] if i > 0 else 0
+            bucket_value = (boundary_upper_bound + boundary_lower_bound) / 2
+
+            # Calculate how many observations to add to the metric
+            diff = (
+                self._bucket_counts[i] - previous_bucket_counts[i]
+                if previous_bucket_counts is not None
+                else self._bucket_counts[i]
+            )
+            for _ in range(diff):
+                metric.observe(bucket_value, tags)
+
+        metric.last_applied_bucket_counts_for_tags[
+            tags_key
+        ] = self._bucket_counts.copy()
+
+    def __repr__(self):
+        if self._memoized_avg is None:
+            self._memoized_avg = self._calculate_average_value()
+        total_samples, average = self._memoized_avg
+        return f"(samples: {total_samples}, avg: {average:.2f})"
+
+    def _calculate_average_value(self) -> Tuple[int, float]:
+        """
+        Calculate the average value of all samples.
+
+        Used to show a representative value for the histogram when
+        printing the histogram as a string.
+        """
+        total_samples = sum(self._bucket_counts)
+        if total_samples == 0:
+            return total_samples, 0
+
+        weighted_sum = 0.0
+        for i, count in enumerate(self._bucket_counts):
+            if count > 0:
+                # Calculate representative value for this bucket
+                if i == 0:
+                    # First bucket: 0 to first boundary
+                    bucket_value = self._boundaries[0] / 2
+                elif i == len(self._bucket_counts) - 1:
+                    # Last bucket: last boundary to +inf
+                    bucket_value = self._boundaries[-1] * 1.5
+                else:
+                    # Middle buckets: between boundaries
+                    bucket_value = (self._boundaries[i - 1] + self._boundaries[i]) / 2
+
+                weighted_sum += bucket_value * count
+
+        average = weighted_sum / total_samples
+        return total_samples, average
+
+    def _find_bucket_index(self, value: float):
+        return bisect.bisect_left(self._boundaries, value)

python/ray/data/_internal/execution/interfaces/distribution_tracker.py

@@ -0,0 +1,178 @@
+import math
+from typing import Dict, Optional, Union
+
+try:
+    from datasketches import kll_doubles_sketch
+
+    _DATASKETCHES_AVAILABLE = True
+except ImportError:
+    _DATASKETCHES_AVAILABLE = False
+
+
+class DistributionTracker:
+    """Tracks the running mean, variance, min, max, and approximate percentiles of a
+    stream of values using Welford's algorithm for moments and a KLL sketch for
+    quantiles.
+
+    More on Welford's algorithm:
+    https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
+    """
+
+    def __init__(self):
+        self._count = 0
+        self._mean = 0.0
+        self._m2 = 0.0
+        self._min = float("inf")
+        self._max = float("-inf")
+        self._sketch = kll_doubles_sketch(200) if _DATASKETCHES_AVAILABLE else None
+
+    def add_sample(self, value: float) -> None:
+        self._count += 1
+
+        delta = value - self._mean
+        self._mean += delta / self._count
+        delta2 = value - self._mean
+        self._m2 += delta * delta2
+
+        if value < self._min:
+            self._min = value
+        if value > self._max:
+            self._max = value
+
+        if self._sketch is not None:
+            self._sketch.update(value)
+
+    def merge(self, other: "DistributionTracker") -> None:
+        """Merge another tracker into this one (associative, commutative).
+
+        Uses Chan's parallel variant of Welford's algorithm for moments.
+        See: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford:~:text=Parallel%20algorithm%5Bedit%5D
+        """
+        if other is self:
+            # Merging an accumulator into itself would double its samples
+            # (count, m2, and the sketch), so treat it as a no-op.
+            return
+        if other._count == 0:
+            return
+        if self._count == 0:
+            self._count = other._count
+            self._mean = other._mean
+            self._m2 = other._m2
+            self._min = other._min
+            self._max = other._max
+        else:
+            delta = other._mean - self._mean
+            total = self._count + other._count
+            self._m2 += other._m2 + (delta**2) * self._count * other._count / total
+            self._mean = (self._count * self._mean + other._count * other._mean) / total
+            self._count = total
+            self._min = min(self._min, other._min)
+            self._max = max(self._max, other._max)
+        if self._sketch is None or other._sketch is None:
+            # Moments above still merged; quantile detail is lost for the
+            # side(s) without a sketch.
+            self._sketch = None
+        else:
+            try:
+                self._sketch.merge(other._sketch)
+            except Exception:
+                self._sketch = None
+
+    @property
+    def num_samples(self) -> int:
+        return self._count
+
+    @property
+    def mean(self) -> float:
+        return self._mean
+
+    @property
+    def variance(self) -> float:
+        if self._count < 2:
+            return 0.0
+        return self._m2 / (self._count - 1)
+
+    @property
+    def stddev(self) -> float:
+        return math.sqrt(self.variance)
+
+    @property
+    def min(self) -> Optional[float]:
+        if self._count == 0:
+            return None
+        return self._min
+
+    @property
+    def max(self) -> Optional[float]:
+        if self._count == 0:
+            return None
+        return self._max
+
+    def _quantile(self, q: float) -> Optional[float]:
+        if self._sketch is None or self._count == 0:
+            return None
+        return self._sketch.get_quantiles([q])[0]
+
+    @property
+    def p25(self) -> Optional[float]:
+        return self._quantile(0.25)
+
+    @property
+    def p50(self) -> Optional[float]:
+        return self._quantile(0.5)
+
+    @property
+    def p75(self) -> Optional[float]:
+        return self._quantile(0.75)
+
+    @property
+    def p90(self) -> Optional[float]:
+        return self._quantile(0.9)
+
+    @property
+    def p95(self) -> Optional[float]:
+        return self._quantile(0.95)
+
+    @property
+    def p99(self) -> Optional[float]:
+        return self._quantile(0.99)
+
+    def as_dict(self) -> Dict[str, Optional[Union[int, float]]]:
+        return {
+            "num_samples": self.num_samples,
+            "mean": self.mean,
+            "variance": self.variance,
+            "min": self.min,
+            "max": self.max,
+            "p25": self.p25,
+            "p50": self.p50,
+            "p75": self.p75,
+            "p90": self.p90,
+            "p95": self.p95,
+            "p99": self.p99,
+        }
+
+    # ``kll_doubles_sketch`` is a C++-backed object that does not
+    # pickle natively. DistributionTracker rides on DatasetStats
+    # (via Timer), which is cloudpickled when Datasets cross actor /
+    # process boundaries — without these hooks any such transfer
+    # raises ``TypeError: cannot pickle 'kll_doubles_sketch' object``.
+    # The sketch exposes its own byte serialization, so we round-trip
+    # through that.
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        if self._sketch is not None:
+            state["_sketch"] = self._sketch.serialize()
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        # If the source had datasketches but this side doesn't, drop
+        # the sketch (percentiles will return None — same fallback as a
+        # default construction without datasketches installed).
+        if self._sketch is not None and not _DATASKETCHES_AVAILABLE:
+            self._sketch = None
+        elif self._sketch is not None and not isinstance(
+            self._sketch, kll_doubles_sketch
+        ):
+            self._sketch = kll_doubles_sketch.deserialize(self._sketch)

python/ray/data/_internal/execution/interfaces/execution_options.py

@@ -0,0 +1,512 @@
+import functools
+import math
+import operator
+import os
+import warnings
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
+
+from .common import NodeIdStr
+from ray.data._internal.execution.util import memory_string
+from ray.util.annotations import DeveloperAPI
+
+
+class ExecutionResources:
+    """Specifies resources usage or resource limits for execution.
+
+    By default this class represents resource usage. Use `for_limits` or
+    set `default_to_inf` to True to create an object that represents resource limits.
+    """
+
+    # ``__slots__`` keeps instances small and makes attribute access go through
+    # slot descriptors instead of a per-instance ``__dict__``. The scheduler
+    # constructs many of these per iteration (every ``add``/``subtract``/
+    # ``max``/``copy`` returns a new object), so this is a hot-path win.
+    __slots__ = ("_cpu", "_gpu", "_object_store_memory", "_memory")
+
+    def __init__(
+        self,
+        cpu: Optional[float] = None,
+        gpu: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+        memory: Optional[float] = None,
+    ):
+        """Initializes ExecutionResources.
+        Args:
+            cpu: Amount of logical CPU slots.
+            gpu: Amount of logical GPU slots.
+            object_store_memory: Amount of object store memory.
+            memory: Amount of logical memory in bytes.
+        """
+
+        # NOTE: Ray Core allocates fractional resources in up to 5th decimal
+        #       digit, hence we round the values here up to it
+        self._cpu: Optional[float] = safe_round(cpu, 5)
+        self._gpu: Optional[float] = safe_round(gpu, 5)
+        self._object_store_memory: Optional[float] = safe_round(object_store_memory, 0)
+        self._memory: Optional[float] = safe_round(memory, 0)
+
+    @classmethod
+    def from_resource_dict(
+        cls,
+        resource_dict: Dict[str, float],
+    ):
+        """Create an ExecutionResources object from a resource dict."""
+        return ExecutionResources(
+            cpu=resource_dict.get("CPU", None) or resource_dict.get("num_cpus", None),
+            gpu=resource_dict.get("GPU", None) or resource_dict.get("num_gpus", None),
+            object_store_memory=resource_dict.get("object_store_memory", None),
+            memory=resource_dict.get("memory", None),
+        )
+
+    def to_resource_dict(self) -> Dict[str, float]:
+        """Convert this ExecutionResources object to a resource dict."""
+        return {
+            "CPU": self.cpu,
+            "GPU": self.gpu,
+            "object_store_memory": self.object_store_memory,
+            "memory": self.memory,
+        }
+
+    @classmethod
+    def for_limits(
+        cls,
+        cpu: Optional[float] = None,
+        gpu: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+        memory: Optional[float] = None,
+    ) -> "ExecutionResources":
+        """Create an ExecutionResources object that represents resource limits.
+
+        Args:
+            cpu: Amount of logical CPU slots.
+            gpu: Amount of logical GPU slots.
+            object_store_memory: Amount of object store memory.
+            memory: Amount of logical memory in bytes.
+
+        Returns:
+            An ``ExecutionResources`` with the given limits (defaulting to
+            infinity for any unspecified field).
+        """
+        return ExecutionResources(
+            cpu=safe_or(cpu, float("inf")),
+            gpu=safe_or(gpu, float("inf")),
+            object_store_memory=safe_or(object_store_memory, float("inf")),
+            memory=safe_or(memory, float("inf")),
+        )
+
+    @property
+    def cpu(self) -> float:
+        return self._cpu or 0.0
+
+    @property
+    def gpu(self) -> float:
+        return self._gpu or 0.0
+
+    @property
+    def object_store_memory(self) -> float:
+        return self._object_store_memory or 0
+
+    @property
+    def memory(self) -> float:
+        return self._memory or 0
+
+    def __repr__(self):
+        return (
+            f"ExecutionResources(cpu={self.cpu}, gpu={self.gpu}, "
+            f"object_store_memory={self.object_store_memory_str()}, "
+            f"memory={self.memory_str()})"
+        )
+
+    def __eq__(self, other: "ExecutionResources") -> bool:
+        return (
+            self.cpu == other.cpu
+            and self.gpu == other.gpu
+            and self.object_store_memory == other.object_store_memory
+            and self.memory == other.memory
+        )
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.cpu,
+                self.gpu,
+                self.object_store_memory,
+                self.memory,
+            )
+        )
+
+    @classmethod
+    @functools.cache
+    def zero(cls) -> "ExecutionResources":
+        """Returns an ExecutionResources object with zero resources.
+
+        Returns a cached, shared singleton (``functools.cache`` keyed on ``cls``)
+        -- ``zero()`` is called all over the scheduler hot path (e.g.
+        ``.max(zero())``) and instances are immutable in practice (every
+        arithmetic op returns a new object and there are no setters), so sharing
+        one instance is safe and avoids the per-call allocation.
+        """
+        return ExecutionResources(0.0, 0.0, 0.0, 0.0)
+
+    @classmethod
+    @functools.cache
+    def inf(cls) -> "ExecutionResources":
+        """Returns an ExecutionResources object with infinite resources.
+
+        Returns a cached, shared singleton (see :meth:`zero` for why this is
+        safe).
+        """
+        return ExecutionResources.for_limits()
+
+    @classmethod
+    def combine(
+        cls,
+        resources: Iterable["ExecutionResources"],
+        fn: Callable[[float, float], float],
+    ) -> Optional["ExecutionResources"]:
+        """Fold an iterable of ``ExecutionResources`` per dimension with ``fn``.
+
+        ``fn(acc, value)`` combines two per-dimension floats -- e.g.
+        ``operator.add`` for a sum, or ``max``/``min`` for an element-wise
+        max/min. Accumulates raw floats in a single pass and allocates a single
+        result object, instead of one intermediate per element as
+        ``reduce(lambda a, b: a.<op>(b), resources)`` would.
+
+        Seeds with the first element (so no per-``fn`` identity is needed) and
+        returns ``None`` for an empty iterable, which may be a one-shot
+        generator (so it's consumed exactly once).
+        """
+        iterator = iter(resources)
+        first = next(iterator, None)
+        if first is None:
+            return None
+        cpu = first.cpu
+        gpu = first.gpu
+        object_store_memory = first.object_store_memory
+        memory = first.memory
+        for r in iterator:
+            cpu = fn(cpu, r.cpu)
+            gpu = fn(gpu, r.gpu)
+            object_store_memory = fn(object_store_memory, r.object_store_memory)
+            memory = fn(memory, r.memory)
+        return ExecutionResources(cpu, gpu, object_store_memory, memory)
+
+    @classmethod
+    def combine_sum(
+        cls, resources: Iterable["ExecutionResources"]
+    ) -> "ExecutionResources":
+        """Sum an iterable of ``ExecutionResources`` in a single pass.
+
+        Thin wrapper over :meth:`combine` with addition. Empty folds are common
+        (e.g. completed-ops / downstream-ineligible usage rollups on most
+        iterations), so an empty input reuses the shared ``zero()`` singleton
+        instead of allocating.
+        """
+        result = cls.combine(resources, operator.add)
+        return result if result is not None else cls.zero()
+
+    def is_zero(self) -> bool:
+        """Returns True if all resources are zero."""
+        return (
+            self.cpu == 0.0
+            and self.gpu == 0.0
+            and self.object_store_memory == 0.0
+            and self.memory == 0.0
+        )
+
+    def is_non_negative(self) -> bool:
+        """Returns True if all resources are non-negative."""
+        return (
+            self.cpu >= 0
+            and self.gpu >= 0
+            and self.object_store_memory >= 0
+            and self.memory >= 0
+        )
+
+    def object_store_memory_str(self) -> str:
+        """Returns a human-readable string for the object store memory field."""
+        if self.object_store_memory == float("inf"):
+            return "inf"
+        return memory_string(self.object_store_memory)
+
+    def memory_str(self) -> str:
+        """Returns a human-readable string for the memory field."""
+        if self.memory == float("inf"):
+            return "inf"
+        return memory_string(self.memory)
+
+    def copy(
+        self,
+        cpu: Optional[float] = None,
+        gpu: Optional[float] = None,
+        memory: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+    ) -> "ExecutionResources":
+        """Returns a copy of this ExecutionResources object allowing to override
+        specific resources as necessary"""
+        return ExecutionResources(
+            cpu=safe_or(cpu, self.cpu),
+            gpu=safe_or(gpu, self.gpu),
+            object_store_memory=safe_or(object_store_memory, self.object_store_memory),
+            memory=safe_or(memory, self.memory),
+        )
+
+    def add(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Adds execution resources.
+
+        Args:
+            other: The other ``ExecutionResources`` to add to this one.
+
+        Returns:
+            A new ExecutionResource object with summed resources.
+        """
+        return ExecutionResources(
+            cpu=self.cpu + other.cpu,
+            gpu=self.gpu + other.gpu,
+            object_store_memory=self.object_store_memory + other.object_store_memory,
+            memory=self.memory + other.memory,
+        )
+
+    def subtract(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Subtracts execution resources.
+
+        Args:
+            other: The other ``ExecutionResources`` to subtract from this one.
+
+        Returns:
+            A new ExecutionResource object with subtracted resources.
+        """
+        return ExecutionResources(
+            cpu=self.cpu - other.cpu,
+            gpu=self.gpu - other.gpu,
+            object_store_memory=self.object_store_memory - other.object_store_memory,
+            memory=self.memory - other.memory,
+        )
+
+    def max(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Returns the maximum for each resource type."""
+        return ExecutionResources(
+            cpu=max(self.cpu, other.cpu),
+            gpu=max(self.gpu, other.gpu),
+            object_store_memory=max(
+                self.object_store_memory, other.object_store_memory
+            ),
+            memory=max(self.memory, other.memory),
+        )
+
+    def min(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Returns the minimum for each resource type."""
+        return ExecutionResources(
+            cpu=min(self.cpu, other.cpu),
+            gpu=min(self.gpu, other.gpu),
+            object_store_memory=min(
+                self.object_store_memory, other.object_store_memory
+            ),
+            memory=min(self.memory, other.memory),
+        )
+
+    def satisfies_limit(
+        self,
+        limit: "ExecutionResources",
+        *,
+        ignore_object_store_memory: bool = False,
+    ) -> bool:
+        """Return if this resource struct meets the specified limits.
+
+        Note that None for a field means no limit.
+
+        Args:
+            limit: The resource limits to check against.
+            ignore_object_store_memory: If True, ignore the object store memory
+                limit when checking if this resource struct meets the limits.
+
+        Returns:
+            ``True`` if every resource is within the corresponding limit.
+        """
+        return (
+            self.cpu <= limit.cpu
+            and self.gpu <= limit.gpu
+            and (
+                ignore_object_store_memory
+                or self.object_store_memory <= limit.object_store_memory
+            )
+            and self.memory <= limit.memory
+        )
+
+    def scale(self, f: float) -> "ExecutionResources":
+        """Return copy with all set values scaled by `f`."""
+        if f < 0:
+            raise ValueError("Scaling factor must be non-negative.")
+        if f == 0:
+            # Explicitly handle the zero case, because `0 * inf` is undefined.
+            return ExecutionResources.zero()
+
+        return ExecutionResources(
+            cpu=self.cpu * f,
+            gpu=self.gpu * f,
+            object_store_memory=self.object_store_memory * f,
+            memory=self.memory * f,
+        )
+
+    def floordiv(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Returns the floor division of resources."""
+
+        def _div(a, b):
+            if b == 0:
+                return float("inf")
+            if a == float("inf"):
+                return float("inf")
+            return math.floor(a / b)
+
+        return ExecutionResources(
+            cpu=_div(self.cpu, other.cpu),
+            gpu=_div(self.gpu, other.gpu),
+            object_store_memory=_div(
+                self.object_store_memory, other.object_store_memory
+            ),
+            memory=_div(self.memory, other.memory),
+        )
+
+
+@DeveloperAPI
+class ExecutionOptions:
+    """Common options for execution.
+
+    Some options may not be supported on all executors (e.g., resource limits).
+
+    Attributes:
+        resource_limits: Set a limit on the logical resources a Dataset can use.
+            Autodetected by default.
+        exclude_resources: Amount of resources to exclude from Ray Data.
+            Set this if you have other workloads running on the same cluster.
+            Note,
+            - If using Ray Data with Ray Train, training resources are
+            automatically reserved and you don't need to set exclude_resources
+            for them.
+            - For each resource type, resource_limits and exclude_resources can
+            not be both set.
+        preserve_order: Set this to preserve the ordering between blocks processed by
+            operators. Off by default.
+        actor_locality_enabled: Whether to enable locality-aware task dispatch to
+            actors (off by default). This parameter applies to both stateful map and
+            streaming_split operations.
+        verbose_progress: Whether to report progress individually per operator. By
+            default, only AllToAll operators and global progress is reported. This
+            option is useful for performance debugging. On by default.
+        label_selector: A mapping of label key to label value. When set, every task
+            and actor launched by this Dataset (including shuffle, sort, and
+            aggregator actors) carries this label selector in its remote args,
+            constraining placement to nodes whose labels satisfy the selector.
+            Used to scope a Dataset to a labeled subset of the cluster (e.g.
+            ``{"ray-subcluster": "training"}``). Operator-level ``label_selector``
+            entries in ``ray_remote_args`` take precedence on key conflicts so
+            existing node-pin selectors are preserved.
+    """
+
+    def __init__(
+        self,
+        resource_limits: Optional[ExecutionResources] = None,
+        exclude_resources: Optional[ExecutionResources] = None,
+        preserve_order: bool = False,
+        actor_locality_enabled: bool = True,
+        verbose_progress: Optional[bool] = None,
+        label_selector: Optional[Dict[str, str]] = None,
+    ):
+        """Initialize execution options.
+
+        Args:
+            resource_limits: Limit on logical resources a Dataset can use.
+                Defaults to auto-detected limits.
+            exclude_resources: Resources to exclude from Ray Data.
+            preserve_order: Whether to preserve block processing order.
+            actor_locality_enabled: Whether to enable locality-aware dispatch for
+                stateful map and streaming split operations.
+            verbose_progress: Whether to report progress per operator. If None,
+                read from ``RAY_DATA_VERBOSE_PROGRESS``.
+            label_selector: Per-Dataset label selector applied to every task and
+                actor launched by Ray Data. ``None`` means no selector is added.
+        """
+        if resource_limits is None:
+            resource_limits = ExecutionResources.for_limits()
+        self.resource_limits = resource_limits
+        if exclude_resources is None:
+            exclude_resources = ExecutionResources.zero()
+        self.exclude_resources = exclude_resources
+        self.preserve_order = preserve_order
+        self.actor_locality_enabled = actor_locality_enabled
+        if verbose_progress is None:
+            verbose_progress = bool(
+                int(os.environ.get("RAY_DATA_VERBOSE_PROGRESS", "1"))
+            )
+        self.verbose_progress = verbose_progress
+        self.label_selector = label_selector
+
+    def __repr__(self) -> str:
+        return (
+            f"ExecutionOptions(resource_limits={self.resource_limits}, "
+            f"exclude_resources={self.exclude_resources}, "
+            f"preserve_order={self.preserve_order}, "
+            f"actor_locality_enabled={self.actor_locality_enabled}, "
+            f"verbose_progress={self.verbose_progress}, "
+            f"label_selector={self.label_selector})"
+        )
+
+    @property
+    def resource_limits(self) -> ExecutionResources:
+        return self._resource_limits
+
+    @resource_limits.setter
+    def resource_limits(self, value: ExecutionResources) -> None:
+        self._resource_limits = ExecutionResources.for_limits(
+            cpu=value._cpu,
+            gpu=value._gpu,
+            object_store_memory=value._object_store_memory,
+            memory=value._memory,
+        )
+
+    def is_resource_limits_default(self):
+        """Returns True if resource_limits is the default value."""
+        return self._resource_limits == ExecutionResources.for_limits()
+
+    def validate(self) -> None:
+        """Validate the options."""
+        for attr in ["cpu", "gpu", "object_store_memory"]:
+            if (
+                getattr(self.resource_limits, attr) != float("inf")
+                and getattr(self.exclude_resources, attr, 0) > 0
+            ):
+                raise ValueError(
+                    "resource_limits and exclude_resources cannot "
+                    f" both be set for {attr} resource."
+                )
+
+    @property
+    def locality_with_output(self) -> bool:
+        return False
+
+    @locality_with_output.setter
+    def locality_with_output(self, value: Union[bool, List[NodeIdStr]]) -> None:
+        if value:
+            warnings.warn(
+                "`ExecutionOptions.locality_with_output` has been removed and is now "
+                "a no-op. We don't recommend using it anymore, but if you still want "
+                "to replicate its behavior, follow the instructions in this gist: "
+                "https://gist.github.com/bveeramani/51e0383bb3680dd78fdfb92d76ea22a8.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+
+def safe_or(value: Optional[Any], alt: Any) -> Any:
+    return value if value is not None else alt
+
+
+def safe_round(
+    value: Optional[float], ndigits: Optional[int] = None
+) -> Optional[float]:
+    if value is None:
+        return None
+    elif ndigits is None or math.isinf(value):
+        return value
+    else:
+        return round(value, ndigits)

python/ray/data/_internal/execution/interfaces/executor.py

@@ -0,0 +1,106 @@
+from abc import ABC, abstractmethod
+from typing import ContextManager, Iterator, List, Optional
+
+from .execution_options import ExecutionOptions
+from .physical_operator import PhysicalOperator
+from .ref_bundle import RefBundle
+from ray.data._internal.stats import DatasetStats
+
+
+class OutputIterator(Iterator[RefBundle], ABC):
+    """Iterator used to access the output of an Executor execution.
+
+    This is a blocking iterator. Datasets guarantees that all its iterators are
+    thread-safe (i.e., multiple threads can block on them at the same time).
+    """
+
+    @abstractmethod
+    def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
+        """Can be used to pull outputs by a specified output index.
+
+        This is used to support the streaming_split() API, where the output of a
+        streaming execution is to be consumed by multiple processes.
+
+        Args:
+            output_split_idx: The output split index to get results for. This arg is
+                only allowed for iterators created by `Dataset.streaming_split()`.
+
+        Returns:
+            The next ``RefBundle`` of outputs for the given split index.
+
+        Raises:
+            StopIteration: If there are no more outputs to return.
+        """
+        ...
+
+    def __next__(self) -> RefBundle:
+        return self.get_next()
+
+
+class Executor(ContextManager, ABC):
+    """Abstract class for executors, which implement physical operator execution.
+
+    Subclasses:
+        StreamingExecutor
+    """
+
+    def __init__(self, options: ExecutionOptions):
+        """Create the executor."""
+        options.validate()
+        self._options = options
+
+    @abstractmethod
+    def execute(
+        self,
+        dag: PhysicalOperator,
+        initial_stats: Optional[DatasetStats] = None,
+        callbacks: Optional[List] = None,
+    ) -> OutputIterator:
+        """Start execution.
+
+        Args:
+            dag: The operator graph to execute.
+            initial_stats: The DatasetStats to prepend to the stats returned by the
+                executor. These stats represent actions done to compute inputs.
+            callbacks: A list of ExecutionCallbacks to run during execution.
+                This method keeps and uses the exact list you pass in, so do not
+                pass an empty list like ``[]`` directly. Create the list first,
+                then pass it.
+
+        Returns:
+            An ``OutputIterator`` over the execution's output ref bundles.
+        """
+        ...
+
+    def shutdown(self, force: bool, exception: Optional[Exception] = None):
+        """Shutdown an executor, which may still be running.
+
+        This should interrupt execution and clean up any used resources.
+
+        Args:
+            force: Controls whether shutdown should forcefully terminate all execution
+                   activity (making sure that upon returning from this method all
+                   activities are stopped). When force=False, some activities could be
+                   terminated asynchronously (ie this method won't provide guarantee
+                   that they stop executing before returning from this method)
+            exception: The exception that causes the executor to shut down, or None if
+                the executor finishes successfully.
+        """
+        pass
+
+    @abstractmethod
+    def get_stats(self) -> DatasetStats:
+        """Return stats for the execution so far.
+
+        This is generally called after `execute` has completed, but may be called
+        while iterating over `execute` results for streaming execution.
+        """
+        ...
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback, /):
+        # NOTE: ``ContextManager`` semantic must guarantee that executor
+        #       fully shutdown upon returning from this method
+        self.shutdown(force=True, exception=exc_value)

python/ray/data/_internal/execution/interfaces/ref_bundle.py

@@ -0,0 +1,481 @@
+import itertools
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Dict, Iterable, Iterator, List, Optional, Tuple
+
+import ray
+from .common import NodeIdStr
+from ray.data._internal.memory_tracing import trace_deallocation
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockMetadata,
+    Schema,
+    _take_first_non_empty_schema,
+)
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+
+@dataclass(frozen=True)
+class BlockSlice:
+    """A slice of a block."""
+
+    # Starting row offset (inclusive) within the block.
+    start_offset: int
+    # Ending row offset (exclusive) within the block.
+    end_offset: int
+
+    @property
+    def num_rows(self) -> int:
+        return self.end_offset - self.start_offset
+
+
+@dataclass(frozen=True, slots=True)
+class BlockEntry:
+    """One block delivery: the ref + the block's measured metadata.
+
+    Used as the element type of ``RefBundle.blocks`` (replaces the legacy
+    ``(ObjectRef, BlockMetadata)`` 2-tuple shape). Naming the fields makes
+    every call site self-describing and reserves room for the bundle entry
+    to grow without disturbing the surrounding shape.
+    """
+
+    ref: ObjectRef[Block]
+    metadata: BlockMetadata
+
+
+@dataclass(frozen=True)
+class RefBundle:
+    """A group of data block references and their metadata.
+
+    Operators take in and produce streams of RefBundles.
+
+    Most commonly a RefBundle consists of a single block object reference.
+    In some cases, e.g., due to block splitting, or for a reduce task, there may
+    be more than one block.
+
+    Block bundles have ownership semantics, i.e., shared ownership (similar to C++
+    shared_ptr, multiple operators share the same block bundle), or unique ownership
+    (similar to C++ unique_ptr, only one operator owns the block bundle). This
+    allows operators to know whether they can destroy blocks when they don't need
+    them. Destroying blocks eagerly is more efficient than waiting for Python GC /
+    Ray reference counting to kick in.
+    """
+
+    # Per-block entries. The size_bytes must be known in the metadata,
+    # num_rows is optional. Legacy ``(ref, metadata)`` 2-tuples are no longer
+    # accepted at construction and must be explicitly wrapped in ``BlockEntry``
+    # (``__post_init__`` rejects anything else with an actionable assertion).
+    blocks: Tuple[BlockEntry, ...]
+
+    # The schema of the blocks in this bundle. This is optional, and may be None
+    # if blocks are empty.
+    schema: Optional["Schema"]
+
+    # Whether we own the blocks (can safely destroy them).
+    owns_blocks: bool
+
+    # The slices of the blocks in this bundle. After __post_init__, this is always
+    # a list with length equal to len(blocks). Individual entries can be None to
+    # represent a full block (equivalent to BlockSlice(0, num_rows)).
+    # Pass None during construction to initialize all slices as None (full blocks).
+    slices: Optional[Tuple[Optional[BlockSlice], ...]] = None
+
+    # This attribute is used by the split() operator to assign bundles to logical
+    # output splits. It is otherwise None.
+    output_split_idx: Optional[int] = None
+
+    # Object metadata (size, locations, spilling status)
+    _cached_object_meta: Optional[Dict[ObjectRef, "_ObjectMetadata"]] = None
+
+    # Preferred locations for this bundle determined based on the locations
+    # of individual objects and their corresponding size, ie location with the
+    # largest total number of bytes present there has the highest preference.
+    _cached_preferred_locations: Optional[Dict[NodeIdStr, int]] = None
+
+    def __post_init__(self):
+        if self.schema is not None:
+            import pyarrow as pa
+
+            from ray.data._internal.pandas_block import PandasBlockSchema
+
+            assert isinstance(
+                self.schema, (pa.lib.Schema, PandasBlockSchema)
+            ), f"Schema must be a pyarrow or PandasBlockSchema, got {type(self.schema)}"
+
+        if not isinstance(self.blocks, tuple):
+            object.__setattr__(self, "blocks", tuple(self.blocks))
+        for entry in self.blocks:
+            assert isinstance(entry, BlockEntry), (
+                f"RefBundle.blocks must contain BlockEntry instances; got {type(entry).__name__}. "
+                "Construct entries with `BlockEntry(ref=..., metadata=...)`."
+            )
+
+        if self.slices is None:
+            object.__setattr__(self, "slices", (None,) * len(self.blocks))
+        else:
+            if not isinstance(self.slices, tuple):
+                object.__setattr__(self, "slices", tuple(self.slices))
+
+            assert len(self.blocks) == len(
+                self.slices
+            ), "Number of blocks and slices must match"
+            # Validate slice ranges
+            for entry, block_slice in zip(self.blocks, self.slices):
+                if block_slice is not None:
+                    assert (
+                        block_slice.start_offset >= 0
+                    ), f"Slice start_offset must be non-negative: {block_slice.start_offset}"
+                    assert (
+                        block_slice.end_offset >= block_slice.start_offset
+                    ), f"Slice end_offset must be >= start_offset: [{block_slice.start_offset}, {block_slice.end_offset})"
+                    if entry.metadata.num_rows is not None:
+                        assert (
+                            block_slice.end_offset <= entry.metadata.num_rows
+                        ), f"Slice range [{block_slice.start_offset}, {block_slice.end_offset}) exceeds block num_rows: {entry.metadata.num_rows}"
+
+        for entry in self.blocks:
+            if entry.metadata.size_bytes is None:
+                raise ValueError(
+                    "The size in bytes of the block must be known: {}".format(entry)
+                )
+
+    @property
+    def block_refs(self) -> List[ObjectRef[Block]]:
+        """List of block references in this bundle."""
+        return [entry.ref for entry in self.blocks]
+
+    @property
+    def metadata(self) -> List[BlockMetadata]:
+        """List of block metadata in this bundle."""
+        return [entry.metadata for entry in self.blocks]
+
+    def num_rows(self) -> Optional[int]:
+        """Number of rows present in this bundle, if known.
+
+        Iterates through blocks and their corresponding slices to calculate the total.
+        Note: Block metadata always refers to the full block, not the slice.
+
+        - If block_slice is None, uses the full block's metadata.num_rows
+        - If block_slice is present, uses the slice's num_rows (partial block portion)
+        - Returns None if any full block has unknown row count (metadata.num_rows is None)
+        """
+        total = 0
+        for metadata, block_slice in zip(self.metadata, self.slices):
+            if block_slice is None:
+                if metadata.num_rows is None:
+                    return None
+                total += metadata.num_rows
+            else:
+                total += block_slice.num_rows
+        return total
+
+    def size_bytes(self) -> int:
+        """Size of the blocks of this bundle in bytes.
+
+        Iterates through blocks and their corresponding slices to calculate the total size.
+        Note: Block metadata always refers to the full block, not the slice.
+
+        - If block_slice is None, uses the full block's metadata.size_bytes
+        - If block_slice is present but num_rows is unknown or zero, uses full metadata.size_bytes
+        - If block_slice represents a partial block, estimates size proportionally based on
+          (metadata.size_bytes / metadata.num_rows) * block_slice.num_rows
+        - Otherwise, uses the full metadata.size_bytes
+        """
+        total = 0
+        for entry, block_slice in zip(self.blocks, self.slices):
+            metadata = entry.metadata
+            if block_slice is None:
+                # Full block
+                total += metadata.size_bytes
+            elif metadata.num_rows is None or metadata.num_rows == 0:
+                # Unknown num_rows or empty block - use full metadata size
+                total += metadata.size_bytes
+            elif metadata.num_rows != block_slice.num_rows:
+                # Partial block - estimate size based on rows
+                per_row = metadata.size_bytes / metadata.num_rows
+                total += max(1, round(per_row * block_slice.num_rows))
+            else:
+                total += metadata.size_bytes
+        return total
+
+    def destroy_if_owned(self) -> int:
+        """Clears the object store memory for these blocks if owned.
+
+        Returns:
+            The number of bytes freed.
+        """
+        should_free = self.owns_blocks and DataContext.get_current().eager_free
+        for block_ref in self.block_refs:
+            trace_deallocation(
+                block_ref, "RefBundle.destroy_if_owned", free=should_free
+            )
+        return self.size_bytes() if should_free else 0
+
+    def get_preferred_object_locations(self) -> Dict[NodeIdStr, int]:
+        """Returns a mapping of node IDs to total bytes stored on each node.
+
+        Returns:
+            Dict mapping node ID to total bytes stored on that node
+        """
+        meta = self._get_cached_metadata()
+
+        if self._cached_preferred_locations is None:
+            preferred_locs: Dict[NodeIdStr, int] = defaultdict(int)
+
+            for ref, obj_meta in meta.items():
+                for loc in obj_meta.locs:
+                    preferred_locs[loc] += obj_meta.size
+
+            # NOTE: We're working around object being immutable to update cached
+            #       values (safe)
+            object.__setattr__(self, "_cached_preferred_locations", preferred_locs)
+
+        return self._cached_preferred_locations
+
+    def _get_cached_metadata(self) -> Dict[ObjectRef, "_ObjectMetadata"]:
+        if self._cached_object_meta is None:
+            # This call is pretty fast for owned objects (~5k/s), so we don't need to
+            # batch it for now.
+            meta = ray.experimental.get_local_object_locations(self.block_refs)
+            # Extract locations
+            object_metas: Dict[ObjectRef, _ObjectMetadata] = {
+                ref: _ObjectMetadata(
+                    size=meta[ref]["object_size"],
+                    spilled=meta[ref]["did_spill"],
+                    locs=meta[ref]["node_ids"],
+                )
+                for ref in self.block_refs
+            }
+
+            # NOTE: We're working around object being immutable to update cached
+            #       values (safe)
+            object.__setattr__(self, "_cached_object_meta", object_metas)
+
+        return self._cached_object_meta
+
+    def slice(self, needed_rows: int) -> Tuple["RefBundle", "RefBundle"]:
+        """Slice a Ref Bundle into the first bundle containing the first `needed_rows` rows and the remaining bundle containing the remaining rows.
+
+        Args:
+            needed_rows: Number of rows to take from the head of the bundle.
+
+        Returns:
+            A tuple of (sliced_bundle, remaining_bundle). The needed rows must be less than the number of rows in the bundle.
+        """
+        assert needed_rows > 0, "needed_rows must be positive."
+        assert (
+            self.num_rows() is not None
+        ), "Cannot slice a RefBundle with unknown number of rows."
+        assert (
+            needed_rows < self.num_rows()
+        ), f"To slice a RefBundle, the number of requested rows must be less than the number of rows in the bundle. Requested {needed_rows} rows but bundle only has {self.num_rows()} rows."
+
+        block_slices = []
+        for metadata, block_slice in zip(self.metadata, self.slices):
+            if block_slice is None:
+                # None represents a full block, convert to explicit BlockSlice
+                assert (
+                    metadata.num_rows is not None
+                ), "Cannot derive block slice for a RefBundle with unknown block row counts."
+                block_slices.append(
+                    BlockSlice(start_offset=0, end_offset=metadata.num_rows)
+                )
+            else:
+                block_slices.append(block_slice)
+
+        consumed_blocks: List[BlockEntry] = []
+        consumed_slices: List[BlockSlice] = []
+        remaining_blocks: List[BlockEntry] = []
+        remaining_slices: List[BlockSlice] = []
+
+        rows_to_take = needed_rows
+
+        for entry, block_slice in zip(self.blocks, block_slices):
+            block_rows = block_slice.num_rows
+            if rows_to_take >= block_rows:
+                consumed_blocks.append(entry)
+                consumed_slices.append(block_slice)
+                rows_to_take -= block_rows
+            else:
+                if rows_to_take == 0:
+                    remaining_blocks.append(entry)
+                    remaining_slices.append(block_slice)
+                    continue
+                consume_slice = BlockSlice(
+                    start_offset=block_slice.start_offset,
+                    end_offset=block_slice.start_offset + rows_to_take,
+                )
+                consumed_blocks.append(entry)
+                consumed_slices.append(consume_slice)
+
+                leftover_rows = block_rows - rows_to_take
+                if leftover_rows > 0:
+                    remainder_slice = BlockSlice(
+                        start_offset=consume_slice.end_offset,
+                        end_offset=block_slice.end_offset,
+                    )
+                    remaining_blocks.append(entry)
+                    remaining_slices.append(remainder_slice)
+
+                rows_to_take = 0
+
+        sliced_bundle = RefBundle(
+            blocks=tuple(consumed_blocks),
+            schema=self.schema,
+            owns_blocks=False,
+            slices=tuple(consumed_slices) if consumed_slices else None,
+        )
+
+        remaining_bundle = RefBundle(
+            blocks=tuple(remaining_blocks),
+            schema=self.schema,
+            owns_blocks=False,
+            slices=tuple(remaining_slices) if remaining_slices else None,
+        )
+
+        return sliced_bundle, remaining_bundle
+
+    @classmethod
+    def merge_ref_bundles(cls, bundles: Iterable["RefBundle"]) -> "RefBundle":
+        """Merge multiple RefBundles into a single RefBundle.
+
+        Args:
+            bundles: An iterable of RefBundles to merge.
+
+        Returns:
+            A single RefBundle containing all blocks from the input bundles.
+            owns_blocks is True only if all input bundles own their blocks.
+            schema is the first non-empty schema found.
+        """
+
+        bundles = list(bundles)
+        if not bundles:
+            return cls(blocks=(), owns_blocks=True, schema=None)
+
+        merged_blocks = list(
+            itertools.chain.from_iterable(bundle.blocks for bundle in bundles)
+        )
+        merged_slices = list(
+            itertools.chain.from_iterable(bundle.slices for bundle in bundles)
+        )
+        # Ray Data uses the `owns_blocks` flag to determine if the system can eagerly
+        # destroy blocks when they're no longer needed. To be safe, we only set this
+        # to True if all input bundles own their blocks.
+        owns_blocks = all(bundle.owns_blocks for bundle in bundles)
+        # TODO: Reconcile the schemas rather than taking the first non-empty schema.
+        schema = _take_first_non_empty_schema(bundle.schema for bundle in bundles)
+        return cls(
+            blocks=tuple(merged_blocks),
+            schema=schema,
+            owns_blocks=owns_blocks,
+            slices=merged_slices,
+        )
+
+    def __eq__(self, other: "RefBundle"):
+        if self is other:
+            return True
+        elif not isinstance(other, RefBundle):
+            return False
+
+        return (
+            self.blocks == other.blocks
+            and self.slices == other.slices
+            # NOTE: We're establishing a requirement of schemas for `RefBundle`
+            #       to be exactly the same object for it to be considered equal.
+            #
+            #       This is necessary to avoid a full schema equality check that
+            #       is computationally intensive.
+            and self.schema is other.schema
+            and self.owns_blocks == other.owns_blocks
+            and self.output_split_idx == other.output_split_idx
+        )
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                # Only hash block refs
+                *[entry.ref for entry in self.blocks],
+                *self.slices,
+                # Check out comment in ``__eq__``
+                id(self.schema),
+                self.owns_blocks,
+                self.output_split_idx,
+            )
+        )
+
+    def __len__(self) -> int:
+        return len(self.blocks)
+
+    def __str__(self) -> str:
+        lines = [
+            f"RefBundle({len(self.blocks)} blocks,",
+            f"  {self.num_rows()} rows,",
+            f"  schema={self.schema},",
+            f"  owns_blocks={self.owns_blocks},",
+            "  blocks=(",
+        ]
+
+        # Loop through each block and show details
+        for i, (entry, block_slice) in enumerate(zip(self.blocks, self.slices)):
+            metadata = entry.metadata
+            row_str = (
+                f"{metadata.num_rows} rows"
+                if metadata.num_rows is not None
+                else "unknown rows"
+            )
+            bytes_str = f"{metadata.size_bytes} bytes"
+            slice_str = (
+                f"slice={block_slice}"
+                if block_slice is not None
+                else "slice=None (full block)"
+            )
+
+            lines.append(f"    {i}: {row_str}, {bytes_str}, {slice_str}")
+
+        lines.append("  )")
+        lines.append(")")
+
+        return "\n".join(lines)
+
+
+@dataclass
+class _ObjectMetadata:
+    # Object size in bytes
+    size: int
+    # Flag whether object has been spilled
+    spilled: bool
+    # List of nodes object exists on
+    locs: List[NodeIdStr] = None
+
+
+def _ref_bundles_iterator_to_block_refs_list(
+    ref_bundles: Iterator[RefBundle],
+) -> List[ObjectRef[Block]]:
+    """Convert an iterator of RefBundles to a list of Block object references."""
+    return [
+        block_ref for ref_bundle in ref_bundles for block_ref in ref_bundle.block_refs
+    ]
+
+
+def _iter_sliced_blocks(
+    blocks: Iterable[Block],
+    slices: List[Optional[BlockSlice]],
+) -> Iterator[Block]:
+    blocks_list = list(blocks)
+    for block, block_slice in zip(blocks_list, slices):
+        if block_slice is None:
+            # None represents a full block - yield it as is
+            yield block
+        else:
+            accessor = BlockAccessor.for_block(block)
+            start = block_slice.start_offset
+            end = block_slice.end_offset
+            assert start <= end, "start must be less than end"
+            assert start >= 0, "start must be non-negative"
+            assert (
+                end <= accessor.num_rows()
+            ), "end must be less than or equal to the number of rows in the block"
+
+            yield accessor.slice(start, end, copy=False)

python/ray/data/_internal/execution/interfaces/task_context.py

@@ -0,0 +1,88 @@
+import contextlib
+import threading
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.operators.map_transformer import MapTransformer
+    from ray.data._internal.progress.base_progress import BaseProgressBar
+
+
+_thread_local = threading.local()
+
+
+@dataclass
+class TaskContext:
+    """This describes the information of a task running block transform."""
+
+    # The index of task. Each task has a unique task index within the same
+    # operator.
+    task_idx: int
+
+    # Name of the operator that this task belongs to.
+    op_name: str
+
+    # The dictionary of sub progress bar to update. The key is name of sub progress
+    # bar. Note this is only used on driver side.
+    # TODO(chengsu): clean it up from TaskContext with new optimizer framework.
+    sub_progress_bar_dict: Optional[Dict[str, "BaseProgressBar"]] = None
+
+    # NOTE(hchen): `upstream_map_transformer` and `upstream_map_ray_remote_args`
+    # are only used for `RandomShuffle`. DO NOT use them for other operators.
+    # Ideally, they should be handled by the optimizer, and should be transparent
+    # to the specific operators.
+    # But for `RandomShuffle`, the AllToAllOperator doesn't do the shuffle itself.
+    # It uses `ExchangeTaskScheduler` to launch new tasks to do the shuffle.
+    # That's why we need to pass them to `ExchangeTaskScheduler`.
+    # TODO(hchen): Use a physical operator to do the shuffle directly.
+
+    # The underlying function called in a MapOperator; this is used when fusing
+    # an AllToAllOperator with an upstream MapOperator.
+    upstream_map_transformer: Optional["MapTransformer"] = None
+
+    # The Ray remote arguments of the fused upstream MapOperator.
+    # This should be set if upstream_map_transformer is set.
+    upstream_map_ray_remote_args: Optional[Dict[str, Any]] = None
+
+    # Override of the target max-block-size for the task
+    target_max_block_size_override: Optional[int] = None
+
+    # Additional keyword arguments passed to the task.
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def get_current(cls) -> Optional["TaskContext"]:
+        """Get the TaskContext for the current thread.
+        Returns None if no TaskContext has been set.
+        """
+
+        return getattr(_thread_local, "task_context", None)
+
+    @classmethod
+    def set_current(cls, context: "TaskContext") -> None:
+        """Set the TaskContext for the current thread.
+
+        Args:
+            context: The TaskContext instance to set for this thread
+        """
+
+        _thread_local.task_context = context
+
+    @classmethod
+    def reset_current(cls):
+        """Clear the current thread's TaskContext."""
+
+        if hasattr(_thread_local, "task_context"):
+            delattr(_thread_local, "task_context")
+
+    @classmethod
+    @contextlib.contextmanager
+    def current(cls, context: "TaskContext") -> Iterator["TaskContext"]:
+        """Sets this TaskContext as current for the scope
+        of the context block and resets it on exit.
+        """
+        cls.set_current(context)
+        try:
+            yield context
+        finally:
+            cls.reset_current()

python/ray/data/_internal/execution/interfaces/transform_fn.py

@@ -0,0 +1,14 @@
+from typing import Callable, List, Tuple
+
+from .ref_bundle import RefBundle
+from .task_context import TaskContext
+from ray.data._internal.stats import StatsDict
+
+# Result type of AllToAllTransformFn.
+AllToAllTransformFnResult = Tuple[List[RefBundle], StatsDict]
+
+# Block transform function applied in AllToAllOperator.
+AllToAllTransformFn = Callable[
+    [List[RefBundle], TaskContext],
+    AllToAllTransformFnResult,
+]

python/ray/data/_internal/execution/metadata_fetcher.py

@@ -0,0 +1,461 @@
+"""Metadata-fetch strategy for the streaming executor.
+
+``DataOpTask.on_data_ready`` pulls ``(block_ref, meta_ref)`` pairs from a task's
+streaming generator; a ``MetadataFetcher`` turns each pair into an emitted
+``RefBundle``. Two modes, selected by ``RAY_DATA_METADATA_PREFETCH_ON_THREAD``
+(default on):
+
+- :class:`ThreadedMetadataFetcher` (default): defer every pair and fetch its
+  metadata on a dedicated background thread, so the scheduling loop never blocks
+  on ``ray.get(meta_ref)``. The output-budget size comes from the block's local
+  ``object_size`` (no RPC); completion is postponed until the task's deferred
+  pairs have emitted, and the per-operator FIFO preserves emission order.
+- :class:`InlineMetadataFetcher`: fetch each pair's metadata inline with
+  ``ray.get`` and emit the ``RefBundle`` immediately, budgeting off
+  ``meta.size_bytes``; completion and task-failure are handled inline by
+  ``on_data_ready``.
+"""
+
+import logging
+import pickle
+import queue as queue_module
+import threading
+from abc import ABC, abstractmethod
+from collections import defaultdict, deque
+from collections.abc import Hashable
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+
+import ray
+import ray.exceptions
+from ray._common.utils import env_bool
+from ray.data._internal.execution.interfaces.physical_operator import (
+    METADATA_GET_TIMEOUT_S,
+    METADATA_WAIT_TIMEOUT_S,
+    DataOpTask,
+    DeferredEmit,
+)
+from ray.data.block import BlockMetadataWithSchema
+from ray.exceptions import GetTimeoutError
+from ray.experimental.locations import get_local_object_locations
+from ray.util.debug import log_once
+
+logger = logging.getLogger(__name__)
+
+# How long ``ThreadedMetadataFetcher.stop`` waits for the fetch thread to exit.
+_FETCH_THREAD_JOIN_TIMEOUT_S = 5.0
+
+# How long the fetch thread's ``ray.wait`` blocks each pass — bounds the
+# busy-wait when nothing is ready, and how long a straggler can delay a batch.
+_FETCH_WAIT_TIMEOUT_S = 0.1
+
+# Selects the mode (see module docstring). Threaded by default; set to 0/false to
+# fall back to the synchronous inline path.
+_PREFETCH_ON_THREAD = env_bool("RAY_DATA_METADATA_PREFETCH_ON_THREAD", True)
+
+
+class MetadataFetcher(ABC):
+    def start(self) -> None:
+        """Start any background machinery."""
+
+    def stop(self) -> None:
+        """Stop any background machinery."""
+
+    @abstractmethod
+    def in_data_ready_get_object_size(self, task: DataOpTask) -> Optional[int]:
+        """Handle one pulled pair inside the ``on_data_ready`` loop. The pair's
+        refs are read off the task (``task.pending_block_ref`` /
+        ``task.pending_meta_ref``).
+
+        Returns the output-budget bytes for this pair (0 if the size is
+        unknown), or ``None`` to mean "the metadata isn't available yet — stop
+        and retry next iteration" (the caller breaks, leaving the refs set).
+
+        ``None`` must be returned ONLY when the pair was NOT consumed: the
+        caller will hand the same pair back on the next call, so returning
+        ``None`` after emitting/deferring it would emit the block twice.
+        """
+
+    def in_data_ready_done(self, task: DataOpTask) -> None:
+        """Called once a task is drained (generator exhausted/failed)."""
+
+    def submit(self, op_key: Hashable, tasks: List[DataOpTask]) -> None:
+        """Hand the operator's deferred pairs off for processing, and record
+        any end-of-stream/failed tasks whose completion is postponed."""
+
+    def emit_ready_and_fire_done_callbacks(self) -> List[Tuple[str, BaseException]]:
+        """Run once at the end of ``process_completed_tasks``. Returns
+        ``(operator_name, exception)`` for each block-level fetch failure, for
+        the caller's ``max_errored_blocks`` accounting. Default: nothing to do."""
+        return []
+
+
+class InlineMetadataFetcher(MetadataFetcher):
+    """Synchronous mode: fetch metadata inline and emit immediately. Holds no
+    state and starts no thread."""
+
+    def in_data_ready_get_object_size(self, task: DataOpTask) -> Optional[int]:
+        # The ref resolves to pickled metadata bytes, not a BlockMetadata.
+        meta_ref: "ray.ObjectRef[Any]" = task.pending_meta_ref
+        try:
+            # The timeout includes the time to ship the metadata to this node,
+            # so a 0 timeout could cancel an in-flight download. Use a small
+            # non-zero value to avoid that.
+            meta_bytes: bytes = ray.get(meta_ref, timeout=METADATA_GET_TIMEOUT_S)
+        except ray.exceptions.GetTimeoutError:
+            # We have refs to the block and its metadata, but the metadata
+            # object isn't available. This can happen if the node dies. Leave
+            # the pair pending and retry next iteration.
+            logger.warning(
+                f"Timed out ({METADATA_GET_TIMEOUT_S}s) waiting for metadata from "
+                f"operator '{task.operator_name}' "
+                f"(metadata_ref={meta_ref.hex()}). "
+                f"Possible causes include a worker crash, node preemption, or an "
+                f"overloaded worker or head node. Will retry next iteration. "
+                f"If this repeats, check the Ray dashboard and logs for worker "
+                f"crashes, node preemption, or overload."
+            )
+            return None
+        return task.produce_block(task.pending_block_ref, meta_bytes)
+
+    def in_data_ready_done(self, task: DataOpTask) -> None:
+        # Inline mode fires the done-callback the moment the generator drains:
+        # all of the task's pairs have already emitted inline, so
+        # ``_pending_emit_count`` is 0. A task failure is re-raised after the
+        # callback.
+        task.mark_done()
+        if task.task_error is not None:
+            raise task.task_error from None
+
+
+class _Signal(Enum):
+    """Sentinels used by :class:`ThreadedMetadataFetcher`.
+
+    ``STOP`` is enqueued on the request queue to tell the fetch thread to exit;
+    ``NOT_READY`` marks "ref not fetched yet" in the result store. Members of a
+    single enum so identity checks narrow cleanly under type checkers.
+    """
+
+    STOP = "stop"
+    NOT_READY = "not_ready"
+
+
+# A request-queue item: a batch of meta_refs to fetch, or the stop sentinel.
+_Request = Union[List["ray.ObjectRef"], _Signal]
+
+
+class ThreadedMetadataFetcher(MetadataFetcher):
+    """Asynchronous mode: defer every pulled pair and fetch its metadata on a
+    dedicated background thread, so the scheduling (executor) thread never blocks
+    on ``ray.get(meta_refs)``.
+
+    The two threads communicate through one thread-safe queue (``_request_q``);
+    fetched bytes come back via ``_results``. The background thread fetches the
+    refs ``ray.wait(fetch_local=True)`` reports ready; a ref stuck on a bad node
+    merely stays pending instead of wedging the thread.
+
+    Data flow (for a single operator)::
+
+        Executor thread                              Fetch thread
+        ---------------                              ------------
+        on_data_ready  --defer-->  _pending_deferred
+        submit(op)     --meta_refs-->  _request_q  ----->  ray.wait(ready)
+                                                                + ray.get
+                                                                   |
+                            _results  <----- fetched bytes --------+
+        emit_ready_and_fire_done_callbacks():
+            _fifos[op]:  head [d0] -> [d1] -> [d2] tail   (append = yield order)
+                              |
+                              `- emit front-first while its bytes are in
+                                 _results; stop at the first pair not back yet,
+                                 so this op's RefBundle order is preserved.
+
+    Operators each get their own FIFO and are independent, so one operator
+    waiting on metadata never blocks another.
+    """
+
+    def __init__(
+        self,
+        *,
+        get_objects: Optional[Callable] = None,
+        wait_for_objects: Optional[Callable] = None,
+        get_object_locations: Optional[Callable] = None,
+    ):
+        """Create a ThreadedMetadataFetcher.
+
+        Args:
+            get_objects: Fetches object values by ref, like ``ray.get``
+                (called as ``get_objects(refs, timeout=...)``). Injectable so
+                tests can drive the fetch path without a real cluster.
+            wait_for_objects: Reports which refs are locally available, like
+                ``ray.wait`` (called with ``fetch_local=True``).
+            get_object_locations: Returns per-ref location info (including
+                ``object_size``), like ``get_local_object_locations``.
+        """
+        self._get_objects: Callable = get_objects or ray.get
+        self._wait_for_objects: Callable = wait_for_objects or ray.wait
+        self._get_object_locations: Callable = (
+            get_object_locations or get_local_object_locations
+        )
+
+        self._request_q: "queue_module.Queue[_Request]" = queue_module.Queue()
+        # fetch thread -> executor: meta_ref -> bytes (or captured Exception).
+        self._results: Dict["ray.ObjectRef", Any] = {}
+        self._results_lock = threading.Lock()
+
+        # Executor-thread-only state below.
+        # Pairs deferred by ``in_data_ready_get_object_size`` for the current operator,
+        # flushed into the FIFOs by ``submit``.
+        self._pending_deferred: List[DeferredEmit] = []
+        # Per-operator (keyed by the caller's op key) FIFO of pairs awaiting
+        # metadata, in append (= emission) order. Each op's deque is drained
+        # front-first so that op's RefBundle emission order is preserved.
+        self._fifos: "defaultdict[Hashable, deque[DeferredEmit]]" = defaultdict(deque)
+        # Drained (end-of-stream/failed) tasks whose done-callback is postponed
+        # until all of their deferred pairs have been emitted. A set so a task
+        # re-seen on a later iteration (still pending) isn't registered — or
+        # fired — twice.
+        self._drained_tasks: Set[DataOpTask] = set()
+
+        self._thread = threading.Thread(
+            target=self._run, name="ray-data-metadata-prefetch", daemon=True
+        )
+        self._started = False
+        self._stopped = False
+
+    def start(self) -> None:
+        if not self._started:
+            self._started = True
+            self._thread.start()
+
+    def stop(self) -> None:
+        if self._stopped:
+            return
+        self._stopped = True
+        self._request_q.put(_Signal.STOP)
+        if self._started:
+            try:
+                self._thread.join(timeout=_FETCH_THREAD_JOIN_TIMEOUT_S)
+                if self._thread.is_alive():
+                    logger.warning(
+                        "Metadata-fetch thread did not exit within "
+                        f"{_FETCH_THREAD_JOIN_TIMEOUT_S}s; leaving the daemon "
+                        "thread behind."
+                    )
+            except Exception:
+                logger.warning(
+                    "Failed to join the metadata-fetch thread.", exc_info=True
+                )
+
+    def in_data_ready_get_object_size(self, task: DataOpTask) -> Optional[int]:
+        block_ref = task.pending_block_ref
+        meta_ref = task.pending_meta_ref
+        # Output-budget size from the block's local object_size (no RPC).
+        # Normally known: the driver owns the just-yielded block ref, so the
+        # value (which matches ``meta.size_bytes``) is in the local object
+        # directory.
+        # TODO: ``object_size`` is the object-store size of the block, which can
+        # differ from ``meta.size_bytes`` (the in-memory/logical size). We should
+        # add an explicit ``object_size_bytes`` to ``BlockMetadata`` and use it
+        # directly, so the fallback below doesn't conflate the two.
+        info: Optional[Dict[str, Any]] = self._get_object_locations([block_ref]).get(
+            block_ref
+        )
+        object_size: Optional[int] = (
+            info.get("object_size") if info is not None else None
+        )
+        if object_size is None:
+            # Rare: no local size record. Fall back to a short metadata
+            # ``ray.get`` for the size. Log once to flag the path without
+            # spamming if it recurs.
+            if log_once(f"data_object_size_unavailable_{task.operator_name}"):
+                logger.warning(
+                    "Local object_size unavailable for a block from operator "
+                    "'%s'; falling back to its metadata for the output-budget "
+                    "size.",
+                    task.operator_name,
+                )
+            try:
+                meta_with_schema: BlockMetadataWithSchema = pickle.loads(
+                    self._get_objects(meta_ref, timeout=METADATA_WAIT_TIMEOUT_S)
+                )
+            except ray.exceptions.GetTimeoutError:
+                # Metadata isn't local yet either. Leave this pair pending and
+                # retry next iteration.
+                return None
+            # Coalesce a missing size to 0: None is reserved for the "pair not
+            # consumed, retry" signal above, and this pair IS consumed
+            # (deferred) below — returning None here would defer it twice.
+            object_size = meta_with_schema.metadata.size_bytes or 0
+
+        self._pending_deferred.append(
+            DeferredEmit(task=task, block_ref=block_ref, meta_ref=meta_ref)
+        )
+        return object_size
+
+    def submit(self, op_key: Hashable, tasks: List[DataOpTask]) -> None:
+        """Queue the current operator's deferred pairs for metadata fetch +
+        emission — append them to the op's FIFO (preserving emission order) and
+        hand their ``meta_ref``s to the fetch thread — and record any drained
+        (end-of-stream/failed) tasks so their done-callback fires once all of
+        their deferred pairs have emitted. Must run on the executor thread."""
+        deferred = self._pending_deferred
+        self._pending_deferred = []
+        if deferred:
+            fifo = self._fifos[op_key]
+            for d in deferred:
+                d.task.add_pending_metadata_ref()
+                fifo.append(d)
+            self._request_q.put([d.meta_ref for d in deferred])
+        for task in tasks:
+            if task.is_drained():
+                self._drained_tasks.add(task)
+
+    def emit_ready_and_fire_done_callbacks(self) -> List[Tuple[str, BaseException]]:
+        """Emit whatever's ready (per-op order) then fire postponed done
+        callbacks. Returns ``(operator_name, exception)`` for each pair whose
+        metadata fetch failed, for the caller's ``max_errored_blocks``
+        accounting. Must run on the executor thread."""
+        return self._emit_ready() + self._fire_done_callbacks()
+
+    def _emit_ready(self) -> List[Tuple[str, BaseException]]:
+        # Emit every pair whose metadata is now available, in per-op append
+        # order. A failed fetch is accounted as emitted (so the task can still
+        # complete) but its block is dropped and the error is surfaced to the
+        # caller rather than raised.
+        failures: List[Tuple[str, BaseException]] = []
+        for fifo in self._fifos.values():
+            while fifo:
+                d = fifo[0]
+                result = self._pop_result(d.meta_ref)
+                if result is _Signal.NOT_READY:
+                    # Preserve order: stop at the first pair still in flight;
+                    # this operator is retried next call.
+                    # TODO: order only needs to be preserved when
+                    # ``DataContext.get_current().execution_options.preserve_order``
+                    # is True; otherwise we could skip past in-flight pairs and
+                    # emit any ready ones.
+                    break
+                fifo.popleft()
+                d.task.mark_emitted()
+                if isinstance(result, BaseException):
+                    failures.append((d.task.operator_name, result))
+                    continue
+                try:
+                    d.task.produce_block(d.block_ref, result)
+                except Exception as e:
+                    # Deserializing/emitting the fetched metadata can also fail
+                    # (e.g. ``pickle.loads`` raising on a corrupt object). Treat
+                    # it as a block-level error and route it through the same
+                    # accounting, rather than letting it escape.
+                    failures.append((d.task.operator_name, e))
+        return failures
+
+    def _fire_done_callbacks(self) -> List[Tuple[str, BaseException]]:
+        # Fire postponed done-callbacks for drained tasks whose pairs have all
+        # emitted. A failed task fires with its error, which is also surfaced for
+        # ``max_errored_blocks`` accounting.
+        if not self._drained_tasks:
+            return []
+        failures: List[Tuple[str, BaseException]] = []
+        to_mark_done = [t for t in self._drained_tasks if not t.has_pending_emits()]
+        for task in to_mark_done:
+            if task.task_error is not None:
+                failures.append((task.operator_name, task.task_error))
+            task.mark_done()
+        self._drained_tasks.difference_update(to_mark_done)
+        return failures
+
+    def _pop_result(self, ref: "ray.ObjectRef") -> Any:
+        with self._results_lock:
+            return self._results.pop(ref, _Signal.NOT_READY)
+
+    def _run(self) -> None:
+        # Fetch-thread loop: accumulate requested meta_refs into a pending set
+        # and hand them to ``_fetch``, which fetches the locally-available ones
+        # and returns those still in flight.
+        pending: List["ray.ObjectRef"] = []
+        while True:
+            # Block on the queue only when idle; while refs are in flight,
+            # don't block here — get back to ``_fetch`` to keep them moving.
+            try:
+                item = self._request_q.get(block=not pending)
+            except queue_module.Empty:
+                item = None
+            # Drain whatever else is already queued into a single fetch batch.
+            while item is not None:
+                if isinstance(item, _Signal):
+                    assert item is _Signal.STOP
+                    # ``stop()`` enqueued the STOP sentinel: fast teardown —
+                    # drop any in-flight refs and exit. ``stop`` runs after the
+                    # scheduling loop (which feeds us) is joined, so there's
+                    # nothing left to emit.
+                    return
+                pending.extend(item)
+                try:
+                    item = self._request_q.get_nowait()
+                except queue_module.Empty:
+                    item = None
+
+            if pending:
+                pending = self._fetch(pending)
+
+    def _fetch(self, pending: List["ray.ObjectRef"]) -> List["ray.ObjectRef"]:
+        """One fetch pass over ``pending``:
+
+        1. ``ray.wait(fetch_local=True)`` pulls the metadata objects to this
+           (driver) node and reports which are locally available.
+        2. ``ray.get`` the ready refs in one batch (``timeout=0`` — they're
+           local) and publish the bytes to ``_results``.
+        3. If the batched get raises (it hides which ref failed), fall back to
+           per-ref gets to isolate the failure and keep the rest.
+        4. Return the refs to retry next pass: the not-yet-local ones, plus any
+           that raced out of the local store. A ref that resolved to an error is
+           published as that exception for ``_emit_ready`` to surface.
+        """
+        ready, not_ready = self._wait_for_objects(
+            pending,
+            num_returns=len(pending),
+            timeout=_FETCH_WAIT_TIMEOUT_S,
+            fetch_local=True,
+        )
+        if not ready:
+            return not_ready
+
+        retry: List["ray.ObjectRef"] = []
+        try:
+            values = self._get_objects(ready, timeout=0)
+            results: Dict["ray.ObjectRef", Any] = dict(zip(ready, values))
+        except Exception:
+            # A batched get raises on the first error and hides which ref
+            # failed; retry per-ref to isolate it and keep the rest.
+            results = {}
+            for ref in ready:
+                try:
+                    results[ref] = self._get_objects(ref, timeout=0)
+                except GetTimeoutError:
+                    # ray.wait reported it ready but it's no longer local (e.g. a
+                    # raced eviction). Re-queue rather than treating it as a
+                    # block-level error. Shouldn't be common — log once.
+                    if log_once("ray_data_metadata_prefetch_not_local"):
+                        logger.warning(
+                            "A metadata object reported ready by ray.wait was "
+                            "not locally available for ray.get; re-queuing it. "
+                            "If this repeats, the object store may be under "
+                            "memory pressure (objects evicted/spilled)."
+                        )
+                    retry.append(ref)
+                except Exception as e:
+                    results[ref] = e
+        if results:
+            with self._results_lock:
+                self._results.update(results)
+        return not_ready + retry
+
+
+def make_metadata_fetcher() -> MetadataFetcher:
+    """Build the metadata fetcher for the configured mode (see module
+    docstring)."""
+    if _PREFETCH_ON_THREAD:
+        return ThreadedMetadataFetcher()
+    return InlineMetadataFetcher()

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

@@ -0,0 +1,3 @@
+from .actor_location import ActorLocationTracker, get_or_create_actor_location_tracker
+
+__all__ = ["get_or_create_actor_location_tracker", "ActorLocationTracker"]

python/ray/data/_internal/execution/node_trackers/actor_location.py

@@ -0,0 +1,39 @@
+import threading
+from typing import List
+
+import ray
+
+
+@ray.remote(num_cpus=0, max_restarts=-1, max_task_retries=-1)
+class ActorLocationTracker:
+    def __init__(self):
+        self._actor_locations = {}
+        self._actor_locations_lock = threading.Lock()
+
+    def update_actor_location(self, logical_actor_id: str, node_id: str):
+        with self._actor_locations_lock:
+            self._actor_locations[logical_actor_id] = node_id
+
+    def get_actor_locations(self, logical_actor_ids: List[str]):
+        return {
+            logical_actor_id: self._actor_locations.get(logical_actor_id, None)
+            for logical_actor_id in logical_actor_ids
+        }
+
+
+def get_or_create_actor_location_tracker():
+
+    # Pin the actor location tracker to the local node so it fate-shares with the driver.
+    # NOTE: for Ray Client, the ray.get_runtime_context().get_node_id() should
+    # point to the head node.
+    label_selector = {
+        ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
+    }
+    return ActorLocationTracker.options(
+        name="ActorLocationTracker",
+        namespace="ActorLocationTracker",
+        get_if_exists=True,
+        lifetime="detached",
+        label_selector=label_selector,
+        max_concurrency=8,
+    ).remote()

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()

python/ray/data/_internal/execution/ranker.py

@@ -0,0 +1,100 @@
+"""Ranker component for operator selection in streaming executor."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, List, Protocol, Tuple, TypeVar
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+# Protocol for comparable ranking values
+class Comparable(Protocol):
+    """Protocol for types that can be compared for ranking."""
+
+    def __lt__(self, other: "Comparable") -> bool:
+        ...
+
+    def __le__(self, other: "Comparable") -> bool:
+        ...
+
+    def __gt__(self, other: "Comparable") -> bool:
+        ...
+
+    def __ge__(self, other: "Comparable") -> bool:
+        ...
+
+    def __eq__(self, other: "Comparable") -> bool:
+        ...
+
+
+# Generic type for comparable ranking values
+RankingValue = TypeVar("RankingValue", bound=Comparable)
+
+
+class Ranker(ABC, Generic[RankingValue]):
+    """Abstract base class for operator ranking strategies."""
+
+    @abstractmethod
+    def rank_operator(
+        self,
+        op: PhysicalOperator,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ) -> RankingValue:
+        """Rank operator for execution priority.
+
+        Operator to run next is selected as the one with the *smallest* value
+        of the lexicographically ordered ranks composed of (in order):
+
+        Args:
+            op: Operator to rank
+            topology: Current execution topology
+            resource_manager: Resource manager for usage information
+
+        Returns:
+            Rank (tuple) for operator
+        """
+        pass
+
+    def rank_operators(
+        self,
+        ops: List[PhysicalOperator],
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ) -> List[RankingValue]:
+
+        assert len(ops) > 0
+        return [self.rank_operator(op, topology, resource_manager) for op in ops]
+
+
+class DefaultRanker(Ranker[Tuple[int, int]]):
+    """Ranker implementation."""
+
+    def rank_operator(
+        self,
+        op: PhysicalOperator,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ) -> Tuple[int, int]:
+        """Computes rank for op. *Lower means better rank*
+
+            1. Whether operator's could be throttled (int)
+            2. Operators' object store utilization
+
+        Args:
+            op: Operator to rank
+            topology: Current execution topology
+            resource_manager: Resource manager for usage information
+
+        Returns:
+            Rank (tuple) for operator
+        """
+
+        throttling_disabled = 0 if op.throttling_disabled() else 1
+
+        return (
+            throttling_disabled,
+            resource_manager.get_op_usage(op).object_store_memory,
+        )

python/ray/data/_internal/execution/resource_manager.py

@@ -0,0 +1,942 @@
+import logging
+import math
+import time
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional
+
+from ray._common.utils import env_bool, env_float
+from ray.data._internal.execution import create_resource_allocator
+from ray.data._internal.execution.block_ref_counter import BlockRefCounter
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionOptions,
+    ExecutionResources,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    PhysicalOperator,
+    ReportsExtraResourceUsage,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    AllToAllOperator,
+)
+from ray.data._internal.execution.operators.hash_shuffle import (
+    HashShufflingOperatorBase,
+)
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (  # noqa: E501
+    ShuffleMapOp,
+)
+from ray.data._internal.execution.operators.zip_operator import ZipOperator
+from ray.data._internal.execution.util import memory_string
+from ray.data.context import DataContext
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor_state import OpState, Topology
+
+
+logger = logging.getLogger(__name__)
+
+
+LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE: Optional[bool] = env_bool(
+    "RAY_DATA_DEBUG_RESOURCE_MANAGER", None
+)
+
+# Only warn that the cluster can't run any task once the operator has been starved of
+# its minimum resources for this long. This avoids spurious warnings while the cluster
+# is still scaling up or waiting for a response from the autoscaling coordinator.
+#
+# I arbitrarily chose the default delay.
+STARVATION_WARNING_DELAY_S = env_float("RAY_DATA_STARVATION_WARNING_DELAY_S", 60)
+
+
+# Following list is a list of *blocking* materializing operators, that prevent
+# operators downstream from them from starting execution until these operators
+# finish executing.
+_BLOCKING_MATERIALIZING_OPERATORS = (
+    HashShufflingOperatorBase,
+    AllToAllOperator,
+    ShuffleMapOp,
+    # TODO remove after zip made fully streaming
+    ZipOperator,
+)
+
+
+def terminal_operator_from_topology(topology: "Topology") -> PhysicalOperator:
+    """Return the executor sink: the unique op with no in-DAG downstream consumers.
+
+    ``build_streaming_topology`` is rooted at the same node passed to
+    ``StreamingExecutor``; that root is the only operator whose
+    ``output_dependencies`` is empty.
+    """
+    if not topology:
+        raise ValueError("topology must be non-empty")
+    sinks = [op for op in topology if not op.output_dependencies]
+    if len(sinks) == 1:
+        return sinks[0]
+    if not sinks:
+        raise ValueError(
+            "No terminal operator found in topology (expected exactly one "
+            "operator with empty output_dependencies)"
+        )
+    raise ValueError(
+        "Expected exactly one terminal operator in topology, found "
+        f"{len(sinks)}: {sinks!r}"
+    )
+
+
+class ResourceManager:
+    """A class that manages the resource usage of a streaming executor."""
+
+    # The interval in seconds at which the global resource limits are refreshed.
+    GLOBAL_LIMITS_UPDATE_INTERVAL_S = 1
+
+    # The fraction of the object store capacity that will be used as the default object
+    # store memory limit for the streaming executor,
+    # when `OpResourceAllocator` is enabled.
+    DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION = env_float(
+        "RAY_DATA_OBJECT_STORE_MEMORY_LIMIT_FRACTION", 0.5
+    )
+
+    # The fraction of the object store capacity that will be used as the default object
+    # store memory limit for the streaming executor,
+    # when `OpResourceAllocator` is not enabled.
+    DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION = 0.25
+
+    def __init__(
+        self,
+        topology: "Topology",
+        options: ExecutionOptions,
+        get_total_resources: Callable[[], ExecutionResources],
+        data_context: DataContext,
+        block_ref_counter: BlockRefCounter,
+    ):
+        self._topology = topology
+        self._options = options
+        self._get_total_resources = get_total_resources
+        self._global_limits = ExecutionResources.zero()
+        self._global_limits_last_update_time = 0
+        self._global_usage = ExecutionResources.zero()
+        self._global_running_usage = ExecutionResources.zero()
+        self._global_pending_usage = ExecutionResources.zero()
+        self._op_usages: Dict[PhysicalOperator, ExecutionResources] = {}
+        self._op_running_usages: Dict[PhysicalOperator, ExecutionResources] = {}
+        self._op_pending_usages: Dict[PhysicalOperator, ExecutionResources] = {}
+        # Object store memory usage of pending task outputs (blocks being
+        # generated by running tasks but not yet yielded).
+        self._mem_op_internal: Dict[PhysicalOperator, int] = defaultdict(int)
+        # Object store memory usage of the operator's outputs, including:
+        # internal output queue, external output buffer in OpState, and
+        # downstream operators' input buffers (inqueue + pending task inputs).
+        self._mem_op_outputs: Dict[PhysicalOperator, int] = defaultdict(int)
+
+        # Bytes buffered by external consumers (iterators) consuming Batches
+        # (including the prefetched blocks). For example,
+        # - ds.iter_batches -> one iterator
+        # - streaming_split -> multiple iterators
+        self._external_consumer_bytes: int = 0
+        self._has_external_consumer: bool = False
+
+        # Executor sink (DAG root: unique op with no output_dependencies).
+        # Iterator/streaming_split prefetch bytes are charged on this
+        # operator's output usage.
+        self._output_operator = terminal_operator_from_topology(topology)
+
+        self._block_ref_counter = block_ref_counter
+
+        self._op_resource_allocator: Optional[
+            "OpResourceAllocator"
+        ] = create_resource_allocator(self, data_context)
+
+        self._object_store_memory_limit_fraction = (
+            data_context.override_object_store_memory_limit_fraction
+            if data_context.override_object_store_memory_limit_fraction is not None
+            else (
+                self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION
+                if self.op_resource_allocator_enabled()
+                else self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION
+            )
+        )
+
+    @property
+    def has_external_consumer(self) -> bool:
+        """Return whether there is any external consumer."""
+        return self._has_external_consumer
+
+    def set_external_consumer_bytes(self, num_bytes: int) -> None:
+        """Set the bytes buffered by external consumers."""
+        assert (
+            num_bytes >= 0
+        ), f"external consumer bytes must be non-negative, got {num_bytes}"
+        self._external_consumer_bytes = num_bytes
+        self._has_external_consumer = True
+
+    def get_external_consumer_bytes(self) -> int:
+        """Get the bytes buffered by external consumers."""
+        return self._external_consumer_bytes
+
+    def _estimate_object_store_memory_usage(
+        self, op: "PhysicalOperator", state: "OpState"
+    ) -> int:
+        # Don't count input refs towards dynamic memory usage, as they have been
+        # pre-created already outside this execution.
+        if isinstance(op, InputDataBuffer):
+            if op is self._output_operator:
+                self._mem_op_internal[op] = 0
+                self._mem_op_outputs[op] = self._external_consumer_bytes
+                return self._external_consumer_bytes
+            return 0
+
+        usage = op.estimate_object_store_usage(state)
+        self._mem_op_internal[op] = usage.internal
+        self._mem_op_outputs[op] = usage.outputs
+
+        # Attribute iterator / streaming_split prefetch to the executor sink only.
+        if op is self._output_operator:
+            self._mem_op_outputs[op] += self._external_consumer_bytes
+
+        return self._mem_op_outputs[op] + self._mem_op_internal[op]
+
+    def update_usages(self):
+        """Recalculate resource usages."""
+        # TODO(hchen): This method will be called frequently during the execution loop.
+        # And some computations are redundant. We should either remove redundant
+        # computations or remove this method entirely and compute usages on demand.
+        self._op_usages.clear()
+        self._op_running_usages.clear()
+        self._op_pending_usages.clear()
+
+        # Iterate from last to first operator.
+        for op, state in reversed(self._topology.items()):
+            # Update `self._op_usages`, `self._op_running_usages`,
+            # and `self._op_pending_usages`.
+            op_usage = op.current_logical_usage()
+            op_running_usage = op.running_logical_usage()
+            op_pending_usage = op.pending_logical_usage()
+
+            assert not op_usage.object_store_memory
+            assert not op_running_usage.object_store_memory
+            assert not op_pending_usage.object_store_memory
+
+            used_object_store = self._estimate_object_store_memory_usage(op, state)
+
+            op_usage = op_usage.copy(object_store_memory=used_object_store)
+            op_running_usage = op_running_usage.copy(
+                object_store_memory=used_object_store
+            )
+
+            if isinstance(op, ReportsExtraResourceUsage):
+                op_usage.add(op.extra_resource_usage())
+
+            self._op_usages[op] = op_usage
+            self._op_running_usages[op] = op_running_usage
+            self._op_pending_usages[op] = op_pending_usage
+
+            # Update operator's object store usage, which is used by
+            # DatasetStats and updated on the Ray Data dashboard.
+            op._metrics.obj_store_mem_used = op_usage.object_store_memory
+
+        # Roll the per-op usages up into the global totals in a single pass
+        # each (one allocation per total instead of one per operator).
+        self._global_usage = ExecutionResources.combine_sum(self._op_usages.values())
+        self._global_running_usage = ExecutionResources.combine_sum(
+            self._op_running_usages.values()
+        )
+        self._global_pending_usage = ExecutionResources.combine_sum(
+            self._op_pending_usages.values()
+        )
+
+        if self._op_resource_allocator is not None:
+            self._update_allocated_budgets()
+
+    def _update_allocated_budgets(self):
+        completed_ops_usage = self._get_completed_ops_usage()
+
+        available_limits = (
+            self.get_global_limits()
+            .subtract(completed_ops_usage)
+            .max(ExecutionResources.zero())
+        )
+
+        self._op_resource_allocator.update_budgets(limits=available_limits)
+
+    def get_global_usage(self) -> ExecutionResources:
+        """Return the global resource usage at the current time."""
+        assert (
+            self._global_usage.is_non_negative()
+        ), f"Global usage should be non-negative, got {self._global_usage}"
+        return self._global_usage
+
+    def get_global_running_usage(self) -> ExecutionResources:
+        """Return the global running resource usage at the current time."""
+        return self._global_running_usage
+
+    def get_global_pending_usage(self) -> ExecutionResources:
+        """Return the global pending resource usage at the current time."""
+        return self._global_pending_usage
+
+    def get_global_limits(self) -> ExecutionResources:
+        """Return the global resource limits at the current time.
+
+        This method autodetects any unspecified execution resource limits based on the
+        current cluster size, refreshing these values periodically to support cluster
+        autoscaling.
+        """
+        if (
+            time.time() - self._global_limits_last_update_time
+            < self.GLOBAL_LIMITS_UPDATE_INTERVAL_S
+        ):
+            return self._global_limits
+
+        self._global_limits_last_update_time = time.time()
+        default_limits = self._options.resource_limits
+        exclude = self._options.exclude_resources
+        total_resources = self._get_total_resources()
+        default_mem_fraction = self._object_store_memory_limit_fraction
+        total_resources = total_resources.copy(
+            object_store_memory=total_resources.object_store_memory
+            * default_mem_fraction
+        )
+        # Clamp to non-negative because exclude_resources (e.g., training worker
+        # CPUs) can exceed the total resources reported by the cluster autoscaler,
+        # such as when Ray Train reserves more CPUs than are visible to Ray Data.
+        self._global_limits = (
+            default_limits.min(total_resources)
+            .subtract(exclude)
+            .max(ExecutionResources.zero())
+        )
+        return self._global_limits
+
+    def get_op_usage(
+        self, op: PhysicalOperator, include_ineligible_downstream: bool = False
+    ) -> ExecutionResources:
+        """Return the resource usage of the given operator at the current time."""
+        own_usage = self._op_usages[op]
+
+        if not include_ineligible_downstream:
+            return own_usage
+
+        return own_usage.add(self._get_downstream_ineligible_ops_usage(op))
+
+    def _get_downstream_ineligible_ops_usage(
+        self, op: PhysicalOperator
+    ) -> ExecutionResources:
+        return ExecutionResources.combine_sum(
+            self.get_op_usage(op) for op in self._get_downstream_ineligible_ops(op)
+        )
+
+    def get_mem_op_internal(self, op: PhysicalOperator) -> int:
+        """Return the memory usage of pending task outputs for the given operator."""
+        return self._mem_op_internal[op]
+
+    def get_mem_op_outputs(
+        self,
+        op: PhysicalOperator,
+        include_ineligible_downstream: bool = False,
+    ) -> int:
+        """Return the memory usage of the outputs of the given operator."""
+        # Outputs usage of the current operator.
+        op_outputs_usage = self._mem_op_outputs[op]
+
+        if not include_ineligible_downstream:
+            return op_outputs_usage
+
+        # Also account the downstream ineligible operators' memory usage.
+        return (
+            op_outputs_usage
+            + self._get_downstream_ineligible_ops_usage(op).object_store_memory
+        )
+
+    def get_op_usage_str(self, op: PhysicalOperator, *, verbose: bool) -> str:
+        """Return a human-readable string representation of the resource usage of
+        the given operator."""
+        # Handle case where operator is not in _op_running_usages dict
+        if op not in self._op_running_usages:
+            usage_str = "n/a"
+        else:
+            usage_str = f"{self._op_running_usages[op].cpu:.1f} CPU"
+            if self._op_running_usages[op].memory:
+                usage_str += f", {self._op_running_usages[op].memory_str()} memory"
+            if self._op_running_usages[op].gpu:
+                usage_str += f", {self._op_running_usages[op].gpu:.1f} GPU"
+            usage_str += f", {self._op_running_usages[op].object_store_memory_str()} object store"
+
+        # NOTE: Config can override requested verbosity level
+        if LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE is not None:
+            verbose = LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE
+
+        if verbose:
+            usage_str += (
+                f" (in={memory_string(self.get_mem_op_internal(op))},"
+                f"out={memory_string(self.get_mem_op_outputs(op))}"
+            )
+            # External-consumer bytes (iterator / streaming_split prefetch) are
+            # only attached to the output operator. Surface them in its line so
+            # users can see how much of `out` is held by the downstream iterator
+            # vs. the operator's own output queues.
+            if op is self._output_operator and self._has_external_consumer:
+                usage_str += (
+                    f",external_consumer="
+                    f"{memory_string(self._external_consumer_bytes)}"
+                )
+            usage_str += ")"
+            if self._op_resource_allocator is not None:
+                allocation = self._op_resource_allocator.get_allocation(op)
+                if allocation:
+                    usage_str += f", alloc=(cpu={allocation.cpu:.1f}"
+                    usage_str += f",mem={allocation.memory_str()}"
+                    usage_str += f",gpu={allocation.gpu:.1f}"
+                    usage_str += f",obj_store={allocation.object_store_memory_str()})"
+
+                budget = self._op_resource_allocator.get_budget(op)
+                if budget:
+                    usage_str += f", budget=(cpu={budget.cpu:.1f}"
+                    usage_str += f",mem={budget.memory_str()}"
+                    usage_str += f",gpu={budget.gpu:.1f}"
+                    usage_str += f",obj_store={budget.object_store_memory_str()}"
+
+                    # Remaining memory budget for producing new task outputs.
+                    if isinstance(
+                        self._op_resource_allocator, ReservationOpResourceAllocator
+                    ):
+                        reserved_for_output = memory_string(
+                            self._op_resource_allocator._output_budgets.get(op, 0)
+                        )
+                        usage_str += f",out={reserved_for_output})"
+
+        return usage_str
+
+    def op_resource_allocator_enabled(self) -> bool:
+        """Return whether OpResourceAllocator is enabled."""
+        return self._op_resource_allocator is not None
+
+    @property
+    def op_resource_allocator(self) -> "OpResourceAllocator":
+        """Return the OpResourceAllocator."""
+        assert self._op_resource_allocator is not None
+        return self._op_resource_allocator
+
+    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        """Return the budget for the given operator, or None if the operator
+        has unlimited budget."""
+        if self._op_resource_allocator is None:
+            return None
+        return self._op_resource_allocator.get_budget(op)
+
+    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        """Return the allocation of the given operator, or None if the operator
+        doesn't have a designated allocation."""
+        if self._op_resource_allocator is None:
+            return None
+        return self._op_resource_allocator.get_allocation(op)
+
+    def is_op_eligible(self, op: PhysicalOperator) -> bool:
+        """Whether the op is eligible for memory reservation."""
+        return (
+            not op.throttling_disabled()
+            # As long as the op has finished execution, even if there are still
+            # non-taken outputs, we don't need to allocate resources for it.
+            and not op.has_execution_finished()
+        )
+
+    def _get_downstream_ineligible_ops(
+        self, op: PhysicalOperator
+    ) -> Iterable[PhysicalOperator]:
+        """Get the downstream ineligible operators of the given operator.
+
+        E.g.,
+          - "cur_map->downstream_map" will return an empty list.
+          - "cur_map->limit1->limit2->downstream_map" will return [limit1, limit2].
+        """
+        for next_op in op.output_dependencies:
+            if not self.is_op_eligible(next_op):
+                yield next_op
+                yield from self._get_downstream_ineligible_ops(next_op)
+
+    def get_downstream_eligible_ops(
+        self, op: PhysicalOperator
+    ) -> Iterable[PhysicalOperator]:
+        """Get the downstream eligible operators of the given operator, ignoring
+        intermediate ineligible operators.
+
+        E.g.,
+          - "cur_map->downstream_map" will return [downstream_map].
+          - "cur_map->limit1->limit2->downstream_map" will return [downstream_map].
+        """
+        for next_op in op.output_dependencies:
+            if self.is_op_eligible(next_op):
+                yield next_op
+            else:
+                yield from self.get_downstream_eligible_ops(next_op)
+
+    def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]:
+        if self._op_resource_allocator is not None:
+            return self._op_resource_allocator.max_task_output_bytes_to_read(op)
+        return None
+
+    def _get_completed_ops_usage(self) -> ExecutionResources:
+        """
+        Resource reservation is based on the number of eligible operators.
+        However, there might be completed operators that still have blocks in their output queue, which we need to exclude them from the reservation.
+        And we also need to exclude the downstream ineligible operators.
+
+        E.g., for the following pipeline:
+        ```
+        map1 (completed, but still has blocks in its output queue) -> limit1 (ineligible, not completed) -> map2 (not completed) -> limit2 -> map3
+        ```
+
+        The reservation is based on the number of eligible operators (map2 and map3), but we need to exclude map1 and limit1 from the reservation.
+        """
+
+        last_completed_ops = []
+        ops_to_exclude = []
+        # Traverse operator tree collecting all operators that have already finished
+        for op in self._topology:
+            if not op.has_execution_finished():
+                for dep in op.input_dependencies:
+                    if dep.has_execution_finished():
+                        last_completed_ops.append(dep)
+
+        # In addition to completed operators,
+        # filter out downstream ineligible operators since they are omitted from reservation calculations.
+        for op in last_completed_ops:
+            ops_to_exclude.extend(list(self._get_downstream_ineligible_ops(op)))
+            ops_to_exclude.append(op)
+
+        completed_ops = list(set(ops_to_exclude))
+        completed_ops_usage = ExecutionResources.combine_sum(
+            self.get_op_usage(op) for op in completed_ops
+        )
+
+        return completed_ops_usage
+
+    def _is_blocking_materializing_op(self, op: PhysicalOperator) -> bool:
+        """This method checks whether either
+
+        1. Operator itself or
+        2. One of operator's immediate *ineligible* downstream dependencies
+
+        Are blocking, materializing operators.
+
+        NOTE: That downstream ineligible operators are considered an "extension" of their
+              first preceding eligible operator from resource allocation standpoint.
+        """
+
+        # Check if Op itself is a blocking, materializing operator
+        if isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS):
+            return True
+
+        # Check if any of its direct *ineligible* downstream dependencies are
+        # blocking, materializing operators.
+        #
+        # NOTE: We only check ineligible downstream deps, since eligible downstream
+        #       deps will have their own allocation that is adjusted appropriately
+        return any(
+            isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS)
+            for op in self._get_downstream_ineligible_ops(op)
+        )
+
+
+def _get_first_pending_materializing_op(topology: "Topology") -> int:
+    for idx, op in enumerate(topology):
+        if isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS) and not op.has_completed():
+            return idx
+
+    return -1
+
+
+class OpResourceAllocator(ABC):
+    """An interface for dynamic operator resource allocation.
+
+    This interface allows dynamically allocating available resources to each operator,
+    limiting how many tasks each operator can submit, and how much data each operator
+    can read from its running tasks.
+    """
+
+    def __init__(self, resource_manager: "ResourceManager"):
+        self._resource_manager = resource_manager
+        self._topology = resource_manager._topology
+
+    @abstractmethod
+    def update_budgets(
+        self,
+        *,
+        limits: ExecutionResources,
+    ):
+        """Callback to update resource usages."""
+        ...
+
+    @abstractmethod
+    def can_submit_new_task(self, op: PhysicalOperator) -> bool:
+        """Return whether the given operator can submit a new task."""
+        ...
+
+    @abstractmethod
+    def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]:
+        """Return the maximum bytes of pending task outputs can be read for
+        the given operator. None means no limit."""
+        ...
+
+    @abstractmethod
+    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        """Returns the budget for the given operator or `None` if the operator
+        has unlimited budget. Operator's budget is defined as:
+
+            Budget = Allocation - Usage
+        """
+        ...
+
+    @abstractmethod
+    def get_output_budget(self, op: PhysicalOperator) -> Optional[int]:
+        """Returns the budget for operator's outputs (in object store bytes) or
+        `None` if there's no limit.
+        """
+        ...
+
+    @abstractmethod
+    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        """Returns allocation for the given operator or `None` if operator's
+        allocation is unlimited."""
+        ...
+
+    def _get_eligible_ops(self) -> List[PhysicalOperator]:
+        """Returns a list of operators eligible for allocation.
+
+        Only operators upstream of the first non-completed *materializing* Operator,
+        (like shuffle, etc) are able to receive inputs and start execution.
+
+        Therefore only these operators are eligible for resource allocation.
+        """
+        first_pending_materializing_op_idx = _get_first_pending_materializing_op(
+            self._topology
+        )
+        return [
+            op
+            for idx, op in enumerate(self._topology)
+            if self._is_op_eligible(op)
+            and (
+                first_pending_materializing_op_idx == -1
+                or idx <= first_pending_materializing_op_idx
+            )
+        ]
+
+    @staticmethod
+    def _is_op_eligible(op: PhysicalOperator) -> bool:
+        """Whether the op is eligible for memory reservation."""
+        return (
+            not op.throttling_disabled()
+            # As long as the op has finished execution, even if there are still
+            # non-taken outputs, we don't need to allocate resources for it.
+            and not op.has_execution_finished()
+        )
+
+
+class ReservationOpResourceAllocator(OpResourceAllocator):
+    """An OpResourceAllocator implementation that reserves resources for each operator.
+
+    This class reserves memory and CPU resources for eligible operators, and considers
+    runtime resource usages to limit the resources that each operator can use.
+
+    It works in the following way:
+    1. An operator is eligible for resource reservation, if it has enabled throttling
+       and hasn't completed. Ineligible operators are not throttled, but
+       their usage will be accounted for their upstream eligible operators. E.g., for
+       such a dataset "map1->limit->map2->streaming_split", we'll treat "map1->limit" as
+       a group and "map2->streaming_split" as another group.
+    2. For each eligible operator, we reserve `reservation_ratio * global_resources /
+        num_eligible_ops` resources, half of which is reserved only for the operator
+        outputs, excluding pending task outputs.
+    3. Non-reserved resources are shared among all operators.
+    4. In each scheduling iteration, each eligible operator will get "remaining of their
+       own reserved resources" + "remaining of shared resources / num_eligible_ops"
+       resources.
+
+    The `reservation_ratio` is set to 50% by default. Users can tune this value to
+    adjust how aggressive or conservative the resource allocation is. A higher value
+    will make the resource allocation more even, but may lead to underutilization and
+    worse performance. And vice versa.
+    """
+
+    def __init__(self, resource_manager: ResourceManager, reservation_ratio: float):
+        super().__init__(resource_manager)
+
+        self._reservation_ratio = reservation_ratio
+        assert 0.0 <= self._reservation_ratio <= 1.0
+        # Per-op reserved resources, excluding `_reserved_for_op_outputs`.
+        self._op_reserved: Dict[PhysicalOperator, ExecutionResources] = {}
+        # Memory reserved exclusively for the outputs of each operator.
+        # "Op outputs" refer to blocks that have been taken out of an operator,
+        # i.e., `RessourceManager._mem_op_outputs`.
+        #
+        # Note, if we don't reserve memory for op outputs, all the budget may be used by
+        # the pending task outputs, and/or op's internal output buffers (the latter can
+        # happen when `preserve_order=True`).
+        # Then we'll have no budget to pull blocks from the op.
+        self._reserved_for_op_outputs: Dict[PhysicalOperator, float] = {}
+        # Total shared resources.
+        self._total_shared = ExecutionResources.zero()
+        # Resource budgets for each operator, excluding `_reserved_for_op_outputs`.
+        self._op_budgets: Dict[PhysicalOperator, ExecutionResources] = {}
+        # Remaining memory budget for generating new task outputs, per operator.
+        self._output_budgets: Dict[PhysicalOperator, float] = {}
+        # Whether each operator has reserved the minimum resources to run
+        # at least one task.
+        # This is used to avoid edge cases where the entire resource limits are not
+        # enough to run one task of each op.
+        # See `test_no_deadlock_on_small_cluster_resources` as an example.
+        self._reserved_min_resources: Dict[PhysicalOperator, bool] = {}
+        # `time.monotonic()` timestamp at which each operator most recently became
+        # starved of its minimum resources, or None if it currently has them.
+        self._op_starved_since: Dict[PhysicalOperator, Optional[float]] = {}
+
+    def _update_reservation(self, limits: ExecutionResources):
+        eligible_ops = self._get_eligible_ops()
+
+        self._op_reserved.clear()
+        self._reserved_for_op_outputs.clear()
+        self._reserved_min_resources.clear()
+
+        if len(eligible_ops) == 0:
+            return
+
+        remaining = limits.copy()
+
+        # Reserve `reservation_ratio * global_limits / num_ops` resources for each
+        # operator.
+        default_reserved = limits.scale(self._reservation_ratio / (len(eligible_ops)))
+        for index, op in enumerate(eligible_ops):
+            # Reserve at least half of the default reserved resources for the outputs.
+            # This makes sure that we will have enough budget to pull blocks from the
+            # op.
+            reserved_for_outputs = ExecutionResources(
+                0, 0, max(default_reserved.object_store_memory / 2, 1)
+            )
+
+            reserved_for_tasks = default_reserved.subtract(reserved_for_outputs)
+
+            min_resource_usage, max_resource_usage = op.min_max_resource_requirements()
+
+            if min_resource_usage is not None:
+                reserved_for_tasks = reserved_for_tasks.max(min_resource_usage)
+            if max_resource_usage is not None:
+                reserved_for_tasks = reserved_for_tasks.min(max_resource_usage)
+
+            # Check if the remaining resources are enough for both reserved_for_tasks
+            # and reserved_for_outputs. Note, we only consider CPU and GPU, but not
+            # object_store_memory, because object_store_memory can be oversubscribed,
+            # but CPU/GPU cannot.
+            if reserved_for_tasks.add(reserved_for_outputs).satisfies_limit(
+                remaining, ignore_object_store_memory=True
+            ):
+                self._reserved_min_resources[op] = True
+                self._op_starved_since[op] = None
+            else:
+                self._reserved_min_resources[op] = False
+                if self._op_starved_since.get(op) is None:
+                    self._op_starved_since[op] = time.monotonic()
+
+                # If the remaining resources are not enough to reserve the minimum
+                # resources for this operator, we'll only reserve the minimum object
+                # store memory, but not the CPU and GPU resources.
+                # Because Ray Core doesn't allow CPU/GPU resources to be oversubscribed.
+                # NOTE: we prioritize upstream operators for minimum resource reservation.
+                # ops. It's fine that downstream ops don't get the minimum reservation,
+                # because they can wait for upstream ops to finish and release resources.
+                reserved_for_tasks = ExecutionResources(
+                    0, 0, min_resource_usage.object_store_memory
+                )
+
+            # Log a warning if even the first operator cannot reserve the minimum
+            # resources.
+            if index == 0:
+                self._warn_if_op_starved_too_long(op)
+
+            self._op_reserved[op] = reserved_for_tasks
+            self._reserved_for_op_outputs[op] = reserved_for_outputs.object_store_memory
+
+            op_total_reserved = reserved_for_tasks.add(reserved_for_outputs)
+            remaining = remaining.subtract(op_total_reserved)
+            remaining = remaining.max(ExecutionResources.zero())
+
+        self._total_shared = remaining
+
+    def _warn_if_op_starved_too_long(self, op: PhysicalOperator) -> None:
+        # The operator isn't starved. Return early.
+        if self._op_starved_since.get(op) is None:
+            return
+
+        op_starved_duration = time.monotonic() - self._op_starved_since[op]
+        if (
+            op_starved_duration >= STARVATION_WARNING_DELAY_S
+            # Add `id(self)` to the log_once key so that it will be logged once per
+            # execution.
+            and log_once(f"starvation_warning_{id(self)}")
+        ):
+            logger.warning(
+                f"Cluster resources are not enough to run any task from {op}."
+                " The job may hang forever unless the cluster scales up."
+            )
+
+    def can_submit_new_task(self, op: PhysicalOperator) -> bool:
+        """Return whether the given operator can submit a new task based on budget."""
+        budget = self.get_budget(op)
+
+        if budget is None:
+            return True
+
+        return (
+            op.incremental_resource_usage().satisfies_limit(budget)
+            and
+            # Avoid scheduling if there's no more Object Store budget (for
+            # task outputs)
+            budget.object_store_memory
+            >= (op.metrics.obj_store_mem_max_pending_output_per_task or 0)
+        )
+
+    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        return self._op_budgets.get(op)
+
+    def get_output_budget(self, op: PhysicalOperator) -> Optional[int]:
+        return self._output_budgets.get(op)
+
+    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        budget = self.get_budget(op)
+        if budget is None:
+            return None
+        return budget.add(self._resource_manager.get_op_usage(op))
+
+    def _get_total_reserved(self, op: PhysicalOperator) -> ExecutionResources:
+        """Get total reserved resources for an operator, including outputs reservation."""
+        op_reserved = self._op_reserved[op]
+        reserved_for_outputs = self._reserved_for_op_outputs[op]
+        return op_reserved.copy(
+            object_store_memory=op_reserved.object_store_memory + reserved_for_outputs
+        )
+
+    def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]:
+        if op not in self._op_budgets:
+            return None
+
+        res = self._op_budgets[op].object_store_memory
+        # Add the remaining of `_reserved_for_op_outputs`.
+        op_outputs_usage = self._resource_manager.get_mem_op_outputs(
+            op, include_ineligible_downstream=True
+        )
+
+        res += max(self._reserved_for_op_outputs[op] - op_outputs_usage, 0)
+        if math.isinf(res):
+            self._output_budgets[op] = res
+            return None
+
+        res = int(res)
+        assert res >= 0
+        self._output_budgets[op] = res
+        return res
+
+    def update_budgets(
+        self,
+        *,
+        limits: ExecutionResources,
+    ):
+        # Remaining resources to be distributed across operators
+        remaining_shared = self._update_reservation(limits)
+
+        self._op_budgets.clear()
+        eligible_ops = self._get_eligible_ops()
+        if len(eligible_ops) == 0:
+            return
+
+        # Remaining of shared resources.
+        remaining_shared = self._total_shared
+        for op in eligible_ops:
+            # Calculate the memory usage of the operator.
+            op_mem_usage = 0
+            # Add the memory usage of the operator itself,
+            # excluding `_reserved_for_op_outputs`.
+            op_mem_usage += self._resource_manager.get_mem_op_internal(op)
+            # Add the portion of op outputs usage that has
+            # exceeded `_reserved_for_op_outputs`.
+            op_outputs_usage = self._resource_manager.get_mem_op_outputs(
+                op, include_ineligible_downstream=True
+            )
+            op_mem_usage += max(op_outputs_usage - self._reserved_for_op_outputs[op], 0)
+
+            op_usage = self._resource_manager.get_op_usage(op).copy(
+                object_store_memory=op_mem_usage
+            )
+
+            op_reserved = self._op_reserved[op]
+            # How much of the reserved resources are remaining.
+            op_reserved_remaining = op_reserved.subtract(op_usage).max(
+                ExecutionResources.zero()
+            )
+
+            self._op_budgets[op] = op_reserved_remaining
+            # How much of the reserved resources are exceeded.
+            # If exceeded, we need to subtract from the remaining shared resources.
+            op_reserved_exceeded = op_usage.subtract(op_reserved).max(
+                ExecutionResources.zero()
+            )
+            remaining_shared = remaining_shared.subtract(op_reserved_exceeded)
+
+        remaining_shared = remaining_shared.max(ExecutionResources.zero())
+
+        # Allocate the remaining shared resources to each operator.
+        for i, op in enumerate(reversed(eligible_ops)):
+            # By default, divide the remaining shared resources equally.
+            op_shared = remaining_shared.scale(1.0 / (len(eligible_ops) - i))
+            # But if the op's budget is less than `min_scheduling_resources`,
+            # it will be useless. So we'll let the downstream operator
+            # borrow some resources from the upstream operator, if remaining_shared
+            # is still enough.
+            to_borrow = (
+                op.min_scheduling_resources()
+                .subtract(self._op_budgets[op].add(op_shared))
+                .max(ExecutionResources.zero())
+            )
+            if not to_borrow.is_zero() and op_shared.add(to_borrow).satisfies_limit(
+                remaining_shared
+            ):
+                op_shared = op_shared.add(to_borrow)
+
+            # Cap op_shared so that total allocation doesn't exceed max_resource_usage.
+            # Total allocation = max(total_reserved, op_usage) + op_shared
+            # This ensures excess resources stay in remaining_shared for other operators.
+            _, max_resource_usage = op.min_max_resource_requirements()
+            if max_resource_usage != ExecutionResources.inf():
+                total_reserved = self._get_total_reserved(op)
+                op_usage = self._resource_manager.get_op_usage(op)
+                current_allocation = total_reserved.max(op_usage)
+                max_shared = max_resource_usage.subtract(current_allocation).max(
+                    ExecutionResources.zero()
+                )
+                op_shared = op_shared.min(max_shared)
+
+            remaining_shared = remaining_shared.subtract(op_shared)
+            assert remaining_shared.is_non_negative(), (
+                remaining_shared,
+                op,
+                op_shared,
+                to_borrow,
+            )
+
+            self._op_budgets[op] = self._op_budgets[op].add(op_shared)
+
+        # Give any remaining shared resources to the most downstream uncapped op.
+        # This can happen when some ops have their shared allocation capped.
+        if eligible_ops and not remaining_shared.is_zero():
+            for op in reversed(eligible_ops):
+                _, max_resource_usage = op.min_max_resource_requirements()
+                if max_resource_usage == ExecutionResources.inf():
+                    self._op_budgets[op] = self._op_budgets[op].add(remaining_shared)
+                    break
+
+        # A materializing operator like `AllToAllOperator` waits for all its input
+        # operator's outputs before processing data. This often forces the input
+        # operator to exceed its object store memory budget. To prevent deadlock, we
+        # disable object store memory backpressure for the input operator.
+        for op in eligible_ops:
+            if self._resource_manager._is_blocking_materializing_op(op):
+                self._op_budgets[op] = self._op_budgets[op].copy(
+                    object_store_memory=float("inf")
+                )

python/ray/data/_internal/execution/streaming_executor.py

@@ -0,0 +1,796 @@
+import logging
+import os
+import threading
+import time
+import typing
+from typing import Dict, List, Optional, Tuple
+
+from ray.data._internal.actor_autoscaler import (
+    create_actor_autoscaler,
+)
+from ray.data._internal.cluster_autoscaler import create_cluster_autoscaler
+from ray.data._internal.execution import create_ranker
+from ray.data._internal.execution.backpressure_policy import (
+    BackpressurePolicy,
+    get_backpressure_policies,
+)
+from ray.data._internal.execution.block_ref_counter import BlockRefCounter
+from ray.data._internal.execution.dataset_state import DatasetState
+from ray.data._internal.execution.execution_callback import ExecutionCallback
+from ray.data._internal.execution.interfaces import (
+    Executor,
+    OutputIterator,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.metadata_fetcher import make_metadata_fetcher
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+)
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.resource_manager import (
+    ResourceManager,
+)
+from ray.data._internal.execution.streaming_executor_state import (
+    OpState,
+    OutputBackpressureGuard,
+    Topology,
+    build_streaming_topology,
+    format_op_state_summary,
+    process_completed_tasks,
+    select_operator_to_run,
+    update_operator_states,
+)
+from ray.data._internal.logging import (
+    get_log_directory,
+    register_dataset_logger,
+    unregister_dataset_logger,
+)
+from ray.data._internal.metadata_exporter import (
+    Topology as TopologyMetadata,
+    sanitize_for_struct,
+)
+from ray.data._internal.operator_schema_exporter import (
+    OperatorSchema,
+    get_operator_schema_exporter,
+)
+from ray.data._internal.progress import get_progress_manager
+from ray.data._internal.stats import DatasetStats, Timer, _StatsManager
+from ray.data.context import OK_PREFIX, WARN_PREFIX, DataContext
+from ray.util.debug import log_once
+from ray.util.metrics import Gauge
+
+if typing.TYPE_CHECKING:
+    from ray.data._internal.issue_detection.issue_detector_manager import (
+        IssueDetectorManager,
+    )
+    from ray.data._internal.progress.base_progress import BaseExecutionProgressManager
+    from ray.data.block import Schema
+
+logger = logging.getLogger(__name__)
+
+# Interval for logging execution progress updates and operator metrics.
+DEBUG_LOG_INTERVAL_SECONDS = 5
+
+# Maximum string/sequence length for DataContext logging. Set high to avoid truncation
+# while still protecting against pathological cases.
+DATA_CONTEXT_LOG_TRUNCATE_LENGTH = 10000
+
+# Visible for testing.
+_num_shutdown = 0
+
+
+# Extra environment variables to log that don't start with RAY_DATA.
+_EXTRA_ENV_VARS_TO_LOG = (
+    # We historically recommended users configure this value. If a Ray Data job uses
+    # more object store memory than expected, it's worth checking how this environment
+    # variable has been configured.
+    "RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION",
+)
+
+
+def _log_ray_data_env_vars() -> None:
+    env_vars = {
+        k: v
+        for k, v in os.environ.items()
+        if k.startswith("RAY_DATA") or k in _EXTRA_ENV_VARS_TO_LOG
+    }
+    if env_vars:
+        formatted = ", ".join(f"{k}={v}" for k, v in sorted(env_vars.items()))
+        logger.debug(f"RAY_DATA environment variables: {formatted}")
+    else:
+        logger.debug("No RAY_DATA environment variables set.")
+
+
+class StreamingExecutor(Executor, threading.Thread):
+    """A streaming Dataset executor.
+
+    This implementation executes Dataset DAGs in a fully streamed way. It runs
+    by setting up the operator topology, and then routing blocks through operators in
+    a way that maximizes throughput under resource constraints.
+    """
+
+    UPDATE_METRICS_INTERVAL_S: float = 5.0
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        dataset_id: str = "unknown_dataset",
+    ):
+        self._data_context = data_context
+        self._ranker = create_ranker()
+        self._start_time: Optional[float] = None
+        self._initial_stats: Optional[DatasetStats] = None
+        self._final_stats: Optional[DatasetStats] = None
+        self._progress_manager: Optional["BaseExecutionProgressManager"] = None
+        self._callbacks: List["ExecutionCallback"] = []
+
+        # The executor can be shutdown while still running.
+        self._shutdown_lock = threading.RLock()
+        self._execution_started = False
+        self._shutdown = False
+
+        # Internal execution state shared across thread boundaries. We run the control
+        # loop on a separate thread so that it doesn't become stalled between
+        # generator `yield`s.
+        self._topology: Optional[Topology] = None
+        self._output_node: Optional[Tuple[PhysicalOperator, OpState]] = None
+        self._backpressure_policies: List[BackpressurePolicy] = []
+        self._op_schema: Dict[PhysicalOperator, Schema] = {}
+
+        self._dataset_id = dataset_id
+        # Set by IssueDetectionExecutionCallback when issue detection is registered;
+        # otherwise remains None. Access via the issue_detector_manager property.
+        self._issue_detector_manager: Optional["IssueDetectorManager"] = None
+        # Stores if an operator is completed,
+        # used for marking when an op has just completed.
+        self._has_op_completed: Optional[Dict[PhysicalOperator, bool]] = None
+        self._max_errored_blocks = self._data_context.max_errored_blocks
+        self._num_errored_blocks = 0
+
+        self._last_debug_log_time = 0
+
+        self._data_context.set_dataset_logger_id(
+            register_dataset_logger(self._dataset_id)
+        )
+
+        # This stores the last time we updated the metrics.
+        # This allows us to update metrics on some interval,
+        # by comparing it with the current timestamp.
+        self._metrics_last_updated: float = 0.0
+
+        self._sched_loop_duration_s = Gauge(
+            "data_sched_loop_duration_s",
+            description="Duration of the scheduling loop in seconds",
+            tag_keys=("dataset",),
+        )
+
+        # Resolves pulled (block_ref, meta_ref) pairs into emitted RefBundles.
+        # The threaded fetcher (default) fetches metadata on a background thread
+        # so the scheduling loop never blocks on ``ray.get(meta_refs)``; the
+        # inline fetcher reproduces the synchronous, master-identical path.
+        self._metadata_fetcher = make_metadata_fetcher()
+
+        Executor.__init__(self, self._data_context.execution_options)
+        thread_name = f"StreamingExecutor-{self._dataset_id}"
+        threading.Thread.__init__(self, daemon=True, name=thread_name)
+
+    @property
+    def issue_detector_manager(self) -> Optional["IssueDetectorManager"]:
+        """The issue detector manager, or None if issue detection isn't registered."""
+        return self._issue_detector_manager
+
+    def execute(
+        self,
+        dag: PhysicalOperator,
+        initial_stats: Optional[DatasetStats] = None,
+        callbacks: Optional[List] = None,
+    ) -> OutputIterator:
+        """Executes the DAG using a streaming execution strategy.
+
+        We take an event-loop approach to scheduling. We block on the next scheduling
+        event using `ray.wait`, updating operator state and dispatching new tasks.
+        """
+
+        if callbacks is not None:
+            self._callbacks = callbacks
+        else:
+            self._callbacks = []
+
+        self._initial_stats = initial_stats
+        self._start_time = time.perf_counter()
+
+        if logger.isEnabledFor(logging.DEBUG):
+            _log_ray_data_env_vars()
+
+        if not isinstance(dag, InputDataBuffer):
+            if self._data_context.print_on_execution_start:
+                message = f"Starting execution of Dataset {self._dataset_id}."
+                log_path = get_log_directory()
+                if log_path is not None:
+                    message += f" Full logs are in {log_path}"
+                logger.info(message)
+                logger.info(
+                    f"Execution plan of Dataset {self._dataset_id}: {dag.dag_str}"
+                )
+
+            # Log the full DataContext for traceability
+            if logger.isEnabledFor(logging.DEBUG) and log_once(
+                f"ray_data_log_context_{self._dataset_id}"
+            ):
+                logger.debug(
+                    f"Data Context for dataset {self._dataset_id}:\n%s",
+                    sanitize_for_struct(
+                        self._data_context,
+                        truncate_length=DATA_CONTEXT_LOG_TRUNCATE_LENGTH,
+                    ),
+                )
+
+        # Setup the streaming DAG topology and start the runner thread.
+        self._block_ref_counter = BlockRefCounter()
+        self._topology = build_streaming_topology(
+            dag, self._options, self._block_ref_counter
+        )
+
+        self._resource_manager = ResourceManager(
+            self._topology,
+            self._options,
+            lambda: self._cluster_autoscaler.get_total_resources(),
+            self._data_context,
+            self._block_ref_counter,
+        )
+
+        # Constructed once per executor (not per scheduling iteration) so the
+        # guard's idle-detection state accumulates across scheduling iterations.
+        self._output_backpressure_guard = OutputBackpressureGuard(
+            self._topology, self._resource_manager
+        )
+
+        # Setup progress manager
+        self._progress_manager = get_progress_manager(
+            self._data_context,
+            self._dataset_id,
+            self._topology,
+            self._options.verbose_progress,
+        )
+        self._progress_manager.start()
+
+        self._backpressure_policies = get_backpressure_policies(
+            self._data_context, self._topology, self._resource_manager
+        )
+        self._cluster_autoscaler = create_cluster_autoscaler(
+            self._topology,
+            self._resource_manager,
+            self._data_context,
+            execution_id=self._dataset_id,
+        )
+        self._actor_autoscaler = create_actor_autoscaler(
+            self._topology,
+            self._resource_manager,
+            config=self._data_context.autoscaling_config,
+        )
+
+        self._has_op_completed = dict.fromkeys(self._topology, False)
+
+        self._output_node = dag, self._topology[dag]
+
+        op_to_id = {
+            op: self._get_operator_id(op, i) for i, op in enumerate(self._topology)
+        }
+        _StatsManager.register_dataset_to_stats_actor(
+            self._dataset_id,
+            self._get_operator_tags(),
+            TopologyMetadata.create_topology_metadata(dag, op_to_id),
+            self._data_context,
+        )
+        for callback in self._callbacks:
+            callback.before_execution_starts(self)
+
+        self.start()
+        self._execution_started = True
+
+        return _ClosingIterator(self)
+
+    def __del__(self):
+        # NOTE: Upon garbage-collection we're allowing running tasks
+        #       to be terminated asynchronously (ie avoid unnecessary
+        #       synchronization on their completion)
+        self.shutdown(force=False)
+
+    def shutdown(self, force: bool, exception: Optional[Exception] = None):
+        global _num_shutdown
+
+        with self._shutdown_lock:
+            if not self._execution_started or self._shutdown:
+                return
+
+            start = time.perf_counter()
+
+            status_detail = (
+                f"failed with {exception}" if exception else "completed successfully"
+            )
+
+            logger.debug(
+                f"Shutting down executor for dataset {self._dataset_id} "
+                f"({status_detail})"
+            )
+
+            _num_shutdown += 1
+            self._shutdown = True
+            # Give the scheduling loop some time to finish processing.
+            self.join(timeout=2.0)
+            # Stop the metadata fetcher (after the loop thread that feeds it has
+            # been joined). No-op for the inline fetcher.
+            self._metadata_fetcher.stop()
+            self._update_stats_metrics(
+                state=DatasetState.FINISHED.name
+                if exception is None
+                else DatasetState.FAILED.name,
+                force_update=True,
+            )
+            # Freeze the stats and save it.
+            self._final_stats = self._generate_stats()
+            stats_summary_string = self._final_stats.to_summary().to_string(
+                include_parent=False
+            )
+            # Reset the scheduling loop duration gauge + resource manager budgets/usages.
+            self._resource_manager.update_usages()
+            self.update_metrics(0)
+            if self._data_context.enable_auto_log_stats:
+                logger.info(stats_summary_string)
+            # Close the progress manager with a finishing message.
+            if exception is None:
+                desc = (
+                    f"{OK_PREFIX} Dataset {self._dataset_id} execution finished in "
+                    f"{self._final_stats.time_total_s:.2f} seconds"
+                )
+            else:
+                desc = f"{WARN_PREFIX} Dataset {self._dataset_id} execution failed"
+
+            self._progress_manager.close_with_finishing_description(
+                desc, exception is None
+            )
+            logger.info(desc)
+
+            timer = Timer()
+
+            for op in self._topology.keys():
+                op.shutdown(timer, force=force)
+
+            self._clear_topology_queues_post_shutdown(force, exception)
+            # Queues have been drained; any remaining Ray Core callbacks that fire
+            # after this point should be no-ops.
+            self._block_ref_counter.clear()
+
+            min_ = round(timer.min(), 3)
+            max_ = round(timer.max(), 3)
+            total = round(timer.get(), 3)
+            logger.debug(
+                f"Shut down operator hierarchy for dataset {self._dataset_id}"
+                f" (min/max/total={min_}/{max_}/{total}s)"
+            )
+
+            if exception is None:
+                for callback in self._callbacks:
+                    callback.after_execution_succeeds(self)
+            else:
+                for callback in self._callbacks:
+                    callback.after_execution_fails(self, exception)
+
+            self._cluster_autoscaler.on_executor_shutdown()
+
+            dur = time.perf_counter() - start
+
+            logger.debug(
+                f"Shut down executor for dataset {self._dataset_id} "
+                f"(took {round(dur, 3)}s)"
+            )
+
+            # Unregister should be called after all operators are shut down to
+            # capture as many logs as possible.
+            self._data_context.set_dataset_logger_id(
+                unregister_dataset_logger(self._dataset_id)
+            )
+
+    def _clear_topology_queues_post_shutdown(
+        self, force: bool, exception: Optional[Exception] = None
+    ) -> None:
+        """Drain topology queues after operator shutdown (releases block refs)."""
+        for op, state in self._topology.items():
+            if isinstance(op, InternalQueueOperatorMixin):
+                op.clear_internal_input_queue()
+                op.clear_internal_output_queue()
+            # Input queues alias upstream output queues; clears the DAG except the sink.
+            for inqueue in state.input_queues:
+                inqueue.clear()
+
+        output_op, _ = self._output_node
+        # Clear sink output unless cooperative multi-split success (splits may still read).
+        is_live_multi_split_sink = (
+            output_op.num_output_splits() > 1 and not force and exception is None
+        )
+        if not is_live_multi_split_sink:
+            self._topology[output_op].output_queue.clear()
+
+    def run(self):
+        """Run the control loop in a helper thread.
+
+        Results are returned via the output node's outqueue.
+        """
+        exc: Optional[Exception] = None
+        self._metadata_fetcher.start()
+        try:
+            # Run scheduling loop until complete.
+            while True:
+                # Use `perf_counter` rather than `process_time` to ensure we include
+                # time spent on IO, like RPCs to Ray Core.
+                t_start = time.perf_counter()
+                continue_sched = self._scheduling_loop_step(self._topology)
+
+                sched_loop_duration = time.perf_counter() - t_start
+
+                self.update_metrics(sched_loop_duration)
+                if self._initial_stats:
+                    self._initial_stats.streaming_exec_schedule_s.add(
+                        sched_loop_duration
+                    )
+
+                for callback in self._callbacks:
+                    callback.on_execution_step(self)
+                if not continue_sched or self._shutdown:
+                    break
+        except Exception as e:
+            # Propagate it to the result iterator.
+            exc = e
+        finally:
+            # Mark state of outputting operator as finished
+            _, state = self._output_node
+            state.mark_finished(exc)
+
+    def update_metrics(self, sched_loop_duration: int):
+        self._sched_loop_duration_s.set(
+            sched_loop_duration, tags={"dataset": self._dataset_id}
+        )
+
+    def get_stats(self):
+        """Return the stats object for the streaming execution.
+
+        The stats object will be updated as streaming execution progresses.
+        """
+        if self._final_stats:
+            return self._final_stats
+        else:
+            return self._generate_stats()
+
+    def set_external_consumer_bytes(self, num_bytes: int) -> None:
+        """Set the bytes buffered by external consumers."""
+        if self._resource_manager is not None:
+            self._resource_manager.set_external_consumer_bytes(num_bytes)
+
+    def _generate_stats(self) -> DatasetStats:
+        """Create a new stats object reflecting execution status so far."""
+        stats = self._initial_stats or DatasetStats(metadata={}, parent=None)
+        for op in self._topology:
+            if isinstance(op, InputDataBuffer):
+                continue
+            builder = stats.child_builder(op.name, override_start_time=self._start_time)
+            stats = builder.build_multioperator(op.get_stats())
+            stats.extra_metrics = op.metrics.as_dict(skip_internal_metrics=True)
+        # Always assign a ``Timer`` so downstream consumers can call
+        # ``.get()`` / ``.avg()`` / ``.max()`` / ``.percentile()``
+        # unconditionally. When ``_initial_stats`` is absent we hand
+        # back an empty Timer; zero-sample semantics yield 0 across all
+        # four.
+        stats.streaming_exec_schedule_s = (
+            self._initial_stats.streaming_exec_schedule_s
+            if self._initial_stats
+            else Timer()
+        )
+        return stats
+
+    def _scheduling_loop_step(self, topology: Topology) -> bool:
+        """Run one step of the scheduling loop.
+
+        This runs a few general phases:
+            1. Waiting for the next task completion using `ray.wait()`.
+            2. Pulling completed refs into operator outqueues.
+            3. Selecting and dispatching new inputs to operators.
+
+        Args:
+            topology: The :class:`Topology` of operators being executed.
+
+        Returns:
+            True if we should continue running the scheduling loop.
+        """
+        self._resource_manager.update_usages()
+        # Note: calling process_completed_tasks() is expensive since it incurs
+        # ray.wait() overhead, so make sure to allow multiple dispatch per call for
+        # greater parallelism.
+        num_errored_blocks = process_completed_tasks(
+            topology,
+            self._backpressure_policies,
+            self._max_errored_blocks,
+            output_backpressure_guard=self._output_backpressure_guard,
+            metadata_fetcher=self._metadata_fetcher,
+        )
+        if self._max_errored_blocks > 0:
+            self._max_errored_blocks -= num_errored_blocks
+        self._num_errored_blocks += num_errored_blocks
+
+        self._resource_manager.update_usages()
+        # Dispatch as many operators as we can for completed tasks.
+        self._report_current_usage()
+
+        i = 0
+        while True:
+            op = select_operator_to_run(
+                topology,
+                self._resource_manager,
+                self._backpressure_policies,
+                # If consumer is idling (there's nothing for it to consume)
+                # enforce liveness, ie that at least a single task gets scheduled
+                ensure_liveness=self._consumer_idling(),
+                ranker=self._ranker,
+            )
+
+            if op is None:
+                break
+
+            topology[op].dispatch_next_task()
+
+            self._resource_manager.update_usages()
+
+            i += 1
+            if i % self._progress_manager.TOTAL_PROGRESS_REFRESH_EVERY_N_STEPS == 0:
+                self._refresh_progress_manager(topology)
+
+        # Trigger autoscaling
+        self._cluster_autoscaler.try_trigger_scaling()
+        self._actor_autoscaler.try_trigger_scaling()
+
+        update_operator_states(topology)
+        self._refresh_progress_manager(topology)
+
+        self._update_stats_metrics(state=DatasetState.RUNNING.name)
+        if time.time() - self._last_debug_log_time >= DEBUG_LOG_INTERVAL_SECONDS:
+            _log_op_metrics(topology)
+            _debug_dump_topology(topology, self._resource_manager)
+            self._last_debug_log_time = time.time()
+
+        for op, state in topology.items():
+            # Export operator schema if it's updated
+            if state._schema is not None and self._op_schema.get(op) != state._schema:
+                self._op_schema[op] = state._schema
+                self._export_operator_schema(op)
+
+            # Log metrics of newly completed operators.
+            if not op.has_completed():
+                op.refresh_state()
+            elif not self._has_op_completed[op]:
+                log_str = (
+                    f"Operator {op} completed. "
+                    f"Operator Metrics:\n{op._metrics.as_dict(skip_internal_metrics=True)}"
+                )
+                logger.debug(log_str)
+                self._has_op_completed[op] = True
+                self._validate_operator_queues_empty(op, state)
+
+        # Keep going until all operators run to completion.
+        return not all(op.has_completed() for op in topology)
+
+    def _refresh_progress_manager(self, topology: Topology):
+        # Update the progress manager to reflect scheduling decisions.
+        if self._progress_manager:
+            for op_state in topology.values():
+                if not isinstance(op_state.op, InputDataBuffer):
+                    self._progress_manager.update_operator_progress(
+                        op_state, self._resource_manager
+                    )
+            self._progress_manager.refresh()
+
+    def _consumer_idling(self) -> bool:
+        """Returns whether the user thread is blocked on topology execution."""
+        _, state = self._output_node
+        return len(state.output_queue) == 0
+
+    def _export_operator_schema(self, op: PhysicalOperator) -> None:
+        schema = self._op_schema.get(op)
+        operator_schema_exporter = get_operator_schema_exporter()
+        if (
+            operator_schema_exporter is not None
+            and hasattr(schema, "names")
+            and hasattr(schema, "types")
+        ):
+            names = [str(n) for n in schema.names]
+            types = [str(t) for t in schema.types]
+            operator_schema = OperatorSchema(
+                operator_uuid=op.id,
+                schema_fields=dict(zip(names, types)),
+            )
+            operator_schema_exporter.export_operator_schema(operator_schema)
+
+    def _validate_operator_queues_empty(
+        self, op: PhysicalOperator, state: OpState
+    ) -> None:
+        """Validate that all queues are empty when an operator completes.
+
+        Args:
+            op: The completed operator to validate.
+            state: The operator's execution state.
+        """
+        error_msg = "Expected {} Queue for {} to be empty, but found {} bundles"
+
+        if isinstance(op, InternalQueueOperatorMixin):
+            # 1) Check Internal Input Queue is empty
+            assert op.internal_input_queue_num_blocks() == 0, error_msg.format(
+                "Internal Input", op.name, op.internal_input_queue_num_blocks()
+            )
+
+            # 2) Check Internal Output Queue is empty
+            assert op.internal_output_queue_num_blocks() == 0, error_msg.format(
+                "Internal Output",
+                op.name,
+                op.internal_output_queue_num_blocks(),
+            )
+
+        # 3) Check that External Input Queue is empty
+        for input_q in state.input_queues:
+            assert len(input_q) == 0, error_msg.format(
+                "External Input", op.name, len(input_q)
+            )
+
+    def _report_current_usage(self) -> None:
+        # running_usage is the amount of resources that have been requested but
+        # not necessarily available
+        # TODO(sofian) https://github.com/ray-project/ray/issues/47520
+        # We need to split the reported resources into running, pending-scheduling,
+        # pending-node-assignment.
+        running_usage = self._resource_manager.get_global_running_usage()
+        pending_usage = self._resource_manager.get_global_pending_usage()
+        limits = self._resource_manager.get_global_limits()
+        resources_status = (
+            f"Active & requested resources: "
+            f"{running_usage.cpu:.4g}/{limits.cpu:.4g} CPU, "
+        )
+        if running_usage.memory > 0:
+            resources_status += (
+                f"{running_usage.memory_str()}/{limits.memory_str()} memory, "
+            )
+        if running_usage.gpu > 0:
+            resources_status += f"{running_usage.gpu:.4g}/{limits.gpu:.4g} GPU, "
+        resources_status += (
+            f"{running_usage.object_store_memory_str()}/"
+            f"{limits.object_store_memory_str()} object store"
+        )
+
+        # Only include pending section when there are pending resources.
+        pending_parts = []
+        if pending_usage.cpu:
+            pending_parts.append(f"{pending_usage.cpu:.4g} CPU")
+        if pending_usage.memory:
+            pending_parts.append(f"{pending_usage.memory_str()} memory")
+        if pending_usage.gpu:
+            pending_parts.append(f"{pending_usage.gpu:.4g} GPU")
+        if pending_parts:
+            resources_status += f" (pending: {', '.join(pending_parts)})"
+
+        self._progress_manager.update_total_resource_status(resources_status)
+
+    def _get_operator_id(self, op: PhysicalOperator, topology_index: int) -> str:
+        return f"{op.name}_{topology_index}"
+
+    def _get_operator_tags(self):
+        """Returns a list of operator tags."""
+        return [
+            f"{self._get_operator_id(op, i)}" for i, op in enumerate(self._topology)
+        ]
+
+    def _get_state_dict(self, state):
+        last_op, last_state = list(self._topology.items())[-1]
+        return {
+            "state": state,
+            "progress": last_state.num_completed_tasks,
+            "total": last_op.num_outputs_total(),
+            "total_rows": last_op.num_output_rows_total(),
+            "end_time": time.time()
+            if state in (DatasetState.FINISHED.name, DatasetState.FAILED.name)
+            else None,
+            "operators": {
+                f"{self._get_operator_id(op, i)}": {
+                    "name": op.name,
+                    "progress": op_state.num_completed_tasks,
+                    "total": op.num_outputs_total(),
+                    "total_rows": op.num_output_rows_total(),
+                    "queued_blocks": op_state.total_enqueued_input_blocks(),
+                    "state": DatasetState.FINISHED.name
+                    if op.has_execution_finished()
+                    else state,
+                }
+                for i, (op, op_state) in enumerate(self._topology.items())
+            },
+        }
+
+    def _update_stats_metrics(self, state: str, force_update: bool = False):
+        now = time.time()
+        if (
+            force_update
+            or (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S
+        ):
+            _StatsManager.update_execution_metrics(
+                self._dataset_id,
+                [op.metrics for op in self._topology],
+                self._get_operator_tags(),
+                self._get_state_dict(state=state),
+            )
+            self._metrics_last_updated = now
+
+
+def _debug_dump_topology(topology: Topology, resource_manager: ResourceManager) -> None:
+    """Log current execution state for the topology for debugging.
+
+    Args:
+        topology: The topology to debug.
+        resource_manager: The resource manager for this topology.
+    """
+    logger.debug("Execution Progress:")
+    for i, (op, state) in enumerate(topology.items()):
+        summary_str = format_op_state_summary(state, resource_manager, verbose=True)
+        logger.debug(
+            f"{i}: {op.name} - {summary_str}, "
+            f"Blocks Outputted: {state.num_completed_tasks}/{op.num_outputs_total()}"
+        )
+
+
+def _log_op_metrics(topology: Topology) -> None:
+    """Logs the metrics of each operator.
+
+    Args:
+        topology: The topology to debug.
+    """
+    log_str = "Operator Metrics:\n"
+    for op in topology:
+        metrics_dict = op.metrics.as_dict(skip_internal_metrics=True)
+        log_str += f"{op.name}: {metrics_dict}\n"
+    logger.debug(log_str)
+
+
+class _ClosingIterator(OutputIterator):
+    """Iterator automatically shutting down executor upon exhausting the
+    iterable sequence.
+
+    NOTE: If this iterator isn't fully exhausted, executor still have to
+          be closed manually by the caller!
+    """
+
+    def __init__(self, executor: StreamingExecutor):
+        self._executor = executor
+
+    def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
+        try:
+            op, state = self._executor._output_node
+            bundle = state.get_output_blocking(output_split_idx)
+
+            # Update progress-bars
+            if self._executor._progress_manager:
+                self._executor._progress_manager.update_total_progress(
+                    bundle.num_rows() or 0, op.num_output_rows_total()
+                )
+
+            return bundle
+
+        # Have to be BaseException to catch ``KeyboardInterrupt``
+        #
+        # NOTE: This also handles ``StopIteration``
+        except BaseException as e:
+            # Asynchronously shutdown the executor (ie avoid unnecessary
+            # synchronization on tasks termination)
+            self._executor.shutdown(
+                force=False, exception=e if not isinstance(e, StopIteration) else None
+            )
+            raise
+
+    def __del__(self):
+        # NOTE: Upon garbage-collection we're allowing running tasks
+        #       to be terminated asynchronously (ie avoid unnecessary
+        #       synchronization on their completion)
+        self._executor.shutdown(force=False)

python/ray/data/_internal/execution/util.py

@@ -0,0 +1,125 @@
+import pickle
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Optional, Union
+
+import ray
+from ray.data.block import Block, BlockAccessor, CallableClass
+
+if TYPE_CHECKING:
+    from ray._raylet import StreamingGeneratorStats
+    from ray.data._internal.execution.interfaces import RefBundle
+    from ray.data.block import BlockMetadataWithSchema
+
+
+def merge_label_selector(
+    ray_remote_args: Dict[str, Any],
+    ctx_label_selector: Optional[Dict[str, str]],
+) -> Dict[str, Any]:
+    """Merge a DataContext-level label_selector into ``ray_remote_args``.
+
+    Operator-level keys (already in ``ray_remote_args["label_selector"]``) win on
+    conflict so existing node-pin selectors are preserved. Returns a new dict;
+    the input is not mutated. If ``ctx_label_selector`` is falsy, returns the
+    input unchanged.
+    """
+    if not ctx_label_selector:
+        return ray_remote_args
+    op_selector = ray_remote_args.get("label_selector") or {}
+    merged = {**ctx_label_selector, **op_selector}
+    out = dict(ray_remote_args)
+    out["label_selector"] = merged
+    return out
+
+
+def make_ref_bundles(simple_data: List[List[Any]]) -> List["RefBundle"]:
+    """Create ref bundles from a list of block data.
+
+    One bundle is created for each input block.
+    """
+    import pandas as pd
+    import pyarrow as pa
+
+    from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
+
+    output = []
+    for block in simple_data:
+        block = pd.DataFrame({"id": block})
+        output.append(
+            RefBundle(
+                [
+                    BlockEntry(
+                        ray.put(block),
+                        BlockAccessor.for_block(block).get_metadata(),
+                    )
+                ],
+                owns_blocks=True,
+                schema=pa.lib.Schema.from_pandas(block, preserve_index=False),
+            )
+        )
+    return output
+
+
+memory_units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
+
+
+def memory_string(num_bytes: float) -> str:
+    """Return a human-readable memory string for the given amount of bytes."""
+    k = 0
+    while num_bytes >= 1024 and k < len(memory_units) - 1:
+        num_bytes /= 1024
+        k += 1
+    return f"{num_bytes:.1f}{memory_units[k]}"
+
+
+def locality_string(locality_hits: int, locality_misses) -> str:
+    """Return a human-readable string for object locality stats."""
+    if not locality_misses:
+        return "[all objects local]"
+    return f"[{locality_hits}/{locality_hits + locality_misses} objects local]"
+
+
+def yield_block_with_stats(
+    block: Block,
+    build_metadata: "Callable[[Optional[float]], BlockMetadataWithSchema]",
+) -> Generator[Union[Block, bytes], "StreamingGeneratorStats", None]:
+    """Yield a block then its pickled metadata, per the streaming-gen protocol.
+
+    Args:
+        block: The block to emit.
+        build_metadata: Given the block serialization time in seconds (or ``None``
+            if Ray didn't report it), returns the block's metadata to pickle.
+
+    Yields:
+        Union[Block, bytes]: The block, followed by its pickled
+        ``BlockMetadataWithSchema``.
+    """
+    gen_stats: "StreamingGeneratorStats" = yield block
+    block_ser_time_s = gen_stats.object_creation_dur_s if gen_stats else None
+    yield pickle.dumps(build_metadata(block_ser_time_s))
+
+
+def make_callable_class_single_threaded(callable_cls: CallableClass) -> CallableClass:
+    """Returns a thread-safe CallableClass with the same logic as the provided
+    `callable_cls`.
+
+    This function allows the usage of concurrent actors by safeguarding user logic
+    behind a separate thread.
+
+    This allows batch slicing and formatting to occur concurrently, to overlap with the
+    user provided UDF.
+    """
+
+    class _SingleThreadedWrapper(callable_cls):
+        def __init__(self, *args, **kwargs):
+            self.thread_pool_executor = ThreadPoolExecutor(max_workers=1)
+            super().__init__(*args, **kwargs)
+
+        def __repr__(self):
+            return super().__repr__()
+
+        def __call__(self, *args, **kwargs):
+            # ThreadPoolExecutor will reuse the same thread for every submit call.
+            future = self.thread_pool_executor.submit(super().__call__, *args, **kwargs)
+            return future.result()
+
+    return _SingleThreadedWrapper