chore: import upstream snapshot with attribution

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