chore: import upstream snapshot with attribution

python/ray/data/.claude/CLAUDE.md

@@ -0,0 +1,10 @@
+<!-- Loaded on-demand when Claude works on Ray Data files. -->
+<!-- Keep under 50 lines. Multi-step procedures → skills. Code style → rules/. -->
+
+# Ray Data
+
+## Key Modules
+<!-- Entry points, important abstractions, non-obvious dependencies -->
+
+## Gotchas
+<!-- Non-obvious behaviors, common mistakes, things that break silently -->

python/ray/data/.claude/rules/README.md

@@ -0,0 +1,9 @@
+<!-- Add Ray Data team-specific rules here as .md files. -->
+<!-- Rules with paths: frontmatter only load when matching files are edited. -->
+<!-- Example:
+---
+paths:
+  - "python/ray/data/**/*.py"
+---
+- Use logical operators from ray.data._internal.logical.operators
+-->

python/ray/data/__init__.py

@@ -0,0 +1,209 @@
+# Short term workaround for https://github.com/ray-project/ray/issues/32435
+# Dataset has a hard dependency on pandas, so it doesn't need to be delayed.
+import pandas  # noqa
+from packaging.version import parse as parse_version
+
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+
+from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    ExecutionResources,
+    NodeIdStr,
+)
+from ray.data._internal.logging import configure_logging
+from ray.data._internal.random_config import RandomSeedConfig
+from ray.data.context import DataContext, DatasetContext
+from ray.data.dataset import (
+    Dataset,
+    Schema,
+    SinkMode,
+    ClickHouseTableSettings,
+    SaveMode,
+)
+from ray.data._internal.logical.operators.n_ary_operator import (
+    MixStoppingCondition,
+)
+from ray.data.stats import DatasetSummary
+from ray.data.datasource import (
+    BlockBasedFileDatasink,
+    Datasink,
+    Datasource,
+    FileShuffleConfig,
+    ReadTask,
+    RowBasedFileDatasink,
+)
+from ray.data.iterator import DataIterator, DatasetIterator
+from ray.data.preprocessor import Preprocessor
+from ray.data.read_api import (  # noqa: F401
+    KafkaAuthConfig,  # noqa: F401
+    from_arrow,
+    from_arrow_refs,
+    from_blocks,
+    from_daft,
+    from_dask,
+    from_huggingface,
+    from_items,
+    from_mars,
+    from_modin,
+    from_numpy,
+    from_numpy_refs,
+    from_pandas,
+    from_pandas_refs,
+    from_spark,
+    from_tf,
+    from_torch,
+    range,
+    range_tensor,
+    read_audio,
+    read_avro,
+    read_bigquery,
+    read_binary_files,
+    read_clickhouse,
+    read_csv,
+    read_databricks_tables,
+    read_datasource,
+    read_delta,
+    read_delta_sharing_tables,
+    read_kafka,
+    read_hudi,
+    read_iceberg,
+    read_images,
+    read_json,
+    read_lance,
+    read_mcap,
+    read_mongo,
+    read_numpy,
+    read_parquet,
+    read_snowflake,
+    read_sql,
+    read_text,
+    read_tfrecords,
+    read_unity_catalog,
+    read_videos,
+    read_webdataset,
+    read_zarr,
+)
+from ray.data.catalog import (
+    Catalog,
+    ReaderFormat,
+    ResolvedSource,
+    DatabricksUnityCatalog,
+)
+
+# Module-level cached global functions for callable classes. It needs to be defined here
+# since it has to be process-global across cloudpickled funcs.
+_map_actor_context = None
+
+configure_logging()
+
+try:
+    import pyarrow as pa
+
+    # Import these arrow extension types to ensure that they are registered.
+    from ray.data._internal.tensor_extensions.arrow import (  # noqa
+        ArrowTensorType,
+        ArrowVariableShapedTensorType,
+    )
+
+    # https://github.com/apache/arrow/pull/38608 deprecated `PyExtensionType`, and
+    # disabled it's deserialization by default. To ensure that users can load data
+    # written with earlier version of Ray Data, we enable auto-loading of serialized
+    # tensor extensions.
+    #
+    # NOTE: `PyExtensionType` is deleted from Arrow >= 21.0
+    pyarrow_version = get_pyarrow_version()
+    if pyarrow_version is None or pyarrow_version >= parse_version("21.0.0"):
+        pass
+    else:
+        from ray._common.utils import env_bool
+
+        RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE = env_bool(
+            "RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE", False
+        )
+
+        if (
+            pyarrow_version >= parse_version("14.0.1")
+            and RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE
+        ):
+            pa.PyExtensionType.set_auto_load(True)
+
+except ModuleNotFoundError:
+    pass
+
+
+__all__ = [
+    "ActorPoolStrategy",
+    "BlockBasedFileDatasink",
+    "ClickHouseTableSettings",
+    "Dataset",
+    "DataContext",
+    "DatasetContext",  # Backwards compatibility alias.
+    "DatasetSummary",
+    "DataIterator",
+    "DatasetIterator",  # Backwards compatibility alias.
+    "Datasink",
+    "Datasource",
+    "ExecutionOptions",
+    "ExecutionResources",
+    "FileShuffleConfig",
+    "MixStoppingCondition",
+    "NodeIdStr",
+    "RandomSeedConfig",
+    "ReadTask",
+    "RowBasedFileDatasink",
+    "Schema",
+    "SinkMode",
+    "SaveMode",
+    "TaskPoolStrategy",
+    "from_daft",
+    "from_dask",
+    "from_items",
+    "from_arrow",
+    "from_arrow_refs",
+    "from_blocks",
+    "from_mars",
+    "from_modin",
+    "from_numpy",
+    "from_numpy_refs",
+    "from_pandas",
+    "from_pandas_refs",
+    "from_spark",
+    "from_tf",
+    "from_torch",
+    "from_huggingface",
+    "range",
+    "range_tensor",
+    "read_audio",
+    "read_avro",
+    "read_text",
+    "read_binary_files",
+    "read_clickhouse",
+    "read_csv",
+    "read_datasource",
+    "read_delta",
+    "read_delta_sharing_tables",
+    "read_kafka",
+    "read_hudi",
+    "read_iceberg",
+    "read_images",
+    "read_json",
+    "read_lance",
+    "read_mcap",
+    "read_numpy",
+    "read_mongo",
+    "read_parquet",
+    "read_snowflake",
+    "read_sql",
+    "read_tfrecords",
+    "read_unity_catalog",
+    "read_videos",
+    "read_zarr",
+    "read_webdataset",
+    "Catalog",
+    "ReaderFormat",
+    "ResolvedSource",
+    "DatabricksUnityCatalog",
+    "KafkaAuthConfig",
+    "Preprocessor",
+]

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

@@ -0,0 +1,31 @@
+from typing import TYPE_CHECKING
+
+from .autoscaling_actor_pool import ActorPoolScalingRequest, AutoscalingActorPool
+from .base_actor_autoscaler import ActorAutoscaler
+from .default_actor_autoscaler import DefaultActorAutoscaler, _get_max_scale_up
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+    from ray.data.context import AutoscalingConfig
+
+
+def create_actor_autoscaler(
+    topology: "Topology",
+    resource_manager: "ResourceManager",
+    config: "AutoscalingConfig",
+) -> ActorAutoscaler:
+    return DefaultActorAutoscaler(
+        topology,
+        resource_manager,
+        config=config,
+    )
+
+
+__all__ = [
+    "ActorAutoscaler",
+    "ActorPoolScalingRequest",
+    "AutoscalingActorPool",
+    "create_actor_autoscaler",
+    "_get_max_scale_up",
+]

python/ray/data/_internal/actor_autoscaler/autoscaling_actor_pool.py

@@ -0,0 +1,284 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from ray import ObjectRef
+from ray.actor import ActorHandle
+from ray.data._internal.execution.interfaces.common import NodeIdStr
+from ray.data._internal.execution.interfaces.execution_options import ExecutionResources
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.util.annotations import DeveloperAPI
+
+
+@dataclass(frozen=True)
+class ActorPoolScalingRequest:
+
+    delta: int
+    force: bool = field(default=False)
+    reason: Optional[str] = field(default=None)
+
+    @classmethod
+    def no_op(cls, *, reason: Optional[str] = None) -> "ActorPoolScalingRequest":
+        return ActorPoolScalingRequest(delta=0, reason=reason)
+
+    @classmethod
+    def upscale(cls, *, delta: int, reason: Optional[str] = None):
+        assert delta > 0
+        return ActorPoolScalingRequest(delta=delta, reason=reason)
+
+    @classmethod
+    def downscale(
+        cls, *, delta: int, force: bool = False, reason: Optional[str] = None
+    ):
+        assert delta < 0, "For scale down delta is expected to be negative!"
+        return ActorPoolScalingRequest(delta=delta, force=force, reason=reason)
+
+
+@dataclass(frozen=True)
+class AutoscalingActorConfig:
+    """
+    per_actor_resource_usage: The resource usage per actor.
+    min_size: The minimum number of running actors to be maintained
+        in the pool. Note, that this constraint could be violated when
+        no new work is available for scheduling in the actor pool (ie
+        when operator completes execution).
+    max_size: The maximum number of running actors to be maintained
+        in the pool.
+    initial_size: The initial number of actors to start with.
+    max_actor_concurrency: The maximum number of concurrent tasks a
+        single actor can execute (derived from `ray_remote_args`
+        passed to the operator).
+    max_tasks_in_flight_per_actor: The maximum number of tasks that can
+        be submitted to a single actor at any given time.
+    """
+
+    min_size: int
+    max_size: int
+    initial_size: int
+    max_tasks_in_flight_per_actor: int
+    max_actor_concurrency: int
+    per_actor_resource_usage: ExecutionResources
+
+    def __post_init__(self):
+        assert self.min_size >= 1
+        assert self.max_size >= self.min_size
+        assert self.initial_size <= self.max_size
+        assert self.initial_size >= self.min_size
+        assert self.max_tasks_in_flight_per_actor >= 1
+
+
+@dataclass(frozen=True)
+class ActorPoolInfo:
+    """Breakdown of the state of the actors used by the ``PhysicalOperator``"""
+
+    running: int
+    pending: int
+    restarting: int
+    active: int = 0
+    idle: int = 0
+    pool_utilization: float = 0.0
+    tasks_in_flight: int = 0
+
+    def __str__(self):
+        return (
+            f"running={self.running}, restarting={self.restarting}, "
+            f"pending={self.pending}, active={self.active}, idle={self.idle}, "
+            f"util={self.pool_utilization:.3f}, tasks_in_flight={self.tasks_in_flight}"
+        )
+
+
+@DeveloperAPI
+class AutoscalingActorPool(ABC):
+    """Abstract interface of an autoscaling actor pool.
+
+    A `PhysicalOperator` can manage one or more `AutoscalingActorPool`s.
+    `Autoscaler` is responsible for deciding autoscaling of these actor
+    pools.
+    """
+
+    _LOGICAL_ACTOR_ID_LABEL_KEY = "__ray_data_logical_actor_id"
+    _DEFAULT_POOL_UTILIZATION = 0
+
+    def __init__(self, config: AutoscalingActorConfig):
+        self._config = config
+
+    @abstractmethod
+    def num_running_actors(self) -> int:
+        """Number of running actors."""
+        ...
+
+    @abstractmethod
+    def num_restarting_actors(self) -> int:
+        """Number of restarting actors"""
+        ...
+
+    @abstractmethod
+    def num_active_actors(self) -> int:
+        """Number of actors with at least one active task."""
+        ...
+
+    @abstractmethod
+    def num_pending_actors(self) -> int:
+        """Number of actors pending creation."""
+        ...
+
+    @abstractmethod
+    def num_tasks_in_flight(self) -> int:
+        """Number of current in-flight tasks (ie total nubmer of tasks that have been
+        submitted to the actor pool)."""
+        ...
+
+    def can_schedule_task(self) -> bool:
+        """Returns `True` iff the actor pool has an available actor that can run a task."""
+        return self.select_actors() is not None
+
+    @abstractmethod
+    def scale(self, req: ActorPoolScalingRequest):
+        """Applies autoscaling action"""
+        ...
+
+    @abstractmethod
+    def refresh_actor_state(self):
+        """Refreshes the actor pool state (for, example, running, restarting, pending)"""
+        ...
+
+    @abstractmethod
+    def on_task_submitted(self, actor: ActorHandle):
+        """Callback when an actor is picked for running a task"""
+        ...
+
+    @abstractmethod
+    def on_task_completed(self, actor: ActorHandle):
+        """Called when a task completes. Returns the provided actor to the pool."""
+        ...
+
+    @abstractmethod
+    def select_actors(
+        self,
+        bundle: Optional[RefBundle] = None,
+        actor_locality_enabled: bool = False,
+    ) -> Optional[ActorHandle]:
+        """Select an actor to process the given bundle.
+
+        When ``bundle`` is ``None``, returns any available actor with spare
+        capacity (used by ``can_schedule_task`` to probe schedulability).
+        When ``bundle`` is provided, returns the best actor for that bundle
+        (considering locality when ``actor_locality_enabled`` is True).
+
+        Args:
+            bundle: The bundle to find an actor for. If ``None``, returns any
+                available actor with spare capacity.
+            actor_locality_enabled: Whether to consider locality when selecting
+                an actor.
+
+        Returns:
+            An actor handle if an actor with capacity is available, otherwise
+            ``None``.
+        """
+        ...
+
+    @abstractmethod
+    def get_pending_actor_refs(self) -> List[ObjectRef]:
+        """Return the list of object refs for actors that are pending creation."""
+        ...
+
+    @abstractmethod
+    def pending_to_running(self, ready_ref: ObjectRef) -> Optional[ActorHandle]:
+        """Mark the actor corresponding to the provided ready future as running.
+
+        Args:
+            ready_ref: The ready future for the actor to mark as running.
+
+        Returns:
+            The actor handle if the actor is still alive, otherwise ``None``.
+        """
+        ...
+
+    @abstractmethod
+    def get_actor_location(self, actor: ActorHandle) -> NodeIdStr:
+        """Get the node_id of the actor"""
+        ...
+
+    @abstractmethod
+    def shutdown(self, force: bool = False):
+        """Kills all actors, including running/active actors.
+
+        This is called once the operator is shutting down.
+        """
+        ...
+
+    def get_logical_id_label_key(self) -> str:
+        """Get the label key for the logical actor ID.
+
+        Actors launched by this pool should have this label.
+        """
+        return self._LOGICAL_ACTOR_ID_LABEL_KEY
+
+    def get_actor_info(self) -> ActorPoolInfo:
+        """Returns current snapshot of actors' being used in the pool"""
+        pool_util = self.get_pool_util()
+        # Handle infinite utilization case (no actors)
+        if pool_util == float("inf"):
+            pool_util = self._DEFAULT_POOL_UTILIZATION
+
+        return ActorPoolInfo(
+            running=self.num_alive_actors(),
+            pending=self.num_pending_actors(),
+            restarting=self.num_restarting_actors(),
+            active=self.num_active_actors(),
+            idle=self.num_idle_actors(),
+            pool_utilization=pool_util,
+            tasks_in_flight=self.num_tasks_in_flight(),
+        )
+
+    def num_alive_actors(self) -> int:
+        """Alive actors are all the running actors in ALIVE state."""
+        return self.num_running_actors() - self.num_restarting_actors()
+
+    def num_idle_actors(self) -> int:
+        """Return the number of idle actors in the pool."""
+        return self.num_running_actors() - self.num_active_actors()
+
+    def per_actor_resource_usage(self) -> ExecutionResources:
+        """Per actor resource usage."""
+        return self._config.per_actor_resource_usage
+
+    def max_actor_concurrency(self) -> int:
+        """Returns max number of tasks single actor could run concurrently."""
+        return self._config.max_actor_concurrency
+
+    def max_tasks_in_flight_per_actor(self) -> int:
+        """Max number of in-flight tasks per actor."""
+        return self._config.max_tasks_in_flight_per_actor
+
+    def initial_size(self) -> int:
+        return self._config.initial_size
+
+    def current_size(self) -> int:
+        return self.num_pending_actors() + self.num_running_actors()
+
+    def min_size(self) -> int:
+        """Min size of the actor pool."""
+        return self._config.min_size
+
+    def max_size(self) -> int:
+        """Max size of the actor pool."""
+        return self._config.max_size
+
+    def get_pool_util(self) -> float:
+        """Calculate the utilization of the given actor pool."""
+
+        # If there are no running actors, we set the utilization to indicate that the pool should be scaled up immediately.
+        if self.current_size() == 0:
+            return float("inf")
+        else:
+            # We compute utilization as a ratio of
+            #  - Number of submitted tasks over
+            #  - Max number of tasks that Actor Pool could currently run
+            #
+            # This value could exceed 100%, since by default actors are allowed
+            # to queue tasks (to pipeline task execution by overlapping block
+            # fetching with the execution of the previous task)
+            return self.num_tasks_in_flight() / (
+                self.max_actor_concurrency() * self.current_size()
+            )

python/ray/data/_internal/actor_autoscaler/base_actor_autoscaler.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+
+@DeveloperAPI
+class ActorAutoscaler(ABC):
+    """Abstract interface for Ray Data actor autoscaler."""
+
+    def __init__(
+        self,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ):
+        self._topology = topology
+        self._resource_manager = resource_manager
+
+    @abstractmethod
+    def try_trigger_scaling(self):
+        """Try trigger autoscaling.
+
+        This method will be called each time when StreamingExecutor makes
+        a scheduling decision. A subclass should override this method to
+        handle the autoscaling of `AutoscalingActorPool`s.
+        """
+        ...

python/ray/data/_internal/actor_autoscaler/default_actor_autoscaler.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+import logging
+import math
+import sys
+from typing import TYPE_CHECKING, Optional
+
+from .autoscaling_actor_pool import ActorPoolScalingRequest, AutoscalingActorPool
+from .base_actor_autoscaler import ActorAutoscaler
+from ray.data._internal.execution.interfaces.execution_options import ExecutionResources
+from ray.data.context import WARN_PREFIX, AutoscalingConfig
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import PhysicalOperator
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import OpState, Topology
+
+logger = logging.getLogger(__name__)
+
+
+class DefaultActorAutoscaler(ActorAutoscaler):
+    def __init__(
+        self,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+        *,
+        config: AutoscalingConfig,
+    ):
+        super().__init__(topology, resource_manager)
+
+        self._actor_pool_scaling_up_threshold: float = (
+            config.actor_pool_util_upscaling_threshold
+        )
+        self._actor_pool_scaling_down_threshold: float = (
+            config.actor_pool_util_downscaling_threshold
+        )
+        self._actor_pool_max_upscaling_delta: Optional[
+            int
+        ] = config.actor_pool_max_upscaling_delta
+
+        self._validate_autoscaling_config()
+
+    def try_trigger_scaling(self):
+        for op, state in self._topology.items():
+            actor_pools = op.get_autoscaling_actor_pools()
+            for actor_pool in actor_pools:
+                # Trigger auto-scaling
+                actor_pool.scale(
+                    self._derive_target_scaling_config(actor_pool, op, state)
+                )
+
+    def _derive_target_scaling_config(
+        self,
+        actor_pool: AutoscalingActorPool,
+        op: "PhysicalOperator",
+        op_state: "OpState",
+    ) -> ActorPoolScalingRequest:
+        # If all inputs have been consumed, short-circuit
+        if op.has_completed() or (
+            op._inputs_complete and op_state.total_enqueued_input_blocks() == 0
+        ):
+            num_to_scale_down = self._compute_downscale_delta(actor_pool)
+            return ActorPoolScalingRequest.downscale(
+                delta=-num_to_scale_down, force=True, reason="consumed all inputs"
+            )
+
+        if actor_pool.current_size() < actor_pool.min_size():
+            # Scale up, if the actor pool is below min size.
+            return ActorPoolScalingRequest.upscale(
+                delta=actor_pool.min_size() - actor_pool.current_size(),
+                reason="pool below min size",
+            )
+        elif actor_pool.current_size() > actor_pool.max_size():
+            return ActorPoolScalingRequest.downscale(
+                delta=-(actor_pool.current_size() - actor_pool.max_size()),
+                reason="pool exceeding max size",
+            )
+
+        allocation = self._resource_manager.get_allocation(op)
+        op_usage = self._resource_manager.get_op_usage(op)
+        if allocation is not None and op_usage is not None:
+            over_budget_scale_down = _get_required_scale_down(
+                actor_pool, allocation.subtract(op_usage)
+            )
+            if over_budget_scale_down > 0:
+                max_can_release = actor_pool.current_size() - actor_pool.min_size()
+                num_to_scale_down = min(over_budget_scale_down, max_can_release)
+                if num_to_scale_down > 0:
+                    return ActorPoolScalingRequest.downscale(
+                        delta=-num_to_scale_down,
+                        reason="actor pool exceeds resource allocation",
+                    )
+                return ActorPoolScalingRequest.no_op(
+                    reason="actor pool exceeds resource allocation "
+                    "but cannot scale below min size",
+                )
+
+        # To prevent unexpected downscaling from the initial size, short-circuit if
+        # the operator hasn't received any inputs.
+        if op.metrics.num_inputs_received == 0:
+            return ActorPoolScalingRequest.no_op(reason="no inputs received")
+
+        # Determine whether to scale up based on the actor pool utilization.
+        util = actor_pool.get_pool_util()
+
+        if util >= self._actor_pool_scaling_up_threshold:
+            # Do not scale up if either
+            #   - Actor Pool is at max size already
+            #   - Op is throttled (ie exceeding allocated resource quota)
+            if actor_pool.current_size() >= actor_pool.max_size():
+                return ActorPoolScalingRequest.no_op(reason="reached max size")
+            if not op_state._scheduling_status.under_resource_limits:
+                return ActorPoolScalingRequest.no_op(
+                    reason="operator exceeding resource quota"
+                )
+
+            budget = self._resource_manager.get_budget(op)
+            budget_max_scale_up = (
+                _get_max_scale_up(actor_pool, budget) if budget else sys.maxsize
+            )
+
+            # Determine maximum available scale up based on
+            #   - Maximum available resource budget
+            #   - Configured max scale-up delta (or "+inf" if not configured)
+            #   - Total # of actors needed to reach `max_size`
+            max_scale_up: int = min(
+                budget_max_scale_up,
+                self._get_actor_pool_max_upscaling_delta(),
+                actor_pool.max_size() - actor_pool.current_size(),
+            )
+
+            if max_scale_up == 0:
+                return ActorPoolScalingRequest.no_op(reason="exceeded resource limits")
+
+            if util == float("inf"):
+                return ActorPoolScalingRequest.upscale(
+                    delta=1, reason="no running actors, scale up immediately"
+                )
+
+            delta = self._compute_upscale_delta(actor_pool, op_state)
+            # At least scale up by 1
+            delta = max(1, delta)
+            # Cap delta
+            delta = min(delta, max_scale_up)
+
+            return ActorPoolScalingRequest.upscale(
+                delta=delta,
+                reason=(
+                    f"utilization of {util} >= "
+                    f"{self._actor_pool_scaling_up_threshold}"
+                ),
+            )
+        elif util <= self._actor_pool_scaling_down_threshold:
+            if actor_pool.num_pending_actors() > 0:
+                return ActorPoolScalingRequest.no_op(
+                    reason="no downscaling while actors are pending"
+                )
+            if actor_pool.current_size() <= actor_pool.min_size():
+                return ActorPoolScalingRequest.no_op(reason="reached min size")
+
+            max_can_release = actor_pool.current_size() - actor_pool.min_size()
+            num_to_scale_down = min(
+                self._compute_downscale_delta(actor_pool), max_can_release
+            )
+
+            return ActorPoolScalingRequest.downscale(
+                delta=-num_to_scale_down,
+                reason=(
+                    f"utilization of {util} <= "
+                    f"{self._actor_pool_scaling_down_threshold}"
+                ),
+            )
+        else:
+            return ActorPoolScalingRequest.no_op(
+                reason=(
+                    f"utilization of {util} w/in limits "
+                    f"[{self._actor_pool_scaling_down_threshold}, "
+                    f"{self._actor_pool_scaling_up_threshold}]"
+                )
+            )
+
+    def _get_actor_pool_max_upscaling_delta(self) -> int:
+        return (
+            self._actor_pool_max_upscaling_delta
+            if self._actor_pool_max_upscaling_delta is not None
+            else sys.maxsize
+        )
+
+    def _validate_autoscaling_config(self):
+        # Validate that max upscaling delta is positive to prevent override by safeguard
+        if (
+            self._actor_pool_max_upscaling_delta is not None
+            and self._actor_pool_max_upscaling_delta <= 0
+        ):
+            raise ValueError(
+                f"actor_pool_max_upscaling_delta must be positive, "
+                f"got {self._actor_pool_max_upscaling_delta}"
+            )
+        # Validate that upscaling threshold is positive to prevent division by zero
+        # and incorrect scaling calculations
+        if self._actor_pool_scaling_up_threshold <= 0:
+            raise ValueError(
+                f"actor_pool_util_upscaling_threshold must be positive, "
+                f"got {self._actor_pool_scaling_up_threshold}"
+            )
+
+        for op, state in self._topology.items():
+            for actor_pool in op.get_autoscaling_actor_pools():
+                self._validate_actor_pool_autoscaling_config(actor_pool, op)
+
+    def _validate_actor_pool_autoscaling_config(
+        self,
+        actor_pool: AutoscalingActorPool,
+        op: "PhysicalOperator",
+    ) -> None:
+        """Validate autoscaling configuration.
+
+        Args:
+            actor_pool: Actor pool to validate configuration thereof.
+            op: ``PhysicalOperator`` using target actor pool.
+        """
+        # Fixed-size pools don't autoscale by design
+        if actor_pool.min_size() == actor_pool.max_size():
+            return
+
+        max_tasks_in_flight_per_actor = actor_pool.max_tasks_in_flight_per_actor()
+        max_concurrency = actor_pool.max_actor_concurrency()
+
+        if (
+            max_tasks_in_flight_per_actor / max_concurrency
+            < self._actor_pool_scaling_up_threshold
+        ):
+            logger.warning(
+                f"{WARN_PREFIX} Actor Pool configuration of the {op} will not allow it to scale up: "
+                f"configured utilization threshold ({self._actor_pool_scaling_up_threshold * 100}%) "
+                f"couldn't be reached with configured max_concurrency={max_concurrency} "
+                f"and max_tasks_in_flight_per_actor={max_tasks_in_flight_per_actor} "
+                f"(max utilization will be max_tasks_in_flight_per_actor / max_concurrency = {(max_tasks_in_flight_per_actor / max_concurrency) * 100:g}%)"
+            )
+
+    def _compute_upscale_delta(
+        self, actor_pool: AutoscalingActorPool, op_state: OpState
+    ) -> int:
+        # Calculate desired delta based on utilization
+        return math.ceil(
+            actor_pool.current_size()
+            * (actor_pool.get_pool_util() / self._actor_pool_scaling_up_threshold - 1)
+        )
+
+    def _compute_downscale_delta(self, actor_pool: "AutoscalingActorPool") -> int:
+        return 1
+
+
+def _estimate_total_available_task_slots(actor_pool: "AutoscalingActorPool") -> int:
+    # Estimates number of available task slots to schedule new tasks
+    #
+    # NOTE: This must include pending actors for estimation to make sure
+    #       autoscaler appropriately accounts task slots that will be available
+    #       once pending actors become running.
+    return (
+        actor_pool.max_tasks_in_flight_per_actor() * actor_pool.current_size()
+        - actor_pool.num_tasks_in_flight()
+    )
+
+
+def _get_max_scale_up(
+    actor_pool: AutoscalingActorPool,
+    budget: ExecutionResources,
+) -> int:
+    """Get the maximum number of actors that can be scaled up.
+
+    Args:
+        actor_pool: The actor pool to scale up.
+        budget: The budget to scale up.
+
+    Returns:
+        The maximum number of actors that can be scaled up, or `None` if you can
+        scale up infinitely.
+    """
+    assert budget.cpu >= 0 and budget.gpu >= 0 and budget.memory >= 0
+
+    per_actor = actor_pool.per_actor_resource_usage()
+    assert per_actor.cpu >= 0 and per_actor.gpu >= 0 and per_actor.memory >= 0
+
+    # floordiv handles per_actor.x == 0 → inf (no constraint from that resource)
+    # and budget.x == inf → inf. We ignore object_store_memory since it is not
+    # a per-actor declared resource.
+    divisions = budget.floordiv(per_actor)
+    max_scale_up = min(divisions.cpu, divisions.gpu, divisions.memory)
+    if math.isinf(max_scale_up):
+        return sys.maxsize
+    return int(max_scale_up)
+
+
+def _get_required_scale_down(
+    actor_pool: AutoscalingActorPool,
+    budget: ExecutionResources,
+) -> int:
+    """Get the number of actors that must be removed to fit within budget.
+
+    Args:
+        actor_pool: The actor pool to scale down.
+        budget: The net remaining budget (allocation - usage). Can be negative
+            if the operator is over its allocation.
+
+    Returns:
+        The number of actors that need to be removed, or 0 if the pool
+        is within budget.
+    """
+    per_actor = actor_pool.per_actor_resource_usage()
+
+    required_cpu_scale_down = 0
+    if per_actor.cpu > 0 and budget.cpu < 0:
+        required_cpu_scale_down = math.ceil(abs(budget.cpu) / per_actor.cpu)
+
+    required_gpu_scale_down = 0
+    if per_actor.gpu > 0 and budget.gpu < 0:
+        required_gpu_scale_down = math.ceil(abs(budget.gpu) / per_actor.gpu)
+
+    required_memory_scale_down = 0
+    if per_actor.memory > 0 and budget.memory < 0:
+        required_memory_scale_down = math.ceil(abs(budget.memory) / per_actor.memory)
+
+    return max(
+        required_cpu_scale_down, required_gpu_scale_down, required_memory_scale_down
+    )

python/ray/data/_internal/arrow_block.py

@@ -0,0 +1,694 @@
+import logging
+import random
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import numpy as np
+import pandas as pd
+from packaging.version import parse as parse_version
+
+from ray._common.utils import env_integer
+from ray.data._internal.arrow_ops import transform_polars, transform_pyarrow
+from ray.data._internal.arrow_ops.transform_pyarrow import shuffle
+from ray.data._internal.row import row_repr, row_repr_pretty, row_str
+from ray.data._internal.table_block import TableBlockAccessor, TableBlockBuilder
+from ray.data._internal.tensor_extensions.arrow import (
+    convert_to_pyarrow_array,
+    pyarrow_table_from_pydict,
+)
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockColumn,
+    BlockColumnAccessor,
+    BlockExecStats,
+    BlockMetadataWithSchema,
+    BlockType,
+    U,
+)
+from ray.data.context import DEFAULT_TARGET_MAX_BLOCK_SIZE, DataContext
+from ray.data.expressions import Expr
+
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
+
+if TYPE_CHECKING:
+    import pandas
+
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+
+T = TypeVar("T")
+logger = logging.getLogger(__name__)
+
+
+_MIN_PYARROW_VERSION_TO_NUMPY_ZERO_COPY_ONLY = parse_version("13.0.0")
+_BATCH_SIZE_PRESERVING_STUB_COL_NAME = "__bsp_stub"
+
+
+def _is_user_visible_column(name: str) -> bool:
+    return name != _BATCH_SIZE_PRESERVING_STUB_COL_NAME
+
+
+# Set the max chunk size in bytes for Arrow to Batches conversion in
+# ArrowBlockAccessor.iter_rows(). Default to 4MB, to optimize for image
+# datasets in parquet format.
+ARROW_MAX_CHUNK_SIZE_BYTES = env_integer(
+    "RAY_DATA_ARROW_MAX_CHUNK_SIZE_BYTES",
+    int(DEFAULT_TARGET_MAX_BLOCK_SIZE / 32),
+)
+
+
+# We offload some transformations to polars for performance.
+def get_sort_transform(context: DataContext) -> Callable:
+    if context.use_polars or context.use_polars_sort:
+        return transform_polars.sort
+    else:
+        return transform_pyarrow.sort
+
+
+def get_concat_and_sort_transform(context: DataContext) -> Callable:
+    if context.use_polars or context.use_polars_sort:
+        return transform_polars.concat_and_sort
+    else:
+        return transform_pyarrow.concat_and_sort
+
+
+class ArrowRow(Mapping):
+    """
+    Row of a tabular Dataset backed by a Arrow Table block.
+    """
+
+    def __init__(self, row: Any):
+        self._row = row
+
+    def __getitem__(self, key: Union[str, List[str]]) -> Any:
+        from ray.data.extensions import get_arrow_extension_tensor_types
+
+        tensor_arrow_extension_types = get_arrow_extension_tensor_types()
+
+        def get_item(keys: List[str]) -> Any:
+            schema = self._row.schema
+            if isinstance(schema.field(keys[0]).type, tensor_arrow_extension_types):
+                # Build a tensor row.
+                return tuple(
+                    [
+                        ArrowBlockAccessor._build_tensor_row(
+                            self._row, col_name=key, row_idx=0
+                        )
+                        for key in keys
+                    ]
+                )
+
+            table = self._row.select(keys)
+            if len(table) == 0:
+                return None
+
+            items = [col[0] for col in table.columns]
+            try:
+                # Try to interpret this as a pyarrow.Scalar value.
+                return tuple([item.as_py() for item in items])
+
+            except AttributeError:
+                # Assume that this row is an element of an extension array, and
+                # that it is bypassing pyarrow's scalar model for Arrow < 8.0.0.
+                return items
+
+        is_single_item = isinstance(key, str)
+        keys = [key] if is_single_item else key
+
+        items = get_item(keys)
+
+        if items is None:
+            return None
+        elif is_single_item:
+            return items[0]
+        else:
+            return items
+
+    def __iter__(self) -> Iterator:
+        for k in self._row.column_names:
+            yield k
+
+    def __len__(self):
+        return self._row.num_columns
+
+    def as_pydict(self) -> Dict[str, Any]:
+        return dict(self.items())
+
+    def __str__(self):
+        return row_str(self)
+
+    def __repr__(self):
+        return row_repr(self)
+
+    def _repr_pretty_(self, p, cycle):
+        return row_repr_pretty(self, p, cycle)
+
+
+class ArrowBlockBuilder(TableBlockBuilder):
+    def __init__(self):
+        if pyarrow is None:
+            raise ImportError("Run `pip install pyarrow` for Arrow support")
+        super().__init__((pyarrow.Table, bytes))
+
+    @staticmethod
+    def _table_from_pydict(columns: Dict[str, List[Any]]) -> Block:
+        return pyarrow_table_from_pydict(
+            {
+                column_name: convert_to_pyarrow_array(column_values, column_name)
+                for column_name, column_values in columns.items()
+            }
+        )
+
+    @staticmethod
+    def _combine_tables(tables: List[Block]) -> Block:
+        if len(tables) > 1:
+            return transform_pyarrow.concat(
+                tables, promote_types=True, preserve_order=True
+            )
+        else:
+            return tables[0]
+
+    @staticmethod
+    def _concat_would_copy() -> bool:
+        return False
+
+    @staticmethod
+    def _empty_table() -> "pyarrow.Table":
+        return pyarrow_table_from_pydict({})
+
+    def block_type(self) -> BlockType:
+        return BlockType.ARROW
+
+
+def _get_max_chunk_size(
+    table: "pyarrow.Table", max_chunk_size_bytes: int
+) -> Optional[int]:
+    """
+    Calculate the max chunk size in rows for Arrow to Batches conversion in
+    ArrowBlockAccessor.iter_rows().
+    Args:
+        table: The pyarrow table to calculate the max chunk size for.
+        max_chunk_size_bytes: The max chunk size in bytes.
+    Returns:
+        The max chunk size in rows, or None if the table is empty.
+    """
+    if table.nbytes == 0:
+        return None
+    else:
+        avg_row_size = table.nbytes / table.num_rows
+        return max(1, int(max_chunk_size_bytes / avg_row_size))
+
+
+class ArrowBlockAccessor(TableBlockAccessor):
+    ROW_TYPE = ArrowRow
+
+    def __init__(self, table: "pyarrow.Table"):
+        if pyarrow is None:
+            raise ImportError("Run `pip install pyarrow` for Arrow support")
+        super().__init__(table)
+        self._max_chunk_size: Optional[int] = None
+
+    def _get_row(self, index: int) -> ArrowRow:
+        base_row = self.slice(index, index + 1, copy=False)
+        return ArrowRow(base_row)
+
+    def column_names(self) -> List[str]:
+        return self._table.column_names
+
+    def fill_column(self, name: str, value: Any) -> Block:
+        import pyarrow.compute as pc
+
+        # Check if value is array-like - if so, use upsert_column logic
+        if isinstance(value, (pyarrow.Array, pyarrow.ChunkedArray)):
+            return self.upsert_column(name, value)
+        else:
+            # Scalar value - use original fill_column logic
+            if isinstance(value, pyarrow.Scalar):
+                type = value.type
+            else:
+                type = pyarrow.infer_type([value])
+
+            array = pyarrow.nulls(len(self._table), type=type)
+            array = pc.fill_null(array, value)
+            return self.upsert_column(name, array)
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> "ArrowBlockAccessor":
+        reader = pyarrow.ipc.open_stream(data)
+        return cls(reader.read_all())
+
+    @staticmethod
+    def _build_tensor_row(row: ArrowRow, row_idx: int, col_name: str) -> np.ndarray:
+
+        element = row[col_name][row_idx]
+        arr = element.as_py()
+
+        assert isinstance(arr, np.ndarray), type(arr)
+        return arr
+
+    def slice(self, start: int, end: int, copy: bool = False) -> "pyarrow.Table":
+        view = self._table.slice(start, end - start)
+        if copy:
+            view = transform_pyarrow.combine_chunks(view, copy)
+        return view
+
+    def random_shuffle(self, random_seed: Optional[int]) -> "pyarrow.Table":
+        return shuffle(self._table, random_seed)
+
+    def schema(self) -> "pyarrow.lib.Schema":
+        return self._table.schema
+
+    def to_pandas(self) -> "pandas.DataFrame":
+        from ray.data.util.data_batch_conversion import _cast_tensor_columns_to_ndarrays
+
+        # We specify ignore_metadata=True because pyarrow will use the metadata
+        # to build the Table. This is handled incorrectly for older pyarrow versions
+        ctx = DataContext.get_current()
+
+        # types_mapper preserves Arrow dtypes through the pandas round-trip:
+        # - Standard Arrow types become pd.ArrowDtype, so pa.Table.from_pandas()
+        #   can reconstruct them exactly without lossy numpy conversion.
+        # - Extension types (Ray's ArrowTensorType / ArrowPythonObjectType and
+        #   pyarrow's native FixedShapeTensorType) return None, falling back to
+        #   their own to_pandas_dtype() hooks. Note: native FixedShapeTensorType
+        #   subclasses BaseExtensionType but not ExtensionType, so we check the
+        #   broader BaseExtensionType.
+        def _types_mapper(t):
+            if isinstance(t, pyarrow.BaseExtensionType) or pyarrow.types.is_dictionary(
+                t
+            ):
+                return None
+            return pd.ArrowDtype(t)
+
+        df = self._table.to_pandas(
+            ignore_metadata=ctx.pandas_block_ignore_metadata,
+            types_mapper=_types_mapper,
+        )
+        if ctx.enable_tensor_extension_casting:
+            df = _cast_tensor_columns_to_ndarrays(df, arrow_schema=self._table.schema)
+        return df
+
+    def to_numpy(
+        self, columns: Optional[Union[str, List[str]]] = None
+    ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
+        if columns is None:
+            columns = self._table.column_names
+            should_be_single_ndarray = False
+        elif isinstance(columns, list):
+            should_be_single_ndarray = False
+        else:
+            columns = [columns]
+            should_be_single_ndarray = True
+
+        column_names_set = set(self._table.column_names)
+        for column in columns:
+            if column not in column_names_set:
+                raise ValueError(
+                    f"Cannot find column {column}, available columns: "
+                    f"{column_names_set}"
+                )
+
+        column_values_ndarrays = []
+
+        for col_name in columns:
+            col = self._table[col_name]
+
+            # Combine columnar values arrays to make these contiguous
+            # (making them compatible with numpy format)
+            combined_array = transform_pyarrow.combine_chunked_array(col)
+
+            column_values_ndarrays.append(
+                transform_pyarrow.to_numpy(combined_array, zero_copy_only=False)
+            )
+
+        if should_be_single_ndarray:
+            assert len(columns) == 1
+            return column_values_ndarrays[0]
+        else:
+            return dict(zip(columns, column_values_ndarrays))
+
+    def to_arrow(self) -> "pyarrow.Table":
+        return self._table
+
+    def num_rows(self) -> int:
+        # Arrow may represent an empty table via an N > 0 row, 0-column table, e.g. when
+        # slicing an empty table, so we return 0 if num_columns == 0.
+        return self._table.num_rows if self._table.num_columns > 0 else 0
+
+    def size_bytes(self) -> int:
+        return self._table.nbytes
+
+    def _zip(self, acc: BlockAccessor) -> "Block":
+        r = self.to_arrow()
+        s = acc.to_arrow()
+        for col_name in s.column_names:
+            col = s.column(col_name)
+            # Ensure the column names are unique after zip.
+            if col_name in r.column_names:
+                i = 1
+                new_name = col_name
+                while new_name in r.column_names:
+                    new_name = "{}_{}".format(col_name, i)
+                    i += 1
+                col_name = new_name
+            r = r.append_column(col_name, col)
+        return r
+
+    def upsert_column(
+        self, column_name: str, column_data: BlockColumn
+    ) -> "pyarrow.Table":
+        assert isinstance(
+            column_data, (pyarrow.Array, pyarrow.ChunkedArray)
+        ), f"Expected either a pyarrow.Array or pyarrow.ChunkedArray, got: {type(column_data)}"
+
+        column_idx = self._table.schema.get_field_index(column_name)
+        if column_idx == -1:
+            return self._table.append_column(column_name, column_data)
+        else:
+            return self._table.set_column(column_idx, column_name, column_data)
+
+    @staticmethod
+    def builder() -> ArrowBlockBuilder:
+        return ArrowBlockBuilder()
+
+    @staticmethod
+    def _empty_table() -> "pyarrow.Table":
+        return ArrowBlockBuilder._empty_table()
+
+    def take(
+        self,
+        indices: Union[List[int], "pyarrow.Array", "pyarrow.ChunkedArray"],
+    ) -> "pyarrow.Table":
+        """Select rows from the underlying table.
+
+        This method is an alternative to pyarrow.Table.take(), which breaks for
+        extension arrays.
+        """
+        return transform_pyarrow.take_table(self._table, indices)
+
+    def drop(self, columns: List[str]) -> Block:
+        return self._table.drop(columns)
+
+    def select(self, columns: List[str]) -> "pyarrow.Table":
+        if not all(isinstance(col, str) for col in columns):
+            raise ValueError(
+                "Columns must be a list of column name strings when aggregating on "
+                f"Arrow blocks, but got: {columns}."
+            )
+        if len(columns) == 0:
+            # Empty projection (e.g. count or ``select_columns([])``).
+            # Drop every existing column, then append the stub so row
+            # counts survive downstream ``pa.concat_tables`` calls (which
+            # collapse num_rows to 0 when all inputs have 0 columns).
+            # ``pa.Table`` tracks num_rows as metadata independent of
+            # columns, so ``select([])`` preserves it here. The stub is
+            # filtered out of the user-visible schema; it's a physical
+            # placeholder only.
+            narrowed = self._table.select([])
+            return ArrowBlockAccessor(narrowed).fill_column(
+                _BATCH_SIZE_PRESERVING_STUB_COL_NAME, None
+            )
+        return self._table.select(columns)
+
+    def rename_columns(self, columns_rename: Dict[str, str]) -> "pyarrow.Table":
+        return self._table.rename_columns(columns_rename)
+
+    def hstack(self, other_block: "pyarrow.Table") -> "pyarrow.Table":
+
+        result_table = self._table
+        for name, column in zip(other_block.column_names, other_block.columns):
+            result_table = result_table.append_column(name, column)
+
+        return result_table
+
+    def _sample(self, n_samples: int, sort_key: "SortKey") -> "pyarrow.Table":
+        indices = random.sample(range(self._table.num_rows), n_samples)
+        table = self._table.select(sort_key.get_columns())
+        return transform_pyarrow.take_table(table, indices)
+
+    def sort(self, sort_key: "SortKey") -> Block:
+        assert (
+            sort_key.get_columns()
+        ), f"Sorting columns couldn't be empty (got {sort_key.get_columns()})"
+
+        if self._table.num_rows == 0:
+            # If the pyarrow table is empty we may not have schema
+            # so calling sort_indices() will raise an error.
+            return self._empty_table()
+
+        context = DataContext.get_current()
+        sort = get_sort_transform(context)
+
+        return sort(self._table, sort_key)
+
+    def sort_and_partition(
+        self, boundaries: List[T], sort_key: "SortKey"
+    ) -> List["Block"]:
+        table = self.sort(sort_key)
+
+        if table.num_rows == 0:
+            return [self._empty_table() for _ in range(len(boundaries) + 1)]
+        elif len(boundaries) == 0:
+            return [table]
+
+        return BlockAccessor.for_block(table)._find_partitions_sorted(
+            boundaries, sort_key
+        )
+
+    @staticmethod
+    def merge_sorted_blocks(
+        blocks: List[Block], sort_key: "SortKey"
+    ) -> Tuple[Block, BlockMetadataWithSchema]:
+        stats = BlockExecStats.builder()
+        blocks = [b for b in blocks if b.num_rows > 0]
+        if len(blocks) == 0:
+            ret = ArrowBlockAccessor._empty_table()
+        else:
+            # Handle blocks of different types.
+            blocks = TableBlockAccessor.normalize_block_types(blocks, BlockType.ARROW)
+            concat_and_sort = get_concat_and_sort_transform(DataContext.get_current())
+            ret = concat_and_sort(blocks, sort_key, promote_types=True)
+        return ret, BlockMetadataWithSchema.from_block(
+            ret, block_exec_stats=stats.build()
+        )
+
+    def block_type(self) -> BlockType:
+        return BlockType.ARROW
+
+    def iter_rows(
+        self, public_row_format: bool
+    ) -> Iterator[Union[Mapping, np.ndarray]]:
+        table = self._table
+        if public_row_format:
+            from ray.data._internal.utils.transform_pyarrow import (
+                _is_native_tensor_type,
+            )
+
+            if self._max_chunk_size is None:
+                # Calling _get_max_chunk_size in constructor makes it slow, so we
+                # are calling it here only when needed.
+                self._max_chunk_size = _get_max_chunk_size(
+                    table, ARROW_MAX_CHUNK_SIZE_BYTES
+                )
+            contains_native_tensor_columns = any(
+                _is_native_tensor_type(column.type) for column in table.columns
+            )
+            for batch in table.to_batches(max_chunksize=self._max_chunk_size):
+                if contains_native_tensor_columns:
+                    # HACK: For v1 and v2 tensors, we can control what is returned
+                    # by overriding ExtensionScalar.as_py (see ArrowTensorScalar).
+                    # For pyarrow native FixedShapeTensorArrays we cannot, so we
+                    # use _iter_rows_from_batch_with_tensors to handle conversion.
+                    yield from _iter_rows_from_batch_with_tensors(batch)
+                else:
+                    yield from batch.to_pylist()
+        else:
+            num_rows = self.num_rows()
+            for i in range(num_rows):
+                yield self._get_row(i)
+
+    def filter(self, predicate_expr: "Expr") -> "pyarrow.Table":
+        """Filter rows based on a predicate expression."""
+        if self._table.num_rows == 0:
+            return self._table
+
+        from ray.data._internal.planner.plan_expression.expression_evaluator import (
+            eval_expr,
+        )
+
+        # Evaluate the expression to get a boolean mask
+        mask = eval_expr(predicate_expr, self._table)
+
+        # Use PyArrow's built-in filter method
+        return self._table.filter(mask)
+
+
+def _iter_rows_from_batch_with_tensors(
+    batch: "pyarrow.RecordBatch",
+) -> Iterator[Dict[str, Any]]:
+    """Iterate over rows in a batch that may contain native tensor columns.
+
+    For pyarrow native FixedShapeTensorArrays, we must manually convert them
+    to ndarrays which preserve shape/ndim. Without this, FixedShapeTensorArrays
+    would be translated to contiguous 1d arrays.
+
+    See: https://arrow.apache.org/docs/python/generated/pyarrow.FixedShapeTensorArray.html
+
+    Args:
+        batch: A PyArrow RecordBatch that may contain tensor columns.
+
+    Yields:
+        Dict[str, Any]: Dictionaries mapping column names to values for each row.
+    """
+    from ray.data._internal.utils.transform_pyarrow import _is_native_tensor_type
+
+    col_values = []
+    for column in batch.columns:
+        if _is_native_tensor_type(column.type):
+            col_values.append(column.to_numpy_ndarray())
+        else:
+            col_values.append(column.to_pylist())
+
+    for idx in range(batch.num_rows):
+        yield {name: col[idx] for name, col in zip(batch.column_names, col_values)}
+
+
+class ArrowBlockColumnAccessor(BlockColumnAccessor):
+    def __init__(self, col: Union["pyarrow.Array", "pyarrow.ChunkedArray"]):
+        super().__init__(col)
+
+    def count(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.count(self._column, mode="only_valid" if ignore_nulls else "all")
+        return res.as_py() if as_py else res
+
+    def sum(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.sum(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def min(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.min(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def max(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.max(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def mean(self, *, ignore_nulls: bool, as_py: bool = True) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        res = pac.mean(self._column, skip_nulls=ignore_nulls)
+        return res.as_py() if as_py else res
+
+    def sum_of_squared_diffs_from_mean(
+        self, ignore_nulls: bool, mean: Optional[U] = None, as_py: bool = True
+    ) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        # Calculate mean if not provided
+        if mean is None:
+            mean = self.mean(ignore_nulls=ignore_nulls)
+
+        if mean is None:
+            return None
+
+        res = pac.sum(
+            pac.power(pac.subtract(self._column, mean), 2), skip_nulls=ignore_nulls
+        )
+        return res.as_py() if as_py else res
+
+    def quantile(
+        self, *, q: float, ignore_nulls: bool, as_py: bool = True
+    ) -> Optional[U]:
+        import pyarrow.compute as pac
+
+        array = pac.quantile(self._column, q=q, skip_nulls=ignore_nulls)
+        # NOTE: That quantile method still returns an array
+        res = array[0]
+        return res.as_py() if as_py else res
+
+    def unique(self) -> BlockColumn:
+        import pyarrow.compute as pac
+
+        if self.is_composed_of_lists():
+            # NOTE: Arrow doesn't provide unique kernels for `ListArray`s and
+            #       such, so we rely on Polars to encode and compute unique
+            #       values instead
+            import polars
+
+            return polars.from_arrow(self._column).unique().to_arrow()
+
+        return pac.unique(self._column)
+
+    def value_counts(self) -> Optional[Dict[str, List]]:
+        import pyarrow.compute as pac
+
+        value_counts: pyarrow.StructArray = pac.value_counts(self._column)
+        if len(value_counts) == 0:
+            return None
+        return {
+            "values": value_counts.field("values").to_pylist(),
+            "counts": value_counts.field("counts").to_pylist(),
+        }
+
+    def hash(self) -> BlockColumn:
+        import polars as pl
+
+        df = pl.DataFrame({"col": self._column})
+        hashes = df.hash_rows().cast(pl.Int64, wrap_numerical=True)
+        return hashes.to_arrow()
+
+    def flatten(self) -> BlockColumn:
+        import pyarrow.compute as pac
+
+        return pac.list_flatten(self._column)
+
+    def dropna(self) -> BlockColumn:
+        import pyarrow.compute as pac
+
+        return pac.drop_null(self._column)
+
+    def is_composed_of_lists(self) -> bool:
+        types = (pyarrow.lib.ListType, pyarrow.lib.LargeListType)
+        return isinstance(self._column.type, types)
+
+    def to_pylist(self) -> List[Any]:
+        return self._column.to_pylist()
+
+    def to_numpy(self, zero_copy_only: bool = False) -> np.ndarray:
+        if get_pyarrow_version() < _MIN_PYARROW_VERSION_TO_NUMPY_ZERO_COPY_ONLY:
+            if isinstance(
+                self._column, pyarrow.ChunkedArray
+            ):  # NOTE: ChunkedArray in Pyarrow < 13.0.0 does not support ``zero_copy_only``
+                return self._column.to_numpy()
+            else:
+                return self._column.to_numpy(zero_copy_only=zero_copy_only)
+
+        return self._column.to_numpy(zero_copy_only=zero_copy_only)
+
+    def _to_arrow_compatible_container(self) -> Union[List[Any], "pyarrow.Array"]:
+        return self._column

python/ray/data/_internal/arrow_ops/transform_polars.py

@@ -0,0 +1,59 @@
+from typing import TYPE_CHECKING, List
+
+from packaging.version import parse as parse_version
+
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
+
+if TYPE_CHECKING:
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+pl = None
+
+# Polars 0.16.8 introduced the `descending` parameter for the `sort` method,
+# replacing `reverse`.
+# See https://github.com/pola-rs/polars/issues/5429 for more details.
+_POLARS_SORT_DESCENDING_MIN_VERSION = parse_version("0.16.8")
+
+
+def check_polars_installed():
+    try:
+        global pl
+        import polars as pl
+    except ImportError:
+        raise ImportError(
+            "polars not installed. Install with `pip install polars` or set "
+            "`DataContext.use_polars_sort = False` to fall back to pyarrow"
+        )
+
+
+def sort(table: "pyarrow.Table", sort_key: "SortKey") -> "pyarrow.Table":
+    check_polars_installed()
+    df = pl.from_arrow(table)
+    if parse_version(pl.__version__) >= _POLARS_SORT_DESCENDING_MIN_VERSION:
+        return df.sort(
+            sort_key.get_columns(), descending=sort_key.get_descending()
+        ).to_arrow()
+    else:
+        return df.sort(
+            sort_key.get_columns(), reverse=sort_key.get_descending()
+        ).to_arrow()
+
+
+def concat_and_sort(
+    blocks: List["pyarrow.Table"], sort_key: "SortKey", *, promote_types: bool = False
+) -> "pyarrow.Table":
+    check_polars_installed()
+    blocks = [pl.from_arrow(block) for block in blocks]
+    if parse_version(pl.__version__) >= _POLARS_SORT_DESCENDING_MIN_VERSION:
+        df = pl.concat(blocks).sort(
+            sort_key.get_columns(), descending=sort_key.get_descending()
+        )
+    else:
+        df = pl.concat(blocks).sort(
+            sort_key.get_columns(), reverse=sort_key.get_descending()
+        )
+    return df.to_arrow()

python/ray/data/_internal/arrow_utils.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import itertools
+from typing import List, Union
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+
+
+def _counts_to_offsets(counts: pa.Array) -> pa.Array:
+    """Convert per-row counts to list offsets via cumulative sum."""
+    cumsum = pc.cumulative_sum(counts)
+    return pa.concat_arrays([pa.array([0], type=cumsum.type), cumsum])
+
+
+def _combine_as_list_array(
+    column_values: List[Union[pa.Array, pa.ChunkedArray]] | None = None,
+    *,
+    offsets: pa.Array | None = None,
+    values: pa.Array | None = None,
+    is_large: bool = False,
+    null_mask: pa.Array | None = None,
+) -> pa.Array:
+    """Combine list arrays or build a list array from offsets and values."""
+    if column_values is None:
+        if offsets is None or values is None:
+            raise ValueError(
+                "Either column_values or both offsets and values must be provided."
+            )
+    else:
+        lens = [len(v) for v in column_values]
+        offsets_type = pa.int64() if is_large else pa.int32()
+        offsets = pa.array(np.concatenate([[0], np.cumsum(lens)]), type=offsets_type)
+        values = pa.concat_arrays(
+            itertools.chain(
+                *[
+                    v.chunks if isinstance(v, pa.ChunkedArray) else [v]
+                    for v in column_values
+                ]
+            )
+        )
+
+    offsets_type = pa.int64() if is_large else pa.int32()
+    offsets = pc.cast(offsets, offsets_type)
+    array_cls = pa.LargeListArray if is_large else pa.ListArray
+    list_type = pa.large_list(values.type) if is_large else pa.list_(values.type)
+    return array_cls.from_arrays(offsets, values, list_type, mask=null_mask)

python/ray/data/_internal/average_calculator.py

@@ -0,0 +1,47 @@
+import time
+from collections import deque
+from typing import Deque, Tuple
+
+
+class TimeWindowAverageCalculator:
+    """A utility class to calculate the average of values reported in a time window."""
+
+    def __init__(
+        self,
+        window_s: float,
+    ):
+        assert window_s > 0
+        # Time window in seconds.
+        self._window_s = window_s
+        # Buffer the values reported in the time window, each value is a
+        # tuple of (time, value).
+        self._values: Deque[Tuple[float, float]] = deque()
+        # Sum of all values in the time window.
+        self._sum: float = 0
+
+    def report(self, value: float):
+        """Report a value to the calculator."""
+        now = time.time()
+        self._values.append((now, value))
+        self._sum += value
+        self._trim(now)
+
+    def get_average(self):
+        """Get the average of values reported in the time window,
+        or None if no values reported in the last time window.
+        """
+        self._trim(time.time())
+        if len(self._values) == 0:
+            return None
+        return self._sum / len(self._values)
+
+    def _trim(self, now):
+        """Remove the values reported outside of the time window."""
+        while len(self._values) > 0 and now - self._values[0][0] > self._window_s:
+            _, value = self._values.popleft()
+            self._sum -= value
+
+        # Set sum to 0 if there are no values to avoid accumulated floating-point error
+        # from repeated += / -= operations.
+        if len(self._values) == 0:
+            self._sum = 0

python/ray/data/_internal/batcher.py

@@ -0,0 +1,375 @@
+import warnings
+from typing import Optional
+
+import numpy as np
+
+from ray.data._internal.arrow_block import ArrowBlockAccessor
+from ray.data._internal.arrow_ops import transform_pyarrow
+from ray.data._internal.arrow_ops.transform_pyarrow import try_combine_chunked_columns
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.util import get_total_obj_store_mem_on_node
+from ray.data.block import Block, BlockAccessor
+from ray.util import log_once
+
+# Delay compaction until the shuffle buffer has reached this ratio over the min
+# shuffle buffer size. Setting this to 1 minimizes memory usage, at the cost of
+# frequent compactions. Setting this to higher values increases memory usage but
+# reduces compaction frequency.
+SHUFFLE_BUFFER_COMPACTION_RATIO = 1.5
+
+# Ratio of remaining compacted rows to shuffle_buffer_min_size at which
+# compaction (and re-shuffling of indices) is triggered. Experiments show 0.5
+# is a good trade-off between throughput and randomness.
+SHUFFLE_BUFFER_COMPACTION_THRESHOLD = 0.5
+
+
+class BatcherInterface:
+    def add(self, block: Block):
+        """Add a block to the block buffer.
+
+        Args:
+            block: Block to add to the block buffer.
+        """
+        raise NotImplementedError()
+
+    def done_adding(self) -> bool:
+        """Indicate to the batcher that no more blocks will be added to the buffer."""
+        raise NotImplementedError()
+
+    def has_batch(self) -> bool:
+        """Whether this Batcher has any full batches."""
+        raise NotImplementedError()
+
+    def has_any(self) -> bool:
+        """Whether this Batcher has any data."""
+        raise NotImplementedError()
+
+    def next_batch(self) -> Block:
+        """Get the next batch from the block buffer.
+
+        Returns:
+            A batch represented as a Block.
+        """
+        raise NotImplementedError()
+
+
+class Batcher(BatcherInterface):
+    """Chunks blocks into batches."""
+
+    # Implementation Note: When there are multiple batches per block, this batcher will
+    # slice off and return each batch and add the remaining block back to the buffer
+    # instead of optimally slicing and returning all batches from the block at once.
+    # This will result in extra (and nested) block slicing. However, since slices are
+    # zero-copy views, we sacrifice what should be a small performance hit for better
+    # readability.
+
+    def __init__(self, batch_size: Optional[int], ensure_copy: bool = False):
+        """
+        Construct a batcher that yields batches of batch_sizes rows.
+
+        Args:
+            batch_size: The size of batches to yield.
+            ensure_copy: Whether batches are always copied from the underlying base
+                blocks (not zero-copy views).
+        """
+        self._batch_size = batch_size
+        self._buffer = []
+        self._buffer_size = 0
+        self._done_adding = False
+        self._ensure_copy = ensure_copy
+
+    def add(self, block: Block):
+        """Add a block to the block buffer.
+
+        Note empty block is not added to buffer.
+
+        Args:
+            block: Block to add to the block buffer.
+        """
+        if BlockAccessor.for_block(block).num_rows() > 0:
+            self._buffer.append(block)
+            self._buffer_size += BlockAccessor.for_block(block).num_rows()
+
+    def done_adding(self) -> bool:
+        """Indicate to the batcher that no more blocks will be added to the batcher."""
+        self._done_adding = True
+
+    def has_batch(self) -> bool:
+        """Whether this Batcher has any full batches."""
+        return self.has_any() and (
+            self._batch_size is None or self._buffer_size >= self._batch_size
+        )
+
+    def has_any(self) -> bool:
+        """Whether this Batcher has any data."""
+        return self._buffer_size > 0
+
+    def next_batch(self) -> Block:
+        """Get the next batch from the block buffer.
+
+        Returns:
+            A batch represented as a Block.
+        """
+        assert self.has_batch() or (self._done_adding and self.has_any())
+        needs_copy = self._ensure_copy
+        # If no batch size, short-circuit.
+        if self._batch_size is None:
+            assert len(self._buffer) == 1
+            block = self._buffer[0]
+            if needs_copy:
+                # Copy block if needing to ensure fresh batch copy.
+                block = BlockAccessor.for_block(block)
+                block = block.slice(0, block.num_rows(), copy=True)
+            self._buffer = []
+            self._buffer_size = 0
+            return block
+        output = DelegatingBlockBuilder()
+        leftover = []
+        needed = self._batch_size
+        for block in self._buffer:
+            accessor = BlockAccessor.for_block(block)
+            if needed <= 0:
+                # We already have a full batch, so add this block to
+                # the leftovers.
+                leftover.append(block)
+            elif accessor.num_rows() <= needed:
+                output.add_block(accessor.to_block())
+                needed -= accessor.num_rows()
+            else:
+                # Try de-fragmenting table in case its columns
+                # have too many chunks (potentially hindering performance of
+                # subsequent slicing operation)
+                if isinstance(accessor, ArrowBlockAccessor):
+                    accessor = BlockAccessor.for_block(
+                        transform_pyarrow.try_combine_chunked_columns(
+                            block, min_chunks_to_combine=1
+                        )
+                    )
+
+                # We only need part of the block to fill out a batch.
+                output.add_block(accessor.slice(0, needed, copy=False))
+                # Add the rest of the block to the leftovers.
+                leftover.append(accessor.slice(needed, accessor.num_rows(), copy=False))
+                needed = 0
+
+        # Move the leftovers into the block buffer so they're the first
+        # blocks consumed on the next batch extraction.
+        self._buffer = leftover
+        self._buffer_size -= self._batch_size
+        needs_copy = needs_copy and not output.will_build_yield_copy()
+        batch = output.build()
+        if needs_copy:
+            # Need to ensure that the batch is a fresh copy.
+            batch = BlockAccessor.for_block(batch)
+            batch = batch.slice(0, batch.num_rows(), copy=True)
+        return batch
+
+
+class ShufflingBatcher(BatcherInterface):
+    """Chunks blocks into shuffled batches, using a local in-memory shuffle buffer.
+
+    Uses an **incremental index** approach: on each compaction a permutation
+    array is generated over the buffer rows, and each ``next_batch()`` call
+    gathers a small slice of that permutation via ``take``.
+
+    Properties of this approach:
+
+    * **Memory-efficient** -- the data buffer is kept as-is; only a
+      lightweight int64 index array is allocated on top.
+    * **Smooth per-batch latency** -- each ``take`` operates on a small
+      slice of indices, so per-batch work is short and uniform, making it
+      easy to hide behind prefetch threads.
+
+    Example with ``batch_size=3`` and a 9-row buffer::
+
+        buffer:  [A, B, C, D, E, F, G, H, I]
+        indices: [4, 7, 1, 0, 8, 3, 6, 2, 5]   # random permutation
+
+        next_batch() -> take([4, 7, 1]) -> [E, H, B]   # batch_head 0 -> 3
+        next_batch() -> take([0, 8, 3]) -> [A, I, D]   # batch_head 3 -> 6
+        next_batch() -> take([6, 2, 5]) -> [G, C, F]   # batch_head 6 -> 9
+    """
+
+    def __init__(
+        self,
+        batch_size: Optional[int],
+        shuffle_buffer_min_size: int,
+        shuffle_seed: Optional[int] = None,
+    ):
+        """Constructs a random-shuffling block batcher.
+
+        Args:
+            batch_size: Record batch size.
+            shuffle_buffer_min_size: Minimum number of rows that must be in the local
+                in-memory shuffle buffer in order to yield a batch. When there are no
+                more rows to be added to the buffer, the number of rows in the buffer
+                *will* decrease below this value while yielding the remaining batches,
+                and the final batch may have less than ``batch_size`` rows. Increasing
+                this will improve the randomness of the shuffle but may increase the
+                latency to the first batch.
+            shuffle_seed: The seed to use for the local random shuffle.
+        """
+        if batch_size is None:
+            raise ValueError("Must specify a batch_size if using a local shuffle.")
+        self._batch_size = batch_size
+        self._rng = np.random.default_rng(shuffle_seed)
+        if shuffle_buffer_min_size < batch_size:
+            # Round it up internally to `batch_size` since our algorithm requires it.
+            # This is harmless since it only offers extra randomization.
+            shuffle_buffer_min_size = batch_size
+        self._shuffle_buffer_min_size = shuffle_buffer_min_size
+
+        self._min_rows_to_yield_batch = max(
+            1, int(shuffle_buffer_min_size * SHUFFLE_BUFFER_COMPACTION_THRESHOLD)
+        )
+        self._min_rows_to_trigger_compaction = int(
+            shuffle_buffer_min_size * SHUFFLE_BUFFER_COMPACTION_RATIO
+        )
+        self._builder = DelegatingBlockBuilder()
+        self._shuffle_buffer: Block = None
+        self._shuffled_indices: Optional[np.ndarray] = None
+        self._batch_head = 0
+        self._done_adding = False
+
+        self._total_object_store_nbytes = get_total_obj_store_mem_on_node()
+        self._total_num_rows_added = 0
+        self._total_nbytes_added = 0
+
+    def add(self, block: Block):
+        """Add a block to the shuffle buffer.
+
+        Note empty block is not added to buffer.
+
+        Args:
+            block: Block to add to the shuffle buffer.
+        """
+        # Because Arrow tables are memory mapped, blocks in the builder reside in object
+        # store memory and not local heap memory. So, if you specify a large buffer size
+        # and there isn't enough object store memory on the node, you encounter
+        # spilling.
+        if (
+            self._estimated_min_nbytes_in_buffers is not None
+            and self._estimated_min_nbytes_in_buffers > self._total_object_store_nbytes
+            and log_once("shuffle_buffer_mem_warning")
+        ):
+            warnings.warn(
+                "The node you're iterating on has "
+                f"{memory_string(self._total_object_store_nbytes)} object "
+                "store memory, but the shuffle buffer is estimated to use "
+                f"{memory_string(self._estimated_min_nbytes_in_buffers)}. If you don't "
+                f"decrease the shuffle buffer size from "
+                f"{self._shuffle_buffer_min_size} rows, you might encounter spilling."
+            )
+
+        block_accessor = BlockAccessor.for_block(block)
+        if block_accessor.num_rows() > 0:
+            self._builder.add_block(block)
+            self._total_num_rows_added += block_accessor.num_rows()
+            self._total_nbytes_added += block_accessor.size_bytes()
+
+    @property
+    def _average_row_nbytes(self) -> Optional[int]:
+        """Return the average number of bytes per row added to this batcher."""
+        return (
+            self._total_nbytes_added // self._total_num_rows_added
+            if self._total_num_rows_added > 0
+            else None
+        )
+
+    @property
+    def _estimated_min_nbytes_in_buffers(self) -> Optional[int]:
+        """Return the estimated minimum number of bytes across all buffers.
+
+        This includes data in both the compacted and uncompacted buffers.
+        """
+        if self._average_row_nbytes is None:
+            return None
+
+        return self._average_row_nbytes * self._min_rows_to_trigger_compaction
+
+    def done_adding(self) -> bool:
+        """Indicate to the batcher that no more blocks will be added to the batcher.
+
+        No more blocks should be added to the batcher after calling this.
+        """
+        self._done_adding = True
+
+    def has_any(self) -> bool:
+        """Whether this batcher has any data."""
+        return self._num_rows() > 0
+
+    def has_batch(self) -> bool:
+        """Whether this batcher has any batches."""
+        num_rows = self._num_rows()
+
+        if not self._done_adding:
+            # Delay pulling of batches until the buffer is large enough in order to
+            # amortize compaction overhead.
+            return num_rows >= self._batch_size and (
+                self._num_compacted_rows() >= self._min_rows_to_yield_batch
+                or num_rows - self._batch_size >= self._min_rows_to_trigger_compaction
+            )
+        else:
+            return num_rows >= self._batch_size
+
+    def _num_rows(self) -> int:
+        """Return the total number of rows that haven't been yielded yet.
+
+        This includes rows in both the compacted and uncompacted buffers.
+        """
+        return self._num_compacted_rows() + self._num_uncompacted_rows()
+
+    def _num_compacted_rows(self) -> int:
+        """Return number of unyielded rows in the compacted buffer."""
+        if self._shuffle_buffer is None:
+            return 0
+        return max(0, len(self._shuffled_indices) - self._batch_head)
+
+    def _num_uncompacted_rows(self) -> int:
+        """Return number of unyielded rows in the uncompacted buffer."""
+        return self._builder.num_rows()
+
+    def next_batch(self) -> Block:
+        """Get the next shuffled batch from the shuffle buffer.
+
+        Returns:
+            A batch represented as a Block.
+        """
+        assert self.has_batch() or (self._done_adding and self.has_any())
+        if self._num_uncompacted_rows() > 0 and (
+            self._done_adding
+            or self._num_compacted_rows() <= self._min_rows_to_yield_batch
+        ):
+            if self._shuffle_buffer is not None and self._batch_head < len(
+                self._shuffled_indices
+            ):
+                remaining_indices = self._shuffled_indices[self._batch_head :]
+                remaining_block = BlockAccessor.for_block(self._shuffle_buffer).take(
+                    remaining_indices
+                )
+                self._builder.add_block(remaining_block)
+            self._shuffle_buffer = self._builder.build()
+
+            accessor = BlockAccessor.for_block(self._shuffle_buffer)
+            if isinstance(accessor, ArrowBlockAccessor):
+                self._shuffle_buffer = try_combine_chunked_columns(
+                    self._shuffle_buffer, min_chunks_to_combine=1
+                )
+                accessor = BlockAccessor.for_block(self._shuffle_buffer)
+
+            num_rows = accessor.num_rows()
+            self._shuffled_indices = self._rng.permutation(num_rows)
+
+            self._builder = DelegatingBlockBuilder()
+            self._batch_head = 0
+
+        assert self._shuffle_buffer is not None
+        assert self._shuffled_indices is not None
+        remaining = len(self._shuffled_indices) - self._batch_head
+        batch_size = min(self._batch_size, remaining)
+        batch_indices = self._shuffled_indices[
+            self._batch_head : self._batch_head + batch_size
+        ]
+        self._batch_head += batch_size
+        return BlockAccessor.for_block(self._shuffle_buffer).take(batch_indices)

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

@@ -0,0 +1,3 @@
+from ray.data._internal.block_batching.block_batching import batch_blocks
+
+__all__ = ["batch_blocks"]

python/ray/data/_internal/block_batching/block_batching.py

@@ -0,0 +1,101 @@
+from typing import Callable, Iterator, Optional, TypeVar
+
+from ray.data._internal.block_batching.interfaces import ResolvedBlock
+from ray.data._internal.block_batching.util import (
+    _MappingIterator,
+    blocks_to_batches,
+    collate,
+    format_batches,
+)
+from ray.data._internal.stats import DatasetStats
+from ray.data.block import Block, DataBatch
+
+T = TypeVar("T")
+
+
+def batch_blocks(
+    blocks: Iterator[Block],
+    *,
+    stats: Optional[DatasetStats] = None,
+    batch_size: Optional[int] = None,
+    batch_format: str = "default",
+    drop_last: bool = False,
+    collate_fn: Optional[Callable[[DataBatch], DataBatch]] = None,
+    shuffle_buffer_min_size: Optional[int] = None,
+    shuffle_seed: Optional[int] = None,
+    ensure_copy: bool = False,
+) -> Iterator[DataBatch]:
+    """Create formatted batches of data from 1 or more blocks.
+
+    This function takes in an iterator of already fetched blocks. Consequently, this
+    function doesn't support block prefetching.
+    """
+    # TODO: make stage timings optional at _BatchingIterator so this
+    # shim can be removed. map() avoids holding block references.
+    wrapped_blocks = map(lambda b: ResolvedBlock(block=b), blocks)
+
+    # Build the processing pipeline
+    batch_iter = format_batches(
+        blocks_to_batches(
+            block_iter=wrapped_blocks,
+            stats=stats,
+            batch_size=batch_size,
+            drop_last=drop_last,
+            shuffle_buffer_min_size=shuffle_buffer_min_size,
+            shuffle_seed=shuffle_seed,
+            ensure_copy=ensure_copy,
+        ),
+        batch_format=batch_format,
+        stats=stats,
+        ensure_copy=ensure_copy,
+    )
+
+    if collate_fn is not None:
+        batch_iter = collate(batch_iter, collate_fn=collate_fn, stats=stats)
+
+    return _UserTimingIterator(
+        _MappingIterator(batch_iter, lambda batch: batch.data), stats
+    )
+
+
+class _UserTimingIterator(Iterator[DataBatch]):
+    def __init__(self, iter: Iterator[DataBatch], stats: Optional[DatasetStats]):
+        self._iter = iter
+        self._stats = stats
+        self._active_timer = None
+
+    def __iter__(self) -> Iterator[DataBatch]:
+        return self
+
+    def __next__(self) -> DataBatch:
+        # Since we're tracking time spent in user-code, we stop
+        # the timer immediately when `__next__` is called
+        self._stop_timer()
+
+        try:
+            res = next(self._iter)
+            # Reset timer and return
+            #
+            # NOTE: It's crucial that we reset the timer only after we
+            #       retrieved the result to avoid starting the timer before
+            #       we retrieve the next value
+            self._reset_timer()
+            return res
+        except StopIteration:
+            self._stop_timer()
+            raise
+
+    def _stop_timer(self):
+        if not self._stats:
+            return
+
+        if self._active_timer:
+            self._active_timer.__exit__(None, None, None)
+            self._active_timer = None
+
+    def _reset_timer(self):
+        if not self._stats:
+            return
+
+        self._active_timer = self._stats.iter_user_s.timer()
+        self._active_timer.__enter__()

python/ray/data/_internal/block_batching/interfaces.py

@@ -0,0 +1,134 @@
+import abc
+from dataclasses import dataclass, field
+from typing import Any, Iterable, List, Optional, Tuple
+
+from ray.data._internal.stats import IterationStage, TimeSpan
+from ray.data.block import Block, DataBatch
+from ray.types import ObjectRef
+
+
+@dataclass
+class BlockStageTimings:
+    """Per-block timing for production_wait + data_transfer.
+
+    Both fields are always populated when ``stage_timings`` is set on a
+    ``ResolvedBlock``; the outer ``ResolvedBlock.stage_timings`` Optional
+    encodes "no timing recorded" (e.g. blocks already resolved before
+    entering the pipeline).
+    """
+
+    production_wait: TimeSpan
+    data_transfer: TimeSpan
+
+
+@dataclass
+class ResolvedBlock:
+    """A resolved block paired with its per-block stage timings.
+
+    ``stage_timings`` is None when no timing was recorded (e.g. blocks
+    already resolved before entering the pipeline).
+    """
+
+    block: Block
+    stage_timings: Optional[BlockStageTimings] = None
+
+
+@dataclass
+class BatchStageTimings:
+    """Per-batch timing windows for each iteration stage.
+
+    Fetch stages (production_wait, data_transfer) accumulate one span per
+    block, so they are ``List[TimeSpan]``. Other stages run at most once
+    per batch, so they are ``Optional[TimeSpan]``. ``stages()`` yields
+    ``List[TimeSpan]`` for all stages (single spans wrapped in a 1-element
+    list) so ``_attribute_blocked_time`` can use uniform overlap logic.
+    """
+
+    production_wait: List[TimeSpan] = field(default_factory=list)
+    data_transfer: List[TimeSpan] = field(default_factory=list)
+    batching: Optional[TimeSpan] = None
+    format: Optional[TimeSpan] = None
+    collate: Optional[TimeSpan] = None
+    finalize: Optional[TimeSpan] = None
+
+    def stages(self) -> Iterable[Tuple[IterationStage, List[TimeSpan]]]:
+        """Yield (stage, spans) pairs, wrapping single spans in a list."""
+        return (
+            (IterationStage.PRODUCTION_WAIT, self.production_wait),
+            (IterationStage.DATA_TRANSFER, self.data_transfer),
+            (
+                IterationStage.BATCHING,
+                [self.batching] if self.batching is not None else [],
+            ),
+            (IterationStage.FORMAT, [self.format] if self.format is not None else []),
+            (
+                IterationStage.COLLATE,
+                [self.collate] if self.collate is not None else [],
+            ),
+            (
+                IterationStage.FINALIZE,
+                [self.finalize] if self.finalize is not None else [],
+            ),
+        )
+
+    def accumulate_block_timings(self, src: BlockStageTimings) -> None:
+        """Accumulate a block's fetch timings into this batch's lists.
+
+        A boundary block whose rows span multiple batches is attributed
+        to the first batch it lands in.
+        """
+        self.production_wait.append(src.production_wait)
+        self.data_transfer.append(src.data_transfer)
+
+
+@dataclass
+class BatchMetadata:
+    """Metadata associated with a batch.
+
+    Attributes:
+        batch_idx: The global index of this batch so that downstream operations can
+            maintain ordering.
+        num_rows: Number of rows in this batch (for ``iter_rows_total``).
+        stage_timings: Per-stage timing windows.
+    """
+
+    batch_idx: int
+    num_rows: int = 0
+    stage_timings: BatchStageTimings = field(default_factory=BatchStageTimings)
+
+
+@dataclass
+class Batch:
+    """A batch of data.
+
+    Attributes:
+        metadata: Metadata associated with this batch.
+        data: The batch of data.
+    """
+
+    metadata: BatchMetadata
+    data: DataBatch
+
+
+class CollatedBatch(Batch):
+    """A batch of collated data.
+
+    Attributes:
+        data: The batch of data which is the output of a user provided collate_fn
+            Therefore, the type of this data can be Any.
+    """
+
+    data: Any
+
+
+class BlockPrefetcher(metaclass=abc.ABCMeta):
+    """Interface for prefetching blocks."""
+
+    @abc.abstractmethod
+    def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
+        """Prefetch the provided blocks to this node."""
+        pass
+
+    def stop(self):
+        """Stop prefetching and release resources."""
+        pass

python/ray/data/_internal/block_batching/iter_batches.py

@@ -0,0 +1,558 @@
+import collections
+import time
+from contextlib import contextmanager, nullcontext
+from typing import Any, Callable, Dict, Iterator, List, Optional
+
+import ray
+from ray._common.utils import env_integer
+from ray.data._internal.block_batching.interfaces import (
+    Batch,
+    BlockPrefetcher,
+)
+from ray.data._internal.block_batching.util import (
+    ActorBlockPrefetcher,
+    WaitBlockPrefetcher,
+    blocks_to_batches,
+    collate,
+    finalize_batches,
+    format_batches,
+    iter_threaded,
+    resolve_block_refs,
+)
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data._internal.memory_tracing import trace_deallocation
+from ray.data._internal.stats import DatasetStats, TimeSpan, _StatsManager
+from ray.data.block import Block, DataBatch
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+DEFAULT_FORMAT_THREADPOOL_NUM_WORKERS = env_integer(
+    "RAY_DATA_MAX_FORMAT_THREADPOOL_NUM_WORKERS", 4
+)
+
+
+def _merged_duration(
+    spans: List["TimeSpan"], blocked_start_s: float, blocked_end_s: float
+) -> float:
+    """Total time ``spans`` overlap with ``[blocked_start_s, blocked_end_s]``,
+    with overlapping spans merged so nothing is double-counted."""
+    intervals = []
+    for s in spans:
+        lo = max(s.start_s, blocked_start_s)
+        hi = min(s.end_s, blocked_end_s)
+        if hi > lo:
+            intervals.append((lo, hi))
+    if not intervals:
+        return 0.0
+    intervals.sort()
+    merged = [intervals[0]]
+    for i in range(1, len(intervals)):
+        lo, hi = intervals[i]
+        if lo <= merged[-1][1]:
+            merged[-1] = (merged[-1][0], max(merged[-1][1], hi))
+        else:
+            merged.append((lo, hi))
+    return sum(hi - lo for lo, hi in merged)
+
+
+class BatchIterator:
+    """Defines an iterator pipeline to convert a stream of block object references
+    into a stream of formatted batches ready to be consumed by the user.
+
+    This takes a block iterator and creates batch_size batches, slicing,
+    unioning, shuffling, prefetching, and formatting blocks as needed.
+
+    This involves both pipeline parallelism (e.g. prefetching)
+    and data parallelism (e.g. threadpool operations):
+
+    If prefetch_batches=2, these are all the batches in flight:
+
+    [User thread] trains on Batch 0
+    - [Fetch thread] Batch 1 finalization + move to output queue
+            - [Worker thread 1] Batch 2 formatting + collating
+            - [Worker thread 2] Batch 3 formatting + collating
+            - [Raylet] Batches 4 + 5 fetched to local object store memory
+
+    At any point in time there are prefetch_batches+1 batches in local heap memory.
+    And the next set of prefetch_batches in local object store memory.
+
+    The actual steps are as follows:
+
+    In a single async thread, do the following:
+        1. Trigger Ray local prefetching of `prefetch_batches` worth of block object
+            references.
+        2. Resolve (i.e. call `ray.get()`) on the block references.
+        3. Perform the necessary batch slicing to construct full batches, possibly
+            shuffling if necessary.
+        4. Then, in a threadpool consisting of `prefetch_batches` threads:
+            a. Format the batches to the provided batch format.
+            b. Apply the collate function.
+        5. If preserve_order, restore the original batch order from the
+            threadpool output.
+        6. Finalize each of the (now ordered) collated batches.
+
+    Args:
+        ref_bundles: An iterator over RefBundles.
+        stats: DatasetStats object to record timing and other statistics.
+        dataset_tag: The tag of the dataset to record timing and other statistics.
+        clear_block_after_read: Whether to clear the block from object store
+            manually (i.e. without waiting for Python's automatic GC) after it
+            is read. Doing so will reclaim memory faster and hence reduce the
+            memory footprint. However, the caller has to ensure the safety, i.e.
+            the block will never be accessed again.
+        batch_size: Record batch size, or None to let the system pick.
+        batch_format: The format in which to return each batch.
+            Specify "default" to use the current block format (promoting
+            Arrow to pandas automatically), "pandas" to
+            select ``pandas.DataFrame`` or "pyarrow" to select
+            ``pyarrow.Table``, or None to use entire blocks
+            as batches. Default is "default".
+        drop_last: Whether to drop the last batch if it's incomplete.
+        collate_fn: A function to apply to each data batch before returning it.
+        finalize_fn: A function to apply to each data batch after it has been collated.
+            This function is not run in a threadpool so it can be used for
+            memory-intensive operations such as GPU preloading.
+        shuffle_buffer_min_size: If non-None, the data will be randomly shuffled using a
+            local in-memory shuffle buffer, and this value will serve as the minimum
+            number of rows that must be in the local in-memory shuffle buffer in order
+            to yield a batch.
+        shuffle_seed: The seed to use for the local random shuffle.
+        ensure_copy: Whether batches are always copied from the underlying base
+            blocks (not zero-copy views).
+        prefetch_batches: The number of batches to fetch ahead of the current batch to
+            process. If set to greater than 0, a separate thread will be used to fetch
+            the specified amount of formatted batches from blocks. This improves
+            performance for non-CPU bound UDFs, allowing batch fetching compute and
+            formatting to be overlapped with the UDF. Defaults to 1.
+        prefetch_bytes_callback: A callback to report prefetched bytes to the executor's
+            resource manager.
+        preserve_order: Whether to maintain the original order that the batches
+            were formed from the blocks (e.g., the input block order).
+            This only takes effect in the case that the format/collate threadpool
+            has more than one thread and the output batches have non-deterministic
+            order.
+    """
+
+    UPDATE_METRICS_INTERVAL_S: float = 5.0
+
+    def __init__(
+        self,
+        ref_bundles: Iterator[RefBundle],
+        *,
+        stats: Optional[DatasetStats] = None,
+        dataset_tag: Optional[str] = None,
+        clear_block_after_read: bool = False,
+        batch_size: Optional[int] = None,
+        batch_format: Optional[str] = "default",
+        drop_last: bool = False,
+        collate_fn: Optional[Callable[[DataBatch], Any]] = None,
+        finalize_fn: Optional[Callable[[Any], Any]] = None,
+        shuffle_buffer_min_size: Optional[int] = None,
+        shuffle_seed: Optional[int] = None,
+        ensure_copy: bool = False,
+        prefetch_batches: int = 1,
+        prefetch_bytes_callback: Optional[Callable[[int], None]] = None,
+        preserve_order: bool = False,
+    ):
+        self._ref_bundles = ref_bundles
+        self._stats = stats
+        self._dataset_tag = dataset_tag
+        self._batch_size = batch_size
+        self._batch_format = batch_format
+        self._drop_last = drop_last
+        self._collate_fn = collate_fn
+        self._finalize_fn = finalize_fn
+        self._shuffle_buffer_min_size = shuffle_buffer_min_size
+        self._shuffle_seed = shuffle_seed
+        self._ensure_copy = ensure_copy
+        self._prefetch_batches = prefetch_batches
+        self._prefetch_bytes_callback = prefetch_bytes_callback
+        self._preserve_order = preserve_order
+        self._eager_free = (
+            clear_block_after_read and DataContext.get_current().eager_free
+        )
+
+        actor_prefetcher_enabled = (
+            prefetch_batches > 0
+            and DataContext.get_current().actor_prefetcher_enabled
+            and not ray.util.client.ray.is_connected()
+        )
+        self._prefetcher = (
+            ActorBlockPrefetcher()
+            if actor_prefetcher_enabled
+            else WaitBlockPrefetcher()
+        )
+        self._yielded_first_batch = False
+
+        # 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
+
+    def _prefetch_blocks(
+        self, ref_bundles: Iterator[RefBundle]
+    ) -> Iterator[ObjectRef[Block]]:
+        return prefetch_batches_locally(
+            ref_bundles=ref_bundles,
+            prefetcher=self._prefetcher,
+            num_batches_to_prefetch=self._prefetch_batches,
+            batch_size=self._batch_size,
+            eager_free=self._eager_free,
+            stats=self._stats,
+        )
+
+    def _resolve_block_refs(
+        self, block_refs: Iterator[ObjectRef[Block]]
+    ) -> Iterator[Any]:
+        return resolve_block_refs(block_ref_iter=block_refs, stats=self._stats)
+
+    def _blocks_to_batches(self, blocks: Iterator[Block]) -> Iterator[Batch]:
+        return blocks_to_batches(
+            block_iter=blocks,
+            stats=self._stats,
+            batch_size=self._batch_size,
+            drop_last=self._drop_last,
+            shuffle_buffer_min_size=self._shuffle_buffer_min_size,
+            shuffle_seed=self._shuffle_seed,
+            ensure_copy=self._ensure_copy,
+        )
+
+    def _format_batches(self, batches: Iterator[Batch]) -> Iterator[Batch]:
+        num_threadpool_workers = min(
+            DEFAULT_FORMAT_THREADPOOL_NUM_WORKERS, self._prefetch_batches
+        )
+        return _format_in_threadpool(
+            batch_iter=batches,
+            stats=self._stats,
+            batch_format=self._batch_format,
+            collate_fn=self._collate_fn,
+            num_threadpool_workers=num_threadpool_workers,
+            ensure_copy=self._ensure_copy,
+        )
+
+    def _finalize_batches(
+        self,
+        batch_iter: Iterator[Batch],
+    ) -> Iterator[Batch]:
+        if self._finalize_fn is None:
+            return batch_iter
+
+        return finalize_batches(
+            batch_iter, finalize_fn=self._finalize_fn, stats=self._stats
+        )
+
+    def _restore_original_batch_order(
+        self, batches: Iterator[Batch]
+    ) -> Iterator[Batch]:
+        return restore_original_order(batches)
+
+    def _pipeline(self, ref_bundles: Iterator[RefBundle]) -> Iterator[Batch]:
+        # Step 1: Prefetch logical batches locally.
+        block_iter = self._prefetch_blocks(ref_bundles)
+
+        # Step 2: Resolve the blocks.
+        block_iter = self._resolve_block_refs(block_iter)
+
+        # Step 3: Batch and shuffle the resolved blocks.
+        batch_iter = self._blocks_to_batches(block_iter)
+
+        # Step 4: Format and collate the batches in a threadpool.
+        batch_iter = self._format_batches(batch_iter)
+
+        # Step 5 (optional): Restore the original order of the batches
+        # if preserve_order is True, in the case that the format/collate threadpool
+        # shuffles around the batches non-deterministically.
+        # NOTE: This should happen before finalize_fn so the reorder buffer
+        # holds CPU batches rather than finalize_fn outputs (e.g., GPU tensors).
+        if self._preserve_order:
+            batch_iter = self._restore_original_batch_order(batch_iter)
+
+        # Step 6: Finalize the batches (e.g., move to GPU).
+        batch_iter = self._finalize_batches(batch_iter)
+
+        yield from batch_iter
+
+    def _iter_batches(self) -> Iterator[DataBatch]:
+        """Pull batches from the pipeline and yield batch data.
+
+        Captures the training thread's blocked window around each ``next()``
+        call and attributes it to pipeline stages via
+        ``_attribute_blocked_time``.
+        """
+        batch_iter = iter_threaded(self._ref_bundles, fn=self._pipeline)
+
+        self.before_epoch_start()
+
+        while True:
+            with self.get_next_batch_context():
+                blocked_start_s = time.perf_counter()
+                try:
+                    batch = next(batch_iter)
+                except StopIteration:
+                    break
+                blocked_end_s = time.perf_counter()
+            self._attribute_blocked_time(batch, blocked_start_s, blocked_end_s)
+            with self.yield_batch_context(batch):
+                yield batch.data
+
+        self.after_epoch_end()
+
+    def _attribute_blocked_time(
+        self, batch: Batch, blocked_start_s: float, blocked_end_s: float
+    ) -> None:
+        """Attribute per-stage blocked time via overlap with the training window.
+
+        Each stage's spans on ``batch.metadata.stage_timings`` are intersected
+        with the training thread's blocked window ``[blocked_start_s,
+        blocked_end_s]``. Overlapping spans are merged first, so the result
+        is the total time the stage was active during the stall (no
+        double-counting).
+
+        Limitation: only the yielded batch's spans are attributed. Other
+        in-flight batches (being processed by background threads) may also
+        overlap with the training stall window but are not counted.
+        TODO: track in-flight batches and union their spans for complete
+        attribution. The current implementation suffices for capturing
+        data-loading bottlenecks.
+
+        TODO: reorder buffer wait under ``preserve_order`` is unattributed
+        (per-stage spans are recorded at format/collate completion, before
+        the batch leaves ``restore_original_order``).
+
+        Args:
+            batch: Batch whose per-stage timings should be attributed.
+            blocked_start_s: perf_counter() just before next().
+            blocked_end_s: perf_counter() just after next() returned.
+        """
+        if self._stats is None:
+            return
+        timings = batch.metadata.stage_timings
+        for stage, spans in timings.stages():
+            overlap_s = _merged_duration(spans, blocked_start_s, blocked_end_s)
+            if overlap_s > 0:
+                self._stats.get_blocked_timer(stage).add(overlap_s)
+        self._stats.iter_batches_total += 1
+        self._stats.iter_rows_total += batch.metadata.num_rows
+
+    def __iter__(self) -> Iterator[DataBatch]:
+        return self._iter_batches()
+
+    def before_epoch_start(self):
+        self._yielded_first_batch = False
+
+    def after_epoch_end(self):
+        # Report 0 prefetched bytes at the end of iteration.
+        if self._prefetch_bytes_callback is not None:
+            self._prefetch_bytes_callback(0)
+
+        if self._stats is None:
+            return
+
+        _StatsManager.update_iteration_metrics(self._stats, self._dataset_tag)
+
+    @contextmanager
+    def get_next_batch_context(self):
+        """Context around ``next(batch_iter)``: tracks total blocked time
+        and time-to-first-batch."""
+        try:
+            if self._stats:
+                # Always track total blocked time
+                total_timer = self._stats.iter_total_blocked_s.timer()
+                # Also track the time until the first batch is ready
+                first_batch_ready_timer = (
+                    self._stats.iter_time_to_first_batch_s.timer()
+                    if not self._yielded_first_batch
+                    else nullcontext()
+                )
+                with total_timer, first_batch_ready_timer:
+                    yield
+            else:
+                yield
+        finally:
+            self._yielded_first_batch = True
+
+    @contextmanager
+    def yield_batch_context(self, batch: Batch):
+        """Context around yielding a batch to the user: tracks user time
+        and periodically flushes metrics."""
+        with self._stats.iter_user_s.timer() if self._stats else nullcontext():
+            yield
+
+        # Report prefetched bytes to the executor's resource manager.
+        if self._prefetch_bytes_callback is not None and self._stats is not None:
+            self._prefetch_bytes_callback(self._stats.iter_prefetched_bytes)
+
+        if self._stats is None:
+            return
+        now = time.time()
+        if (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S:
+            _StatsManager.update_iteration_metrics(self._stats, self._dataset_tag)
+            self._metrics_last_updated = now
+
+
+def _format_in_threadpool(
+    batch_iter: Iterator[Batch],
+    stats: DatasetStats,
+    batch_format: Optional[str],
+    collate_fn: Optional[Callable[[DataBatch], Any]],
+    num_threadpool_workers: int,
+    ensure_copy: bool = False,
+) -> Iterator[Batch]:
+    """Executes the batching, formatting, and collation logic in a threadpool.
+
+    Args:
+        batch_iter: An iterator over logical batches.
+        stats: DatasetStats object to record timing and other statistics.
+        batch_format: The format in which to return each batch.
+            Specify "default" to use the current block format (promoting
+            Arrow to pandas automatically), "pandas" to
+            select ``pandas.DataFrame`` or "pyarrow" to select
+            ``pyarrow.Table``, or None to use entire blocks
+            as batches.
+        collate_fn: A function to apply to each data batch before returning it.
+        num_threadpool_workers: The number of threads to use in the threadpool.
+        ensure_copy: Whether batches are always copied from the underlying base
+            blocks (not zero-copy views).
+
+    Returns:
+        An iterator over batches with formatting and collation applied.
+    """
+
+    def threadpool_computations_format_collate(
+        batch_iter: Iterator[Batch],
+    ) -> Iterator[Batch]:
+        # Step 4a: Format the batches.
+        formatted_batch_iter = format_batches(
+            batch_iter, batch_format=batch_format, stats=stats, ensure_copy=ensure_copy
+        )
+
+        # Step 4b: Apply the collate function if applicable.
+        if collate_fn is not None:
+            formatted_batch_iter = collate(
+                formatted_batch_iter, collate_fn=collate_fn, stats=stats
+            )
+
+        return formatted_batch_iter
+
+    if num_threadpool_workers > 0:
+        # Output order is non-deterministic across workers and is restored
+        # downstream by `restore_original_order`.
+        collated_iter = iter_threaded(
+            base_iterator=batch_iter,
+            fn=threadpool_computations_format_collate,
+            num_workers=num_threadpool_workers,
+            output_buffer_size=num_threadpool_workers,
+        )
+    else:
+        collated_iter = threadpool_computations_format_collate(batch_iter)
+
+    return collated_iter
+
+
+def prefetch_batches_locally(
+    ref_bundles: Iterator[RefBundle],
+    prefetcher: BlockPrefetcher,
+    num_batches_to_prefetch: int,
+    batch_size: Optional[int],
+    eager_free: bool = False,
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[ObjectRef[Block]]:
+    """Given an iterator of batched RefBundles, returns an iterator over the
+    corresponding block references while prefetching `num_batches_to_prefetch`
+    batches in advance.
+
+    Args:
+        ref_bundles: An iterator over batched RefBundles.
+        prefetcher: The prefetcher to use.
+        num_batches_to_prefetch: The number of batches to prefetch ahead of the
+            current batch during the scan.
+        batch_size: User specified batch size, or None to let the system pick.
+        eager_free: Whether to eagerly free the object reference from the object store.
+        stats: Dataset stats object used to store ref bundle retrieval time.
+
+    Yields:
+        Block: Block references (as ObjectRefs), in order.
+    """
+
+    def get_next_ref_bundle() -> RefBundle:
+        with stats.iter_get_ref_bundles_s.timer() if stats else nullcontext():
+            return next(ref_bundles)
+
+    sliding_window = collections.deque()
+    current_window_size = 0
+
+    if num_batches_to_prefetch <= 0:
+        if stats:
+            stats.iter_prefetched_bytes = 0
+        for ref_bundle in ref_bundles:
+            for block_ref in ref_bundle.block_refs:
+                yield block_ref
+        return
+
+    if batch_size is not None:
+        num_rows_to_prefetch = num_batches_to_prefetch * batch_size
+    else:
+        num_rows_to_prefetch = None
+
+    # Create and fetch the initial window.
+    # Stop adding if the number of rows in this window is greater than requested
+    # batch size, or if the batch size is None and the number of blocks in this window
+    # is greater than requested batches to prefetch.
+    while (batch_size is not None and current_window_size < num_rows_to_prefetch) or (
+        batch_size is None and len(sliding_window) < num_batches_to_prefetch
+    ):
+        try:
+            next_ref_bundle = get_next_ref_bundle()
+            sliding_window.extend(next_ref_bundle.blocks)
+            current_window_size += next_ref_bundle.num_rows()
+        except StopIteration:
+            break
+
+    prefetcher.prefetch_blocks([entry.ref for entry in sliding_window])
+    if stats:
+        stats.iter_prefetched_bytes = sum(
+            entry.metadata.size_bytes or 0 for entry in sliding_window
+        )
+
+    while sliding_window:
+        entry = sliding_window.popleft()
+        current_window_size -= entry.metadata.num_rows
+        if batch_size is None or current_window_size < num_rows_to_prefetch:
+            try:
+                next_ref_bundle = get_next_ref_bundle()
+                for next_entry in next_ref_bundle.blocks:
+                    sliding_window.append(next_entry)
+                    current_window_size += next_entry.metadata.num_rows
+                prefetcher.prefetch_blocks([entry.ref for entry in sliding_window])
+            except StopIteration:
+                pass
+        if stats:
+            stats.iter_prefetched_bytes = sum(
+                entry.metadata.size_bytes or 0 for entry in sliding_window
+            )
+        yield entry.ref
+        trace_deallocation(entry.ref, loc="iter_batches", free=eager_free)
+    prefetcher.stop()
+
+
+def restore_original_order(batch_iter: Iterator[Batch]) -> Iterator[Batch]:
+    """Restores the original order of the provided `batch_iter`
+
+    This function will yield items from `base_iterator` in the correct order based on
+    each batch's batch_idx. All indexes are expected to be unique.
+
+    `batch_iter` is expected to not have any missing indexes. All indexes from 0 to len
+    (base_iterator) must be present.
+    """
+    next_index_required = 0
+    buffer: Dict[int, Batch] = {}
+    for batch in batch_iter:
+        assert batch.metadata.batch_idx not in buffer
+        buffer[batch.metadata.batch_idx] = batch
+        while next_index_required in buffer:
+            yield buffer.pop(next_index_required)
+            next_index_required += 1
+
+    while next_index_required in buffer:
+        yield buffer.pop(next_index_required)
+        next_index_required += 1

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

@@ -0,0 +1,561 @@
+import dataclasses
+import functools
+import logging
+import queue
+import threading
+import time
+from typing import (
+    Any,
+    Callable,
+    Generator,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+)
+
+import ray
+from ray.actor import ActorHandle
+from ray.data._internal.batcher import Batcher, ShufflingBatcher
+from ray.data._internal.block_batching.interfaces import (
+    Batch,
+    BatchMetadata,
+    BatchStageTimings,
+    BlockPrefetcher,
+    BlockStageTimings,
+    CollatedBatch,
+    ResolvedBlock,
+)
+from ray.data._internal.stats import DatasetStats, TimeSpan, _maybe_time
+from ray.data.block import Block, BlockAccessor, DataBatch
+from ray.types import ObjectRef
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+U = TypeVar("U")
+I = TypeVar("I")
+O = TypeVar("O")
+
+_SENTINEL = object()
+
+
+def iter_threaded(
+    base_iterator: Iterator[T],
+    fn: Callable[[Iterator[T]], Iterator[U]],
+    num_workers: int = 1,
+    output_buffer_size: int = 1,
+) -> Generator[U, None, None]:
+    """Apply ``fn`` to ``base_iterator`` across ``num_workers`` background
+    threads, yielding results through a bounded queue.
+
+    Workers share ``base_iterator`` under a lock (so it may be a stateful,
+    non-thread-safe generator) and run ``fn`` concurrently. With
+    ``num_workers > 1`` the output order is not preserved and must be restored
+    downstream by the consumer.
+
+    Invariant: the number of output-queue items + items in-flight in workers is
+    bounded by ``output_buffer_size``.
+    Workers reserve an output buffer slot before pulling from ``fn``, ensuring
+    they don't run ``fn`` (and hold the result) while waiting for queue space.
+
+    When the consumer stops early (``break``, ``.close()``, or GC), workers
+    are signaled via a stop event so they don't leak. Note: a hanging
+    ``fn`` cannot be interrupted, so ``fn`` must terminate or raise within
+    bounded time per element. For example, the user function should have
+    timeouts if doing blocking I/O.
+
+    Args:
+        base_iterator: Iterator consumed (under a lock) by the workers.
+        fn: Transform applied by each worker to its view of ``base_iterator``.
+        num_workers: Number of background worker threads.
+        output_buffer_size: Max number of items held by the output-queue
+            + in-flight in the workers.
+    """
+    if num_workers < 1:
+        raise ValueError("num_workers must be at least 1.")
+    if output_buffer_size < 1:
+        raise ValueError("output_buffer_size must be at least 1.")
+
+    stopped = threading.Event()
+    result_queue: queue.Queue = queue.Queue()
+    slots = threading.Semaphore(output_buffer_size)
+    iter_lock = threading.Lock()
+
+    def _locked_iter() -> Iterator[T]:
+        while True:
+            with iter_lock:
+                if stopped.is_set():
+                    return
+                try:
+                    item = next(base_iterator)
+                except StopIteration:
+                    return
+            yield item
+
+    def _acquire_slot() -> bool:
+        # Block until a slot is acquired or the consumer has stopped.
+        while not stopped.is_set():
+            if slots.acquire(timeout=0.1):
+                return True
+        return False
+
+    remaining_workers = num_workers
+    remaining_lock = threading.Lock()
+
+    def _worker():
+        nonlocal remaining_workers
+        slot_acquired = False
+        try:
+            # Construct `fn_iter` inside the try so any exception during
+            # construction propagates to the consumer via the outer except.
+            fn_iter = fn(_locked_iter())
+            while True:
+                slot_acquired = _acquire_slot()
+                if not slot_acquired:
+                    break
+                item = next(fn_iter)
+                result_queue.put(item)
+                # The consumer pulling from the result_queue will release the slot.
+                # Resetting here prevents the finally block from double-releasing.
+                slot_acquired = False
+        except StopIteration:
+            pass
+        except Exception as e:
+            # Handle errors in `fn` by propagating them to the consumer.
+            if not stopped.is_set():
+                result_queue.put(e)
+        finally:
+            if slot_acquired:
+                slots.release()
+            with remaining_lock:
+                remaining_workers -= 1
+                is_last = remaining_workers == 0
+            # Signal the consumer that all thread workers have exhausted their input.
+            if is_last and not stopped.is_set():
+                result_queue.put(_SENTINEL)
+
+    worker_threads = [
+        threading.Thread(target=_worker, name="iter_threaded", daemon=True)
+        for _ in range(num_workers)
+    ]
+    for t in worker_threads:
+        t.start()
+
+    try:
+        while True:
+            item = result_queue.get()
+            if item is _SENTINEL:
+                break
+            if isinstance(item, Exception):
+                raise item
+            # Release one slot at yield time so a worker can run `fn` for the next item.
+            slots.release()
+            yield item
+    finally:
+        stopped.set()
+
+
+class _MappingIterator(Iterator[O], Generic[I, O]):
+    """Iterator that applies a transform function to each element.
+
+    Unlike a generator, local variables in __next__ go out of scope when the method
+    returns, avoiding holding references to yielded values.
+    """
+
+    def __init__(self, input_iter: Iterator[I], transform_fn: Callable[[I], O]):
+        self._input_iter = input_iter
+        self._transform_fn = transform_fn
+
+    def __iter__(self) -> "_MappingIterator[I, O]":
+        return self
+
+    def __next__(self) -> O:
+        return self._transform_fn(next(self._input_iter))
+
+
+def _calculate_ref_hits(refs: List[ObjectRef[Any]]) -> Tuple[int, int, int]:
+    """Given a list of object references, returns how many are already on the local
+    node, how many require fetching from another node, and how many have unknown
+    locations. If `DataContext.get_current().enable_get_object_locations_for_metrics` is
+    False, this will return `(0, 0, 0)` as getting object locations is disabled."""
+    current_node_id = ray.get_runtime_context().get_node_id()
+
+    ctx = ray.data.DataContext.get_current()
+    if ctx.enable_get_object_locations_for_metrics:
+        locs = ray.experimental.get_object_locations(refs)
+        nodes: List[List[str]] = [loc["node_ids"] for loc in locs.values()]
+        hits = sum(current_node_id in node_ids for node_ids in nodes)
+        unknowns = sum(1 for node_ids in nodes if not node_ids)
+        misses = len(nodes) - hits - unknowns
+        return hits, misses, unknowns
+
+    return 0, 0, 0
+
+
+def resolve_block_refs(
+    block_ref_iter: Iterator[ObjectRef[Block]],
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[ResolvedBlock]:
+    """Resolve block references via ``ray.get()`` and attach per-block
+    stage timings.
+
+    production_wait is captured manually (no Timer accumulation) to avoid
+    double-counting with ``prefetch_batches_locally``'s
+    ``iter_get_ref_bundles_s`` timer; data_transfer uses ``_maybe_time``
+    normally (no overlap with other timers).
+
+    Args:
+        block_ref_iter: An iterator over block object references.
+        stats: An optional stats object to record block hits, misses, and
+            cumulative ray.get() time.
+
+    Yields:
+        ResolvedBlock: Each resolved block with its stage timings.
+    """
+    hits = 0
+    misses = 0
+    unknowns = 0
+
+    while True:
+        production_wait_start = time.perf_counter() if stats else 0.0
+        try:
+            block_ref = next(block_ref_iter)
+        except StopIteration:
+            break
+        production_wait_span = (
+            TimeSpan(start_s=production_wait_start, end_s=time.perf_counter())
+            if stats
+            else None
+        )
+
+        current_hit, current_miss, current_unknown = _calculate_ref_hits([block_ref])
+        hits += current_hit
+        misses += current_miss
+        unknowns += current_unknown
+
+        # data_transfer: cross-node transfer via ray.get().
+        # TODO(amogkam): batch multiple references in one ray.get() call.
+        with _maybe_time(stats.iter_get_s if stats else None) as data_transfer_span:
+            block = ray.get(block_ref)
+
+        if stats:
+            assert production_wait_span is not None
+            assert data_transfer_span is not None
+            stage_timings = BlockStageTimings(
+                production_wait=production_wait_span,
+                data_transfer=data_transfer_span,
+            )
+        else:
+            stage_timings = None
+        yield ResolvedBlock(block=block, stage_timings=stage_timings)
+
+    if stats:
+        stats.iter_blocks_local = hits
+        stats.iter_blocks_remote = misses
+        stats.iter_unknown_location = unknowns
+
+
+def blocks_to_batches(
+    block_iter: Iterator[ResolvedBlock],
+    stats: Optional[DatasetStats] = None,
+    batch_size: Optional[int] = None,
+    drop_last: bool = False,
+    shuffle_buffer_min_size: Optional[int] = None,
+    shuffle_seed: Optional[int] = None,
+    ensure_copy: bool = False,
+) -> Iterator[Batch]:
+    """Given an iterator over blocks, returns an iterator over batches."""
+    return _BatchingIterator(
+        block_iter,
+        stats=stats,
+        batch_size=batch_size,
+        drop_last=drop_last,
+        shuffle_buffer_min_size=shuffle_buffer_min_size,
+        shuffle_seed=shuffle_seed,
+        ensure_copy=ensure_copy,
+    )
+
+
+class _BatchingIterator(Iterator[Batch]):
+    """Iterator that converts blocks to batches.
+
+    Unlike a generator, local variables in __next__ go out of scope when the method
+    returns, avoiding holding references to yielded values.
+    """
+
+    def __init__(
+        self,
+        block_iter: Iterator[ResolvedBlock],
+        stats: Optional[DatasetStats] = None,
+        batch_size: Optional[int] = None,
+        drop_last: bool = False,
+        shuffle_buffer_min_size: Optional[int] = None,
+        shuffle_seed: Optional[int] = None,
+        ensure_copy: bool = False,
+    ):
+        self._block_iter = block_iter
+        self._stats = stats
+        self._drop_last = drop_last
+        self._global_counter = 0
+        self._done_adding = False
+        # Accumulates per-block stage timings until a batch is yielded.
+        self._pending_timings = BatchStageTimings()
+
+        if shuffle_buffer_min_size is not None:
+            self._batcher = ShufflingBatcher(
+                batch_size=batch_size,
+                shuffle_buffer_min_size=shuffle_buffer_min_size,
+                shuffle_seed=shuffle_seed,
+            )
+        else:
+            self._batcher = Batcher(batch_size=batch_size, ensure_copy=ensure_copy)
+
+    def __iter__(self) -> "_BatchingIterator":
+        return self
+
+    def __next__(self) -> Batch:
+        # Try to get a batch from current batcher state
+        while True:
+            can_yield = self._batcher.has_batch() or (
+                self._batcher.has_any() and self._done_adding and not self._drop_last
+            )
+
+            if can_yield:
+                with _maybe_time(
+                    self._stats.iter_next_batch_s if self._stats else None
+                ) as span:
+                    next_batch = self._batcher.next_batch()
+                self._pending_timings.batching = span
+
+                res = Batch(
+                    metadata=BatchMetadata(
+                        batch_idx=self._global_counter,
+                        num_rows=BlockAccessor.for_block(next_batch).num_rows(),
+                        stage_timings=self._pending_timings,
+                    ),
+                    data=next_batch,
+                )
+                self._pending_timings = BatchStageTimings()
+
+                self._global_counter += 1
+                return res
+
+            elif not self._done_adding:
+                # If can't yield try adding more blocks
+                try:
+                    # NOTE: Block ref is released immediately
+                    block_result = next(self._block_iter)
+                    if block_result.stage_timings is not None:
+                        self._pending_timings.accumulate_block_timings(
+                            block_result.stage_timings
+                        )
+                    self._batcher.add(block_result.block)
+                except StopIteration:
+                    self._batcher.done_adding()
+                    self._done_adding = True
+            else:
+                # In case when
+                #   - We've exhausted input AND
+                #   - There's nothing to yield anymore
+                #
+                # We stop the iteration
+                raise StopIteration
+
+
+def _format_batch(
+    batch: Batch,
+    batch_format: Optional[str],
+    stats: Optional[DatasetStats],
+    ensure_copy: bool = False,
+) -> Batch:
+    with _maybe_time(stats.iter_format_batch_s if stats else None) as span:
+        formatted_data = BlockAccessor.for_block(batch.data).to_batch_format(
+            batch_format
+        )
+        if ensure_copy:
+            formatted_data = _copy_batch(formatted_data)
+    batch.metadata.stage_timings.format = span
+    return dataclasses.replace(batch, data=formatted_data)
+
+
+def _copy_batch(batch: "DataBatch") -> "DataBatch":
+    """Return a copy of a batch, making it writable.
+
+    ``pa.Array.to_numpy()`` returns read-only arrays by default, so when
+    a caller passes ``ensure_copy=True`` (i.e. ``zero_copy_batch=False``) and the
+    block is Arrow, the numpy-format batch must be explicitly copied to give the UDF
+    writable arrays.
+    """
+    import numpy as np
+
+    if isinstance(batch, dict):
+        # Return a dictionary with the same keys (column names) and values (column numpy arrays),
+        # with the values copied
+        return {
+            k: v.copy() if isinstance(v, np.ndarray) else v for k, v in batch.items()
+        }
+    elif isinstance(batch, np.ndarray):
+        return batch.copy()
+    return batch
+
+
+def format_batches(
+    batch_iter: Iterator[Batch],
+    batch_format: Optional[str],
+    stats: Optional[DatasetStats] = None,
+    ensure_copy: bool = False,
+) -> Iterator[Batch]:
+    """Given an iterator of batches, returns an iterator of formatted batches."""
+    return _MappingIterator(
+        batch_iter,
+        functools.partial(
+            _format_batch,
+            batch_format=batch_format,
+            stats=stats,
+            ensure_copy=ensure_copy,
+        ),
+    )
+
+
+def _collate_batch(
+    batch: Batch,
+    collate_fn: Callable[[DataBatch], Any],
+    stats: Optional[DatasetStats],
+) -> CollatedBatch:
+    with _maybe_time(stats.iter_collate_batch_s if stats else None) as span:
+        collated_data = collate_fn(batch.data)
+    batch.metadata.stage_timings.collate = span
+    return CollatedBatch(metadata=batch.metadata, data=collated_data)
+
+
+def collate(
+    batch_iter: Iterator[Batch],
+    collate_fn: Optional[Callable[[DataBatch], Any]],
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[CollatedBatch]:
+    """Returns an iterator with the provided collate_fn applied to batches."""
+    if not isinstance(batch_iter, Iterator):
+        batch_iter = iter(batch_iter)
+
+    return _MappingIterator(
+        batch_iter,
+        functools.partial(_collate_batch, collate_fn=collate_fn, stats=stats),
+    )
+
+
+def _finalize_batch(
+    batch: CollatedBatch,
+    finalize_fn: Callable[[Any], Any],
+    stats: Optional[DatasetStats],
+) -> CollatedBatch:
+    with _maybe_time(stats.iter_finalize_batch_s if stats else None) as span:
+        finalized_data = finalize_fn(batch.data)
+    batch.metadata.stage_timings.finalize = span
+    return dataclasses.replace(batch, data=finalized_data)
+
+
+def finalize_batches(
+    batch_iter: Iterator[CollatedBatch],
+    finalize_fn: Callable[[Any], Any],
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[CollatedBatch]:
+    """Returns an iterator with finalize_fn applied to batches."""
+    if not isinstance(batch_iter, Iterator):
+        batch_iter = iter(batch_iter)
+
+    return _MappingIterator(
+        batch_iter,
+        functools.partial(_finalize_batch, finalize_fn=finalize_fn, stats=stats),
+    )
+
+
+PREFETCHER_ACTOR_NAMESPACE = "ray.dataset"
+
+
+class WaitBlockPrefetcher(BlockPrefetcher):
+    """Block prefetcher using ray.wait."""
+
+    def __init__(self):
+        self._blocks = []
+        self._stopped = False
+        self._condition = threading.Condition()
+        self._thread = threading.Thread(
+            target=self._run,
+            name="Prefetcher",
+            daemon=True,
+        )
+        self._thread.start()
+
+    def _run(self):
+        while not self._stopped:
+            try:
+                with self._condition:
+                    if len(self._blocks) == 0:
+                        # Park, waiting for notification that prefetching
+                        # should resume
+                        self._condition.wait()
+
+                    blocks_to_fetch, self._blocks = self._blocks[:], []
+
+                if len(blocks_to_fetch) > 0:
+                    ray.wait(
+                        blocks_to_fetch,
+                        num_returns=1,
+                        # NOTE: We deliberately setting timeout to 0 to avoid
+                        #       blocking the fetching thread unnecessarily
+                        timeout=0,
+                        fetch_local=True,
+                    )
+            except Exception:
+                logger.exception("Error in prefetcher thread.")
+
+        logger.debug("Exiting prefetcher's background thread")
+
+    def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
+        with self._condition:
+            if self._stopped:
+                raise RuntimeError("Prefetcher is stopped.")
+            self._blocks = blocks
+            self._condition.notify()
+
+    def stop(self):
+        with self._condition:
+            if self._stopped:
+                return
+            self._stopped = True
+            self._condition.notify()
+
+    def __del__(self):
+        self.stop()
+
+
+class ActorBlockPrefetcher(BlockPrefetcher):
+    """Block prefetcher using a local actor."""
+
+    def __init__(self):
+        self.prefetch_actor = self._get_or_create_actor_prefetcher()
+
+    @staticmethod
+    def _get_or_create_actor_prefetcher() -> "ActorHandle":
+        node_id = ray.get_runtime_context().get_node_id()
+        actor_name = f"dataset-block-prefetcher-{node_id}"
+        return _BlockPretcher.options(
+            label_selector={ray._raylet.RAY_NODE_ID_KEY: node_id},
+            name=actor_name,
+            namespace=PREFETCHER_ACTOR_NAMESPACE,
+            get_if_exists=True,
+        ).remote()
+
+    def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
+        self.prefetch_actor.prefetch.remote(*blocks)
+
+
+@ray.remote(num_cpus=0)
+class _BlockPretcher:
+    """Helper actor that prefetches blocks asynchronously."""
+
+    def prefetch(self, *blocks) -> None:
+        pass

python/ray/data/_internal/block_builder.py

@@ -0,0 +1,39 @@
+from typing import Generic
+
+from ray.data.block import Block, BlockAccessor, BlockType, T
+
+
+class BlockBuilder(Generic[T]):
+    """A builder class for blocks."""
+
+    @staticmethod
+    def for_block(block: Block) -> "BlockBuilder":
+        return BlockAccessor.for_block(block).builder()
+
+    def add(self, item: T) -> None:
+        """Append a single row to the block being built."""
+        raise NotImplementedError
+
+    def add_block(self, block: Block) -> None:
+        """Append an entire block to the block being built."""
+        raise NotImplementedError
+
+    def will_build_yield_copy(self) -> bool:
+        """Whether building this block will yield a new block copy."""
+        raise NotImplementedError
+
+    def build(self) -> Block:
+        """Build the block."""
+        raise NotImplementedError
+
+    def num_rows(self) -> int:
+        """Return the number of rows added in the block."""
+        raise NotImplementedError
+
+    def get_estimated_memory_usage(self) -> int:
+        """Return the estimated memory usage so far in bytes."""
+        raise NotImplementedError
+
+    def block_type(self) -> BlockType:
+        """Return the block type."""
+        raise NotImplementedError

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

@@ -0,0 +1,71 @@
+import enum
+import logging
+import os
+from typing import TYPE_CHECKING
+
+from .base_autoscaling_coordinator import (
+    AutoscalingCoordinator,
+    ResourceDict,
+    ResourceRequestPriority,
+)
+from .base_cluster_autoscaler import ClusterAutoscaler
+from .default_autoscaling_coordinator import (
+    DefaultAutoscalingCoordinator,
+    get_or_create_autoscaling_coordinator,
+)
+from .default_cluster_autoscaler_v2 import DefaultClusterAutoscalerV2
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+CLUSTER_AUTOSCALER_ENV_KEY = "RAY_DATA_CLUSTER_AUTOSCALER"
+DEFAULT_CLUSTER_AUTOSCALER_VERSION = "V2"
+
+
+class ClusterAutoscalerVersion(str, enum.Enum):
+    V2 = "V2"
+
+
+def create_cluster_autoscaler(
+    topology: "Topology",
+    resource_manager: "ResourceManager",
+    data_context: "DataContext",
+    *,
+    execution_id: str,
+) -> ClusterAutoscaler:
+    resource_limits = data_context.execution_options.resource_limits
+    label_selector = data_context.execution_options.label_selector
+    cluster_autoscaler_version = os.environ.get(
+        CLUSTER_AUTOSCALER_ENV_KEY, DEFAULT_CLUSTER_AUTOSCALER_VERSION
+    )
+    logger.debug(f"Using cluster autoscaler version: {cluster_autoscaler_version!r}")
+
+    if cluster_autoscaler_version == ClusterAutoscalerVersion.V2:
+        return DefaultClusterAutoscalerV2(
+            resource_manager,
+            execution_id=execution_id,
+            resource_limits=resource_limits,
+            label_selector=label_selector,
+        )
+
+    else:
+        valid_values = [version.value for version in ClusterAutoscalerVersion]
+        raise ValueError(
+            f"Cluster autoscaler version of {cluster_autoscaler_version} isn't a valid "
+            f"option. Valid options are: {valid_values}."
+        )
+
+
+__all__ = [
+    "ClusterAutoscaler",
+    # Objects related to the `AutoscalingCoordinator`.
+    "AutoscalingCoordinator",
+    "DefaultAutoscalingCoordinator",
+    "get_or_create_autoscaling_coordinator",
+    "ResourceDict",
+    "ResourceRequestPriority",
+]

python/ray/data/_internal/cluster_autoscaler/base_autoscaling_coordinator.py

@@ -0,0 +1,58 @@
+import abc
+from enum import Enum
+from typing import Dict, List, Optional
+
+ResourceDict = Dict[str, float]
+
+
+class ResourceRequestPriority(Enum):
+    """Priority of a resource request."""
+
+    LOW = -10
+    MEDIUM = 0
+    HIGH = 10
+
+
+class AutoscalingCoordinator(abc.ABC):
+    @abc.abstractmethod
+    def request_resources(
+        self,
+        resources: List[ResourceDict],
+        expire_after_s: float,
+        request_remaining: bool = False,
+        priority: ResourceRequestPriority = ResourceRequestPriority.MEDIUM,
+        label_selectors: Optional[List[Dict[str, str]]] = None,
+    ) -> None:
+        """Request cluster resources.
+
+        The requested resources should represent the full set of resources needed,
+        not just the incremental amount.
+
+        Args:
+            resources: The requested resources. This should match the format accepted
+                by `ray.autoscaler.sdk.request_resources`.
+            expire_after_s: Time in seconds after which this request will expire.
+                The requester is responsible for periodically sending new requests
+                to avoid the request being purged.
+            request_remaining: If true, after allocating requested resources to each
+                requester, remaining resources will also be allocated to this requester.
+            priority: The priority of the request. Higher value means higher priority.
+            label_selectors: Optional per-bundle label selectors, one per entry in
+                ``resources``. Forwarded to the autoscaler as
+                ``bundle_label_selectors``.
+        """
+        ...
+
+    @abc.abstractmethod
+    def cancel_request(self) -> None:
+        """Cancel the resource request from the requester."""
+        ...
+
+    @abc.abstractmethod
+    def get_allocated_resources(self) -> List[ResourceDict]:
+        """Get the allocated resources for the requester.
+
+        Returns:
+            A list of dictionaries representing the allocated resources bundles.
+        """
+        ...

python/ray/data/_internal/cluster_autoscaler/base_cluster_autoscaler.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.execution_options import (
+        ExecutionResources,
+    )
+
+
+@DeveloperAPI
+class ClusterAutoscaler(ABC):
+    """Abstract interface for Ray Data cluster autoscaler."""
+
+    @abstractmethod
+    def try_trigger_scaling(self):
+        """Try trigger autoscaling.
+
+        This method will be called each time when StreamingExecutor makes
+        a scheduling decision. A subclass should override this method to
+        handle the autoscaling of the cluster.
+        """
+        ...
+
+    @abstractmethod
+    def on_executor_shutdown(self):
+        """Callback when the StreamingExecutor is shutting down."""
+        ...
+
+    @abstractmethod
+    def get_total_resources(self) -> "ExecutionResources":
+        """Get the total resources that are available to this data execution."""
+        ...

python/ray/data/_internal/cluster_autoscaler/default_autoscaling_coordinator.py

@@ -0,0 +1,577 @@
+import copy
+import functools
+import logging
+import threading
+import time
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional
+
+import ray
+import ray.exceptions
+from .base_autoscaling_coordinator import (
+    AutoscalingCoordinator,
+    ResourceDict,
+    ResourceRequestPriority,
+)
+from ray._common.utils import env_bool
+from ray.data._internal.execution.util import memory_string
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+logger = logging.getLogger(__name__)
+
+HEAD_NODE_RESOURCE_LABEL = "node:__internal_head__"
+_RESOURCE_LOG_KEYS = ("CPU", "GPU", "memory", "object_store_memory")
+_RESOURCE_LOG_MEMORY_KEYS = {"memory", "object_store_memory"}
+# Label key the cluster autoscaler uses to bucket nodes by subcluster.
+# Hardcoded so all components agree without per-Dataset configuration.
+SUBCLUSTER_LABEL_KEY = "ray-subcluster"
+# Sentinel for "no subcluster" — used as both a node-label fallback and
+# the bucket key for unlabeled nodes in ``_cluster_node_resources``.
+DEFAULT_SUBCLUSTER: Optional[str] = None
+
+
+RAY_DATA_AUTOSCALING_COORDINATOR_LOG_TRACEBACK = env_bool(
+    "RAY_DATA_AUTOSCALING_COORDINATOR_LOG_TRACEBACK", True
+)
+
+
+def _format_resource_value_for_log(resource_name: str, value: float) -> str:
+    """Format a numerical resource value to a human-readable string.
+
+    Args:
+        resource_name: The resource name.
+        value: The resource value.
+
+    Returns:
+        A human-readable string.
+    """
+    if resource_name in _RESOURCE_LOG_MEMORY_KEYS:
+        return memory_string(value)
+    if isinstance(value, float) and value.is_integer():
+        return str(int(value))
+    return str(value)
+
+
+def _format_resource_bundle_for_log(bundle: ResourceDict) -> str:
+    """Format a resource bundle to a human-readable string.
+
+    Drops custom resource keys (e.g. ``anyscale/...``, ``node:...``) and
+    zero-valued resources, keeping only the standard keys in ``_RESOURCE_LOG_KEYS``.
+
+    Args:
+        bundle: The resource bundle to format.
+
+    Returns:
+        A human-readable string, e.g. ``"{CPU: 8, memory: 32.0GiB}"``.
+
+    Example:
+        >>> from ray.data._internal.util import GiB
+        >>> _format_resource_bundle_for_log({"CPU": 8, "GPU": 0, "memory": 32 * GiB})
+        '{CPU: 8, memory: 32.0GiB}'
+    """
+    resources = []
+    for resource_name in _RESOURCE_LOG_KEYS:
+        value = bundle.get(resource_name, 0)
+        if value == 0:
+            continue
+        resources.append(
+            f"{resource_name}: {_format_resource_value_for_log(resource_name, value)}"
+        )
+    return "{" + ", ".join(resources) + "}"
+
+
+def _format_resources_for_log(resources: List[ResourceDict]) -> str:
+    """Format and aggregate resource bundles for logging.
+
+    Bundles that format to the same string (after dropping custom/zero-valued
+    resources) are collapsed into a single ``N x {...}`` entry.
+
+    Args:
+        resources: The resource bundles to format.
+
+    Returns:
+        A human-readable string, e.g. ``"[2 x {CPU: 1}, 1 x {GPU: 1}]"``.
+
+    Example:
+        >>> _format_resources_for_log([{"CPU": 1}, {"CPU": 1}, {"GPU": 1}])
+        '[2 x {CPU: 1}, 1 x {GPU: 1}]'
+    """
+    bundle_counts: Dict[str, int] = {}
+    for resource in resources:
+        bundle = _format_resource_bundle_for_log(resource)
+        if bundle == "{}":
+            continue
+        bundle_counts[bundle] = bundle_counts.get(bundle, 0) + 1
+
+    return (
+        "["
+        + ", ".join(f"{count} x {bundle}" for bundle, count in bundle_counts.items())
+        + "]"
+    )
+
+
+@dataclass
+class OngoingRequest:
+    """Represents an ongoing resource request from a requester."""
+
+    # The time when the request was first received.
+    first_request_time: float
+    # Requested resources.
+    requested_resources: List[ResourceDict]
+    # The expiration time of the request.
+    expiration_time: float
+    # If true, after allocating requested resources to each requester,
+    # remaining resources will also be allocated to this requester.
+    request_remaining: bool
+    # The priority of the request, higher value means higher priority.
+    priority: int
+    # Resources that are already allocated to the requester.
+    allocated_resources: List[ResourceDict]
+    # Per-bundle label selectors, parallel to ``requested_resources``.
+    # Empty dicts mean no label constraint on that bundle. Required to have
+    # the same length as ``requested_resources``.
+    requested_label_selectors: List[Dict[str, str]]
+
+    def __lt__(self, other):
+        """Used to sort requests when allocating resources.
+
+        Higher priority first, then earlier first_request_time first.
+        """
+        if self.priority != other.priority:
+            return self.priority > other.priority
+        return self.first_request_time < other.first_request_time
+
+
+class DefaultAutoscalingCoordinator(AutoscalingCoordinator):
+    """Non-blocking client-side proxy for the _AutoscalingCoordinatorActor.
+
+    Not thread-safe; all methods must be called from a single thread.
+
+    Create one instance per requester. Multiple instances sharing the same
+    ``requester_id`` will have diverging caches and break the FIFO ordering
+    guarantee that ``request_resources`` and ``get_allocated_resources`` rely on.
+    """
+
+    def __init__(
+        self,
+        requester_id: str,
+        autoscaling_coordinator_actor=None,  # For testing only: injects an actor instead of using the shared named singleton.
+        subcluster_selector: Optional[Dict[str, str]] = None,
+    ):
+        self._requester_id = requester_id
+        # Label selector keyed by ``SUBCLUSTER_LABEL_KEY`` pinning this
+        # requester to a single subcluster.
+        self._subcluster_selector = subcluster_selector
+        self._cached_allocated_resources: List[ResourceDict] = []
+        # In-flight get_allocated_resources ref, or None if no request is pending.
+        self._pending_allocated_resources: Optional[ray.ObjectRef] = None
+        if autoscaling_coordinator_actor is not None:
+            # Bypass the cached_property by injecting the actor directly.
+            # Used in tests to avoid the shared named actor.
+            self.__dict__["_autoscaling_coordinator"] = autoscaling_coordinator_actor
+
+    @functools.cached_property
+    def _autoscaling_coordinator(self):
+        # Lazy: avoids creating the actor in __init__.
+        return get_or_create_autoscaling_coordinator()
+
+    def request_resources(
+        self,
+        resources: List[ResourceDict],
+        expire_after_s: float,
+        request_remaining: bool = False,
+        priority: ResourceRequestPriority = ResourceRequestPriority.MEDIUM,
+        label_selectors: Optional[List[Dict[str, str]]] = None,
+    ) -> None:
+        """Fire-and-forget: submit a resource request to the coordinator actor.
+
+        Actor-side errors are not surfaced to the caller.
+        """
+        self._autoscaling_coordinator.request_resources.remote(
+            requester_id=self._requester_id,
+            resources=resources,
+            expire_after_s=expire_after_s,
+            request_remaining=request_remaining,
+            priority=priority,
+            label_selectors=label_selectors,
+            subcluster_selector=self._subcluster_selector,
+        )
+
+    def cancel_request(self) -> None:
+        """Fire-and-forget: cancel a resource request on the coordinator actor.
+
+        Also clears client-side state (pending ref and cached allocation) so
+        a subsequent ``get_allocated_resources`` call returns a fresh result
+        rather than stale data from a prior pipeline run.
+        """
+        self._pending_allocated_resources = None
+        self._cached_allocated_resources = []
+        self._autoscaling_coordinator.cancel_request.remote(self._requester_id)
+
+    def get_allocated_resources(self) -> List[ResourceDict]:
+        """Return allocated resources without blocking.
+
+        Submits an async RPC and immediately returns the last cached result.
+        The cache is updated the next time the pending RPC completes.
+
+        Because the actor processes calls in FIFO order, the result always
+        reflects state after all previously submitted ``request_resources`` calls
+        to the same actor.
+
+        On actor errors, returns the cached value and logs a warning; never raises.
+        """
+        ref = self._pending_allocated_resources
+        if ref is not None:
+            ready, _ = ray.wait([ref], timeout=0)
+            if ready:
+                self._pending_allocated_resources = None
+                try:
+                    self._cached_allocated_resources = ray.get(ref, timeout=0)
+                except ray.exceptions.RayError:
+                    logger.warning(
+                        f"Failed to get allocated resources for {self._requester_id};"
+                        " falling back to the cached value."
+                        " If this persists, file a GitHub issue.",
+                        exc_info=RAY_DATA_AUTOSCALING_COORDINATOR_LOG_TRACEBACK,
+                    )
+
+        # Submit a new request if none is currently in-flight
+        # (first call, or the previous request completed or errored).
+        if self._pending_allocated_resources is None:
+            self._pending_allocated_resources = (
+                self._autoscaling_coordinator.get_allocated_resources.remote(
+                    self._requester_id,
+                )
+            )
+
+        return self._cached_allocated_resources
+
+
+def _default_send_resources_request(
+    bundles: List[ResourceDict],
+    label_selectors: Optional[List[Dict[str, str]]] = None,
+) -> None:
+    """Default ``send_resources_request`` implementation for the actor."""
+    ray.autoscaler.sdk.request_resources(
+        bundles=bundles, bundle_label_selectors=label_selectors
+    )
+
+
+class _AutoscalingCoordinatorActor:
+    """An actor to coordinate autoscaling resource requests from different components.
+
+    This actor is responsible for:
+    * Merging received requests and dispatching them to Ray Autoscaler.
+    * Allocating cluster resources to the requesters.
+    """
+
+    TICK_INTERVAL_S = 20
+
+    def __init__(
+        self,
+        get_current_time: Callable[[], float] = time.time,
+        send_resources_request: Callable[
+            [List[ResourceDict], Optional[List[Dict[str, str]]]], None
+        ] = _default_send_resources_request,
+        get_cluster_nodes: Callable[[], List[Dict]] = ray.nodes,
+    ):
+        self._get_current_time = get_current_time
+        self._send_resources_request = send_resources_request
+        self._get_cluster_nodes = get_cluster_nodes
+
+        self._ongoing_reqs: Dict[str, OngoingRequest] = {}
+        # Map from requester id to its subcluster selector.
+        self._subcluster_selectors: Dict[str, Optional[Dict[str, str]]] = {}
+        # Node resources bucketed by their ``SUBCLUSTER_LABEL_KEY`` value.
+        # Nodes without the key fall under ``DEFAULT_SUBCLUSTER``.
+        self._cluster_node_resources: Dict[Optional[str], List[ResourceDict]] = {}
+        # Lock for thread-safe access to shared state from the background
+        self._lock = threading.Lock()
+        self._update_cluster_node_resources()
+
+        # This is an actor, so the following check should always be True.
+        # It's only needed for unit tests.
+        if ray.is_initialized():
+            # Start a thread to perform periodical operations.
+            def tick_thread_run():
+                while True:
+                    time.sleep(self.TICK_INTERVAL_S)
+                    self._tick()
+
+            self._tick_thread = threading.Thread(target=tick_thread_run, daemon=True)
+            self._tick_thread.start()
+
+    def _tick(self):
+        """Used to perform periodical operations, e.g., purge expired requests,
+        merge and send requests, check cluster resource updates, etc."""
+        with self._lock:
+            self._merge_and_send_requests()
+            self._update_cluster_node_resources()
+            self._reallocate_resources()
+
+    def request_resources(
+        self,
+        requester_id: str,
+        resources: List[ResourceDict],
+        expire_after_s: float,
+        request_remaining: bool = False,
+        priority: ResourceRequestPriority = ResourceRequestPriority.MEDIUM,
+        label_selectors: Optional[List[Dict[str, str]]] = None,
+        subcluster_selector: Optional[Dict[str, str]] = None,
+    ) -> None:
+        logger.debug(
+            "Received request from %s: %s "
+            "(label_selectors=%s, subcluster_selector=%s).",
+            requester_id,
+            resources,
+            label_selectors,
+            subcluster_selector,
+        )
+        if label_selectors is None:
+            label_selectors = [{} for _ in resources]
+        elif len(label_selectors) != len(resources):
+            raise ValueError(
+                f"label_selectors length ({len(label_selectors)}) must match "
+                f"resources length ({len(resources)})."
+            )
+        if subcluster_selector and label_selectors:
+            req_subcluster = subcluster_selector.get(SUBCLUSTER_LABEL_KEY)
+            for i, sel in enumerate(label_selectors):
+                bundle_subcluster = sel.get(SUBCLUSTER_LABEL_KEY)
+                if (
+                    bundle_subcluster is not None
+                    and bundle_subcluster != req_subcluster
+                ):
+                    raise ValueError(
+                        f"Bundle {i} label_selector targets subcluster "
+                        f"{bundle_subcluster!r}, but requester is registered to "
+                        f"{req_subcluster!r}. Per-bundle cross-subcluster "
+                        f"allocation is not supported."
+                    )
+        with self._lock:
+            now = self._get_current_time()
+            request_updated = False
+            old_req = self._ongoing_reqs.get(requester_id)
+            if old_req is not None:
+                if request_remaining != old_req.request_remaining:
+                    raise ValueError(
+                        "Cannot change request_remaining flag of an ongoing request."
+                    )
+                if priority.value != old_req.priority:
+                    raise ValueError("Cannot change priority of an ongoing request.")
+                if (
+                    requester_id in self._subcluster_selectors
+                    and self._subcluster_selectors[requester_id] != subcluster_selector
+                ):
+                    raise ValueError(
+                        "Cannot change subcluster_selector of an ongoing request "
+                        f"from {self._subcluster_selectors[requester_id]!r} to "
+                        f"{subcluster_selector!r}."
+                    )
+
+                request_updated = (
+                    resources != old_req.requested_resources
+                    or label_selectors != old_req.requested_label_selectors
+                )
+                old_req.requested_resources = resources
+                old_req.requested_label_selectors = label_selectors
+                old_req.expiration_time = now + expire_after_s
+            else:
+                request_updated = True
+                self._ongoing_reqs[requester_id] = OngoingRequest(
+                    first_request_time=now,
+                    requested_resources=resources,
+                    requested_label_selectors=label_selectors,
+                    request_remaining=request_remaining,
+                    priority=priority.value,
+                    expiration_time=now + expire_after_s,
+                    allocated_resources=[],
+                )
+            # Write subcluster after all validations so a rejected call
+            # never leaves the registry on a new subcluster.
+            self._subcluster_selectors[requester_id] = subcluster_selector
+            if request_updated:
+                # If the request has updated, immediately send
+                # a new request and reallocate resources.
+                self._merge_and_send_requests()
+                self._reallocate_resources()
+
+    def cancel_request(
+        self,
+        requester_id: str,
+    ):
+        logger.debug("Canceling request for %s.", requester_id)
+        with self._lock:
+            if requester_id not in self._ongoing_reqs:
+                return
+            del self._ongoing_reqs[requester_id]
+            self._subcluster_selectors.pop(requester_id, None)
+            self._merge_and_send_requests()
+            self._reallocate_resources()
+
+    def _purge_expired_requests(self):
+        now = self._get_current_time()
+        live = {
+            requester_id: req
+            for requester_id, req in self._ongoing_reqs.items()
+            if req.expiration_time > now
+        }
+        for expired_id in self._ongoing_reqs.keys() - live.keys():
+            self._subcluster_selectors.pop(expired_id, None)
+        self._ongoing_reqs = live
+
+    def _merge_and_send_requests(self):
+        """Merge requests and send them to Ray Autoscaler.
+
+        Each bundle's forwarded selector is the union of its per-bundle
+        ``requested_label_selectors`` entry and the requester's
+        ``subcluster_selector``. The subcluster pin wins on key conflict,
+        so the autoscaler always sees the correct subcluster regardless
+        of what the per-bundle selectors contain.
+        """
+        self._purge_expired_requests()
+        merged_req: List[ResourceDict] = []
+        merged_selectors: List[Dict[str, str]] = []
+        for requester_id, req in self._ongoing_reqs.items():
+            merged_req.extend(req.requested_resources)
+            subcluster_selector = self._subcluster_selectors.get(requester_id) or {}
+            for per_bundle in req.requested_label_selectors:
+                merged_selectors.append({**per_bundle, **subcluster_selector})
+        if any(merged_selectors):
+            self._send_resources_request(merged_req, label_selectors=merged_selectors)
+        else:
+            self._send_resources_request(merged_req)
+
+    def get_allocated_resources(self, requester_id: str) -> List[ResourceDict]:
+        """Get the allocated resources for the requester."""
+        with self._lock:
+            if requester_id not in self._ongoing_reqs:
+                return []
+            return self._ongoing_reqs[requester_id].allocated_resources
+
+    def _maybe_subtract_resources(self, res1: ResourceDict, res2: ResourceDict) -> bool:
+        """If res2<=res1, subtract res2 from res1 in-place, and return True.
+        Otherwise return False."""
+        if any(res1.get(key, 0) < res2[key] for key in res2):
+            return False
+        for key in res2:
+            if key in res1:
+                res1[key] -= res2[key]
+        return True
+
+    def _update_cluster_node_resources(self) -> bool:
+        """Update cluster resources bucketed by subcluster. Return True if changed."""
+
+        def _is_node_eligible(node):
+            # Exclude dead nodes.
+            if not node["Alive"]:
+                return False
+            resources = node["Resources"]
+            # Exclude the head node if it doesn't have CPUs and GPUs,
+            # because the object store is not usable.
+            if HEAD_NODE_RESOURCE_LABEL in resources and (
+                resources.get("CPU", 0) == 0 and resources.get("GPU", 0) == 0
+            ):
+                return False
+            return True
+
+        nodes = list(filter(_is_node_eligible, self._get_cluster_nodes()))
+        nodes = sorted(nodes, key=lambda node: node.get("NodeID", ""))
+        cluster_node_resources: Dict[Optional[str], List[ResourceDict]] = {}
+        for node in nodes:
+            # Safeguard against case where the value of Labels is None.
+            labels = node.get("Labels") or {}
+            subcluster = labels.get(SUBCLUSTER_LABEL_KEY, DEFAULT_SUBCLUSTER)
+            cluster_node_resources.setdefault(subcluster, []).append(node["Resources"])
+        if cluster_node_resources == self._cluster_node_resources:
+            return False
+        logger.debug("Cluster resources updated: %s.", cluster_node_resources)
+        self._cluster_node_resources = cluster_node_resources
+        return True
+
+    def _reallocate_resources(self):
+        """Reallocate cluster resources.
+
+        Each requester's subcluster comes from its ``subcluster_selector``.
+        A requester without one is eligible only for the ``None`` bucket.
+        """
+        now = self._get_current_time()
+        cluster_node_resources: Dict[Optional[str], List[ResourceDict]] = copy.deepcopy(
+            self._cluster_node_resources
+        )
+        live_items = [
+            (req_id, req)
+            for req_id, req in self._ongoing_reqs.items()
+            if req.expiration_time >= now
+        ]
+        live_items.sort(key=lambda item: item[1])
+
+        def _subcluster_of(requester_id: str) -> Optional[str]:
+            selector = self._subcluster_selectors.get(requester_id)
+            return (selector or {}).get(SUBCLUSTER_LABEL_KEY, DEFAULT_SUBCLUSTER)
+
+        # TODO(hchen): Optimize the following triple loop.
+        for requester_id, ongoing_req in live_items:
+            ongoing_req.allocated_resources = []
+            subcluster = _subcluster_of(requester_id)
+            for bundle in ongoing_req.requested_resources:
+                for node_resource in cluster_node_resources.get(subcluster, []):
+                    if self._maybe_subtract_resources(node_resource, bundle):
+                        ongoing_req.allocated_resources.append(bundle)
+                        break
+
+        # Allocate remaining resources. Multiple concurrent requesters in
+        # the same subcluster split that subcluster's leftovers equally.
+        remaining_items = [
+            (req_id, req) for req_id, req in live_items if req.request_remaining
+        ]
+        for subcluster, node_resources in cluster_node_resources.items():
+            eligible = [
+                req
+                for req_id, req in remaining_items
+                if _subcluster_of(req_id) == subcluster
+            ]
+            if not eligible:
+                continue
+            for node_resource in node_resources:
+                # Integer division may leave some resources unallocated.
+                divided = {k: v // len(eligible) for k, v in node_resource.items()}
+                if not any(v > 0 for v in divided.values()):
+                    continue
+                for r in eligible:
+                    r.allocated_resources.append(divided)
+
+        if logger.isEnabledFor(logging.DEBUG):
+            msg = "Allocated resources:\n"
+            for requester_id, ongoing_req in self._ongoing_reqs.items():
+                allocated_resources_log_str = _format_resources_for_log(
+                    ongoing_req.allocated_resources
+                )
+                msg += f"Requester {requester_id}: {allocated_resources_log_str}\n"
+            logger.debug(msg)
+
+
+_get_or_create_lock = threading.Lock()
+
+
+def get_or_create_autoscaling_coordinator():
+    """Get or create the AutoscalingCoordinator actor."""
+    # Create the actor on the local node,
+    # to reduce network overhead.
+    scheduling_strategy = NodeAffinitySchedulingStrategy(
+        ray.get_runtime_context().get_node_id(),
+        soft=False,
+    )
+    actor_cls = ray.remote(num_cpus=0, max_restarts=-1, max_task_retries=-1)(
+        _AutoscalingCoordinatorActor
+    ).options(
+        name="AutoscalingCoordinator",
+        namespace="AutoscalingCoordinator",
+        get_if_exists=True,
+        lifetime="detached",
+        scheduling_strategy=scheduling_strategy,
+    )
+    # NOTE: Need the following lock, because Ray Core doesn't allow creating the same
+    # actor from multiple threads simultaneously.
+    with _get_or_create_lock:
+        return actor_cls.remote()

python/ray/data/_internal/cluster_autoscaler/default_cluster_autoscaler_v2.py

@@ -0,0 +1,428 @@
+import logging
+import math
+import time
+from collections import Counter, defaultdict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+
+import ray
+from .base_autoscaling_coordinator import AutoscalingCoordinator, ResourceDict
+from .default_autoscaling_coordinator import (
+    DEFAULT_SUBCLUSTER,
+    SUBCLUSTER_LABEL_KEY,
+    DefaultAutoscalingCoordinator,
+)
+from .resource_utilization_gauge import (
+    ResourceUtilizationGauge,
+    RollingLogicalUtilizationGauge,
+)
+from .util import cap_resource_request_to_limits, is_autoscaling_enabled
+from ray._common.utils import env_bool, env_float, env_integer
+from ray.data._internal.cluster_autoscaler import ClusterAutoscaler
+from ray.data._internal.execution.interfaces.execution_options import ExecutionResources
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.util import GiB
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class _NodeResourceSpec:
+    cpu: int
+    gpu: int
+    mem: int
+
+    def __post_init__(self):
+        assert isinstance(self.cpu, int)
+        assert self.cpu >= 0
+        assert isinstance(self.gpu, int)
+        assert self.gpu >= 0
+        assert isinstance(self.mem, int)
+        assert self.mem >= 0
+
+    def __str__(self):
+        return (
+            "{"
+            + f"CPU: {self.cpu}, GPU: {self.gpu}, memory: {memory_string(self.mem)}"
+            + "}"
+        )
+
+    @classmethod
+    def of(cls, *, cpu=0, gpu=0, mem=0):
+        cpu = math.floor(cpu)
+        gpu = math.floor(gpu)
+        # Round memory to the nearest 0.1 GiB so that nodes of the same type
+        # with slightly different reported physical memory are grouped together.
+        mem = int(round(mem / GiB, 1) * GiB) if mem > 0 else 0
+        return cls(cpu=cpu, gpu=gpu, mem=mem)
+
+    @classmethod
+    def from_bundle(cls, bundle: Dict[str, Any]) -> "_NodeResourceSpec":
+        return _NodeResourceSpec.of(
+            cpu=bundle.get("CPU", 0),
+            gpu=bundle.get("GPU", 0),
+            mem=bundle.get("memory", 0),
+        )
+
+    def to_bundle(self):
+        return {"CPU": self.cpu, "GPU": self.gpu, "memory": self.mem}
+
+
+def _get_node_resource_spec_and_count(
+    subcluster: Optional[str] = DEFAULT_SUBCLUSTER,
+) -> Dict[_NodeResourceSpec, int]:
+    """Get the unique node resource specs and their count in the cluster,
+    scoped to a single subcluster.
+
+    The returned specs are the scalable worker shapes used to build scale-up
+    requests, so the head node is deliberately excluded:
+
+      * Head node *instances* are not counted (they can't be used to schedule tasks).
+      * A node group *config* dedicated to the head node is dropped as well
+        (``max_count == 1`` and a shape matching the running head node).
+        Otherwise its shape would be requested as an extra node even though the
+        head can't be scaled. Groups that can also host workers
+        (``max_count > 1``) or that have a non-head shape are kept, including
+        worker types with zero running instances (scale-from-zero).
+
+    Quirk: the returned dict also contains a ``node_type: 0`` (ex: `"m5.xlarge": 0`) entry for every
+    node type registered in ``cluster_config.node_group_configs`` that
+    isn't included in this subcluster. ``get_cluster_config()``
+    reports node types but not labels, so the only way to know a
+    shape's subcluster is to inspect live nodes. Harmless: for example,
+    if m5.xlarge nodes only exist in the training subcluster, the validation
+    dataset will emit pending-bundle scale-up demand for foo nodes
+    stamped with the validation label, which the autoscaler can never
+    satisfy.
+    TODO: get labels from cluster config so the catalog can be filtered.
+
+    Args:
+        subcluster: The value at ``SUBCLUSTER_LABEL_KEY`` to match against.
+            The default ``DEFAULT_SUBCLUSTER`` (None) selects nodes with no
+            subcluster label.
+
+    Returns:
+        A mapping from each scalable worker node shape to its current count of
+        running instances (0 for shapes discovered only from the cluster
+        config).
+    """
+    nodes_resource_spec_count = defaultdict(int)
+
+    # Split nodes into the head node and worker nodes. There is exactly one head
+    # node, and it can't be scaled up, so it's excluded from the counts and used
+    # below to identify a node group dedicated to the head. Worker nodes are
+    # further scoped to the requester's subcluster, so foreign subclusters'
+    # shapes and counts don't leak into this requester's active / pending
+    # bundles. Head detection is intentionally not subcluster-scoped: the head
+    # node group is global.
+    head_node_spec = None
+    worker_node_resources = []
+    for node in ray.nodes():
+        if not node["Alive"]:
+            continue
+        if "node:__internal_head__" in node["Resources"]:
+            head_node_spec = _NodeResourceSpec.from_bundle(node["Resources"])
+            continue
+        if (node.get("Labels") or {}).get(
+            SUBCLUSTER_LABEL_KEY, DEFAULT_SUBCLUSTER
+        ) == subcluster:
+            worker_node_resources.append(node["Resources"])
+
+    cluster_config = ray._private.state.state.get_cluster_config()
+    if cluster_config and cluster_config.node_group_configs:
+        for node_group_config in cluster_config.node_group_configs:
+            if not node_group_config.resources or node_group_config.max_count == 0:
+                continue
+
+            node_resource_spec = _NodeResourceSpec.from_bundle(
+                node_group_config.resources
+            )
+            # Skip a node group dedicated to the head node, since it can't be scaled up and thus shouldn't be counted towards the current cluster capacity or used as a template for scaling up.
+            if (
+                node_group_config.max_count == 1
+                and node_resource_spec == head_node_spec
+            ):
+                continue
+
+            nodes_resource_spec_count[node_resource_spec] = 0
+
+    for r in worker_node_resources:
+        node_resource_spec = _NodeResourceSpec.from_bundle(r)
+        nodes_resource_spec_count[node_resource_spec] += 1
+
+    return nodes_resource_spec_count
+
+
+class DefaultClusterAutoscalerV2(ClusterAutoscaler):
+    """Ray Data's second cluster autoscaler implementation.
+
+    It works in the following way:
+
+      * Check the average cluster utilization (CPU and memory)
+        in a time window (by default 10s). If the utilization is above a threshold (by
+        default 0.75), send a request to Ray's autoscaler to scale up the cluster.
+      * Unlike previous implementation, each resource bundle in the request is a node
+        resource spec, rather than an `incremental_resource_usage()`. This allows us
+        to directly scale up nodes.
+      * Cluster scaling down isn't handled here. It depends on the idle node
+        termination.
+    """
+
+    # Default cluster utilization threshold to trigger scaling up.
+    DEFAULT_CLUSTER_SCALING_UP_UTIL_THRESHOLD: float = env_float(
+        "RAY_DATA_CLUSTER_SCALING_UP_UTIL_THRESHOLD",
+        0.75,
+    )
+    # Default time window in seconds to calculate the average of cluster utilization.
+    DEFAULT_CLUSTER_UTIL_AVG_WINDOW_S: int = env_integer(
+        "RAY_DATA_CLUSTER_UTIL_AVG_WINDOW_S",
+        10,
+    )
+    # Default number of nodes to add per node type.
+    DEFAULT_CLUSTER_SCALING_UP_DELTA: int = env_integer(
+        "RAY_DATA_CLUSTER_SCALING_UP_DELTA",
+        1,
+    )
+
+    # Min number of seconds between two autoscaling requests.
+    MIN_GAP_BETWEEN_AUTOSCALING_REQUESTS: int = env_integer(
+        "RAY_DATA_MIN_GAP_BETWEEN_AUTOSCALING_REQUESTS",
+        10,
+    )
+    # The time in seconds after which an autoscaling request will expire.
+    AUTOSCALING_REQUEST_EXPIRE_TIME_S: int = env_integer(
+        "RAY_DATA_AUTOSCALING_REQUEST_EXPIRE_TIME_S",
+        180,
+    )
+    # When utilization drops below the scale-up threshold, keep renewing the last
+    # explicit request for a short time before releasing it.
+    DEFAULT_LOW_UTIL_REQUEST_RELEASE_DELAY_S: float = env_float(
+        "RAY_DATA_LOW_UTIL_REQUEST_RELEASE_DELAY_S",
+        180,
+    )
+    # Whether to disable INFO-level logs.
+    RAY_DATA_DISABLE_AUTOSCALER_LOGGING = env_bool(
+        "RAY_DATA_DISABLE_AUTOSCALER_LOGGING", False
+    )
+
+    def __init__(
+        self,
+        resource_manager: "ResourceManager",
+        execution_id: str,
+        resource_limits: ExecutionResources = ExecutionResources.inf(),
+        resource_utilization_calculator: Optional[ResourceUtilizationGauge] = None,
+        cluster_scaling_up_util_threshold: float = DEFAULT_CLUSTER_SCALING_UP_UTIL_THRESHOLD,  # noqa: E501
+        cluster_scaling_up_delta: float = DEFAULT_CLUSTER_SCALING_UP_DELTA,
+        cluster_util_avg_window_s: float = DEFAULT_CLUSTER_UTIL_AVG_WINDOW_S,
+        min_gap_between_autoscaling_requests_s: float = MIN_GAP_BETWEEN_AUTOSCALING_REQUESTS,  # noqa: E501
+        low_util_request_release_delay_s: float = DEFAULT_LOW_UTIL_REQUEST_RELEASE_DELAY_S,  # noqa: E501
+        autoscaling_coordinator: Optional[AutoscalingCoordinator] = None,
+        get_node_counts: Optional[Callable[[], Dict[_NodeResourceSpec, int]]] = None,
+        get_time: Callable[[], float] = time.time,
+        label_selector: Optional[Dict[str, str]] = None,
+    ):
+        assert cluster_scaling_up_delta > 0
+        assert cluster_util_avg_window_s > 0
+        assert min_gap_between_autoscaling_requests_s >= 0
+        assert low_util_request_release_delay_s >= 0
+
+        if resource_utilization_calculator is None:
+            resource_utilization_calculator = RollingLogicalUtilizationGauge(
+                resource_manager,
+                cluster_util_avg_window_s=cluster_util_avg_window_s,
+                execution_id=execution_id,
+            )
+
+        self._resource_limits = resource_limits
+        self._label_selector = label_selector or {}
+        self._resource_utilization_calculator = resource_utilization_calculator
+        # Threshold of cluster utilization to trigger scaling up.
+        self._cluster_scaling_up_util_threshold = cluster_scaling_up_util_threshold
+        self._cluster_scaling_up_delta = int(math.ceil(cluster_scaling_up_delta))
+        self._min_gap_between_autoscaling_requests_s = (
+            min_gap_between_autoscaling_requests_s
+        )
+        self._low_util_request_release_delay_s = low_util_request_release_delay_s
+        # Last time when a request was sent to Ray's autoscaler.
+        self._last_request_time = 0
+        # Track the last non-empty explicit request so low-utilization heartbeats
+        # can keep it alive briefly without turning allocated remaining-share
+        # resources into explicit autoscaler demand.
+        self._last_non_empty_resource_request: List[ResourceDict] = []
+        self._last_non_empty_request_time: Optional[float] = None
+        # Unique identifier for the cluster autoscaler as a requester for
+        # the autoscaling coordinator.
+        self._requester_id = f"data-{execution_id}"
+        if autoscaling_coordinator is None:
+            autoscaling_coordinator = DefaultAutoscalingCoordinator(
+                requester_id=self._requester_id,
+                subcluster_selector=label_selector,
+            )
+        self._autoscaling_coordinator = autoscaling_coordinator
+        if get_node_counts is None:
+            # Scope node-shape/count discovery to this requester's subcluster
+            # so try_trigger_scaling doesn't pull node shapes / counts from
+            # other subclusters into ``active_bundles`` / ``pending_bundles``.
+            subcluster = self._label_selector.get(
+                SUBCLUSTER_LABEL_KEY, DEFAULT_SUBCLUSTER
+            )
+            get_node_counts = lambda: _get_node_resource_spec_and_count(  # noqa: E731
+                subcluster=subcluster
+            )
+        self._get_node_counts = get_node_counts
+        self._get_time = get_time
+        self._autoscaling_enabled = is_autoscaling_enabled()
+
+        # Register with the coordinator immediately so the actor knows about this
+        # requester before the first ``get_allocated_resources call``. The cached value
+        # returned by ``get_allocated_resources`` (and thus ``get_total_resources``) will
+        # be empty until the actor responds with the first allocation (cold-start).
+        self._send_resource_request([])
+
+    def try_trigger_scaling(self):
+        # Note, should call this method before checking `_last_request_time`,
+        # in order to update the average cluster utilization.
+        self._resource_utilization_calculator.observe()
+
+        # Limit the frequency of autoscaling requests.
+        now = self._get_time()
+        if now - self._last_request_time < self._min_gap_between_autoscaling_requests_s:
+            return
+
+        util = self._resource_utilization_calculator.get()
+        if (
+            util.cpu < self._cluster_scaling_up_util_threshold
+            and util.gpu < self._cluster_scaling_up_util_threshold
+            and util.memory < self._cluster_scaling_up_util_threshold
+            and util.object_store_memory < self._cluster_scaling_up_util_threshold
+        ):
+            logger.debug(
+                "Cluster utilization is below threshold: "
+                f"CPU={util.cpu:.2f}, GPU={util.gpu:.2f}, memory={util.memory:.2f}, "
+                f"object_store_memory={util.object_store_memory:.2f}."
+            )
+            self._send_resource_request(None)
+            return
+
+        # We separate active bundles (existing nodes) from pending bundles (scale-up delta)
+        # to ensure existing nodes' resources are never crowded out by scale-up requests.
+        # TODO(hchen): We scale up all nodes by the same delta for now.
+        # We may want to distinguish different node types based on their individual
+        # utilization.
+        active_bundles = []
+        pending_bundles = []
+        node_resource_spec_count = self._get_node_counts()
+        for node_resource_spec, count in node_resource_spec_count.items():
+            bundle = node_resource_spec.to_bundle()
+            # Bundles for existing nodes -> active (must include)
+            active_bundles.extend([bundle] * count)
+            # Bundles for scale-up delta -> pending (best-effort)
+            pending_bundles.extend([bundle] * self._cluster_scaling_up_delta)
+
+        # Cap the resource request to respect user-configured limits.
+        # Active bundles (existing nodes) are always included; pending bundles
+        # (scale-up requests) are best-effort.
+        resource_request = cap_resource_request_to_limits(
+            active_bundles, pending_bundles, self._resource_limits
+        )
+
+        if resource_request != active_bundles:
+            self._log_resource_request(util, active_bundles, resource_request)
+
+        self._send_resource_request(resource_request)
+
+    def _log_resource_request(
+        self,
+        current_utilization: ExecutionResources,
+        active_bundles: List[Dict[str, float]],
+        resource_request: List[Dict[str, float]],
+    ) -> None:
+        message = (
+            "The utilization of one or more logical resource is higher than the "
+            f"specified threshold of {self._cluster_scaling_up_util_threshold:.0%}: "
+            f"CPU={current_utilization.cpu:.0%}, GPU={current_utilization.gpu:.0%}, "
+            f"memory={current_utilization.memory:.0%}, "
+            f"object_store_memory={current_utilization.object_store_memory:.0%}. "
+            f"Requesting {self._cluster_scaling_up_delta} node(s) of each shape:"
+        )
+
+        current_node_counts = Counter(
+            [_NodeResourceSpec.from_bundle(bundle) for bundle in active_bundles]
+        )
+        requested_node_counts = Counter(
+            [_NodeResourceSpec.from_bundle(bundle) for bundle in resource_request]
+        )
+        for node_spec, requested_count in requested_node_counts.items():
+            current_count = current_node_counts.get(node_spec, 0)
+            message += f" [{node_spec}: {current_count} -> {requested_count}]"
+
+        if self.RAY_DATA_DISABLE_AUTOSCALER_LOGGING or not self._autoscaling_enabled:
+            level = logging.DEBUG
+        else:
+            level = logging.INFO
+
+        logger.log(level, message)
+
+    def _should_keep_non_empty_request(self, now: float) -> bool:
+        return (
+            self._last_non_empty_request_time is not None
+            and now - self._last_non_empty_request_time
+            < self._low_util_request_release_delay_s
+        )
+
+    def _send_resource_request(
+        self,
+        resource_request: Optional[List[ResourceDict]],
+    ):
+        now = self._get_time()
+        update_non_empty_request_state = True
+        if resource_request is None:
+            if self._should_keep_non_empty_request(now):
+                resource_request = self._last_non_empty_resource_request
+                update_non_empty_request_state = False
+            else:
+                # Renew our registration on AutoscalingCoordinator without
+                # keeping explicit autoscaler demand alive.
+                resource_request = []
+
+        # Make autoscaler resource request.
+        self._autoscaling_coordinator.request_resources(
+            resources=resource_request,
+            expire_after_s=self.AUTOSCALING_REQUEST_EXPIRE_TIME_S,
+            request_remaining=True,
+        )
+        if resource_request and update_non_empty_request_state:
+            self._last_non_empty_resource_request = [
+                bundle.copy() for bundle in resource_request
+            ]
+            self._last_non_empty_request_time = now
+        elif not resource_request:
+            self._last_non_empty_resource_request = []
+            self._last_non_empty_request_time = None
+        self._last_request_time = now
+
+    def on_executor_shutdown(self):
+        # Cancel the resource request when the executor is shutting down.
+        try:
+            self._autoscaling_coordinator.cancel_request()
+        except Exception:
+            # cancel_request is fire-and-forget and shouldn't raise, but guard
+            # against unexpected Ray Core errors at submit time. At shutdown
+            # there's nothing useful to do except log and let the request expire.
+            msg = (
+                f"Failed to cancel resource request for {self._requester_id}."
+                " The request will still expire after the timeout of"
+                f" {self._min_gap_between_autoscaling_requests_s} seconds."
+            )
+            logger.warning(msg, exc_info=True)
+
+    def get_total_resources(self) -> ExecutionResources:
+        """Get total resources available from the autoscaling coordinator."""
+        resources = self._autoscaling_coordinator.get_allocated_resources()
+        total = ExecutionResources.zero()
+        for res in resources:
+            total = total.add(ExecutionResources.from_resource_dict(res))
+        return total

python/ray/data/_internal/cluster_autoscaler/fake_autoscaling_coordinator.py

@@ -0,0 +1,83 @@
+import time
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional
+
+from .base_autoscaling_coordinator import (
+    AutoscalingCoordinator,
+    ResourceDict,
+    ResourceRequestPriority,
+)
+
+
+class FakeAutoscalingCoordinator(AutoscalingCoordinator):
+    """A lightweight implementation for testing.
+
+    This implementation always allocates the requested resources to the
+    requester. It doesn't support the `priority` parameter.
+    """
+
+    @dataclass
+    class Allocation:
+        resources: List[ResourceDict]
+        expiration_time_s: float
+        request_remaining: bool
+
+    def __init__(
+        self,
+        get_time: Callable[[], float] = time.time,
+        initial_cluster_resources: Optional[List[ResourceDict]] = None,
+    ):
+        """Initialize the coordinator.
+
+        Args:
+            get_time: A function that returns the current time in seconds. This is a
+                seam for testing.
+            initial_cluster_resources: If the requester sends an empty request and
+                ``request_remaining`` is True, the coordinator allocates these resources
+                to the requester. Otherwise, the coordinator allocates the requested
+                resources.
+        """
+        if initial_cluster_resources is None:
+            initial_cluster_resources = []
+
+        self._get_time = get_time
+        self._initial_cluster_resources = initial_cluster_resources
+        self._allocation: Optional[FakeAutoscalingCoordinator.Allocation] = None
+
+    def request_resources(
+        self,
+        resources: List[ResourceDict],
+        expire_after_s: float,
+        request_remaining: bool = False,
+        priority: ResourceRequestPriority = ResourceRequestPriority.MEDIUM,
+        label_selectors: Optional[List[Dict[str, str]]] = None,
+        subcluster_selector: Optional[Dict[str, str]] = None,
+    ) -> None:
+        if priority != ResourceRequestPriority.MEDIUM:
+            raise NotImplementedError(
+                "This fake implementation doesn't support the `priority` parameter."
+            )
+
+        if not resources and request_remaining:
+            resources = [r.copy() for r in self._initial_cluster_resources]
+
+        # Always accept the request and record it.
+        self._allocation = self.Allocation(
+            resources=resources,
+            expiration_time_s=self._get_time() + expire_after_s,
+            request_remaining=request_remaining,
+        )
+
+    def cancel_request(self) -> None:
+        self._allocation = None
+
+    def get_allocated_resources(self) -> List[ResourceDict]:
+        """Return the allocated resources if they haven't expired."""
+        if self._allocation is None:
+            return []
+
+        if self._allocation.expiration_time_s < self._get_time():
+            self._allocation = None
+            return []
+
+        return [r.copy() for r in self._allocation.resources]

python/ray/data/_internal/cluster_autoscaler/resource_utilization_gauge.py

@@ -0,0 +1,141 @@
+import abc
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+from ray.data._internal.average_calculator import TimeWindowAverageCalculator
+from ray.data._internal.execution.resource_manager import ResourceManager
+from ray.util.metrics import Gauge
+
+
+@dataclass(frozen=True)
+class ClusterUtil:
+    cpu: float = 0.0
+    gpu: float = 0.0
+    memory: float = 0.0
+    object_store_memory: float = 0.0
+
+    def __post_init__(self):
+        # If we overcommit tasks, the logical utilization can exceed 1.0.
+        assert math.isfinite(self.cpu) and 0 <= self.cpu, self.cpu
+        assert math.isfinite(self.gpu) and 0 <= self.gpu, self.gpu
+        assert math.isfinite(self.memory) and 0 <= self.memory, self.memory
+        assert (
+            math.isfinite(self.object_store_memory) and 0 <= self.object_store_memory
+        ), self.object_store_memory
+
+
+class ResourceUtilizationGauge(abc.ABC):
+    @abc.abstractmethod
+    def observe(self):
+        """Observe the cluster utilization."""
+        ...
+
+    @abc.abstractmethod
+    def get(self) -> ClusterUtil:
+        """Get the resource cluster utilization."""
+        ...
+
+
+class RollingLogicalUtilizationGauge(ResourceUtilizationGauge):
+    # Default time window in seconds to calculate the average of cluster utilization.
+    DEFAULT_CLUSTER_UTIL_AVG_WINDOW_S: int = 10
+
+    def __init__(
+        self,
+        resource_manager: ResourceManager,
+        *,
+        cluster_util_avg_window_s: float = DEFAULT_CLUSTER_UTIL_AVG_WINDOW_S,
+        execution_id: Optional[str] = None,
+    ):
+        self._resource_manager = resource_manager
+        self._execution_id = execution_id
+
+        self._cluster_cpu_util_calculator = TimeWindowAverageCalculator(
+            cluster_util_avg_window_s
+        )
+        self._cluster_gpu_util_calculator = TimeWindowAverageCalculator(
+            cluster_util_avg_window_s
+        )
+        self._cluster_mem_util_calculator = TimeWindowAverageCalculator(
+            cluster_util_avg_window_s
+        )
+        self._cluster_obj_mem_util_calculator = TimeWindowAverageCalculator(
+            cluster_util_avg_window_s
+        )
+
+        self._cluster_cpu_utilization_gauge = None
+        self._cluster_gpu_utilization_gauge = None
+        self._cluster_mem_utilization_gauge = None
+        self._cluster_object_store_memory_utilization_gauge = None
+
+        if self._execution_id is not None:
+            self._cluster_cpu_utilization_gauge = Gauge(
+                "data_cluster_cpu_utilization",
+                description="Cluster utilization % (CPU)",
+                tag_keys=("dataset",),
+            )
+            self._cluster_gpu_utilization_gauge = Gauge(
+                "data_cluster_gpu_utilization",
+                description="Cluster utilization % (GPU)",
+                tag_keys=("dataset",),
+            )
+            self._cluster_mem_utilization_gauge = Gauge(
+                "data_cluster_mem_utilization",
+                description="Cluster utilization % (Memory)",
+                tag_keys=("dataset",),
+            )
+            self._cluster_object_store_memory_utilization_gauge = Gauge(
+                "data_cluster_object_store_memory_utilization",
+                description="Cluster utilization % (Object Store Memory)",
+                tag_keys=("dataset",),
+            )
+
+    def observe(self):
+        """Report the cluster utilization based on global usage / global limits."""
+
+        def save_div(numerator, denominator):
+            if not denominator:
+                return 0
+            else:
+                return numerator / denominator
+
+        global_usage = self._resource_manager.get_global_usage()
+        global_limits = self._resource_manager.get_global_limits()
+
+        cpu_util = save_div(global_usage.cpu, global_limits.cpu)
+        gpu_util = save_div(global_usage.gpu, global_limits.gpu)
+        mem_util = save_div(global_usage.memory, global_limits.memory)
+        obj_store_mem_util = save_div(
+            global_usage.object_store_memory, global_limits.object_store_memory
+        )
+
+        self._cluster_cpu_util_calculator.report(cpu_util)
+        self._cluster_gpu_util_calculator.report(gpu_util)
+        self._cluster_mem_util_calculator.report(mem_util)
+        self._cluster_obj_mem_util_calculator.report(obj_store_mem_util)
+
+        if self._execution_id is not None:
+            tags = {"dataset": self._execution_id}
+            if self._cluster_cpu_utilization_gauge is not None:
+                self._cluster_cpu_utilization_gauge.set(cpu_util * 100, tags=tags)
+            if self._cluster_gpu_utilization_gauge is not None:
+                self._cluster_gpu_utilization_gauge.set(gpu_util * 100, tags=tags)
+            if self._cluster_mem_utilization_gauge is not None:
+                self._cluster_mem_utilization_gauge.set(mem_util * 100, tags=tags)
+            if self._cluster_object_store_memory_utilization_gauge is not None:
+                self._cluster_object_store_memory_utilization_gauge.set(
+                    obj_store_mem_util * 100, tags=tags
+                )
+
+    def get(self) -> ClusterUtil:
+        """Get the average cluster utilization based on global usage / global limits."""
+        # Clamp to 0 to handle floating-point drift in the rolling average.
+        return ClusterUtil(
+            cpu=max(0, self._cluster_cpu_util_calculator.get_average() or 0),
+            gpu=max(0, self._cluster_gpu_util_calculator.get_average() or 0),
+            memory=max(0, self._cluster_mem_util_calculator.get_average() or 0),
+            object_store_memory=max(
+                0, self._cluster_obj_mem_util_calculator.get_average() or 0
+            ),
+        )

python/ray/data/_internal/cluster_autoscaler/throughput_solver.py

@@ -0,0 +1,127 @@
+from typing import Dict, TypeVar
+
+from ray.data._internal.execution.interfaces import ExecutionResources
+
+# The math functions defined in this module use a generic type rather than
+# `PhysicalOperator` so it's easier to test. We already pass in all of the necessary
+# inputs, so the actual type doesn't matter.
+T = TypeVar("T")
+
+
+_SCHEDULABLE_RESOURCE_NAMES = ("cpu", "gpu", "memory")
+
+
+def allocate_resources(
+    throughput: float,
+    *,
+    rates: Dict[T, float],
+    resource_requirements: Dict[T, ExecutionResources],
+) -> Dict[T, ExecutionResources]:
+    """Allocate resources for a pipeline to sustain the given throughput.
+
+    Key insight: in a pipeline, all operators must sustain the same throughput T.
+    Operator i with per-task rate r_i needs T/r_i tasks to sustain T. So maximizing
+    throughput is equivalent to finding the largest feasible T, then deriving task
+    counts from it.
+
+    Args:
+        throughput: The throughput for the pipeline in the same units as the rates.
+        rates: The rate at which a task or actor produces outputs for each operator.
+        resource_requirements: The logical resources required to schedule a task or
+            actor for each operator.
+
+    Returns:
+        A dictionary mapping operators to the allocated resources.
+    """
+    assert throughput >= 0, "Throughput must be non-negative"
+    assert all(rate > 0 for rate in rates.values()), "Rates must be positive"
+
+    if not rates:
+        return {}
+
+    if throughput == 0:
+        return {op: ExecutionResources.zero() for op in rates}
+
+    # NOTE: This implementation computes fractional task counts. In practice, you
+    # can't schedule a fractional task or actor, so the allocations might be infeasible.
+    task_counts = {op: throughput / rate for op, rate in rates.items()}
+    return {op: resource_requirements[op].scale(task_counts[op]) for op in rates}
+
+
+def compute_optimal_throughput(
+    *,
+    rates: Dict[T, float],
+    resource_requirements: Dict[T, ExecutionResources],
+    resource_limits: ExecutionResources,
+    concurrency_limits: Dict[T, int | None],
+) -> float:
+    """Compute the optimal throughput for a pipeline.
+
+    The optimal throughput is bounded by two constraints (we take the tightest):
+      1. Resource limits — total resource usage across all operators must fit the
+         budget.
+      2. Concurrency limits — each operator's task count cannot exceed its limit.
+
+    Args:
+        rates: The rate at which a task or actor produces outputs for each operator.
+        resource_requirements: The logical resources required to schedule a task or
+            actor for each operator.
+        resource_limits: The resource limits for the cluster.
+        concurrency_limits: The maximum number of tasks or actors that can be scheduled
+            concurrently for each operator.
+
+    Returns:
+        The optimal throughput for the pipeline in the same units as the rates.
+    """
+    assert rates, "Rates must be non-empty"
+    return min(
+        _max_throughput_from_resources(rates, resource_requirements, resource_limits),
+        _max_throughput_from_concurrency(rates, concurrency_limits),
+    )
+
+
+def _max_throughput_from_resources(
+    rates: Dict[T, float],
+    resource_requirements: Dict[T, ExecutionResources],
+    resource_limits: ExecutionResources,
+) -> float:
+    """For each resource type, compute the max throughput the resource budget allows."""
+    assert rates, "Rates must be non-empty"
+    assert all(rate > 0 for rate in rates.values()), "Rates must be positive"
+    assert (
+        rates.keys() <= resource_requirements.keys()
+    ), "You must provide a resource requirement for each operator with a rate."
+
+    max_throughput = float("inf")
+
+    for resource_name in _SCHEDULABLE_RESOURCE_NAMES:
+        resource_limit = getattr(resource_limits, resource_name)
+        resource_cost_per_unit_throughput = sum(
+            getattr(resource_requirements[op], resource_name) / rates[op]
+            for op in rates
+        )
+        if resource_cost_per_unit_throughput > 0:
+            max_throughput = min(
+                max_throughput, resource_limit / resource_cost_per_unit_throughput
+            )
+
+    assert max_throughput >= 0, "Max throughput must be non-negative"
+    return max_throughput
+
+
+def _max_throughput_from_concurrency(
+    rates: Dict[T, float],
+    concurrency_limits: Dict[T, int | None],
+) -> float:
+    """Each operator's throughput is capped at rate * concurrency_limit."""
+    assert rates, "Rates must be non-empty"
+    assert (
+        rates.keys() <= concurrency_limits.keys()
+    ), "You must provide a concurrency limit for each operator with a rate."
+
+    # Convert `None` to float("inf") for operators with no concurrency limit
+    normalized_concurrency_limits: Dict[T, float] = {
+        op: limit if limit is not None else float("inf")
+        for op, limit in concurrency_limits.items()
+    }
+    return min(rates[op] * normalized_concurrency_limits[op] for op in rates)

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

@@ -0,0 +1,97 @@
+import logging
+from typing import Dict, List
+
+from ray.data._internal.execution.interfaces import ExecutionResources
+
+logger = logging.getLogger(__name__)
+
+
+def is_autoscaling_enabled() -> bool:
+    """Check if any node type has autoscaling enabled (can scale up).
+
+    A node type is autoscalable if max_count == -1 (unlimited) or
+    max_count > min_count. If no cluster config is available or no node type
+    is autoscalable, returns False.
+    """
+    import ray._private.state
+
+    cluster_config = ray._private.state.state.get_cluster_config()
+    if not cluster_config or not cluster_config.node_group_configs:
+        return False
+    return any(
+        ngc.max_count == -1 or ngc.max_count > ngc.min_count
+        for ngc in cluster_config.node_group_configs
+        if ngc.resources and ngc.max_count != 0
+    )
+
+
+def cap_resource_request_to_limits(
+    active_bundles: List[Dict],
+    pending_bundles: List[Dict],
+    resource_limits: ExecutionResources,
+) -> List[Dict]:
+    """Cap the resource request to not exceed user-configured resource limits.
+
+    Active bundles (for running tasks or existing nodes) are always included first
+    since they represent resources already in use. Pending bundles (for future work
+    or scale-up requests) are then added best-effort, sorted smallest-first to
+    maximize packing within limits.
+
+    This ensures that resources for already-running tasks are never crowded out
+    by pending work from smaller operators.
+
+    Args:
+        active_bundles: Bundles for already-running tasks or existing nodes
+            (must include - these represent current resource usage).
+        pending_bundles: Bundles for pending work or scale-up requests
+            (best-effort - only added if within limits).
+        resource_limits: The user-configured resource limits.
+
+    Returns:
+        A list of resource bundles that respects user limits, with active bundles
+        always included first.
+    """
+    # If no explicit limits are set (all infinite), return everything
+    if resource_limits == ExecutionResources.inf():
+        return active_bundles + pending_bundles
+
+    # Always include active bundles first - they're already running/allocated
+    capped_request = list(active_bundles)
+    total = ExecutionResources.zero()
+    for bundle in active_bundles:
+        total = total.add(ExecutionResources.from_resource_dict(bundle))
+
+    # Sort pending bundles by size (smallest first) to maximize packing.
+    # This ensures smaller bundles aren't excluded due to larger bundles
+    # appearing earlier in arbitrary iteration order.
+    def bundle_sort_key(bundle: Dict) -> tuple:
+        return (
+            bundle.get("CPU", 0),
+            bundle.get("GPU", 0),
+            bundle.get("memory", 0),
+        )
+
+    sorted_pending = sorted(pending_bundles, key=bundle_sort_key)
+
+    for bundle in sorted_pending:
+        new_total = total.add(ExecutionResources.from_resource_dict(bundle))
+
+        # Skip bundles that don't fit, continue checking smaller ones
+        if not new_total.satisfies_limit(resource_limits):
+            continue
+
+        capped_request.append(bundle)
+        total = new_total
+
+    total_input = len(active_bundles) + len(pending_bundles)
+    if len(capped_request) < total_input:
+        logger.debug(
+            f"Capped autoscaling resource request from {total_input} "
+            f"bundles to {len(capped_request)} bundles to respect "
+            f"user-configured resource limits: {resource_limits}. "
+            f"({len(active_bundles)} active bundles kept, "
+            f"{len(capped_request) - len(active_bundles)}/{len(pending_bundles)} "
+            f"pending bundles included)."
+        )
+
+    return capped_request

python/ray/data/_internal/collections.py

@@ -0,0 +1,49 @@
+from typing import Dict, TypeVar
+
+K = TypeVar("K")
+
+
+def collapse_transitive_map(d: Dict[K, K]) -> Dict[K, K]:
+    """Collapse transitive mappings in a dictionary. Given a mapping like
+    {a: b, b: c, c: d}, returns {a: d}, removing intermediate b -> c, c -> d.
+
+    Only keeps mappings where the key is NOT a value in another mapping (i.e., chain starting points).
+
+    Args:
+        d: Dictionary representing a mapping
+
+    Returns:
+        Dictionary with all transitive mappings collapsed, keeping only KV-pairs,
+        such that K and V are starting and terminal point of a chain
+
+    Examples:
+        >>> collapse_transitive_map({"a": "b", "b": "c", "c": "d"})
+        {'a': 'd'}
+        >>> collapse_transitive_map({"a": "b", "x": "y"})
+        {'a': 'b', 'x': 'y'}
+    """
+    if not d:
+        return {}
+
+    collapsed = {}
+    values_set = set(d.values())
+    for k in d:
+        # Skip mappings that are in the value-set, meaning that they are
+        # part of the mapping chain (for ex, {a -> b, b -> c})
+        if k in values_set:
+            continue
+
+        cur = k
+        visited = {cur}
+
+        # Follow the chain until we reach a key that's not in the mapping
+        while cur in d:
+            next = d[cur]
+            if next in visited:
+                raise ValueError(f"Detected a cycle in the mapping {d}")
+            visited.add(next)
+            cur = next
+
+        collapsed[k] = cur
+
+    return collapsed

python/ray/data/_internal/compute.py

@@ -0,0 +1,218 @@
+import logging
+from typing import Any, Callable, Iterable, Optional, TypeVar, Union
+
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.block import Block, UserDefinedFunction
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+# Block transform function applied by task and actor pools.
+BlockTransform = Union[
+    # TODO(Clark): Once Ray only supports Python 3.8+, use protocol to constrain block
+    # transform type.
+    # Callable[[Block, ...], Iterable[Block]]
+    # Callable[[Block, UserDefinedFunction, ...], Iterable[Block]],
+    Callable[[Iterable[Block], TaskContext], Iterable[Block]],
+    Callable[[Iterable[Block], TaskContext, UserDefinedFunction], Iterable[Block]],
+    Callable[..., Iterable[Block]],
+]
+
+
+@DeveloperAPI
+class ComputeStrategy:
+    pass
+
+
+@PublicAPI
+class TaskPoolStrategy(ComputeStrategy):
+    """Specify the task-based compute strategy for a Dataset transform.
+
+    TaskPoolStrategy executes dataset transformations using Ray tasks that are
+    scheduled through a pool. Provide ``size`` to cap the number of concurrent
+    tasks; leave it unset to allow Ray Data to scale the task count
+    automatically.
+    """
+
+    def __init__(
+        self,
+        size: Optional[int] = None,
+    ):
+        """Construct TaskPoolStrategy for a Dataset transform.
+
+        Args:
+            size: Specify the maximum size of the task pool.
+        """
+
+        if size is not None and size < 1:
+            raise ValueError("`size` must be >= 1", size)
+        self.size = size
+
+    def __eq__(self, other: Any) -> bool:
+        return (isinstance(other, TaskPoolStrategy) and self.size == other.size) or (
+            other == "tasks" and self.size is None
+        )
+
+    def __repr__(self) -> str:
+        return f"TaskPoolStrategy(size={self.size})"
+
+
+@PublicAPI
+class ActorPoolStrategy(ComputeStrategy):
+    """Specify the actor-based compute strategy for a Dataset transform.
+
+    ActorPoolStrategy specifies that an autoscaling pool of actors should be used
+    for a given Dataset transform. This is useful for stateful setup of callable
+    classes.
+
+    For a fixed-sized pool of size ``n``, use ``ActorPoolStrategy(size=n)``.
+
+    To autoscale from ``m`` to ``n`` actors, use
+    ``ActorPoolStrategy(min_size=m, max_size=n)``.
+
+    To autoscale from ``m`` to ``n`` actors, with an initial size of ``initial``, use
+    ``ActorPoolStrategy(min_size=m, max_size=n, initial_size=initial)``.
+
+    To increase opportunities for pipelining task dependency prefetching with
+    computation and avoiding actor startup delays, set max_tasks_in_flight_per_actor
+    to 2 or greater; to try to decrease the delay due to queueing of tasks on the worker
+    actors, set max_tasks_in_flight_per_actor to 1.
+
+    The `enable_true_multi_threading` argument primarily exists to prevent GPU OOM issues with multi-threaded actors.
+    The life cycle of an actor task involves 3 main steps:
+
+        1. Batching Inputs
+        2. Running actor UDF
+        3. Batching Outputs
+
+    The `enable_true_multi_threading` flag affects step 2. If set to `True`, then the UDF can be run concurrently.
+    By default, it is set to `False`, so at most 1 actor UDF is running at a time per actor. The `max_concurrency`
+    flag on `ray.remote` affects steps 1 and 3. Below is a matrix summary:
+
+    - [`enable_true_multi_threading=False or True`, `max_concurrency=1`] = 1 actor task running per actor. So at most 1
+        of steps 1, 2, or 3 is running at any point in time.
+    - [`enable_true_multi_threading=False`, `max_concurrency>1`] = multiple tasks running per actor
+      (respecting GIL) but UDF runs 1 at a time. This is useful for doing CPU and GPU work,
+      where you want to use a large batch size but want to hide the overhead of *batching*
+      the inputs. In this case, CPU *batching* is done concurrently, while GPU *inference*
+      is done 1 at a time. Concretely, steps 1 and 3 can have multiple threads, while step 2 is done serially.
+    - [`enable_true_multi_threading=True`, `max_concurrency>1`] = multiple tasks running per actor.
+      Unlike bullet #3 ^, the UDF runs concurrently (respecting GIL). No restrictions on steps 1, 2, or 3
+
+    NOTE: `enable_true_multi_threading` does not apply to async actors
+    """
+
+    def __init__(
+        self,
+        *,
+        size: Optional[int] = None,
+        min_size: Optional[int] = None,
+        max_size: Optional[int] = None,
+        initial_size: Optional[int] = None,
+        max_tasks_in_flight_per_actor: Optional[int] = None,
+        enable_true_multi_threading: bool = False,
+    ):
+        """Construct ActorPoolStrategy for a Dataset transform.
+
+        Args:
+            size: Specify a fixed size actor pool of this size. It is an error to
+                specify both `size` and `min_size` or `max_size`.
+            min_size: The minimum size of the actor pool.
+            max_size: The maximum size of the actor pool.
+            initial_size: The initial number of actors to start with. If not specified,
+                defaults to min_size. Must be between min_size and max_size.
+            max_tasks_in_flight_per_actor: The maximum number of tasks to concurrently
+                send to a single actor worker. Increasing this will increase
+                opportunities for pipelining task dependency prefetching with
+                computation and avoiding actor startup delays, but will also increase
+                queueing delay.
+            enable_true_multi_threading: If enable_true_multi_threading=False, no more than 1 UDF
+                runs per actor. Otherwise, respects the `max_concurrency` argument. For more details, see
+                the `ActorPoolStrategy` class docstring.
+        """
+        if size is not None:
+            if size < 1:
+                raise ValueError("size must be >= 1", size)
+            if max_size is not None or min_size is not None or initial_size is not None:
+                raise ValueError(
+                    "min_size, max_size, and initial_size cannot be set at the same time as `size`"
+                )
+            min_size = size
+            max_size = size
+            initial_size = size
+        if min_size is not None and min_size < 1:
+            raise ValueError("min_size must be >= 1", min_size)
+        if max_size is not None:
+            if min_size is None:
+                min_size = 1  # Legacy default.
+            if min_size > max_size:
+                raise ValueError("min_size must be <= max_size", min_size, max_size)
+        if (
+            max_tasks_in_flight_per_actor is not None
+            and max_tasks_in_flight_per_actor < 1
+        ):
+            raise ValueError(
+                "max_tasks_in_flight_per_actor must be >= 1, got: ",
+                max_tasks_in_flight_per_actor,
+            )
+
+        self.min_size = min_size or 1
+        self.max_size = max_size or float("inf")
+
+        # Validate and set initial_size
+        if initial_size is not None:
+            if initial_size < self.min_size:
+                raise ValueError(
+                    f"initial_size ({initial_size}) must be >= min_size ({self.min_size})"
+                )
+            if self.max_size != float("inf") and initial_size > self.max_size:
+                raise ValueError(
+                    f"initial_size ({initial_size}) must be <= max_size ({self.max_size})"
+                )
+
+        self.initial_size = initial_size or self.min_size
+        self.max_tasks_in_flight_per_actor = max_tasks_in_flight_per_actor
+        self.num_workers = 0
+        self.ready_to_total_workers_ratio = 0.8
+        self.enable_true_multi_threading = enable_true_multi_threading
+
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, ActorPoolStrategy) and (
+            self.min_size == other.min_size
+            and self.max_size == other.max_size
+            and self.initial_size == other.initial_size
+            and self.enable_true_multi_threading == other.enable_true_multi_threading
+            and self.max_tasks_in_flight_per_actor
+            == other.max_tasks_in_flight_per_actor
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"ActorPoolStrategy(min_size={self.min_size}, "
+            f"max_size={self.max_size}, "
+            f"initial_size={self.initial_size}, "
+            f"max_tasks_in_flight_per_actor={self.max_tasks_in_flight_per_actor})"
+            f"num_workers={self.num_workers}, "
+            f"enable_true_multi_threading={self.enable_true_multi_threading}, "
+            f"ready_to_total_workers_ratio={self.ready_to_total_workers_ratio})"
+        )
+
+
+def get_compute(compute_spec: Union[str, ComputeStrategy]) -> ComputeStrategy:
+    if not isinstance(compute_spec, (TaskPoolStrategy, ActorPoolStrategy)):
+        raise ValueError(
+            "In Ray 2.5, the compute spec must be either "
+            f"TaskPoolStrategy or ActorPoolStrategy, was: {compute_spec}."
+        )
+    elif not compute_spec or compute_spec == "tasks":
+        return TaskPoolStrategy()
+    elif compute_spec == "actors":
+        return ActorPoolStrategy()
+    elif isinstance(compute_spec, ComputeStrategy):
+        return compute_spec
+    else:
+        raise ValueError("compute must be one of [`tasks`, `actors`, ComputeStrategy]")

python/ray/data/_internal/dataset_repr.py

@@ -0,0 +1,483 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+import ray
+from ray.data._internal.logical.interfaces import SourceOperator
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.exceptions import RayError
+from ray.types import ObjectRef
+
+if TYPE_CHECKING:
+    from ray.data._internal.logical.interfaces.operator import Operator
+    from ray.data.dataset import Dataset, Schema
+
+_DATASET_REPR_ELLIPSIS = "…"  # Ellipsis marker for truncated cells/rows.
+_DATASET_REPR_MAX_ROWS = 10  # Total preview row budget when materialized.
+_DATASET_REPR_HEAD_ROWS = 5  # Number of head rows to show before the gap.
+_DATASET_REPR_MAX_COLUMN_WIDTH = 40  # Max width per column cell in the table.
+_DATASET_REPR_GET_TIMEOUT_S = 30.0  # Timeout for fetching preview blocks.
+
+__all__ = [
+    "build_dataset_ascii_repr",
+    "build_dataset_summary_repr",
+]
+
+
+def build_dataset_ascii_repr(
+    dataset: "Dataset",
+    schema: "Schema",
+    is_materialized: bool,
+) -> str:
+    """Render the dataset as a multi-line tabular string."""
+    columns = list(schema.names)
+    if not columns:
+        return build_dataset_summary_repr(dataset)
+
+    num_rows = dataset._meta_count()
+    head_rows: List[List[str]] = []
+    tail_rows: List[List[str]] = []
+    if is_materialized:
+        try:
+            head_data, tail_data, _ = _collect_materialized_rows_for_repr(
+                dataset, num_rows
+            )
+            head_rows = _format_rows_for_repr(head_data, columns)
+            tail_rows = _format_rows_for_repr(tail_data, columns)
+        except RayError:
+            head_rows = []
+            tail_rows = []
+
+    return _build_dataset_ascii_repr_from_rows(
+        schema=schema,
+        num_rows=num_rows,
+        dataset_name=dataset.name,
+        is_materialized=is_materialized,
+        head_rows=head_rows,
+        tail_rows=tail_rows,
+    )
+
+
+def _build_dataset_ascii_repr_from_rows(
+    *,
+    schema: "Schema",
+    num_rows: Optional[int],
+    dataset_name: Optional[str],
+    is_materialized: bool,
+    head_rows: List[List[str]],
+    tail_rows: List[List[str]],
+) -> str:
+    """Render the dataset repr given schema metadata and preview rows."""
+    columns = list(schema.names)
+    num_cols = len(columns)
+    shape_line = f"shape: ({num_rows if num_rows is not None else '?'}, {num_cols})"
+
+    # Build header rows from schema.
+    dtype_strings = [_repr_format_dtype(t) for t in schema.types]
+    column_headers = [
+        _truncate_to_cell_width(str(col), _DATASET_REPR_MAX_COLUMN_WIDTH)
+        for col in columns
+    ]
+    dtype_headers = [
+        _truncate_to_cell_width(dtype, _DATASET_REPR_MAX_COLUMN_WIDTH)
+        for dtype in dtype_strings
+    ]
+    separator_row = ["---"] * len(columns)
+
+    # Assemble rows, including an ellipsis gap if needed.
+    show_gap = bool(head_rows) and bool(tail_rows)
+    display_rows: List[List[str]] = []
+    display_rows.extend(head_rows)
+    if show_gap:
+        display_rows.append([_DATASET_REPR_ELLIPSIS] * len(columns))
+    display_rows.extend(tail_rows)
+
+    # Render the table with computed column widths.
+    column_widths = _compute_column_widths(
+        column_headers, dtype_headers, separator_row, display_rows
+    )
+
+    table_lines = _render_table_lines(
+        column_headers,
+        dtype_headers,
+        separator_row,
+        display_rows,
+        column_widths,
+    )
+
+    # Append a summary line describing row coverage.
+    num_rows_shown = len(head_rows) + len(tail_rows)
+    summary_line = (
+        f"(Showing {num_rows_shown} of {num_rows} rows)"
+        if is_materialized
+        else "(Dataset isn't materialized)"
+    )
+    if is_materialized and num_rows is None:
+        summary_line = f"(Showing {num_rows_shown} of ? rows)"
+
+    components = []
+    if dataset_name is not None:
+        components.append(f"name: {dataset_name}")
+    components.extend([shape_line, "\n".join(table_lines), summary_line])
+    return "\n".join(components)
+
+
+def _repr_format_dtype(dtype: object) -> str:
+    """Format a dtype into a compact string for the schema row.
+
+    Dtypes may come from PyArrow, pandas/NumPy, or be plain Python types.
+    """
+    if isinstance(dtype, type):
+        return dtype.__name__
+    name = getattr(dtype, "name", None)
+    if isinstance(name, str):
+        return name
+    return str(dtype)
+
+
+def _collect_materialized_rows_for_repr(
+    dataset: "Dataset",
+    num_rows: Optional[int],
+) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], bool]:
+    """Collect head/tail rows for preview and whether to show a gap row."""
+    block_entries: List[Tuple[ObjectRef, BlockMetadata]] = []
+    for ref_bundle in dataset.iter_internal_ref_bundles():
+        block_entries.extend(zip(ref_bundle.block_refs, ref_bundle.metadata))
+
+    if not block_entries:
+        return [], [], False
+
+    # Compute how many head/tail rows to show within the preview budget.
+    head_row_limit, tail_row_limit = _determine_preview_row_targets(num_rows)
+    block_cache: Dict[ObjectRef, Block] = {}
+
+    def _resolve_block(block_ref: ObjectRef) -> Block:
+        if block_ref not in block_cache:
+            block_cache[block_ref] = ray.get(
+                block_ref, timeout=_DATASET_REPR_GET_TIMEOUT_S
+            )
+        return block_cache[block_ref]
+
+    head_rows: List[Dict[str, Any]] = []
+    head_remaining = head_row_limit
+    for block_ref, _ in block_entries:
+        if head_remaining <= 0:
+            break
+        block = _resolve_block(block_ref)
+        accessor = BlockAccessor.for_block(block)
+        for row in accessor.iter_rows(public_row_format=True):
+            head_rows.append(row)
+            head_remaining -= 1
+            if head_remaining <= 0:
+                break
+
+    tail_rows: List[Dict[str, Any]] = []
+    tail_remaining = tail_row_limit
+    tail_parts: List[List[Dict[str, Any]]] = []
+    if tail_remaining > 0:
+        for block_ref, metadata in reversed(block_entries):
+            if tail_remaining <= 0:
+                break
+            block = _resolve_block(block_ref)
+            accessor = BlockAccessor.for_block(block)
+            total_rows = metadata.num_rows
+            if total_rows is None:
+                total_rows = accessor.num_rows()
+            if total_rows == 0:
+                continue
+            start = max(0, total_rows - tail_remaining)
+            sliced_block = accessor.slice(start, total_rows, copy=False)
+            slice_accessor = BlockAccessor.for_block(sliced_block)
+            block_rows = list(slice_accessor.iter_rows(public_row_format=True))
+            tail_parts.append(block_rows)
+            tail_remaining -= len(block_rows)
+            if tail_remaining <= 0:
+                break
+
+    for part in reversed(tail_parts):
+        tail_rows.extend(part)
+
+    show_gap = bool(head_rows) and bool(tail_rows)
+    return head_rows, tail_rows, show_gap
+
+
+def _determine_preview_row_targets(num_rows: Optional[int]) -> Tuple[int, int]:
+    """Compute how many head and tail rows to preview."""
+    max_rows = _DATASET_REPR_MAX_ROWS
+    if num_rows is None or num_rows <= max_rows:
+        head = num_rows if num_rows is not None else max_rows
+        return head, 0
+
+    head = min(_DATASET_REPR_HEAD_ROWS, max_rows)
+    tail = max_rows - head
+    return head, tail
+
+
+def _format_rows_for_repr(
+    rows: List[Dict[str, Any]],
+    column_names: List[str],
+) -> List[List[str]]:
+    """Format row dicts into string cell rows for table rendering."""
+    formatted_rows: List[List[str]] = []
+    for row in rows:
+        formatted_row = []
+        for column in column_names:
+            value = row.get(column)
+            formatted_value = _format_value(value)
+            formatted_row.append(
+                _truncate_to_cell_width(formatted_value, _DATASET_REPR_MAX_COLUMN_WIDTH)
+            )
+        formatted_rows.append(formatted_row)
+    return formatted_rows
+
+
+def _format_value(value: Any) -> str:
+    if isinstance(value, np.generic):
+        value = value.item()
+    return str(value).replace("\n", " ").replace("\r", " ")
+
+
+def _truncate_to_cell_width(value: str, max_width: int) -> str:
+    """Truncate a single cell to the configured max width."""
+    if max_width is None:
+        return value
+    if max_width <= 0:
+        return _DATASET_REPR_ELLIPSIS if value else ""
+    if len(value) <= max_width:
+        return value
+    if max_width == 1:
+        return _DATASET_REPR_ELLIPSIS
+    return value[: max_width - 1] + _DATASET_REPR_ELLIPSIS
+
+
+def _compute_column_widths(
+    headers: List[str],
+    dtype_headers: List[str],
+    separator_row: List[str],
+    data_rows: List[List[str]],
+) -> List[int]:
+    """Compute per-column widths for table rendering."""
+    column_widths: List[int] = []
+    for idx in range(len(headers)):
+        widths = [
+            len(headers[idx]),
+            len(dtype_headers[idx]),
+            len(separator_row[idx]),
+        ]
+        for row in data_rows:
+            widths.append(len(row[idx]))
+        column_widths.append(max(widths))
+    return column_widths
+
+
+def _render_table_lines(
+    headers: List[str],
+    dtype_headers: List[str],
+    separator_row: List[str],
+    data_rows: List[List[str]],
+    column_widths: List[int],
+) -> List[str]:
+    """Render the full table (borders, headers, data) as lines."""
+    lines: List[str] = []
+    top = _render_border("╭", "┬", "╮", "─", column_widths)
+    header_row = _render_row(headers, column_widths)
+    separator_line = _render_row(separator_row, column_widths)
+    dtype_row = _render_row(dtype_headers, column_widths)
+    lines.extend([top, header_row, separator_line, dtype_row])
+
+    if data_rows:
+        middle = _render_border("╞", "╪", "╡", "═", column_widths)
+        lines.append(middle)
+        for row in data_rows:
+            lines.append(_render_row(row, column_widths))
+
+    bottom = _render_border("╰", "┴", "╯", "─", column_widths)
+    lines.append(bottom)
+    return lines
+
+
+def _render_border(
+    left: str, middle: str, right: str, fill: str, column_widths: List[int]
+) -> str:
+    """Render a table border line given column widths."""
+    segments = [fill * (width + 2) for width in column_widths]
+    return f"{left}{middle.join(segments)}{right}"
+
+
+def _render_row(values: List[str], column_widths: List[int]) -> str:
+    """Render a single table row with padding."""
+    cells = []
+    for idx, value in enumerate(values):
+        padded = value.ljust(column_widths[idx])
+        cells.append(f" {padded} ")
+    return f"│{'┆'.join(cells)}│"
+
+
+def _format_operator_dag(
+    op: "Operator",
+    curr_str: str = "",
+    depth: int = 0,
+    including_source: bool = True,
+    show_op_repr: bool = False,
+) -> Tuple[str, int]:
+    """Traverse (DFS) the Plan DAG and
+    return a string representation of the operators."""
+    if not including_source and isinstance(op, SourceOperator):
+        return curr_str, depth
+
+    curr_max_depth = depth
+
+    # For logical plan, only show the operator name like "Aggregate".
+    # But for physical plan, show the operator class name as well like
+    # "AllToAllOperator[Aggregate]".
+    op_str = repr(op) if show_op_repr else op.name
+
+    if depth == 0:
+        curr_str += f"{op_str}\n"
+    else:
+        trailing_space = " " * ((depth - 1) * 3)
+        curr_str += f"{trailing_space}+- {op_str}\n"
+
+    for input in op.input_dependencies:
+        curr_str, input_max_depth = _format_operator_dag(
+            input, curr_str, depth + 1, including_source, show_op_repr
+        )
+        curr_max_depth = max(curr_max_depth, input_max_depth)
+    return curr_str, curr_max_depth
+
+
+def build_dataset_summary_repr(dataset: "Dataset") -> str:
+    """Create a cosmetic string representation of a dataset.
+
+    This is used for Dataset.__repr__ when no tabular preview is available.
+    Must be very cheap — never forces execution.
+    """
+    from ray.data.dataset import MaterializedDataset
+
+    dataset_cls = type(dataset)
+    logical_plan = dataset._logical_plan
+    dataset_name = dataset._dataset_name
+
+    plan_str = ""
+    plan_max_depth = 0
+
+    if not dataset._has_computed_output():
+        plan_str, plan_max_depth = _format_operator_dag(
+            logical_plan.dag, including_source=False
+        )
+
+    schema = dataset._base_schema(fetch_if_missing=False)
+    count = dataset._cache.get_num_rows(logical_plan.dag)
+
+    if schema is None or count is None:
+        has_n_ary_operator = False
+        dag = logical_plan.dag
+
+        while not isinstance(dag, SourceOperator):
+            if len(dag.input_dependencies) > 1:
+                has_n_ary_operator = True
+                break
+
+            dag = dag.input_dependencies[0]
+
+        # TODO(@bveeramani): Handle schemas for n-ary operators like `Union`.
+        if not has_n_ary_operator:
+            assert isinstance(dag, SourceOperator), dag
+            # We infer from logical plan's dag directly as we know that
+            # we don't have any cached values, so inferring is the only
+            # option left.
+            if schema is None:
+                schema = dag.infer_schema()
+            if count is None:
+                count = dag.infer_metadata().num_rows
+
+    if schema is None:
+        schema_str = "Unknown schema"
+    elif isinstance(schema, type):
+        schema_str = str(schema)
+    else:
+        schema_str = []
+        for n, t in zip(schema.names, schema.types):
+            if hasattr(t, "__name__"):
+                t = t.__name__
+            schema_str.append(f"{n}: {t}")
+        schema_str = ", ".join(schema_str)
+        schema_str = "{" + schema_str + "}"
+
+    if count is None:
+        count = "?"
+
+    num_blocks = None
+    if dataset_cls == MaterializedDataset:
+        num_blocks = logical_plan.initial_num_blocks()
+        assert num_blocks is not None
+
+    name_str = "name={}, ".format(dataset_name) if dataset_name is not None else ""
+    num_blocks_str = f"num_blocks={num_blocks}, " if num_blocks else ""
+
+    dataset_str = "{}({}{}num_rows={}, schema={})".format(
+        dataset_cls.__name__,
+        name_str,
+        num_blocks_str,
+        count,
+        schema_str,
+    )
+
+    # If the resulting string representation fits in one line, use it directly.
+    SCHEMA_LINE_CHAR_LIMIT = 80
+    MIN_FIELD_LENGTH = 10
+    INDENT_STR = " " * 3
+    trailing_space = INDENT_STR * plan_max_depth
+
+    if len(dataset_str) > SCHEMA_LINE_CHAR_LIMIT:
+        # If the resulting string representation exceeds the line char limit,
+        # first try breaking up each `Dataset` parameter into its own line
+        # and check if each line fits within the line limit. We check the
+        # `schema` param's length, since this is likely the longest string.
+        schema_str_on_new_line = f"{trailing_space}{INDENT_STR}schema={schema_str}"
+        if len(schema_str_on_new_line) > SCHEMA_LINE_CHAR_LIMIT:
+            # If the schema cannot fit on a single line, break up each field
+            # into its own line.
+            schema_str = []
+            for n, t in zip(schema.names, schema.types):
+                if hasattr(t, "__name__"):
+                    t = t.__name__
+                col_str = f"{trailing_space}{INDENT_STR * 2}{n}: {t}"
+                # If the field line exceeds the char limit, abbreviate
+                # the field name to fit while maintaining the full type
+                if len(col_str) > SCHEMA_LINE_CHAR_LIMIT:
+                    shortened_suffix = f"...: {str(t)}"
+                    # Show at least 10 characters of the field name, even if
+                    # we have already hit the line limit with the type.
+                    chars_left_for_col_name = max(
+                        SCHEMA_LINE_CHAR_LIMIT - len(shortened_suffix),
+                        MIN_FIELD_LENGTH,
+                    )
+                    col_str = f"{col_str[:chars_left_for_col_name]}{shortened_suffix}"
+                schema_str.append(col_str)
+            schema_str = ",\n".join(schema_str)
+            schema_str = "{\n" + schema_str + f"\n{trailing_space}{INDENT_STR}" + "}"
+        name_str = (
+            f"\n{trailing_space}{INDENT_STR}name={dataset_name},"
+            if dataset_name is not None
+            else ""
+        )
+        num_blocks_str = (
+            f"\n{trailing_space}{INDENT_STR}num_blocks={num_blocks},"
+            if num_blocks
+            else ""
+        )
+        dataset_str = (
+            f"{dataset_cls.__name__}("
+            f"{name_str}"
+            f"{num_blocks_str}"
+            f"\n{trailing_space}{INDENT_STR}num_rows={count},"
+            f"\n{trailing_space}{INDENT_STR}schema={schema_str}"
+            f"\n{trailing_space})"
+        )
+
+    if plan_max_depth == 0:
+        plan_str += dataset_str
+    else:
+        plan_str += f"{INDENT_STR * (plan_max_depth - 1)}+- {dataset_str}"
+    return plan_str

python/ray/data/_internal/datasource/audio_datasource.py

@@ -0,0 +1,57 @@
+import io
+from typing import TYPE_CHECKING, Iterator, List, Union
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class AudioDatasource(FileBasedDatasource):
+    _FILE_EXTENSIONS = [
+        "mp3",
+        "wav",
+        "aac",
+        "flac",
+        "ogg",
+        "m4a",
+        "wma",
+        "alac",
+        "aiff",
+        "pcm",
+        "amr",
+        "opus",
+        "ra",
+        "rm",
+        "au",
+        "mid",
+        "midi",
+        "caf",
+    ]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="soundfile", package="soundfile")
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import soundfile
+
+        # `soundfile` doesn't support reading from a `pyarrow.NativeFile` directly, so
+        # we need to read the file into memory first.
+        stream = io.BytesIO(f.read())
+        amplitude, sample_rate = soundfile.read(stream, always_2d=True, dtype="float32")
+
+        # (amplitude, channels) -> (channels, amplitude)
+        amplitude = amplitude.transpose((1, 0))
+
+        builder = DelegatingBlockBuilder()
+        builder.add({"amplitude": amplitude, "sample_rate": sample_rate})
+        yield builder.build()

python/ray/data/_internal/datasource/avro_datasource.py

@@ -0,0 +1,43 @@
+from typing import TYPE_CHECKING, Iterator, List, Union
+
+from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption
+from ray.data._internal.util import _check_import
+from ray.data.block import Block
+from ray.data.context import DataContext
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class AvroDatasource(FileBasedDatasource):
+    """A datasource that reads Avro files."""
+
+    _FILE_EXTENSIONS = ["avro"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="fastavro", package="fastavro")
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import fastavro
+
+        # Read the Avro file. This assumes the Avro file includes its schema.
+        reader = fastavro.reader(f)
+
+        ctx = DataContext.get_current()
+        output_block_size_option = OutputBlockSizeOption.of(
+            target_max_block_size=ctx.target_max_block_size
+        )
+        output_buffer = BlockOutputBuffer(output_block_size_option)
+        for record in reader:
+            output_buffer.add(record)
+            yield from output_buffer.iter_ready_blocks()
+
+        output_buffer.finalize()
+        yield from output_buffer.iter_ready_blocks()

python/ray/data/_internal/datasource/bigquery_datasink.py

@@ -0,0 +1,133 @@
+import logging
+import os
+import tempfile
+import time
+import uuid
+from typing import TYPE_CHECKING, Iterable, Optional
+
+import pyarrow.parquet as pq
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+import ray
+from ray.data._internal.datasource import bigquery_datasource
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MAX_RETRY_CNT = 10
+RATE_LIMIT_EXCEEDED_SLEEP_TIME = 11
+
+
+class BigQueryDatasink(Datasink[None]):
+    def __init__(
+        self,
+        project_id: str,
+        dataset: str,
+        max_retry_cnt: int = DEFAULT_MAX_RETRY_CNT,
+        overwrite_table: Optional[bool] = True,
+    ) -> None:
+        _check_import(self, module="google.cloud", package="bigquery")
+        _check_import(self, module="google.cloud", package="bigquery_storage")
+        _check_import(self, module="google.api_core", package="exceptions")
+
+        self.project_id = project_id
+        self.dataset = dataset
+        self.max_retry_cnt = max_retry_cnt
+        self.overwrite_table = overwrite_table
+
+    def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
+        from google.api_core import exceptions
+
+        if self.project_id is None or self.dataset is None:
+            raise ValueError("project_id and dataset are required args")
+
+        # Set up datasets to write
+        client = bigquery_datasource._create_client(project_id=self.project_id)
+        dataset_id = self.dataset.split(".", 1)[0]
+        try:
+            client.get_dataset(dataset_id)
+        except exceptions.NotFound:
+            client.create_dataset(f"{self.project_id}.{dataset_id}", timeout=30)
+            logger.info("Created dataset " + dataset_id)
+
+        # Delete table if overwrite_table is True
+        if self.overwrite_table:
+            logger.info(
+                f"Attempting to delete table {self.dataset}"
+                + " if it already exists since kwarg overwrite_table = True."
+            )
+            client.delete_table(f"{self.project_id}.{self.dataset}", not_found_ok=True)
+        else:
+            logger.info(
+                f"The write will append to table {self.dataset}"
+                + " if it already exists since kwarg overwrite_table = False."
+            )
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        def _write_single_block(block: Block, project_id: str, dataset: str) -> None:
+            from google.api_core import exceptions
+            from google.cloud import bigquery
+
+            block = BlockAccessor.for_block(block).to_arrow()
+
+            client = bigquery_datasource._create_client(project_id=project_id)
+            job_config = bigquery.LoadJobConfig(autodetect=True)
+            job_config.source_format = bigquery.SourceFormat.PARQUET
+            job_config.write_disposition = bigquery.WriteDisposition.WRITE_APPEND
+
+            with tempfile.TemporaryDirectory() as temp_dir:
+                fp = os.path.join(temp_dir, f"block_{uuid.uuid4()}.parquet")
+                pq.write_table(block, fp, compression="SNAPPY")
+
+                retry_cnt = 0
+                while retry_cnt <= self.max_retry_cnt:
+                    with open(fp, "rb") as source_file:
+                        job = client.load_table_from_file(
+                            source_file, dataset, job_config=job_config
+                        )
+                    try:
+                        logger.info(job.result())
+                        break
+                    except (exceptions.Forbidden, exceptions.TooManyRequests) as e:
+                        retry_cnt += 1
+                        if retry_cnt > self.max_retry_cnt:
+                            break
+                        logger.info(
+                            "A block write encountered a rate limit exceeded error"
+                            + f" {retry_cnt} time(s). Sleeping to try again."
+                        )
+                        logging.debug(e)
+                        time.sleep(RATE_LIMIT_EXCEEDED_SLEEP_TIME)
+
+                # Raise exception if retry_cnt exceeds max_retry_cnt
+                if retry_cnt > self.max_retry_cnt:
+                    logger.info(
+                        f"Maximum ({self.max_retry_cnt}) retry count exceeded. Ray"
+                        " will attempt to retry the block write via fault tolerance."
+                    )
+                    raise RuntimeError(
+                        f"Write failed due to {retry_cnt}"
+                        " repeated API rate limit exceeded responses. Consider"
+                        " specifying the max_retry_cnt kwarg with a higher value."
+                    )
+
+        _write_single_block = cached_remote_fn(_write_single_block)
+
+        # Launch a remote task for each block within this write task
+        ray.get(
+            [
+                _write_single_block.remote(block, self.project_id, self.dataset)
+                for block in blocks
+                if BlockAccessor.for_block(block).num_rows() > 0
+            ]
+        )

python/ray/data/_internal/datasource/bigquery_datasource.py

@@ -0,0 +1,164 @@
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+def _create_user_agent() -> str:
+    import ray
+
+    return f"ray/{ray.__version__}"
+
+
+def _create_client_info():
+    from google.api_core.client_info import ClientInfo
+
+    return ClientInfo(
+        user_agent=_create_user_agent(),
+    )
+
+
+def _create_client_info_gapic():
+    from google.api_core.gapic_v1.client_info import ClientInfo
+
+    return ClientInfo(
+        user_agent=_create_user_agent(),
+    )
+
+
+def _create_client(project_id: str):
+    from google.cloud import bigquery
+
+    return bigquery.Client(
+        project=project_id,
+        client_info=_create_client_info(),
+    )
+
+
+def _create_read_client():
+    from google.cloud import bigquery_storage
+
+    return bigquery_storage.BigQueryReadClient(
+        client_info=_create_client_info_gapic(),
+    )
+
+
+class BigQueryDatasource(Datasource):
+    def __init__(
+        self,
+        project_id: str,
+        dataset: Optional[str] = None,
+        query: Optional[str] = None,
+    ):
+        _check_import(self, module="google.cloud", package="bigquery")
+        _check_import(self, module="google.cloud", package="bigquery_storage")
+        _check_import(self, module="google.api_core", package="exceptions")
+
+        self._project_id = project_id
+        self._dataset = dataset
+        self._query = query
+
+        if query is not None and dataset is not None:
+            raise ValueError(
+                "Query and dataset kwargs cannot both be provided "
+                + "(must be mutually exclusive)."
+            )
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        from google.cloud import bigquery_storage
+
+        def _read_single_partition(stream) -> Block:
+            client = _create_read_client()
+            reader = client.read_rows(stream.name)
+            return reader.to_arrow()
+
+        if self._query:
+            query_client = _create_client(project_id=self._project_id)
+            query_job = query_client.query(self._query)
+            query_job.result()
+            destination = str(query_job.destination)
+            dataset_id = destination.split(".")[-2]
+            table_id = destination.split(".")[-1]
+        else:
+            self._validate_dataset_table_exist(self._project_id, self._dataset)
+            dataset_id = self._dataset.split(".")[0]
+            table_id = self._dataset.split(".")[1]
+
+        bqs_client = _create_read_client()
+        table = f"projects/{self._project_id}/datasets/{dataset_id}/tables/{table_id}"
+
+        if parallelism == -1:
+            parallelism = None
+        requested_session = bigquery_storage.types.ReadSession(
+            table=table,
+            data_format=bigquery_storage.types.DataFormat.ARROW,
+        )
+        read_session = bqs_client.create_read_session(
+            parent=f"projects/{self._project_id}",
+            read_session=requested_session,
+            max_stream_count=parallelism,
+        )
+
+        read_tasks = []
+        logger.info("Created streams: " + str(len(read_session.streams)))
+        if len(read_session.streams) < parallelism:
+            logger.info(
+                "The number of streams created by the "
+                + "BigQuery Storage Read API is less than the requested "
+                + "parallelism due to the size of the dataset."
+            )
+
+        for stream in read_session.streams:
+            # Create a metadata block object to store schema, etc.
+            metadata = BlockMetadata(
+                num_rows=None,
+                size_bytes=None,
+                input_files=None,
+                exec_stats=None,
+            )
+
+            # Create the read task and pass the no-arg wrapper and metadata in
+            read_task = ReadTask(
+                lambda stream=stream: [_read_single_partition(stream)],
+                metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def _validate_dataset_table_exist(self, project_id: str, dataset: str) -> None:
+        from google.api_core import exceptions
+
+        client = _create_client(project_id=project_id)
+        dataset_id = dataset.split(".")[0]
+        try:
+            client.get_dataset(dataset_id)
+        except exceptions.NotFound:
+            raise ValueError(
+                "Dataset {} is not found. Please ensure that it exists.".format(
+                    dataset_id
+                )
+            )
+
+        try:
+            client.get_table(dataset)
+        except exceptions.NotFound:
+            raise ValueError(
+                "Table {} is not found. Please ensure that it exists.".format(dataset)
+            )

python/ray/data/_internal/datasource/binary_datasource.py

@@ -0,0 +1,24 @@
+from typing import TYPE_CHECKING
+
+from ray.data._internal.arrow_block import ArrowBlockBuilder
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class BinaryDatasource(FileBasedDatasource):
+    """Binary datasource, for reading and writing binary files."""
+
+    _COLUMN_NAME = "bytes"
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        data = f.readall()
+
+        builder = ArrowBlockBuilder()
+        item = {self._COLUMN_NAME: data}
+        builder.add(item)
+        yield builder.build()
+
+    def _rows_per_file(self):
+        return 1

python/ray/data/_internal/datasource/clickhouse_datasink.py

@@ -0,0 +1,435 @@
+import logging
+import re
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Optional,
+)
+
+import pyarrow
+import pyarrow as pa
+import pyarrow.types as pat
+
+from ray.data._internal.arrow_ops.transform_pyarrow import (
+    reorder_columns_by_schema,
+)
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink, WriteReturnType
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DECIMAL_PRECISION = 38
+DEFAULT_DECIMAL_SCALE = 10
+
+
+def _pick_best_arrow_field_for_order_by(schema: pyarrow.Schema) -> str:
+    if len(schema) == 0:
+        return "tuple()"
+    # Prefer a timestamp column if available
+    for f in schema:
+        if pat.is_timestamp(f.type):
+            return f.name
+    # Next prefer a non-string column
+    for f in schema:
+        if not (pat.is_string(f.type) or pat.is_large_string(f.type)):
+            return f.name
+    # Otherwise pick the first column
+    return schema[0].name
+
+
+def _arrow_to_clickhouse_type(field: pyarrow.Field) -> str:
+    """Convert a PyArrow field to an appropriate ClickHouse column type."""
+    t = field.type
+    if pat.is_decimal(t):
+        precision = t.precision or DEFAULT_DECIMAL_PRECISION
+        scale = t.scale or DEFAULT_DECIMAL_SCALE
+        return f"Decimal({precision}, {scale})"
+    if pat.is_boolean(t):
+        return "UInt8"
+    if pat.is_int8(t):
+        return "Int8"
+    if pat.is_int16(t):
+        return "Int16"
+    if pat.is_int32(t):
+        return "Int32"
+    if pat.is_int64(t):
+        return "Int64"
+    if pat.is_uint8(t):
+        return "UInt8"
+    if pat.is_uint16(t):
+        return "UInt16"
+    if pat.is_uint32(t):
+        return "UInt32"
+    if pat.is_uint64(t):
+        return "UInt64"
+    if pat.is_float16(t):
+        return "Float32"
+    if pat.is_float32(t):
+        return "Float32"
+    if pat.is_float64(t):
+        return "Float64"
+    if pat.is_timestamp(t):
+        return "DateTime64(3)"
+    return "String"
+
+
+@dataclass
+class ClickHouseTableSettings:
+    """
+    Additional table creation instructions for ClickHouse.
+
+    Attributes:
+        engine: The engine definition for the created table. Defaults
+            to "MergeTree()".
+        order_by: The ORDER BY clause for the table.
+        partition_by: The PARTITION BY clause for the table.
+        primary_key: The PRIMARY KEY clause for the table.
+        settings: Additional SETTINGS clause for the table
+            (comma-separated or any valid string).
+    """
+
+    engine: str = "MergeTree()"
+    order_by: Optional[str] = None
+    partition_by: Optional[str] = None
+    primary_key: Optional[str] = None
+    settings: Optional[str] = None
+
+
+@PublicAPI(stability="alpha")
+class SinkMode(IntEnum):
+    """
+    Enum of possible modes for sinking data
+
+    Attributes:
+        CREATE: Create a new table; fail if that table already exists.
+        APPEND: Use an existing table if present, otherwise create one; then append data.
+        OVERWRITE: Drop the table if it already exists, then re-create it and write.
+    """
+
+    # Create a new table and fail if that table already exists.
+    CREATE = 1
+
+    # Append data to an existing table, or create one if it does not exist.
+    APPEND = 2
+
+    # Drop the table if it already exists, then re-create it and write.
+    OVERWRITE = 3
+
+
+@DeveloperAPI
+class ClickHouseDatasink(Datasink):
+    """ClickHouse Ray Datasink.
+
+    A Ray Datasink for writing data into ClickHouse, with support for distributed
+    writes and mode-based table management (create, append, or overwrite).
+
+    Args:
+        table: Fully qualified table identifier (e.g., "default.my_table").
+        dsn: A string in DSN (Data Source Name) HTTP format
+            (e.g., "clickhouse+http://username:password@host:8123/default").
+            For more information, see `ClickHouse Connection String doc
+            <https://clickhouse.com/docs/en/integrations/sql-clients/cli#connection_string>`_.
+        mode: One of SinkMode.CREATE, SinkMode.APPEND,
+            or SinkMode.OVERWRITE.
+            - **CREATE**: Create a new table; fail if that table already exists.
+              Requires a user-supplied schema if the table doesn’t already exist.
+            - **APPEND**: Use an existing table if present, otherwise create one.
+              If the table does not exist, the user must supply a schema. Data
+              is then appended to the table.
+            - **OVERWRITE**: Drop the table if it exists, then re-create it.
+              **Always requires** a user-supplied schema to define the new table.
+        schema: An optional PyArrow schema object that, if provided, will
+            override any inferred schema for table creation.
+            - If you are creating a new table (CREATE or APPEND when the table
+              doesn’t exist) or overwriting an existing table, you **must**
+              provide a schema.
+            - If you’re appending to an already-existing table, the schema is
+              not strictly required unless you want to cast data or enforce
+              column types. If omitted, the existing table definition is used.
+        client_settings: Optional ClickHouse server settings to be used with the
+            session/every request.
+        client_kwargs: Additional keyword arguments to pass to the
+            ClickHouse client.
+        table_settings: An optional dataclass with additional table creation
+            instructions (e.g., engine, order_by, partition_by, primary_key, settings).
+        max_insert_block_rows: If you have extremely large blocks, specifying
+            a limit here will chunk the insert into multiple smaller insert calls.
+            Defaults to None (no chunking).
+    """
+
+    NAME = "ClickHouse"
+
+    _CREATE_TABLE_TEMPLATE = """
+        CREATE TABLE IF NOT EXISTS {table_name} (
+            {columns}
+        )
+        ENGINE = {engine}
+        ORDER BY {order_by}
+        {additional_props}
+    """
+    _DROP_TABLE_TEMPLATE = """DROP TABLE IF EXISTS {table_name}"""
+    _CHECK_TABLE_EXISTS_TEMPLATE = """EXISTS {table_name}"""
+    _SHOW_CREATE_TABLE_TEMPLATE = """SHOW CREATE TABLE {table_name}"""
+
+    def __init__(
+        self,
+        table: str,
+        dsn: str,
+        mode: SinkMode = SinkMode.CREATE,
+        schema: Optional[pyarrow.Schema] = None,
+        client_settings: Optional[Dict[str, Any]] = None,
+        client_kwargs: Optional[Dict[str, Any]] = None,
+        table_settings: Optional[ClickHouseTableSettings] = None,
+        max_insert_block_rows: Optional[int] = None,
+    ) -> None:
+        self._table = table
+        self._dsn = dsn
+        self._mode = mode
+        self._schema = schema
+        self._client_settings = client_settings or {}
+        self._client_kwargs = client_kwargs or {}
+        self._table_settings = table_settings or ClickHouseTableSettings()
+        self._max_insert_block_rows = max_insert_block_rows
+        self._table_dropped = False
+
+    def _init_client(self):
+        _check_import(self, module="clickhouse_connect", package="clickhouse-connect")
+        import clickhouse_connect
+
+        return clickhouse_connect.get_client(
+            dsn=self._dsn,
+            settings=self._client_settings,
+            **self._client_kwargs,
+        )
+
+    def _generate_create_table_sql(
+        self,
+        schema: pyarrow.Schema,
+    ) -> str:
+        engine = self._table_settings.engine
+        if self._table_settings.order_by is not None:
+            order_by = self._table_settings.order_by
+        else:
+            order_by = _pick_best_arrow_field_for_order_by(schema)
+        additional_clauses = []
+        if self._table_settings.partition_by is not None:
+            additional_clauses.append(
+                f"PARTITION BY {self._table_settings.partition_by}"
+            )
+        if self._table_settings.primary_key is not None:
+            additional_clauses.append(
+                f"PRIMARY KEY ({self._table_settings.primary_key})"
+            )
+        if self._table_settings.settings is not None:
+            additional_clauses.append(f"SETTINGS {self._table_settings.settings}")
+        additional_props = ""
+        if additional_clauses:
+            additional_props = "\n" + "\n".join(additional_clauses)
+        columns_def = []
+        for field in schema:
+            ch_type = _arrow_to_clickhouse_type(field)
+            columns_def.append(f"`{field.name}` {ch_type}")
+        columns_str = ",\n    ".join(columns_def)
+        return self._CREATE_TABLE_TEMPLATE.format(
+            table_name=self._table,
+            columns=columns_str,
+            engine=engine,
+            order_by=order_by,
+            additional_props=additional_props,
+        )
+
+    def _table_exists(self, client) -> bool:
+        try:
+            result = client.query(
+                self._CHECK_TABLE_EXISTS_TEMPLATE.format(table_name=self._table)
+            )
+            if result is None:
+                return False
+            for helper in ("scalar", "first_item", "first_value"):
+                _check_import(
+                    self, module="clickhouse_connect", package="clickhouse-connect"
+                )
+                from clickhouse_connect.driver.exceptions import Error as CHError
+
+                if hasattr(result, helper):
+                    try:
+                        return bool(getattr(result, helper)())
+                    except (TypeError, ValueError, CHError) as exc:
+                        # Helper exists but failed – continue probing.
+                        logger.debug(
+                            "Helper %s failed: %s; will try fallbacks", helper, exc
+                        )
+            # Fallback to inspecting common container attributes.
+            for attr in ("result_rows", "rows", "data"):
+                rows = getattr(result, attr, None)
+                if rows:
+                    first = rows[0]
+                    # Unwrap an extra layer if present (i.e. [[1]] or [(1,)])
+                    if isinstance(first, (list, tuple)):
+                        first = first[0] if first else 0
+                    return bool(first)
+            return False
+        except Exception as e:
+            logger.warning(f"Could not verify if table {self._table} exists: {e}")
+            return False
+
+    def _get_existing_order_by(self, client) -> Optional[str]:
+        logger.debug(
+            f"Retrieving ORDER BY clause from SHOW CREATE TABLE for {self._table}"
+        )
+        try:
+            show_create_sql = self._SHOW_CREATE_TABLE_TEMPLATE.format(
+                table_name=self._table
+            )
+            result = client.command(show_create_sql)
+            ddl_str = str(result)
+            pattern = r"(?is)\border\s+by\s+(.*?)(?=\bengine\b|$)"
+            match = re.search(pattern, ddl_str)
+            if match:
+                return match.group(1).strip()
+            return None
+        except Exception as e:
+            logger.warning(
+                f"Could not retrieve SHOW CREATE TABLE for {self._table}: {e}"
+            )
+            return None
+
+    @property
+    def supports_distributed_writes(self) -> bool:
+        return True
+
+    def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
+        client = None
+        try:
+            client = self._init_client()
+            table_exists = self._table_exists(client)
+            schema_required = (
+                # Overwrite always needs a schema because it recreates the table
+                self._mode == SinkMode.OVERWRITE
+                # For CREATE or APPEND we need a schema only when the table is
+                # absent and will therefore be created in this call.
+                or (
+                    self._mode in (SinkMode.CREATE, SinkMode.APPEND)
+                    and not table_exists
+                )
+            )
+            if schema_required and self._schema is None:
+                if self._mode == SinkMode.OVERWRITE:
+                    raise ValueError(
+                        f"Overwriting table {self._table} requires a user‑provided schema."
+                    )
+                else:
+                    raise ValueError(
+                        f"Table {self._table} does not exist in mode='{self._mode.name}' and "
+                        "no schema was provided. Cannot create the table without a schema."
+                    )
+            # OVERWRITE MODE
+            if self._mode == SinkMode.OVERWRITE:
+                # If we plan to overwrite, drop the table if it exists,
+                # then re-create it using the user-provided schema.
+                if table_exists and self._table_settings.order_by is None:
+                    # Collect existing ORDER BY. This lets us preserve it if the user
+                    # hasn't specified one explicitly.
+                    existing_order_by = self._get_existing_order_by(client)
+                    if existing_order_by is not None:
+                        self._table_settings.order_by = existing_order_by
+                        logger.info(
+                            f"Reusing old ORDER BY from overwritten table:"
+                            f" {existing_order_by}"
+                        )
+                # DROP and CREATE the table.
+                drop_sql = self._DROP_TABLE_TEMPLATE.format(table_name=self._table)
+                logger.info(f"Mode=OVERWRITE => {drop_sql}")
+                client.command(drop_sql)
+                self._table_dropped = True
+                create_sql = self._generate_create_table_sql(self._schema)
+                client.command(create_sql)
+            # CREATE MODE
+            elif self._mode == SinkMode.CREATE:
+                # If table already exists in CREATE mode, fail immediately.
+                if table_exists:
+                    msg = (
+                        f"Table {self._table} already exists in mode='CREATE'. "
+                        "Use mode='APPEND' or 'OVERWRITE' instead."
+                    )
+                    logger.error(msg)
+                    raise ValueError(msg)
+                # Otherwise, create it (requires user-provided schema).
+                create_sql = self._generate_create_table_sql(self._schema)
+                client.command(create_sql)
+            # APPEND MODE
+            elif self._mode == SinkMode.APPEND:
+                if table_exists:
+                    # Table exists – validate or adopt ORDER BY.
+                    existing_order_by = self._get_existing_order_by(client)
+                    user_order_by = self._table_settings.order_by
+                    if user_order_by is not None:
+                        # The user explicitly set an order_by. Check if it conflicts:
+                        if existing_order_by and existing_order_by != user_order_by:
+                            raise ValueError(
+                                f"Conflict with order_by. The existing table {self._table} "
+                                f"has ORDER BY {existing_order_by}, but the user specified "
+                                f"ORDER BY {user_order_by}. Please drop or overwrite the table, "
+                                f"or use the same ordering."
+                            )
+                    elif existing_order_by:
+                        self._table_settings.order_by = existing_order_by
+                        logger.info(
+                            f"Reusing existing ORDER BY for table {self._table}: {existing_order_by}"
+                        )
+                else:
+                    # Table doesn't exist, so create it with user schema
+                    create_sql = self._generate_create_table_sql(self._schema)
+                    client.command(create_sql)
+
+        except Exception as e:
+            logger.error(
+                f"Could not complete on_write_start for table {self._table}: {e}"
+            )
+            raise e
+        finally:
+            if client:
+                client.close()
+
+    def get_name(self) -> str:
+        return self.NAME
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> WriteReturnType:
+        client = self._init_client()
+        total_inserted = 0
+        try:
+            for block in blocks:
+                arrow_table = BlockAccessor.for_block(block).to_arrow()
+                row_count = arrow_table.num_rows
+                if self._schema is not None:
+                    arrow_table = reorder_columns_by_schema(arrow_table, self._schema)
+                    arrow_table = arrow_table.cast(self._schema, safe=True)
+                if self._max_insert_block_rows is not None:
+                    max_chunk_size = self._max_insert_block_rows
+                else:
+                    # If max_insert_block_rows is not set, insert all rows in one go
+                    max_chunk_size = row_count if row_count > 0 else 1
+                offsets = list(range(0, row_count, max_chunk_size))
+                offsets.append(row_count)
+                for i in range(len(offsets) - 1):
+                    start = offsets[i]
+                    end = offsets[i + 1]
+                    chunk = arrow_table.slice(start, end - start)
+                    client.insert_arrow(self._table, chunk)
+                    total_inserted += chunk.num_rows
+        except Exception as e:
+            logger.error(f"Failed to write block(s) to table {self._table}: {e}")
+            raise e
+        finally:
+            client.close()
+        return [total_inserted]

python/ray/data/_internal/datasource/clickhouse_datasource.py

@@ -0,0 +1,363 @@
+import logging
+import math
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple
+
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+def _is_filter_string_safe(filter_str: str) -> bool:
+    in_string = False
+    escape_next = False
+    for c in filter_str:
+        if in_string:
+            # If we're inside a string, check if we're closing it.
+            if c == "'" and not escape_next:
+                in_string = False
+            escape_next = (c == "\\") and not escape_next
+        else:
+            # If we're not in a string, entering one if we see a single quote
+            if c == "'":
+                in_string = True
+                escape_next = False
+            # Disallow semicolon if we're not in a string
+            elif c == ";":
+                return False
+            else:
+                escape_next = False
+    # If we end inside a string, it's suspicious, but let's allow
+    # it to be further validated by the DB. Just return True here.
+    return True
+
+
+@DeveloperAPI
+class ClickHouseDatasource(Datasource):
+    """
+    A Ray datasource for reading from ClickHouse.
+
+    Args:
+        table: Fully qualified table or view identifier (e.g.,
+            "default.table_name").
+        dsn: A string in DSN (Data Source Name) HTTP format (e.g.,
+            "clickhouse+http://username:password@host:8124/default").
+            For more information, see `ClickHouse Connection String doc
+            <https://clickhouse.com/docs/en/integrations/sql-clients/cli#connection_string>`_.
+        columns: Optional List of columns to select from the data source.
+            If no columns are specified, all columns will be selected by default.
+        filter: Optional SQL filter string that will be used in the
+            WHERE statement (e.g., "label = 2 AND text IS NOT NULL").
+            The filter must be valid for use in a ClickHouse SQL WHERE clause.
+            Note: Parallel reads are not currently supported when a filter is set.
+            Specifying a filter forces the parallelism to 1 to ensure deterministic
+            and consistent results. For more information, see
+            `ClickHouse SQL WHERE Clause doc
+            <https://clickhouse.com/docs/en/sql-reference/statements/select/where>`_.
+        order_by: Optional Tuple containing a list of columns to order by
+            and a boolean indicating the order. Note: order_by is required to
+            support parallelism.
+        client_settings: Optional ClickHouse server settings to be used with the
+            session/every request. For more information, see
+            `ClickHouse Client Settings doc
+            <https://clickhouse.com/docs/en/integrations/python#settings-argument>`_.
+        client_kwargs: Optional Additional keyword arguments to pass to the
+            ClickHouse client. For more information,
+            see `ClickHouse Core Settings doc
+            <https://clickhouse.com/docs/en/integrations/python#additional-options>`_.
+    """
+
+    NUM_SAMPLE_ROWS = 100
+    MIN_ROWS_PER_READ_TASK = 50
+    _BASE_QUERY = "SELECT {select_clause} FROM {table}"
+    _EXPLAIN_FILTERS_QUERY = "EXPLAIN SELECT 1 FROM {table} WHERE {filter_clause}"
+    _SIZE_ESTIMATE_QUERY = "SELECT SUM(byteSize(*)) AS estimate FROM ({query})"
+    _COUNT_ESTIMATE_QUERY = "SELECT COUNT(*) AS estimate FROM ({query})"
+    _SAMPLE_BLOCK_QUERY = "{query} LIMIT {limit_row_count}"
+    _FIRST_BLOCK_QUERY = """
+        {query}
+        FETCH FIRST {fetch_row_count} {fetch_row_or_rows} ONLY
+    """
+    _NEXT_BLOCK_QUERY = """
+        {query}
+        OFFSET {offset_row_count} {offset_row_or_rows}
+        FETCH NEXT {fetch_row_count} {fetch_row_or_rows} ONLY
+    """
+
+    def __init__(
+        self,
+        table: str,
+        dsn: str,
+        columns: Optional[List[str]] = None,
+        filter: Optional[str] = None,
+        order_by: Optional[Tuple[List[str], bool]] = None,
+        client_settings: Optional[Dict[str, Any]] = None,
+        client_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        self._table = table
+        self._dsn = dsn
+        self._columns = columns
+        self._filter = filter
+        self._order_by = order_by
+        self._client_settings = client_settings or {}
+        self._client_kwargs = client_kwargs or {}
+        self._query = self._generate_query()
+
+    def _init_client(self):
+        _check_import(self, module="clickhouse_connect", package="clickhouse-connect")
+        import clickhouse_connect
+
+        return clickhouse_connect.get_client(
+            dsn=self._dsn,
+            settings=self._client_settings or {},
+            **self._client_kwargs or {},
+        )
+
+    def _validate_filter(self):
+        if not self._filter:
+            return
+        # Minimal lexical check (regex or manual approach for semicolons, etc.).
+        if not _is_filter_string_safe(self._filter):
+            raise ValueError(
+                f"Invalid characters outside of "
+                f"string literals in filter: {self._filter}"
+            )
+        # Test "EXPLAIN" query to confirm parse-ability and catch expression errors.
+        client = self._init_client()
+        try:
+            test_query = self._EXPLAIN_FILTERS_QUERY.format(
+                table=self._table,
+                filter_clause=self._filter,
+            )
+            client.query(test_query)
+        except Exception as e:
+            raise ValueError(
+                f"Invalid filter expression: {self._filter}. Error: {e}",
+            )
+        finally:
+            client.close()
+
+    def _generate_query(self) -> str:
+        query = self._BASE_QUERY.format(
+            select_clause=", ".join(self._columns) if self._columns else "*",
+            table=self._table,
+        )
+        if self._filter:
+            self._validate_filter()
+            query += f" WHERE {self._filter}"
+        if self._order_by:
+            columns, desc = self._order_by
+            direction = " DESC" if desc else ""
+            if len(columns) == 1:
+                query += f" ORDER BY {columns[0]}{direction}"
+            elif len(columns) > 1:
+                columns_clause = ", ".join(columns)
+                query += f" ORDER BY ({columns_clause}){direction}"
+        return query
+
+    def _build_block_query(self, limit_row_count: int, offset_row_count: int) -> str:
+        if offset_row_count == 0:
+            # The first block query is optimized to use FETCH FIRST clause
+            # with an OFFSET specified.
+            return self._FIRST_BLOCK_QUERY.format(
+                query=self._query,
+                fetch_row_count=limit_row_count,
+                fetch_row_or_rows="ROWS" if limit_row_count > 1 else "ROW",
+            )
+        # Subsequent block queries use OFFSET and FETCH NEXT clauses to read the
+        # next block of data.
+        return self._NEXT_BLOCK_QUERY.format(
+            query=self._query,
+            offset_row_count=offset_row_count,
+            offset_row_or_rows="ROWS" if offset_row_count > 1 else "ROW",
+            fetch_row_count=limit_row_count,
+            fetch_row_or_rows="ROWS" if limit_row_count > 1 else "ROW",
+        )
+
+    def _create_read_fn(
+        self,
+        query: str,
+    ) -> Callable[[], Iterable[Block]]:
+        def read_fn() -> Iterable[Block]:
+            return [self._execute_block_query(query)]
+
+        return read_fn
+
+    def _get_sampled_estimates(self):
+        if self._order_by is not None:
+            # If the query is ordered, we can use a FETCH clause to get a sample.
+            # This reduces the CPU overhead on ClickHouse and speeds up the
+            # estimation query.
+            query = self._FIRST_BLOCK_QUERY.format(
+                query=self._query,
+                fetch_row_count=self.NUM_SAMPLE_ROWS,
+                fetch_row_or_rows="ROWS" if self.NUM_SAMPLE_ROWS > 1 else "ROW",
+            )
+        else:
+            # If the query is not ordered, we need to use a LIMIT clause to
+            # get a sample.
+            query = self._SAMPLE_BLOCK_QUERY.format(
+                query=self._query,
+                limit_row_count=self.NUM_SAMPLE_ROWS,
+            )
+        sample_block_accessor = BlockAccessor.for_block(
+            self._execute_block_query(query)
+        )
+        estimated_size_bytes_per_row = math.ceil(
+            sample_block_accessor.size_bytes() / sample_block_accessor.num_rows()
+        )
+        sample_block_schema = sample_block_accessor.schema()
+        return estimated_size_bytes_per_row, sample_block_schema
+
+    def _get_estimate_count(self) -> Optional[int]:
+        return self._execute_estimate_query(self._COUNT_ESTIMATE_QUERY)
+
+    def _get_estimate_size(self) -> Optional[int]:
+        return self._execute_estimate_query(self._SIZE_ESTIMATE_QUERY)
+
+    def _execute_estimate_query(self, estimate_query: str) -> Optional[int]:
+        client = self._init_client()
+        try:
+            # Estimate queries wrap around the primary query, self._query.
+            # This allows us to use self._query as a sub-query to efficiently
+            # and accurately estimate the size or count of the result set.
+            query = estimate_query.format(query=self._query)
+            result = client.query(query)
+            if result and len(result.result_rows) > 0:
+                estimate = result.result_rows[0][0]
+                return int(estimate) if estimate is not None else None
+        except Exception as e:
+            logger.warning(f"Failed to execute estimate query: {e}")
+        finally:
+            client.close()
+        return None
+
+    def _execute_block_query(self, query: str) -> Block:
+        import pyarrow as pa
+
+        client = self._init_client()
+        try:
+            with client.query_arrow_stream(query) as stream:
+                record_batches = list(stream)  # Collect all record batches
+            return pa.Table.from_batches(record_batches)
+        except Exception as e:
+            raise RuntimeError(f"Failed to execute block query: {e}")
+        finally:
+            client.close()
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        """
+        Estimate the in-memory data size for the query.
+
+        Returns:
+            Estimated in-memory data size in bytes, or
+            None if the estimation cannot be performed.
+        """
+        return self._get_estimate_size()
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        """
+        Create read tasks for the ClickHouse query.
+
+        Args:
+            parallelism: The desired number of partitions to read the data into.
+                - If ``order_by`` is not set, parallelism will be forced to 1.
+                - If ``filter`` is set, parallelism will also be forced to 1
+                    to ensure deterministic results.
+            per_task_row_limit: Maximum number of rows allowed in each emitted
+                task.  Blocks larger than this limit will be sliced before
+                being yielded downstream.
+            data_context: The data context to use to get read tasks. Not used by this
+                datasource.
+
+        Returns:
+            A list of read tasks to be executed.
+        """
+        num_rows_total = self._get_estimate_count()
+        if num_rows_total == 0 or num_rows_total is None:
+            return []
+        parallelism = min(
+            parallelism, math.ceil(num_rows_total / self.MIN_ROWS_PER_READ_TASK)
+        )
+        # To ensure consistent order of query results, self._order_by
+        # must be specified and self.filter must not be specified
+        # in order to support parallelism.
+        if self._filter is not None and parallelism > 1:
+            logger.warning(
+                "ClickHouse datasource does not currently support parallel reads "
+                "when a filter is set; falling back to parallelism of 1."
+            )
+            # When filter is specified and parallelism is greater than 1,
+            # we need to reduce parallelism to 1 to ensure consistent results.
+            parallelism = 1
+        # To ensure consistent order of query results, self._order_by
+        # must be specified in order to support parallelism.
+        if self._order_by is None and parallelism > 1:
+            logger.warning(
+                "ClickHouse datasource requires dataset to be explicitly ordered "
+                "to support parallelism; falling back to parallelism of 1."
+            )
+            # When order_by is not specified and parallelism is greater than 1,
+            # we need to reduce parallelism to 1 to ensure consistent results.
+            parallelism = 1
+        # By reducing parallelism to 1 when either of the conditions above are met,
+        # we ensure the downstream process is treated exactly as a non-parallelized
+        # (single block) process would be, thus ensuring output consistency.
+        num_rows_per_block = num_rows_total // parallelism
+        num_blocks_with_extra_row = num_rows_total % parallelism
+        (
+            estimated_size_bytes_per_row,
+            sample_block_schema,
+        ) = self._get_sampled_estimates()
+
+        def _get_read_task(
+            block_rows: int, offset_rows: int, parallelized: bool
+        ) -> ReadTask:
+            if parallelized:
+                # When parallelized, we need to build a block query with OFFSET
+                # and FETCH clauses.
+                query = self._build_block_query(block_rows, offset_rows)
+            else:
+                # When not parallelized, we can use the original query without
+                # OFFSET and FETCH clauses.
+                query = self._query
+            return ReadTask(
+                self._create_read_fn(query),
+                BlockMetadata(
+                    num_rows=block_rows,
+                    size_bytes=estimated_size_bytes_per_row * block_rows,
+                    input_files=None,
+                    exec_stats=None,
+                ),
+                schema=sample_block_schema,
+                per_task_row_limit=per_task_row_limit,
+            )
+
+        if parallelism == 1:
+            # When parallelism is 1, we can read the entire dataset in a single task.
+            # We then optimize this scenario by using self._query directly without
+            # unnecessary OFFSET and FETCH clauses.
+            return [_get_read_task(num_rows_total, 0, False)]
+
+        # Otherwise we need to split the dataset into multiple tasks.
+        # Each task will include OFFSET and FETCH clauses to efficiently
+        # read a subset of the dataset.
+        read_tasks = []
+        offset = 0
+        for i in range(parallelism):
+            this_block_size = num_rows_per_block
+            if i < num_blocks_with_extra_row:
+                this_block_size += 1
+            read_tasks.append(_get_read_task(this_block_size, offset, True))
+            offset += this_block_size
+        return read_tasks

python/ray/data/_internal/datasource/csv_datasink.py

@@ -0,0 +1,36 @@
+from typing import Any, Callable, Dict, Optional
+
+import pyarrow
+
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_based_datasource import _resolve_kwargs
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class CSVDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        arrow_csv_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        arrow_csv_args: Optional[Dict[str, Any]] = None,
+        file_format="csv",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        if arrow_csv_args_fn is None:
+            arrow_csv_args_fn = lambda: {}  # noqa: E731
+
+        if arrow_csv_args is None:
+            arrow_csv_args = {}
+
+        self.arrow_csv_args_fn = arrow_csv_args_fn
+        self.arrow_csv_args = arrow_csv_args
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        from pyarrow import csv
+
+        writer_args = _resolve_kwargs(self.arrow_csv_args_fn, **self.arrow_csv_args)
+        write_options = writer_args.pop("write_options", None)
+        csv.write_csv(block.to_arrow(), file, write_options, **writer_args)

python/ray/data/_internal/datasource/csv_datasource.py

@@ -0,0 +1,82 @@
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class CSVDatasource(FileBasedDatasource):
+    """CSV datasource, for reading and writing CSV files."""
+
+    _FILE_EXTENSIONS = [
+        "csv",
+        "csv.gz",  # gzip-compressed files
+        "csv.br",  # Brotli-compressed files
+        "csv.zst",  # Zstandard-compressed files
+        "csv.lz4",  # lz4-compressed files
+    ]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        arrow_csv_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        from pyarrow import csv
+
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        if arrow_csv_args is None:
+            arrow_csv_args = {}
+
+        self.read_options = arrow_csv_args.pop(
+            "read_options", csv.ReadOptions(use_threads=False)
+        )
+        self.parse_options = arrow_csv_args.pop("parse_options", csv.ParseOptions())
+        self.arrow_csv_args = arrow_csv_args
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import pyarrow as pa
+        from pyarrow import csv
+
+        # Re-init invalid row handler: https://issues.apache.org/jira/browse/ARROW-17641
+        if hasattr(self.parse_options, "invalid_row_handler"):
+            self.parse_options.invalid_row_handler = (
+                self.parse_options.invalid_row_handler
+            )
+
+        filter_expr = (
+            self._predicate_expr.to_pyarrow()
+            if self._predicate_expr is not None
+            else None
+        )
+
+        try:
+            reader = csv.open_csv(
+                f,
+                read_options=self.read_options,
+                parse_options=self.parse_options,
+                **self.arrow_csv_args,
+            )
+            schema = None
+            while True:
+                try:
+                    batch = reader.read_next_batch()
+                    table = pa.Table.from_batches([batch], schema=schema)
+                    if schema is None:
+                        schema = table.schema
+                    if filter_expr is not None:
+                        table = table.filter(filter_expr)
+
+                    yield table
+                except StopIteration:
+                    return
+        except pa.lib.ArrowInvalid as e:
+            raise ValueError(
+                f"Failed to read CSV file: {path}. "
+                "Please check the CSV file has correct format, or filter out non-CSV "
+                "file with 'partition_filter' field. See read_csv() documentation for "
+                "more details."
+            ) from e

python/ray/data/_internal/datasource/databricks_credentials.py

@@ -0,0 +1,285 @@
+"""Databricks credential providers for Ray Data.
+
+This module provides credential abstraction for Databricks authentication,
+supporting static tokens with extensibility for future credential sources.
+"""
+
+import logging
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Optional, Union
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+# Default environment variable names for Databricks credentials
+DEFAULT_TOKEN_ENV_VAR = "DATABRICKS_TOKEN"
+DEFAULT_HOST_ENV_VAR = "DATABRICKS_HOST"
+
+
+class DatabricksCredentialProvider(ABC):
+    """Abstract base class for Databricks credential providers.
+
+    This abstraction allows different credential sources (static tokens,
+    file-based credentials, etc.) to be used with DatabricksUCDatasource.
+
+    Subclasses must implement:
+        - get_token(): Returns the current authentication token
+        - get_host(): Returns the Databricks host URL (optional)
+        - invalidate(): Clears any cached credentials
+    """
+
+    @abstractmethod
+    def get_token(self) -> str:
+        """Get the current authentication token.
+
+        Returns:
+            The Databricks authentication token string.
+
+        Raises:
+            ValueError: If no valid token is available.
+        """
+        pass
+
+    @abstractmethod
+    def get_host(self) -> str:
+        """Get the Databricks host URL.
+
+        Returns:
+            The Databricks host URL.
+
+        Raises:
+            ValueError: If no valid host is available.
+        """
+        pass
+
+    @abstractmethod
+    def invalidate(self) -> None:
+        """Invalidate any cached credentials.
+
+        This method should be called when credentials need to be refreshed,
+        such as after an authentication error.
+        """
+        pass
+
+
+class StaticCredentialProvider(DatabricksCredentialProvider):
+    """A credential provider that wraps static token and host.
+
+    This is the simplest credential provider, useful when you have a
+    token that doesn't need to be refreshed.
+
+    Args:
+        token: The Databricks authentication token.
+        host: The Databricks host URL.
+
+    Raises:
+        ValueError: If token or host is empty or None.
+    """
+
+    def __init__(self, token: str, host: str):
+        if not token:
+            raise ValueError("Token cannot be empty or None")
+        if not host:
+            raise ValueError("Host cannot be empty or None")
+        self._token = token
+        self._host = host
+
+    def get_token(self) -> str:
+        """Get the static token.
+
+        Returns:
+            The authentication token provided at construction.
+        """
+        return self._token
+
+    def get_host(self) -> str:
+        """Get the host URL.
+
+        Returns:
+            The host URL provided at construction.
+        """
+        return self._host
+
+    def invalidate(self) -> None:
+        """No-op for static credentials.
+
+        Static credentials cannot be refreshed, so this method
+        does nothing.
+        """
+        pass
+
+
+class EnvironmentCredentialProvider(DatabricksCredentialProvider):
+    """A credential provider that reads from environment variables.
+
+    Reads token and host from environment variables.
+    If host env var is not set and running in Databricks runtime,
+    automatically detects the host.
+
+    Args:
+        token_env_var: Environment variable name for the token.
+            Defaults to DEFAULT_TOKEN_ENV_VAR ("DATABRICKS_TOKEN").
+        host_env_var: Environment variable name for the host.
+            Defaults to DEFAULT_HOST_ENV_VAR ("DATABRICKS_HOST").
+
+    Raises:
+        ValueError: If token or host cannot be resolved.
+    """
+
+    def __init__(
+        self,
+        token_env_var: str = DEFAULT_TOKEN_ENV_VAR,
+        host_env_var: str = DEFAULT_HOST_ENV_VAR,
+    ):
+        self._token_env_var = token_env_var
+        self._host_env_var = host_env_var
+
+        # Validate token is set at initialization
+        token = os.environ.get(self._token_env_var)
+        if not token:
+            raise ValueError(
+                f"Environment variable '{self._token_env_var}' is not set. "
+                "Please set it to your Databricks access token."
+            )
+        self._token = token
+
+        # Resolve host: env var > Databricks runtime detection
+        host = os.environ.get(self._host_env_var) or self._detect_databricks_host()
+        if not host:
+            raise ValueError(
+                "You are not in databricks runtime, please set environment variable "
+                f"'{self._host_env_var}' to databricks workspace URL "
+                '(e.g. "adb-<workspace-id>.<random-number>.azuredatabricks.net").'
+            )
+        self._host = host
+
+    def _detect_databricks_host(self) -> Optional[str]:
+        """Detect host from Databricks runtime if available."""
+        try:
+            from ray.util.spark.utils import is_in_databricks_runtime
+
+            if is_in_databricks_runtime():
+                import IPython
+
+                ip_shell = IPython.get_ipython()
+                if ip_shell is not None:
+                    dbutils = ip_shell.ns_table["user_global"]["dbutils"]
+                    ctx = (
+                        dbutils.notebook.entry_point.getDbutils()
+                        .notebook()
+                        .getContext()
+                    )
+                    return ctx.tags().get("browserHostName").get()
+        except Exception as e:
+            logger.warning(f"Failed to detect Databricks host from runtime: {e}")
+        return None
+
+    def get_token(self) -> str:
+        """Get the token from environment variable.
+
+        Returns:
+            The authentication token from the environment.
+        """
+        return self._token
+
+    def get_host(self) -> str:
+        """Get the host from environment variable or Databricks runtime.
+
+        Returns:
+            The host URL.
+        """
+        return self._host
+
+    def invalidate(self) -> None:
+        """Re-read token from environment variable.
+
+        This allows picking up refreshed tokens when the environment
+        variable is updated (e.g., by an external token refresh process).
+        """
+        token = os.environ.get(self._token_env_var)
+        if token:
+            self._token = token
+
+
+@dataclass
+class BaseCredentialConfig:
+    credential_provider: Optional[DatabricksCredentialProvider] = None
+
+
+@dataclass
+class DatabricksTableCredentialConfig(BaseCredentialConfig):
+    pass
+
+
+@dataclass
+class UnityCatalogCredentialConfig(BaseCredentialConfig):
+    url: Optional[str] = None
+    token: Optional[str] = None
+
+
+CredentialConfig = Union[DatabricksTableCredentialConfig, UnityCatalogCredentialConfig]
+
+
+def resolve_credential_provider(
+    config: CredentialConfig,
+) -> DatabricksCredentialProvider:
+    if config.credential_provider is not None:
+        return config.credential_provider
+
+    match config:
+        case DatabricksTableCredentialConfig():
+            return EnvironmentCredentialProvider()
+        case UnityCatalogCredentialConfig(url=str(url), token=str(token)):
+            return StaticCredentialProvider(token=token, host=url)
+
+    raise ValueError(
+        "Either 'credential_provider' or both 'url' and 'token' must be provided."
+    )
+
+
+def build_headers(
+    credential_provider: DatabricksCredentialProvider,
+) -> dict[str, str]:
+    """Build request headers with fresh token from credential provider.
+
+    Args:
+        credential_provider: The credential provider to get the token from.
+
+    Returns:
+        Dictionary containing Content-Type and Authorization headers.
+    """
+    return {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {credential_provider.get_token()}",
+    }
+
+
+def request_with_401_retry(
+    request_fn: Callable[..., requests.Response],
+    url: str,
+    credential_provider: DatabricksCredentialProvider,
+    **kwargs,
+) -> requests.Response:
+    """Make an HTTP request with one retry on 401 after invalidating credentials.
+
+    Args:
+        request_fn: Request function (e.g., requests.get or requests.post)
+        url: Request URL
+        credential_provider: Credential provider for authentication
+        **kwargs: Additional arguments passed to requests
+
+    Returns:
+        Response object (after calling raise_for_status)
+    """
+    response = request_fn(url, headers=build_headers(credential_provider), **kwargs)
+
+    if response.status_code == 401:
+        logger.info("Received 401 response, invalidating credentials and retrying.")
+        credential_provider.invalidate()
+        response = request_fn(url, headers=build_headers(credential_provider), **kwargs)
+
+    response.raise_for_status()
+    return response

python/ray/data/_internal/datasource/databricks_uc_datasource.py

@@ -0,0 +1,233 @@
+import json
+import logging
+import os
+import time
+from typing import TYPE_CHECKING, List, Optional
+from urllib.parse import urljoin
+
+import numpy as np
+import pyarrow
+import requests
+
+from ray.data._internal.datasource.databricks_credentials import (
+    DatabricksCredentialProvider,
+    build_headers,
+    request_with_401_retry,
+)
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+_STATEMENT_EXEC_POLL_TIME_S = 1
+
+
+@PublicAPI(stability="alpha")
+class DatabricksUCDatasource(Datasource):
+    def __init__(
+        self,
+        warehouse_id: str,
+        catalog: str,
+        schema: str,
+        query: str,
+        credential_provider: DatabricksCredentialProvider,
+    ):
+        self._credential_provider = credential_provider
+
+        # Get host from provider (token is fetched fresh for each request)
+        self.host = self._credential_provider.get_host()
+        self.warehouse_id = warehouse_id
+        self.catalog = catalog
+        self.schema_name = schema
+        self.query = query
+
+        if not self.host.startswith(("http://", "https://")):
+            self.host = f"https://{self.host}"
+
+        url_base = f"{self.host}/api/2.0/sql/statements/"
+
+        payload = json.dumps(
+            {
+                "statement": self.query,
+                "warehouse_id": self.warehouse_id,
+                "wait_timeout": "0s",
+                "disposition": "EXTERNAL_LINKS",
+                "format": "ARROW_STREAM",
+                "catalog": self.catalog,
+                "schema": self.schema_name,
+            }
+        )
+
+        response = request_with_401_retry(
+            requests.post,
+            url_base,
+            self._credential_provider,
+            data=payload,
+        )
+        statement_id = response.json()["statement_id"]
+
+        state = response.json()["status"]["state"]
+
+        logger.info(f"Waiting for query {query!r} execution result.")
+        try:
+            while state in ["PENDING", "RUNNING"]:
+                time.sleep(_STATEMENT_EXEC_POLL_TIME_S)
+                response = request_with_401_retry(
+                    requests.get,
+                    urljoin(url_base, statement_id) + "/",
+                    self._credential_provider,
+                )
+                state = response.json()["status"]["state"]
+        except KeyboardInterrupt:
+            # User cancel the command, so we cancel query execution.
+            requests.post(
+                urljoin(url_base, f"{statement_id}/cancel"),
+                headers=build_headers(self._credential_provider),
+            )
+            try:
+                response.raise_for_status()
+            except Exception as e:
+                logger.warning(
+                    f"Canceling query {query!r} execution failed, reason: {repr(e)}."
+                )
+            raise
+
+        if state != "SUCCEEDED":
+            raise RuntimeError(
+                f"Query {self.query!r} execution failed.\n{response.json()}"
+            )
+
+        manifest = response.json()["manifest"]
+        self.is_truncated = manifest.get("truncated", False)
+
+        if self.is_truncated:
+            logger.warning(
+                f"The resulting size of the dataset of '{query!r}' exceeds "
+                "100GiB and it is truncated."
+            )
+
+        chunks = manifest.get("chunks", [])
+
+        # Make chunks metadata are ordered by index.
+        chunks = sorted(chunks, key=lambda x: x["chunk_index"])
+        num_chunks = len(chunks)
+        self.num_chunks = num_chunks
+        self._estimate_inmemory_data_size = sum(chunk["byte_count"] for chunk in chunks)
+
+        # Capture credential provider (not self) to avoid serializing entire datasource
+        credential_provider_for_tasks = self._credential_provider
+
+        def get_read_task(
+            task_index: int, parallelism: int, per_task_row_limit: Optional[int] = None
+        ):
+            # Handle empty chunk list by yielding an empty PyArrow table
+            if num_chunks == 0:
+                import pyarrow as pa
+
+                metadata = BlockMetadata(
+                    num_rows=0,
+                    size_bytes=0,
+                    input_files=None,
+                    exec_stats=None,
+                )
+
+                def empty_read_fn():
+                    yield pa.Table.from_pydict({})
+
+                return ReadTask(read_fn=empty_read_fn, metadata=metadata)
+
+            # get chunk list to be read in this task and preserve original chunk order
+            chunk_index_list = list(
+                np.array_split(range(num_chunks), parallelism)[task_index]
+            )
+
+            num_rows = sum(
+                chunks[chunk_index]["row_count"] for chunk_index in chunk_index_list
+            )
+            size_bytes = sum(
+                chunks[chunk_index]["byte_count"] for chunk_index in chunk_index_list
+            )
+
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=size_bytes,
+                input_files=None,
+                exec_stats=None,
+            )
+
+            def _read_fn():
+                for chunk_index in chunk_index_list:
+                    resolve_external_link_url = urljoin(
+                        url_base, f"{statement_id}/result/chunks/{chunk_index}"
+                    )
+
+                    resolve_response = request_with_401_retry(
+                        requests.get,
+                        resolve_external_link_url,
+                        credential_provider_for_tasks,
+                    )
+                    external_url = resolve_response.json()["external_links"][0][
+                        "external_link"
+                    ]
+                    # NOTE: do _NOT_ send the authorization header to external urls
+                    raw_response = requests.get(external_url, auth=None, headers=None)
+                    raw_response.raise_for_status()
+
+                    with pyarrow.ipc.open_stream(raw_response.content) as reader:
+                        arrow_table = reader.read_all()
+
+                    yield arrow_table
+
+            def read_fn():
+                if mock_setup_fn_path := os.environ.get(
+                    "RAY_DATABRICKS_UC_DATASOURCE_READ_FN_MOCK_TEST_SETUP_FN_PATH"
+                ):
+                    import ray.cloudpickle as pickle
+
+                    # This is for testing.
+                    with open(mock_setup_fn_path, "rb") as f:
+                        mock_setup = pickle.load(f)
+                    with mock_setup():
+                        yield from _read_fn()
+                else:
+                    yield from _read_fn()
+
+            return ReadTask(
+                read_fn=read_fn,
+                metadata=metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+
+        self._get_read_task = get_read_task
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return self._estimate_inmemory_data_size
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        # Handle empty dataset case
+        if self.num_chunks == 0:
+            return [self._get_read_task(0, 1, per_task_row_limit)]
+
+        assert parallelism > 0, f"Invalid parallelism {parallelism}"
+
+        if parallelism > self.num_chunks:
+            parallelism = self.num_chunks
+            logger.info(
+                "The parallelism is reduced to chunk number due to "
+                "insufficient chunk parallelism."
+            )
+
+        return [
+            self._get_read_task(index, parallelism, per_task_row_limit)
+            for index in range(parallelism)
+        ]

python/ray/data/_internal/datasource/delta_sharing_datasource.py

@@ -0,0 +1,138 @@
+import logging
+from json import loads
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import numpy as np
+
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+class DeltaSharingDatasource(Datasource):
+    def __init__(
+        self,
+        url: str,
+        json_predicate_hints: Optional[str] = None,
+        limit: Optional[int] = None,
+        version: Optional[int] = None,
+        timestamp: Optional[str] = None,
+    ):
+        _check_import(self, module="delta_sharing", package="delta-sharing")
+
+        if limit is not None:
+            assert (
+                isinstance(limit, int) and limit >= 0
+            ), "'limit' must be a non-negative int"
+
+        self._url = url
+        self._json_predicate_hints = json_predicate_hints
+        self._limit = limit
+        self._version = version
+        self._timestamp = timestamp
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def _read_files(self, files, converters):
+        """Read files with Delta Sharing."""
+        from delta_sharing.reader import DeltaSharingReader
+
+        for file in files:
+            yield DeltaSharingReader._to_pandas(
+                action=file, converters=converters, for_cdf=False, limit=None
+            )
+
+    def setup_delta_sharing_connections(self, url: str):
+        """
+        Set up delta sharing connections based on the url.
+
+        Args:
+            url: A URL under the format "<profile>#<share>.<schema>.<table>"
+
+        Returns:
+            A tuple of (table, rest_client) where table is a delta_sharing Table
+            object and rest_client is a DataSharingRestClient instance.
+        """
+        from delta_sharing.protocol import DeltaSharingProfile, Table
+        from delta_sharing.rest_client import DataSharingRestClient
+
+        profile_str, share, schema, table_str = _parse_delta_sharing_url(url)
+        table = Table(name=table_str, share=share, schema=schema)
+
+        profile = DeltaSharingProfile.read_from_file(profile_str)
+        rest_client = DataSharingRestClient(profile)
+        return table, rest_client
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        assert parallelism > 0, f"Invalid parallelism {parallelism}"
+        from delta_sharing.converter import to_converters
+
+        self._table, self._rest_client = self.setup_delta_sharing_connections(self._url)
+        self._response = self._rest_client.list_files_in_table(
+            self._table,
+            jsonPredicateHints=self._json_predicate_hints,
+            limitHint=self._limit,
+            version=self._version,
+            timestamp=self._timestamp,
+        )
+
+        if len(self._response.add_files) == 0 or self._limit == 0:
+            logger.warning("No files found from the delta sharing table or limit is 0")
+
+        schema_json = loads(self._response.metadata.schema_string)
+        self._converters = to_converters(schema_json)
+
+        read_tasks = []
+        # get file list to be read in this task and preserve original chunk order
+        for files in np.array_split(self._response.add_files, parallelism):
+            files = files.tolist()
+            metadata = BlockMetadata(
+                num_rows=None,
+                input_files=files,
+                size_bytes=None,
+                exec_stats=None,
+            )
+            converters = self._converters
+            read_task = ReadTask(
+                lambda f=files: self._read_files(f, converters),
+                metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+
+def _parse_delta_sharing_url(url: str) -> Tuple[str, str, str, str]:
+    """
+    Developed from delta_sharing's _parse_url function.
+    https://github.com/delta-io/delta-sharing/blob/main/python/delta_sharing/delta_sharing.py#L36
+
+    Args:
+        url: a url under the format "<profile>#<share>.<schema>.<table>"
+
+    Returns:
+        a tuple with parsed (profile, share, schema, table)
+    """
+    shape_index = url.rfind("#")
+    if shape_index < 0:
+        raise ValueError(f"Invalid 'url': {url}")
+    profile = url[0:shape_index]
+    fragments = url[shape_index + 1 :].split(".")
+    if len(fragments) != 3:
+        raise ValueError(f"Invalid 'url': {url}")
+    share, schema, table = fragments
+    if len(profile) == 0 or len(share) == 0 or len(schema) == 0 or len(table) == 0:
+        raise ValueError(f"Invalid 'url': {url}")
+    return (profile, share, schema, table)

python/ray/data/_internal/datasource/hudi_datasource.py

@@ -0,0 +1,154 @@
+import logging
+import os
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple
+
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+class HudiQueryType(Enum):
+    SNAPSHOT = "snapshot"
+    INCREMENTAL = "incremental"
+
+    @classmethod
+    def supported_types(cls) -> List[str]:
+        return [e.value for e in cls]
+
+
+class HudiDatasource(Datasource):
+    """Hudi datasource, for reading Apache Hudi table."""
+
+    def __init__(
+        self,
+        table_uri: str,
+        query_type: str,
+        filters: Optional[List[Tuple[str, str, str]]] = None,
+        hudi_options: Optional[Dict[str, str]] = None,
+        storage_options: Optional[Dict[str, str]] = None,
+    ):
+        _check_import(self, module="hudi", package="hudi-python")
+
+        self._table_uri = table_uri
+        self._query_type = HudiQueryType(query_type.lower())
+        self._filters = filters or []
+        self._hudi_options = hudi_options or {}
+        self._storage_options = storage_options or {}
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List["ReadTask"]:
+        import numpy as np
+        import pyarrow
+        from hudi import HudiTableBuilder
+
+        def _perform_read(
+            table_uri: str,
+            base_file_paths: List[str],
+            options: Dict[str, str],
+        ) -> Iterator["pyarrow.Table"]:
+            from hudi import HudiFileGroupReader
+
+            for p in base_file_paths:
+                file_group_reader = HudiFileGroupReader(table_uri, options)
+                batch = file_group_reader.read_file_slice_by_base_file_path(p)
+                yield pyarrow.Table.from_batches([batch])
+
+        hudi_table = (
+            HudiTableBuilder.from_base_uri(self._table_uri)
+            .with_hudi_options(self._hudi_options)
+            .with_storage_options(self._storage_options)
+            # Although hudi-rs supports MOR snapshot, we need to add an API in
+            # the next release to allow file group reader to take in a list of
+            # files. Hence, setting this config for now to restrain reading
+            # only on parquet files (read optimized mode).
+            # This won't affect reading COW.
+            .with_hudi_option("hoodie.read.use.read_optimized.mode", "true")
+            .build()
+        )
+
+        logger.info("Collecting file slices for Hudi table at: %s", self._table_uri)
+
+        if self._query_type == HudiQueryType.SNAPSHOT:
+            file_slices_splits = hudi_table.get_file_slices_splits(
+                parallelism, self._filters
+            )
+        elif self._query_type == HudiQueryType.INCREMENTAL:
+            start_ts = self._hudi_options.get("hoodie.read.file_group.start_timestamp")
+            end_ts = self._hudi_options.get("hoodie.read.file_group.end_timestamp")
+            # TODO(xushiyan): add table API to return splits of file slices
+            file_slices = hudi_table.get_file_slices_between(start_ts, end_ts)
+            file_slices_splits = np.array_split(file_slices, parallelism)
+        else:
+            raise ValueError(
+                f"Unsupported query type: {self._query_type}. Supported types are: {HudiQueryType.supported_types()}."
+            )
+
+        logger.info("Creating read tasks for Hudi table at: %s", self._table_uri)
+
+        reader_options = {
+            **hudi_table.storage_options(),
+            **hudi_table.hudi_options(),
+        }
+
+        schema = hudi_table.get_schema()
+        read_tasks = []
+
+        for file_slices_split in file_slices_splits:
+            num_rows = 0
+            relative_paths = []
+            input_files = []
+            size_bytes = 0
+            for file_slice in file_slices_split:
+                # A file slice in a Hudi table is a logical group of data files
+                # within a physical partition. Records stored in a file slice
+                # are associated with a commit on the Hudi table's timeline.
+                # For more info, see https://hudi.apache.org/docs/file_layouts
+                num_rows += file_slice.num_records
+                relative_path = file_slice.base_file_relative_path()
+                relative_paths.append(relative_path)
+                full_path = os.path.join(self._table_uri, relative_path)
+                input_files.append(full_path)
+                size_bytes += file_slice.base_file_size
+
+            if self._query_type == HudiQueryType.SNAPSHOT:
+                metadata = BlockMetadata(
+                    num_rows=num_rows,
+                    input_files=input_files,
+                    size_bytes=size_bytes,
+                    exec_stats=None,
+                )
+            elif self._query_type == HudiQueryType.INCREMENTAL:
+                # need the check due to
+                # https://github.com/apache/hudi-rs/issues/401
+                metadata = BlockMetadata(
+                    num_rows=None,
+                    input_files=input_files,
+                    size_bytes=None,
+                    exec_stats=None,
+                )
+
+            read_task = ReadTask(
+                read_fn=lambda paths=relative_paths: _perform_read(
+                    self._table_uri, paths, reader_options
+                ),
+                metadata=metadata,
+                schema=schema,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # TODO(xushiyan) add APIs to provide estimated in-memory size
+        return None

python/ray/data/_internal/datasource/huggingface_datasource.py

@@ -0,0 +1,194 @@
+import sys
+from typing import TYPE_CHECKING, Iterable, List, Optional, Union
+
+from ray.data._internal.tensor_extensions.arrow import pyarrow_table_from_pydict
+from ray.data._internal.util import _check_pyarrow_version
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.dataset import Dataset
+from ray.data.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import datasets
+
+    from ray.data.context import DataContext
+
+
+TRANSFORMERS_IMPORT_ERROR: Optional[ImportError] = None
+
+try:
+    # Due to HF Dataset's dynamic module system, we need to dynamically import the
+    # datasets_modules module on every actor when training.
+    # We accomplish this by simply running the following bit of code directly
+    # in the module you are currently viewing. This ensures that when we
+    # unpickle the Dataset, it runs before pickle tries to
+    # import datasets_modules and prevents an exception from being thrown.
+    # Same logic is present inside HF Transformers Ray
+    # integration: https://github.com/huggingface/transformers/blob/\
+    # 7d5fde991d598370d961be8cb7add6541e2b59ce/src/transformers/integrations.py#L271
+    # Also see https://github.com/ray-project/ray/issues/28084
+    from transformers.utils import is_datasets_available
+
+    if "datasets_modules" not in sys.modules and is_datasets_available():
+        import importlib
+        import importlib.metadata
+        import os
+
+        import datasets.load
+        from packaging.version import parse
+
+        # Datasets >= 4.0 removed dataset scripts support and the dynamic-modules cache.
+        # Only initialize dynamic modules on <= 3.x where the initializer `init_dynamic_modules` exists.
+        DATASETS_VERSION = parse(importlib.metadata.version("datasets"))
+        DATASETS_VERSION_WITHOUT_SCRIPT_SUPPORT = parse("4.0.0")
+
+        if DATASETS_VERSION < DATASETS_VERSION_WITHOUT_SCRIPT_SUPPORT:
+            dynamic_modules_path = os.path.join(
+                datasets.load.init_dynamic_modules(), "__init__.py"
+            )
+            # load dynamic_modules from path
+            spec = importlib.util.spec_from_file_location(
+                "datasets_modules", dynamic_modules_path
+            )
+            datasets_modules = importlib.util.module_from_spec(spec)
+            sys.modules[spec.name] = datasets_modules
+            spec.loader.exec_module(datasets_modules)
+except ImportError as e:
+    TRANSFORMERS_IMPORT_ERROR = e
+
+
+class HuggingFaceDatasource(Datasource):
+    """Hugging Face Dataset datasource, for reading from a
+    `Hugging Face Datasets Dataset <https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.Dataset/>`_.
+    This Datasource implements a streamed read using a
+    single read task, most beneficial for a
+    `Hugging Face Datasets IterableDataset <https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.IterableDataset/>`_
+    or datasets which are too large to fit in-memory.
+    For an in-memory Hugging Face Dataset (`datasets.Dataset`), use :meth:`~ray.data.from_huggingface`
+    directly for faster performance.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        dataset: Union["datasets.Dataset", "datasets.IterableDataset"],
+        batch_size: int = 4096,
+    ):
+        if TRANSFORMERS_IMPORT_ERROR is not None:
+            raise TRANSFORMERS_IMPORT_ERROR
+
+        self._dataset = dataset
+        self._batch_size = batch_size
+
+    @classmethod
+    def list_parquet_urls_from_dataset(
+        cls, dataset: Union["datasets.Dataset", "datasets.IterableDataset"]
+    ) -> Dataset:
+        """Return list of Hugging Face hosted parquet file URLs if they
+        exist for the data (i.e. if the dataset is a public dataset that
+        has not been transformed) else return an empty list."""
+        import datasets
+
+        # We can use the dataset name, config name, and split name to load
+        # public hugging face datasets from the Hugging Face Hub. More info
+        # here: https://huggingface.co/docs/datasets-server/parquet
+        dataset_name = dataset.info.dataset_name
+        config_name = dataset.info.config_name
+        split_name = str(dataset.split)
+
+        # If a dataset is not an iterable dataset, we will check if the
+        # dataset with the matching dataset name, config name, and split name
+        # on the Hugging Face Hub has the same fingerprint as the
+        # dataset passed into this function. If it is not matching, transforms
+        # or other operations have been performed so we cannot use the parquet
+        # files on the Hugging Face Hub, so we return an empty list.
+        if not isinstance(dataset, datasets.IterableDataset):
+            from datasets import load_dataset
+
+            try:
+                ds = load_dataset(dataset_name, config_name, split=split_name)
+                if ds._fingerprint != dataset._fingerprint:
+                    return []
+            except Exception:
+                # If an exception is thrown when trying to reload the dataset
+                # we should exit gracefully by returning an empty list.
+                return []
+
+        import requests
+
+        public_url = (
+            f"https://huggingface.co/api/datasets/{dataset_name}"
+            f"/parquet/{config_name}/{split_name}"
+        )
+        resp = requests.get(public_url)
+        if resp.status_code == requests.codes["ok"]:
+            # dataset corresponds to a public dataset, return list of parquet_files
+            return resp.json()
+        else:
+            return []
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return self._dataset.dataset_size
+
+    def _read_dataset(self) -> Iterable[Block]:
+        # Note: This is a method instead of a higher level function because
+        # we need to capture `self`. This will trigger the try-import logic at
+        # the top of file to avoid import error of dataset_modules.
+        import numpy as np
+        import pandas as pd
+        import pyarrow
+
+        for batch in self._dataset.with_format("arrow").iter(
+            batch_size=self._batch_size
+        ):
+            # HuggingFace IterableDatasets do not fully support methods like
+            # `set_format`, `with_format`, and `formatted_as`, so the dataset
+            # can return whatever is the default configured batch type, even if
+            # the format is manually overridden before iterating above.
+            # Therefore, we limit support to batch formats which have native
+            # block types in Ray Data (pyarrow.Table, pd.DataFrame),
+            # or can easily be converted to such (dict, np.array).
+            # See: https://github.com/huggingface/datasets/issues/3444
+            if not isinstance(batch, (pyarrow.Table, pd.DataFrame, dict, np.array)):
+                raise ValueError(
+                    f"Batch format {type(batch)} isn't supported. Only the "
+                    f"following batch formats are supported: "
+                    f"dict (corresponds to `None` in `dataset.with_format()`), "
+                    f"pyarrow.Table, np.array, pd.DataFrame."
+                )
+            # Ensure np.arrays are wrapped in a dict
+            # (subsequently converted to a pyarrow.Table).
+            if isinstance(batch, np.ndarray):
+                batch = {"item": batch}
+            if isinstance(batch, dict):
+                batch = pyarrow_table_from_pydict(batch)
+            # Ensure that we return the default block type.
+            block = BlockAccessor.for_block(batch).to_default()
+            yield block
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        # Note: `parallelism` arg is currently not used by HuggingFaceDatasource.
+        # We always generate a single ReadTask to perform the read.
+        _check_pyarrow_version()
+
+        # TODO(scottjlee): IterableDataset doesn't provide APIs
+        # for getting number of rows, byte size, etc., so the
+        # BlockMetadata is currently empty. Properly retrieve
+        # or calculate these so that progress bars have meaning.
+        meta = BlockMetadata(
+            num_rows=None,
+            size_bytes=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        read_tasks: List[ReadTask] = [
+            ReadTask(
+                self._read_dataset,
+                meta,
+                per_task_row_limit=per_task_row_limit,
+            )
+        ]
+        return read_tasks

python/ray/data/_internal/datasource/iceberg_datasink.py

@@ -0,0 +1,968 @@
+"""
+Module to write a Ray Dataset into an iceberg table, by using the Ray Datasink API.
+"""
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Union
+
+import ray
+from ray._common.retry import call_with_retry
+from ray.data._internal.datasource.parquet_datasource import (
+    PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT,
+)
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.savemode import SaveMode
+from ray.data._internal.util import MiB
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.datasink import Datasink, WriteResult
+from ray.data.expressions import Expr
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+    from pyiceberg.catalog import Catalog
+    from pyiceberg.expressions import BooleanExpression
+    from pyiceberg.io import FileIO
+    from pyiceberg.manifest import DataFile
+    from pyiceberg.schema import Schema
+    from pyiceberg.table import DataScan, FileScanTask, Table
+    from pyiceberg.table.metadata import TableMetadata
+    from pyiceberg.table.update.schema import UpdateSchema
+
+logger = logging.getLogger(__name__)
+
+_REWRITE_STALL_TIMEOUT_S = 600
+
+
+@ray.remote
+def _rewrite_iceberg_file(
+    file_scan_task: "FileScanTask",
+    keys_ref: "pa.Table",
+    upsert_cols: List[str],
+    table_metadata: "TableMetadata",
+    io: "FileIO",
+) -> "tuple[Optional[DataFile], List[DataFile]]":
+    """Read one Iceberg file, anti-join against upsert keys, write preserved rows.
+
+    Preserved rows are rows in the file that are not in the upsert batch. The
+    coarse range filter would delete them (see ``IcebergDatasink._build_coarse_range_filter``),
+    so we preserve them by writing them as new data files before the delete.
+
+    The file is read in streaming fashion via ``ArrowScan.to_record_batches()``
+    so the full file is never materialised at once. The anti-join is applied
+    per RecordBatch and preserved rows are accumulated, then concatenated and
+    written as a single output once the stream is exhausted.
+
+    Returns (original DataFile to delete, list of new preserved DataFiles).
+    If the entire file is matched (no preserved rows), returns (file, []).
+    If the file has no matched rows at all, returns (None, []), leave it untouched.
+    """
+    import hashlib
+    import time as _time
+    import uuid as _uuid
+
+    import numpy as np
+    import pyarrow as pa
+    from pyiceberg.expressions import AlwaysTrue
+    from pyiceberg.io.pyarrow import ArrowScan, _dataframe_to_data_files
+
+    file_path = file_scan_task.file.file_path
+    file_name = file_path.split("/")[-1]
+    file_size_mb = file_scan_task.file.file_size_in_bytes / MiB
+    t_start = _time.perf_counter()
+
+    # Cast target pulled from keys_ref once.  Applied per batch so PyArrow's join
+    # doesn't raise ArrowInvalid on utf8/large_utf8 or similar width mismatches.
+    target_key_schema = pa.schema([keys_ref.schema.field(c) for c in upsert_cols])
+
+    record_batches = ArrowScan(
+        table_metadata=table_metadata,
+        io=io,
+        projected_schema=table_metadata.schema(),
+        row_filter=AlwaysTrue(),
+    ).to_record_batches(tasks=[file_scan_task])
+
+    preserved_rows: Optional["pa.Table"] = None
+    total_in_rows = 0
+    total_preserved_rows = 0
+    n_batches = 0
+
+    for rb in record_batches:
+        n_batches += 1
+        batch_table = pa.Table.from_batches([rb])
+        if len(batch_table) == 0:
+            continue
+        total_in_rows += len(batch_table)
+
+        batch_keys = batch_table.select(upsert_cols).cast(target_key_schema)
+
+        idx_col = pa.array(np.arange(len(batch_table), dtype=np.int64))
+        preserved_keys = batch_keys.append_column("__row_idx__", idx_col).join(
+            keys_ref, keys=upsert_cols, join_type="left anti"
+        )
+
+        if len(preserved_keys) > 0:
+            new_rows = batch_table.take(preserved_keys["__row_idx__"])
+            if preserved_rows is None:
+                preserved_rows = new_rows
+            else:
+                preserved_rows = pa.concat_tables(
+                    [preserved_rows, new_rows], promote_options="permissive"
+                )
+            total_preserved_rows += len(preserved_keys)
+
+    t_read = _time.perf_counter()
+    logger.debug(
+        "[rewrite] stream-read+join %d rows / %.1f MB (compressed) from %s "
+        "across %d batch(es) in %.2fs",
+        total_in_rows,
+        file_size_mb,
+        file_name,
+        n_batches,
+        t_read - t_start,
+    )
+
+    if total_in_rows == 0:
+        return (None, [])
+
+    if total_preserved_rows == 0:
+        # Every row in this file is being upserted — delete the whole file, no preserved file needed.
+        logger.debug(
+            "[rewrite] %s: all %d rows matched -> whole-file delete",
+            file_name,
+            total_in_rows,
+        )
+        return (file_scan_task.file, [])
+
+    if total_preserved_rows == total_in_rows:
+        # No rows in this file match any upsert key — leave it alone entirely.
+        logger.debug("[rewrite] %s: 0 rows matched -> untouched", file_name)
+        return (None, [])
+
+    # Derive a deterministic write_uuid from the source file path so that
+    # task retries overwrite the same object rather than leaking orphan files.
+    preserved_write_uuid = _uuid.UUID(hashlib.md5(file_path.encode()).hexdigest())
+    preserved_files = list(
+        _dataframe_to_data_files(
+            table_metadata=table_metadata,
+            df=preserved_rows,
+            io=io,
+            write_uuid=preserved_write_uuid,
+        )
+    )
+    logger.debug(
+        "[rewrite] %s: %d/%d rows preserved -> wrote %d preserved file(s) in %.2fs",
+        file_name,
+        total_preserved_rows,
+        total_in_rows,
+        len(preserved_files),
+        _time.perf_counter() - t_read,
+    )
+    return (file_scan_task.file, preserved_files)
+
+
+@dataclass
+class IcebergWriteResult:
+    """Result from writing blocks to Iceberg storage.
+
+    Attributes:
+        data_files: List of DataFile objects containing metadata about written Parquet files.
+        upsert_keys: PyArrow table containing key columns for upsert operations.
+        schemas: List of PyArrow schemas from all non-empty blocks.
+    """
+
+    data_files: List["DataFile"] = field(default_factory=list)
+    upsert_keys: Optional["pa.Table"] = None
+    schemas: List["pa.Schema"] = field(default_factory=list)
+
+
+_UPSERT_COLS_ID = "join_cols"
+
+
+@DeveloperAPI
+class IcebergDatasink(Datasink[IcebergWriteResult]):
+    """
+    Iceberg datasink to write a Ray Dataset into an existing Iceberg table.
+    This datasink handles concurrent writes by:
+    - Each worker writes Parquet files to storage and returns DataFile metadata
+    - The driver collects all DataFile objects and performs a single commit
+
+    Schema evolution is supported:
+    - New columns in incoming data are automatically added to the table schema
+    - Type promotion across blocks is handled via schema reconciliation on the driver
+    """
+
+    def __init__(
+        self,
+        table_identifier: str,
+        catalog_kwargs: Optional[Dict[str, Any]] = None,
+        snapshot_properties: Optional[Dict[str, str]] = None,
+        mode: SaveMode = SaveMode.APPEND,
+        overwrite_filter: Optional["Expr"] = None,
+        upsert_kwargs: Optional[Dict[str, Any]] = None,
+        overwrite_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize the IcebergDatasink
+
+        Args:
+            table_identifier: The identifier of the table such as `default.taxi_dataset`
+            catalog_kwargs: Optional arguments to use when setting up the Iceberg catalog
+            snapshot_properties: Custom properties to write to snapshot summary
+            mode: Write mode - APPEND, UPSERT, or OVERWRITE. Defaults to APPEND.
+                - APPEND: Add new data without checking for duplicates
+                - UPSERT: Update existing rows or insert new ones based on a join condition
+                - OVERWRITE: Replace table data (all data or filtered subset)
+            overwrite_filter: Optional filter for OVERWRITE mode to perform partial overwrites.
+                Must be a Ray Data expression from `ray.data.expressions`. Only rows matching
+                this filter are replaced. If None with OVERWRITE mode, replaces all table data.
+            upsert_kwargs: Optional arguments for upsert operations.
+                Supported parameters: join_cols (List[str]), case_sensitive (bool),
+                branch (str). Note: This implementation uses a copy-on-write strategy
+                that always updates all columns for matched keys and inserts all new keys.
+            overwrite_kwargs: Optional arguments to pass through to PyIceberg's table.overwrite()
+                method. Supported parameters include case_sensitive (bool) and branch (str).
+                See PyIceberg documentation for details.
+
+        Note:
+            Schema evolution is automatically enabled. New columns in the incoming data
+            are automatically added to the table schema. The schema is extracted from
+            the first input bundle when on_write_start is called.
+        """
+        self.table_identifier = table_identifier
+        self._catalog_kwargs = (catalog_kwargs or {}).copy()
+        self._snapshot_properties = (snapshot_properties or {}).copy()
+        self._mode = mode
+        self._overwrite_filter = overwrite_filter
+        self._upsert_kwargs = (upsert_kwargs or {}).copy()
+        self._overwrite_kwargs = (overwrite_kwargs or {}).copy()
+
+        # Validate kwargs are only set for relevant modes
+        if self._upsert_kwargs and self._mode != SaveMode.UPSERT:
+            raise ValueError(
+                f"upsert_kwargs can only be specified when mode is SaveMode.UPSERT, but mode is {self._mode}"
+            )
+        if self._overwrite_kwargs and self._mode != SaveMode.OVERWRITE:
+            raise ValueError(
+                f"overwrite_kwargs can only be specified when mode is SaveMode.OVERWRITE, but mode is {self._mode}"
+            )
+        if self._overwrite_filter and self._mode != SaveMode.OVERWRITE:
+            raise ValueError(
+                f"overwrite_filter can only be specified when mode is SaveMode.OVERWRITE, but mode is {self._mode}"
+            )
+
+        # Remove invalid parameters from overwrite_kwargs if present
+        for invalid_param, reason in [
+            (
+                "overwrite_filter",
+                "should be passed as a separate parameter to write_iceberg()",
+            ),
+            (
+                "delete_filter",
+                "is an internal PyIceberg parameter; use 'overwrite_filter' instead",
+            ),
+        ]:
+            if self._overwrite_kwargs.pop(invalid_param, None) is not None:
+                logger.warning(
+                    f"Removed '{invalid_param}' from overwrite_kwargs: {reason}"
+                )
+
+        if "name" in self._catalog_kwargs:
+            self._catalog_name = self._catalog_kwargs.pop("name")
+        else:
+            self._catalog_name = "default"
+
+        self._table: "Table" = None
+        self._io: "FileIO" = None
+        self._table_metadata: "TableMetadata" = None
+        self._data_context = DataContext.get_current()
+
+    def __getstate__(self) -> dict:
+        """Exclude `_table` during pickling."""
+        state = self.__dict__.copy()
+        state.pop("_table", None)
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        self._table = None
+
+    def _with_retry(self, func: Callable, description: str) -> Any:
+        """Execute a function with retry logic.
+
+        This helper encapsulates the common retry pattern for Iceberg catalog
+        operations, using the configured retry parameters from DataContext.
+
+        Args:
+            func: The callable to execute with retry logic.
+            description: Human-readable description for logging/error messages.
+
+        Returns:
+            The result of calling func.
+        """
+        iceberg_config = self._data_context.iceberg_config
+        return call_with_retry(
+            func,
+            description=description,
+            match=iceberg_config.catalog_retried_errors,
+            max_attempts=iceberg_config.catalog_max_attempts,
+            max_backoff_s=iceberg_config.catalog_retry_max_backoff_s,
+        )
+
+    def _get_catalog(self) -> "Catalog":
+        from pyiceberg import catalog
+
+        return self._with_retry(
+            lambda: catalog.load_catalog(self._catalog_name, **self._catalog_kwargs),
+            description=f"load Iceberg catalog '{self._catalog_name}'",
+        )
+
+    def _reload_table(self) -> None:
+        """Reload the Iceberg table from the catalog."""
+        cat = self._get_catalog()
+        self._table = self._with_retry(
+            lambda: cat.load_table(self.table_identifier),
+            description=f"load Iceberg table '{self.table_identifier}'",
+        )
+        self._io = self._table.io
+        self._table_metadata = self._table.metadata
+
+    def _get_upsert_cols(self) -> List[str]:
+        """Get join columns for upsert, using table identifier fields as fallback."""
+        upsert_cols = self._upsert_kwargs.get(_UPSERT_COLS_ID, [])
+        if not upsert_cols:
+            # Use table's identifier fields as fallback
+            identifier_cols = []
+            schema = self._table_metadata.schema()
+            for field_id in schema.identifier_field_ids:
+                col_name = schema.find_column_name(field_id)
+                if col_name:
+                    identifier_cols.append(col_name)
+            return identifier_cols
+
+        case_sensitive = self._upsert_kwargs.get("case_sensitive", True)
+
+        # To support case insensitivity, we need to define a mapping of
+        # provided (possibly case-modified) names to their original names in the schema
+        if not case_sensitive:
+            schema = self._table_metadata.schema()
+            lower_to_original_mapping = {
+                col.name.lower(): col.name for col in schema.fields
+            }
+            resolved_upsert_cols = []
+            for upsert_col in upsert_cols:
+                resolved_col = lower_to_original_mapping.get(upsert_col.lower())
+                if resolved_col is None:
+                    raise ValueError(
+                        f"Upsert join column {upsert_col!r} does not match any column in "
+                        f"table schema (case-insensitive)."
+                    )
+                resolved_upsert_cols.append(resolved_col)
+            upsert_cols = resolved_upsert_cols
+
+        return upsert_cols
+
+    def _build_coarse_range_filter(
+        self,
+        keys_table: "pa.Table",
+        upsert_cols: List[str],
+    ) -> "BooleanExpression":
+        """Build an O(1) coarse range filter covering all upsert key values.
+
+        For each upsert column computes AND(GTE(col, min), LTE(col, max)).
+        The filter may match rows outside the upsert batch (filter overshoot);
+        callers must anti-join to identify and preserve those rows.
+        """
+        import pyarrow.compute as pc
+        from pyiceberg.expressions import (
+            AlwaysTrue,
+            And,
+            GreaterThanOrEqual,
+            LessThanOrEqual,
+        )
+
+        expr = None
+        for col_name in upsert_cols:
+            mm = pc.min_max(keys_table[col_name])
+            min_val = mm["min"].as_py()
+            max_val = mm["max"].as_py()
+            if min_val is None:
+                continue
+            col_expr = And(
+                GreaterThanOrEqual(col_name, min_val),
+                LessThanOrEqual(col_name, max_val),
+            )
+            expr = col_expr if expr is None else And(expr, col_expr)
+
+        return expr if expr is not None else AlwaysTrue()
+
+    def _commit_upsert_scan_merge(
+        self,
+        txn: "Table.transaction",
+        data_files: List["DataFile"],
+        keys_table: "pa.Table",
+        upsert_cols: List[str],
+    ) -> None:
+        """Upsert commit using coarse range filter + per-file distributed anti-join.
+
+        ┌─────────────────────────────────────────────────────────────┐
+        │  Stage 1: Build coarse filter (driver)                      │
+        │    keys_table ──► min/max per col ──► coarse_filter         │
+        └─────────────────────────────────────────────────────────────┘
+                                    │
+                                    ▼
+        ┌─────────────────────────────────────────────────────────────┐
+        │  Stage 2: Plan candidate files (driver)                     │
+        │    table.scan(coarse_filter).plan_files()                   │
+        │        ──► file_scan_tasks                                  │
+        └─────────────────────────────────────────────────────────────┘
+                                    │
+                                    ▼
+        ┌─────────────────────────────────────────────────────────────┐
+        │  Stage 3: Rewrite (one _rewrite_iceberg_file task per file) │
+        │    read file ─► anti-join keys ─► write preserved rows      │
+        │    returns (old_file, preserved_files)                      │
+        └─────────────────────────────────────────────────────────────┘
+                                    │
+                                    ▼
+        ┌─────────────────────────────────────────────────────────────┐
+        │  Stage 4: Atomic overwrite (driver)                         │
+        │    delete  old_file         (each rewritten candidate)      │
+        │    append  preserved_files  (preserved rows kept)           │
+        │    append  data_files       (new upsert payload)            │
+        └─────────────────────────────────────────────────────────────┘
+                                    │
+                                    ▼
+                            commit_transaction
+
+        1. Build an O(1) coarse range filter using min-max covering upsert key values (for each column).
+        2. plan_files() on the driver to find candidate files that could be updated
+        3. Dispatch one Ray task per candidate file. Each task reads its file,
+           anti-joins against the upsert keys to find preserved rows (rows that
+           the coarse delete would remove but that are NOT being upserted), and
+           writes them as new data files directly to storage.
+        4. Commit atomically via txn.update_snapshot().overwrite(): delete each
+           original candidate file and append preserved files + new upsert data files.
+        """
+        import time
+
+        case_sensitive = self._upsert_kwargs.get("case_sensitive", True)
+        branch = self._upsert_kwargs.get("branch", "main")
+        unknown = set(self._upsert_kwargs) - {
+            _UPSERT_COLS_ID,
+            "case_sensitive",
+            "branch",
+        }
+        if unknown:
+            logger.warning(
+                "[scan-merge] ignoring unsupported upsert_kwargs: %s", sorted(unknown)
+            )
+
+        # Dedup keys to minimise per-task anti-join hash table size.
+        keys_table = keys_table.group_by(upsert_cols).aggregate([])
+
+        coarse_filter = self._build_coarse_range_filter(keys_table, upsert_cols)
+        logger.debug("[scan-merge] coarse_filter=%s", coarse_filter)
+
+        # plan_files() reads only manifest metadata, no Parquet data on the driver.
+        t0 = time.perf_counter()
+        scan: "DataScan" = self._table.scan(
+            row_filter=coarse_filter, case_sensitive=case_sensitive
+        )
+        # Use the specific branch for the scan
+        scan = scan.use_ref(branch)
+        file_scan_tasks: List["FileScanTask"] = list(scan.plan_files())
+
+        logger.info(
+            "[scan-merge] planned %d candidate file(s) in %.2fs",
+            len(file_scan_tasks),
+            time.perf_counter() - t0,
+        )
+
+        if not file_scan_tasks:
+            # No existing files match the coarse filter, so it's a pure insert.
+            self._append_and_commit(txn, data_files, branch=branch)
+            return
+
+        # Put the deduped keys in the object store once; all tasks share one copy.
+        keys_ref = ray.put(keys_table)
+
+        t0 = time.perf_counter()
+        refs = [
+            _rewrite_iceberg_file.options(
+                memory=int(
+                    task.file.file_size_in_bytes
+                    * PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT
+                    * 3  # Bump memory estimate to account for the anti-join and the preserved rows (also since to_record_batches materializes the entire table in memory, see https://github.com/apache/iceberg-python/issues/3036)
+                ),
+                num_cpus=1,
+            ).remote(task, keys_ref, upsert_cols, self._table_metadata, self._io)
+            for task in file_scan_tasks
+        ]
+        logger.info("[scan-merge] dispatched %d rewrite task(s)", len(refs))
+
+        # Collect results with periodic progress logs so long rewrites aren't silent.
+        results = []
+        pending = list(refs)
+        _LOG_INTERVAL = max(1, len(refs) // 10)  # log ~10 times total
+        while pending:
+            done, pending = ray.wait(
+                pending,
+                num_returns=min(_LOG_INTERVAL, len(pending)),
+                timeout=_REWRITE_STALL_TIMEOUT_S,
+                fetch_local=True,
+            )
+            results.extend(ray.get(done))
+            logger.debug(
+                "[scan-merge] rewrite progress: %d/%d file(s) done (%.1fs elapsed)",
+                len(results),
+                len(refs),
+                time.perf_counter() - t0,
+            )
+
+        logger.info(
+            "[scan-merge] all %d file(s) rewritten in %.2fs",
+            len(refs),
+            time.perf_counter() - t0,
+        )
+
+        # Count how many files were wholly deleted vs partially rewritten.
+        n_whole_delete = n_partial = n_untouched = 0
+        for old, preserved_files in results:
+            if old is None:
+                n_untouched += 1
+            elif preserved_files:
+                n_partial += 1
+            else:
+                n_whole_delete += 1
+        logger.info(
+            "[scan-merge] files: %d whole-delete, %d partial-rewrite, %d untouched",
+            n_whole_delete,
+            n_partial,
+            n_untouched,
+        )
+
+        # Single atomic commit: schema update (already staged in txn), and overwrite.
+        # _OverwriteFiles handles both file-level deletes and appends in one snapshot.
+        t0 = time.perf_counter()
+        with txn.update_snapshot(
+            snapshot_properties=self._snapshot_properties, branch=branch
+        ).overwrite() as snap:
+            for old_file, preserved_files in results:
+                if old_file is not None:
+                    snap.delete_data_file(old_file)
+                for preserved_file in preserved_files:
+                    snap.append_data_file(preserved_file)
+            for df in data_files:
+                snap.append_data_file(df)
+
+        self._with_retry(
+            txn.commit_transaction,
+            description=f"commit upsert transaction to Iceberg table '{self.table_identifier}'",
+        )
+        logger.info("[scan-merge] committed in %.2fs", time.perf_counter() - t0)
+
+    def _append_and_commit(
+        self,
+        txn: "Table.transaction",
+        data_files: List["DataFile"],
+        branch: str = "main",
+    ) -> None:
+        """Append data files to a transaction and commit.
+
+        Args:
+            txn: PyIceberg transaction object
+            data_files: List of DataFile objects to append
+            branch: Iceberg branch to commit the snapshot to. Defaults to "main"
+                to match pyiceberg's default
+        """
+        with txn._append_snapshot_producer(
+            self._snapshot_properties, branch=branch
+        ) as append_files:
+            for data_file in data_files:
+                append_files.append_data_file(data_file)
+
+        self._with_retry(
+            txn.commit_transaction,
+            description=f"commit transaction to Iceberg table '{self.table_identifier}'",
+        )
+
+    def _commit_upsert(
+        self,
+        txn: "Table.transaction",
+        data_files: List["DataFile"],
+        upsert_keys: Optional["pa.Table"],
+    ) -> None:
+        """
+        Commit upsert transaction with copy-on-write strategy.
+
+        Args:
+            txn: PyIceberg transaction object
+            data_files: List of DataFile objects to commit
+            upsert_keys: PyArrow table containing upsert key columns
+        """
+        import functools
+        import time
+
+        import pyarrow as pa
+
+        # Create delete filter if we have join keys
+        if upsert_keys is not None and len(upsert_keys) > 0:
+            # Filter out rows with any NULL values in join columns
+            # (NULL != NULL in SQL semantics)
+            upsert_cols = self._get_upsert_cols()
+            logger.info(
+                "[upsert commit] Filtering NULL keys from %d rows on cols %s",
+                len(upsert_keys),
+                upsert_cols,
+            )
+            t0 = time.perf_counter()
+            masks = (pa.compute.is_valid(upsert_keys[col]) for col in upsert_cols)
+            mask = functools.reduce(pa.compute.and_, masks)
+            keys_table = upsert_keys.filter(mask)
+            logger.info(
+                "[upsert commit] NULL filter done in %.2fs: %d -> %d rows (dropped %d NULLs)",
+                time.perf_counter() - t0,
+                len(upsert_keys),
+                len(keys_table),
+                len(upsert_keys) - len(keys_table),
+            )
+
+            # Only delete if we have non-NULL keys
+            if len(keys_table) > 0:
+                self._commit_upsert_scan_merge(txn, data_files, keys_table, upsert_cols)
+                return
+        else:
+            logger.info("[upsert commit] No upsert keys — skipping delete phase")
+
+        # No non-NULL keys — just append new data files and commit
+        logger.info(
+            "[upsert commit] Appending %d data files and committing ...",
+            len(data_files),
+        )
+        t0 = time.perf_counter()
+        branch = self._upsert_kwargs.get("branch", "main")
+        self._append_and_commit(txn, data_files, branch=branch)
+        logger.info(
+            "[upsert commit] Append+commit done in %.2fs",
+            time.perf_counter() - t0,
+        )
+
+    def _preserve_identifier_field_requirements(
+        self, update: "UpdateSchema", table_schema: "Schema"
+    ) -> None:
+        """Ensure identifier fields remain required after schema union.
+
+        When union_by_name is called with a schema that has nullable fields,
+        PyIceberg may make identifier fields optional. Since identifier fields
+        must be required, this helper ensures they remain required after union.
+
+        Example:
+            Table schema:   id: int (required, identifier), val: string
+            Input schema:   id: int (optional),             val: string
+
+            `union_by_name` merges them to:
+                            id: int (optional),             val: string
+
+            This violates the identifier constraint. This function forces `id`
+            back to required in the pending update.
+
+        Args:
+            update: The UpdateSchema object from update_schema() context manager
+            table_schema: The current table schema to get identifier field IDs from
+        """
+        from pyiceberg.types import NestedField
+
+        identifier_field_ids = table_schema.identifier_field_ids
+        for field_id in identifier_field_ids:
+            # Check if this field has a pending update
+            if field_id in update._updates:
+                updated_field = update._updates[field_id]
+                # If it was made optional (likely by union_by_name), force it back to required
+                if not updated_field.required:
+                    # Directly update the pending change to enforce required=True.
+                    # We create a new NestedField because it might be immutable.
+                    # We bypass _set_column_requirement because it has a check that
+                    # incorrectly returns early if the original field is already required,
+                    # ignoring the fact that we are overwriting a pending update.
+                    update._updates[field_id] = NestedField(
+                        field_id=updated_field.field_id,
+                        name=updated_field.name,
+                        field_type=updated_field.field_type,
+                        doc=updated_field.doc,
+                        required=True,
+                        initial_default=updated_field.initial_default,
+                        write_default=updated_field.write_default,
+                    )
+
+    def _update_schema_with_union(
+        self,
+        update: "UpdateSchema",
+        new_schema: Union["pa.Schema", "Schema"],
+        table_schema: "Schema",
+    ) -> None:
+        """Update schema using union_by_name while preserving identifier field requirements.
+
+        Args:
+            update: The UpdateSchema object.
+            new_schema: The new schema to union with the table schema.
+            table_schema: The current table schema.
+        """
+        update.union_by_name(new_schema)
+        self._preserve_identifier_field_requirements(update, table_schema)
+
+    def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
+        """Initialize table for writing and create a shared write UUID.
+
+        Args:
+            schema: The PyArrow schema of the data being written. This is
+                automatically extracted from the first input bundle by the
+                Write operator. Used to evolve the table schema before writing
+                to avoid PyIceberg name mapping errors.
+        """
+        self._reload_table()
+
+        # Evolve schema BEFORE any files are written
+        # This prevents PyIceberg name mapping errors when incoming data has new columns
+        if schema is not None:
+            table_schema = self._table.metadata.schema()
+
+            def _update_schema():
+                with self._table.update_schema() as update:
+                    self._update_schema_with_union(update, schema, table_schema)
+
+            self._with_retry(
+                _update_schema,
+                description=f"update schema for Iceberg table '{self.table_identifier}'",
+            )
+            # Succeeded, reload to get latest table version and exit.
+            self._reload_table()
+
+        # Validate join_cols for UPSERT mode before writing any files
+        if self._mode == SaveMode.UPSERT:
+            upsert_cols = self._upsert_kwargs.get(_UPSERT_COLS_ID, [])
+            if not upsert_cols:
+                # Check if table has identifier fields as fallback
+                identifier_field_ids = (
+                    self._table.metadata.schema().identifier_field_ids
+                )
+                if not identifier_field_ids:
+                    raise ValueError(
+                        "join_cols must be specified in upsert_kwargs for UPSERT mode "
+                        "when table has no identifier fields"
+                    )
+
+    def write(self, blocks: Iterable[Block], ctx: TaskContext) -> IcebergWriteResult:
+        """
+        Write blocks to Parquet files in storage and return DataFile metadata with schemas.
+
+        This runs on each worker in parallel. Files are written directly to storage
+        (S3, HDFS, etc.) and only metadata is returned to the driver.
+        Schema updates are NOT performed here - they happen on the driver.
+
+        Args:
+            blocks: Iterable of Ray Data blocks to write
+            ctx: TaskContext object containing task-specific information
+
+        Returns:
+            IcebergWriteResult containing DataFile objects, upsert keys, and schemas.
+        """
+        from pyiceberg.io.pyarrow import _dataframe_to_data_files
+
+        all_data_files = []
+        upsert_keys_tables = []
+        block_schemas = []
+        use_copy_on_write_upsert = self._mode == SaveMode.UPSERT
+
+        for block in blocks:
+            pa_table = BlockAccessor.for_block(block).to_arrow()
+            if pa_table.num_rows > 0:
+                block_schemas.append(pa_table.schema)
+
+                # Extract join key values for copy-on-write upsert
+                if use_copy_on_write_upsert:
+                    upsert_cols = self._get_upsert_cols()
+                    if len(upsert_cols) > 0:
+                        upsert_keys_tables.append(pa_table.select(upsert_cols))
+
+                # Write data files to storage with retry for transient errors
+                def _write_data_files():
+                    return list(
+                        _dataframe_to_data_files(
+                            table_metadata=self._table_metadata,
+                            df=pa_table,
+                            io=self._io,
+                        )
+                    )
+
+                iceberg_config = self._data_context.iceberg_config
+                data_files = call_with_retry(
+                    _write_data_files,
+                    description=f"write data files to Iceberg table '{self.table_identifier}'",
+                    match=self._data_context.retried_io_errors,
+                    max_attempts=iceberg_config.write_file_max_attempts,
+                    max_backoff_s=iceberg_config.write_file_retry_max_backoff_s,
+                )
+                all_data_files.extend(data_files)
+
+        # Combine all upsert key tables into one
+        from ray.data._internal.arrow_ops.transform_pyarrow import concat
+
+        upsert_keys = concat(upsert_keys_tables) if upsert_keys_tables else None
+
+        return IcebergWriteResult(
+            data_files=all_data_files,
+            upsert_keys=upsert_keys,
+            schemas=block_schemas,
+        )
+
+    def _commit_overwrite(
+        self, txn: "Table.transaction", data_files: List["DataFile"]
+    ) -> None:
+        """Commit data files using OVERWRITE mode."""
+        from pyiceberg.expressions import AlwaysTrue
+
+        # Default - Full overwrite - delete all
+        pyi_filter = AlwaysTrue()
+
+        # Delete matching data if filter provided
+        if self._overwrite_filter is not None:
+            from ray.data._internal.datasource.iceberg_datasource import (
+                _IcebergExpressionVisitor,
+            )
+
+            visitor = _IcebergExpressionVisitor()
+            pyi_filter = visitor.visit(self._overwrite_filter)
+
+        txn.delete(
+            delete_filter=pyi_filter,
+            snapshot_properties=self._snapshot_properties,
+            **self._overwrite_kwargs,
+        )
+
+        # Append on the same branch the delete targeted (defaults to "main").
+        branch = self._overwrite_kwargs.get("branch", "main")
+        self._append_and_commit(txn, data_files, branch=branch)
+
+    def on_write_complete(self, write_result: WriteResult) -> None:
+        """
+        Complete the write by reconciling schemas and committing all data files.
+
+        This runs on the driver after all workers finish writing files.
+        Collects all DataFile objects and schemas from all workers, reconciles schemas
+        (allowing type promotion), updates table schema if needed, then performs a single
+        atomic commit.
+        """
+        import time
+
+        t_start = time.perf_counter()
+        logger.info("[on_write_complete] Starting commit phase (mode=%s)", self._mode)
+
+        # Collect all data files and schemas from all workers
+        all_data_files: List["DataFile"] = []
+        all_schemas: List["pa.Schema"] = []
+        upsert_keys_tables: List["pa.Table"] = []
+
+        for write_return in write_result.write_returns:
+            if not write_return:
+                continue
+
+            if write_return.data_files:  # Only add schema if we have data files
+                all_data_files.extend(write_return.data_files)
+                all_schemas.extend(write_return.schemas)
+                if write_return.upsert_keys is not None:
+                    upsert_keys_tables.append(write_return.upsert_keys)
+
+        logger.info(
+            "[on_write_complete] Collected results: %d data files, %d schema blocks, "
+            "%d upsert key batches from workers (%.2fs)",
+            len(all_data_files),
+            len(all_schemas),
+            len(upsert_keys_tables),
+            time.perf_counter() - t_start,
+        )
+
+        if not all_data_files:
+            logger.info("[on_write_complete] No data files written, nothing to commit")
+            return
+
+        # Concatenate all upsert keys from all workers into a single table
+        from ray.data._internal.arrow_ops.transform_pyarrow import concat
+
+        if upsert_keys_tables:
+            total_key_rows = sum(len(t) for t in upsert_keys_tables)
+            logger.info(
+                "[on_write_complete] Concatenating %d upsert key batches (%d total rows) ...",
+                len(upsert_keys_tables),
+                total_key_rows,
+            )
+            t0 = time.perf_counter()
+            upsert_keys = concat(upsert_keys_tables)
+            logger.info(
+                "[on_write_complete] upsert key concat done in %.2fs: %d rows, cols=%s",
+                time.perf_counter() - t0,
+                len(upsert_keys),
+                upsert_keys.column_names,
+            )
+        else:
+            upsert_keys = None
+
+        # Reconcile all schemas from all blocks across all workers
+        # Get table schema and union with reconciled schema using unify_schemas with promotion
+        from pyiceberg.io import pyarrow as pyi_pa_io
+
+        from ray.data._internal.arrow_ops.transform_pyarrow import unify_schemas
+
+        logger.info("[on_write_complete] Reconciling %d schemas ...", len(all_schemas))
+        t0 = time.perf_counter()
+        table_schema = pyi_pa_io.schema_to_pyarrow(self._table.schema())
+        final_reconciled_schema = unify_schemas(
+            [table_schema] + all_schemas, promote_types=True
+        )
+        logger.info(
+            "[on_write_complete] Schema reconciliation done in %.2fs",
+            time.perf_counter() - t0,
+        )
+
+        # Create transaction and commit schema update + data files atomically
+        txn = self._table.transaction()
+
+        # Update table schema within the transaction if it differs
+        if not final_reconciled_schema.equals(table_schema):
+            logger.info(
+                "[on_write_complete] Schema changed — updating table schema ..."
+            )
+            t0 = time.perf_counter()
+            current_table_schema = self._table.metadata.schema()
+            with txn.update_schema() as update:
+                self._update_schema_with_union(
+                    update, final_reconciled_schema, current_table_schema
+                )
+            logger.info(
+                "[on_write_complete] Schema update done in %.2fs",
+                time.perf_counter() - t0,
+            )
+        else:
+            logger.info("[on_write_complete] Schema unchanged, skipping update")
+
+        # Create transaction and commit based on mode
+        logger.info(
+            "[on_write_complete] Starting %s commit for %d data files ...",
+            self._mode,
+            len(all_data_files),
+        )
+        t0 = time.perf_counter()
+        if self._mode == SaveMode.APPEND:
+            self._append_and_commit(txn, all_data_files)
+        elif self._mode == SaveMode.OVERWRITE:
+            self._commit_overwrite(txn, all_data_files)
+        elif self._mode == SaveMode.UPSERT:
+            self._commit_upsert(txn, all_data_files, upsert_keys)
+        else:
+            raise ValueError(f"Unsupported mode: {self._mode}")
+        logger.info(
+            "[on_write_complete] Commit complete in %.2fs (total on_write_complete=%.2fs)",
+            time.perf_counter() - t0,
+            time.perf_counter() - t_start,
+        )

python/ray/data/_internal/datasource/iceberg_datasource.py

@@ -0,0 +1,514 @@
+"""
+Module to read an iceberg table into a Ray Dataset, by using the Ray Datasource API.
+"""
+
+import heapq
+import itertools
+import logging
+from functools import partial
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+
+import pyarrow as pa
+from packaging import version
+
+from ray.data._internal.planner.plan_expression.expression_visitors import _ExprVisitor
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    DownloadExpr,
+    LiteralExpr,
+    MonotonicallyIncreasingIdExpr,
+    Operation,
+    RandomExpr,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+    UUIDExpr,
+)
+from ray.util import log_once
+from ray.util.annotations import DeveloperAPI
+
+try:
+    from pyiceberg.expressions import (
+        And,
+        EqualTo,
+        GreaterThan,
+        GreaterThanOrEqual,
+        In,
+        IsNull,
+        LessThan,
+        LessThanOrEqual,
+        Literal,
+        Not,
+        NotEqualTo,
+        NotIn,
+        NotNull,
+        Or,
+        Reference,
+        UnboundTerm,
+        literal,
+    )
+
+    RAY_DATA_OPERATION_TO_ICEBERG = {
+        Operation.EQ: EqualTo,
+        Operation.NE: NotEqualTo,
+        Operation.GT: GreaterThan,
+        Operation.GE: GreaterThanOrEqual,
+        Operation.LT: LessThan,
+        Operation.LE: LessThanOrEqual,
+        Operation.AND: And,
+        Operation.OR: Or,
+        Operation.IN: In,
+        Operation.NOT_IN: NotIn,
+        Operation.IS_NULL: IsNull,
+        Operation.IS_NOT_NULL: NotNull,
+        Operation.NOT: Not,
+    }
+except ImportError:
+    log_once("pyiceberg.expressions not found. Please install pyiceberg >= 0.9.0")
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
+    from pyiceberg.expressions import BooleanExpression
+    from pyiceberg.io import FileIO
+    from pyiceberg.manifest import DataFile
+    from pyiceberg.schema import Schema
+    from pyiceberg.table import DataScan, FileScanTask, Table
+    from pyiceberg.table.metadata import TableMetadata
+
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+class _IcebergExpressionVisitor(
+    _ExprVisitor["BooleanExpression | UnboundTerm | Literal"]
+):
+    """
+    Visitor that converts Ray Data expressions to PyIceberg expressions.
+
+    This enables Ray Data users to write filters using the familiar col() syntax
+    while leveraging Iceberg's native filtering capabilities.
+
+    Example:
+        >>> from ray.data.expressions import col
+        >>> ray_expr = (col("date") >= "2024-01-01") & (col("status") == "active")
+        >>> iceberg_expr = _IcebergExpressionVisitor().visit(ray_expr)
+        >>> # iceberg_expr can now be used with PyIceberg's filter APIs
+    """
+
+    def visit_column(self, expr: "ColumnExpr") -> "UnboundTerm":
+        """Convert a column reference to an Iceberg reference."""
+        return Reference(expr.name)
+
+    def visit_literal(self, expr: "LiteralExpr") -> "Literal":
+        """Convert a literal value to an Iceberg literal."""
+        return literal(expr.value)
+
+    def visit_binary(self, expr: "BinaryExpr") -> "BooleanExpression":
+        """Convert a binary operation to an Iceberg expression."""
+        # Handle IN/NOT_IN specially since they don't visit the right operand
+        # (the right operand is a list literal that can't be converted)
+        if expr.op in (Operation.IN, Operation.NOT_IN):
+            left = self.visit(expr.left)
+            if not isinstance(expr.right, LiteralExpr):
+                raise ValueError(
+                    f"{expr.op.name} operation requires right operand to be a literal list, "
+                    f"got {type(expr.right).__name__}"
+                )
+            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](left, expr.right.value)
+
+        # For all other operations, visit both operands
+        left = self.visit(expr.left)
+        right = self.visit(expr.right)
+
+        if expr.op in RAY_DATA_OPERATION_TO_ICEBERG:
+            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](left, right)
+        else:
+            # Arithmetic operations are not supported in filter expressions
+            raise ValueError(
+                f"Unsupported binary operation for Iceberg filters: {expr.op}. "
+                f"Iceberg filters support: {RAY_DATA_OPERATION_TO_ICEBERG.keys()}. "
+                f"Arithmetic operations (ADD, SUB, MUL, DIV) cannot be used in filters."
+            )
+
+    def visit_unary(self, expr: "UnaryExpr") -> "BooleanExpression":
+        """Convert a unary operation to an Iceberg expression."""
+        operand = self.visit(expr.operand)
+
+        if expr.op in RAY_DATA_OPERATION_TO_ICEBERG:
+            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](operand)
+        else:
+            raise ValueError(
+                f"Unsupported unary operation for Iceberg: {expr.op}. "
+                f"Supported operations: {RAY_DATA_OPERATION_TO_ICEBERG.keys()}"
+            )
+
+    def visit_alias(
+        self, expr: "AliasExpr"
+    ) -> "BooleanExpression | UnboundTerm | Literal":
+        """Convert an aliased expression (just unwrap the alias)."""
+        return self.visit(expr.expr)
+
+    def visit_udf(self, expr: "UDFExpr") -> "BooleanExpression | UnboundTerm | Literal":
+        """UDF expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "UDF expressions cannot be converted to Iceberg expressions. "
+            "Iceberg filters must use simple column comparisons and boolean operations."
+        )
+
+    def visit_download(
+        self, expr: "DownloadExpr"
+    ) -> "BooleanExpression | UnboundTerm | Literal":
+        """Download expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "Download expressions cannot be converted to Iceberg expressions."
+        )
+
+    def visit_star(
+        self, expr: "StarExpr"
+    ) -> "BooleanExpression | UnboundTerm | Literal":
+        """Star expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "Star expressions cannot be converted to Iceberg filter expressions."
+        )
+
+    def visit_monotonically_increasing_id(
+        self, expr: "MonotonicallyIncreasingIdExpr"
+    ) -> "BooleanExpression | UnboundTerm | Literal":
+        """Monotonically increasing ID expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "monotonically_increasing_id expressions cannot be converted to Iceberg filter expressions."
+        )
+
+    def visit_random(
+        self, expr: "RandomExpr"
+    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
+        """Random expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "Random expressions cannot be converted to Iceberg filter expressions."
+        )
+
+    def visit_uuid(
+        self, expr: "UUIDExpr"
+    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
+        """UUID expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "UUID expressions cannot be converted to Iceberg filter expressions."
+        )
+
+
+def _get_read_task(
+    tasks: Iterable["FileScanTask"],
+    table_io: "FileIO",
+    table_metadata: "TableMetadata",
+    row_filter: "BooleanExpression",
+    case_sensitive: bool,
+    limit: Optional[int],
+    schema: "Schema",
+) -> Iterable[Block]:
+    # Determine the PyIceberg version to handle backward compatibility
+    import pyiceberg
+
+    def _generate_tables() -> Iterable[pa.Table]:
+        if version.parse(pyiceberg.__version__) >= version.parse("0.9.0"):
+            # Modern implementation using ArrowScan (PyIceberg 0.9.0+)
+            from pyiceberg.io.pyarrow import ArrowScan
+
+            # Initialize scanner with Iceberg metadata and query parameters
+            scanner = ArrowScan(
+                table_metadata=table_metadata,
+                io=table_io,
+                row_filter=row_filter,
+                projected_schema=schema,
+                case_sensitive=case_sensitive,
+                limit=limit,
+            )
+
+            # Convert scanned data to Arrow Table format
+            result_table = scanner.to_table(tasks=tasks)
+
+            # Stream results as RecordBatches for memory efficiency
+            for batch in result_table.to_batches():
+                yield pa.Table.from_batches([batch])
+
+        else:
+            # Legacy implementation using project_table (PyIceberg <0.9.0)
+            from pyiceberg.io import pyarrow as pyi_pa_io
+
+            # Use the PyIceberg API to read only a single task (specifically, a
+            # FileScanTask) - note that this is not as simple as reading a single
+            # parquet file, as there might be delete files, etc. associated, so we
+            # must use the PyIceberg API for the projection.
+            table = pyi_pa_io.project_table(
+                tasks=tasks,
+                table_metadata=table_metadata,
+                io=table_io,
+                row_filter=row_filter,
+                projected_schema=schema,
+                case_sensitive=case_sensitive,
+                limit=limit,
+            )
+            yield table
+
+    yield from _generate_tables()
+
+
+@DeveloperAPI
+class IcebergDatasource(Datasource):
+    """
+    Iceberg datasource to read Iceberg tables into a Ray Dataset. This module heavily
+    uses PyIceberg to read iceberg tables. All the routines in this class override
+    `ray.data.Datasource`.
+    """
+
+    def __init__(
+        self,
+        table_identifier: str,
+        row_filter: Union[str, "BooleanExpression"] = None,
+        selected_fields: Tuple[str, ...] = ("*",),
+        snapshot_id: Optional[int] = None,
+        scan_kwargs: Optional[Dict[str, Any]] = None,
+        catalog_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize an IcebergDatasource.
+
+        Args:
+            table_identifier: Fully qualified table identifier (i.e.,
+                "db_name.table_name")
+            row_filter: A PyIceberg BooleanExpression to use to filter the data *prior*
+                 to reading
+            selected_fields: Which columns from the data to read, passed directly to
+                PyIceberg's load functions
+            snapshot_id: Optional snapshot ID for the Iceberg table
+            scan_kwargs: Optional arguments to pass to PyIceberg's Table.scan()
+                function
+            catalog_kwargs: Optional arguments to use when setting up the Iceberg
+                catalog
+        """
+        # Initialize parent class to set up predicate pushdown mixin
+        super().__init__()
+
+        _check_import(self, module="pyiceberg", package="pyiceberg")
+        from pyiceberg.expressions import AlwaysTrue
+
+        self._scan_kwargs = scan_kwargs if scan_kwargs is not None else {}
+        self._catalog_kwargs = catalog_kwargs if catalog_kwargs is not None else {}
+
+        if "name" in self._catalog_kwargs:
+            self._catalog_name = self._catalog_kwargs.pop("name")
+        else:
+            self._catalog_name = "default"
+
+        self.table_identifier = table_identifier
+
+        self._row_filter = row_filter if row_filter is not None else AlwaysTrue()
+        # Convert selected_fields to projection_map (identity mapping if specified)
+        # Note: Empty tuple () means no columns, None/"*" means all columns
+        if selected_fields is None or selected_fields == ("*",):
+            self._projection_map = None
+        else:
+            self._projection_map = {col: col for col in selected_fields}
+
+        if snapshot_id:
+            self._scan_kwargs["snapshot_id"] = snapshot_id
+
+        self._plan_files = None
+        self._table = None
+
+    def _get_catalog(self) -> "Catalog":
+        from pyiceberg import catalog
+
+        return catalog.load_catalog(self._catalog_name, **self._catalog_kwargs)
+
+    @property
+    def table(self) -> "Table":
+        """
+        Return the table reference from the catalog
+        """
+        if self._table is None:
+            catalog = self._get_catalog()
+            self._table = catalog.load_table(self.table_identifier)
+        return self._table
+
+    @property
+    def plan_files(self) -> List["FileScanTask"]:
+        """
+        Return the plan files specified by this query
+        """
+        # Calculate and cache the plan_files if they don't already exist
+        if self._plan_files is None:
+            data_scan = self._get_data_scan()
+            self._plan_files = data_scan.plan_files()
+
+        return self._plan_files
+
+    def _get_combined_filter(self) -> "BooleanExpression":
+        """Get the combined filter including both row_filter and pushed-down predicates."""
+        combined_filter = self._row_filter
+
+        if self._predicate_expr is not None:
+            # Convert Ray Data expression to PyIceberg expression using internal visitor
+            visitor = _IcebergExpressionVisitor()
+            iceberg_filter = visitor.visit(self._predicate_expr)
+
+            # Combine with existing row_filter using AND
+            from pyiceberg.expressions import AlwaysTrue, And
+
+            if not isinstance(combined_filter, AlwaysTrue):
+                combined_filter = And(combined_filter, iceberg_filter)
+            else:
+                combined_filter = iceberg_filter
+
+        return combined_filter
+
+    def _get_data_scan(self) -> "DataScan":
+        # Get the combined filter
+        combined_filter = self._get_combined_filter()
+
+        # Convert back to tuple for PyIceberg API (None -> ("*",))
+        data_columns = self._get_data_columns()
+        selected_fields = ("*",) if data_columns is None else tuple(data_columns)
+
+        data_scan = self.table.scan(
+            row_filter=combined_filter,
+            selected_fields=selected_fields,
+            **self._scan_kwargs,
+        )
+
+        return data_scan
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # Approximate the size by using the plan files - this will not
+        # incorporate the deletes, but that's a reasonable approximation
+        # task
+        return sum(task.file.file_size_in_bytes for task in self.plan_files)
+
+    def supports_predicate_pushdown(self) -> bool:
+        """Returns True to indicate this datasource supports predicate pushdown."""
+        return True
+
+    def supports_projection_pushdown(self) -> bool:
+        """Returns True to indicate this datasource supports projection pushdown."""
+        return True
+
+    @staticmethod
+    def _distribute_tasks_into_equal_chunks(
+        plan_files: Iterable["FileScanTask"], n_chunks: int
+    ) -> List[List["FileScanTask"]]:
+        """
+        Implement a greedy knapsack algorithm to distribute the files in the scan
+        across tasks, based on their file size, as evenly as possible
+        """
+        chunks = [list() for _ in range(n_chunks)]
+
+        chunk_sizes = [(0, chunk_id) for chunk_id in range(n_chunks)]
+        heapq.heapify(chunk_sizes)
+
+        # From largest to smallest, add the plan files to the smallest chunk one at a
+        # time
+        for plan_file in sorted(
+            plan_files, key=lambda f: f.file.file_size_in_bytes, reverse=True
+        ):
+            smallest_chunk = heapq.heappop(chunk_sizes)
+            chunks[smallest_chunk[1]].append(plan_file)
+            heapq.heappush(
+                chunk_sizes,
+                (
+                    smallest_chunk[0] + plan_file.file.file_size_in_bytes,
+                    smallest_chunk[1],
+                ),
+            )
+
+        return chunks
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        from pyiceberg.io import pyarrow as pyi_pa_io
+        from pyiceberg.manifest import DataFileContent
+
+        # Get the PyIceberg scan
+        data_scan = self._get_data_scan()
+        # Get the plan files in this query
+        plan_files = self.plan_files
+
+        # Get the projected schema for this scan, given all the row filters,
+        # snapshot ID, etc.
+        projected_schema = data_scan.projection()
+        # Get the arrow schema, to set in the metadata
+        pya_schema = pyi_pa_io.schema_to_pyarrow(projected_schema)
+
+        # Set the n_chunks to the min of the number of plan files and the actual
+        # requested n_chunks, so that there are no empty tasks
+        if parallelism > len(list(plan_files)):
+            parallelism = len(list(plan_files))
+            logger.warning(
+                f"Reducing the parallelism to {parallelism}, as that is the number of files"
+            )
+
+        # Get required properties for reading tasks - table IO, table metadata,
+        # row filter, case sensitivity,limit and projected schema to pass
+        # them directly to `_get_read_task` to avoid capture of `self` reference
+        # within the closure carrying substantial overhead invoking these tasks
+        #
+        # See https://github.com/ray-project/ray/issues/49107 for more context
+        table_io = self.table.io
+        table_metadata = self.table.metadata
+        row_filter = self._get_combined_filter()
+        case_sensitive = self._scan_kwargs.get("case_sensitive", True)
+        limit = self._scan_kwargs.get("limit")
+
+        get_read_task = partial(
+            _get_read_task,
+            table_io=table_io,
+            table_metadata=table_metadata,
+            row_filter=row_filter,
+            case_sensitive=case_sensitive,
+            limit=limit,
+            schema=projected_schema,
+        )
+
+        read_tasks = []
+        # Chunk the plan files based on the requested parallelism
+        for chunk_tasks in IcebergDatasource._distribute_tasks_into_equal_chunks(
+            plan_files, parallelism
+        ):
+            unique_deletes: Set[DataFile] = set(
+                itertools.chain.from_iterable(
+                    [task.delete_files for task in chunk_tasks]
+                )
+            )
+            # Get a rough estimate of the number of deletes by just looking at
+            # position deletes. Equality deletes are harder to estimate, as they
+            # can delete multiple rows.
+            position_delete_count = sum(
+                delete.record_count
+                for delete in unique_deletes
+                if delete.content == DataFileContent.POSITION_DELETES
+            )
+            metadata = BlockMetadata(
+                num_rows=sum(task.file.record_count for task in chunk_tasks)
+                - position_delete_count,
+                size_bytes=sum(task.file.file_size_in_bytes for task in chunk_tasks),
+                input_files=[task.file.file_path for task in chunk_tasks],
+                exec_stats=None,
+            )
+            read_tasks.append(
+                ReadTask(
+                    read_fn=lambda tasks=chunk_tasks: get_read_task(tasks),
+                    metadata=metadata,
+                    schema=pya_schema,
+                    per_task_row_limit=per_task_row_limit,
+                )
+            )
+
+        return read_tasks

python/ray/data/_internal/datasource/image_datasink.py

@@ -0,0 +1,24 @@
+import io
+from typing import Any, Dict
+
+import pyarrow
+
+from ray.data.datasource.file_datasink import RowBasedFileDatasink
+
+
+class ImageDatasink(RowBasedFileDatasink):
+    def __init__(
+        self, path: str, column: str, file_format: str, **file_datasink_kwargs
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        self.column = column
+        self.file_format = file_format
+
+    def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
+        from PIL import Image
+
+        image = Image.fromarray(row[self.column])
+        buffer = io.BytesIO()
+        image.save(buffer, format=self.file_format)
+        file.write(buffer.getvalue())

python/ray/data/_internal/datasource/image_datasource.py

@@ -0,0 +1,177 @@
+import io
+import logging
+import time
+from dataclasses import replace
+from typing import TYPE_CHECKING, Iterator, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+from ray.data.datasource.file_meta_provider import DefaultFileMetadataProvider
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+logger = logging.getLogger(__name__)
+
+# The default size multiplier for reading image data source.
+# This essentially is using image on-disk file size to estimate
+# in-memory data size.
+IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT = 1
+
+# The lower bound value to estimate image encoding ratio.
+IMAGE_ENCODING_RATIO_ESTIMATE_LOWER_BOUND = 0.5
+
+
+class ImageDatasource(FileBasedDatasource):
+    """A datasource that lets you read images."""
+
+    _WRITE_FILE_PER_ROW = True
+    _FILE_EXTENSIONS = ["png", "jpg", "jpeg", "tif", "tiff", "bmp", "gif"]
+    # Use 8 threads per task to read image files.
+    _NUM_THREADS_PER_TASK = 8
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        size: Optional[Tuple[int, int]] = None,
+        mode: Optional[str] = None,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="PIL", package="Pillow")
+
+        if size is not None and len(size) != 2:
+            raise ValueError(
+                "Expected `size` to contain two integers for height and width, "
+                f"but got {len(size)} integers instead."
+            )
+
+        if size is not None and (size[0] < 0 or size[1] < 0):
+            raise ValueError(
+                f"Expected `size` to contain positive integers, but got {size} instead."
+            )
+
+        self.size = size
+        self.mode = mode
+
+        meta_provider = file_based_datasource_kwargs.get("meta_provider", None)
+        if isinstance(meta_provider, ImageFileMetadataProvider):
+            self._encoding_ratio = self._estimate_files_encoding_ratio()
+            meta_provider._set_encoding_ratio(self._encoding_ratio)
+        else:
+            self._encoding_ratio = IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+    def _read_stream(
+        self,
+        f: "pyarrow.NativeFile",
+        path: str,
+    ) -> Iterator[Block]:
+        from PIL import Image, UnidentifiedImageError
+
+        data = f.readall()
+
+        try:
+            image = Image.open(io.BytesIO(data))
+        except UnidentifiedImageError as e:
+            raise ValueError(f"PIL couldn't load image file at path '{path}'.") from e
+
+        if self.size is not None and image.size != tuple(reversed(self.size)):
+            height, width = self.size
+            image = image.resize((width, height), resample=Image.BILINEAR)
+        if self.mode is not None and image.mode != self.mode:
+            image = image.convert(self.mode)
+
+        builder = DelegatingBlockBuilder()
+        array = np.asarray(image)
+        item = {"image": array}
+        builder.add(item)
+        block = builder.build()
+
+        yield block
+
+    def _rows_per_file(self):
+        return 1
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        total_size = 0
+        for file_size in self._file_sizes():
+            # NOTE: check if file size is not None, because some metadata providers
+            # may not provide file size information.
+            if file_size is not None:
+                total_size += file_size
+        return total_size * self._encoding_ratio
+
+    def _estimate_files_encoding_ratio(self) -> float:
+        """Return an estimate of the image files encoding ratio."""
+        start_time = time.perf_counter()
+        # Filter out empty file to avoid noise.
+        non_empty_path_and_size = list(
+            filter(lambda p: p[1] > 0, zip(self._paths(), self._file_sizes()))
+        )
+        num_files = len(non_empty_path_and_size)
+        if num_files == 0:
+            logger.warning(
+                "All input image files are empty. "
+                "Use on-disk file size to estimate images in-memory size."
+            )
+            return IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+        if self.size is not None and self.mode is not None:
+            # Use image size and mode to calculate data size for all images,
+            # because all images are homogeneous with same size after resizing.
+            # Resizing is enforced when reading every image in `ImageDatasource`
+            # when `size` argument is provided.
+            if self.mode in ["1", "L", "P"]:
+                dimension = 1
+            elif self.mode in ["RGB", "YCbCr", "LAB", "HSV"]:
+                dimension = 3
+            elif self.mode in ["RGBA", "CMYK", "I", "F"]:
+                dimension = 4
+            else:
+                logger.warning(f"Found unknown image mode: {self.mode}.")
+                return IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+            height, width = self.size
+            single_image_size = height * width * dimension
+            total_estimated_size = single_image_size * num_files
+            total_file_size = sum(p[1] for p in non_empty_path_and_size)
+            ratio = total_estimated_size / total_file_size
+        else:
+            # TODO(chengsu): sample images to estimate data size
+            ratio = IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+        sampling_duration = time.perf_counter() - start_time
+        if sampling_duration > 5:
+            logger.warning(
+                "Image input size estimation took "
+                f"{round(sampling_duration, 2)} seconds."
+            )
+        logger.debug(f"Estimated image encoding ratio from sampling is {ratio}.")
+        return max(ratio, IMAGE_ENCODING_RATIO_ESTIMATE_LOWER_BOUND)
+
+
+class ImageFileMetadataProvider(DefaultFileMetadataProvider):
+    def _set_encoding_ratio(self, encoding_ratio: int):
+        """Set image file encoding ratio, to provide accurate size in bytes metadata."""
+        self._encoding_ratio = encoding_ratio
+
+    def _get_block_metadata(
+        self,
+        paths: List[str],
+        *,
+        rows_per_file: Optional[int],
+        file_sizes: List[Optional[int]],
+    ) -> BlockMetadata:
+        metadata = super()._get_block_metadata(
+            paths, rows_per_file=rows_per_file, file_sizes=file_sizes
+        )
+        if metadata.size_bytes is not None:
+            metadata = replace(
+                metadata, size_bytes=int(metadata.size_bytes * self._encoding_ratio)
+            )
+        return metadata

python/ray/data/_internal/datasource/json_datasink.py

@@ -0,0 +1,36 @@
+from typing import Any, Callable, Dict, Optional
+
+import pyarrow
+
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_based_datasource import _resolve_kwargs
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class JSONDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        pandas_json_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        pandas_json_args: Optional[Dict[str, Any]] = None,
+        file_format: str = "json",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        if pandas_json_args_fn is None:
+            pandas_json_args_fn = lambda: {}  # noqa: E731
+
+        if pandas_json_args is None:
+            pandas_json_args = {}
+
+        self.pandas_json_args_fn = pandas_json_args_fn
+        self.pandas_json_args = pandas_json_args
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        writer_args = _resolve_kwargs(self.pandas_json_args_fn, **self.pandas_json_args)
+        orient = writer_args.pop("orient", "records")
+        lines = writer_args.pop("lines", True)
+
+        block.to_pandas().to_json(file, orient=orient, lines=lines, **writer_args)

python/ray/data/_internal/datasource/json_datasource.py

@@ -0,0 +1,293 @@
+import io
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import pandas as pd
+
+from ray.data._internal.pandas_block import PandasBlockAccessor
+from ray.data._internal.tensor_extensions.arrow import pyarrow_table_from_pydict
+from ray.data.context import DataContext
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+logger = logging.getLogger(__name__)
+
+JSON_FILE_EXTENSIONS = [
+    "json",
+    "jsonl",
+    # gzip-compressed files
+    "json.gz",
+    "jsonl.gz",
+    # Brotli-compressed fi;es
+    "json.br",
+    "jsonl.br",
+    # Zstandard-compressed files
+    "json.zst",
+    "jsonl.zst",
+    # lz4-compressed files
+    "json.lz4",
+    "jsonl.lz4",
+]
+
+
+class ArrowJSONDatasource(FileBasedDatasource):
+    """JSON datasource, for reading and writing JSON and JSONL files."""
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        *,
+        arrow_json_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        from pyarrow import json
+
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        if arrow_json_args is None:
+            arrow_json_args = {}
+
+        self.read_options = arrow_json_args.pop(
+            "read_options", json.ReadOptions(use_threads=False)
+        )
+        self.arrow_json_args = arrow_json_args
+
+    def _read_with_pyarrow_read_json(self, buffer: "pyarrow.lib.Buffer"):
+        """Read with PyArrow JSON reader, trying to auto-increase the
+        read block size in the case of the read object
+        straddling block boundaries."""
+        import pyarrow as pa
+        import pyarrow.json as pajson
+
+        # When reading large files, the default block size configured in PyArrow can be
+        # too small, resulting in the following error: `pyarrow.lib.ArrowInvalid:
+        # straddling object straddles two block boundaries (try to increase block
+        # size?)`. More information on this issue can be found here:
+        # https://github.com/apache/arrow/issues/25674
+        # The read will be retried with geometrically increasing block size
+        # until the size reaches `DataContext.get_current().target_max_block_size`.
+        # The initial block size will start at the PyArrow default block size
+        # or it can be manually set through the `read_options` parameter as follows.
+        # >>> import pyarrow.json as pajson
+        # >>> block_size = 10 << 20 # Set block size to 10MB
+        # >>> ray.data.read_json(  # doctest: +SKIP
+        # ...     "s3://anonymous@ray-example-data/log.json",
+        # ...     read_options=pajson.ReadOptions(block_size=block_size)
+        # ... )
+
+        init_block_size = self.read_options.block_size
+        max_block_size = DataContext.get_current().target_max_block_size
+        while True:
+            try:
+                yield pajson.read_json(
+                    io.BytesIO(buffer),
+                    read_options=self.read_options,
+                    **self.arrow_json_args,
+                )
+                self.read_options.block_size = init_block_size
+                break
+            except pa.ArrowInvalid as e:
+                if "straddling object straddles two block boundaries" in str(e):
+                    if (
+                        max_block_size is None
+                        or self.read_options.block_size < max_block_size
+                    ):
+                        # Increase the block size in case it was too small.
+                        logger.debug(
+                            f"JSONDatasource read failed with "
+                            f"block_size={self.read_options.block_size}. Retrying with "
+                            f"block_size={self.read_options.block_size * 2}."
+                        )
+                        self.read_options.block_size *= 2
+                    else:
+                        raise pa.ArrowInvalid(
+                            f"{e} - Auto-increasing block size to "
+                            f"{self.read_options.block_size} bytes failed. "
+                            f"Please try manually increasing the block size through "
+                            f"the `read_options` parameter to a larger size. "
+                            f"For example: `read_json(..., read_options="
+                            f"pyarrow.json.ReadOptions(block_size=10 << 25))`"
+                            f"More information on this issue can be found here: "
+                            f"https://github.com/apache/arrow/issues/25674"
+                        )
+                else:
+                    # unrelated error, simply reraise
+                    raise e
+
+    def _read_with_python_json(self, buffer: "pyarrow.lib.Buffer"):
+        """Fallback method to read JSON files with Python's native json.load(),
+        in case the default pyarrow json reader fails."""
+        import json
+
+        import pyarrow as pa
+
+        # Check if the buffer is empty
+        if buffer.size == 0:
+            return
+
+        parsed_json = json.load(io.BytesIO(buffer))
+        try:
+            yield pa.Table.from_pylist(parsed_json)
+        except AttributeError as e:
+            # For PyArrow < 7.0.0, `pa.Table.from_pylist()` is not available.
+            # Construct a dict from the list and call
+            # `pa.Table.from_pydict()` instead.
+            assert "no attribute 'from_pylist'" in str(e), str(e)
+            from collections import defaultdict
+
+            dct = defaultdict(list)
+            for row in parsed_json:
+                for k, v in row.items():
+                    dct[k].append(v)
+            yield pyarrow_table_from_pydict(dct)
+
+    # TODO(ekl) The PyArrow JSON reader doesn't support streaming reads.
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        import pyarrow as pa
+
+        buffer: pa.lib.Buffer = f.read_buffer()
+
+        try:
+            yield from self._read_with_pyarrow_read_json(buffer)
+        except pa.ArrowInvalid as e:
+            # If read with PyArrow fails, try falling back to native json.load().
+            logger.warning(
+                f"Error reading with pyarrow.json.read_json(). "
+                f"Falling back to native json.load(), which may be slower. "
+                f"PyArrow error was:\n{e}"
+            )
+            yield from self._read_with_python_json(buffer)
+
+
+class PandasJSONDatasource(FileBasedDatasource):
+
+    # Buffer size in bytes for reading files. Default is 1MB.
+    #
+    # pandas reads data in small chunks (~8 KiB), which leads to many costly
+    # small read requests when accessing cloud storage. To reduce overhead and
+    # improve performance, we wrap the file in a larger buffered reader that
+    # reads bigger blocks at once.
+    _BUFFER_SIZE = 1024**2
+
+    # In the case of zipped json files, we cannot infer the chunk_size.
+    _DEFAULT_CHUNK_SIZE = 10000
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        target_output_size_bytes: int,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self._target_output_size_bytes = target_output_size_bytes
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        chunksize = self._estimate_chunksize(f)
+
+        with StrictBufferedReader(f, buffer_size=self._BUFFER_SIZE) as stream:
+            if chunksize is None:
+                # When chunksize=None, pandas returns DataFrame directly
+                # (no context manager).
+                df = pd.read_json(stream, chunksize=chunksize, lines=True)
+                yield _cast_range_index_to_string(df)
+            else:
+                # When chunksize is a number, pandas returns JsonReader
+                # (supports context manager).
+                with pd.read_json(stream, chunksize=chunksize, lines=True) as reader:
+                    for df in reader:
+                        yield _cast_range_index_to_string(df)
+
+    def _estimate_chunksize(self, f: "pyarrow.NativeFile") -> Optional[int]:
+        """Estimate the chunksize by sampling the first row.
+
+        This is necessary to avoid OOMs while reading the file.
+        """
+
+        if not f.seekable():
+            return self._DEFAULT_CHUNK_SIZE
+
+        # ``_read_stream`` can be recreated on the same file handle when
+        # ``FileBasedDatasource`` retries a transient read error.
+        f.seek(0)
+
+        if self._target_output_size_bytes is None:
+            return None
+
+        try:
+            with StrictBufferedReader(f, buffer_size=self._BUFFER_SIZE) as stream:
+                with pd.read_json(stream, chunksize=1, lines=True) as reader:
+                    try:
+                        df = _cast_range_index_to_string(next(reader))
+                    except StopIteration:
+                        return 1
+
+            block_accessor = PandasBlockAccessor.for_block(df)
+            if block_accessor.num_rows() == 0:
+                chunksize = 1
+            else:
+                bytes_per_row = block_accessor.size_bytes() / block_accessor.num_rows()
+                chunksize = max(
+                    round(self._target_output_size_bytes / bytes_per_row), 1
+                )
+
+            return chunksize
+        finally:
+            # Reset file pointer to the beginning for the actual read and for any
+            # subsequent retry that reuses the same file handle.
+            f.seek(0)
+
+    def _open_input_source(
+        self,
+        filesystem: "pyarrow.fs.FileSystem",
+        path: str,
+        **open_args,
+    ) -> "pyarrow.NativeFile":
+
+        compression = self.resolve_compression(path, open_args)
+
+        if compression is None:
+            # We use a seekable file to estimate chunksize.
+            return filesystem.open_input_file(path)
+
+        return super()._open_input_source(filesystem, path, **open_args)
+
+
+def _cast_range_index_to_string(df: pd.DataFrame):
+    # NOTE: PandasBlockAccessor doesn't support RangeIndex, so we need to convert
+    # to string.
+    if isinstance(df.columns, pd.RangeIndex):
+        df.columns = df.columns.astype(str)
+    return df
+
+
+class StrictBufferedReader(io.RawIOBase):
+    """Wrapper that prevents premature file closure and ensures full-buffered reads.
+
+    This is necessary for two reasons:
+    1. The datasource reads the file twice -- first to sample and determine the chunk size,
+       and again to load the actual data. Since pandas assumes ownership of the file and
+       may close it, we prevent that by explicitly detaching the underlying file before
+       closing the buffer.
+
+    2. pandas wraps the file in a TextIOWrapper to decode bytes into text. TextIOWrapper
+       prefers calling read1(), which doesn't prefetch for random-access files
+       (e.g., from PyArrow). This wrapper forces all reads through the full buffer to
+       avoid inefficient small-range S3 GETs.
+    """
+
+    def __init__(self, file: io.RawIOBase, buffer_size: int):
+        self._file = io.BufferedReader(file, buffer_size=buffer_size)
+
+    def read(self, size=-1, /):
+        return self._file.read(size)
+
+    def readable(self) -> bool:
+        return True
+
+    def close(self):
+        if not self.closed:
+            self._file.detach()
+            super().close()

python/ray/data/_internal/datasource/kafka_datasink.py

@@ -0,0 +1,296 @@
+"""Kafka datasink
+
+This module provides a Kafka datasink implementation for Ray Data.
+
+Requires:
+    - confluent-kafka: https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/
+"""
+
+import json
+import logging
+from collections.abc import Iterable, Mapping
+from enum import Enum
+from typing import Any, Optional
+
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+logger = logging.getLogger(__name__)
+
+# Polling/flush constants.  These are intentionally conservative defaults
+# that work well for typical workloads.  All three can be tuned indirectly
+# through producer_config (e.g. queue.buffering.max.messages, message.timeout.ms)
+# or overridden by subclassing if needed.
+
+# Number of messages to produce before polling for delivery reports.
+# At ~1.5 KB per message (a common upper bound), 10 000 messages ≈ 15 MB
+# of buffered data — well within librdkafka's default queue limits while
+# keeping Python→C crossing overhead low.
+_POLL_BATCH_SIZE = 10000
+# Timeout (seconds) for the final flush that waits for all in-flight messages
+_FLUSH_TIMEOUT_S = 30
+# Timeout (seconds) when polling to drain the queue after a BufferError
+_BUFFER_FULL_POLL_TIMEOUT_S = 10
+
+
+class SerializerFormat(str, Enum):
+    """Supported serialization formats for Kafka message keys and values."""
+
+    JSON = "json"
+    STRING = "string"
+    BYTES = "bytes"
+
+
+def _serialize(data: Any, serializer: SerializerFormat) -> bytes:
+    """Serialize *data* according to *serializer*.
+
+    This is a standalone function so it can be used without a class instance.
+    """
+    if serializer == SerializerFormat.JSON:
+        return json.dumps(data).encode("utf-8")
+    elif serializer == SerializerFormat.STRING:
+        return str(data).encode("utf-8")
+    else:  # BYTES
+        return data if isinstance(data, bytes) else str(data).encode("utf-8")
+
+
+class KafkaDatasink(Datasink):
+    """
+    Ray Data sink for writing to Apache Kafka topics using confluent-kafka.
+
+    Writes blocks of data to Kafka with configurable serialization
+    and producer settings.
+
+    Delivery guarantees:
+        This sink provides best-effort delivery. Partial writes can occur if a
+        task fails midway (already-flushed messages are not rolled back), and
+        duplicates are possible when the system retries a failed task (e.g., on
+        node failure), since each attempt re-sends all messages from scratch
+        without Kafka transactions or cross-task deduplication.
+    """
+
+    def __init__(
+        self,
+        topic: str,
+        bootstrap_servers: str,
+        key_field: Optional[str] = None,
+        key_serializer: SerializerFormat = SerializerFormat.STRING,
+        value_serializer: SerializerFormat = SerializerFormat.JSON,
+        producer_config: Optional[dict[str, Any]] = None,
+    ):
+        """
+        Initialize Kafka sink.
+
+        Args:
+            topic: Kafka topic name
+            bootstrap_servers: Comma-separated Kafka broker addresses (e.g., 'localhost:9092')
+            key_field: Optional field name to use as message key
+            key_serializer: Key serialization format ('json', 'string', or 'bytes')
+            value_serializer: Value serialization format ('json', 'string', or 'bytes')
+            producer_config: Additional Kafka producer configuration.
+                Uses confluent-kafka / librdkafka configuration keys
+                (see https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md).
+                Example: ``{"linger.ms": 5, "acks": "all"}``.
+                The ``bootstrap.servers`` option is derived from ``bootstrap_servers``
+                and cannot be overridden here.
+        """
+        _check_import(self, module="confluent_kafka", package="confluent-kafka")
+
+        try:
+            key_serializer = SerializerFormat(key_serializer)
+        except ValueError:
+            raise ValueError(
+                f"key_serializer must be one of "
+                f"{[s.value for s in SerializerFormat]}, "
+                f"got '{key_serializer}'"
+            )
+        try:
+            value_serializer = SerializerFormat(value_serializer)
+        except ValueError:
+            raise ValueError(
+                f"value_serializer must be one of "
+                f"{[s.value for s in SerializerFormat]}, "
+                f"got '{value_serializer}'"
+            )
+
+        self.topic = topic
+        self.bootstrap_servers = bootstrap_servers
+        self.key_field = key_field
+        self.key_serializer = key_serializer
+        self.value_serializer = value_serializer
+        self.producer_config = producer_config or {}
+
+    @staticmethod
+    def _row_to_dict(row: Any) -> Any:
+        """Convert a Ray data row to a plain dict if possible.
+
+        Handles Ray's internal row types (ArrowRow, PandasRow), namedtuples,
+        and generic Mappings.  Returns the input unchanged for primitives.
+        """
+        if isinstance(row, dict):
+            return row
+        if hasattr(row, "as_pydict"):
+            return row.as_pydict()
+        if hasattr(row, "_asdict"):
+            return row._asdict()
+        if isinstance(row, Mapping):
+            return dict(row)
+        return row
+
+    def _serialize_value(self, value: Any) -> bytes:
+        """Serialize value based on configured format."""
+        return _serialize(value, self.value_serializer)
+
+    def _serialize_key(self, key: Any) -> bytes:
+        """Serialize key based on configured format."""
+        return _serialize(key, self.key_serializer)
+
+    def _extract_key(self, row_dict: Any) -> Optional[bytes]:
+        """Extract and serialize the message key from a row dict.
+
+        Returns ``None`` when no ``key_field`` is configured, when the row is
+        not a dict, or when the key field is absent/``None`` in the row.  A
+        ``None`` key tells the Kafka producer to use the default partitioner
+        (round-robin or sticky partitioning depending on librdkafka version),
+        distributing messages evenly across partitions.
+        """
+        if self.key_field and isinstance(row_dict, dict):
+            key_value = row_dict.get(self.key_field)
+            if key_value is not None:
+                return self._serialize_key(key_value)
+        return None
+
+    def _produce_with_retry(self, producer, value, key, on_delivery):
+        """Produce a single message, retrying once on ``BufferError``.
+
+        ``producer.produce()`` is asynchronous — it enqueues the message into
+        librdkafka's internal buffer and returns immediately.  Actual delivery
+        happens in a background thread; results are reported via the
+        *on_delivery* callback when ``producer.poll()`` or ``producer.flush()``
+        is called.
+
+        If the internal buffer is full, a ``BufferError`` is raised.  We handle
+        this by polling to drain completed deliveries (which frees buffer
+        space) and retrying once.
+        """
+        try:
+            producer.produce(
+                self.topic,
+                value=value,
+                key=key,
+                on_delivery=on_delivery,
+            )
+        except BufferError:
+            # Internal queue is full — poll to serve delivery callbacks and
+            # free space, then retry.  The poll timeout caps how long we block
+            # waiting for the broker to acknowledge in-flight messages.
+            producer.poll(_BUFFER_FULL_POLL_TIMEOUT_S)
+            try:
+                producer.produce(
+                    self.topic,
+                    value=value,
+                    key=key,
+                    on_delivery=on_delivery,
+                )
+            except BufferError:
+                raise RuntimeError(
+                    f"Kafka producer queue is still full after "
+                    f"{_BUFFER_FULL_POLL_TIMEOUT_S}s of polling "
+                    f"for topic '{self.topic}'. "
+                    f"Consider increasing queue.buffering.max.messages "
+                    f"in producer_config."
+                )
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> Any:
+        """
+        Write blocks of data to Kafka.
+
+        Args:
+            blocks: Iterable of Ray data blocks
+            ctx: Ray data context
+
+        Returns:
+            Dict with ``total_records`` and ``failed_messages`` counts.
+        """
+        from confluent_kafka import KafkaException, Producer
+
+        # Build confluent config
+        config: dict[str, Any] = {
+            "bootstrap.servers": self.bootstrap_servers,
+        }
+        for k, v in self.producer_config.items():
+            if k == "bootstrap.servers":
+                logger.warning(
+                    "Ignoring 'bootstrap.servers' from producer_config; "
+                    "use bootstrap_servers parameter instead."
+                )
+                continue
+            config[k] = v
+
+        producer = Producer(config)
+        total_records = 0
+        remaining = 0
+        # Mutable container so on_delivery callback can update without nonlocal
+        delivery_state = {"failed": 0, "first_exception": None}
+
+        def on_delivery(err, msg):
+            if err is not None:
+                delivery_state["failed"] += 1
+                if delivery_state["first_exception"] is None:
+                    delivery_state["first_exception"] = KafkaException(err)
+
+        try:
+            for block in blocks:
+                block_accessor = BlockAccessor.for_block(block)
+
+                for row in block_accessor.iter_rows(public_row_format=False):
+                    row_dict = self._row_to_dict(row)
+                    key = self._extract_key(row_dict)
+                    value = self._serialize_value(row_dict)
+
+                    self._produce_with_retry(producer, value, key, on_delivery)
+                    total_records += 1
+
+                    # Periodically poll to serve delivery report callbacks
+                    # and avoid unbounded internal queue growth.
+                    if total_records % _POLL_BATCH_SIZE == 0:
+                        producer.poll(0)
+
+            # Final flush: blocks until all in-flight messages are delivered
+            # or the timeout expires.  Returns the count of messages still
+            # queued (0 = everything delivered).  Does NOT raise on timeout.
+            remaining = producer.flush(timeout=_FLUSH_TIMEOUT_S)
+
+        except KafkaException as e:
+            raise RuntimeError(
+                f"Failed to write to Kafka topic '{self.topic}': {e}"
+            ) from e
+
+        if remaining > 0:
+            raise RuntimeError(
+                f"{remaining} out of {total_records} messages were still "
+                f"in-flight after flush timeout for topic '{self.topic}'. "
+                f"This usually means the broker is unreachable."
+            )
+
+        failed_messages = delivery_state["failed"]
+        if failed_messages > 0:
+            raise RuntimeError(
+                f"Failed to write {failed_messages} out of {total_records} "
+                f"messages to Kafka topic '{self.topic}'."
+            ) from delivery_state["first_exception"]
+
+        # Logged once per write task (one task per data block partition).
+        logger.debug(
+            "Wrote %d records to Kafka topic '%s'.",
+            total_records,
+            self.topic,
+        )
+
+        return {"total_records": total_records, "failed_messages": failed_messages}

python/ray/data/_internal/datasource/kafka_datasource.py

@@ -0,0 +1,817 @@
+"""Kafka datasource for bounded data reads.
+
+This module provides a Kafka datasource implementation for Ray Data that supports
+bounded reads with offset-based range queries.
+
+Message keys and values are returned as raw bytes to support any serialization format
+(JSON, Avro, Protobuf, etc.). Users can decode them using map operations.
+
+Requires:
+    - confluent-kafka: https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/
+"""
+
+import logging
+import time
+import warnings
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
+
+import pyarrow as pa
+
+if TYPE_CHECKING:
+    from confluent_kafka import Consumer, TopicPartition
+
+from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource import Datasource, ReadTask
+
+PartitionOffsets = Dict[int, Union[int, str]]
+PerPartitionOffsets = Dict[str, PartitionOffsets]
+
+logger = logging.getLogger(__name__)
+
+# Mapping from kafka-python style KafkaAuthConfig fields to Confluent/librdkafka config.
+# TODO(youcheng): Remove this mapping and use consumer_config directly.
+_KAFKA_AUTH_TO_CONFLUENT: Dict[str, str] = {
+    "security_protocol": "security.protocol",
+    "sasl_mechanism": "sasl.mechanism",
+    "sasl_plain_username": "sasl.username",
+    "sasl_plain_password": "sasl.password",
+    "sasl_kerberos_service_name": "sasl.kerberos.service.name",
+    "sasl_kerberos_name": "sasl.kerberos.principal",
+    "ssl_cafile": "ssl.ca.location",
+    "ssl_certfile": "ssl.certificate.location",
+    "ssl_keyfile": "ssl.key.location",
+    "ssl_password": "ssl.key.password",
+    "ssl_ciphers": "ssl.cipher.suites",
+    "ssl_crlfile": "ssl.crl.location",
+    # Note: ssl_check_hostname is intentionally NOT mapped due to semantics mismatch.
+}
+
+
+KAFKA_TOPIC_METADATA_TIMEOUT_S = 10
+KAFKA_QUERY_OFFSET_TIMEOUT_S = 10
+# Cap each consume timeout to keep responsiveness of timeout/position checks.
+KAFKA_CONSUME_TIMEOUT_MAX_S = 10  # 10 seconds per consume call
+
+KAFKA_MSG_SCHEMA = pa.schema(
+    [
+        ("offset", pa.int64()),
+        ("key", pa.binary()),
+        ("value", pa.binary()),
+        ("topic", pa.string()),
+        ("partition", pa.int32()),
+        ("timestamp", pa.int64()),  # Kafka timestamp in milliseconds
+        (
+            "timestamp_type",
+            pa.int32(),
+        ),  # 0=TIMESTAMP_NOT_AVAILABLE, 1=TIMESTAMP_CREATE_TIME, 2=TIMESTAMP_LOG_APPEND_TIME
+        ("headers", pa.map_(pa.string(), pa.binary())),  # Message headers
+    ]
+)
+
+
+@dataclass
+class KafkaAuthConfig:
+    """Authentication configuration for Kafka connections.
+
+    Uses standard kafka-python parameter names. See kafka-python documentation
+    for full details: https://kafka-python.readthedocs.io/
+
+    Note: Ray Data maps these options to Confluent/librdkafka config under the hood.
+    Some options have different semantics or are unsupported by the Confluent client;
+    see notes below for those fields. Prefer passing Confluent options directly via
+    ``consumer_config`` where possible.
+
+    security_protocol: Protocol used to communicate with brokers.
+        Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
+        Default: PLAINTEXT.
+    sasl_mechanism: Authentication mechanism when security_protocol
+        is configured for SASL_PLAINTEXT or SASL_SSL. Valid values are:
+        PLAIN, GSSAPI, OAUTHBEARER, SCRAM-SHA-256, SCRAM-SHA-512.
+    sasl_plain_username: username for sasl PLAIN and SCRAM authentication.
+        Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms.
+    sasl_plain_password: password for sasl PLAIN and SCRAM authentication.
+        Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms.
+    sasl_kerberos_name: Constructed gssapi.Name for use with
+        sasl mechanism handshake. If provided, sasl_kerberos_service_name and
+        sasl_kerberos_domain name are ignored. Default: None.
+    sasl_kerberos_service_name: Service name to include in GSSAPI
+        sasl mechanism handshake. Default: 'kafka'
+    sasl_kerberos_domain_name: kerberos domain name to use in GSSAPI
+        sasl mechanism handshake. Default: one of bootstrap servers.
+        Note (Confluent): This option is not supported by Confluent/librdkafka
+        and will be ignored when building the client configuration. Prefer specifying
+        an explicit principal via ``sasl_kerberos_name`` or rely on defaults.
+    sasl_oauth_token_provider: OAuthBearer token provider instance. Default: None.
+        Note (Confluent): Not supported directly; use ``consumer_config`` with
+        ``sasl.oauthbearer.*`` options instead.
+    ssl_context: Pre-configured SSLContext for wrapping socket connections. If provided,
+        all other ssl_* configurations will be ignored. Default: None.
+        Note (Confluent): Passing an SSLContext object is not supported and will be
+        ignored. Use ``ssl_cafile``, ``ssl_certfile``, and ``ssl_keyfile`` instead.
+    ssl_check_hostname: Flag to configure whether ssl handshake should verify that the
+        certificate matches the broker's hostname. Default: True.
+        Note (Confluent): There is no 1:1 equivalent; disabling hostname verification
+        via ``enable.ssl.certificate.verification=False`` would also disable the entire
+        certificate chain verification. To avoid weakening security, this flag is not
+        mapped when False. If you need to disable only hostname verification, set
+        ``ssl.endpoint.identification.algorithm=none`` via ``consumer_config`` (if supported
+        by your librdkafka version).
+    ssl_cafile: Optional filename of ca file to use in certificate verification. Default: None.
+    ssl_certfile: Optional filename of file in pem format containing the client certificate,
+        as well as any ca certificates needed to establish the certificate's authenticity. Default: None.
+    ssl_keyfile: Optional filename containing the client private key. Default: None.
+    ssl_password: Optional password to be used when loading the certificate chain. Default: None.
+    ssl_crlfile: Optional filename containing the CRL to check for certificate expiration. By default,
+        no CRL check is done. When providing a file, only the leaf certificate will be checked against
+        this CRL. The CRL can only be checked with Python 3.4+ or 2.7.9+. Default: None.
+    ssl_ciphers: optionally set the available ciphers for ssl connections. It should be a string in the
+        OpenSSL cipher list format. If no cipher can be selected (because compile-time options or other
+        configuration forbids use of all the specified ciphers), an ssl.SSLError will be raised.
+        See ssl.SSLContext.set_ciphers.
+    """
+
+    # Security protocol
+    security_protocol: Optional[str] = None
+
+    # SASL configuration
+    sasl_mechanism: Optional[str] = None
+    sasl_plain_username: Optional[str] = None
+    sasl_plain_password: Optional[str] = None
+    sasl_kerberos_name: Optional[str] = None
+    sasl_kerberos_service_name: Optional[str] = None
+    sasl_kerberos_domain_name: Optional[str] = None
+    sasl_oauth_token_provider: Optional[Any] = None
+
+    # SSL configuration
+    ssl_context: Optional[Any] = None
+    ssl_check_hostname: Optional[bool] = None
+    ssl_cafile: Optional[str] = None
+    ssl_certfile: Optional[str] = None
+    ssl_keyfile: Optional[str] = None
+    ssl_password: Optional[str] = None
+    ssl_ciphers: Optional[str] = None
+    ssl_crlfile: Optional[str] = None
+
+
+def _handle_deprecated_configs(kafka_auth_config: KafkaAuthConfig) -> None:
+    # Handle special fields with warnings
+    if kafka_auth_config.ssl_context is not None:
+        logger.warning(
+            "ssl_context is not supported by Confluent. Skipping. "
+            "Use KafkaAuthConfig fields ssl_cafile, ssl_certfile, ssl_keyfile instead."
+        )
+    if kafka_auth_config.sasl_oauth_token_provider is not None:
+        logger.warning(
+            "sasl_oauth_token_provider is not supported by Confluent. Skipping. "
+            "Use consumer_config with sasl.oauthbearer.* options instead."
+        )
+    if kafka_auth_config.sasl_kerberos_domain_name is not None:
+        logger.warning(
+            "sasl_kerberos_domain_name is not supported by Confluent and will be ignored. "
+            "Set sasl_kerberos_name (principal) or rely on defaults."
+        )
+    if kafka_auth_config.ssl_check_hostname is False:
+        logger.warning(
+            "ssl_check_hostname=False cannot be mapped safely to Confluent; "
+            "setting enable.ssl.certificate.verification=False would disable all certificate verification. "
+            "Ignoring ssl_check_hostname. If you need to disable only hostname verification, "
+            "configure the client directly via consumer_config (e.g., ssl.endpoint.identification.algorithm=none)."
+        )
+
+
+def _add_authentication_to_config(
+    config: Dict[str, Any], kafka_auth_config: Optional[KafkaAuthConfig]
+) -> None:
+    """Map KafkaAuthConfig (kafka-python style) into Confluent/librdkafka config.
+
+    Special cases:
+      - ssl_context: unsupported; warn and ignore
+      - sasl_oauth_token_provider: unsupported; warn and ignore
+      - sasl_kerberos_domain_name: unsupported; warn and ignore
+      - ssl_check_hostname: not mapped due to semantics; if False, warn and ignore
+    """
+    if not kafka_auth_config:
+        return
+    warnings.warn(
+        "kafka_auth_config (kafka-python style) is deprecated and will be removed in a future release. "
+        "Please provide Confluent/librdkafka options via consumer_config instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+    _handle_deprecated_configs(kafka_auth_config)
+
+    # Map directly compatible fields
+    for key, confluent_key in _KAFKA_AUTH_TO_CONFLUENT.items():
+        val = getattr(kafka_auth_config, key, None)
+        if val is not None:
+            config[confluent_key] = val
+
+
+def _build_confluent_config(
+    bootstrap_servers: List[str],
+    kafka_auth_config: Optional[KafkaAuthConfig] = None,
+    extra: Optional[Dict[str, Any]] = None,
+    user_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Build Confluent config with bootstrap servers and auth.
+
+    Args:
+        bootstrap_servers: List of Kafka broker addresses.
+        kafka_auth_config: Authentication configuration (kafka-python style). Deprecated; prefer consumer_config with Confluent keys. Mutually exclusive with consumer_config.
+        extra: Additional config options.
+        user_config: User-provided config options.
+
+    Returns:
+        Confluent configuration dict.
+    """
+    config: Dict[str, Any] = {
+        "bootstrap.servers": ",".join(bootstrap_servers),
+    }
+    # Map kafka-python-style auth if provided
+    _add_authentication_to_config(config, kafka_auth_config)
+    if extra:
+        config.update(extra)
+    if user_config:
+        if (
+            "bootstrap.servers" in user_config
+            and user_config["bootstrap.servers"] != config["bootstrap.servers"]
+        ):
+            logger.warning(
+                "Ignoring 'bootstrap.servers' from consumer_config; use bootstrap_servers parameter instead."
+            )
+        for k, v in user_config.items():
+            if k == "bootstrap.servers":
+                continue
+            config[k] = v
+    return config
+
+
+def _build_consumer_config_for_read(
+    bootstrap_servers: List[str],
+    kafka_auth_config: Optional[KafkaAuthConfig] = None,
+    consumer_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Build Consumer config for reading messages (Confluent)."""
+    return _build_confluent_config(
+        bootstrap_servers,
+        extra={
+            "enable.auto.commit": False,
+            # Confluent requires a group.id even when using manual assign.
+            "group.id": "ray-data-kafka-reader",
+        },
+        user_config=consumer_config,
+        kafka_auth_config=kafka_auth_config,
+    )
+
+
+def _datetime_to_ms(dt: datetime) -> int:
+    """Convert a datetime to milliseconds since epoch (UTC).
+
+    If the datetime has no timezone info (i.e., ``tzinfo is None``),
+    it is assumed to be UTC. Timezone-aware datetimes are converted to
+    UTC automatically via ``datetime.timestamp()``.
+
+    Args:
+        dt: A datetime object, with or without timezone info.
+
+    Returns:
+        Milliseconds since Unix epoch.
+    """
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    return int(dt.timestamp() * 1000)
+
+
+def _validate_offsets(
+    start_offset: Union[int, datetime, Literal["earliest"], PerPartitionOffsets],
+    end_offset: Union[int, datetime, Literal["latest"], PerPartitionOffsets],
+) -> None:
+    if isinstance(start_offset, dict):
+        for topic, partition_map in start_offset.items():
+            if not isinstance(partition_map, dict):
+                raise ValueError(
+                    f"start_offset[{topic!r}] must be a dict mapping "
+                    f"partition_id (int) to offset (int or str)."
+                )
+            for partition_id, offset in partition_map.items():
+                if not isinstance(partition_id, int):
+                    raise ValueError(
+                        f"start_offset[{topic!r}] keys must be integers "
+                        f"(partition IDs), got {type(partition_id).__name__!r}."
+                    )
+                if isinstance(offset, str) and offset == "latest":
+                    raise ValueError(
+                        f"start_offset[{topic!r}][{partition_id}] cannot be 'latest'."
+                    )
+    else:
+        if isinstance(start_offset, int) and isinstance(end_offset, int):
+            if start_offset > end_offset:
+                raise ValueError("start_offset must be less than end_offset")
+        if isinstance(start_offset, datetime) and isinstance(end_offset, datetime):
+            if _datetime_to_ms(start_offset) > _datetime_to_ms(end_offset):
+                raise ValueError("start_offset must be less than end_offset")
+        if isinstance(start_offset, str) and start_offset == "latest":
+            raise ValueError("start_offset cannot be 'latest'")
+
+    if isinstance(end_offset, dict):
+        for topic, partition_map in end_offset.items():
+            if not isinstance(partition_map, dict):
+                raise ValueError(
+                    f"end_offset[{topic!r}] must be a dict mapping "
+                    f"partition_id (int) to offset (int or str)."
+                )
+            for partition_id, offset in partition_map.items():
+                if not isinstance(partition_id, int):
+                    raise ValueError(
+                        f"end_offset[{topic!r}] keys must be integers "
+                        f"(partition IDs), got {type(partition_id).__name__!r}."
+                    )
+                if isinstance(offset, str) and offset == "earliest":
+                    raise ValueError(
+                        f"end_offset[{topic!r}][{partition_id}] cannot be 'earliest'."
+                    )
+    else:
+        if isinstance(end_offset, str) and end_offset == "earliest":
+            raise ValueError("end_offset cannot be 'earliest'")
+
+
+def _resolve_datetime_to_offset(
+    consumer: "Consumer",
+    topic_partition: "TopicPartition",
+    dt: datetime,
+    fallback_offset: int,
+) -> int:
+    """Resolve a datetime to an integer offset via offsets_for_times.
+
+    Returns fallback_offset when offsets_for_times returns empty or offset < 0
+    (e.g., datetime in future or no messages at that time).
+    """
+    from confluent_kafka import TopicPartition
+
+    timestamp_ms = _datetime_to_ms(dt)
+    tp_with_ts = TopicPartition(
+        topic_partition.topic, topic_partition.partition, timestamp_ms
+    )
+    result = consumer.offsets_for_times(
+        [tp_with_ts], timeout=KAFKA_QUERY_OFFSET_TIMEOUT_S
+    )
+    if result and result[0].offset >= 0:
+        return result[0].offset
+    return fallback_offset
+
+
+def _resolve_offsets(
+    consumer: "Consumer",
+    topic_partition: "TopicPartition",
+    start_offset: Union[int, datetime, Literal["earliest"]],
+    end_offset: Union[int, datetime, Literal["latest"]],
+) -> Tuple[int, int]:
+    """Resolve start and end offsets to actual integer offsets.
+
+    Handles int offsets, "earliest"/"latest" strings, and datetime objects.
+    For datetime objects, uses ``consumer.offsets_for_times()`` to find the
+    earliest offset whose timestamp is >= the given datetime.
+
+    Args:
+        consumer: Confluent Kafka consumer instance.
+        topic_partition: TopicPartition to resolve offsets for.
+        start_offset: Start offset (int, datetime, or "earliest").
+        end_offset: End offset (int, datetime, or "latest").
+
+    Returns:
+        Tuple of (resolved_start_offset, resolved_end_offset).
+    """
+    # TODO(youcheng): add retry logic for this call.
+    low, high = consumer.get_watermark_offsets(
+        topic_partition, timeout=KAFKA_QUERY_OFFSET_TIMEOUT_S
+    )
+    earliest_offset = low
+    latest_offset = high
+
+    # Keep original values for error messages
+    original_start = start_offset
+    original_end = end_offset
+
+    if start_offset == "earliest" or start_offset is None:
+        start_offset = earliest_offset
+    elif isinstance(start_offset, datetime):
+        # fallback to latest_offset if the start_offset is in the future, so the read range is empty (start == end).
+        start_offset = _resolve_datetime_to_offset(
+            consumer, topic_partition, start_offset, latest_offset
+        )
+
+    if end_offset == "latest" or end_offset is None:
+        end_offset = latest_offset
+    elif isinstance(end_offset, datetime):
+        end_offset = _resolve_datetime_to_offset(
+            consumer, topic_partition, end_offset, latest_offset
+        )
+
+    # Clamp end_offset to the high watermark so we never try to read beyond
+    # what is currently available.  This prevents the read loop from polling
+    # indefinitely when a user-supplied integer end_offset exceeds the number
+    # of messages in the partition.
+    end_offset = min(end_offset, latest_offset)
+
+    if isinstance(start_offset, int) and start_offset < earliest_offset:
+        logger.warning(
+            f"start_offset ({start_offset}) is below the earliest available offset "
+            f"({earliest_offset}) for partition {topic_partition.partition} in topic "
+            f"{topic_partition.topic} (data may have been deleted by Kafka retention). "
+            f"Falling back to earliest available offset ({earliest_offset})."
+        )
+        start_offset = earliest_offset
+
+    if start_offset > end_offset:
+        start_str = (
+            f"{original_start}"
+            if original_start == start_offset
+            else f"{original_start} (resolved to {start_offset})"
+        )
+        end_str = (
+            f"{original_end}"
+            if original_end == end_offset
+            else f"{original_end} (resolved to {end_offset})"
+        )
+        raise ValueError(
+            f"start_offset ({start_str}) > end_offset ({end_str}) "
+            f"for partition {topic_partition.partition} in topic {topic_partition.topic}"
+        )
+
+    return start_offset, end_offset
+
+
+class KafkaDatasource(Datasource):
+    """Kafka datasource for reading from Kafka topics with bounded reads."""
+
+    # Batch size for incremental block yielding
+    BATCH_SIZE_FOR_YIELD = 1000
+
+    def __init__(
+        self,
+        topics: Union[str, List[str]],
+        bootstrap_servers: Union[str, List[str]],
+        start_offset: Union[int, datetime, Literal["earliest"], PerPartitionOffsets],
+        end_offset: Union[int, datetime, Literal["latest"], PerPartitionOffsets],
+        kafka_auth_config: Optional[KafkaAuthConfig] = None,
+        consumer_config: Optional[Dict[str, Any]] = None,
+        timeout_ms: Optional[int] = None,
+    ):
+        """Initialize Kafka datasource.
+
+        Args:
+            topics: Kafka topic name(s) to read from.
+            bootstrap_servers: Kafka broker addresses (string or list of strings).
+            start_offset: Starting position. Can be:
+                - int: Offset number
+                - datetime: Read from the first message at or after this time.
+                  datetimes with no timezone info are treated as UTC.
+                - str: "earliest"
+            end_offset: Ending position (exclusive). Can be:
+                - int: Offset number
+                - datetime: Read up to (but not including) the first message
+                  at or after this time. datetimes with no timezone info are treated as UTC.
+                - str: "latest"
+            kafka_auth_config: Authentication configuration (kafka-python style). Deprecated; prefer consumer_config with Confluent keys. Mutually exclusive with consumer_config.
+            consumer_config: Confluent/librdkafka consumer configuration dict.
+                Keys and values are passed through to the underlying client. The
+                `bootstrap.servers` option is derived from `bootstrap_servers` and
+                cannot be overridden here.
+            timeout_ms: Optional timeout in milliseconds for every read task to poll until reaching end_offset.
+                If None (default), no task-level timeout is applied and each read task
+                will poll until it reaches end_offset. If set, the read task will stop
+                polling after the timeout and return the messages it has read so far.
+
+        Raises:
+            ValueError: If required configuration is missing.
+            ImportError: If confluent-kafka is not installed.
+        """
+        _check_import(self, module="confluent_kafka", package="confluent-kafka")
+
+        if not topics:
+            raise ValueError("topics cannot be empty")
+
+        if not bootstrap_servers:
+            raise ValueError("bootstrap_servers cannot be empty")
+
+        if timeout_ms is not None and timeout_ms <= 0:
+            raise ValueError("timeout_ms must be positive")
+
+        _validate_offsets(start_offset, end_offset)
+
+        # Validate bootstrap_servers format
+        if isinstance(bootstrap_servers, str):
+            if not bootstrap_servers or ":" not in bootstrap_servers:
+                raise ValueError(
+                    f"Invalid bootstrap_servers format: {bootstrap_servers}. "
+                    "Expected 'host:port' or list of 'host:port' strings."
+                )
+        elif isinstance(bootstrap_servers, list):
+            if not bootstrap_servers:
+                raise ValueError("bootstrap_servers cannot be empty list")
+            for server in bootstrap_servers:
+                if not isinstance(server, str) or ":" not in server:
+                    raise ValueError(
+                        f"Invalid bootstrap_servers format: {server}. "
+                        "Expected 'host:port' string."
+                    )
+
+        # Disallow specifying both config styles at once to avoid ambiguity.
+        if kafka_auth_config is not None and consumer_config is not None:
+            raise ValueError(
+                "Provide only one of kafka_auth_config (deprecated) or consumer_config, not both."
+            )
+
+        self._topics = topics if isinstance(topics, list) else [topics]
+        self._bootstrap_servers = (
+            bootstrap_servers
+            if isinstance(bootstrap_servers, list)
+            else [bootstrap_servers]
+        )
+        self._start_offset = start_offset
+        self._end_offset = end_offset
+        self._kafka_auth_config = kafka_auth_config
+        self._consumer_config = consumer_config
+        self._timeout_ms = timeout_ms
+        self._target_max_block_size = DataContext.get_current().target_max_block_size
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        """Return an estimate of the in-memory data size, or None if unknown."""
+        return None
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        """Create read tasks for Kafka partitions.
+
+        Creates one read task per partition.
+        Each task reads from a single partition of a single topic.
+
+        Args:
+            parallelism: This argument is deprecated.
+            per_task_row_limit: Maximum number of rows per read task.
+            data_context: The data context to use to get read tasks. This is not used by this datasource.
+
+        Returns:
+            List of ReadTask objects, one per partition.
+        """
+        from confluent_kafka import Consumer
+
+        consumer_config = _build_consumer_config_for_read(
+            self._bootstrap_servers, self._kafka_auth_config, self._consumer_config
+        )
+        discovery_consumer = Consumer(consumer_config)
+        try:
+            metadata = discovery_consumer.list_topics(
+                timeout=KAFKA_TOPIC_METADATA_TIMEOUT_S
+            )
+
+            topic_partitions: List[Tuple[str, int]] = []
+            for topic in self._topics:
+                if topic not in metadata.topics:
+                    raise ValueError(
+                        f"Topic {topic} has no partitions or doesn't exist"
+                    )
+                topic_meta = metadata.topics[topic]
+                if not topic_meta.partitions:
+                    raise ValueError(
+                        f"Topic {topic} has no partitions or doesn't exist"
+                    )
+                for partition_id in topic_meta.partitions.keys():
+                    topic_partitions.append((topic, partition_id))
+        finally:
+            discovery_consumer.close()
+
+        bootstrap_servers = self._bootstrap_servers
+        start_offset = self._start_offset
+        end_offset = self._end_offset
+        timeout_ms = self._timeout_ms
+        target_max_block_size = self._target_max_block_size
+
+        # Validate that any partitions referenced in a per-partition dict
+        # actually exist on the broker. Check once per topic before the loop.
+        actual_partition_ids: Dict[str, Set[int]] = {}
+        for topic_name, partition_id in topic_partitions:
+            actual_partition_ids.setdefault(topic_name, set()).add(partition_id)
+
+        for param_name, offset in (
+            ("start_offset", start_offset),
+            ("end_offset", end_offset),
+        ):
+            if isinstance(offset, dict):
+                for topic, partition_map in offset.items():
+                    existing_partitions = actual_partition_ids.get(topic, set())
+                    for pid in partition_map:
+                        if pid not in existing_partitions:
+                            raise ValueError(
+                                f"{param_name} references partition {pid} in topic "
+                                f"{topic!r}, but that partition does not exist."
+                            )
+
+        tasks = []
+        for topic_name, partition_id in topic_partitions:
+
+            def create_kafka_read_fn(
+                topic_name: str = topic_name,
+                partition_id: int = partition_id,
+                bootstrap_servers: List[str] = bootstrap_servers,
+                start_offset: Optional[
+                    Union[int, datetime, Literal["earliest"]]
+                ] = start_offset,
+                end_offset: Optional[
+                    Union[int, datetime, Literal["latest"]]
+                ] = end_offset,
+                kafka_auth_config: Optional[KafkaAuthConfig] = self._kafka_auth_config,
+                user_consumer_config: Optional[Dict[str, Any]] = self._consumer_config,
+                timeout_ms: Optional[int] = timeout_ms,
+                target_max_block_size: int = target_max_block_size,
+            ):
+                """Create a Kafka read function with captured variables."""
+
+                def kafka_read_fn() -> Iterable[Block]:
+                    """Read function for a single Kafka partition using confluent-kafka."""
+                    from confluent_kafka import (
+                        Consumer,
+                        KafkaError,
+                        KafkaException,
+                        TopicPartition,
+                    )
+
+                    built_consumer_config = _build_consumer_config_for_read(
+                        bootstrap_servers, kafka_auth_config, user_consumer_config
+                    )
+
+                    consumer = Consumer(built_consumer_config)
+                    try:
+                        topic_partition = TopicPartition(topic_name, partition_id)
+                        resolved_start, resolved_end = _resolve_offsets(
+                            consumer, topic_partition, start_offset, end_offset
+                        )
+
+                        records = []
+
+                        output_buffer = BlockOutputBuffer(
+                            OutputBlockSizeOption.of(
+                                target_max_block_size=target_max_block_size
+                            )
+                        )
+
+                        if resolved_start < resolved_end:
+                            start_time = time.perf_counter()
+                            timeout_seconds = (
+                                timeout_ms / 1000.0 if timeout_ms is not None else None
+                            )
+
+                            tp_with_offset = TopicPartition(
+                                topic_name, partition_id, resolved_start
+                            )
+                            consumer.assign([tp_with_offset])
+
+                            next_offset = resolved_start
+
+                            partition_done = False
+                            while not partition_done:
+                                if next_offset >= resolved_end:
+                                    break
+
+                                if timeout_seconds is not None:
+                                    elapsed_time_s = time.perf_counter() - start_time
+                                    if elapsed_time_s >= timeout_seconds:
+                                        logger.warning(
+                                            f"Kafka read task timed out after {timeout_ms}ms while reading partition {partition_id} of topic {topic_name}; "
+                                            f"end_offset {resolved_end} was not reached. Returning {len(records)} messages collected in this read task so far."
+                                        )
+                                        break
+                                    remaining_timeout_s = (
+                                        timeout_seconds - elapsed_time_s
+                                    )
+                                    consume_timeout_s = min(
+                                        remaining_timeout_s, KAFKA_CONSUME_TIMEOUT_MAX_S
+                                    )
+                                else:
+                                    consume_timeout_s = KAFKA_CONSUME_TIMEOUT_MAX_S
+
+                                remaining = resolved_end - next_offset
+                                batch_size = min(
+                                    remaining,
+                                    KafkaDatasource.BATCH_SIZE_FOR_YIELD,
+                                )
+                                msgs = consumer.consume(
+                                    num_messages=batch_size,
+                                    timeout=consume_timeout_s,
+                                )
+
+                                if not msgs:
+                                    continue
+
+                                for msg in msgs:
+                                    if msg.error():
+                                        # In confluent-kafka, errors are delivered
+                                        # as messages. Only skip partition EOF
+                                        # events, raise others.
+                                        err = msg.error()
+                                        if err.code() == KafkaError._PARTITION_EOF:
+                                            continue
+                                        raise KafkaException(err)
+
+                                    # Stop once we reached the end offset
+                                    # (exclusive).
+                                    if msg.offset() >= resolved_end:
+                                        partition_done = True
+                                        break
+
+                                    ts_type, ts_ms = msg.timestamp()
+                                    headers_list = msg.headers() or []
+                                    headers_dict = dict(headers_list)
+                                    records.append(
+                                        {
+                                            "offset": msg.offset(),
+                                            "key": msg.key(),
+                                            "value": msg.value(),
+                                            "topic": msg.topic(),
+                                            "partition": msg.partition(),
+                                            "timestamp": ts_ms,
+                                            "timestamp_type": ts_type,
+                                            "headers": headers_dict,
+                                        }
+                                    )
+
+                                    next_offset = msg.offset() + 1
+
+                                    if (
+                                        len(records)
+                                        >= KafkaDatasource.BATCH_SIZE_FOR_YIELD
+                                    ):
+                                        table = pa.Table.from_pylist(records)
+                                        output_buffer.add_block(table)
+                                        yield from output_buffer.iter_ready_blocks()
+                                        records = []
+
+                        # Yield any remaining records
+                        if records:
+                            table = pa.Table.from_pylist(records)
+                            output_buffer.add_block(table)
+
+                        output_buffer.finalize()
+                        yield from output_buffer.iter_ready_blocks()
+
+                    finally:
+                        consumer.close()
+
+                return kafka_read_fn
+
+            # TODO: We could output the offset range for every partition after the read is done.
+            metadata = BlockMetadata(
+                num_rows=None,
+                size_bytes=None,
+                input_files=[f"kafka://{topic_name}/{partition_id}"],
+                exec_stats=None,
+            )
+
+            effective_start = (
+                start_offset.get(topic_name, {}).get(partition_id, "earliest")
+                if isinstance(start_offset, dict)
+                else start_offset
+            )
+            effective_end = (
+                end_offset.get(topic_name, {}).get(partition_id, "latest")
+                if isinstance(end_offset, dict)
+                else end_offset
+            )
+            kafka_read_fn = create_kafka_read_fn(
+                topic_name,
+                partition_id,
+                start_offset=effective_start,
+                end_offset=effective_end,
+            )
+
+            # Create read task
+            task = ReadTask(
+                read_fn=kafka_read_fn,
+                metadata=metadata,
+                schema=KAFKA_MSG_SCHEMA,
+                per_task_row_limit=per_task_row_limit,
+            )
+            tasks.append(task)
+
+        return tasks

python/ray/data/_internal/datasource/lance_datasink.py

@@ -0,0 +1,460 @@
+import pickle
+from itertools import chain
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import pyarrow as pa
+
+from ray._common.retry import call_with_retry
+from ray.data._internal.arrow_ops.transform_pyarrow import (
+    reorder_columns_by_schema,
+)
+from ray.data._internal.datasource.lance_utils import (
+    create_storage_options_provider,
+    get_or_create_namespace,
+)
+from ray.data._internal.savemode import SaveMode
+from ray.data._internal.util import _check_import, unify_schemas_with_validation
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.datasink import Datasink
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from lance import LanceDataset
+    from lance.fragment import FragmentMetadata
+
+
+_WRITE_LANCE_FRAGMENTS_DESCRIPTION = "write lance fragments"
+
+
+def _declare_table_with_fallback(
+    namespace, table_id: List[str]
+) -> Tuple[str, Optional[Dict[str, str]]]:
+    """Declare a table using declare_table, falling back to create_empty_table."""
+    try:
+        from lance_namespace import DeclareTableRequest
+
+        declare_request = DeclareTableRequest(id=table_id, location=None)
+        declare_response = namespace.declare_table(declare_request)
+        return declare_response.location, declare_response.storage_options
+    except (AttributeError, NotImplementedError):
+        # Fallback for older namespace implementations without declare_table.
+        from lance_namespace import CreateEmptyTableRequest
+
+        create_request = CreateEmptyTableRequest(id=table_id)
+        create_response = namespace.create_empty_table(create_request)
+        return create_response.location, create_response.storage_options
+
+
+def _make_stream_factory(
+    stream: Iterable[Block], replayable: bool
+) -> Tuple[Optional[Callable[[], Iterator[Block]]], Optional[Block]]:
+    """Return a reusable stream factory and the first block, or (None, None)."""
+    if replayable:
+        blocks = list(stream)
+        if not blocks:
+            return None, None
+
+        def stream_factory() -> Iterator[Block]:
+            return iter(blocks)
+
+        return stream_factory, blocks[0]
+
+    stream_iter = iter(stream)
+    first = next(stream_iter, None)
+    if first is None:
+        return None, None
+
+    def stream_factory() -> Iterator[Block]:
+        return chain([first], stream_iter)
+
+    return stream_factory, first
+
+
+def _write_fragment(
+    stream: Iterable[Block],
+    uri: str,
+    *,
+    schema: Optional["pa.Schema"] = None,
+    max_rows_per_file: int = 64 * 1024 * 1024,
+    max_bytes_per_file: Optional[int] = None,
+    max_rows_per_group: int = 1024,  # Only useful for v1 writer.
+    data_storage_version: Optional[str] = None,
+    storage_options: Optional[Dict[str, Any]] = None,
+    namespace_impl: Optional[str] = None,
+    namespace_properties: Optional[Dict[str, str]] = None,
+    table_id: Optional[List[str]] = None,
+    retry_params: Optional[Dict[str, Any]] = None,
+) -> List[Tuple["FragmentMetadata", "pa.Schema"]]:
+    import pandas as pd
+    from lance.fragment import DEFAULT_MAX_BYTES_PER_FILE, write_fragments
+
+    if retry_params is None:
+        retry_params = {
+            "description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
+            "match": [],
+            "max_attempts": 1,
+            "max_backoff_s": 0,
+        }
+
+    max_attempts = retry_params.get("max_attempts", 1)
+    stream_factory, first = _make_stream_factory(stream, replayable=max_attempts > 1)
+    if stream_factory is None or first is None:
+        return []
+
+    if schema is None:
+        if isinstance(first, pd.DataFrame):
+            schema = pa.Schema.from_pandas(first).remove_metadata()
+        else:
+            schema = first.schema
+        if len(schema.names) == 0:
+            # Empty table.
+            schema = None
+
+    def record_batch_converter(block_stream):
+        for block in block_stream:
+            tbl = BlockAccessor.for_block(block).to_arrow()
+            # `RecordBatchReader.from_batches(schema, ...)` is positional.
+            if schema is not None:
+                tbl = reorder_columns_by_schema(tbl, schema)
+            yield from tbl.to_batches()
+
+    max_bytes_per_file = (
+        DEFAULT_MAX_BYTES_PER_FILE if max_bytes_per_file is None else max_bytes_per_file
+    )
+
+    storage_options_provider = create_storage_options_provider(
+        namespace_impl,
+        namespace_properties,
+        table_id,
+    )
+
+    def _write_once():
+        reader = pa.RecordBatchReader.from_batches(
+            schema, record_batch_converter(stream_factory())
+        )
+        return write_fragments(
+            reader,
+            uri,
+            schema=schema,
+            max_rows_per_file=max_rows_per_file,
+            max_rows_per_group=max_rows_per_group,
+            max_bytes_per_file=max_bytes_per_file,
+            data_storage_version=data_storage_version,
+            storage_options=storage_options,
+            storage_options_provider=storage_options_provider,
+        )
+
+    fragments = call_with_retry(_write_once, **retry_params)
+    return [(fragment, schema) for fragment in fragments]
+
+
+class _BaseLanceDatasink(Datasink):
+    """Base class for Lance Datasink."""
+
+    def __init__(
+        self,
+        uri: Optional[str] = None,
+        schema: Optional[pa.Schema] = None,
+        mode: SaveMode = SaveMode.CREATE,
+        storage_options: Optional[Dict[str, Any]] = None,
+        table_id: Optional[List[str]] = None,
+        namespace_impl: Optional[str] = None,
+        namespace_properties: Optional[Dict[str, str]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        super().__init__(*args, **kwargs)
+
+        if mode not in {SaveMode.CREATE, SaveMode.APPEND, SaveMode.OVERWRITE}:
+            raise ValueError(
+                f"Unsupported Lance write mode: {mode!r}. "
+                "Supported modes are SaveMode.CREATE, SaveMode.APPEND, and SaveMode.OVERWRITE."
+            )
+
+        merged_storage_options: Dict[str, Any] = {}
+        if storage_options:
+            merged_storage_options.update(storage_options)
+
+        self._namespace_impl = namespace_impl
+        self._namespace_properties = namespace_properties
+
+        namespace = get_or_create_namespace(namespace_impl, namespace_properties)
+
+        if namespace is not None and table_id is not None:
+            if uri is not None:
+                import warnings
+
+                warnings.warn(
+                    "The 'uri' argument is ignored when namespace parameters are "
+                    "provided. The resolved namespace location will be used instead.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+            self.table_id = table_id
+            if mode != SaveMode.CREATE:
+                raise ValueError(
+                    "Namespace writes currently only support mode='create'. "
+                    "Use mode='create' for now."
+                )
+
+            uri, ns_storage_options = _declare_table_with_fallback(namespace, table_id)
+            self.uri = uri
+            if ns_storage_options:
+                merged_storage_options.update(ns_storage_options)
+            self._has_namespace_storage_options = True
+        else:
+            self.table_id = None
+            if uri is None:
+                raise ValueError(
+                    "Must provide either 'uri' or ('namespace_impl' and 'table_id')."
+                )
+            self.uri = uri
+            self._has_namespace_storage_options = False
+
+        self.schema = schema
+        self.mode = mode
+
+        self.read_version: Optional[int] = None
+        self.storage_options = merged_storage_options
+
+    @property
+    def storage_options_provider(self):
+        """Lazily create storage options provider using namespace_impl/properties."""
+        if not self._has_namespace_storage_options:
+            return None
+        return create_storage_options_provider(
+            self._namespace_impl,
+            self._namespace_properties,
+            self.table_id,
+        )
+
+    @property
+    def supports_distributed_writes(self) -> bool:
+        return True
+
+    def _open_dataset(self) -> "LanceDataset":
+        """Open the Lance dataset at ``self.uri``.
+
+        Raises whatever Lance raises if the dataset can't be opened (missing,
+        bad ``storage_options``, etc.). Opening natively honors
+        ``storage_options``/``storage_options_provider``.
+        """
+        import lance
+
+        return lance.LanceDataset(
+            self.uri,
+            storage_options=self.storage_options,
+            storage_options_provider=self.storage_options_provider,
+        )
+
+    def _dataset_exists(self) -> bool:
+        """Whether a Lance dataset already exists at ``self.uri``.
+
+        A *successful open* is the authoritative existence signal, and it honors
+        ``storage_options`` for every backend. We intentionally do not try to
+        classify failures (e.g. by matching Lance error strings, which drift
+        across versions): if the dataset can't be opened we report it as
+        not-existing and let the subsequent write surface any real error (such
+        as invalid ``storage_options``) with Lance's own message.
+        """
+        try:
+            self._open_dataset()
+            return True
+        except Exception:
+            return False
+
+    def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
+        _check_import(self, module="lance", package="pylance")
+
+        if self.mode == SaveMode.CREATE:
+            # CREATE must not clobber an existing dataset. Users who want to
+            # replace existing data should use SaveMode.OVERWRITE. Namespace
+            # writes manage table creation separately (the table location is
+            # declared/created up front), so skip the check in that case.
+            if self.table_id is None and self._dataset_exists():
+                raise ValueError(
+                    f"Dataset at {self.uri} already exists. "
+                    "Use mode=SaveMode.OVERWRITE to replace it, or "
+                    "mode=SaveMode.APPEND to add to it."
+                )
+        elif self.mode == SaveMode.APPEND:
+            # APPEND needs the existing dataset's version/schema. Let Lance
+            # raise its own error (e.g. dataset not found) if it can't open.
+            ds = self._open_dataset()
+            self.read_version = ds.version
+            if self.schema is None:
+                self.schema = ds.schema
+
+    def on_write_complete(
+        self,
+        write_results: List[List[Tuple[str, str]]],
+    ):
+        import warnings
+
+        import lance
+
+        if not write_results:
+            warnings.warn(
+                "write_results is empty.",
+                DeprecationWarning,
+            )
+            return
+        if hasattr(write_results, "write_returns"):
+            write_results = write_results.write_returns
+
+        if len(write_results) == 0:
+            warnings.warn(
+                "write results is empty. please check ray version or internal error",
+                DeprecationWarning,
+            )
+            return
+
+        fragments = []
+        schemas = []
+        for batch in write_results:
+            for fragment_str, schema_str in batch:
+                fragment = pickle.loads(fragment_str)
+                fragments.append(fragment)
+                schema = pickle.loads(schema_str)
+                if schema is not None:
+                    schemas.append(schema)
+        # Skip commit when there are no fragments/schemas to commit.
+        if not schemas:
+            return
+
+        unified_schema = unify_schemas_with_validation(schemas)
+        if unified_schema is None:
+            return
+
+        if self.mode in {SaveMode.CREATE, SaveMode.OVERWRITE}:
+            op = lance.LanceOperation.Overwrite(unified_schema, fragments)
+        elif self.mode == SaveMode.APPEND:
+            op = lance.LanceOperation.Append(fragments)
+        lance.LanceDataset.commit(
+            self.uri,
+            op,
+            read_version=self.read_version,
+            storage_options=self.storage_options,
+            storage_options_provider=self.storage_options_provider,
+        )
+
+
+class LanceDatasink(_BaseLanceDatasink):
+    """Lance Ray Datasink.
+
+    Write a Ray dataset to lance.
+
+    If we expect to write larger-than-memory files,
+    we can use `LanceFragmentWriter` and `LanceCommitter`.
+
+    Args:
+        uri: The base URI of the dataset.
+        schema: The schema of the dataset.
+        mode: The write mode. Default is SaveMode.CREATE. Choices are
+            SaveMode.CREATE, SaveMode.APPEND, SaveMode.OVERWRITE. Namespace-backed
+            writes currently support only SaveMode.CREATE.
+        min_rows_per_file: The minimum number of rows per file. Default is 1024 * 1024.
+        max_rows_per_file: The maximum number of rows per file. Default is 64 * 1024 * 1024.
+        data_storage_version: The version of the data storage format to use. Newer versions are more
+            efficient but require newer versions of lance to read. The default is "legacy",
+            which will use the legacy v1 version. See the user guide for more details.
+        storage_options: The storage options for the writer. Default is None.
+        table_id: The table identifier as a list of strings, used with namespace params.
+        namespace_impl: The namespace implementation type (e.g., "rest", "dir").
+            Used together with namespace_properties and table_id for credentials vending.
+        namespace_properties: Properties for connecting to the namespace.
+            Used together with namespace_impl and table_id for credentials vending.
+            When namespace params are provided, only SaveMode.CREATE is currently
+            supported.
+        *args: Additional positional arguments forwarded to the base class.
+        **kwargs: Additional keyword arguments forwarded to the base class.
+    """
+
+    NAME = "Lance"
+
+    def __init__(
+        self,
+        uri: Optional[str] = None,
+        schema: Optional[pa.Schema] = None,
+        mode: SaveMode = SaveMode.CREATE,
+        min_rows_per_file: int = 1024 * 1024,
+        max_rows_per_file: int = 64 * 1024 * 1024,
+        data_storage_version: Optional[str] = None,
+        storage_options: Optional[Dict[str, Any]] = None,
+        table_id: Optional[List[str]] = None,
+        namespace_impl: Optional[str] = None,
+        namespace_properties: Optional[Dict[str, str]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            uri,
+            schema=schema,
+            mode=mode,
+            storage_options=storage_options,
+            table_id=table_id,
+            namespace_impl=namespace_impl,
+            namespace_properties=namespace_properties,
+            *args,
+            **kwargs,
+        )
+
+        self.min_rows_per_file = min_rows_per_file
+        self.max_rows_per_file = max_rows_per_file
+        self.data_storage_version = data_storage_version
+        # if mode is append, read_version is read from existing dataset.
+        self.read_version: Optional[int] = None
+
+        data_context = DataContext.get_current()
+        lance_config = data_context.lance_config
+        match = []
+        match.extend(lance_config.write_fragments_errors_to_retry)
+        match.extend(data_context.retried_io_errors)
+        self._retry_params = {
+            "description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
+            "match": match,
+            "max_attempts": lance_config.write_fragments_max_attempts,
+            "max_backoff_s": lance_config.write_fragments_retry_max_backoff_s,
+        }
+
+    @property
+    def min_rows_per_write(self) -> int:
+        return self.min_rows_per_file
+
+    def get_name(self) -> str:
+        return self.NAME
+
+    def write(
+        self,
+        blocks: Iterable[Union[pa.Table, "pd.DataFrame"]],
+        _ctx,
+    ):
+        fragments_and_schema = _write_fragment(
+            blocks,
+            self.uri,
+            schema=self.schema,
+            max_rows_per_file=self.max_rows_per_file,
+            data_storage_version=self.data_storage_version,
+            storage_options=self.storage_options,
+            namespace_impl=self._namespace_impl,
+            namespace_properties=self._namespace_properties,
+            table_id=self.table_id,
+            retry_params=self._retry_params,
+        )
+        return [
+            (pickle.dumps(fragment), pickle.dumps(schema))
+            for fragment, schema in fragments_and_schema
+        ]

python/ray/data/_internal/datasource/lance_datasource.py

@@ -0,0 +1,160 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+
+import numpy as np
+
+from ray._common.retry import call_with_retry
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+logger = logging.getLogger(__name__)
+
+
+class LanceDatasource(Datasource):
+    """Lance datasource, for reading Lance dataset."""
+
+    def __init__(
+        self,
+        uri: str,
+        version: Optional[Union[int, str]] = None,
+        columns: Optional[List[str]] = None,
+        filter: Optional[str] = None,
+        storage_options: Optional[Dict[str, str]] = None,
+        scanner_options: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__()
+        _check_import(self, module="lance", package="pylance")
+
+        import lance
+
+        self._projection_map = None
+        self.uri = uri
+        self.scanner_options = scanner_options or {}
+        if columns is not None:
+            self.scanner_options["columns"] = columns
+        if filter is not None:
+            self.scanner_options["filter"] = filter
+        self.storage_options = storage_options
+        self.lance_ds = lance.dataset(
+            uri=uri, version=version, storage_options=storage_options
+        )
+
+        data_context = DataContext.get_current()
+        lance_config = data_context.lance_config
+        match = []
+        match.extend(lance_config.read_fragments_errors_to_retry)
+        match.extend(data_context.retried_io_errors)
+        self._retry_params = {
+            "description": "read lance fragments",
+            "match": match,
+            "max_attempts": lance_config.read_fragments_max_attempts,
+            "max_backoff_s": lance_config.read_fragments_retry_max_backoff_s,
+        }
+
+    def supports_predicate_pushdown(self) -> bool:
+        return True
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        read_tasks = []
+        ds_fragments = self.scanner_options.get("fragments")
+        if ds_fragments is None:
+            ds_fragments = self.lance_ds.get_fragments()
+
+        # Lance scanner's filter attr accepts only a string (SQL).
+        # See: https://github.com/lance-format/lance/blob/aac74b441cdb6df7d78700dbba33c521e6379ca5/python/python/lance/lance/__init__.pyi#L230
+        filter_expr = (
+            str(self._predicate_expr.to_pyarrow())
+            if self._predicate_expr is not None
+            else None
+        )
+        filter_from_arg = self.scanner_options.get("filter")
+        if filter_from_arg is not None:
+            filter_expr = (
+                filter_from_arg
+                if filter_expr is None
+                else f"({filter_expr}) AND ({filter_from_arg})"
+            )
+
+        for fragments in np.array_split(ds_fragments, parallelism):
+            if len(fragments) <= 0:
+                continue
+
+            fragment_ids = [f.metadata.id for f in fragments]
+            num_rows = sum(f.count_rows() for f in fragments)
+            input_files = [
+                data_file.path() for f in fragments for data_file in f.data_files()
+            ]
+
+            # TODO(chengsu): Take column projection into consideration for schema.
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=None,
+                input_files=input_files,
+                exec_stats=None,
+            )
+            # Use a copy per task to avoid mutation races when tasks run in parallel
+            task_scanner_options = dict(self.scanner_options)
+            if filter_expr is not None:
+                task_scanner_options["filter"] = filter_expr
+            lance_ds = self.lance_ds
+            retry_params = self._retry_params
+
+            read_task = ReadTask(
+                lambda f=fragment_ids, opts=task_scanner_options: _read_fragments_with_retry(
+                    f,
+                    lance_ds,
+                    opts,
+                    retry_params,
+                ),
+                metadata,
+                schema=fragments[0].schema,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+        return read_tasks
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # TODO(chengsu): Add memory size estimation to improve auto-tune of parallelism.
+        return None
+
+
+def _read_fragments_with_retry(
+    fragment_ids,
+    lance_ds,
+    scanner_options,
+    retry_params,
+) -> Iterator["pyarrow.Table"]:
+    return call_with_retry(
+        lambda: _read_fragments(fragment_ids, lance_ds, scanner_options),
+        **retry_params,
+    )
+
+
+def _read_fragments(
+    fragment_ids,
+    lance_ds,
+    scanner_options,
+) -> Iterator["pyarrow.Table"]:
+    """Read Lance fragments in batches.
+
+    NOTE: Use fragment ids, instead of fragments as parameter, because pickling
+    LanceFragment is expensive.
+    """
+    import pyarrow
+
+    fragments = [lance_ds.get_fragment(id) for id in fragment_ids]
+    scanner_options["fragments"] = fragments
+    scanner = lance_ds.scanner(**scanner_options)
+    for batch in scanner.to_reader():
+        yield pyarrow.Table.from_batches([batch])

python/ray/data/_internal/datasource/lance_utils.py

@@ -0,0 +1,58 @@
+import os
+from functools import lru_cache
+from typing import Dict, List, Optional, Tuple
+
+_NAMESPACE_CACHE_SIZE = int(os.environ.get("LANCE_RAY_NAMESPACE_CACHE_SIZE", "16"))
+
+
+def _has_namespace_params(
+    namespace_impl: Optional[str],
+    table_id: Optional[List[str]],
+) -> bool:
+    return namespace_impl is not None and table_id is not None
+
+
+@lru_cache(maxsize=_NAMESPACE_CACHE_SIZE)
+def _get_cached_namespace(
+    namespace_impl: str,
+    namespace_properties_tuple: Optional[Tuple[Tuple[str, str], ...]],
+):
+    import lance_namespace as ln
+
+    namespace_properties = (
+        dict(namespace_properties_tuple) if namespace_properties_tuple else {}
+    )
+    return ln.connect(namespace_impl, namespace_properties)
+
+
+def get_or_create_namespace(
+    namespace_impl: Optional[str],
+    namespace_properties: Optional[Dict[str, str]],
+):
+    """Get or create a cached namespace client for the worker."""
+    if namespace_impl is None:
+        return None
+
+    namespace_properties_tuple = (
+        tuple(sorted(namespace_properties.items())) if namespace_properties else None
+    )
+    return _get_cached_namespace(namespace_impl, namespace_properties_tuple)
+
+
+def create_storage_options_provider(
+    namespace_impl: Optional[str],
+    namespace_properties: Optional[Dict[str, str]],
+    table_id: Optional[List[str]],
+):
+    """Create a LanceNamespaceStorageOptionsProvider if namespace params exist."""
+    if not _has_namespace_params(namespace_impl, table_id):
+        return None
+
+    assert table_id is not None
+    namespace = get_or_create_namespace(namespace_impl, namespace_properties)
+    if namespace is None:
+        return None
+
+    from lance import LanceNamespaceStorageOptionsProvider
+
+    return LanceNamespaceStorageOptionsProvider(namespace=namespace, table_id=table_id)

python/ray/data/_internal/datasource/mcap_datasource.py

@@ -0,0 +1,258 @@
+"""MCAP (Message Capture) datasource for Ray Data.
+
+MCAP is a standardized format for storing timestamped messages from robotics and
+autonomous systems, commonly used for sensor data, control commands, and other
+time-series data.
+"""
+
+import json
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Set, Union
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import pyarrow
+    from mcap.reader import Channel, Message, Schema
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TimeRange:
+    """Time range for filtering MCAP messages.
+
+    Attributes:
+        start_time: Start time in nanoseconds (inclusive).
+        end_time: End time in nanoseconds (exclusive).
+    """
+
+    start_time: int
+    end_time: int
+
+    def __post_init__(self):
+        """Validate time range after initialization."""
+        if self.start_time >= self.end_time:
+            raise ValueError(
+                f"start_time ({self.start_time}) must be less than "
+                f"end_time ({self.end_time})"
+            )
+        if self.start_time < 0 or self.end_time < 0:
+            raise ValueError(
+                f"time values must be non-negative, got start_time={self.start_time}, "
+                f"end_time={self.end_time}"
+            )
+
+
+@DeveloperAPI
+class MCAPDatasource(FileBasedDatasource):
+    """MCAP (Message Capture) datasource for Ray Data.
+
+    This datasource provides reading of MCAP files with predicate pushdown
+    optimization for filtering by topics, time ranges, and message types.
+
+    MCAP is a standardized format for storing timestamped messages from robotics and
+    autonomous systems, commonly used for sensor data, control commands, and other
+    time-series data.
+
+    Examples:
+        Basic usage:
+
+        >>> import ray  # doctest: +SKIP
+        >>> ds = ray.data.read_mcap("/path/to/data.mcap")  # doctest: +SKIP
+
+        With topic filtering and time range:
+
+        >>> from ray.data.datasource import TimeRange  # doctest: +SKIP
+        >>> ds = ray.data.read_mcap(  # doctest: +SKIP
+        ...     "/path/to/data.mcap",
+        ...     topics={"/camera/image_raw", "/lidar/points"},
+        ...     time_range=TimeRange(start_time=1000000000, end_time=2000000000)
+        ... )  # doctest: +SKIP
+
+        With multiple files and metadata:
+
+        >>> ds = ray.data.read_mcap(  # doctest: +SKIP
+        ...     ["file1.mcap", "file2.mcap"],
+        ...     topics={"/camera/image_raw", "/lidar/points"},
+        ...     message_types={"sensor_msgs/Image", "sensor_msgs/PointCloud2"},
+        ...     include_metadata=True
+        ... )  # doctest: +SKIP
+    """
+
+    _FILE_EXTENSIONS = ["mcap"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        topics: Optional[Union[List[str], Set[str]]] = None,
+        time_range: Optional[TimeRange] = None,
+        message_types: Optional[Union[List[str], Set[str]]] = None,
+        include_metadata: bool = True,
+        **file_based_datasource_kwargs,
+    ):
+        """Initialize MCAP datasource.
+
+        Args:
+            paths: Path or list of paths to MCAP files.
+            topics: Optional list/set of topic names to include. If specified,
+                only messages from these topics will be read.
+            time_range: Optional TimeRange for filtering messages by timestamp.
+                TimeRange contains start_time and end_time in nanoseconds, where
+                both values must be non-negative and start_time < end_time.
+            message_types: Optional list/set of message type names (schema names)
+                to include. Only messages with matching schema names will be read.
+            include_metadata: Whether to include MCAP metadata fields in the output.
+                Defaults to True. When True, includes schema, channel, and message
+                metadata.
+            **file_based_datasource_kwargs: Additional arguments for FileBasedDatasource.
+        """
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="mcap", package="mcap")
+
+        # Convert to sets for faster lookup
+        self._topics = set(topics) if topics else None
+        self._message_types = set(message_types) if message_types else None
+        self._time_range = time_range
+        self._include_metadata = include_metadata
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        """Read MCAP file and yield blocks of message data.
+
+        This method implements efficient MCAP reading with predicate pushdown.
+        It uses MCAP's built-in filtering capabilities for optimal performance
+        and applies additional filters when needed.
+
+        Args:
+            f: File-like object to read from. Must be seekable for MCAP reading.
+            path: Path to the MCAP file being processed.
+
+        Yields:
+            Block: Blocks of MCAP message data as pyarrow Tables.
+
+        Raises:
+            ValueError: If the MCAP file cannot be read or has invalid format.
+        """
+        from mcap.reader import make_reader
+
+        reader = make_reader(f)
+        # Note: MCAP summaries are optional and iter_messages works without them
+        # We don't need to validate the summary since it's not required
+
+        # Use MCAP's built-in filtering for topics and time range
+        messages = reader.iter_messages(
+            topics=list(self._topics) if self._topics else None,
+            start_time=self._time_range.start_time if self._time_range else None,
+            end_time=self._time_range.end_time if self._time_range else None,
+            log_time_order=True,
+            reverse=False,
+        )
+
+        builder = DelegatingBlockBuilder()
+
+        for schema, channel, message in messages:
+            # Apply filters that couldn't be pushed down to MCAP level
+            if not self._should_include_message(schema, channel, message):
+                continue
+
+            # Convert message to dictionary format
+            message_data = self._message_to_dict(schema, channel, message, path)
+            builder.add(message_data)
+
+        # Yield the block if we have any messages
+        if builder.num_rows() > 0:
+            yield builder.build()
+
+    def _should_include_message(
+        self, schema: "Schema", channel: "Channel", message: "Message"
+    ) -> bool:
+        """Check if a message should be included based on filters.
+
+        This method applies Python-level filtering that cannot be pushed down
+        to the MCAP library level. Topic filters are already handled by the
+        MCAP reader, so only message_types filtering is needed here.
+
+        Args:
+            schema: MCAP schema object containing message type information.
+            channel: MCAP channel object containing topic and metadata.
+            message: MCAP message object containing the actual data.
+
+        Returns:
+            True if the message should be included, False otherwise.
+        """
+        # Message type filter (cannot be pushed down to MCAP reader)
+        if self._message_types and schema and schema.name not in self._message_types:
+            return False
+
+        return True
+
+    def _message_to_dict(
+        self, schema: "Schema", channel: "Channel", message: "Message", path: str
+    ) -> Dict[str, Any]:
+        """Convert MCAP message to dictionary format.
+
+        This method converts MCAP message objects into a standardized dictionary
+        format suitable for Ray Data processing.
+
+        Args:
+            schema: MCAP schema object containing message type and encoding info.
+            channel: MCAP channel object containing topic and channel metadata.
+            message: MCAP message object containing the actual message data.
+            path: Path to the source file (for include_paths functionality).
+
+        Returns:
+            Dictionary containing message data in Ray Data format.
+        """
+        # Decode message data based on encoding
+        decoded_data = message.data
+        if channel.message_encoding == "json" and isinstance(message.data, bytes):
+            try:
+                decoded_data = json.loads(message.data.decode("utf-8"))
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                # Keep raw bytes if decoding fails
+                decoded_data = message.data
+
+        # Core message data
+        message_data = {
+            "data": decoded_data,
+            "topic": channel.topic,
+            "log_time": message.log_time,
+            "publish_time": message.publish_time,
+            "sequence": message.sequence,
+        }
+
+        # Add metadata if requested
+        if self._include_metadata:
+            message_data.update(
+                {
+                    "channel_id": message.channel_id,
+                    "message_encoding": channel.message_encoding,
+                    "schema_name": schema.name if schema else None,
+                    "schema_encoding": schema.encoding if schema else None,
+                    "schema_data": schema.data if schema else None,
+                }
+            )
+
+        # Add file path if include_paths is enabled (from FileBasedDatasource)
+        if getattr(self, "include_paths", False):
+            message_data["path"] = path
+
+        return message_data
+
+    def get_name(self) -> str:
+        """Return a human-readable name for this datasource."""
+        return "MCAP"
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        """Whether this datasource supports distributed reads.
+
+        MCAP files can be read in parallel across multiple files.
+        """
+        return True

python/ray/data/_internal/datasource/mongo_datasink.py

@@ -0,0 +1,48 @@
+import logging
+from typing import Iterable
+
+from ray.data._internal.datasource.mongo_datasource import (
+    _validate_database_collection_exist,
+)
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDatasink(Datasink[None]):
+    def __init__(self, uri: str, database: str, collection: str) -> None:
+        _check_import(self, module="pymongo", package="pymongo")
+        _check_import(self, module="pymongoarrow", package="pymongoarrow")
+
+        self.uri = uri
+        self.database = database
+        self.collection = collection
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        import pymongo
+
+        _validate_database_collection_exist(
+            pymongo.MongoClient(self.uri), self.database, self.collection
+        )
+
+        def write_block(uri: str, database: str, collection: str, block: Block):
+            from pymongoarrow.api import write
+
+            block = BlockAccessor.for_block(block).to_arrow()
+            client = pymongo.MongoClient(uri)
+            write(client[database][collection], block)
+
+        builder = DelegatingBlockBuilder()
+        for block in blocks:
+            builder.add_block(block)
+        block = builder.build()
+
+        write_block(self.uri, self.database, self.collection, block)

python/ray/data/_internal/datasource/mongo_datasource.py

@@ -0,0 +1,147 @@
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import pymongoarrow.api
+
+    from ray.data.context import DataContext
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDatasource(Datasource):
+    """Datasource for reading from and writing to MongoDB."""
+
+    def __init__(
+        self,
+        uri: str,
+        database: str,
+        collection: str,
+        pipeline: Optional[List[Dict]] = None,
+        schema: Optional["pymongoarrow.api.Schema"] = None,
+        **mongo_args,
+    ):
+        self._uri = uri
+        self._database = database
+        self._collection = collection
+        self._pipeline = pipeline
+        self._schema = schema
+        self._mongo_args = mongo_args
+        # If pipeline is unspecified, read the entire collection.
+        if not pipeline:
+            self._pipeline = [{"$match": {"_id": {"$exists": "true"}}}]
+        # Initialize Mongo client lazily later when creating read tasks.
+        self._client = None
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # TODO(jian): Add memory size estimation to improve auto-tune of parallelism.
+        return None
+
+    def _get_match_query(self, pipeline: List[Dict]) -> Dict:
+        if len(pipeline) == 0 or "$match" not in pipeline[0]:
+            return {}
+        return pipeline[0]["$match"]
+
+    def _get_or_create_client(self):
+        import pymongo
+
+        if self._client is None:
+            self._client = pymongo.MongoClient(self._uri)
+            _validate_database_collection_exist(
+                self._client, self._database, self._collection
+            )
+            self._avg_obj_size = self._client[self._database].command(
+                "collStats", self._collection
+            )["avgObjSize"]
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        from bson.objectid import ObjectId
+
+        self._get_or_create_client()
+        coll = self._client[self._database][self._collection]
+        match_query = self._get_match_query(self._pipeline)
+        partitions_ids = list(
+            coll.aggregate(
+                [
+                    {"$match": match_query},
+                    {"$bucketAuto": {"groupBy": "$_id", "buckets": parallelism}},
+                ],
+                allowDiskUse=True,
+            )
+        )
+
+        def make_block(
+            uri: str,
+            database: str,
+            collection: str,
+            pipeline: List[Dict],
+            min_id: ObjectId,
+            max_id: ObjectId,
+            right_closed: bool,
+            schema: "pymongoarrow.api.Schema",
+            kwargs: dict,
+        ) -> Block:
+            import pymongo
+            from pymongoarrow.api import aggregate_arrow_all
+
+            # A range query over the partition.
+            match = [
+                {
+                    "$match": {
+                        "_id": {
+                            "$gte": min_id,
+                            "$lte" if right_closed else "$lt": max_id,
+                        }
+                    }
+                }
+            ]
+            client = pymongo.MongoClient(uri)
+            return aggregate_arrow_all(
+                client[database][collection], match + pipeline, schema=schema, **kwargs
+            )
+
+        read_tasks: List[ReadTask] = []
+
+        for i, partition in enumerate(partitions_ids):
+            metadata = BlockMetadata(
+                num_rows=partition["count"],
+                size_bytes=partition["count"] * self._avg_obj_size,
+                input_files=None,
+                exec_stats=None,
+            )
+            make_block_args = (
+                self._uri,
+                self._database,
+                self._collection,
+                self._pipeline,
+                partition["_id"]["min"],
+                partition["_id"]["max"],
+                i == len(partitions_ids) - 1,
+                self._schema,
+                self._mongo_args,
+            )
+            read_task = ReadTask(
+                lambda args=make_block_args: [make_block(*args)],
+                metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+
+def _validate_database_collection_exist(client, database: str, collection: str):
+    db_names = client.list_database_names()
+    if database not in db_names:
+        raise ValueError(f"The destination database {database} doesn't exist.")
+    collection_names = client[database].list_collection_names()
+    if collection not in collection_names:
+        raise ValueError(f"The destination collection {collection} doesn't exist.")

python/ray/data/_internal/datasource/numpy_datasink.py

@@ -0,0 +1,23 @@
+import numpy as np
+import pyarrow
+
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class NumpyDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        column: str,
+        *,
+        file_format: str = "npy",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        self.column = column
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        value = block.to_numpy(self.column)
+        np.save(file, value)

python/ray/data/_internal/datasource/numpy_datasource.py

@@ -0,0 +1,41 @@
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+
+import numpy as np
+
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class NumpyDatasource(FileBasedDatasource):
+    """Numpy datasource, for reading and writing Numpy files."""
+
+    _COLUMN_NAME = "data"
+    _FILE_EXTENSIONS = ["npy"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        numpy_load_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        if numpy_load_args is None:
+            numpy_load_args = {}
+
+        self.numpy_load_args = numpy_load_args
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        # TODO(ekl) Ideally numpy can read directly from the file, but it
+        # seems like it requires the file to be seekable.
+        buf = BytesIO()
+        data = f.readall()
+        buf.write(data)
+        buf.seek(0)
+        yield BlockAccessor.batch_to_block(
+            {"data": np.load(buf, allow_pickle=True, **self.numpy_load_args)}
+        )

python/ray/data/_internal/datasource/parquet_datasink.py

@@ -0,0 +1,390 @@
+import logging
+from collections import defaultdict
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional
+
+from ray._common.retry import call_with_retry
+from ray._private.utils import INT32_MAX
+from ray.data._internal.arrow_ops.transform_pyarrow import (
+    reorder_columns_by_schema,
+)
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.planner.plan_write_op import WRITE_UUID_KWARG_NAME
+from ray.data._internal.savemode import SaveMode
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.file_based_datasource import _resolve_kwargs
+from ray.data.datasource.file_datasink import _FileDatasink
+from ray.data.datasource.filename_provider import FilenameProvider
+
+if TYPE_CHECKING:
+    import pyarrow
+
+WRITE_FILE_MAX_ATTEMPTS = 10
+WRITE_FILE_RETRY_MAX_BACKOFF_SECONDS = 32
+FILE_FORMAT = "parquet"
+
+# These args are part of https://arrow.apache.org/docs/python/generated/pyarrow.fs.FileSystem.html#pyarrow.fs.FileSystem.open_output_stream
+# and are not supported by ParquetDatasink.
+UNSUPPORTED_OPEN_STREAM_ARGS = {"path", "buffer", "metadata"}
+
+# https://arrow.apache.org/docs/python/generated/pyarrow.dataset.write_dataset.html
+ARROW_DEFAULT_MAX_ROWS_PER_GROUP = 1024 * 1024
+
+DEFAULT_PARTITIONING_FLAVOR = "hive"
+
+logger = logging.getLogger(__name__)
+
+
+def choose_row_group_limits(
+    row_group_size: Optional[int],
+    min_rows_per_file: Optional[int],
+    max_rows_per_file: Optional[int],
+) -> tuple[Optional[int], Optional[int], Optional[int]]:
+    """Configure row-group limits for Pyarrow's ``write_dataset`` API.
+
+    Configures the ``min_rows_per_group``, ``max_rows_per_group``, and
+    ``max_rows_per_file`` parameters based on Ray Data's configuration.
+
+    Args:
+        row_group_size: The requested row-group size.
+        min_rows_per_file: The minimum number of rows per file.
+        max_rows_per_file: The maximum number of rows per file.
+
+    Returns:
+        A tuple of (min_rows_per_group, max_rows_per_group, max_rows_per_file).
+    """
+
+    if (
+        row_group_size is None
+        and min_rows_per_file is None
+        and max_rows_per_file is None
+    ):
+        return None, None, None
+
+    elif row_group_size is None:
+        # No explicit row group size provided. We are defaulting to
+        # either the caller's min_rows_per_file or max_rows_per_file limits
+        # or Arrow's defaults
+        min_rows_per_group, max_rows_per_group, max_rows_per_file = (
+            min_rows_per_file,
+            max_rows_per_file,
+            max_rows_per_file,
+        )
+
+        # If min_rows_per_group is provided and max_rows_per_group is not,
+        # and min_rows_per_group is greater than Arrow's default max_rows_per_group,
+        # we set max_rows_per_group to min_rows_per_group to avoid creating too many row groups.
+        if (
+            min_rows_per_group is not None
+            and max_rows_per_group is None
+            and min_rows_per_group > ARROW_DEFAULT_MAX_ROWS_PER_GROUP
+        ):
+            max_rows_per_group, max_rows_per_file = (
+                min_rows_per_group,
+                min_rows_per_group,
+            )
+
+        return min_rows_per_group, max_rows_per_group, max_rows_per_file
+
+    elif row_group_size is not None and (
+        min_rows_per_file is None or max_rows_per_file is None
+    ):
+        return row_group_size, row_group_size, max_rows_per_file
+
+    else:
+        # Clamp the requested `row_group_size` so that it is
+        # * no smaller than `min_rows_per_file` (`lower`)
+        # * no larger than `max_rows_per_file` (or Arrow's default cap) (`upper`)
+        # This keeps each row-group within the per-file limits while staying
+        # as close as possible to the requested size.
+        clamped_group_size = max(
+            min_rows_per_file, min(row_group_size, max_rows_per_file)
+        )
+        return clamped_group_size, clamped_group_size, max_rows_per_file
+
+
+def _widen_offset_overflowing_columns(
+    tables: List["pyarrow.Table"], schema: "pyarrow.Schema"
+) -> "pyarrow.Schema":
+    """Promote `string`/`binary` columns to 64-bit-offset variants when needed.
+
+    Arrow addresses the data of `string` and `binary` columns with int32
+    offsets, so a single contiguous array can hold at most `INT32_MAX` bytes.
+    When the writer coalesces blocks into a large row group (e.g. because
+    `min_rows_per_file` set `min_rows_per_group`), pyarrow must materialize each
+    column of that row group as one contiguous array. If a `string`/`binary`
+    column's combined size across the blocks in this write exceeds the int32
+    limit, `write_dataset` fails with an offset-overflow error that surfaces as
+    a column-length mismatch (the column is truncated to what fits).
+
+    Promoting such columns to `large_string`/`large_binary` (int64 offsets)
+    removes the ceiling. The promotion is invisible on disk -- parquet stores
+    both as `BYTE_ARRAY` -- so only the in-memory Arrow type changes.
+
+    Only top-level `string`/`binary` columns are promoted. Other int32-offset
+    types are not handled: `string`/`binary` nested inside `list`/`struct`/`map`
+    (whose inner offsets carry the same byte ceiling) and `list`/`map` columns
+    themselves (which overflow on cumulative child-element count rather than
+    bytes). These require recursive type rewriting and are far rarer in practice.
+
+    Args:
+        tables: The blocks about to be written together in one task.
+        schema: The unified output schema for those blocks.
+
+    Returns:
+        ``schema`` unchanged when no column overflows, otherwise a copy with the
+        overflowing variable-width columns promoted to their ``large_*`` type.
+    """
+    import pyarrow as pa
+
+    candidate_types = {}
+    for field in schema:
+        if pa.types.is_string(field.type):
+            candidate_types[field.name] = pa.large_string()
+        elif pa.types.is_binary(field.type):
+            candidate_types[field.name] = pa.large_binary()
+    if not candidate_types:
+        return schema
+
+    # `.nbytes` is O(1) buffer metadata, so summing across blocks is cheap.
+    overflowing: set = set()
+    combined_nbytes: Dict[str, int] = defaultdict(int)
+    for table in tables:
+        for name in candidate_types:
+            idx = table.schema.get_field_index(name)
+            if idx != -1:
+                combined_nbytes[name] += table.column(idx).nbytes
+                if combined_nbytes[name] > INT32_MAX:
+                    overflowing.add(name)
+
+    if not overflowing:
+        return schema
+
+    new_fields = [
+        field.with_type(candidate_types[field.name])
+        if field.name in overflowing
+        else field
+        for field in schema
+    ]
+    return pa.schema(new_fields, metadata=schema.metadata)
+
+
+class ParquetDatasink(_FileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        partition_cols: Optional[List[str]] = None,
+        arrow_parquet_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        arrow_parquet_args: Optional[Dict[str, Any]] = None,
+        min_rows_per_file: Optional[int] = None,
+        max_rows_per_file: Optional[int] = None,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        try_create_dir: bool = True,
+        open_stream_args: Optional[Dict[str, Any]] = None,
+        filename_provider: Optional[FilenameProvider] = None,
+        dataset_uuid: Optional[str] = None,
+        mode: SaveMode = SaveMode.APPEND,
+    ):
+        if arrow_parquet_args_fn is None:
+            arrow_parquet_args_fn = lambda: {}  # noqa: E731
+
+        if arrow_parquet_args is None:
+            arrow_parquet_args = {}
+
+        self.arrow_parquet_args_fn = arrow_parquet_args_fn
+        self.arrow_parquet_args = arrow_parquet_args
+        self.min_rows_per_file = min_rows_per_file
+        self.max_rows_per_file = max_rows_per_file
+        self.partition_cols = partition_cols
+
+        if self.min_rows_per_file is not None and self.max_rows_per_file is not None:
+            if self.min_rows_per_file > self.max_rows_per_file:
+                raise ValueError(
+                    "min_rows_per_file must be less than or equal to max_rows_per_file"
+                )
+
+        if open_stream_args is not None:
+            intersecting_keys = UNSUPPORTED_OPEN_STREAM_ARGS.intersection(
+                set(open_stream_args.keys())
+            )
+            if intersecting_keys:
+                logger.warning(
+                    "open_stream_args contains unsupported arguments: %s. These arguments "
+                    "are not supported by ParquetDatasink. They will be ignored.",
+                    intersecting_keys,
+                )
+
+            if "compression" in open_stream_args:
+                self.arrow_parquet_args["compression"] = open_stream_args["compression"]
+
+        if ("partitioning_flavor" in self.arrow_parquet_args) or (
+            self.arrow_parquet_args_fn is not None
+            and "partitioning_flavor" in self.arrow_parquet_args_fn()
+        ):
+            if self.partition_cols is None:
+                raise ValueError(
+                    "partition_cols must be provided when partitioning_flavor is set."
+                )
+
+        super().__init__(
+            path,
+            filesystem=filesystem,
+            try_create_dir=try_create_dir,
+            open_stream_args=open_stream_args,
+            filename_provider=filename_provider,
+            dataset_uuid=dataset_uuid,
+            file_format=FILE_FORMAT,
+            mode=mode,
+        )
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        import pyarrow as pa
+
+        blocks = list(blocks)
+
+        if all(BlockAccessor.for_block(block).num_rows() == 0 for block in blocks):
+            return
+
+        blocks = [
+            block for block in blocks if BlockAccessor.for_block(block).num_rows() > 0
+        ]
+
+        filename = self.filename_provider.get_filename_for_task(
+            ctx.kwargs[WRITE_UUID_KWARG_NAME], ctx.task_idx
+        )
+        write_kwargs = _resolve_kwargs(
+            self.arrow_parquet_args_fn, **self.arrow_parquet_args
+        )
+        user_schema = write_kwargs.pop("schema", None)
+
+        # For partitioning_flavor, if it's not provided, the default is "hive"
+        # Otherwise, it follows pyarrow's behavior: None for directory,
+        # "hive" for hive, and "filename" for FilenamePartitioning.
+        partitioning_flavor = write_kwargs.pop(
+            "partitioning_flavor", DEFAULT_PARTITIONING_FLAVOR
+        )
+
+        def write_blocks_to_path():
+            tables = [BlockAccessor.for_block(block).to_arrow() for block in blocks]
+            if user_schema is None:
+                output_schema = pa.unify_schemas([table.schema for table in tables])
+                # Coalescing many blocks into one row group can push a
+                # `string`/`binary` column past Arrow's 2 GiB int32-offset
+                # limit; promote such columns to their `large_*` variant so the
+                # contiguous row-group array can address all of its bytes.
+                output_schema = _widen_offset_overflowing_columns(tables, output_schema)
+            else:
+                output_schema = user_schema
+
+            self._write_parquet_files(
+                tables,
+                filename,
+                output_schema,
+                ctx.kwargs[WRITE_UUID_KWARG_NAME],
+                write_kwargs,
+                partitioning_flavor,
+            )
+
+        logger.debug(f"Writing {filename} file to {self.path}.")
+
+        call_with_retry(
+            write_blocks_to_path,
+            description=f"write '{filename}' to '{self.path}'",
+            match=self._data_context.retried_io_errors,
+            max_attempts=WRITE_FILE_MAX_ATTEMPTS,
+            max_backoff_s=WRITE_FILE_RETRY_MAX_BACKOFF_SECONDS,
+        )
+
+    def _get_basename_template(self, filename: str, write_uuid: str) -> str:
+        # Check if write_uuid is present in filename, add if missing
+        if write_uuid not in filename and self.mode == SaveMode.APPEND:
+            raise ValueError(
+                f"Write UUID '{write_uuid}' missing from filename template '{filename}'. This could result in files being overwritten."
+                f"Modify your FileNameProvider implementation to include the `write_uuid` into the filename template or change your write mode to SaveMode.OVERWRITE. "
+            )
+        # Check if filename is already templatized
+        if "{i}" in filename:
+            # Filename is already templatized, but may need file extension
+            if FILE_FORMAT not in filename:
+                # Add file extension to templatized filename
+                basename_template = f"{filename}.{FILE_FORMAT}"
+            else:
+                # Already has extension, use as-is
+                basename_template = filename
+        elif FILE_FORMAT not in filename:
+            # No extension and not templatized, add extension and template
+            basename_template = f"{filename}-{{i}}.{FILE_FORMAT}"
+        else:
+            # TODO(@goutamvenkat-anyscale): Add a warning if you pass in a custom
+            # filename provider and it isn't templatized.
+            # Use pathlib.Path to properly handle filenames with dots
+            filename_path = Path(filename)
+            stem = filename_path.stem  # filename without extension
+            assert "." not in stem, "Filename should not contain a dot"
+            suffix = filename_path.suffix  # extension including the dot
+            basename_template = f"{stem}-{{i}}{suffix}"
+        return basename_template
+
+    def _write_parquet_files(
+        self,
+        tables: List["pyarrow.Table"],
+        filename: str,
+        output_schema: "pyarrow.Schema",
+        write_uuid: str,
+        write_kwargs: Dict[str, Any],
+        partitioning_flavor: Optional[str],
+    ) -> None:
+        import pyarrow.dataset as ds
+
+        # Make every incoming batch conform to the final schema *before* writing.
+        # `pa.unify_schemas` above fixed column order from the first block.
+        for idx, table in enumerate(tables):
+            if output_schema and not table.schema.equals(output_schema):
+                table = reorder_columns_by_schema(table, output_schema)
+                table = table.cast(output_schema)
+            tables[idx] = table
+
+        row_group_size = write_kwargs.pop("row_group_size", None)
+
+        # We set this to "overwrite_or_ignore", to avoid the race condition seen in parallel writes when this is set to "error". The driver already handles the save mode check in on_write_start.
+        existing_data_behavior = "overwrite_or_ignore"
+
+        (
+            min_rows_per_group,
+            max_rows_per_group,
+            max_rows_per_file,
+        ) = choose_row_group_limits(
+            row_group_size,
+            min_rows_per_file=self.min_rows_per_file,
+            max_rows_per_file=self.max_rows_per_file,
+        )
+
+        basename_template = self._get_basename_template(filename, write_uuid)
+
+        # Note that the driver already handles the save mode logic, checking if the directory exists and raising an error if it does on SaveMode.ERROR
+        ds.write_dataset(
+            data=tables,
+            base_dir=self.path,
+            schema=output_schema,
+            basename_template=basename_template,
+            filesystem=self.filesystem,
+            partitioning=self.partition_cols,
+            format=FILE_FORMAT,
+            existing_data_behavior=existing_data_behavior,
+            partitioning_flavor=partitioning_flavor,
+            use_threads=True,
+            min_rows_per_group=min_rows_per_group,
+            max_rows_per_group=max_rows_per_group,
+            max_rows_per_file=max_rows_per_file,
+            file_options=ds.ParquetFileFormat().make_write_options(**write_kwargs),
+            create_dir=self.try_create_dir,
+        )
+
+    @property
+    def min_rows_per_write(self) -> Optional[int]:
+        return self.min_rows_per_file

python/ray/data/_internal/datasource/range_datasource.py

@@ -0,0 +1,146 @@
+import builtins
+import functools
+from typing import Iterable, List, Optional, Tuple
+
+import numpy as np
+
+from ray.data._internal.util import _check_pyarrow_version
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource import Datasource, ReadTask
+
+
+class RangeDatasource(Datasource):
+    """An example datasource that generates ranges of numbers from [0..n)."""
+
+    def __init__(
+        self,
+        n: int,
+        block_format: str = "arrow",
+        tensor_shape: Tuple = (1,),
+        column_name: Optional[str] = None,
+    ):
+        self._n = int(n)
+        self._block_format = block_format
+        self._tensor_shape = tensor_shape
+        self._column_name = column_name
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        if self._block_format == "tensor":
+            element_size = int(np.prod(self._tensor_shape))
+        else:
+            element_size = 1
+        return 8 * self._n * element_size
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        if self._n == 0:
+            return []
+
+        read_tasks: List[ReadTask] = []
+        n = self._n
+        block_format = self._block_format
+        tensor_shape = self._tensor_shape
+        block_size = max(1, n // parallelism)
+        # TODO(swang): This target block size may not match the driver's
+        # context if it was overridden. Set target max block size during
+        # optimizer stage to fix this.
+        ctx = DataContext.get_current()
+        if ctx.target_max_block_size is None:
+            # If target_max_block_size is ``None``, treat it as unlimited and
+            # avoid further splitting.
+            target_rows_per_block = n  # whole block in one shot
+        else:
+            row_size_bytes = self.estimate_inmemory_data_size() // self._n
+            row_size_bytes = max(row_size_bytes, 1)
+            target_rows_per_block = max(1, ctx.target_max_block_size // row_size_bytes)
+
+        # Example of a read task. In a real datasource, this would pull data
+        # from an external system instead of generating dummy data.
+        def make_block(start: int, count: int) -> Block:
+            if block_format == "arrow":
+                import pyarrow as pa
+
+                return pa.Table.from_arrays(
+                    [np.arange(start, start + count)],
+                    names=[self._column_name or "value"],
+                )
+            elif block_format == "tensor":
+                import pyarrow as pa
+
+                tensor = np.ones(tensor_shape, dtype=np.int64) * np.expand_dims(
+                    np.arange(start, start + count),
+                    tuple(range(1, 1 + len(tensor_shape))),
+                )
+                return BlockAccessor.batch_to_block(
+                    {self._column_name: tensor} if self._column_name else tensor
+                )
+            else:
+                return list(builtins.range(start, start + count))
+
+        def make_blocks(
+            start: int, count: int, target_rows_per_block: int
+        ) -> Iterable[Block]:
+            while count > 0:
+                num_rows = min(count, target_rows_per_block)
+                yield make_block(start, num_rows)
+                start += num_rows
+                count -= num_rows
+
+        if block_format == "tensor":
+            element_size = int(np.prod(tensor_shape))
+        else:
+            element_size = 1
+
+        i = 0
+        while i < n:
+            count = min(block_size, n - i)
+            meta = BlockMetadata(
+                num_rows=count,
+                size_bytes=8 * count * element_size,
+                input_files=None,
+                exec_stats=None,
+            )
+            read_tasks.append(
+                ReadTask(
+                    lambda i=i, count=count: make_blocks(
+                        i, count, target_rows_per_block
+                    ),
+                    meta,
+                    schema=self._schema,
+                    per_task_row_limit=per_task_row_limit,
+                )
+            )
+            i += block_size
+
+        return read_tasks
+
+    @functools.cached_property
+    def _schema(self):
+        if self._n == 0:
+            return None
+
+        if self._block_format == "arrow":
+            _check_pyarrow_version()
+            import pyarrow as pa
+
+            schema = pa.Table.from_pydict({self._column_name or "value": [0]}).schema
+        elif self._block_format == "tensor":
+            _check_pyarrow_version()
+            import pyarrow as pa
+
+            tensor = np.ones(self._tensor_shape, dtype=np.int64) * np.expand_dims(
+                np.arange(0, 10), tuple(range(1, 1 + len(self._tensor_shape)))
+            )
+            schema = BlockAccessor.batch_to_block(
+                {self._column_name: tensor} if self._column_name else tensor
+            ).schema
+        elif self._block_format == "list":
+            schema = int
+        else:
+            raise ValueError("Unsupported block type", self._block_format)
+        return schema

python/ray/data/_internal/datasource/sql_datasink.py

@@ -0,0 +1,35 @@
+from typing import Callable, Iterable
+
+from ray.data._internal.datasource.sql_datasource import Connection, _connect
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+
+class SQLDatasink(Datasink[None]):
+
+    _MAX_ROWS_PER_WRITE = 128
+
+    def __init__(self, sql: str, connection_factory: Callable[[], Connection]):
+        self.sql = sql
+        self.connection_factory = connection_factory
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        with _connect(self.connection_factory) as cursor:
+            for block in blocks:
+                block_accessor = BlockAccessor.for_block(block)
+
+                values = []
+                for row in block_accessor.iter_rows(public_row_format=False):
+                    values.append(tuple(row.values()))
+                    assert len(values) <= self._MAX_ROWS_PER_WRITE, len(values)
+                    if len(values) == self._MAX_ROWS_PER_WRITE:
+                        cursor.executemany(self.sql, values)
+                        values = []
+
+                if values:
+                    cursor.executemany(self.sql, values)

python/ray/data/_internal/datasource/sql_datasource.py

@@ -0,0 +1,211 @@
+import logging
+import math
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, List, Optional
+
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+Connection = Any  # A Python DB API2-compliant `Connection` object.
+Cursor = Any  # A Python DB API2-compliant `Cursor` object.
+
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+
+logger = logging.getLogger(__name__)
+
+
+def _cursor_to_block(cursor) -> Block:
+    import pyarrow as pa
+
+    rows = cursor.fetchall()
+    # Each `column_description` is a 7-element sequence. The first element is the column
+    # name. To learn more, read https://peps.python.org/pep-0249/#description.
+    columns = [column_description[0] for column_description in cursor.description]
+    pydict = {column: [row[i] for row in rows] for i, column in enumerate(columns)}
+    return pa.Table.from_pydict(pydict)
+
+
+def _check_connection_is_dbapi2_compliant(connection) -> None:
+    for attr in "close", "commit", "cursor":
+        if not hasattr(connection, attr):
+            raise ValueError(
+                "Your `connection_factory` created a `Connection` object without a "
+                f"{attr!r} method, but this method is required by the Python DB API2 "
+                "specification. Check that your database connector is DB API2-"
+                "compliant. To learn more, read https://peps.python.org/pep-0249/."
+            )
+
+
+def _check_cursor_is_dbapi2_compliant(cursor) -> None:
+    # These aren't all the methods required by the specification, but it's all the ones
+    # we care about.
+    for attr in "execute", "executemany", "fetchone", "fetchall", "description":
+        if not hasattr(cursor, attr):
+            raise ValueError(
+                "Your database connector created a `Cursor` object without a "
+                f"{attr!r} method, but this method is required by the Python DB API2 "
+                "specification. Check that your database connector is DB API2-"
+                "compliant. To learn more, read https://peps.python.org/pep-0249/."
+            )
+
+
+@contextmanager
+def _connect(connection_factory: Callable[[], Connection]) -> Iterator[Cursor]:
+    connection = connection_factory()
+    _check_connection_is_dbapi2_compliant(connection)
+
+    try:
+        cursor = connection.cursor()
+        _check_cursor_is_dbapi2_compliant(cursor)
+        yield cursor
+        connection.commit()
+    except Exception:
+        # `rollback` is optional since not all databases provide transaction support.
+        try:
+            connection.rollback()
+        except Exception as e:
+            # Each connector implements its own `NotSupportError` class, so we check
+            # the exception's name instead of using `isinstance`.
+            if (
+                isinstance(e, AttributeError)
+                or e.__class__.__name__ == "NotSupportedError"
+            ):
+                pass
+        raise
+    finally:
+        connection.close()
+
+
+def _execute(cursor: Cursor, sql: str, params: Optional[Any]) -> None:
+    if params is None:
+        cursor.execute(sql)
+    else:
+        cursor.execute(sql, params)
+
+
+class SQLDatasource(Datasource):
+    MIN_ROWS_PER_READ_TASK = 50
+
+    def __init__(
+        self,
+        sql: str,
+        connection_factory: Callable[[], Connection],
+        shard_hash_fn: str,
+        shard_keys: Optional[List[str]] = None,
+        sql_params: Optional[Any] = None,
+    ):
+        self.sql = sql
+        if shard_keys and len(shard_keys) > 1:
+            self.shard_keys = f"CONCAT({','.join(shard_keys)})"
+        elif shard_keys and len(shard_keys) == 1:
+            self.shard_keys = f"{shard_keys[0]}"
+        else:
+            self.shard_keys = None
+        self.shard_hash_fn = shard_hash_fn
+        self.connection_factory = connection_factory
+        self.sql_params = sql_params
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def supports_sharding(self, parallelism: int) -> bool:
+        """Check if database supports sharding with MOD/ABS/CONCAT operations.
+
+        Args:
+            parallelism: The number of shards to split the read into.
+
+        Returns:
+            bool: True if sharding is supported, False otherwise.
+        """
+        if parallelism <= 1 or self.shard_keys is None:
+            return False
+
+        # Test if database supports required operations (MOD, ABS, MD5, CONCAT)
+        # by executing a sample query
+        hash_fn = self.shard_hash_fn
+        query = (
+            f"SELECT COUNT(1) FROM ({self.sql}) as T"
+            f" WHERE MOD(ABS({hash_fn}({self.shard_keys})), {parallelism}) = 0"
+        )
+        try:
+            with _connect(self.connection_factory) as cursor:
+                _execute(cursor, query, self.sql_params)
+            return True
+        except Exception as e:
+            logger.info(f"Database does not support sharding: {str(e)}.")
+            return False
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        def fallback_read_fn() -> Iterable[Block]:
+            """Read all data in a single block when sharding is not supported."""
+            with _connect(self.connection_factory) as cursor:
+                _execute(cursor, self.sql, self.sql_params)
+                return [_cursor_to_block(cursor)]
+
+        # Check if sharding is supported by the database first
+        # If not, fall back to reading all data in a single task without counting rows
+        if not self.supports_sharding(parallelism):
+            logger.info(
+                "Sharding is not supported. "
+                "Falling back to reading all data in a single task."
+            )
+            metadata = BlockMetadata(None, None, None, None)
+            return [ReadTask(fallback_read_fn, metadata)]
+
+        # Only perform the expensive COUNT(*) query if sharding is supported
+        num_rows_total = self._get_num_rows()
+
+        if num_rows_total == 0:
+            return []
+
+        parallelism = min(
+            parallelism, math.ceil(num_rows_total / self.MIN_ROWS_PER_READ_TASK)
+        )
+        num_rows_per_block = num_rows_total // parallelism
+        num_blocks_with_extra_row = num_rows_total % parallelism
+
+        tasks = []
+        for i in range(parallelism):
+            num_rows = num_rows_per_block
+            if i < num_blocks_with_extra_row:
+                num_rows += 1
+            read_fn = self._create_parallel_read_fn(i, parallelism)
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=None,
+                input_files=None,
+                exec_stats=None,
+            )
+            tasks.append(
+                ReadTask(read_fn, metadata, per_task_row_limit=per_task_row_limit)
+            )
+
+        return tasks
+
+    def _get_num_rows(self) -> int:
+        with _connect(self.connection_factory) as cursor:
+            _execute(cursor, f"SELECT COUNT(*) FROM ({self.sql}) as T", self.sql_params)
+            return cursor.fetchone()[0]
+
+    def _create_parallel_read_fn(self, task_id: int, parallelism: int):
+        hash_fn = self.shard_hash_fn
+        query = (
+            f"SELECT * FROM ({self.sql}) as T "
+            f"WHERE MOD(ABS({hash_fn}({self.shard_keys})), {parallelism}) = {task_id}"
+        )
+
+        def read_fn() -> Iterable[Block]:
+            with _connect(self.connection_factory) as cursor:
+                _execute(cursor, query, self.sql_params)
+                block = _cursor_to_block(cursor)
+                return [block]
+
+        return read_fn

python/ray/data/_internal/datasource/text_datasource.py

@@ -0,0 +1,42 @@
+from typing import TYPE_CHECKING, Iterator, List
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class TextDatasource(FileBasedDatasource):
+    """Text datasource, for reading and writing text files."""
+
+    _COLUMN_NAME = "text"
+
+    def __init__(
+        self,
+        paths: List[str],
+        *,
+        drop_empty_lines: bool = False,
+        encoding: str = "utf-8",
+        **file_based_datasource_kwargs
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self.drop_empty_lines = drop_empty_lines
+        self.encoding = encoding
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        data = f.readall()
+
+        builder = DelegatingBlockBuilder()
+
+        lines = data.decode(self.encoding).splitlines()
+        for line in lines:
+            if self.drop_empty_lines and line.strip() == "":
+                continue
+            item = {self._COLUMN_NAME: line}
+            builder.add(item)
+
+        block = builder.build()
+        yield block

python/ray/data/_internal/datasource/tfrecords_datasink.py

@@ -0,0 +1,237 @@
+import struct
+from typing import TYPE_CHECKING, Dict, Iterable, Optional, Union
+
+import numpy as np
+
+from .tfrecords_datasource import _get_single_true_type
+from ray.data._internal.tensor_extensions.arrow import (
+    get_arrow_extension_tensor_types,
+)
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+if TYPE_CHECKING:
+    import pyarrow
+    import tensorflow as tf
+    from tensorflow_metadata.proto.v0 import schema_pb2
+
+
+class TFRecordDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        tf_schema: Optional["schema_pb2.Schema"] = None,
+        file_format: str = "tar",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        _check_import(self, module="crc32c", package="crc32c")
+
+        self.tf_schema = tf_schema
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        arrow_table = block.to_arrow()
+
+        # It seems like TFRecords are typically row-based,
+        # https://www.tensorflow.org/tutorials/load_data/tfrecord#writing_a_tfrecord_file_2
+        # so we must iterate through the rows of the block,
+        # serialize to tf.train.Example proto, and write to file.
+
+        examples = _convert_arrow_table_to_examples(arrow_table, self.tf_schema)
+
+        # Write each example to the arrow file in the TFRecord format.
+        for example in examples:
+            _write_record(file, example)
+
+
+def _convert_arrow_table_to_examples(
+    arrow_table: "pyarrow.Table",
+    tf_schema: Optional["schema_pb2.Schema"] = None,
+) -> Iterable["tf.train.Example"]:
+    import tensorflow as tf
+
+    schema_dict = {}
+    # Convert user-specified schema into dict for convenient mapping
+    if tf_schema is not None:
+        for schema_feature in tf_schema.feature:
+            schema_dict[schema_feature.name] = schema_feature.type
+
+    column_names = arrow_table.column_names
+
+    # Validate schema once up-front rather than per row.
+    if tf_schema is not None:
+        for name in column_names:
+            if name not in schema_dict:
+                raise ValueError(
+                    f"Found extra unexpected feature {name} "
+                    f"not in specified schema: {tf_schema}"
+                )
+
+    # Hoist per-column lookups out of the row loop. Iterating columns in
+    # lockstep with zip uses ChunkedArray.__iter__ (a C-level loop) instead
+    # of per-row __getitem__ calls, which avoids Python-side method dispatch
+    # on every element.
+    columns = arrow_table.columns
+    schema_feature_types = [schema_dict.get(name) for name in column_names]
+
+    for row_values in zip(*columns):
+        features: Dict[str, "tf.train.Feature"] = {
+            name: _value_to_feature(value, sft)
+            for name, value, sft in zip(column_names, row_values, schema_feature_types)
+        }
+        proto = tf.train.Example(features=tf.train.Features(feature=features))
+        yield proto
+
+
+def _value_to_feature(
+    value: Union["pyarrow.Scalar", "pyarrow.Array"],
+    schema_feature_type: Optional["schema_pb2.FeatureType"] = None,
+) -> "tf.train.Feature":
+    import pyarrow as pa
+    import tensorflow as tf
+
+    from ray.data._internal.utils.transform_pyarrow import _is_native_tensor_type
+
+    if isinstance(value, pa.ListScalar):
+        # Use the underlying type of the ListScalar's value in
+        # determining the output feature's data type.
+        value_type = value.type.value_type
+        value = value.as_py()
+    elif isinstance(value.type, get_arrow_extension_tensor_types()):
+        value_type = value.type
+        py_val = value.as_py()
+        if _is_native_tensor_type(value_type):
+            # PyArrow's native FixedShapeTensorType returns a flat list from
+            # as_py(), so use to_numpy()
+            value = value.to_numpy()
+        else:
+            value = py_val
+    else:
+        value_type = value.type
+        value = value.as_py()
+        if value is None:
+            value = []
+        else:
+            value = [value]
+
+    underlying_value_type = {
+        "bytes": pa.types.is_binary(value_type),
+        "string": pa.types.is_string(value_type),
+        "float": pa.types.is_floating(value_type),
+        "int": pa.types.is_integer(value_type),
+        "tensor": isinstance(value_type, get_arrow_extension_tensor_types()),
+    }
+    assert sum(bool(value) for value in underlying_value_type.values()) <= 1
+
+    if schema_feature_type is not None:
+        try:
+            from tensorflow_metadata.proto.v0 import schema_pb2
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "To use TensorFlow schemas, please install "
+                "the tensorflow-metadata package."
+            )
+        specified_feature_type = {
+            # We default anything that is not a string or tensor to be
+            # a byte array, mostly to deal with the case when we have
+            # null input, but we specify a schema.
+            "bytes": schema_feature_type == schema_pb2.FeatureType.BYTES
+            and not underlying_value_type["string"]
+            and not underlying_value_type["tensor"],
+            "string": schema_feature_type == schema_pb2.FeatureType.BYTES
+            and underlying_value_type["string"],
+            "float": schema_feature_type == schema_pb2.FeatureType.FLOAT,
+            "int": schema_feature_type == schema_pb2.FeatureType.INT,
+            "tensor": schema_feature_type == schema_pb2.FeatureType.BYTES
+            and underlying_value_type["tensor"],
+        }
+
+        und_type = _get_single_true_type(underlying_value_type)
+        spec_type = _get_single_true_type(specified_feature_type)
+        if und_type is not None and und_type != spec_type:
+            raise ValueError(
+                "Schema field type mismatch during write: specified type is "
+                f"{spec_type}, but underlying type is {und_type}",
+            )
+        # Override the underlying value type with the type in the user-specified schema.
+        underlying_value_type = specified_feature_type
+
+    if underlying_value_type["int"]:
+        return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
+    if underlying_value_type["float"]:
+        return tf.train.Feature(float_list=tf.train.FloatList(value=value))
+    if underlying_value_type["bytes"]:
+        return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+    if underlying_value_type["string"]:
+        value = [v.encode() for v in value]  # casting to bytes
+        return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+    if underlying_value_type["tensor"]:
+        if value is None:
+            value = []
+        else:
+            value = [tf.io.serialize_tensor(value).numpy()]
+        return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+    if pa.types.is_null(value_type):
+        raise ValueError(
+            "Unable to infer type from partially missing column. "
+            "Try setting read parallelism = 1, or use an input data source which "
+            "explicitly specifies the schema."
+        )
+    raise ValueError(
+        f"Value is of type {value_type}, "
+        "which we cannot convert to a supported tf.train.Feature storage type "
+        "(bytes, float, or int)."
+    )
+
+
+# Adapted from https://github.com/vahidk/tfrecord/blob/74b2d24a838081356d993ec0e147eaf59ccd4c84/tfrecord/writer.py#L57-L72  # noqa: E501
+#
+# MIT License
+#
+# Copyright (c) 2020 Vahid Kazemi
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+def _write_record(
+    file: "pyarrow.NativeFile",
+    example: "tf.train.Example",
+) -> None:
+    record = example.SerializeToString()
+    length = len(record)
+    length_bytes = struct.pack("<Q", length)
+    file.write(length_bytes)
+    file.write(_masked_crc(length_bytes))
+    file.write(record)
+    file.write(_masked_crc(record))
+
+
+def _masked_crc(data: bytes) -> bytes:
+    """CRC checksum."""
+    import crc32c
+
+    mask = 0xA282EAD8
+    crc = crc32c.crc32(data)
+    masked = ((crc >> 15) | (crc << 17)) + mask
+    masked = np.uint32(masked & np.iinfo(np.uint32).max)
+    masked_bytes = struct.pack("<I", masked)
+    return masked_bytes

python/ray/data/_internal/datasource/tfrecords_datasource.py

@@ -0,0 +1,259 @@
+import logging
+import struct
+from typing import TYPE_CHECKING, Dict, Iterable, Iterator, List, Optional, Union
+
+import pyarrow
+
+from ray.data._internal.tensor_extensions.arrow import pyarrow_table_from_pydict
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import tensorflow as tf
+    from tensorflow_metadata.proto.v0 import schema_pb2
+
+logger = logging.getLogger(__name__)
+
+
+class TFRecordDatasource(FileBasedDatasource):
+    """TFRecord datasource, for reading and writing TFRecord files."""
+
+    _FILE_EXTENSIONS = ["tfrecords"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        tf_schema: Optional["schema_pb2.Schema"] = None,
+        **file_based_datasource_kwargs,
+    ):
+        """Initialize the TFRecord datasource.
+
+        Args:
+            paths: One or more file paths to read TFRecord data from.
+            tf_schema: Optional TensorFlow Schema which is used to explicitly set
+                the schema of the underlying Dataset.
+            **file_based_datasource_kwargs: Additional keyword arguments forwarded
+                to the underlying :class:`FileBasedDatasource`.
+        """
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self._tf_schema = tf_schema
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import tensorflow as tf
+        from google.protobuf.message import DecodeError
+
+        for record in _read_records(f, path):
+            example = tf.train.Example()
+            try:
+                example.ParseFromString(record)
+            except DecodeError as e:
+                raise ValueError(
+                    "`TFRecordDatasource` failed to parse `tf.train.Example` "
+                    f"record in '{path}'. This error can occur if your TFRecord "
+                    f"file contains a message type other than `tf.train.Example`: {e}"
+                )
+
+            yield pyarrow_table_from_pydict(
+                _convert_example_to_dict(example, self._tf_schema)
+            )
+
+
+def _convert_example_to_dict(
+    example: "tf.train.Example",
+    tf_schema: Optional["schema_pb2.Schema"],
+) -> Dict[str, pyarrow.Array]:
+    record = {}
+    schema_dict = {}
+    # Convert user-specified schema into dict for convenient mapping
+    if tf_schema is not None:
+        for schema_feature in tf_schema.feature:
+            schema_dict[schema_feature.name] = schema_feature.type
+
+    for feature_name, feature in example.features.feature.items():
+        if tf_schema is not None and feature_name not in schema_dict:
+            raise ValueError(
+                f"Found extra unexpected feature {feature_name} "
+                f"not in specified schema: {tf_schema}"
+            )
+        schema_feature_type = schema_dict.get(feature_name)
+        record[feature_name] = _get_feature_value(feature, schema_feature_type)
+    return record
+
+
+def _get_single_true_type(dct) -> str:
+    """Utility function for getting the single key which has a `True` value in
+    a dict. Used to filter a dict of `{field_type: is_valid}` to get
+    the field type from a schema or data source."""
+    filtered_types = iter([_type for _type in dct if dct[_type]])
+    # In the case where there are no keys with a `True` value, return `None`
+    return next(filtered_types, None)
+
+
+def _get_feature_value(
+    feature: "tf.train.Feature",
+    schema_feature_type: Optional["schema_pb2.FeatureType"] = None,
+) -> pyarrow.Array:
+    import pyarrow as pa
+
+    underlying_feature_type = {
+        "bytes": feature.HasField("bytes_list"),
+        "float": feature.HasField("float_list"),
+        "int": feature.HasField("int64_list"),
+    }
+    # At most one of `bytes_list`, `float_list`, and `int64_list`
+    # should contain values. If none contain data, this indicates
+    # an empty feature value.
+    assert sum(bool(value) for value in underlying_feature_type.values()) <= 1
+
+    if schema_feature_type is not None:
+        try:
+            from tensorflow_metadata.proto.v0 import schema_pb2
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "To use TensorFlow schemas, please install "
+                "the tensorflow-metadata package."
+            )
+        # If a schema is specified, compare to the underlying type
+        specified_feature_type = {
+            "bytes": schema_feature_type == schema_pb2.FeatureType.BYTES,
+            "float": schema_feature_type == schema_pb2.FeatureType.FLOAT,
+            "int": schema_feature_type == schema_pb2.FeatureType.INT,
+        }
+        und_type = _get_single_true_type(underlying_feature_type)
+        spec_type = _get_single_true_type(specified_feature_type)
+        if und_type is not None and und_type != spec_type:
+            raise ValueError(
+                "Schema field type mismatch during read: specified type is "
+                f"{spec_type}, but underlying type is {und_type}",
+            )
+        # Override the underlying value type with the type in the user-specified schema.
+        underlying_feature_type = specified_feature_type
+
+    if underlying_feature_type["bytes"]:
+        value = feature.bytes_list.value
+        type_ = pa.binary()
+    elif underlying_feature_type["float"]:
+        value = feature.float_list.value
+        type_ = pa.float32()
+    elif underlying_feature_type["int"]:
+        value = feature.int64_list.value
+        type_ = pa.int64()
+    else:
+        value = []
+        type_ = pa.null()
+    value = list(value)
+    if len(value) == 1 and schema_feature_type is None:
+        # Use the value itself if the features contains a single value.
+        # This is to give better user experience when writing preprocessing UDF on
+        # these single-value lists.
+        value = value[0]
+    else:
+        # If the feature value is empty and no type is specified in the user-provided
+        # schema, set the type to null for now to allow pyarrow to construct a valid
+        # Array; later, infer the type from other records which have non-empty values
+        # for the feature.
+        if len(value) == 0:
+            type_ = pa.null()
+        type_ = pa.list_(type_)
+    return pa.array([value], type=type_)
+
+
+# Adapted from https://github.com/vahidk/tfrecord/blob/74b2d24a838081356d993ec0e147eaf59ccd4c84/tfrecord/reader.py#L16-L96  # noqa: E501
+#
+# MIT License
+#
+# Copyright (c) 2020 Vahid Kazemi
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+def _read_records(
+    file: "pyarrow.NativeFile",
+    path: str,
+) -> Iterable[memoryview]:
+    """
+    Read records from TFRecord file.
+
+    A TFRecord file contains a sequence of records. The file can only be read
+    sequentially. Each record is stored in the following formats:
+        uint64 length
+        uint32 masked_crc32_of_length
+        byte   data[length]
+        uint32 masked_crc32_of_data
+
+    See https://www.tensorflow.org/tutorials/load_data/tfrecord#tfrecords_format_details
+    for more details.
+    """
+    length_bytes = bytearray(8)
+    crc_bytes = bytearray(4)
+    datum_bytes = bytearray(1024 * 1024)
+    row_count = 0
+    while True:
+        try:
+            # Read "length" field.
+            num_length_bytes_read = file.readinto(length_bytes)
+            if num_length_bytes_read == 0:
+                break
+            elif num_length_bytes_read != 8:
+                raise ValueError(
+                    "Failed to read the length of record data. Expected 8 bytes but "
+                    "got {num_length_bytes_read} bytes."
+                )
+
+            # Read "masked_crc32_of_length" field.
+            num_length_crc_bytes_read = file.readinto(crc_bytes)
+            if num_length_crc_bytes_read != 4:
+                raise ValueError(
+                    "Failed to read the length of CRC-32C hashes. Expected 4 bytes "
+                    "but got {num_length_crc_bytes_read} bytes."
+                )
+
+            # Read "data[length]" field.
+            (data_length,) = struct.unpack("<Q", length_bytes)
+            if data_length > len(datum_bytes):
+                datum_bytes = datum_bytes.zfill(int(data_length * 1.5))
+            datum_bytes_view = memoryview(datum_bytes)[:data_length]
+            num_datum_bytes_read = file.readinto(datum_bytes_view)
+            if num_datum_bytes_read != data_length:
+                raise ValueError(
+                    f"Failed to read the record. Exepcted {data_length} bytes but got "
+                    f"{num_datum_bytes_read} bytes."
+                )
+
+            # Read "masked_crc32_of_data" field.
+            # TODO(chengsu): ideally we should check CRC-32C against the actual data.
+            num_crc_bytes_read = file.readinto(crc_bytes)
+            if num_crc_bytes_read != 4:
+                raise ValueError(
+                    "Failed to read the CRC-32C hashes. Expected 4 bytes but got "
+                    f"{num_crc_bytes_read} bytes."
+                )
+
+            # Return the data.
+            yield datum_bytes_view
+
+            row_count += 1
+            data_length = None
+        except Exception as e:
+            error_message = (
+                f"Failed to read TFRecord file {path}. Please ensure that the "
+                f"TFRecord file has correct format. Already read {row_count} rows."
+            )
+            if data_length is not None:
+                error_message += f" Byte size of current record data is {data_length}."
+            raise RuntimeError(error_message) from e

python/ray/data/_internal/datasource/torch_datasource.py

@@ -0,0 +1,82 @@
+from typing import TYPE_CHECKING, Optional
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import torch
+
+    from ray.data.context import DataContext
+
+
+TORCH_DATASOURCE_READER_BATCH_SIZE = 32
+
+
+class TorchDatasource(Datasource):
+    """Torch datasource, for reading from `Torch
+    datasets <https://pytorch.org/docs/stable/data.html/>`_.
+    This datasource implements a streaming read using a single read task.
+    """
+
+    def __init__(
+        self,
+        dataset: "torch.utils.data.Dataset",
+    ):
+        self._dataset = dataset
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ):
+        assert parallelism == 1
+
+        meta = BlockMetadata(
+            # Note: avoid len(self._dataset) because it will trigger
+            # iterating through IterableDataset, which can cause OOM.
+            num_rows=None,
+            size_bytes=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        read_task = ReadTask(
+            lambda subset=self._dataset: _read_subset(
+                subset,
+            ),
+            metadata=meta,
+            per_task_row_limit=per_task_row_limit,
+        )
+
+        return [read_task]
+
+    def estimate_inmemory_data_size(self):
+        return None
+
+
+def _read_subset(subset: "torch.utils.data.Subset"):
+    batch = []
+
+    # Get items from dataset based on its type
+    if hasattr(subset, "__iter__"):
+        # IterableDataset: Use the iterator directly
+        items = subset
+    else:
+        # Map-style dataset: Respect __len__
+        items = (subset[i] for i in range(len(subset)))
+
+    # Process items in batches
+    for item in items:
+        batch.append(item)
+        if len(batch) == TORCH_DATASOURCE_READER_BATCH_SIZE:
+            builder = DelegatingBlockBuilder()
+            builder.add_batch({"item": batch})
+            yield builder.build()
+            batch.clear()
+
+    # Handle any remaining items
+    if len(batch) > 0:
+        builder = DelegatingBlockBuilder()
+        builder.add_batch({"item": batch})
+        yield builder.build()

python/ray/data/_internal/datasource/turbopuffer_datasink.py

@@ -0,0 +1,427 @@
+"""
+TurbopufferDatasink - Ray Data datasink for Turbopuffer vector database
+
+Implementation following the pattern of MongoDatasink and Daft's Turbopuffer sink.
+
+This is based on [Turbopuffer Write API](https://turbopuffer.com/docs/write)
+"""
+
+import logging
+import os
+from typing import TYPE_CHECKING, Iterable, Literal, Optional, Union
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from ray._common.retry import call_with_retry
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+if TYPE_CHECKING:
+    import turbopuffer
+
+logger = logging.getLogger(__name__)
+
+# Reserved column names for Turbopuffer
+_ID_COLUMN = "id"
+_VECTOR_COLUMN = "vector"
+TURBOPUFFER_API_KEY_ENV_VAR = "TURBOPUFFER_API_KEY"
+
+
+class TurbopufferDatasink(Datasink):
+    """Turbopuffer Ray Datasink.
+
+    A Ray :class:`~ray.data.datasource.Datasink` for writing data into the
+    `Turbopuffer <https://turbopuffer.com/>`_ vector database.
+
+    Supports two modes of operation:
+
+    * **Single namespace** -- provide ``namespace`` to write all rows into one
+      Turbopuffer namespace.
+    * **Multi-namespace** -- provide ``namespace_column`` to route each row to
+      the namespace whose name is stored in that column.  The column is
+      automatically dropped before the data is sent to Turbopuffer.
+
+    Exactly one of ``namespace`` or ``namespace_column`` must be supplied.
+
+    Args:
+        namespace: Name of the Turbopuffer namespace to write into.
+            Mutually exclusive with ``namespace_column``.
+        namespace_column: Name of a column whose values determine the
+            target namespace for each row.  Rows are grouped by this column
+            and each group is written to its corresponding namespace.  The
+            column is removed from the data before writing.  Mutually
+            exclusive with ``namespace``.
+        region: Turbopuffer region identifier (for example,
+            ``"gcp-us-central1"``).  Mutually exclusive with ``base_url``.
+            Exactly one of ``region`` or ``base_url`` must be supplied.
+        base_url: Base URL for the Turbopuffer API (for example,
+            ``"https://gcp-us-central1.turbopuffer.com"``).  Mutually
+            exclusive with ``region``.  Exactly one of ``region`` or
+            ``base_url`` must be supplied.
+        api_key: Turbopuffer API key. If omitted, the value is read from the
+            ``TURBOPUFFER_API_KEY`` environment variable.
+        schema: Optional Turbopuffer schema definition to pass along with
+            writes. If provided, it is forwarded as the ``schema`` argument
+            to ``namespace.write``.
+        id_column: Name of the column to treat as the document identifier.
+            Rows with null IDs are dropped before writing. Defaults to ``"id"``.
+        vector_column: Name of the column containing embedding vectors.
+            If this differs from ``"vector"``, it is renamed to ``"vector"``
+            before writing. Defaults to ``"vector"``.
+        batch_size: Maximum number of rows to include in a single Turbopuffer
+            write call (logical row batching; subject to Turbopuffer's
+            256MiB request-size limit). Defaults to ``10000``.
+        distance_metric: Distance metric for the namespace. Passed to
+            ``namespace.write`` as the ``distance_metric`` argument.
+            Defaults to ``"cosine_distance"``.
+        concurrency: Unused; Ray Data controls write parallelism via
+            :meth:`~ray.data.Dataset.write_datasink` ``concurrency``.
+
+    Examples:
+        Write to a single namespace using a region:
+
+        .. testcode::
+           :skipif: True
+
+           import ray
+           from ray.data._internal.datasource.turbopuffer_datasink import (
+               TurbopufferDatasink,
+           )
+
+           ds = ray.data.range(100)
+           ds = ds.map_batches(lambda batch: {"id": batch["id"], "vector": ...})
+
+           ds.write_datasink(
+               TurbopufferDatasink(
+                   namespace="my-namespace",
+                   api_key="<YOUR_API_KEY>",
+                   region="gcp-us-central1",
+               )
+           )
+
+        Write using a base URL instead of a region:
+
+        .. testcode::
+           :skipif: True
+
+           ds.write_datasink(
+               TurbopufferDatasink(
+                   namespace="my-namespace",
+                   api_key="<YOUR_API_KEY>",
+                   base_url="https://gcp-us-central1.turbopuffer.com",
+               )
+           )
+
+        Write to multiple namespaces driven by a column:
+
+        .. testcode::
+           :skipif: True
+
+           ds.write_datasink(
+               TurbopufferDatasink(
+                   namespace_column="tenant",
+                   api_key="<YOUR_API_KEY>",
+                   region="gcp-us-central1",
+               )
+           )
+    """
+
+    def __init__(
+        self,
+        namespace: Optional[str] = None,
+        *,
+        namespace_column: Optional[str] = None,
+        region: Optional[str] = None,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        schema: Optional[dict] = None,
+        id_column: str = "id",
+        vector_column: str = "vector",
+        batch_size: int = 10000,
+        distance_metric: Literal[
+            "cosine_distance", "euclidean_distance"
+        ] = "cosine_distance",
+        concurrency: Optional[int] = None,
+    ):
+        _check_import(self, module="turbopuffer", package="turbopuffer")
+
+        # Validate namespace / namespace_column mutual exclusivity.
+        if namespace and namespace_column:
+            raise ValueError(
+                "Specify exactly one of 'namespace' or 'namespace_column', " "not both."
+            )
+        if not namespace and not namespace_column:
+            raise ValueError(
+                "Either 'namespace' or 'namespace_column' must be provided."
+            )
+
+        # Validate region / base_url mutual exclusivity.
+        if region is not None and base_url is not None:
+            raise ValueError("Specify exactly one of 'region' or 'base_url', not both.")
+        if region is None and base_url is None:
+            raise ValueError("Either 'region' or 'base_url' must be provided.")
+
+        # Store configuration
+        self.namespace = namespace
+        self.namespace_column = namespace_column
+        self.api_key = api_key or os.getenv(TURBOPUFFER_API_KEY_ENV_VAR)
+        self.region = region
+        self.base_url = base_url
+        self.schema = schema
+        self.id_column = id_column
+        self.vector_column = vector_column
+        self.batch_size = batch_size
+        self.distance_metric = distance_metric
+
+        # Validate column configuration
+        if self.id_column == self.vector_column:
+            raise ValueError(
+                "id_column and vector_column refer to the same column "
+                f"'{self.id_column}'. They must be distinct."
+            )
+
+        if self.namespace_column and self.namespace_column in (
+            self.id_column,
+            self.vector_column,
+        ):
+            raise ValueError(
+                f"namespace_column '{self.namespace_column}' must not be the "
+                f"same as id_column ('{self.id_column}') or vector_column "
+                f"('{self.vector_column}')."
+            )
+
+        # Validate API key
+        if not self.api_key:
+            raise ValueError(
+                "API key is required. Provide via api_key parameter or "
+                "TURBOPUFFER_API_KEY environment variable"
+            )
+
+        # Initialize client
+        self._client = None
+
+    def __getstate__(self) -> dict:
+        """Exclude `_client` during pickling."""
+        state = self.__dict__.copy()
+        state.pop("_client", None)
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        self._client = None
+
+    def _get_client(self):
+        """Lazy initialize Turbopuffer client."""
+        if self._client is None:
+            import turbopuffer
+
+            kwargs = {"api_key": self.api_key}
+            if self.region is not None:
+                kwargs["region"] = self.region
+            else:
+                kwargs["base_url"] = self.base_url
+            self._client = turbopuffer.Turbopuffer(**kwargs)
+        return self._client
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        """
+        Write blocks to Turbopuffer in a streaming fashion.
+
+        For memory efficiency, blocks are processed one at a time rather than
+        concatenating all blocks into a single large table. This follows the
+        pattern used by ClickHouseDatasink.
+
+        Each block is prepared (columns renamed, null IDs filtered), then
+        written in batches of ``batch_size``.
+
+        When ``namespace_column`` is set, each block is grouped by the
+        namespace column and each group is written to its corresponding
+        Turbopuffer namespace.
+        """
+        client = self._get_client()
+
+        for block in blocks:
+            accessor = BlockAccessor.for_block(block)
+            table = accessor.to_arrow()
+
+            if table.num_rows == 0:
+                continue
+
+            if self.namespace_column:
+                # Multi-namespace: group by namespace column, write to each.
+                self._write_multi_namespace(client, table)
+            else:
+                # Single namespace.
+                table = self._prepare_arrow_table(table)
+
+                if table.num_rows == 0:
+                    continue
+
+                ns = client.namespace(self.namespace)
+                for batch in table.to_batches(max_chunksize=self.batch_size):
+                    self._write_batch_with_retry(ns, batch, self.namespace)
+
+    def _rename_column_if_needed(
+        self,
+        table: pa.Table,
+        source_column: str,
+        target_column: str,
+        column_type: str,
+    ) -> pa.Table:
+        """
+        Rename a column in the table if it differs from the target name.
+
+        Args:
+            table: The Arrow table to modify.
+            source_column: The current column name in the table.
+            target_column: The required column name for Turbopuffer.
+            column_type: Human-readable type for error messages (e.g., "ID", "Vector").
+
+        Returns:
+            The table with the column renamed, or the original table if no rename needed.
+
+        Raises:
+            ValueError: If source column is missing or target column already exists.
+        """
+        if source_column not in table.column_names:
+            raise ValueError(
+                f"{column_type} column '{source_column}' not found in table"
+            )
+
+        # No rename needed if source and target are the same
+        if source_column == target_column:
+            return table
+
+        if target_column in table.column_names:
+            raise ValueError(
+                f"Table already has a '{target_column}' column; cannot also rename "
+                f"'{source_column}' to '{target_column}'. Please disambiguate your schema."
+            )
+
+        return BlockAccessor.for_block(table).rename_columns(
+            {source_column: target_column}
+        )
+
+    def _prepare_arrow_table(self, table: pa.Table) -> pa.Table:
+        """
+        Prepare Arrow table for Turbopuffer write.
+
+        1. Rename ID column to "id" if needed
+        2. Rename vector column to "vector" if needed
+        3. Filter out rows with null IDs
+        """
+        table = self._rename_column_if_needed(table, self.id_column, _ID_COLUMN, "ID")
+        table = self._rename_column_if_needed(
+            table, self.vector_column, _VECTOR_COLUMN, "Vector"
+        )
+
+        # Filter out rows with null IDs
+        if _ID_COLUMN in table.column_names:
+            table = table.filter(pc.is_valid(table.column(_ID_COLUMN)))
+
+        return table
+
+    def _write_multi_namespace(
+        self, client: "turbopuffer.Turbopuffer", table: pa.Table
+    ) -> None:
+        """Group rows by ``namespace_column`` and write each group to its namespace.
+
+        Uses :meth:`BlockAccessor._iter_groups_sorted` for efficient
+        zero-copy slicing by group.
+        """
+        group_col_name = self.namespace_column
+
+        if group_col_name not in table.column_names:
+            raise ValueError(
+                f"Namespace column '{group_col_name}' not found in table. "
+                f"Available columns: {table.column_names}"
+            )
+
+        # Reject null namespace values early -- we cannot route them.
+        ns_col = table.column(group_col_name)
+        if pc.any(pc.is_null(ns_col)).as_py():
+            raise ValueError(
+                f"Namespace column '{group_col_name}' contains null values; "
+                "fill or drop them before writing with namespace_column."
+            )
+
+        # Sort by the namespace column so _iter_groups_sorted can yield
+        # contiguous zero-copy slices for each unique namespace value.
+        sort_key = SortKey(key=group_col_name, descending=False)
+        block_accessor = BlockAccessor.for_block(
+            BlockAccessor.for_block(table).sort(sort_key)
+        )
+
+        for (namespace_name,), group_table in block_accessor._iter_groups_sorted(
+            sort_key
+        ):
+            # Drop the namespace column -- it is routing metadata, not data.
+            group_table = group_table.drop(group_col_name)
+
+            # Prepare (rename id/vector columns, filter null IDs).
+            group_table = self._prepare_arrow_table(group_table)
+            if group_table.num_rows == 0:
+                continue
+
+            ns = client.namespace(namespace_name)
+            for batch in group_table.to_batches(max_chunksize=self.batch_size):
+                self._write_batch_with_retry(ns, batch, namespace_name)
+
+    def _transform_to_turbopuffer_format(
+        self, table: Union[pa.Table, pa.RecordBatch]
+    ) -> dict:
+        if _ID_COLUMN not in table.column_names:
+            raise ValueError(f"Table must have '{_ID_COLUMN}' column")
+
+        # Cast 16-byte binary ID column to native UUID type for Turbopuffer performance.
+        # Native UUIDs are 16 bytes vs 36 bytes for string-encoded UUIDs.
+        # See: https://turbopuffer.com/docs/performance
+        id_col = table.column(_ID_COLUMN)
+        if pa.types.is_fixed_size_binary(id_col.type) and id_col.type.byte_width == 16:
+            # Cast fixed_size_binary(16) to uuid type
+            uuid_col = id_col.cast(pa.uuid())
+            table = table.set_column(
+                table.schema.get_field_index(_ID_COLUMN), _ID_COLUMN, uuid_col
+            )
+
+        # to_pydict() on UuidArray automatically returns uuid.UUID objects
+        return table.to_pydict()
+
+    def _write_batch_with_retry(
+        self,
+        namespace: "turbopuffer.Namespace",
+        batch: pa.Table,
+        namespace_name: Optional[str] = None,
+    ):
+        """Write a single batch with exponential backoff retry.
+
+        Args:
+            namespace: The Turbopuffer namespace object to write to.
+            batch: Arrow table or record-batch to write.
+            namespace_name: Human-readable namespace name for log messages.
+                Falls back to ``self.namespace`` when not provided.
+        """
+        ns_label = namespace_name or self.namespace
+        try:
+            batch_data = self._transform_to_turbopuffer_format(batch)
+            call_with_retry(
+                lambda: namespace.write(
+                    upsert_columns=batch_data,
+                    schema=self.schema,
+                    distance_metric=self.distance_metric,
+                ),
+                description=f"write batch to namespace '{ns_label}'",
+                max_attempts=5,
+                max_backoff_s=32,
+            )
+        except Exception as e:
+            logger.error(f"Write failed for namespace '{ns_label}' after retries: {e}")
+            raise

python/ray/data/_internal/datasource/video_datasource.py

@@ -0,0 +1,70 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+logger = logging.getLogger(__name__)
+
+
+class VideoDatasource(FileBasedDatasource):
+    _FILE_EXTENSIONS = [
+        "mp4",
+        "mkv",
+        "mov",
+        "avi",
+        "wmv",
+        "flv",
+        "webm",
+        "m4v",
+        "3gp",
+        "mpeg",
+        "mpg",
+        "ts",
+        "ogv",
+        "rm",
+        "rmvb",
+        "vob",
+        "asf",
+        "f4v",
+        "m2ts",
+        "mts",
+        "divx",
+        "xvid",
+        "mxf",
+    ]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        include_timestamps=False,
+        decord_load_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="decord", package="decord")
+
+        self.include_timestamps = include_timestamps
+        if decord_load_args is None:
+            self.decord_load_args = {}
+        else:
+            self.decord_load_args = decord_load_args
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        from decord import VideoReader
+
+        reader = VideoReader(f, **self.decord_load_args)
+
+        for frame_index, frame in enumerate(reader):
+            item = {"frame": frame.asnumpy(), "frame_index": frame_index}
+            if self.include_timestamps is True:
+                item["frame_timestamp"] = reader.get_frame_timestamp(frame_index)
+
+            builder = DelegatingBlockBuilder()
+            builder.add(item)
+            yield builder.build()

python/ray/data/_internal/datasource/webdataset_datasink.py

@@ -0,0 +1,53 @@
+import io
+import tarfile
+import time
+import uuid
+from typing import Optional, Union
+
+import pyarrow
+
+from ray.data._internal.datasource.webdataset_datasource import (
+    _apply_list,
+    _default_encoder,
+    _make_iterable,
+)
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class WebDatasetDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        encoder: Optional[Union[bool, str, callable, list]] = True,
+        *,
+        file_format: str = "tar",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format="tar", **file_datasink_kwargs)
+
+        self.encoder = encoder
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        stream = tarfile.open(fileobj=file, mode="w|")
+        samples = _make_iterable(block)
+        for sample in samples:
+            if not isinstance(sample, dict):
+                sample = sample.as_pydict()
+            if self.encoder is not None:
+                sample = _apply_list(self.encoder, sample, default=_default_encoder)
+            if "__key__" not in sample:
+                sample["__key__"] = uuid.uuid4().hex
+            key = sample["__key__"]
+            for k, v in sample.items():
+                if v is None or k.startswith("__"):
+                    continue
+                assert isinstance(v, bytes) or isinstance(v, str)
+                if not isinstance(v, bytes):
+                    v = v.encode("utf-8")
+                ti = tarfile.TarInfo(f"{key}.{k}")
+                ti.size = len(v)
+                ti.mtime = time.time()
+                ti.mode, ti.uname, ti.gname = 0o644, "data", "data"
+                stream.addfile(ti, io.BytesIO(v))
+        stream.close()

python/ray/data/_internal/datasource/webdataset_datasource.py

@@ -0,0 +1,442 @@
+# Copyright NVIDIA Corporation 2023
+# SPDX-License-Identifier: Apache-2.0
+
+import fnmatch
+import io
+import json
+import re
+import tarfile
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Union
+
+from ray._common.utils import env_bool
+from ray.data._internal.util import iterate_with_retry
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR = (
+    "RAY_DATA_WEBDATASET_ALLOW_UNSAFE_DESERIALIZATION"
+)
+
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+def _base_plus_ext(path: str):
+    """Split off all file extensions.
+
+    Returns base, allext.
+
+    Args:
+        path: path with extensions
+
+    Returns:
+        str: path with all extensions removed
+    """
+    match = re.match(r"^((?:.*/|)[^.]+)[.]([^/]*)$", path)
+    if not match:
+        return None, None
+    return match.group(1), match.group(2)
+
+
+def _valid_sample(sample: Dict[str, Any]):
+    """Check whether a sample is valid.
+
+    Args:
+        sample: sample to be checked
+
+    Returns:
+        ``True`` if the sample is a non-empty dict without the ``__bad__`` flag.
+    """
+    return (
+        sample is not None
+        and isinstance(sample, dict)
+        and len(list(sample.keys())) > 0
+        and not sample.get("__bad__", False)
+    )
+
+
+def _apply_list(
+    f: Union[Callable, List[Callable]], sample: Dict[str, Any], default: Callable = None
+):
+    """Apply a list of functions to a sample.
+
+    Args:
+        f: function or list of functions
+        sample: sample to be modified
+        default: default function to be applied to all keys.
+            Defaults to None.
+
+    Returns:
+        modified sample
+    """
+    if f is None:
+        return sample
+    if not isinstance(f, list):
+        f = [f]
+    for g in f:
+        if default is not None and not callable(g):
+            g = partial(default, format=g)
+        sample = g(sample)
+    return sample
+
+
+def _check_suffix(suffix: str, suffixes: Union[list, callable]):
+    """Check whether a suffix is valid.
+
+    Suffixes can be either None (=accept everything), a callable,
+    or a list of patterns. If the pattern contains */? it is treated
+    as a glob pattern, otherwise it is treated as a literal.
+
+    Args:
+        suffix: suffix to be checked
+        suffixes: list of valid suffixes
+
+    Returns:
+        ``True`` if the suffix matches the allowed patterns.
+    """
+    if suffixes is None:
+        return True
+    if callable(suffixes):
+        return suffixes(suffix)
+    for pattern in suffixes:
+        if "*" in pattern or "?" in pattern:
+            if fnmatch.fnmatch("." + suffix, pattern):
+                return True
+        elif suffix == pattern or "." + suffix == pattern:
+            return True
+    return False
+
+
+def _tar_file_iterator(
+    fileobj: Any,
+    fileselect: Optional[Union[bool, callable, list]] = None,
+    filerename: Optional[Union[bool, callable, list]] = None,
+    verbose_open: bool = False,
+    meta: dict = None,
+) -> Iterator[Dict[str, Any]]:
+    """Iterate over tar file, yielding filename, content pairs for the given tar stream.
+
+    Args:
+        fileobj: file object
+        fileselect: patterns or function selecting
+            files to be selected
+        filerename: patterns or function used to rename selected files
+            before yielding them.
+        verbose_open: if ``True``, print progress messages when starting
+            and finishing iteration over the tar stream.
+        meta: metadata to be added to each sample
+
+    Yields:
+        Dict[str, Any]: Dictionaries with ``fname`` and ``data`` keys for each
+        selected file in the tar stream.
+    """
+    meta = meta or {}
+    stream = tarfile.open(fileobj=fileobj, mode="r|*")
+    if verbose_open:
+        print(f"start {meta}")
+    for tarinfo in stream:
+        fname = tarinfo.name
+        if not tarinfo.isreg() or fname is None:
+            continue
+        data = stream.extractfile(tarinfo).read()
+        fname = _apply_list(filerename, fname)
+        assert isinstance(fname, str)
+        if not _check_suffix(fname, fileselect):
+            continue
+        result = dict(fname=fname, data=data)
+        yield result
+    if verbose_open:
+        print(f"done {meta}")
+
+
+def _group_by_keys(
+    data: List[Dict[str, Any]],
+    keys: callable = _base_plus_ext,
+    suffixes: Optional[Union[list, callable]] = None,
+    meta: dict = None,
+) -> Iterator[Dict[str, Any]]:
+    """Return function over iterator that groups key, value pairs into samples.
+
+    Args:
+        data: iterator over key, value pairs
+        keys: function that returns key, suffix for a given key
+        suffixes: list of suffixes to be included in the sample
+        meta: metadata to be added to each sample
+
+    Yields:
+        Dict[str, Any]: Grouped samples, where files sharing the same key prefix are
+        combined into a single dictionary.
+    """
+    meta = meta or {}
+    current_sample = None
+    for filesample in data:
+        assert isinstance(filesample, dict)
+        fname, value = filesample["fname"], filesample["data"]
+        prefix, suffix = keys(fname)
+        if prefix is None:
+            continue
+        if current_sample is None or prefix != current_sample["__key__"]:
+            if _valid_sample(current_sample):
+                current_sample.update(meta)
+                yield current_sample
+            current_sample = dict(__key__=prefix)
+            if "__url__" in filesample:
+                current_sample["__url__"] = filesample["__url__"]
+        if suffix in current_sample:
+            raise ValueError(
+                f"{fname}: duplicate file name in tar file "
+                + f"{suffix} {current_sample.keys()}, tar is {meta['__url__']}"
+            )
+        if suffixes is None or _check_suffix(suffix, suffixes):
+            current_sample[suffix] = value
+    if _valid_sample(current_sample):
+        current_sample.update(meta)
+        yield current_sample
+
+
+def _default_decoder(
+    sample: Dict[str, Any],
+    format: Optional[Union[bool, str]] = True,
+    allow_unsafe: bool = False,
+):
+    """A default decoder for webdataset.
+
+    This handles common file extensions: .txt, .cls, .cls2,
+        .jpg, .png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl.
+    These are the most common extensions used in webdataset.
+    For other extensions, users can provide their own decoder.
+
+    Args:
+        sample: sample, modified in place
+        format: optional image format hint (e.g. ``"PIL"`` to return PIL
+            images instead of numpy arrays).
+        allow_unsafe: if True, allow pickle/torch deserialization
+
+    Returns:
+        The sample with values decoded according to their key extension.
+    """
+    sample = dict(sample)
+    for key, value in sample.items():
+        extension = key.split(".")[-1]
+        if key.startswith("__"):
+            continue
+        elif extension in ["txt", "text"]:
+            sample[key] = value.decode("utf-8")
+        elif extension in ["cls", "cls2"]:
+            sample[key] = int(value.decode("utf-8"))
+        elif extension in ["jpg", "png", "ppm", "pgm", "pbm", "pnm"]:
+            import numpy as np
+            import PIL.Image
+
+            if format == "PIL":
+                sample[key] = PIL.Image.open(io.BytesIO(value))
+            else:
+                sample[key] = np.asarray(PIL.Image.open(io.BytesIO(value)))
+        elif extension == "json":
+            sample[key] = json.loads(value)
+        elif extension == "npy":
+            import numpy as np
+
+            sample[key] = np.load(io.BytesIO(value))
+        elif extension == "mp":
+            import msgpack
+
+            sample[key] = msgpack.unpackb(value, raw=False)
+        elif extension in ["pt", "pth"]:
+            if not allow_unsafe:
+                raise ValueError(
+                    f"Refusing to load .{extension} member {key!r} from "
+                    f"WebDataset with weights_only=False (arbitrary code "
+                    f"execution risk). Provide a custom decoder or set "
+                    f"{ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR}=1 "
+                    f"for trusted sources."
+                )
+            import torch
+
+            sample[key] = torch.load(io.BytesIO(value), weights_only=False)
+        elif extension in ["pickle", "pkl"]:
+            if not allow_unsafe:
+                raise ValueError(
+                    f"Refusing to unpickle WebDataset member {key!r} "
+                    f"(arbitrary code execution risk). Provide a custom "
+                    f"decoder or set "
+                    f"{ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR}=1 "
+                    f"for trusted sources."
+                )
+            import pickle
+
+            sample[key] = pickle.loads(value)
+    return sample
+
+
+extension_to_format = {"jpg": "jpeg"}
+
+
+def _default_encoder(sample: Dict[str, Any], format: Optional[Union[str, bool]] = True):
+    """A default encoder for webdataset.
+
+    This handles common file extensions: .txt, .cls, .cls2, .jpg,
+        .png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl
+    These are the most common extensions used in webdataset.
+    For other extensions, users can provide their own encoder.
+
+    Args:
+        sample: sample to encode.
+        format: optional image format hint forwarded to the underlying
+            image encoder.
+
+    Returns:
+        The sample with values encoded according to their key extension.
+    """
+    sample = dict(sample)
+    for key, value in sample.items():
+        extension = key.split(".")[-1]
+        if key.startswith("__"):
+            continue
+        elif extension in ["txt"]:
+            sample[key] = value.encode("utf-8")
+        elif extension in ["cls", "cls2"]:
+            sample[key] = str(value).encode("utf-8")
+        elif extension in ["jpg", "jpeg", "png", "ppm", "pgm", "pbm", "pnm"]:
+            import numpy as np
+            import PIL.Image
+
+            if isinstance(value, np.ndarray):
+                value = PIL.Image.fromarray(value)
+            assert isinstance(value, PIL.Image.Image)
+            stream = io.BytesIO()
+            value.save(
+                stream, format=extension_to_format.get(extension.lower(), extension)
+            )
+            sample[key] = stream.getvalue()
+        elif extension == "json":
+            sample[key] = json.dumps(value).encode("utf-8")
+        elif extension == "npy":
+            import numpy as np
+
+            stream = io.BytesIO()
+            np.save(stream, value)
+            sample[key] = stream.getvalue()
+        elif extension == "mp":
+            import msgpack
+
+            sample[key] = msgpack.dumps(value)
+        elif extension in ["pt", "pth"]:
+            import torch
+
+            stream = io.BytesIO()
+            torch.save(value, stream)
+            sample[key] = stream.getvalue()
+        elif extension in ["pickle", "pkl"]:
+            import pickle
+
+            stream = io.BytesIO()
+            pickle.dump(value, stream)
+            sample[key] = stream.getvalue()
+    return sample
+
+
+def _make_iterable(block: BlockAccessor):
+    """Make a block iterable.
+
+    This is a placeholder for dealing with more complex blocks.
+
+    Args:
+        block: Ray Dataset block
+
+    Returns:
+        Iterable[Dict[str,Any]]: Iterable of samples
+    """
+    return block.iter_rows(public_row_format=False)
+
+
+class WebDatasetDatasource(FileBasedDatasource):
+    """A Datasource for WebDataset datasets (tar format with naming conventions)."""
+
+    _FILE_EXTENSIONS = ["tar"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        decoder: Optional[Union[bool, str, callable, list]] = True,
+        fileselect: Optional[Union[bool, callable, list]] = None,
+        filerename: Optional[Union[bool, callable, list]] = None,
+        suffixes: Optional[Union[bool, callable, list]] = None,
+        verbose_open: bool = False,
+        expand_json: bool = False,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self.decoder = decoder
+        self.fileselect = fileselect
+        self.filerename = filerename
+        self.suffixes = suffixes
+        self.verbose_open = verbose_open
+        self.expand_json = expand_json
+
+        self._allow_unsafe_deserialization = env_bool(
+            ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR, False
+        )
+
+    def _read_stream(self, stream: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        """Read and decode samples from a stream.
+
+        Note that fileselect selects files during reading, while suffixes
+        selects files during the grouping step.
+
+        Args:
+            stream: File descriptor to read from.
+            path: Path to the data.
+
+        Yields:
+            Block: Single-row blocks (one per WebDataset sample).
+        """
+
+        import pandas as pd
+
+        def get_tar_file_iterator():
+            return _tar_file_iterator(
+                stream,
+                fileselect=self.fileselect,
+                filerename=self.filerename,
+                verbose_open=self.verbose_open,
+            )
+
+        # S3 can raise transient errors during iteration
+        files = iterate_with_retry(
+            get_tar_file_iterator,
+            "iterate tar file",
+            match=self._data_context.retried_io_errors,
+        )
+
+        samples = _group_by_keys(files, meta=dict(__url__=path), suffixes=self.suffixes)
+        default_decoder = partial(
+            _default_decoder, allow_unsafe=self._allow_unsafe_deserialization
+        )
+        for sample in samples:
+            if self.decoder is not None:
+                sample = _apply_list(self.decoder, sample, default=default_decoder)
+            if self.expand_json:
+                if isinstance(sample["json"], bytes):
+                    parsed_json = json.loads(sample["json"].decode("utf-8"))
+                elif isinstance(sample["json"], str):
+                    parsed_json = json.loads(sample["json"])
+                elif isinstance(sample["json"], dict):
+                    parsed_json = sample["json"]
+                else:
+                    raise TypeError(
+                        f"Unsupported data type" f" {type(sample['json'])} for sample"
+                    )
+                for k, v in parsed_json.items():
+                    if k not in sample:
+                        sample[k] = []
+                    sample[k].append(v)
+            yield pd.DataFrame(
+                {
+                    k: v if isinstance(v, list) and len(v) == 1 else [v]
+                    for k, v in sample.items()
+                }
+            )

python/ray/data/_internal/datasource/zarrv2_datasource.py

@@ -0,0 +1,685 @@
+from __future__ import annotations
+
+import logging
+import math
+import numbers
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, List, Optional
+
+import numpy as np
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import (
+    _check_import,
+    _is_local_scheme,
+    iterate_with_retry,
+)
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from fsspec.spec import AbstractFileSystem
+    from pyarrow import fs as pyarrow_fs
+    from zarr import Array as ZarrArray
+    from zarr.hierarchy import Group as ZarrGroup
+
+    from ray.data.context import DataContext
+
+    ZarrRoot = ZarrGroup | ZarrArray
+
+
+@dataclass(frozen=True)
+class ZarrArrayMeta:
+    """``shape``/``chunks``/``dtype`` for a single Zarr v2 array."""
+
+    shape: tuple[int, ...]
+    chunks: tuple[int, ...]
+    dtype: str
+
+    @classmethod
+    def from_zarr_array(cls, arr: "ZarrArray") -> ZarrArrayMeta:
+        return cls(
+            shape=tuple(int(s) for s in arr.shape),
+            chunks=tuple(int(c) for c in arr.chunks),
+            dtype=str(arr.dtype),
+        )
+
+    @property
+    def rank(self) -> int:
+        return len(self.shape)
+
+    @property
+    def itemsize(self) -> int:
+        return np.dtype(self.dtype).itemsize
+
+    def effective_chunks(
+        self,
+        array_name: str,
+        user_chunk_shape: tuple[int, ...] | dict[str, tuple[int, ...]] | None,
+    ) -> tuple[int, ...]:
+        """Resolve the user's ``chunk_shapes`` override(s) against this array's chunks.
+
+        A single sequence overrides the leading axes (trailing axes keep the
+        native chunks), so one ``chunk_shapes=[16]`` applies across arrays of
+        different ranks. A dict maps array path → that array's override prefix;
+        arrays absent from it keep native chunks. ``None`` keeps native chunks;
+        an override longer than the array's rank raises ``ValueError``.
+        """
+        if user_chunk_shape is None:
+            return self.chunks
+
+        if isinstance(user_chunk_shape, dict):
+            user_chunk_shape = user_chunk_shape.get(array_name)
+            if user_chunk_shape is None:
+                return self.chunks
+        if len(user_chunk_shape) > self.rank:
+            raise ValueError(
+                f"chunk_shapes override for array {array_name!r} has "
+                f"{len(user_chunk_shape)} axes but array of shape "
+                f"{self.shape!r} has rank {self.rank}. Each chunk_shapes "
+                f"override may not be longer than its target array's rank."
+            )
+        return user_chunk_shape + self.chunks[len(user_chunk_shape) :]
+
+    def grid_shape(self, chunks: tuple[int, ...]) -> tuple[int, ...]:
+        """Number of chunks along each axis under the given chunk shape."""
+        return tuple(math.ceil(s / c) for s, c in zip(self.shape, chunks))
+
+    def chunk_slices(
+        self, chunk_index: tuple[int, ...], chunks: tuple[int, ...]
+    ) -> tuple[tuple[int, int], ...]:
+        """Per-axis ``(start, stop)`` for ``array[chunk_index]`` under ``chunks``.
+
+        Trailing-edge chunks are clamped to ``shape[i]``, so they may be
+        shorter than ``chunks[i]``. No padding is applied.
+        """
+        return tuple(
+            (i * c, min((i + 1) * c, s))
+            for i, c, s in zip(chunk_index, chunks, self.shape)
+        )
+
+
+# ---------------------------------------------------------------------------
+# Chunk reading
+# ---------------------------------------------------------------------------
+
+
+def _read_chunk(
+    root: ZarrRoot,
+    array_name: str,
+    chunk_slices: tuple[tuple[int, int], ...],
+    retry_match: Optional[List[str]] = None,
+) -> np.ndarray:
+    """Read ``array[chunk_slices]`` as an ndarray.
+
+    The underlying filesystem's own retry policy still applies underneath.
+    """
+
+    def _read() -> np.ndarray:
+        indexer = tuple(slice(s, e) for s, e in chunk_slices)
+        arr = root if array_name == "" else root[array_name]
+        return np.asarray(arr[indexer])
+
+    if not retry_match:
+        return _read()
+    # TODO(Artur): This would be more elegant with a general retry helper for non-iterables.
+    return next(
+        iterate_with_retry(
+            lambda: [_read()], description="read a Zarr chunk", match=retry_match
+        )
+    )
+
+
+@dataclass(frozen=True)
+class _ChunkRange:
+    """A contiguous slice ``[flat_start, flat_stop)`` of an array's chunk grid.
+
+    The flat indices address the row-major flattening of the chunk grid; the
+    read fn unravels each to an N-D ``chunk_index`` lazily on the worker. Keeping
+    a range (not a materialized per-chunk list) makes read-task planning
+    O(parallelism) rather than O(total chunks) -- important for stores with very
+    many chunks.
+    """
+
+    array_name: str
+    meta: ZarrArrayMeta
+    chunks: tuple[int, ...]
+    grid: tuple[int, ...]
+    flat_start: int
+    flat_stop: int
+
+
+@dataclass(frozen=True)
+class _AlignedChunkDescriptor:
+    """One wide row: a global axis-0 range ``[t_start, t_stop)`` across the
+    aligned arrays. With ``overlap > 0`` the row's data extends to
+    ``t_stop_data = min(t_stop + overlap, shape[0])`` (lookahead so windows
+    starting in this row reach their tail without crossing a row boundary).
+    """
+
+    chunk_index: int
+    t_start: int
+    t_stop: int
+    t_stop_data: int
+
+
+def _create_read_fn(
+    chunk_range: _ChunkRange,
+    root: ZarrRoot,
+    per_task_row_limit: Optional[int],
+    retry_match: Optional[List[str]],
+) -> Callable[[], Iterable[Block]]:
+    """Build a callable that materializes one block for a chunk-grid range.
+
+    This is the case where arrays are not aligned. Chunks are enumerated lazily
+    (on the worker) from ``chunk_range``. ``per_task_row_limit`` caps how many
+    chunks this task reads so a downstream ``limit`` reads only what it needs
+    (``None`` reads the whole range).
+    """
+    cr = chunk_range
+    stop = cr.flat_stop
+    if per_task_row_limit is not None:
+        stop = min(stop, cr.flat_start + per_task_row_limit)
+
+    def read_fn() -> Iterable[Block]:
+        builder = DelegatingBlockBuilder()
+        for flat_index in range(cr.flat_start, stop):
+            chunk_index = tuple(int(i) for i in np.unravel_index(flat_index, cr.grid))
+            chunk_slices = cr.meta.chunk_slices(chunk_index, cr.chunks)
+            builder.add(
+                {
+                    "array": cr.array_name,
+                    "chunk_index": chunk_index,
+                    "chunk_slices": chunk_slices,
+                    "chunk": _read_chunk(
+                        root, cr.array_name, chunk_slices, retry_match
+                    ),
+                }
+            )
+        yield builder.build()
+
+    return read_fn
+
+
+def _create_aligned_read_fn(
+    batch: list[_AlignedChunkDescriptor],
+    aligned_array_names: list[str],
+    root: ZarrRoot,
+    per_task_row_limit: Optional[int],
+    retry_match: Optional[List[str]],
+) -> Callable[[], Iterable[Block]]:
+    """Build a callable for aligned (wide-row) reads.
+
+    Each output row carries ``t_start``, ``t_stop``, and one column per
+    aligned array holding that array's ``[t_start:t_stop, ...]`` slice at
+    its natural shape (edge rows may be shorter). All arrays in one row
+    share the same axis-0 range.
+
+    This is the case where arrays are aligned on axis 0. ``per_task_row_limit``
+    caps how many rows this task reads (``None`` reads the whole batch).
+    """
+    batch = batch[:per_task_row_limit]
+
+    def read_fn() -> Iterable[Block]:
+        builder = DelegatingBlockBuilder()
+        for d in batch:
+            row: dict[str, Any] = {"t_start": d.t_start, "t_stop": d.t_stop}
+            for name in aligned_array_names:
+                row[name] = _read_chunk(
+                    root, name, ((d.t_start, d.t_stop_data),), retry_match
+                )
+            builder.add(row)
+        yield builder.build()
+
+    return read_fn
+
+
+def _is_positive_int(x) -> bool:
+    """True for a positive integer, including NumPy integers; False for bool."""
+    return not isinstance(x, bool) and isinstance(x, numbers.Integral) and int(x) > 0
+
+
+def _validate_chunk_shapes_dict(chunk_shapes: dict) -> dict[str, tuple[int, ...]]:
+    """Normalize chunk_shapes keys to store paths and validate their values."""
+    from zarr.util import normalize_storage_path
+
+    normalized: dict[str, tuple[int, ...]] = {}
+    for k, v in chunk_shapes.items():
+        if (
+            not isinstance(v, (tuple, list))
+            or not v
+            or not all(_is_positive_int(x) for x in v)
+        ):
+            raise ValueError(
+                f"chunk_shapes[{k!r}] must be a non-empty sequence of positive "
+                f"integers (list or tuple), got {v!r}"
+            )
+        normalized[normalize_storage_path(k)] = tuple(int(x) for x in v)
+    return normalized
+
+
+# ---------------------------------------------------------------------------
+# Datasource
+# ---------------------------------------------------------------------------
+
+
+class ZarrV2Datasource(Datasource):
+    """Reads one or more Zarr v2 arrays into a Ray Data ``Dataset``.
+
+    Emits long-form rows (one per chunk per array) or, with
+    ``align_axis_0=True``, wide rows (one per axis-0 chunk, one column per
+    array). See :func:`ray.data.read_zarr` for the row schemas and full API.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        filesystem: pyarrow_fs.FileSystem | AbstractFileSystem | None = None,
+        chunk_shapes: dict[str, list] | list | None = None,
+        array_paths: list[str] | None = None,
+        allow_full_metadata_scan: bool = False,
+        align_axis_0: bool = False,
+        overlap: int = 0,
+    ) -> None:
+        super().__init__()
+        _check_import(self, module="zarr", package="zarr")
+        import zarr
+
+        _check_import(self, module="fsspec", package="fsspec")
+        from fsspec.spec import AbstractFileSystem
+
+        if int(zarr.__version__.split(".")[0]) >= 3:
+            raise ImportError(
+                f"read_zarr supports zarr-python 2.x (Zarr v2 stores), but found "
+                f"zarr=={zarr.__version__}. Install a compatible version with "
+                f"`pip install 'zarr<3'`."
+            )
+
+        self.allow_full_metadata_scan = allow_full_metadata_scan
+        self.paths = [str(path)]
+        # ``local://`` stores live on the driver's local disk, so pin reads to
+        # the driver node (workers on other nodes can't see those files).
+        self._supports_distributed_reads = not _is_local_scheme(self.paths)
+
+        # Resolve filesystem + store path. The order of precedence:
+        #   1. Explicit ``filesystem=`` always wins.
+        #   2. ``.zip`` URL/path: auto-wrap with fsspec's ZipFileSystem.
+        #   3. Otherwise delegate to Ray Data's standard URL to filesystem
+        #      helper (the same one every other ``read_*`` API uses).
+        # "store path" is the path to the Zarr store, relative to the filesystem root.
+        # It is used to construct the Zarr root object.
+        if filesystem is None and self.paths[0].endswith(".zip"):
+            import fsspec
+
+            self._fs = fsspec.filesystem("zip", fo=self.paths[0])
+            self._store_path = ""
+        elif filesystem is None:
+            from fsspec.implementations.arrow import ArrowFSWrapper
+
+            from ray.data.datasource.path_util import (
+                _resolve_paths_and_filesystem,
+            )
+
+            resolved_paths, pa_fs = _resolve_paths_and_filesystem([self.paths[0]])
+            self._fs = ArrowFSWrapper(pa_fs)
+            self._store_path = resolved_paths[0].rstrip("/")
+        else:
+            from pyarrow.fs import FileSystem
+
+            if isinstance(filesystem, AbstractFileSystem):
+                self._fs = filesystem
+            elif isinstance(filesystem, FileSystem):
+                from fsspec.implementations.arrow import ArrowFSWrapper
+
+                self._fs = ArrowFSWrapper(filesystem)
+            else:
+                raise TypeError(
+                    f"filesystem must be pyarrow.fs.FileSystem or "
+                    f"fsspec.spec.AbstractFileSystem, got "
+                    f"{type(filesystem).__name__}"
+                )
+            from fsspec.implementations.zip import ZipFileSystem
+
+            if isinstance(self._fs, ZipFileSystem) and self.paths[0].endswith(".zip"):
+                # An explicit archive filesystem: the store is the archive root,
+                # not a ``.zip``-named entry inside it.
+                self._store_path = ""
+            else:
+                from fsspec.core import split_protocol
+
+                _, store_path = split_protocol(self.paths[0])
+                self._store_path = store_path.rstrip("/")
+
+        if chunk_shapes is not None and not isinstance(
+            chunk_shapes, (tuple, list, dict)
+        ):
+            raise ValueError(
+                f"chunk_shapes must be a non-empty sequence of positive "
+                f"integers (list or tuple), or a dict, got {chunk_shapes!r}"
+            )
+
+        self.chunk_shapes: tuple[int, ...] | dict[str, tuple[int, ...]] | None = None
+        if chunk_shapes is not None:
+            if isinstance(chunk_shapes, dict):
+                self.chunk_shapes = _validate_chunk_shapes_dict(chunk_shapes)
+            else:
+                if not chunk_shapes or not all(
+                    _is_positive_int(x) for x in chunk_shapes
+                ):
+                    raise ValueError(
+                        "chunk_shapes must be a non-empty sequence of positive integers "
+                        f"(list or tuple), got {chunk_shapes!r}"
+                    )
+
+                self.chunk_shapes = tuple(int(x) for x in chunk_shapes)
+
+        # Open the store with zarr (consolidated metadata when available).
+        # Detect consolidation by *trying* ``open_consolidated``.
+        store = self._fs.get_mapper(self._store_path)
+        try:
+            self.root = zarr.open_consolidated(store, mode="r")
+            self._consolidated = True
+        except KeyError:
+            self.root = zarr.open(store, mode="r")
+            self._consolidated = False
+
+        self._metadata_by_path = self._load_metadata(array_paths)
+        if not self._metadata_by_path:
+            raise ValueError(
+                f"No arrays discovered in Zarr store at {self.paths[0]!r}."
+            )
+
+        # Reject per-array overrides that do not correspond to any selected
+        # array in this read.
+        if isinstance(self.chunk_shapes, dict):
+            unknown_chunk_shape_keys = sorted(
+                set(self.chunk_shapes) - set(self._metadata_by_path)
+            )
+            if unknown_chunk_shape_keys:
+                raise ValueError(
+                    f"Unknown array path(s) in chunk_shapes: {unknown_chunk_shape_keys}"
+                )
+
+        if not align_axis_0:
+            self._aligned_array_names = None
+        else:
+            scalar_arrays = sorted(
+                name for name, meta in self._metadata_by_path.items() if not meta.shape
+            )
+            if scalar_arrays:
+                raise ValueError(
+                    f"align_axis_0=True requires every selected array to have "
+                    f"at least one axis, but these are 0-D (scalar): "
+                    f"{scalar_arrays}. Drop them with array_paths=[...]."
+                )
+            shape0_by_array = {
+                name: meta.shape[0] for name, meta in self._metadata_by_path.items()
+            }
+            if len(set(shape0_by_array.values())) > 1:
+                raise ValueError(
+                    f"All selected arrays must share shape[0] when "
+                    f"align_axis_0=True. Got: {shape0_by_array}. Pass a "
+                    f"shape-compatible subset via array_paths=[...]."
+                )
+            self._aligned_array_names = list(self._metadata_by_path.keys())
+
+        # Validate overlap. Only meaningful when arrays are co-iterated as
+        # wide rows, since the trailing lookahead is exposed via the
+        # per-array column being longer than ``t_stop - t_start``.
+        if not isinstance(overlap, int) or overlap < 0:
+            raise ValueError(f"overlap must be a non-negative integer, got {overlap!r}")
+        if overlap and self._aligned_array_names is None:
+            raise ValueError(
+                "overlap requires align_axis_0=True. In the default long-form "
+                "(chunk-per-row) mode, there's no wide row to extend forward — "
+                "the ``chunk_slices`` column on each chunk row already exposes "
+                "the global axis-0 range."
+            )
+        self.overlap = overlap
+
+        # Resolve per-array chunk geometry. ``effective_chunks`` raises a
+        # ``ValueError`` if a shared ``chunk_shapes`` prefix or any per-array
+        # ``chunk_shapes`` override is longer than the target array's rank —
+        # so this loop is also where rank validation happens.
+        self._array_chunks: dict[str, tuple[int, ...]] = {}
+        self._array_grids: dict[str, tuple[int, ...]] = {}
+        for name, meta in self._metadata_by_path.items():
+            chunks = meta.effective_chunks(name, self.chunk_shapes)
+            self._array_chunks[name] = chunks
+            self._array_grids[name] = meta.grid_shape(chunks)
+
+        # If aligned, all listed arrays must share the same axis-0 chunk size
+        # so each wide row corresponds to one axis-0 step across every array.
+        if self._aligned_array_names is not None:
+            axis_0_chunks = {
+                name: self._array_chunks[name][0] for name in self._aligned_array_names
+            }
+            unique = set(axis_0_chunks.values())
+            if len(unique) > 1:
+                raise ValueError(
+                    f"Aligned arrays must share the same axis-0 chunk size. "
+                    f"Got: {axis_0_chunks}. Pass chunk_shapes=[N] (or a "
+                    f"per-array chunk_shapes dict that resolves all aligned "
+                    f"arrays to the same axis-0 prefix) to re-tile them."
+                )
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        return self._supports_distributed_reads
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        """Total bytes = sum over selected arrays of ``prod(shape) * itemsize``."""
+        return sum(
+            math.prod(meta.shape) * meta.itemsize
+            for meta in self._metadata_by_path.values()
+        )
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        """Enumerate every chunk and wrap it (or batches of chunks) in ReadTasks."""
+        from ray.data.context import DataContext
+
+        retry_match = (data_context or DataContext.get_current()).retried_io_errors
+        if self._aligned_array_names is not None:
+            return self._get_aligned_read_tasks(
+                parallelism, per_task_row_limit, retry_match
+            )
+        return self._get_long_form_read_tasks(
+            parallelism, per_task_row_limit, retry_match
+        )
+
+    def _get_long_form_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int],
+        retry_match: Optional[List[str]],
+    ) -> List[ReadTask]:
+        read_tasks: List[ReadTask] = []
+        for name, meta in self._metadata_by_path.items():
+            chunks = self._array_chunks[name]
+            grid = self._array_grids[name]
+            n_chunks = math.prod(grid)
+            if n_chunks == 0:
+                continue
+            # Split the chunk grid into contiguous flat-index ranges. This is
+            # O(n_tasks), not O(n_chunks): we never materialize a per-chunk list
+            # on the driver -- the read fn unravels chunks lazily on the worker.
+            n_tasks = max(1, min(parallelism, n_chunks))
+            batch_size = math.ceil(n_chunks / n_tasks)
+            for flat_start in range(0, n_chunks, batch_size):
+                flat_stop = min(flat_start + batch_size, n_chunks)
+                chunk_range = _ChunkRange(
+                    name, meta, chunks, grid, flat_start, flat_stop
+                )
+                read_tasks.append(
+                    ReadTask(
+                        _create_read_fn(
+                            chunk_range, self.root, per_task_row_limit, retry_match
+                        ),
+                        BlockMetadata(
+                            num_rows=flat_stop - flat_start,
+                            size_bytes=self._estimate_range_mem_size(chunk_range),
+                            input_files=(self.paths[0],),
+                            exec_stats=None,
+                        ),
+                        per_task_row_limit=per_task_row_limit,
+                    )
+                )
+        return read_tasks
+
+    def _estimate_range_mem_size(self, chunk_range: _ChunkRange) -> int:
+        """Upper-bound in-memory bytes for a chunk-grid range.
+
+        Assumes a full-size chunk per index; trailing-edge chunks are smaller,
+        so this slightly over-estimates. O(1) -- it does not enumerate the range.
+        """
+        n = chunk_range.flat_stop - chunk_range.flat_start
+        return n * math.prod(chunk_range.chunks) * chunk_range.meta.itemsize
+
+    def _get_aligned_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int],
+        retry_match: Optional[List[str]],
+    ) -> List[ReadTask]:
+        """Aligned read tasks. See :meth:`get_read_tasks` for semantics."""
+        assert self._aligned_array_names is not None
+        # All aligned arrays share the same axis-0 chunk size (validated in
+        # ``__init__``) and the same shape[0]. Read the geometry off the first.
+        first_name = self._aligned_array_names[0]
+        axis_0_chunk = self._array_chunks[first_name][0]
+        shape0 = self._metadata_by_path[first_name].shape[0]
+
+        descriptors = [
+            _AlignedChunkDescriptor(
+                chunk_index=i,
+                t_start=i * axis_0_chunk,
+                t_stop=min((i + 1) * axis_0_chunk, shape0),
+                t_stop_data=min((i + 1) * axis_0_chunk + self.overlap, shape0),
+            )
+            for i in range(math.ceil(shape0 / axis_0_chunk))
+        ]
+        if not descriptors:
+            return []
+
+        n_tasks = max(1, min(parallelism, len(descriptors)))
+        batch_size = math.ceil(len(descriptors) / n_tasks)
+
+        read_tasks: List[ReadTask] = []
+        for start in range(0, len(descriptors), batch_size):
+            batch = descriptors[start : start + batch_size]
+            read_tasks.append(
+                ReadTask(
+                    _create_aligned_read_fn(
+                        batch,
+                        self._aligned_array_names,
+                        self.root,
+                        per_task_row_limit,
+                        retry_match,
+                    ),
+                    BlockMetadata(
+                        num_rows=len(batch),
+                        size_bytes=self._estimate_aligned_batch_mem_size(batch),
+                        input_files=(self.paths[0],),
+                        exec_stats=None,
+                    ),
+                    per_task_row_limit=per_task_row_limit,
+                )
+            )
+        return read_tasks
+
+    def _estimate_aligned_batch_mem_size(
+        self, batch: list[_AlignedChunkDescriptor]
+    ) -> int:
+        """Sum bytes across all (row, aligned-array) pairs in a wide-row batch.
+
+        Accounts for the trailing overlap data each row carries: the row's
+        per-array slice covers ``[t_start, t_stop_data)``, not just
+        ``[t_start, t_stop)``.
+        """
+        assert self._aligned_array_names is not None
+        return sum(
+            (desc.t_stop_data - desc.t_start)
+            * (math.prod(meta.shape[1:]) if len(meta.shape) > 1 else 1)
+            * meta.itemsize
+            for desc in batch
+            for meta in (
+                self._metadata_by_path[name] for name in self._aligned_array_names
+            )
+        )
+
+    def _load_metadata(self, array_paths) -> dict[str, ZarrArrayMeta]:
+        """Read ``shape``/``chunks``/``dtype`` for the selected arrays off ``self.root``.
+
+        zarr validated the store's metadata when it was opened, so this only
+        adapts the resulting ``zarr.Array`` objects. Discovery uses consolidated
+        metadata when present, then explicit ``array_paths``, then an optional
+        full scan (``allow_full_metadata_scan``). If ``array_paths`` is given,
+        the discovered set is filtered down to it.
+        """
+        import zarr
+        from zarr.util import normalize_storage_path
+
+        root = self.root
+        requested = (
+            {normalize_storage_path(p) for p in array_paths} if array_paths else None
+        )
+
+        if isinstance(root, zarr.Array):
+            # A store that is itself an array exposes exactly one path: "" (root).
+            # Reject any requested path that isn't the root so a bad ``array_paths``
+            # fails loudly here instead of silently returning the root array.
+            if requested is not None and requested != {""}:
+                raise ValueError(
+                    f"This Zarr store is a single root-level array (path ''), "
+                    f"but array_paths={array_paths!r} requested other path(s). "
+                    f"Pass array_paths=[''] or omit it."
+                )
+            return {"": ZarrArrayMeta.from_zarr_array(root)}
+
+        if not self._consolidated and not self.allow_full_metadata_scan:
+            if requested is None:
+                raise ValueError(
+                    "No array_paths were provided and this Zarr store does not "
+                    "contain .zmetadata. Pass array_paths=[...] or set "
+                    "allow_full_metadata_scan=True."
+                )
+            out: dict[str, ZarrArrayMeta] = {}
+            for raw in array_paths:
+                name = normalize_storage_path(raw)
+                try:
+                    arr = root[name]
+                except KeyError as e:
+                    raise ValueError(
+                        f"Array path {raw!r} not found in Zarr store."
+                    ) from e
+                if not isinstance(arr, zarr.Array):
+                    raise ValueError(f"Array path {raw!r} is a group, not an array.")
+                out[name] = ZarrArrayMeta.from_zarr_array(arr)
+            return out
+
+        all_arrays: dict[str, ZarrArrayMeta] = {}
+
+        def _collect(name: str, obj) -> None:
+            if isinstance(obj, zarr.Array):
+                all_arrays[name] = ZarrArrayMeta.from_zarr_array(obj)
+
+        root.visititems(_collect)
+
+        if requested is not None:
+            missing = sorted(requested - all_arrays.keys())
+            if missing:
+                raise ValueError(
+                    f"Array(s) not found: {', '.join(repr(m) for m in missing)}. "
+                    f"Available: {', '.join(repr(a) for a in sorted(all_arrays))}"
+                )
+            all_arrays = {k: v for k, v in all_arrays.items() if k in requested}
+
+        return all_arrays

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

@@ -0,0 +1,3 @@
+from typing import TypeVar
+
+InputSplit = TypeVar("InputSplit")

python/ray/data/_internal/datasource_v2/chunkers/file_chunker.py

@@ -0,0 +1,286 @@
+"""File chunkers for DataSourceV2.
+
+A ``FileChunker`` decides how a single listed file is split into one or
+more parallel-read units. The indexer drives the chunker once per file
+and emits one manifest row per chunk; downstream the partitioner /
+reader carry the per-chunk metadata through to the read task.
+"""
+
+import abc
+import logging
+import math
+from typing import (
+    TYPE_CHECKING,
+    Iterable,
+    Optional,
+    Tuple,
+    Type,
+    TypedDict,
+    TypeVar,
+    cast,
+    get_type_hints,
+)
+
+from ray.data._internal.util import MiB, infer_compression
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from pyarrow.fs import FileSystem
+
+logger = logging.getLogger(__name__)
+
+
+class ChunkMetadata(TypedDict):
+    """Base interface for chunk metadata types."""
+
+    pass
+
+
+_ChunkMetadataT = TypeVar("_ChunkMetadataT", bound=ChunkMetadata)
+
+
+def create_chunk_metadata(cls: Type[_ChunkMetadataT], **kwargs) -> _ChunkMetadataT:
+    """Create a metadata instance with validation, ensure the keys are correct."""
+    required_keys = list(get_type_hints(cls).keys())
+
+    missing_keys = [key for key in required_keys if key not in kwargs]
+    if missing_keys:
+        raise ValueError(f"Missing required keys: {missing_keys}")
+
+    extra_keys = [key for key in kwargs if key not in required_keys]
+    if extra_keys:
+        raise ValueError(f"Unexpected keys: {extra_keys}")
+
+    return cast(_ChunkMetadataT, kwargs)
+
+
+class LineDelimitedFileChunkMetadata(ChunkMetadata):
+    """Metadata for line-delimited file chunks."""
+
+    chunk_byte_start_idx: int
+    chunk_byte_end_idx: int
+
+
+class ParquetFileChunkMetadata(ChunkMetadata):
+    """Metadata for Parquet file chunks.
+
+    A chunk is an explicit, half-open range of consecutive row groups
+    ``[row_group_start, row_group_end)`` within a single file, computed at
+    listing time from the file's footer. The reader slices the fragment to
+    exactly this range — no estimation or read-time reconciliation.
+    """
+
+    row_group_start: int  # inclusive
+    row_group_end: int  # exclusive
+
+
+@DeveloperAPI
+class FileChunker(abc.ABC):
+    """Abstract base class for chunking files into smaller pieces for parallel processing.
+
+    File chunkers determine how large files should be split into chunks that can be
+    processed in parallel. Different file formats may require different chunking strategies.
+
+    For example:
+    - Line-delimited files (JSONL, CSV) can be chunked by byte ranges
+    - Parquet files can be chunked by row groups
+    """
+
+    # Whether ``generate_chunk_metadatas`` performs file I/O (e.g. reading a
+    # Parquet footer). When True, the indexer fans chunking across its thread
+    # pool so the per-file reads parallelize even for a single input
+    # directory. When False, the indexer chunks inline (no thread hand-off).
+    reads_file_metadata: bool = False
+
+    @abc.abstractmethod
+    def generate_chunk_metadatas(
+        self,
+        path: str,
+        file_size: int,
+        filesystem: Optional["FileSystem"] = None,
+    ) -> Iterable[Tuple[Optional[ChunkMetadata], int]]:
+        """Generate metadata for file chunks.
+
+        Args:
+            path: The file path being chunked.
+            file_size: The total size in bytes of the file to be chunked.
+            filesystem: PyArrow filesystem used to read per-file metadata
+                (e.g. the Parquet footer). Ignored by chunkers that do not
+                read file metadata.
+
+        Returns:
+            An iterable of tuples containing (metadata, chunk_size) where metadata
+            describes the chunk and chunk_size is the size of the chunk in bytes.
+            Metadata can be None for chunks that don't require metadata
+            (e.g., whole file processing).
+        """
+        ...
+
+
+@DeveloperAPI
+class WholeFileChunker(FileChunker):
+    """File chunker that treats the whole file as a single chunk.
+
+    This chunker is used when files should be processed as a single unit,
+    typically for smaller files or when the file format doesn't support
+    efficient chunking (e.g., compressed files).
+
+    Yields a single chunk with no metadata, indicating the entire file
+    should be processed as one unit.
+    """
+
+    def generate_chunk_metadatas(
+        self,
+        path: str,
+        file_size: int,
+        filesystem: Optional["FileSystem"] = None,
+    ) -> Iterable[Tuple[Optional[ChunkMetadata], int]]:
+        yield None, file_size
+
+
+@DeveloperAPI
+class LineDelimitedFileChunker(FileChunker):
+    """File chunker for line-delimited files (JSONL, CSV, TSV, etc.).
+
+    This chunker splits files into fixed-size byte chunks (default: 256 MiB)
+    and provides metadata about the byte ranges for each chunk. The actual
+    line boundaries are handled by the reader to ensure complete records.
+    """
+
+    _CHUNK_BYTE_SIZE = 256 * MiB  # 256 MiB
+
+    def generate_chunk_metadatas(
+        self,
+        path: str,
+        file_size: int,
+        filesystem: Optional["FileSystem"] = None,
+    ) -> Iterable[Tuple[Optional[ChunkMetadata], int]]:
+        compression = infer_compression(path)
+        if compression is not None:
+            yield None, file_size
+        else:
+            num_chunks = math.ceil(file_size / self._CHUNK_BYTE_SIZE)
+            for chunk_idx in range(num_chunks):
+                chunk_start = self._CHUNK_BYTE_SIZE * chunk_idx
+                chunk_end = min(self._CHUNK_BYTE_SIZE * (chunk_idx + 1), file_size)
+                chunk_size = chunk_end - chunk_start
+                yield (
+                    create_chunk_metadata(
+                        LineDelimitedFileChunkMetadata,
+                        chunk_byte_start_idx=chunk_start,
+                        chunk_byte_end_idx=chunk_end,
+                    ),
+                    chunk_size,
+                )
+
+
+@DeveloperAPI
+class ParquetFileChunker(FileChunker):
+    """File chunker for Parquet files.
+
+    Reads each file's footer at listing time and chunks on **true row-group
+    boundaries**: consecutive row groups are bundled into a chunk until the
+    bundle's on-disk size reaches ``target_chunk_size`` (always at least one
+    row group per chunk). Each chunk carries an explicit half-open row-group
+    range, so the reader slices to exactly those row groups with no
+    estimation or read-time reconciliation, and the listing stage never
+    produces empty read tasks.
+
+    The row group is Parquet's atomic read unit, so a chunk can never be
+    smaller than a single row group. With the default target (which falls
+    back to ``DataContext.target_min_block_size``), a file's row groups map
+    1:1 to chunks unless they are smaller than the target, in which case
+    consecutive small row groups are bundled to avoid an excessive number of
+    tiny chunks.
+    """
+
+    # Hard fallback used only when neither an explicit target nor the
+    # DataContext size knobs are set.
+    _FALLBACK_TARGET_CHUNK_SIZE = 1 * MiB
+
+    # Footer reads are file I/O — let the indexer parallelize them.
+    reads_file_metadata: bool = True
+
+    def __init__(self, target_chunk_size: Optional[int] = None):
+        from ray.data.context import DataContext
+
+        ctx = DataContext.get_current()
+        # Resolve with explicit ``is not None`` checks rather than ``or`` so an
+        # explicit ``0`` (e.g. to force one row group per chunk) isn't treated as
+        # "unset" and silently overridden by a falsy-coalescing fallback.
+        if target_chunk_size is not None:
+            self._target_chunk_size = target_chunk_size
+        elif ctx.parquet_chunker_target_chunk_size is not None:
+            self._target_chunk_size = ctx.parquet_chunker_target_chunk_size
+        elif ctx.target_min_block_size is not None:
+            self._target_chunk_size = ctx.target_min_block_size
+        else:
+            self._target_chunk_size = self._FALLBACK_TARGET_CHUNK_SIZE
+
+    def generate_chunk_metadatas(
+        self,
+        path: str,
+        file_size: int,
+        filesystem: Optional["FileSystem"] = None,
+    ) -> Iterable[Tuple[Optional[ChunkMetadata], int]]:
+        import pyarrow.parquet as pq
+
+        try:
+            # Reads only the Parquet footer (file metadata), not data.
+            metadata = pq.read_metadata(path, filesystem=filesystem)
+        except Exception as e:
+            # Corrupt / unreadable footer (or a non-Parquet file that slipped
+            # through). Fall back to a single whole-file chunk so the file is
+            # still read rather than dropped.
+            logger.debug(
+                "Could not read Parquet footer for chunking (%s): %s; "
+                "falling back to a whole-file chunk.",
+                path,
+                e,
+            )
+            yield None, file_size
+            return
+
+        num_row_groups = metadata.num_row_groups
+        if num_row_groups == 0:
+            yield None, file_size
+            return
+
+        # Greedily bundle consecutive row groups until the running on-disk
+        # size reaches the target. Always emit at least one row group per
+        # chunk (the atomic read unit).
+        start = 0
+        running_size = 0
+        for rg_idx in range(num_row_groups):
+            rg_meta = metadata.row_group(rg_idx)
+            # On-disk (compressed) row-group size. ``RowGroupMetaData`` exposes
+            # only the *uncompressed* ``total_byte_size``; the on-disk size lives
+            # on each ``ColumnChunkMetaData``, so sum the per-column compressed
+            # sizes. Keeping chunk sizes in on-disk units matches the manifest's
+            # ``file_sizes`` and the ``×encoding_ratio`` in-memory estimator.
+            rg_size = sum(
+                rg_meta.column(c).total_compressed_size
+                for c in range(rg_meta.num_columns)
+            )
+            if running_size > 0 and running_size + rg_size > self._target_chunk_size:
+                yield (
+                    create_chunk_metadata(
+                        ParquetFileChunkMetadata,
+                        row_group_start=start,
+                        row_group_end=rg_idx,
+                    ),
+                    running_size,
+                )
+                start = rg_idx
+                running_size = 0
+            running_size += rg_size
+
+        # Flush the final bundle.
+        yield (
+            create_chunk_metadata(
+                ParquetFileChunkMetadata,
+                row_group_start=start,
+                row_group_end=num_row_groups,
+            ),
+            running_size,
+        )

python/ray/data/_internal/datasource_v2/chunkers/parquet_file_chunking_utils.py

@@ -0,0 +1,49 @@
+"""Parquet file-level chunking helpers for DataSourceV2.
+
+Maps planner chunk metadata (``ParquetFileChunkMetadata``) to PyArrow
+``ParquetFileFragment`` subsets for parallel reads. Chunk metadata carries
+an explicit half-open row-group range computed at listing time from the
+file's footer, so no estimation or reconciliation is needed here.
+"""
+from typing import List, Tuple
+
+import pyarrow.dataset as pds
+
+from ray.data._internal.datasource_v2.chunkers.file_chunker import (
+    ParquetFileChunkMetadata,
+)
+
+
+def _fragments_from_chunk_metadata(
+    fragment: pds.ParquetFileFragment,
+    chunk_metadata: ParquetFileChunkMetadata,
+) -> List[Tuple[pds.ParquetFileFragment, int]]:
+    """Slice ``fragment`` into per-row-group sub-fragments for the chunk's range.
+
+    The chunk carries an explicit ``[row_group_start, row_group_end)`` range.
+    Returns one ``(ParquetFileFragment, file_row_offset)`` pair per row group
+    in that range, where ``file_row_offset`` is the sum of ``num_rows`` across
+    all row groups that precede the sub-fragment in the underlying file.
+    Callers seed per-fragment hashing offsets with this value so sub-fragments
+    of the same file don't collide on ``(path, 0, n)``.
+
+    The range is defensively clamped to the file's actual row-group count;
+    since ranges are computed from the same footer the reader sees, the clamp
+    is a no-op in practice and never drops real row groups.
+    """
+    start = chunk_metadata["row_group_start"]
+    end = chunk_metadata["row_group_end"]
+    metadata = fragment.metadata
+    total_row_groups = metadata.num_row_groups
+
+    start = min(start, total_row_groups)
+    end = min(end, total_row_groups)
+
+    file_row_offset = sum(metadata.row_group(i).num_rows for i in range(start))
+    sub_fragments: List[Tuple[pds.ParquetFileFragment, int]] = []
+    for row_group_index in range(start, end):
+        sub_fragments.append(
+            (fragment.subset(row_group_ids=[row_group_index]), file_row_offset)
+        )
+        file_row_offset += metadata.row_group(row_group_index).num_rows
+    return sub_fragments

python/ray/data/_internal/datasource_v2/datasource_v2.py

@@ -0,0 +1,191 @@
+"""
+DataSourceV2 API - Unified Abstraction for Reading Data Sources
+
+This module defines a unified, extensible API for reading data from diverse sources
+in Ray Data. The API provides a common abstraction layer that enables datasources to
+declaratively expose their capabilities—such as filter pushdown, projection pruning,
+and parallel reads—while allowing the execution engine to leverage these capabilities
+transparently.
+
+Core Principles:
+- Modularity: Separate concerns (indexing, scanning, reading)
+- Expressivity: Declarative capability exposure via mixins
+- Extensibility: Easy to add new datasources with custom optimizations
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Optional,
+)
+
+import pyarrow as pa
+
+from ray.data._internal.datasource_v2 import InputSplit
+from ray.data._internal.datasource_v2.listing.file_indexer import FileIndexer
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from pyarrow.fs import FileSystem
+
+    from ray.data._internal.datasource_v2.readers.in_memory_size_estimator import (
+        InMemorySizeEstimator,
+    )
+    from ray.data._internal.datasource_v2.scanners.scanner import Scanner
+
+
+@DeveloperAPI
+class DatasourceCategory(Enum):
+    """Categories of datasources with different capability profiles.
+
+    Each category has a distinct set of applicable optimizations:
+    - FILE_BASED: Local/cloud files (parquet, csv, json, images)
+    - DATABASE: SQL databases (postgres, mysql, snowflake)
+    - DATA_LAKE: Table formats (iceberg, delta, hudi)
+    - IN_MEMORY: In-process data (pandas, numpy, arrow)
+    - SYNTHETIC: Generated data (range, range_tensor)
+    - STREAMING: Unbounded sources (kafka, kinesis)
+    """
+
+    FILE_BASED = "file_based"
+    DATABASE = "database"
+    DATA_LAKE = "data_lake"
+    IN_MEMORY = "in_memory"
+    SYNTHETIC = "synthetic"
+    STREAMING = "streaming"
+
+
+@DeveloperAPI
+class DataSourceV2(ABC, Generic[InputSplit]):
+    """Abstract base class for V2 datasources.
+
+    DataSourceV2 is the entry point for reading data from a source. It provides:
+    1. File listing (for file-based sources) - via _get_file_indexer()
+    2. Schema inference
+    3. Size estimation
+    4. Scanner creation
+
+    Subclasses should implement the abstract methods and can optionally
+    override _get_file_indexer() and get_size_estimator() for file-based sources.
+
+    Example::
+
+        datasource = ParquetDatasourceV2()
+        indexer = datasource._get_file_indexer()
+        # List files with optional sampling
+        for manifest in indexer.list_files(paths, filesystem=fs):
+            schema = datasource.infer_schema(manifest)
+            break  # Just need first manifest for schema
+        scanner = datasource.create_scanner(schema)
+        scanner = scanner.prune_columns(["col1", "col2"])
+        reader = scanner.create_reader()
+        for table in reader.read(manifest):
+            process(table)
+    """
+
+    def __init__(self, name: str, category: DatasourceCategory):
+        """Initialize the datasource.
+
+        Args:
+            name: Human-readable name for this datasource.
+            category: Category of this datasource.
+        """
+        self._name = name
+        self._category = category
+        # File-based subclasses set this to ``False`` in their ``__init__``
+        # when the user-supplied paths are in the ``local://`` scheme —
+        # the driver node is the only one that can read those files.
+        # ``_read_datasource_v2`` consults the flag to decide whether to
+        # pin read tasks via a ``label_selector``.
+        self._supports_distributed_reads: bool = True
+
+    @property
+    def name(self) -> str:
+        """Human-readable name for this datasource."""
+        return self._name
+
+    @property
+    def category(self) -> DatasourceCategory:
+        """Category of this datasource."""
+        return self._category
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        """Whether read tasks may run on any cluster node.
+
+        Defaults to ``True``. File-based subclasses (e.g.
+        :class:`ParquetDatasourceV2`) flip this to ``False`` when the
+        user supplies ``local://``-scheme paths so ``_read_datasource_v2``
+        can pin reads to the driver node via a ``ray.io/node-id``
+        label selector. Mirrors V1 ``Datasource.supports_distributed_reads``.
+        """
+        return self._supports_distributed_reads
+
+    def _get_file_indexer(self) -> Optional[FileIndexer]:
+        """Return FileIndexer component if applicable.
+
+        Override this for file-based datasources to provide file discovery.
+
+        Returns:
+            FileIndexer instance, or None for non-file-based sources.
+        """
+        return None
+
+    def get_size_estimator(self) -> Optional[InMemorySizeEstimator]:
+        """Return size estimator for this datasource.
+
+        Override this to provide format-specific size estimation.
+
+        Returns:
+            InMemorySizeEstimator instance, or None if not supported.
+        """
+        return None
+
+    @abstractmethod
+    def infer_schema(self, sample: InputSplit) -> pa.Schema:
+        """Infer schema from a sample of data.
+
+        Args:
+            sample: Sample data to infer schema from.
+
+        Returns:
+            PyArrow Schema inferred from the sample.
+
+        Raises:
+            ValueError: If schema cannot be inferred from the sample.
+        """
+        ...
+
+    @abstractmethod
+    def create_scanner(
+        self,
+        schema: pa.Schema,
+        filesystem: Optional["FileSystem"] = None,
+        **options: Any,
+    ) -> Scanner[InputSplit]:
+        """Create a Scanner for reading data.
+
+        Args:
+            schema: Schema for the data to read.
+            filesystem: Optional filesystem for file-based sources.
+            **options: Additional datasource-specific options.
+
+        Returns:
+            Configured Scanner instance.
+        """
+        ...
+
+    def resolve_partitioning(self, sample: InputSplit) -> Optional[Any]:
+        """Return a partitioning descriptor derived from ``sample``, or ``None``.
+
+        Override this for file-based sources whose partition keys must be
+        discovered from a sample path (e.g. hive layouts where field names
+        are not known up front). The resolved descriptor is passed into
+        :meth:`create_scanner`.
+        """
+        return None

python/ray/data/_internal/datasource_v2/listing/file_indexer.py

@@ -0,0 +1,376 @@
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union
+
+from pyarrow.fs import FileSystem
+
+from ray._common.utils import env_integer
+from ray.data._internal.datasource_v2.chunkers.file_chunker import (
+    ChunkMetadata,
+    FileChunker,
+    WholeFileChunker,
+)
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.listing.file_pruners import FilePruner
+from ray.data._internal.datasource_v2.listing.indexing_utils import (
+    _get_file_infos,
+    _get_path_contents,
+)
+from ray.data._internal.dynamic_work_queue import parallel_process_work_stealing
+from ray.data._internal.util import make_async_gen
+from ray.data.block import BlockColumn
+from ray.data.datasource.path_util import _resolve_paths_and_filesystem
+
+logger = logging.getLogger(__name__)
+
+
+class FileIndexer(ABC):
+    @property
+    @abstractmethod
+    def file_chunker(self) -> FileChunker:
+        """The file chunker that this indexer uses."""
+        ...
+
+    @abstractmethod
+    def list_files(
+        self,
+        paths: "BlockColumn",
+        *,
+        filesystem: "FileSystem",
+        pruners: Optional[List[FilePruner]] = None,
+        preserve_order: bool = False,
+    ) -> Iterable[FileManifest]:
+        """List files and their on-disk sizes for the given path.
+
+        Args:
+            paths: A column of paths pointing to files or directories.
+            filesystem: A PyArrow filesystem object.
+            pruners: A list of file pruners to apply.
+            preserve_order: Whether to preserve order in file listing.
+
+        Returns:
+            An iterator of `FileManifest` objects, each of which contains a file path
+            and the on-disk size of the file in bytes.
+        """
+        ...
+
+
+@dataclass(frozen=True)
+class FileInfo:
+    """File information for file listing."""
+
+    path: str
+    size: Optional[int]
+
+
+@dataclass(frozen=True)
+class _TraversalWorkItem:
+    """Work item for parallel directory traversal. Distinguishes seed paths
+    (user-provided, need resolution) from subdir paths (from filesystem listing,
+    use directly to avoid redundant resolution that breaks non-local
+    filesystems)."""
+
+    # Could be a file path or a directory path.
+    path: str
+    # True for subdirectories discovered during traversal; False for seed input paths.
+    is_discovered_subdir: bool = False
+    # Original seed-path index used to restore deterministic ordering when requested.
+    input_path_index: Optional[int] = None
+    # Top-level path the traversal started from, used to scope hidden-prefix
+    # exclusion to entries whose path relative to the root is hidden.
+    root_path: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class OrderedFileResult:
+    """File result with order information for sorting when preserve_order is True."""
+
+    input_path_index: int
+    # The leaf file path.
+    file_path: str
+    file_info: FileInfo
+
+
+class NonSamplingFileIndexer(FileIndexer):
+    """A file indexer that exhaustively lists files.
+
+    This implementation works with paths that point to files or directories,
+    although it's slow if you try to list lots of paths pointing to files
+    rather than a single directory.
+    """
+
+    _DEFAULT_MAX_PATHS_PER_OUTPUT = env_integer(
+        "RAY_DATA_MAX_PATHS_PER_LIST_FILES_OUTPUT", 1000
+    )
+
+    _DEFAULT_NUM_WORKERS = env_integer("RAY_DATA_LIST_FILES_THREADED_NUM_WORKERS", 8)
+
+    def __init__(
+        self,
+        *,
+        ignore_missing_paths: bool,
+        num_workers: Optional[int] = None,
+        max_paths_per_output: Optional[int] = None,
+        file_chunker: Optional[FileChunker] = None,
+    ):
+        self._ignore_missing_paths = ignore_missing_paths
+        self._max_paths_per_output = (
+            max_paths_per_output
+            if max_paths_per_output is not None
+            else self._DEFAULT_MAX_PATHS_PER_OUTPUT
+        )
+        self._num_workers = (
+            num_workers if num_workers is not None else self._DEFAULT_NUM_WORKERS
+        )
+        self._queue_size_per_thread = env_integer(
+            "RAY_DATA_LIST_FILES_QUEUE_SIZE_PER_THREAD",
+            self._max_paths_per_output * 4,
+        )
+        self._file_chunker: FileChunker = (
+            file_chunker if file_chunker is not None else WholeFileChunker()
+        )
+
+    @property
+    def file_chunker(self) -> FileChunker:
+        """The file chunker that this indexer uses.
+
+        Exposed primarily for tests and shuffle-aware planning code that needs
+        to introspect or override the chunking strategy.
+        """
+        return self._file_chunker
+
+    def list_files(
+        self,
+        paths: "BlockColumn",
+        *,
+        filesystem: "FileSystem",
+        pruners: Optional[List[FilePruner]] = None,
+        preserve_order: bool = False,
+    ) -> Iterable[FileManifest]:
+        file_info_iterator = (
+            self._get_file_info_iterator_threaded(paths, filesystem, preserve_order)
+            if self._num_workers > 1
+            else self._get_file_info_iterator_sequential(paths, filesystem)
+        )
+
+        # Stage pipeline: list → prune (cheap, inline) → chunk (may read
+        # per-file metadata) → batch into manifests. Pruning runs *before*
+        # chunking so we never read a footer for a file we'd discard.
+        pruned = self._filter_file_infos(file_info_iterator, pruners or [])
+        chunk_records = self._generate_chunk_records(pruned, filesystem, preserve_order)
+        yield from self._batch_chunk_records_to_manifests(chunk_records)
+
+    def _get_file_info_iterator_sequential(
+        self,
+        paths: "BlockColumn",
+        filesystem: "FileSystem",
+    ) -> Iterable[FileInfo]:
+        for input_path in paths.to_pylist():
+            resolved_paths, _ = _resolve_paths_and_filesystem(input_path, filesystem)
+            assert len(resolved_paths) == 1
+
+            for path, file_size in _get_file_infos(
+                resolved_paths[0], filesystem, self._ignore_missing_paths
+            ):
+                yield FileInfo(path=path, size=file_size)
+
+    def _get_file_info_iterator_threaded(
+        self,
+        paths: "BlockColumn",
+        filesystem: "FileSystem",
+        preserve_order: bool = False,
+    ) -> Iterable[FileInfo]:
+        """Threaded file info iterator with work stealing for parallel directory
+        traversal. Subdirectories are added as work items for idle workers to
+        process."""
+        paths_list = paths.to_pylist()
+        if len(paths_list) == 0:
+            return
+
+        num_workers = self._num_workers
+
+        seed_items = [
+            _TraversalWorkItem(
+                path=p,
+                is_discovered_subdir=False,
+                input_path_index=i if preserve_order else None,
+            )
+            for i, p in enumerate(paths_list)
+        ]
+
+        def process_fn(
+            item: _TraversalWorkItem,
+            add_work: Callable[[_TraversalWorkItem], None],
+            add_result: Callable[[Union[OrderedFileResult, FileInfo]], None],
+        ) -> None:
+            """Process a single item, adding discovered subdirs as work and
+            files as results."""
+            input_path_index = item.input_path_index
+
+            if item.is_discovered_subdir:
+                # Subdir paths from filesystem listing: use directly. Re-resolution
+                # would infer LocalFileSystem for scheme-less paths on S3/GCS,
+                # and add redundant overhead.
+                path = item.path
+                root_path = item.root_path
+            else:
+                # Seed paths from user: resolve once to get normalized path + fs.
+                resolved_paths, _ = _resolve_paths_and_filesystem(item.path, filesystem)
+                assert len(resolved_paths) == 1
+                path = resolved_paths[0]
+                root_path = path
+
+            contents = _get_path_contents(
+                path, filesystem, self._ignore_missing_paths, root_path=root_path
+            )
+            for file_path, file_size in contents.files:
+                file_info_result = FileInfo(path=file_path, size=file_size)
+                if preserve_order:
+                    add_result(
+                        OrderedFileResult(
+                            input_path_index=input_path_index,
+                            file_path=file_path,
+                            file_info=file_info_result,
+                        )
+                    )
+                else:
+                    add_result(file_info_result)
+            for subdir_path in contents.subdirs:
+                add_work(
+                    _TraversalWorkItem(
+                        path=subdir_path,
+                        is_discovered_subdir=True,
+                        input_path_index=input_path_index,
+                        root_path=root_path,
+                    )
+                )
+
+        def _ordered_result_key(result: OrderedFileResult) -> Tuple[int, str]:
+            return (result.input_path_index, result.file_path)
+
+        if preserve_order:
+            for result in parallel_process_work_stealing(
+                seed_items=seed_items,
+                process_fn=process_fn,
+                num_workers=num_workers,
+                preserve_order=True,
+                order_key=_ordered_result_key,
+            ):
+                # Ordered mode returns `OrderedFileResult` for sorting, so unwrap.
+                yield result.file_info
+        else:
+            yield from parallel_process_work_stealing(
+                seed_items=seed_items,
+                process_fn=process_fn,
+                num_workers=num_workers,
+            )
+
+    def _filter_file_infos(
+        self,
+        file_infos: Iterable[FileInfo],
+        pruners: List[FilePruner],
+    ) -> Iterator[FileInfo]:
+        """Drop zero-size and pruned files before any per-file metadata read."""
+        for file_info in file_infos:
+            if file_info.size is None or file_info.size == 0:
+                logger.warning(f"Skipping zero-size file: {file_info.path!r}")
+                continue
+            if not all(pruner.should_include(file_info.path) for pruner in pruners):
+                continue
+            yield file_info
+
+    def _generate_chunk_records(
+        self,
+        file_infos: Iterable[FileInfo],
+        filesystem: "FileSystem",
+        preserve_order: bool,
+    ) -> Iterator[Tuple[str, int, Optional[ChunkMetadata]]]:
+        """Drive the chunker per file, yielding ``(path, chunk_size, metadata)``.
+
+        When the chunker reads per-file metadata (e.g. ``ParquetFileChunker``
+        reading footers), fan the work across the indexer's thread pool so the
+        I/O parallelizes even for a single input directory — ``make_async_gen``
+        over the *discovered files*, not the input paths. Chunkers that don't
+        read metadata (whole-file / line-delimited) are driven inline to avoid
+        a pointless thread hand-off.
+        """
+        chunker = self._file_chunker
+
+        def chunk_file(
+            fi: FileInfo,
+        ) -> List[Tuple[str, int, Optional[ChunkMetadata]]]:
+            return [
+                (fi.path, chunk_size, chunk_metadata)
+                for chunk_metadata, chunk_size in chunker.generate_chunk_metadatas(
+                    fi.path, fi.size, filesystem
+                )
+            ]
+
+        if chunker.reads_file_metadata and self._num_workers > 1:
+            # Fan per-file footer reads across the thread pool. ``make_async_gen``
+            # only preserves ordering for a 1:1 map (one output per input item), so
+            # emit ONE record list per file and flatten here. Yielding chunk rows
+            # individually would let its round-robin merge interleave chunks from
+            # the files processed concurrently -- breaking per-file contiguity and
+            # discovery order under ``preserve_order=True``.
+            def chunk_files(
+                infos: Iterator[FileInfo],
+            ) -> Iterator[List[Tuple[str, int, Optional[ChunkMetadata]]]]:
+                for fi in infos:
+                    yield chunk_file(fi)
+
+            for records in make_async_gen(
+                # ``iter(...)`` so a non-iterator iterable (e.g. a list from a
+                # test or subclass) is still consumed correctly by the helper.
+                base_iterator=iter(file_infos),
+                fn=chunk_files,
+                preserve_ordering=preserve_order,
+                num_workers=self._num_workers,
+                buffer_size=self._queue_size_per_thread,
+            ):
+                yield from records
+        else:
+            for fi in file_infos:
+                yield from chunk_file(fi)
+
+    def _batch_chunk_records_to_manifests(
+        self,
+        chunk_records: Iterable[Tuple[str, int, Optional[ChunkMetadata]]],
+    ) -> Iterable[FileManifest]:
+        """Batch chunk records into ``FileManifest`` blocks of bounded size."""
+        running_paths: List[str] = []
+        running_file_sizes: List[int] = []
+        running_chunk_metadatas: List[Optional[ChunkMetadata]] = []
+        manifests_count = 0
+        chunks_count = 0
+
+        for path, chunk_size, chunk_metadata in chunk_records:
+            running_paths.append(path)
+            running_file_sizes.append(chunk_size)
+            running_chunk_metadatas.append(chunk_metadata)
+            chunks_count += 1
+
+            if len(running_paths) >= self._max_paths_per_output:
+                manifests_count += 1
+                yield FileManifest.construct_manifest(
+                    running_paths,
+                    running_file_sizes,
+                    running_chunk_metadatas,
+                )
+                running_paths = []
+                running_file_sizes = []
+                running_chunk_metadatas = []
+
+        if running_paths:
+            manifests_count += 1
+            yield FileManifest.construct_manifest(
+                running_paths,
+                running_file_sizes,
+                running_chunk_metadatas,
+            )
+
+        logger.debug(
+            f"Listing files: constructed {manifests_count} manifests "
+            f"with {chunks_count} file chunks"
+        )

python/ray/data/_internal/datasource_v2/listing/file_manifest.py

@@ -0,0 +1,141 @@
+from functools import cached_property
+from typing import List, Optional
+
+import numpy as np
+import pyarrow as pa
+
+from ray.data._internal.datasource_v2.chunkers.file_chunker import ChunkMetadata
+from ray.data.block import Block, BlockAccessor, BlockColumnAccessor
+
+# File manifest column names
+PATH_COLUMN_NAME = "__path"
+FILE_SIZE_COLUMN_NAME = "__file_size"
+FILE_CHUNK_METADATA_COLUMN_NAME = "__file_chunk_metadata"
+
+
+class FileManifest:
+    """Structured view over file paths, sizes, and per-chunk metadata.
+
+    Provides structured access to file paths, sizes, and chunk metadata. This avoids
+    making implicit assumptions about block structure as data moves between file
+    listing, partitioning, and reading stages.
+
+    All extracted views (i.e., `paths`, `file_sizes`, `file_chunk_metadatas`) share
+    the same row order as the underlying block. Any transformation must preserve this.
+
+    Each row represents a single chunk of a file. For unchunked files (whole-file
+    reads), the chunk-metadata entry is ``None`` and ``file_sizes`` equals the
+    on-disk file size. For chunked files, multiple rows can share the same path
+    but carry different chunk metadata.
+    """
+
+    def __init__(self, block: Block):
+        """Create a new `FileManifest` from a block.
+
+        Args:
+            block: Block with `PATH_COLUMN_NAME`, `FILE_SIZE_COLUMN_NAME`, and
+                `FILE_CHUNK_METADATA_COLUMN_NAME` columns. Any other columns are
+                optional and treated as input data.
+        """
+        column_names = BlockAccessor.for_block(block).column_names()
+        assert FILE_SIZE_COLUMN_NAME in column_names
+        assert PATH_COLUMN_NAME in column_names
+        assert FILE_CHUNK_METADATA_COLUMN_NAME in column_names
+
+        self._block = block
+
+        self._paths = block[PATH_COLUMN_NAME]
+        self._file_sizes = block[FILE_SIZE_COLUMN_NAME]
+        self._file_chunk_metadatas = block[FILE_CHUNK_METADATA_COLUMN_NAME]
+
+    def __len__(self) -> int:
+        return len(self._block)
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} length={len(self._block)}>"
+
+    # TODO Use arrow arrays instead of numpy for these properties.
+    @cached_property
+    def paths(self) -> np.ndarray:
+        return BlockColumnAccessor.for_column(self._paths).to_numpy()
+
+    @cached_property
+    def file_sizes(self) -> np.ndarray:
+        return BlockColumnAccessor.for_column(self._file_sizes).to_numpy()
+
+    @cached_property
+    def file_chunk_metadatas(self) -> np.ndarray:
+        return BlockColumnAccessor.for_column(self._file_chunk_metadatas).to_numpy()
+
+    def as_block(self) -> Block:
+        """Return the underlying block for the `FileManifest`.
+
+        This doesn't make a copy of the underlying data.
+        """
+        return self._block
+
+    @classmethod
+    def concat(cls, manifests: List["FileManifest"]) -> "FileManifest":
+        """Return a new `FileManifest` whose rows are the concatenation of
+        ``manifests`` in order.
+
+        Row alignment of ``paths`` / ``file_sizes`` is preserved because
+        each input already satisfies it.
+        """
+        assert len(manifests) > 0, "concat requires at least one manifest"
+        if len(manifests) == 1:
+            return manifests[0]
+
+        merged = pa.concat_tables(
+            [
+                BlockAccessor.for_block(manifest._block).to_arrow()
+                for manifest in manifests
+            ]
+        )
+        return cls(merged)
+
+    def shuffle(self, seed: Optional[int]) -> "FileManifest":
+        """Return a new `FileManifest` with rows permuted.
+
+        Args:
+            seed: Random seed. ``None`` for non-deterministic shuffling.
+                When set, input rows are first sorted by path so the shuffle
+                is reproducible regardless of upstream listing order
+                (the threaded ``FileIndexer`` doesn't preserve order).
+
+        Returns:
+            A new `FileManifest` with the same rows in a shuffled order. The
+            underlying row alignment between `paths` and `file_sizes` is
+            preserved because the permutation is applied to the block as a
+            whole.
+        """
+        n = len(self)
+        if n <= 1:
+            return self
+        block = self._block
+        if seed is not None:
+            sort_indices = pa.compute.sort_indices(
+                BlockAccessor.for_block(block).to_arrow(),
+                sort_keys=[(PATH_COLUMN_NAME, "ascending")],
+            )
+            block = block.take(sort_indices)
+        permutation = np.random.default_rng(seed).permutation(n)
+        return FileManifest(block.take(permutation))
+
+    @classmethod
+    def construct_manifest(
+        cls,
+        paths: List[str],
+        sizes: List[int],
+        chunk_metadatas: List[Optional[ChunkMetadata]],
+    ) -> "FileManifest":
+        assert len(paths) == len(sizes) == len(chunk_metadatas)
+
+        block = pa.table(
+            {
+                PATH_COLUMN_NAME: paths,
+                FILE_SIZE_COLUMN_NAME: sizes,
+                FILE_CHUNK_METADATA_COLUMN_NAME: chunk_metadatas,
+            }
+        )
+        return cls(block)

python/ray/data/_internal/datasource_v2/listing/file_pruners.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+from ray.data.datasource import PathPartitionFilter
+from ray.data.datasource.path_util import _has_file_extension
+
+
+class FilePruner(ABC):
+    """Generic file-level filter applied during listing."""
+
+    @abstractmethod
+    def should_include(self, path: str) -> bool:
+        """Return True if this file should be included, False to skip it."""
+        ...
+
+
+class FileExtensionPruner(FilePruner):
+    """Skip files that don't match the expected extensions."""
+
+    def __init__(self, file_extensions: List[str]):
+        self._file_extensions = file_extensions
+
+    def should_include(self, path: str) -> bool:
+        return _has_file_extension(path, self._file_extensions)
+
+
+class PartitionPruner(FilePruner):
+    """Skip files based on partition column predicates (e.g., hive partitioning)."""
+
+    def __init__(self, partition_filter: PathPartitionFilter):
+        self._filter = partition_filter
+
+    def should_include(self, path: str) -> bool:
+        return self._filter.apply(path)

python/ray/data/_internal/datasource_v2/listing/indexing_utils.py

@@ -0,0 +1,123 @@
+import logging
+from dataclasses import dataclass
+from typing import Iterable, List, Optional, Tuple
+
+import pyarrow as pa
+from pyarrow.fs import FileSelector, FileType
+
+from ray.data.datasource.file_meta_provider import _handle_read_os_error
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class PathContents:
+    """Contents of a path: files (path, size) and subdirectories to expand."""
+
+    files: List[Tuple[str, Optional[int]]]
+    subdirs: List[str]
+
+
+def _expand_directory(
+    base_path: str,
+    filesystem: pa.fs.FileSystem,
+    ignore_missing_path: bool,
+    *,
+    root_path: Optional[str] = None,
+) -> PathContents:
+    """List one level of a directory.
+
+    Hidden-prefix (``.``/``_``) exclusion is applied relative to ``root_path``
+    (the top-level path the traversal started from), not the immediate parent,
+    so a nested entry is only excluded when its path *relative to the root*
+    begins with an excluded prefix. When ``root_path`` is ``None`` it defaults
+    to ``base_path``.
+    """
+    exclude_prefixes = [".", "_"]
+
+    if root_path is None:
+        root_path = base_path
+
+    selector = FileSelector(
+        base_path, recursive=False, allow_not_found=ignore_missing_path
+    )
+    children = filesystem.get_file_info(selector)
+
+    # Lineage reconstruction doesn't work if tasks aren't deterministic, and
+    # `filesystem.get_file_info` might return files in a non-deterministic order. So, we
+    # sort the files.
+    assert isinstance(children, list), type(children)
+    children.sort(key=lambda file_: file_.path)
+
+    files: List[Tuple[str, Optional[int]]] = []
+    subdirs: List[str] = []
+
+    for child in children:
+        if not child.path.startswith(root_path):
+            continue
+
+        relative = child.path[len(root_path) :].lstrip("/")
+        if any(relative.startswith(prefix) for prefix in exclude_prefixes):
+            continue
+
+        if child.type == FileType.File:
+            files.append((child.path, child.size))
+        elif child.type == FileType.Directory:
+            subdirs.append(child.path)
+        elif child.type == FileType.UNKNOWN:
+            logger.warning(f"Discovered file with unknown type: '{child.path}'")
+            continue
+        else:
+            assert child.type == FileType.NotFound
+            raise FileNotFoundError(child.path)
+
+    return PathContents(files=files, subdirs=subdirs)
+
+
+def _get_path_contents(
+    path: str,
+    filesystem: pa.fs.FileSystem,
+    ignore_missing_path: bool,
+    *,
+    root_path: Optional[str] = None,
+) -> PathContents:
+    """Get files and subdirs for a path. Handles File, Directory, and NotFound.
+
+    Only one level of a directory is expanded; discovered subdirectories are
+    returned in :attr:`PathContents.subdirs` for the caller to expand.
+    """
+    try:
+        file_info = filesystem.get_file_info(path)
+    except OSError as e:
+        _handle_read_os_error(e, path)
+
+    if file_info.type == FileType.File:
+        return PathContents(files=[(path, file_info.size)], subdirs=[])
+    elif file_info.type == FileType.Directory:
+        return _expand_directory(
+            path, filesystem, ignore_missing_path, root_path=root_path
+        )
+    elif file_info.type == FileType.NotFound and ignore_missing_path:
+        return PathContents(files=[], subdirs=[])
+    else:
+        raise FileNotFoundError(path)
+
+
+def _get_file_infos(
+    path: str,
+    filesystem: pa.fs.FileSystem,
+    ignore_missing_path: bool,
+    *,
+    _root_path: Optional[str] = None,
+) -> Iterable[Tuple[str, Optional[int]]]:
+    """Recursively expand a path (file or directory) into ``(path, size)`` tuples."""
+    if _root_path is None:
+        _root_path = path
+    contents = _get_path_contents(
+        path, filesystem, ignore_missing_path, root_path=_root_path
+    )
+    yield from contents.files
+    for subdir in contents.subdirs:
+        yield from _get_file_infos(
+            subdir, filesystem, ignore_missing_path, _root_path=_root_path
+        )

python/ray/data/_internal/datasource_v2/listing/listing_utils.py

@@ -0,0 +1,159 @@
+from typing import TYPE_CHECKING, Iterable, List, Optional
+
+import pyarrow as pa
+
+from ray.data._internal.datasource_v2.listing.file_manifest import (
+    PATH_COLUMN_NAME,
+    FileManifest,
+)
+from ray.data._internal.datasource_v2.listing.file_pruners import (
+    FileExtensionPruner,
+    FilePruner,
+    PartitionPruner,
+)
+from ray.data._internal.datasource_v2.partitioners.file_partitioner import (
+    FilePartitioner,
+)
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data.block import Block, BlockAccessor
+
+if TYPE_CHECKING:
+    from pyarrow.fs import FileSystem
+
+    from ray.data._internal.datasource_v2.listing.file_indexer import FileIndexer
+    from ray.data.datasource.file_based_datasource import FileShuffleConfig
+    from ray.data.datasource.partitioning import PathPartitionFilter
+
+
+def partition_files(
+    blocks: Iterable[Block],
+    _: TaskContext,
+    partitioner: FilePartitioner,
+) -> Iterable[Block]:
+    for block in blocks:
+        partitioner.add_input(FileManifest(block))
+        while partitioner.has_partition():
+            yield partitioner.next_partition().as_block()
+
+    partitioner.finalize()
+    while partitioner.has_partition():
+        yield partitioner.next_partition().as_block()
+
+
+def _build_pruners(
+    file_extensions: Optional[List[str]],
+    partition_filter: Optional["PathPartitionFilter"],
+) -> List[FilePruner]:
+    pruners: List[FilePruner] = []
+    if file_extensions is not None:
+        pruners.append(FileExtensionPruner(file_extensions))
+    if partition_filter is not None:
+        pruners.append(PartitionPruner(partition_filter))
+    return pruners
+
+
+def list_files_for_each_block(
+    blocks: Iterable[Block],
+    _: TaskContext,
+    *,
+    indexer: "FileIndexer",
+    filesystem: "FileSystem",
+    file_extensions: Optional[List[str]] = None,
+    partition_filter: Optional["PathPartitionFilter"] = None,
+    preserve_order: bool = False,
+) -> Iterable[Block]:
+    """Expand path blocks into ``FileManifest`` blocks.
+
+    Each input block carries a single ``__path`` column of path strings.
+    For every path, the indexer is invoked to produce a stream of
+    ``FileManifest`` objects; each manifest's backing block is yielded.
+    Pruners are constructed once per task from ``file_extensions`` /
+    ``partition_filter`` — keeps pruner construction out of the
+    ``_read_datasource_v2`` entry point.
+    """
+    pruners = _build_pruners(file_extensions, partition_filter)
+    for block in blocks:
+        for manifest in indexer.list_files(
+            block[PATH_COLUMN_NAME],
+            filesystem=filesystem,
+            pruners=pruners,
+            preserve_order=preserve_order,
+        ):
+            if len(manifest) > 0:
+                yield manifest.as_block()
+
+
+def shuffle_files(
+    blocks: Iterable[Block],
+    _: TaskContext,
+    *,
+    shuffle_config: "FileShuffleConfig",
+    execution_idx: int,
+) -> Iterable[Block]:
+    """Concatenate manifest blocks and shuffle rows with the seeded RNG.
+
+    Runs in a single task (`plan_list_files_op` sets `should_parallelize=False`
+    when shuffle is requested) so we have the full manifest before
+    shuffling. Emits one merged manifest block. Determinism comes from
+    ``FileManifest.shuffle`` which sorts by path before applying the
+    permutation — this protects against non-deterministic upstream
+    indexer yield order.
+    """
+    builder = DelegatingBlockBuilder()
+    for block in blocks:
+        if len(block) > 0:
+            builder.add_block(block)
+
+    combined = builder.build()
+    if len(combined) == 0:
+        return
+
+    seed = shuffle_config.get_seed(execution_idx)
+    yield FileManifest(combined).shuffle(seed).as_block()
+
+
+def sample_files(
+    indexer: "FileIndexer",
+    paths: List[str],
+    filesystem: "FileSystem",
+    pruners: Optional[List[FilePruner]] = None,
+    max_files: int = 16,
+) -> FileManifest:
+    """Drive the indexer until up to ``max_files`` files arrive; return them.
+
+    Used for driver-side schema inference in ``_read_datasource_v2``.
+    Sampling more than one file lets callers unify schemas (e.g., if the
+    first file has an all-null column, later files' non-null types can
+    promote it). No caching — the returned manifest is discarded after
+    schema inference, and the ``ListFiles`` op lists the same paths
+    again on workers at execution time.
+    """
+    assert max_files >= 1
+    paths_column = pa.array(paths, type=pa.string())
+    collected: List[FileManifest] = []
+    collected_rows = 0
+    for manifest in indexer.list_files(
+        paths_column,
+        filesystem=filesystem,
+        pruners=pruners or [],
+        preserve_order=True,
+    ):
+        if len(manifest) == 0:
+            continue
+        remaining = max_files - collected_rows
+        if len(manifest) <= remaining:
+            collected.append(manifest)
+            collected_rows += len(manifest)
+        else:
+            collected.append(
+                FileManifest(
+                    BlockAccessor.for_block(manifest.as_block()).slice(0, remaining)
+                )
+            )
+            collected_rows = max_files
+        if collected_rows >= max_files:
+            break
+    if not collected:
+        return FileManifest.construct_manifest(paths=[], sizes=[], chunk_metadatas=[])
+    return FileManifest.concat(collected)

python/ray/data/_internal/datasource_v2/logical_optimizers.py

@@ -0,0 +1,122 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, List, Optional, Set, Tuple
+
+from ray.data.expressions import Expr
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.data._internal.datasource_v2.scanners.scanner import Scanner
+
+
+@DeveloperAPI
+class SupportsFilterPushdown(ABC):
+    """Mixin for scanners that support filter/predicate pushdown.
+
+    Filter pushdown allows predicates to be evaluated at the data source level,
+    reducing the amount of data that needs to be read and transferred.
+    """
+
+    @abstractmethod
+    def push_filters(self, predicate: "Expr") -> Tuple["Scanner", Optional["Expr"]]:
+        """Push a filter predicate down to the scanner.
+
+        Args:
+            predicate: Expression representing the filter condition.
+
+        Returns:
+            Tuple of (new_scanner, residual_predicate) where:
+            - new_scanner: New Scanner instance with the filter applied
+            - residual_predicate: Any part of the predicate that couldn't be
+              pushed down and must be applied post-scan. None if fully pushed.
+        """
+        ...
+
+
+@DeveloperAPI
+class SupportsColumnPruning(ABC):
+    """Mixin for scanners that support column pruning/projection pushdown.
+
+    Column pruning allows reading only the columns needed by the query,
+    which is especially beneficial for columnar formats like Parquet.
+    """
+
+    @abstractmethod
+    def prune_columns(self, columns: List[str]) -> "Scanner":
+        """Prune the scanner to only read the specified columns.
+
+        Args:
+            columns: List of column names to read.
+
+        Returns:
+            New Scanner instance configured to read only the specified columns.
+        """
+        ...
+
+    @abstractmethod
+    def pruned_column_names(self) -> Optional[Tuple[str, ...]]:
+        """Physical column names selected after pruning, if any.
+
+        Returns:
+            ``None`` when no pruning has been applied (read all columns).
+            A tuple (possibly empty) after :meth:`prune_columns` has been
+            applied, listing on-disk / reader column names in read order.
+        """
+        ...
+
+
+@DeveloperAPI
+class SupportsLimitPushdown(ABC):
+    """Mixin for scanners that support limit pushdown.
+
+    Limit pushdown allows the scanner to stop early once the required number
+    of rows has been read.
+    """
+
+    @abstractmethod
+    def push_limit(self, limit: int) -> "Scanner":
+        """Push a row limit down to the scanner.
+
+        Args:
+            limit: Maximum number of rows to read.
+
+        Returns:
+            New Scanner instance with the limit applied.
+        """
+        ...
+
+
+@DeveloperAPI
+class SupportsPartitionPruning(ABC):
+    """Mixin for scanners that support partition pruning.
+
+    Partition pruning allows skipping entire files/partitions based on
+    predicates that reference partition columns.
+    """
+
+    @property
+    @abstractmethod
+    def partition_columns(self) -> Set[str]:
+        """Names of columns that are partition keys.
+
+        Callers (e.g. the predicate-pushdown rule) use this to decide
+        whether a predicate should be routed through :meth:`push_filters`
+        (data columns) or :meth:`prune_partitions` (partition columns).
+        Must be fully populated by schema inference at planning time.
+        """
+        ...
+
+    @abstractmethod
+    def prune_partitions(self, predicate: "Expr") -> "Scanner":
+        """Prune partitions based on a predicate.
+
+        The scanner determines its partition columns from its
+        ``Partitioning`` configuration, which is fully populated
+        by schema inference at planning time.
+
+        Args:
+            predicate: Expression to evaluate against partition values.
+
+        Returns:
+            New Scanner instance with partition pruning applied.
+        """
+        ...

python/ray/data/_internal/datasource_v2/parquet_datasource_v2.py

@@ -0,0 +1,309 @@
+"""Concrete ``DataSourceV2`` for Parquet files.
+
+Wires the V2 listing (`NonSamplingFileIndexer`, driven by the upstream
+`ListFiles` op), scanning (`ParquetScanner`), and reading
+(`ParquetFileReader`) components against a user-supplied path set.
+Constructed from `read_api.read_parquet` when
+`DataContext.use_datasource_v2` is set.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union
+
+import pyarrow as pa
+from typing_extensions import override
+
+from ray.data._internal.datasource.parquet_datasource import (
+    ParquetDatasource,
+    check_for_legacy_tensor_type,
+)
+from ray.data._internal.datasource_v2.chunkers.file_chunker import (
+    FileChunker,
+    ParquetFileChunker,
+)
+from ray.data._internal.datasource_v2.datasource_v2 import (
+    DatasourceCategory,
+    DataSourceV2,
+)
+from ray.data._internal.datasource_v2.listing.file_indexer import (
+    FileIndexer,
+    NonSamplingFileIndexer,
+)
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.readers.file_reader import (
+    INCLUDE_PATHS_COLUMN_NAME,
+)
+from ray.data._internal.datasource_v2.readers.in_memory_size_estimator import (
+    ParquetInMemorySizeEstimator,
+)
+from ray.data._internal.datasource_v2.scanners.parquet_scanner import ParquetScanner
+from ray.data._internal.util import _is_local_scheme
+from ray.data.context import DataContext
+from ray.data.datasource.partitioning import (
+    Partitioning,
+    PartitionStyle,
+    PathPartitionParser,
+    _partition_field_types_to_pa_schema,
+)
+from ray.data.datasource.path_util import _resolve_paths_and_filesystem
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from pyarrow.fs import FileSystem
+
+    from ray.data.datasource.file_based_datasource import FileShuffleConfig
+
+
+@DeveloperAPI
+class ParquetDatasourceV2(DataSourceV2[FileManifest]):
+    """V2 Parquet datasource.
+
+    Listing is delegated to :class:`NonSamplingFileIndexer` driven by the
+    ``ListFiles`` logical op; scanning and reading are delegated to
+    :class:`ParquetScanner` and :class:`ParquetFileReader`. Schema
+    inference reads the first file's footer only and augments with
+    partition/path columns as configured.
+    """
+
+    def __init__(
+        self,
+        paths: List[str],
+        *,
+        filesystem: Optional["FileSystem"] = None,
+        partitioning: Optional[Partitioning] = Partitioning(PartitionStyle.HIVE),
+        file_extensions: Optional[List[str]] = None,
+        ignore_missing_paths: bool = False,
+        include_paths: bool = False,
+        include_row_hash: bool = False,
+        shuffle: Optional[Union[Literal["files"], "FileShuffleConfig"]] = None,
+        arrow_parquet_args: Optional[dict] = None,
+        schema: Optional[pa.Schema] = None,
+        parquet_format_kwargs: Optional[dict] = None,
+        file_chunker: Optional[FileChunker] = None,
+    ):
+        super().__init__(name="ParquetV2", category=DatasourceCategory.FILE_BASED)
+        # Capture the ``local://`` check against the *original* paths;
+        # ``_resolve_paths_and_filesystem`` below strips the scheme, so
+        # introspecting ``self._paths`` after construction can't tell a
+        # plain local path from a ``local://`` one. ``_supports_distributed_reads``
+        # is exposed by the base-class ``supports_distributed_reads``
+        # property and consumed by ``_read_datasource_v2``.
+        self._supports_distributed_reads = not _is_local_scheme(paths)
+        resolved_paths, resolved_filesystem = _resolve_paths_and_filesystem(
+            paths, filesystem
+        )
+        self._paths: List[str] = resolved_paths
+        self._filesystem = resolved_filesystem
+        self._partitioning = partitioning
+        self._file_extensions = file_extensions or ParquetDatasource._FILE_EXTENSIONS
+        self._ignore_missing_paths = ignore_missing_paths
+        self._include_paths = include_paths
+        self._include_row_hash = include_row_hash
+        self._shuffle = shuffle
+        self._arrow_parquet_args = arrow_parquet_args or {}
+        # ``pds.ParquetFileFormat`` kwargs forwarded from the deprecated
+        # ``read_parquet(dataset_kwargs=...)`` arg. Spread into the format
+        # built by ``ParquetFileReader._make_format``.
+        self._parquet_format_kwargs = parquet_format_kwargs or {}
+        # User-supplied schema override. When set, ``infer_schema`` returns
+        # it verbatim (plus partition/path augmentation) rather than reading
+        # footers, and the scanner pins it on the pyarrow dataset so files
+        # are cast to these types at scan time.
+        self._user_schema = schema
+        # Chunker that splits each listed Parquet file into one or more
+        # row-group-aligned read units. Defaults to ``ParquetFileChunker``
+        # (1 GiB target chunk size, or whatever ``DataContext`` configures).
+        # Callers can inject an alternative for tests or shuffle-aware
+        # planning code that wants whole-file reads.
+        self._file_chunker: FileChunker = (
+            file_chunker if file_chunker is not None else ParquetFileChunker()
+        )
+
+    @property
+    def paths(self) -> List[str]:
+        return self._paths
+
+    @property
+    def filesystem(self) -> Optional["FileSystem"]:
+        return self._filesystem
+
+    @property
+    def partitioning(self) -> Optional[Partitioning]:
+        return self._partitioning
+
+    @property
+    def file_extensions(self) -> List[str]:
+        return self._file_extensions
+
+    @property
+    def ignore_missing_paths(self) -> bool:
+        return self._ignore_missing_paths
+
+    @property
+    def include_paths(self) -> bool:
+        return self._include_paths
+
+    @property
+    def shuffle(self) -> Optional[Union[Literal["files"], "FileShuffleConfig"]]:
+        return self._shuffle
+
+    def _get_file_indexer(self) -> FileIndexer:
+        return NonSamplingFileIndexer(
+            ignore_missing_paths=self._ignore_missing_paths,
+            file_chunker=self._file_chunker,
+        )
+
+    def get_size_estimator(self) -> ParquetInMemorySizeEstimator:
+        return ParquetInMemorySizeEstimator()
+
+    @override
+    def resolve_partitioning(self, sample: FileManifest) -> Optional[Partitioning]:
+        """Return ``self._partitioning`` with path-discovered field names.
+
+        Hive partitioning ships with ``field_names=None`` by default and
+        discovers keys from the file path at plan time. Directory
+        partitioning already carries ``field_names`` at construction and
+        needs no discovery. Returns a fresh ``Partitioning`` rather than
+        mutating ``self`` so schema inference stays side-effect-free.
+        """
+        import copy
+
+        if self._partitioning is None or len(sample) == 0:
+            return copy.deepcopy(self._partitioning)
+        if self._partitioning.field_names:
+            return copy.deepcopy(self._partitioning)
+
+        first_path = sample.paths.tolist()[0]
+        parser = PathPartitionParser(self._partitioning)
+        partition_kv = parser(first_path)
+        if not partition_kv:
+            return copy.deepcopy(self._partitioning)
+        return Partitioning(
+            style=self._partitioning.style,
+            base_dir=self._partitioning.base_dir,
+            field_names=list(partition_kv.keys()),
+            field_types=self._partitioning.field_types,
+            filesystem=self._partitioning.filesystem,
+        )
+
+    def infer_schema(self, sample: FileManifest) -> pa.Schema:
+        """Read Parquet footers from the sample manifest; unify and augment.
+
+        When the sample has multiple files, their schemas are unified via
+        ``unify_schemas_with_validation`` so a first file with all-null
+        columns doesn't lock in ``null`` types that can't be cast to the
+        actual types in later files.
+
+        Pure: does not mutate ``self``. Partitioning field-name discovery
+        is delegated to :meth:`resolve_partitioning` so the discovered
+        ``Partitioning`` can flow through ``_read_datasource_v2`` into
+        :meth:`create_scanner` without side effects.
+        """
+        from concurrent.futures import ThreadPoolExecutor
+
+        import pyarrow.parquet as pq
+
+        from ray.data._internal.util import unify_schemas_with_validation
+
+        # Empty sample — typically means the user pointed ``read_parquet``
+        # at an empty directory. Return an empty schema so the rest of
+        # the plan stays valid; downstream ops produce zero blocks and
+        # the executor runs through without error (matches V1).
+        if len(sample) == 0:
+            return self._user_schema if self._user_schema is not None else pa.schema([])
+
+        sample_paths: List[str] = sample.paths.tolist()
+        # Parquet footer reads against high-latency object stores
+        # (S3, GCS) are ~50-100 ms each. Reading the sample's footers in
+        # parallel keeps driver-side schema inference bounded by the
+        # slowest single read rather than the sum. ``executor.map``
+        # preserves input order, which matters because the unified
+        # schema's field order follows the first schema's and
+        # ``sample_paths[0]`` drives partition discovery below.
+        #
+        # NOTE: ``pq.read_schema`` only accepts a ``filesystem=`` kwarg in
+        # recent pyarrow releases; older wheels in CI don't have it. Open
+        # the file through the configured filesystem and hand the file
+        # handle to ``read_schema`` for cross-version compatibility.
+        filesystem = self._filesystem
+
+        if self._user_schema is not None:
+            # Caller pinned the schema — skip footer reads. Partition/path
+            # augmentation below still applies so downstream ops see the
+            # synthesized columns.
+            schema = self._user_schema
+        else:
+
+            def _read_schema(path: str):
+                if filesystem is None:
+                    return pq.read_schema(path)
+                with filesystem.open_input_file(path) as handle:
+                    return pq.read_schema(handle)
+
+            with ThreadPoolExecutor(max_workers=min(len(sample_paths), 16)) as executor:
+                per_file_schemas = list(executor.map(_read_schema, sample_paths))
+            schema = (
+                unify_schemas_with_validation(per_file_schemas) or per_file_schemas[0]
+            )
+            assert isinstance(schema, pa.Schema)
+
+        resolved_partitioning = self.resolve_partitioning(sample)
+        if resolved_partitioning is not None:
+            first_path = sample_paths[0]
+            parser = PathPartitionParser(resolved_partitioning)
+            partition_kv = parser(first_path)
+            # For hive partitioning the parser discovers key names from the
+            # path itself; for directory partitioning it uses ``field_names``.
+            # In both cases ``partition_kv`` is the authoritative list of
+            # partition columns for the first sample file.
+            partition_pa_schema = _partition_field_types_to_pa_schema(
+                field_names=list(partition_kv.keys()),
+                field_types=resolved_partitioning.field_types or {},
+            )
+            for field_name in partition_kv.keys():
+                if schema.get_field_index(field_name) == -1:
+                    pa_type = partition_pa_schema.field(field_name).type
+                    schema = schema.append(pa.field(field_name, pa_type))
+
+        if (
+            self._include_paths
+            and schema.get_field_index(INCLUDE_PATHS_COLUMN_NAME) == -1
+        ):
+            schema = schema.append(pa.field(INCLUDE_PATHS_COLUMN_NAME, pa.string()))
+
+        if self._include_row_hash:
+            # ``row_hash`` is synthesized post-read as ``uint64``. Replace
+            # the field type when the file already has a ``row_hash``
+            # column (matches V1 ``_derive_schema``); otherwise append.
+            idx = schema.get_field_index("row_hash")
+            if idx == -1:
+                schema = schema.append(pa.field("row_hash", pa.uint64()))
+            elif schema.field(idx).type != pa.uint64():
+                schema = schema.set(idx, pa.field("row_hash", pa.uint64()))
+
+        check_for_legacy_tensor_type(schema)
+        return schema
+
+    def create_scanner(
+        self,
+        schema: pa.Schema,
+        filesystem: Optional["FileSystem"] = None,
+        **options: Any,
+    ) -> ParquetScanner:
+        # Callers (``_read_datasource_v2``) supply the sample-resolved
+        # ``Partitioning`` via ``options["partitioning"]`` so the
+        # datasource itself stays immutable — fall back to the
+        # constructor-provided one for direct users of this API.
+        partitioning = options.get("partitioning", self._partitioning)
+        return ParquetScanner(
+            schema=schema,
+            filesystem=filesystem or self._filesystem,
+            partitioning=partitioning,
+            include_paths=self._include_paths,
+            include_row_hash=self._include_row_hash,
+            shuffle=self._shuffle,
+            ignore_prefixes=options.get("ignore_prefixes"),
+            target_block_size=DataContext.get_current().target_max_block_size,
+            parquet_format_kwargs=dict(self._parquet_format_kwargs),
+        )

python/ray/data/_internal/datasource_v2/partitioners/file_partitioner.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+
+
+class FilePartitioner(ABC):
+    """Abstract base class for partitioning file manifests.
+
+    A ``FilePartitioner`` groups file paths and their associated metadata into new
+    file manifests based on a specific partitioning strategy.
+
+    Implementations must be deterministic to ensure consistent partitioning across
+    retries.
+    """
+
+    @abstractmethod
+    def add_input(self, input_manifest: FileManifest):
+        """Add a file manifest to be partitioned.
+
+        Args:
+            input_manifest: A ``FileManifest`` containing paths and metadata to partition.
+        """
+        ...
+
+    @abstractmethod
+    def has_partition(self) -> bool:
+        """Check if there are any partitions available.
+
+        Returns:
+            ``True`` if there are partitions ready to be retrieved via
+            ``next_partition()``, ``False`` otherwise.
+        """
+        ...
+
+    @abstractmethod
+    def next_partition(self) -> FileManifest:
+        """Get the next available partition.
+
+        Returns:
+            A ``FileManifest`` containing the paths and metadata for the next partition.
+        """
+        ...
+
+    @abstractmethod
+    def finalize(self):
+        """Process any remaining files and complete the partitioning.
+
+        This method is called after all inputs have been added via ``add_input()`` to
+        ensure any buffered files are properly partitioned.
+        """
+        ...

python/ray/data/_internal/datasource_v2/partitioners/round_robin_partitioner.py

@@ -0,0 +1,89 @@
+import logging
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.partitioners.file_partitioner import (
+    FilePartitioner,
+)
+from ray.data._internal.datasource_v2.readers.in_memory_size_estimator import (
+    InMemorySizeEstimator,
+)
+from ray.data._internal.weighted_round_robin import WeightedRoundRobinPartitioner
+
+logger = logging.getLogger(__name__)
+
+
+class RoundRobinPartitioner(FilePartitioner):
+    """Partitions input paths into blocks based on the in-memory size of files.
+
+    This partitioning ensures read tasks effectively utilize the cluster and
+    produce appropriately-sized blocks
+
+    **Steps:**
+        1. Initialize empty buckets.
+        2. Iterate through input blocks and add paths to buckets. For each path:
+            - If the current bucket falls below `min_bucket_size`, add the path and don't move
+              to the next bucket.
+            - If the current bucket exceeds `min_bucket_size` but not `max_bucket_size`,
+              add the path and move to the next bucket.
+            - If the current bucket exceeds `max_bucket_size`, yield the paths as a block, clear
+              the bucket, and move to the next bucket.
+        3. Yield any remaining paths in the buckets as blocks.
+
+    This algorithm ensures that each block contains [min_bucket_size, max_bucket_size]
+    worth of files.  It's a deterministic algorithm, but it doesn't maintain the order
+    of the input paths.
+    """
+
+    def __init__(
+        self,
+        in_memory_size_estimator: InMemorySizeEstimator,
+        *,
+        min_bucket_size: int,
+        max_bucket_size: int,
+        num_buckets: int,
+    ):
+        self._in_memory_size_estimator = in_memory_size_estimator
+        self._partitioner = WeightedRoundRobinPartitioner(
+            min_bucket_size=min_bucket_size,
+            max_bucket_size=max_bucket_size,
+            num_buckets=num_buckets,
+        )
+
+    def add_input(self, input_manifest: FileManifest):
+        in_memory_size_estimates = (
+            self._in_memory_size_estimator.estimate_in_memory_sizes(input_manifest)
+        )
+        for (
+            file_path,
+            file_size,
+            file_chunk_metadata,
+            in_memory_size_estimate,
+        ) in zip(
+            input_manifest.paths,
+            input_manifest.file_sizes,
+            input_manifest.file_chunk_metadatas,
+            in_memory_size_estimates,
+        ):
+            self._partitioner.add_item(
+                (file_path, file_size, file_chunk_metadata),
+                in_memory_size_estimate,
+            )
+
+    def has_partition(self) -> bool:
+        return self._partitioner.has_partition()
+
+    @property
+    def num_buckets(self) -> int:
+        return self._partitioner.num_buckets
+
+    def next_partition(self) -> FileManifest:
+        partition = self._partitioner.next_partition()
+        paths, file_sizes, file_chunk_metadatas = zip(*partition)
+        return FileManifest.construct_manifest(
+            list(paths),
+            list(file_sizes),
+            list(file_chunk_metadatas),
+        )
+
+    def finalize(self):
+        self._partitioner.finalize()

python/ray/data/_internal/datasource_v2/readers/base_reader.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import Generic, Iterator
+
+import pyarrow as pa
+
+from ray.data._internal.datasource_v2 import InputSplit
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class Reader(ABC, Generic[InputSplit]):
+    """Abstract base class for reading data from input buckets.
+
+    Readers execute on workers to actually read data. They receive an InputSplit
+    (e.g., FileManifest for file-based sources) and yield Arrow tables.
+
+    The Reader is created by Scanner.create_reader() and is configured with all
+    pushdown optimizations (columns, predicates, limits) that were applied.
+    """
+
+    @abstractmethod
+    def read(self, input_split: InputSplit) -> Iterator[pa.Table]:
+        """Read data from the input bucket and yield Arrow tables.
+
+        This method is called on workers to perform the actual read operation.
+        It should respect all pushdowns configured on this reader.
+
+        Args:
+            input_split: Work unit describing what data to read.
+
+        Returns:
+            Iterator[pa.Table]: Iterator of PyArrow Tables containing the read data.
+        """
+        ...

python/ray/data/_internal/datasource_v2/readers/file_reader.py

@@ -0,0 +1,513 @@
+from enum import Enum
+from functools import cached_property, partial
+from typing import Any, Iterator, List, Optional, Set, Tuple
+
+import pyarrow as pa
+import pyarrow.dataset as pds
+from pyarrow.fs import FileSystem, LocalFileSystem
+
+from ray._common.utils import env_integer
+from ray.data._internal.arrow_block import _BATCH_SIZE_PRESERVING_STUB_COL_NAME
+from ray.data._internal.datasource.parquet_datasource import _compute_row_hashes
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.readers.base_reader import Reader
+from ray.data._internal.util import iterate_with_retry, make_async_gen
+from ray.data.context import DataContext
+from ray.data.datasource.partitioning import Partitioning, PathPartitionParser
+from ray.data.expressions import Expr
+from ray.util.annotations import DeveloperAPI
+
+# Synthetic column name produced when ``include_paths=True``. Shared with
+# the V2 datasource and scanner layers so all references are spelled the
+# same way.
+INCLUDE_PATHS_COLUMN_NAME = "path"
+
+# https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Scanner.html#pyarrow.dataset.Scanner.from_batches
+# Default is specified by PyArrow.
+_ARROW_DEFAULT_BATCH_SIZE = 131_072
+
+# Number of batches read ahead per scanner. PyArrow's default is 16,
+# which can retain a multi-GB working set when scanning jumbo tensor
+# columns. 8 keeps I/O pipelined on remote filesystems for typical
+# Parquet workloads without doubling memory peak. Drop to 1 via the
+# env var when reading wide tensor columns.
+_ARROW_SCANNER_BATCH_READAHEAD = env_integer(
+    "RAY_DATA_ARROW_SCANNER_BATCH_READAHEAD", 8
+)
+
+# Number of worker threads used to read fragments concurrently per task.
+# Defaults to 4 to overlap remote-filesystem I/O latency across multiple
+# fragments. ``_read_fragment_batches`` caps this to ``len(fragments)``
+# at runtime so single-fragment tasks don't spin up extra workers, and
+# falls back to the sequential path entirely when
+# ``DataContext.execution_options.preserve_order`` is set.
+_DEFAULT_NUM_THREADS = env_integer("RAY_DATA_READ_FILES_NUM_THREADS", 4)
+
+ROW_HASH_COLUMN_NAME = "row_hash"
+
+
+class FileFormat(str, Enum):
+    PARQUET = "parquet"
+    CSV = "csv"
+    FEATHER = "feather"
+    JSON = "json"
+    ARROW = "arrow"
+    IPC = "ipc"
+
+
+@DeveloperAPI
+class FileReader(Reader[FileManifest]):
+    """Reader for file-based sources.
+
+    This reader uses PyArrow's Dataset API which automatically handles:
+    - Column pruning
+    - Filter pushdown (row group pruning)
+    - Batch-level filtering
+    """
+
+    def __init__(
+        self,
+        format: FileFormat,
+        batch_size: int = _ARROW_DEFAULT_BATCH_SIZE,
+        columns: Optional[List[str]] = None,
+        predicate: Optional[Expr] = None,
+        limit: Optional[int] = None,
+        filesystem: Optional[FileSystem] = None,
+        partitioning: Optional[Partitioning] = None,
+        ignore_prefixes: Optional[List[str]] = None,
+        include_paths: bool = False,
+        include_row_hash: bool = False,
+        schema: Optional[pa.Schema] = None,
+    ):
+        """Initialize the reader.
+        Refer to https://arrow.apache.org/docs/python/generated/pyarrow.dataset.dataset.html for more details.
+
+        Args:
+            format: Format of the files to read.
+            batch_size: Number of rows per batch.
+            columns: Columns to read. None means all columns.
+            predicate: Ray Data expression for filtering. Converted to a
+                PyArrow expression at the scanner-kwargs boundary.
+            limit: Maximum number of rows to read.
+            filesystem: Filesystem for reading files.
+            partitioning: Ray ``Partitioning`` object. Partition columns are
+                synthesized per-path via ``PathPartitionParser`` after each
+                batch is read, producing string-typed columns (V1 parity).
+            ignore_prefixes: Prefixes to ignore when reading files. Default is ['.', '_'] set by PyArrow.
+            include_paths: If True, include the source file path in a
+                ``'path'`` column for each row.
+            include_row_hash: If True, include a deterministic uint64 hash
+                per row in a ``'row_hash'`` column. The hash is derived from
+                the source file path and the row's post-filter output
+                position within the fragment, matching V1 semantics. If a
+                ``'row_hash'`` column already exists in the file, it is
+                overwritten.
+            schema: Caller-supplied unified schema used both to override
+                pyarrow's per-fragment inference (so a file whose column
+                is all-null doesn't pin the type to ``null``) and to cast
+                path-derived partition values to their target types when
+                ``Partitioning(field_types=...)`` is set.
+
+        """
+        self._format = format
+        self._columns = columns
+        self._predicate = predicate
+        self._batch_size = batch_size
+        self._limit = limit
+        self._filesystem = filesystem
+        self._partition_parser: Optional[PathPartitionParser] = (
+            PathPartitionParser(partitioning) if partitioning is not None else None
+        )
+        self._ignore_prefixes = ignore_prefixes
+        self._include_paths = include_paths
+        self._include_row_hash = include_row_hash
+        self._schema = schema
+
+    @cached_property
+    def _file_dataset_schema(self) -> Optional[pa.Schema]:
+        """Schema passed to ``pds.dataset`` — partition keys and ``path``
+        stripped out since those are synthesized post-read.
+
+        Pinning the caller-supplied schema at the pyarrow layer is how
+        we cover the "first file has an all-null column, later files
+        have the real type" case (e.g.
+        ``test_read_null_data_in_first_file``): without the pin,
+        pyarrow locks column X to ``null`` across the fragment group
+        and the later string-typed file fails the cast.
+
+        But pyarrow refuses extension-to-extension casts (e.g.
+        ``ArrowTensorTypeV2(shape=X)`` → ``ArrowVariableShapedTensor``),
+        and files with different per-file tensor shapes only unify
+        through ``ArrowVariableShapedTensor``. When the caller schema
+        contains *any* extension column we skip the pin entirely and
+        let pyarrow infer per-file — downstream concat handles the
+        heterogeneous blocks. Losing the all-null promotion in this
+        narrow case is acceptable; the combination of an all-null
+        first file *and* an extension column is uncommon, whereas
+        reading multiple files with variable-shape tensors is a
+        supported V1 feature.
+        """
+        if self._schema is None:
+            return None
+        if any(isinstance(f.type, pa.ExtensionType) for f in self._schema):
+            return None
+        partition_keys = (
+            set(self._partition_parser._scheme.field_names or [])
+            if self._partition_parser is not None
+            else set()
+        )
+        synthesized = {INCLUDE_PATHS_COLUMN_NAME}
+        if self._include_row_hash:
+            # ``row_hash`` is synthesized post-read, and the schema's type
+            # (``uint64``) may not match the on-disk column's type when a
+            # file already carries a ``row_hash`` column. Strip it from the
+            # dataset schema so pyarrow doesn't try to cast.
+            synthesized.add(ROW_HASH_COLUMN_NAME)
+        fields = [
+            f
+            for f in self._schema
+            if f.name not in partition_keys and f.name not in synthesized
+        ]
+        return pa.schema(fields) if fields else None
+
+    def _broadcast_partition_value(
+        self, name: str, value: Any, num_rows: int
+    ) -> pa.Array:
+        """Broadcast a single path-derived partition value to ``num_rows``,
+        casting to the caller-supplied schema's field type if set.
+
+        Values are stringified first (``PathPartitionParser`` in
+        ``explicit`` mode can return arrow-scalar-like non-strings) and
+        then cast to the target type, so ``Partitioning(field_types=
+        {"year": int})`` still promotes them correctly.
+        """
+        str_val = None if value is None else str(value)
+        arr = pa.repeat(pa.scalar(str_val, type=pa.string()), num_rows)
+        if self._schema is not None:
+            idx = self._schema.get_field_index(name)
+            if idx != -1 and self._schema.field(idx).type != pa.string():
+                arr = arr.cast(self._schema.field(idx).type)
+        return arr
+
+    def read(self, input_split: FileManifest) -> Iterator[pa.Table]:
+        """Read data from the input bucket and yield Arrow tables.
+
+        This method is called on workers to perform the actual read operation.
+        It should respect all pushdowns configured on this reader.
+
+        Args:
+            input_split: Work unit describing what data to read.
+
+        Yields:
+            pa.Table: PyArrow Tables containing the read data.
+        """
+        if len(input_split) == 0:
+            return
+
+        # Dedupe paths before handing them to pyarrow. When chunking is on,
+        # a manifest can carry multiple rows per file (each describing a
+        # different row-group slice); pyarrow only needs one fragment per
+        # file, and ``_get_fragments_to_read`` then fans out chunk-level
+        # sub-fragments using the per-row chunk metadata.
+        paths = list(dict.fromkeys(list(input_split.paths)))
+        filesystem = self._filesystem or LocalFileSystem()
+        # Build a ``pds.Dataset`` over *all* manifest paths so pyarrow's
+        # listing + column metadata is shared, but then iterate its
+        # fragments one at a time. ``dataset.scanner(fragments=...)``
+        # at the aggregate level would force a cross-fragment cast —
+        # which breaks variable-shape tensor extensions where each
+        # file has its own ``ArrowTensorTypeV2(shape=...)``. Per-
+        # fragment scanners let pyarrow use the native per-file type,
+        # and downstream concat handles unification.
+        dataset = pds.dataset(
+            source=paths,
+            format=self._make_format(),
+            filesystem=filesystem,
+            schema=self._file_dataset_schema,
+            ignore_prefixes=self._ignore_prefixes,
+        )
+
+        # Split the requested columns into ones the on-disk file has
+        # (pyarrow reads these) and ones we need to synthesize post-read
+        # (hive partition keys, "path"). ``self._columns is None`` means
+        # "no projection" — read every file column and synthesize every
+        # available partition/path column.
+        on_disk_column_names = set(dataset.schema.names)
+        if self._columns is None:
+            columns_to_read_from_file: Optional[List[str]] = None
+            columns_to_synthesize: Optional[Set[str]] = None
+        else:
+            columns_to_read_from_file = [
+                c for c in self._columns if c in on_disk_column_names
+            ]
+            columns_to_synthesize = set(self._columns) - on_disk_column_names
+
+        scanner_kwargs = {
+            "columns": columns_to_read_from_file,
+            "filter": (
+                self._predicate.to_pyarrow() if self._predicate is not None else None
+            ),
+            "batch_size": self._resolve_batch_size(dataset),
+            "batch_readahead": _ARROW_SCANNER_BATCH_READAHEAD,
+        }
+        scanner_kwargs.update(self._arrow_scanner_kwargs())
+
+        rows_read = 0
+        for table, fragment_path, fragment_row_offset in self._read_fragment_batches(
+            dataset, scanner_kwargs, input_split
+        ):
+            if self._limit is not None:
+                if rows_read >= self._limit:
+                    break
+                if len(table) > self._limit - rows_read:
+                    table = table.slice(0, self._limit - rows_read)
+
+            # Build the list of (name, value) pairs to synthesize from
+            # the fragment path: hive partitions + optional ``path``.
+            derived_items: List[Tuple[str, Any]] = []
+            if self._partition_parser is not None:
+                derived_items.extend(self._partition_parser(fragment_path).items())
+            if self._include_paths:
+                derived_items.append((INCLUDE_PATHS_COLUMN_NAME, fragment_path))
+
+            for name, value in derived_items:
+                if (
+                    columns_to_synthesize is not None
+                    and name not in columns_to_synthesize
+                ):
+                    continue
+                if name in table.column_names:
+                    # When the caller schema names a partition key, pyarrow
+                    # expects it in every file and fills it with nulls when
+                    # absent (the hive-typical case). Drop that placeholder
+                    # so the path-derived value below replaces it.
+                    table = table.drop([name])
+                table = table.append_column(
+                    name,
+                    self._broadcast_partition_value(name, value, table.num_rows),
+                )
+
+            # Skip when projection pushdown has narrowed ``columns`` to
+            # exclude ``row_hash`` — the projection below would just drop it.
+            if self._include_row_hash and (
+                columns_to_synthesize is None
+                or ROW_HASH_COLUMN_NAME in columns_to_synthesize
+            ):
+                hashes = _compute_row_hashes(
+                    fragment_path, fragment_row_offset, table.num_rows
+                )
+                if ROW_HASH_COLUMN_NAME in table.column_names:
+                    table = table.drop([ROW_HASH_COLUMN_NAME])
+                table = table.append_column(
+                    ROW_HASH_COLUMN_NAME, pa.array(hashes, type=pa.uint64())
+                )
+
+            if self._columns is not None:
+                # Project/reorder to the caller's requested column order;
+                # drop any that weren't produced (matches V1's lenient
+                # behavior). Always select — an empty projection must
+                # narrow the table to zero columns so the stub-column
+                # guard below handles row preservation.
+                produced = set(table.column_names)
+                projected = [c for c in self._columns if c in produced]
+                table = table.select(projected)
+
+            if table.num_columns == 0 and table.num_rows > 0:
+                # Guards against ``pa.concat_tables`` collapsing rows
+                # when a batch has zero columns (e.g., empty projection
+                # for a count query). The stub column is dropped by
+                # downstream projections.
+                table = table.append_column(
+                    _BATCH_SIZE_PRESERVING_STUB_COL_NAME,
+                    pa.nulls(table.num_rows),
+                )
+
+            self._on_batch_read(table)
+            rows_read += len(table)
+            yield table
+
+    def _resolve_batch_size(self, dataset: pds.Dataset) -> int:
+        """Return the batch size to use for scanning.
+
+        Subclasses can override this to implement adaptive batch sizing.
+        """
+        return self._batch_size
+
+    def _on_batch_read(self, table: pa.Table) -> None:
+        """Hook called after each batch is read.
+
+        Subclasses can override this to update internal state (e.g., refine
+        batch size estimates from actual data).
+        """
+        pass
+
+    def _arrow_scanner_kwargs(self) -> dict:
+        """Additional keyword arguments passed to ``pds.Dataset.scanner()``.
+
+        Subclasses override this to inject format-specific options.
+        """
+        return {}
+
+    def _make_format(self) -> Any:
+        """Format passed to ``pds.dataset(format=...)``.
+
+        Defaults to the format string (e.g. ``"parquet"``); subclasses
+        override to return a configured ``pds.FileFormat`` instance when
+        format-specific options (read options, fragment scan options) need
+        to be threaded through.
+        """
+        return self._format.value
+
+    def _get_fragments_to_read(
+        self,
+        dataset: pds.Dataset,
+        manifest: FileManifest,
+    ) -> List[Tuple[pds.Fragment, int]]:
+        """Return ``(fragment, file_row_offset)`` pairs to scan for this
+        manifest.
+
+        ``file_row_offset`` is the cumulative pre-filter row count of all
+        rows in the underlying file that precede this fragment. It seeds
+        the per-fragment hashing offset so chunked sub-fragments of the
+        same file produce unique ``_compute_row_hashes`` keys instead of
+        colliding on ``(path, 0, n)``.
+
+        Default impl returns one ``(fragment, 0)`` per file in the dataset
+        (paths are deduped in :meth:`read` before the dataset is built).
+        Subclasses that support per-row chunk metadata
+        (e.g. :class:`ParquetFileReader`) override this to fan a single
+        file fragment out into N sub-fragments — one per row-group slice —
+        based on :attr:`FileManifest.file_chunk_metadatas`, each paired
+        with its starting row offset in the file.
+        """
+        return [(fragment, 0) for fragment in dataset.get_fragments()]
+
+    def _read_fragment_batches(
+        self,
+        dataset: pds.Dataset,
+        scanner_kwargs: dict,
+        manifest: FileManifest,
+    ) -> Iterator[Tuple[pa.Table, str, int]]:
+        """Yield non-empty (table, fragment_path, fragment_row_offset) triples.
+
+        ``fragment_row_offset`` is the post-filter row position of the first
+        row of ``table`` within its fragment. ``iterate_with_retry`` skips
+        already-yielded items on retry, so ``offset`` reflects only the
+        rows that actually surface to the caller — matching V1 row-hash
+        semantics even when a fragment fails partway through.
+
+        Retry is scoped per-fragment: if a fragment fails mid-read, only
+        that fragment is re-read (skipping batches already yielded).
+        Wrapping the whole manifest in a single retry would re-iterate
+        fragments that already succeeded and double-emit their batches.
+
+        Each fragment gets its own scanner so pyarrow uses the native
+        per-file schema. A cross-fragment scanner would force a unified
+        schema cast, which refuses extension-to-extension conversion
+        (e.g. variable-shape tensors). V1 ``ParquetDatasource`` follows
+        the same per-fragment pattern via ``fragment.to_batches``.
+
+        When ``RAY_DATA_READ_FILES_NUM_THREADS > 1`` and
+        ``execution_options.preserve_order`` is False, fragments are
+        read concurrently via :func:`make_async_gen`. We still pass
+        ``preserve_ordering=True`` so concurrent reads emit blocks in
+        fragment order; otherwise Ray Data task retries (block
+        reconstruction) could produce a different block sequence.
+
+        ``make_async_gen`` consumes the whole input iterator up front
+        when preserving order. That is acceptable here because the input
+        is the finite fragment manifest from ``_get_fragments_to_read``,
+        which we materialize below anyway. File data is still read lazily
+        by the worker threads.
+        """
+        ctx = DataContext.get_current()
+
+        # ``preserve_ordering=True`` would drain the input iterator
+        # eagerly anyway, so materialize once here to (a) cap
+        # ``num_workers`` at the actual fragment count and (b) avoid
+        # an early-fallback when the manifest has a single fragment.
+        # Subclasses (e.g. ``ParquetFileReader``) override
+        # ``_get_fragments_to_read`` to fan out chunk-level
+        # sub-fragments from the manifest's chunk metadata.
+        fragments_with_offsets = self._get_fragments_to_read(dataset, manifest)
+        if not fragments_with_offsets:
+            return
+
+        num_workers = min(_DEFAULT_NUM_THREADS, len(fragments_with_offsets))
+        if num_workers <= 1 or ctx.execution_options.preserve_order:
+            yield from self._read_fragments_sequential(
+                iter(fragments_with_offsets), scanner_kwargs
+            )
+            return
+
+        # Set `preserve_ordering=True` to ensure deterministic output ordering.
+        # This is required so that Ray Data task retries (block reconstruction)
+        yield from make_async_gen(
+            base_iterator=iter(fragments_with_offsets),
+            fn=partial(self._read_fragments_sequential, scanner_kwargs=scanner_kwargs),
+            preserve_ordering=True,
+            num_workers=num_workers,
+        )
+
+    def _read_fragments_sequential(
+        self,
+        fragments_with_offsets: Iterator[Tuple[pds.Fragment, int]],
+        scanner_kwargs: dict,
+    ) -> Iterator[Tuple[pa.Table, str, int]]:
+        """Read each fragment in ``fragments_with_offsets`` in order, yielding
+        ``(table, fragment_path, fragment_row_offset)`` triples.
+
+        Each input pair is ``(fragment, file_row_offset)``. The yielded
+        ``fragment_row_offset`` starts at ``file_row_offset`` (the row
+        position of the fragment's first row within its underlying file)
+        and accumulates per yielded batch, so the per-fragment row-hash
+        math in :meth:`read` keys off the right window even when chunking
+        fans one file into multiple sub-fragments sharing ``fragment.path``.
+
+        ``iterate_with_retry`` is scoped to a single fragment so a
+        transient I/O failure only re-reads the failing file (skipping
+        batches already yielded), not the whole input.
+
+        This is the per-worker body for the threaded path in
+        :meth:`_read_fragment_batches` (one thread per call, each
+        consuming a disjoint slice of fragments via ``make_async_gen``)
+        and is also the entire read loop for the sequential path.
+        """
+        ctx = DataContext.get_current()
+        for fragment, file_row_offset in fragments_with_offsets:
+            offset = file_row_offset
+            for table in iterate_with_retry(
+                partial(self._iter_fragment_tables, fragment, scanner_kwargs),
+                f"read fragment {fragment.path}",
+                match=ctx.retried_io_errors,
+            ):
+                if table.num_rows > 0:
+                    yield table, fragment.path, offset
+                    offset += table.num_rows
+
+    def _iter_fragment_tables(
+        self,
+        fragment: pds.Fragment,
+        scanner_kwargs: dict,
+    ) -> Iterator[pa.Table]:
+        """Yield Arrow tables for a single fragment.
+
+        Subclasses override this to swap in a format-specific reader for
+        fragments that don't fit the default scanner-based path (e.g.
+        Parquet's ARROW-5030 nested-type fallback).
+
+        When a non-extension caller schema is available we pin it at the
+        scanner so pyarrow null-fills any column the unified schema names
+        but the fragment lacks (V1 parity — ``ParquetDatasource`` passes
+        ``read_schema`` to ``fragment.to_batches``). Falling back to the
+        per-fragment ``physical_schema`` preserves the variable-shape
+        tensor escape hatch already encoded in ``_file_dataset_schema``.
+        """
+        fragment_schema = (
+            self._file_dataset_schema
+            if self._file_dataset_schema is not None
+            else fragment.physical_schema
+        )
+        scanner = fragment.scanner(**scanner_kwargs, schema=fragment_schema)
+        for tagged in scanner.scan_batches():
+            yield pa.Table.from_batches(batches=[tagged.record_batch])

python/ray/data/_internal/datasource_v2/readers/in_memory_size_estimator.py

@@ -0,0 +1,146 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import numpy as np
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.readers.file_reader import FileReader
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data.block import BlockAccessor
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class InMemorySizeEstimator(ABC):
+    @abstractmethod
+    def estimate_in_memory_sizes(self, manifest: FileManifest) -> np.ndarray:
+        """Estimate the in-memory sizes of the paths in the given manifest.
+
+        Some `FilePartitioner` implementations use this method to ensure that each
+        read task receives an appropriate amount of data. To ensure that file listing
+        is efficient, this method must be cheap to call, on average.
+
+        Args:
+            manifest: A manifest containing the paths and on-disk sizes of the files.
+
+        Returns:
+            The estimated in-memory sizes of the data in bytes.
+        """
+        ...
+
+
+@DeveloperAPI
+class SamplingInMemorySizeEstimator(InMemorySizeEstimator):
+    """Estimates in-memory sizes by reading files.
+
+    This class estimates the in-memory size of files by multiplying the on-disk
+    size by an estimated encoding ratio. If an instance hasn't estimated an encoding
+    ratio yet, it'll read a file to estimate it. Otherwise, it'll use the previously
+    estimated encoding ratio.
+
+    TODO: This approach doesn't work well for formats that produce multiple batches
+    (because we assume a 1:1 encoding ratio) or for formats that vary in encoding
+    ratios (e.g. videos).
+    """
+
+    def __init__(self, reader: "FileReader"):
+        self._reader = reader
+
+        self._encoding_ratio = None
+
+    def estimate_in_memory_sizes(self, manifest: FileManifest) -> np.ndarray:
+        assert np.all(manifest.file_sizes >= 0)
+
+        for path, file_size in zip(manifest.paths, manifest.file_sizes):
+            if self._encoding_ratio is None:
+                # Estimating the encoding ratio can be expensive since it requires
+                # reading the file. So, we only estimate the encoding ratio if we don't
+                # already have one.
+                self._encoding_ratio = self._estimate_encoding_ratio(path, file_size)
+                break
+
+        if self._encoding_ratio is None:
+            # If we couldn't estimate the encoding ratio, assume a 1:1 encoding ratio.
+            return manifest.file_sizes
+        else:
+            return manifest.file_sizes * self._encoding_ratio
+
+    def _estimate_encoding_ratio(
+        self,
+        path: str,
+        file_size: int,
+    ) -> Optional[float]:
+        """
+        Estimate the encoding ratio (in-memory size / on-disk size) for a file.
+
+        Args:
+            path: The path to the file.
+            file_size: The on-disk size of the file/chunk in bytes.
+
+        Returns:
+            The estimated encoding ratio of the file, or `None` if the ratio can't
+            be estimated.
+        """
+        # If the file is empty, we can't estimate the encoding ratio.
+        if not file_size:
+            return None
+
+        # Use ``None`` chunk metadata: the size estimator reads the file whole
+        # to estimate the encoding ratio; chunk-level splitting is irrelevant here.
+        manifest = FileManifest.construct_manifest(
+            [path],
+            [file_size],
+            [None],
+        )
+        batches = self._reader.read(manifest)
+
+        try:
+            first_batch = next(batches)
+        except StopIteration:
+            # If there's no data, we can't estimate the encoding ratio.
+            return None
+
+        try:
+            # Try to read a second batch. If it succeeds, it means the file contains
+            # multiple batches.
+            next(batches)
+        except StopIteration:
+            # Each file contains exactly one batch.
+            builder = DelegatingBlockBuilder()
+            builder.add_batch(first_batch)
+            block = builder.build()
+
+            in_memory_size = BlockAccessor.for_block(block).size_bytes()
+        else:
+            # Each file contains multiple batches.
+            #
+            # NOTE: To avoid reading the entire file to estimate the encoding ratio,
+            # we assume the file is 1:1 encoded. We can't return `None` because if
+            # all files contain multiple batches, then we'd try to re-estimate the
+            # encoding ratio for every file, and that'd be very expensive.
+            in_memory_size = file_size
+
+        return in_memory_size / file_size
+
+
+# Default Parquet encoding ratio: in-memory is ~5x on-disk size.
+# Parquet uses columnar compression and encoding, so Arrow in-memory
+# representation is significantly larger than the on-disk format.
+PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT = 5
+
+
+@DeveloperAPI
+class ParquetInMemorySizeEstimator(InMemorySizeEstimator):
+    """Estimates in-memory sizes for Parquet files using a fixed encoding ratio.
+
+    Parquet files are typically much smaller on disk than in memory due to
+    columnar compression and encoding. This estimator applies a constant
+    ratio (default 5x) to avoid the overhead of reading file metadata or
+    sampling data, which can be slow for Parquet files and hurt startup time.
+    """
+
+    def __init__(self, encoding_ratio: float = PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT):
+        self._encoding_ratio = encoding_ratio
+
+    def estimate_in_memory_sizes(self, manifest: FileManifest) -> np.ndarray:
+        return self._encoding_ratio * manifest.file_sizes