chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/databricks_test_utils.py

@@ -0,0 +1,69 @@
+"""Test utilities for Databricks datasource tests."""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from ray.data._internal.datasource.databricks_credentials import (
+    DatabricksCredentialProvider,
+)
+
+
+@dataclass
+class MockResponse:
+    """Mock HTTP response for testing.
+
+    Args:
+        status_code: HTTP status code. Defaults to 200.
+        content: Response content as bytes. Defaults to None.
+        _json_data: JSON response data. Defaults to None.
+        raise_on_error: If True, raise_for_status() raises for status >= 400.
+            Defaults to True.
+    """
+
+    status_code: int = 200
+    content: Optional[bytes] = None
+    _json_data: Optional[dict] = None
+    raise_on_error: bool = field(default=True, repr=False)
+
+    def raise_for_status(self):
+        """Raise an exception if status code indicates an error."""
+        if self.raise_on_error and self.status_code >= 400:
+            raise Exception(f"HTTP Error {self.status_code}")
+
+    def json(self):
+        """Return the JSON data."""
+        return self._json_data
+
+
+class RefreshableCredentialProvider(DatabricksCredentialProvider):
+    """A credential provider that simulates token refresh on invalidate.
+
+    Useful for testing 401 retry logic. When invalidate() is called,
+    the token changes from initial_token to "refreshed_token".
+
+    Args:
+        initial_token: The initial token value. Defaults to "expired_token".
+        host: The host URL to return. Defaults to "https://test-host.databricks.com".
+    """
+
+    def __init__(
+        self,
+        initial_token: str = "expired_token",
+        host: str = "https://test-host.databricks.com",
+    ):
+        self.current_token = initial_token
+        self.invalidate_count = 0
+        self._host = host
+
+    def get_token(self) -> str:
+        """Get the current token."""
+        return self.current_token
+
+    def get_host(self) -> str:
+        """Get the host URL."""
+        return self._host
+
+    def invalidate(self) -> None:
+        """Simulate token refresh by changing to 'refreshed_token'."""
+        self.invalidate_count += 1
+        self.current_token = "refreshed_token"

python/ray/data/tests/datasource/test_arrow.py

@@ -0,0 +1,134 @@
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data.tests.conftest import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def sample_dataframes():
+    """Fixture providing sample pandas DataFrames for testing.
+
+    Returns:
+        tuple: (df1, df2) where df1 has 3 rows and df2 has 3 rows
+    """
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+    return df1, df2
+
+
+def test_from_arrow(ray_start_regular_shared, sample_dataframes):
+    """Test basic from_arrow functionality with single and multiple tables."""
+    df1, df2 = sample_dataframes
+
+    ds = ray.data.from_arrow([pa.Table.from_pandas(df1), pa.Table.from_pandas(df2)])
+    values = [(r["one"], r["two"]) for r in ds.take(6)]
+    rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
+    assert values == rows
+    # Check that metadata fetch is included in stats.
+    assert "FromArrow" in ds.stats()
+
+    # test from single pyarrow table
+    ds = ray.data.from_arrow(pa.Table.from_pandas(df1))
+    values = [(r["one"], r["two"]) for r in ds.take(3)]
+    rows = [(r.one, r.two) for _, r in df1.iterrows()]
+    assert values == rows
+    # Check that metadata fetch is included in stats.
+    assert "FromArrow" in ds.stats()
+
+
+@pytest.mark.parametrize(
+    "tables,override_num_blocks,expected_blocks,expected_rows",
+    [
+        # Single table scenarios
+        ("single", 1, 1, 3),  # Single table, 1 block
+        ("single", 2, 2, 3),  # Single table split into 2 blocks
+        ("single", 5, 5, 3),  # Single table, more blocks than rows
+        (
+            "single",
+            10,
+            10,
+            3,
+        ),  # Edge case: 3 rows split into 10 blocks (creates empty blocks)
+        # Multiple tables scenarios
+        ("multiple", 3, 3, 6),  # Multiple tables split into 3 blocks
+        ("multiple", 10, 10, 6),  # Multiple tables, more blocks than rows
+        # Empty table scenarios
+        ("empty", 1, 1, 0),  # Empty table, 1 block
+        ("empty", 5, 5, 0),  # Empty table, more blocks than rows
+    ],
+)
+def test_from_arrow_override_num_blocks(
+    ray_start_regular_shared,
+    sample_dataframes,
+    tables,
+    override_num_blocks,
+    expected_blocks,
+    expected_rows,
+):
+    """Test from_arrow with override_num_blocks parameter."""
+    df1, df2 = sample_dataframes
+    empty_df = pd.DataFrame({"one": [], "two": []})
+
+    # Prepare tables based on test case
+    if tables == "single":
+        arrow_tables = pa.Table.from_pandas(df1)
+        expected_data = [(r.one, r.two) for _, r in df1.iterrows()]
+    elif tables == "multiple":
+        arrow_tables = [pa.Table.from_pandas(df1), pa.Table.from_pandas(df2)]
+        expected_data = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
+    elif tables == "empty":
+        arrow_tables = pa.Table.from_pandas(empty_df)
+        expected_data = []
+
+    # Create dataset with override_num_blocks
+    ds = ray.data.from_arrow(arrow_tables, override_num_blocks=override_num_blocks)
+
+    # Verify number of blocks
+    assert ds.num_blocks() == expected_blocks
+
+    # Verify row count
+    assert ds.count() == expected_rows
+
+    # Verify data integrity (only for non-empty datasets)
+    if expected_rows > 0:
+        values = [(r["one"], r["two"]) for r in ds.take_all()]
+        assert values == expected_data
+
+
+def test_from_arrow_refs(ray_start_regular_shared, sample_dataframes):
+    df1, df2 = sample_dataframes
+    ds = ray.data.from_arrow_refs(
+        [ray.put(pa.Table.from_pandas(df1)), ray.put(pa.Table.from_pandas(df2))]
+    )
+    values = [(r["one"], r["two"]) for r in ds.take(6)]
+    rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
+    assert values == rows
+    # Check that metadata fetch is included in stats.
+    assert "FromArrow" in ds.stats()
+
+    # test from single pyarrow table ref
+    ds = ray.data.from_arrow_refs(ray.put(pa.Table.from_pandas(df1)))
+    values = [(r["one"], r["two"]) for r in ds.take(3)]
+    rows = [(r.one, r.two) for _, r in df1.iterrows()]
+    assert values == rows
+    # Check that metadata fetch is included in stats.
+    assert "FromArrow" in ds.stats()
+
+
+def test_to_arrow_refs(ray_start_regular_shared):
+    n = 5
+    df = pd.DataFrame({"id": list(range(n))})
+    ds = ray.data.range(n)
+    dfds = pd.concat(
+        [t.to_pandas() for t in ray.get(ds.to_arrow_refs())], ignore_index=True
+    )
+    assert df.equals(dfds)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_audio.py

@@ -0,0 +1,38 @@
+import numpy as np
+import pytest
+
+import ray
+from ray.tests.conftest import *  # noqa
+
+NUM_AUDIO_FILES = 10
+
+
+@pytest.fixture
+def audio_uri():
+    root = "s3://anonymous@air-example-data-2/6G-audio-data-LibriSpeech-train-clean-100-flac"  # noqa: E501
+    return [
+        f"{root}/train-clean-100/5022/29411/5022-29411-{n:04}.flac"
+        for n in range(NUM_AUDIO_FILES)
+    ]
+
+
+def test_read_audio(ray_start_regular_shared, audio_uri):
+    ds = ray.data.read_audio(audio_uri)
+
+    # Verify basic audio properties
+    assert ds.count() == NUM_AUDIO_FILES, ds.count()
+    assert ds.schema().names == ["amplitude", "sample_rate"], ds.schema()
+
+    # Check the sample rate
+    assert all(row["sample_rate"] == 16000 for row in ds.take_all())
+
+    for row in ds.take_all():
+        assert row["amplitude"].ndim == 2
+        assert row["amplitude"].shape[0] == 1
+        assert row["amplitude"].dtype == np.float32
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_avro.py

@@ -0,0 +1,40 @@
+import os
+
+import fastavro
+import pytest
+
+import ray
+
+schema = {
+    "type": "record",
+    "name": "TestRecord",
+    "fields": [{"name": "test_field", "type": "string"}],
+}
+
+
+def test_read_basic_avro_file(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "sample.avro")
+    records = [{"test_field": "test_value1"}, {"test_field": "test_value2"}]
+    with open(path, "wb") as out:
+        fastavro.writer(out, schema, records)
+
+    ds = ray.data.read_avro(path)
+
+    expected = [{"test_field": "test_value1"}, {"test_field": "test_value2"}]
+    assert ds.take_all() == expected
+
+
+def test_read_empty_avro_files(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "empty.avro")
+    # Write an empty Avro file with the schema
+    with open(path, "wb") as out:
+        # Write the schema with no records
+        fastavro.writer(out, schema, [])
+
+    ds = ray.data.read_avro(path)
+
+    assert ds.count() == 0
+
+
+if __name__ == "__main__":
+    pytest.main(["-v", __file__])

python/ray/data/tests/datasource/test_bigquery.py

@@ -0,0 +1,284 @@
+from typing import Iterator
+from unittest import mock
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+from google.api_core import exceptions, operation
+from google.cloud import bigquery, bigquery_storage
+from google.cloud.bigquery import job
+from google.cloud.bigquery_storage_v1.types import stream as gcbqs_stream
+
+import ray
+from ray.data._internal.datasource.bigquery_datasink import BigQueryDatasink
+from ray.data._internal.datasource.bigquery_datasource import BigQueryDatasource
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.planner.plan_write_op import generate_collect_write_stats_fn
+from ray.data.block import Block
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+_TEST_GCP_PROJECT_ID = "mock-test-project-id"
+_TEST_BQ_DATASET_ID = "mockdataset"
+_TEST_BQ_TABLE_ID = "mocktable"
+_TEST_BQ_DATASET = _TEST_BQ_DATASET_ID + "." + _TEST_BQ_TABLE_ID
+_TEST_BQ_TEMP_DESTINATION = _TEST_GCP_PROJECT_ID + ".tempdataset.temptable"
+
+
+@pytest.fixture(autouse=True)
+def bq_client_full_mock(monkeypatch):
+    client_mock = mock.create_autospec(bigquery.Client)
+    client_mock.return_value = client_mock
+
+    def bq_get_dataset_mock(dataset_id):
+        if dataset_id != _TEST_BQ_DATASET_ID:
+            raise exceptions.NotFound(
+                "Dataset {} is not found. Please ensure that it exists.".format(
+                    _TEST_BQ_DATASET
+                )
+            )
+
+    def bq_get_table_mock(table_id):
+        if table_id != _TEST_BQ_DATASET:
+            raise exceptions.NotFound(
+                "Table {} is not found. Please ensure that it exists.".format(
+                    _TEST_BQ_DATASET
+                )
+            )
+
+    def bq_create_dataset_mock(dataset_id, **kwargs):
+        if dataset_id == "existingdataset":
+            raise exceptions.Conflict("Dataset already exists")
+        return mock.Mock(operation.Operation)
+
+    def bq_delete_table_mock(table, **kwargs):
+        return None
+
+    def bq_query_mock(query):
+        fake_job_ref = job._JobReference(
+            "fake_job_id", _TEST_GCP_PROJECT_ID, "us-central1"
+        )
+        fake_query_job = job.QueryJob(fake_job_ref, query, None)
+        fake_query_job.configuration.destination = _TEST_BQ_TEMP_DESTINATION
+        return fake_query_job
+
+    client_mock.get_dataset = bq_get_dataset_mock
+    client_mock.get_table = bq_get_table_mock
+    client_mock.create_dataset = bq_create_dataset_mock
+    client_mock.delete_table = bq_delete_table_mock
+    client_mock.query = bq_query_mock
+
+    monkeypatch.setattr(bigquery, "Client", client_mock)
+    return client_mock
+
+
+@pytest.fixture(autouse=True)
+def bqs_client_full_mock(monkeypatch):
+    client_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
+    client_mock.return_value = client_mock
+
+    def bqs_create_read_session(max_stream_count=0, **kwargs):
+        read_session_proto = gcbqs_stream.ReadSession()
+        read_session_proto.streams = [
+            gcbqs_stream.ReadStream() for _ in range(max_stream_count)
+        ]
+        return read_session_proto
+
+    client_mock.create_read_session = bqs_create_read_session
+
+    monkeypatch.setattr(bigquery_storage, "BigQueryReadClient", client_mock)
+    client_mock.reset_mock()
+    return client_mock
+
+
+@pytest.fixture
+def bq_query_result_mock():
+    with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock:
+        yield query_result_mock
+
+
+@pytest.fixture
+def bq_query_result_mock_fail():
+    with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock_fail:
+        query_result_mock_fail.side_effect = exceptions.BadRequest("400 Syntax error")
+        yield query_result_mock_fail
+
+
+@pytest.fixture
+def ray_get_mock():
+    with mock.patch.object(ray, "get") as ray_get:
+        ray_get.return_value = None
+        yield ray_get
+
+
+class TestReadBigQuery:
+    """Tests for BigQuery Read."""
+
+    @pytest.mark.parametrize(
+        "parallelism",
+        [1, 2, 3, 4, 10, 100],
+    )
+    def test_create_read_tasks(self, parallelism):
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset=_TEST_BQ_DATASET,
+        )
+        read_tasks_list = bq_ds.get_read_tasks(parallelism)
+        assert len(read_tasks_list) == parallelism
+
+    @pytest.mark.parametrize(
+        "parallelism",
+        [1, 2, 3, 4, 10, 100],
+    )
+    def test_create_reader_query(self, parallelism, bq_query_result_mock):
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            query="SELECT * FROM mockdataset.mocktable",
+        )
+        read_tasks_list = bq_ds.get_read_tasks(parallelism)
+        bq_query_result_mock.assert_called_once()
+        assert len(read_tasks_list) == parallelism
+
+    @pytest.mark.parametrize(
+        "parallelism",
+        [1, 2, 3, 4, 10, 100],
+    )
+    def test_create_reader_query_bad_request(
+        self,
+        parallelism,
+        bq_query_result_mock_fail,
+    ):
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            query="SELECT * FROM mockdataset.mocktable",
+        )
+        with pytest.raises(exceptions.BadRequest):
+            bq_ds.get_read_tasks(parallelism)
+        bq_query_result_mock_fail.assert_called()
+
+    def test_dataset_query_kwargs_provided(self):
+        with pytest.raises(ValueError) as exception:
+            BigQueryDatasource(
+                project_id=_TEST_GCP_PROJECT_ID,
+                dataset=_TEST_BQ_DATASET,
+                query="SELECT * FROM mockdataset.mocktable",
+            )
+        expected_message = (
+            "Query and dataset kwargs cannot both be provided"
+            + " (must be mutually exclusive)."
+        )
+        assert str(exception.value) == expected_message
+
+    def test_create_reader_dataset_not_found(self):
+        parallelism = 4
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset="nonexistentdataset.mocktable",
+        )
+        with pytest.raises(ValueError) as exception:
+            bq_ds.get_read_tasks(parallelism)
+        expected_message = (
+            "Dataset nonexistentdataset is not found. Please ensure that it exists."
+        )
+        assert str(exception.value) == expected_message
+
+    def test_create_reader_table_not_found(self):
+        parallelism = 4
+        bq_ds = BigQueryDatasource(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset="mockdataset.nonexistenttable",
+        )
+        with pytest.raises(ValueError) as exception:
+            bq_ds.get_read_tasks(parallelism)
+        expected_message = (
+            "Table mockdataset.nonexistenttable is not found."
+            + " Please ensure that it exists."
+        )
+        assert str(exception.value) == expected_message
+
+
+class TestWriteBigQuery:
+    """Tests for BigQuery Write."""
+
+    def _extract_write_result(self, stats: Iterator[Block]):
+        return dict(next(stats).iloc[0])
+
+    def test_write(self, ray_get_mock):
+        bq_datasink = BigQueryDatasink(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset=_TEST_BQ_DATASET,
+        )
+        arr = pa.array([2, 4, 5, 100])
+        block = pa.Table.from_arrays([arr], names=["data"])
+        ctx = TaskContext(1, "")
+        bq_datasink.write(
+            blocks=[block],
+            ctx=ctx,
+        )
+
+        collect_stats_fn = generate_collect_write_stats_fn()
+        stats = collect_stats_fn([block], ctx)
+        pd.testing.assert_frame_equal(
+            next(stats),
+            pd.DataFrame(
+                {
+                    "num_rows": [4],
+                    "size_bytes": [32],
+                    "write_return": [None],
+                }
+            ),
+        )
+
+    def test_write_dataset_exists(self, ray_get_mock):
+        bq_datasink = BigQueryDatasink(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset="existingdataset" + "." + _TEST_BQ_TABLE_ID,
+        )
+        arr = pa.array([2, 4, 5, 100])
+        block = pa.Table.from_arrays([arr], names=["data"])
+        ctx = TaskContext(1, "")
+        bq_datasink.write(
+            blocks=[block],
+            ctx=ctx,
+        )
+        collect_stats_fn = generate_collect_write_stats_fn()
+        stats = collect_stats_fn([block], ctx)
+        pd.testing.assert_frame_equal(
+            next(stats),
+            pd.DataFrame(
+                {
+                    "num_rows": [4],
+                    "size_bytes": [32],
+                    "write_return": [None],
+                }
+            ),
+        )
+
+    def test_write_empty_block(self, ray_get_mock):
+        """Test that writing a zero-sized block doesn't crash.
+
+        See https://github.com/ray-project/ray/issues/51892
+        """
+        bq_datasink = BigQueryDatasink(
+            project_id=_TEST_GCP_PROJECT_ID,
+            dataset=_TEST_BQ_DATASET,
+        )
+        # Create an empty block with schema but no rows
+        block = pa.Table.from_arrays([pa.array([], type=pa.int64())], names=["data"])
+        ctx = TaskContext(1, "")
+        # This should not raise an error - empty blocks should be skipped
+        bq_datasink.write(
+            blocks=[block],
+            ctx=ctx,
+        )
+
+        # write() always calls ray.get(), but with an empty list since the
+        # zero-row block is filtered out (no remote write tasks launched).
+        ray_get_mock.assert_called_once_with([])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_binary.py

@@ -0,0 +1,53 @@
+import os
+from io import BytesIO
+
+import pytest
+import snappy
+
+import ray
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.util import extract_values, gen_bin_files
+from ray.tests.conftest import *  # noqa
+
+
+def test_read_binary_files(ray_start_regular_shared):
+    with gen_bin_files(10) as (_, paths):
+        ds = ray.data.read_binary_files(paths)
+        for i, item in enumerate(ds.iter_rows()):
+            expected = open(paths[i], "rb").read()
+            assert expected == item["bytes"]
+        # Test metadata ops.
+        assert ds.count() == 10
+        assert "bytes" in str(ds.schema()), ds
+        assert "bytes" in str(ds), ds
+
+
+def test_read_binary_snappy(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "test_binary_snappy")
+    os.mkdir(path)
+    with open(os.path.join(path, "file"), "wb") as f:
+        byte_str = "hello, world".encode()
+        bytes = BytesIO(byte_str)
+        snappy.stream_compress(bytes, f)
+    ds = ray.data.read_binary_files(
+        path,
+        arrow_open_stream_args=dict(compression="snappy"),
+    )
+    assert sorted(extract_values("bytes", ds.take())) == [byte_str]
+
+
+def test_read_binary_snappy_inferred(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "test_binary_snappy_inferred")
+    os.mkdir(path)
+    with open(os.path.join(path, "file.snappy"), "wb") as f:
+        byte_str = "hello, world".encode()
+        bytes = BytesIO(byte_str)
+        snappy.stream_compress(bytes, f)
+    ds = ray.data.read_binary_files(path)
+    assert sorted(extract_values("bytes", ds.take())) == [byte_str]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_clickhouse.py

@@ -0,0 +1,833 @@
+import re
+from unittest import mock
+from unittest.mock import MagicMock, patch
+
+import pyarrow as pa
+import pytest
+from clickhouse_connect.driver.summary import QuerySummary
+
+from ray.data._internal.datasource.clickhouse_datasink import (
+    ClickHouseDatasink,
+    ClickHouseTableSettings,
+    SinkMode,
+)
+from ray.data._internal.datasource.clickhouse_datasource import ClickHouseDatasource
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+
+
+@pytest.fixture(autouse=True)
+def patch_clickhouse_get_client():
+    with patch("clickhouse_connect.get_client") as mock_factory:
+        mock_instance = MagicMock()
+        mock_instance.insert_arrow.return_value = QuerySummary({"written_rows": 3})
+        mock_factory.return_value = mock_instance
+        yield mock_instance
+
+
+@pytest.fixture
+def mock_clickhouse_client():
+    client_mock = mock.MagicMock()
+    client_mock.return_value = client_mock
+    return client_mock
+
+
+class TestClickHouseDatasource:
+    """Tests for ClickHouseDatasource."""
+
+    @pytest.fixture
+    def datasource(self, mock_clickhouse_client):
+        datasource = ClickHouseDatasource(
+            table="default.table_name",
+            dsn="clickhouse://user:password@localhost:8123/default",
+            columns=["column1", "column2"],
+            order_by=(["column1"], False),
+            client_settings={"setting1": "value1"},
+            client_kwargs={"client_name": "test-client"},
+        )
+        datasource._client = mock_clickhouse_client
+        return datasource
+
+    def test_init(self, datasource):
+        expected_query = (
+            "SELECT column1, column2 FROM default.table_name ORDER BY column1"
+        )
+        assert datasource._query == expected_query
+
+    @mock.patch.object(ClickHouseDatasource, "_init_client")
+    def test_init_with_filter(self, mock_init_client):
+        mock_client = MagicMock()
+        mock_init_client.return_value = mock_client
+        mock_client.query.return_value = MagicMock()
+        ds_with_filter = ClickHouseDatasource(
+            table="default.table_name",
+            dsn="clickhouse://user:password@localhost:8123/default",
+            columns=["column1", "column2"],
+            filter="label = 2 AND text IS NOT NULL",
+            order_by=(["column1"], False),
+        )
+        assert (
+            ds_with_filter._query == "SELECT column1, column2 FROM default.table_name "
+            "WHERE label = 2 AND text IS NOT NULL "
+            "ORDER BY column1"
+        )
+
+    def test_estimate_inmemory_data_size(self, datasource):
+        mock_client = mock.MagicMock()
+        datasource._init_client = MagicMock(return_value=mock_client)
+        mock_client.query.return_value.result_rows = [[12345]]
+        size = datasource.estimate_inmemory_data_size()
+        assert size == 12345
+        mock_client.query.assert_called_once_with(
+            f"SELECT SUM(byteSize(*)) AS estimate FROM ({datasource._query})"
+        )
+
+    @pytest.mark.parametrize(
+        "limit_row_count, offset_row_count, expected_query",
+        [
+            (
+                10,
+                0,
+                """
+                SELECT column1, column2 FROM default.table_name ORDER BY column1
+                FETCH FIRST 10 ROWS ONLY
+                """.strip(),
+            ),
+            (
+                1,
+                0,
+                """
+                SELECT column1, column2 FROM default.table_name ORDER BY column1
+                FETCH FIRST 1 ROW ONLY
+                """.strip(),
+            ),
+            (
+                10,
+                5,
+                """
+                SELECT column1, column2 FROM default.table_name ORDER BY column1
+                OFFSET 5 ROWS
+                FETCH NEXT 10 ROWS ONLY
+                """.strip(),
+            ),
+            (
+                1,
+                1,
+                """
+                SELECT column1, column2 FROM default.table_name ORDER BY column1
+                OFFSET 1 ROW
+                FETCH NEXT 1 ROW ONLY
+                """.strip(),
+            ),
+        ],
+    )
+    def test_build_block_query(
+        self, datasource, limit_row_count, offset_row_count, expected_query
+    ):
+        generated_query = datasource._build_block_query(
+            limit_row_count, offset_row_count
+        )
+        clean_generated_query = re.sub(r"\s+", " ", generated_query.strip())
+        clean_expected_query = re.sub(r"\s+", " ", expected_query.strip())
+        assert clean_generated_query == clean_expected_query
+
+    @pytest.mark.parametrize(
+        "columns, expected_query_part",
+        [
+            (
+                ["field1"],
+                "SELECT field1 FROM default.table_name",
+            ),
+            (["field1", "field2"], "SELECT field1, field2 FROM default.table_name"),
+            (None, "SELECT * FROM default.table_name"),
+        ],
+    )
+    def test_generate_query_columns(self, datasource, columns, expected_query_part):
+        datasource._columns = columns
+        generated_query = datasource._generate_query()
+        assert expected_query_part in generated_query
+
+    @pytest.mark.parametrize(
+        "order_by, expected_query_part",
+        [
+            ((["field1"], False), "ORDER BY field1"),
+            ((["field2"], True), "ORDER BY field2 DESC"),
+            ((["field1", "field2"], False), "ORDER BY (field1, field2)"),
+        ],
+    )
+    def test_generate_query_with_order_by(
+        self, datasource, order_by, expected_query_part
+    ):
+        datasource._order_by = order_by
+        generated_query = datasource._generate_query()
+        assert expected_query_part in generated_query
+
+    @mock.patch.object(ClickHouseDatasource, "_init_client")
+    @pytest.mark.parametrize(
+        "query_params, expected_query",
+        [
+            (
+                {},
+                "SELECT * FROM default.table_name",
+            ),
+            (
+                {
+                    "columns": ["field1"],
+                },
+                "SELECT field1 FROM default.table_name",
+            ),
+            (
+                {
+                    "columns": ["field1"],
+                    "order_by": (["field1"], False),
+                },
+                "SELECT field1 FROM default.table_name ORDER BY field1",
+            ),
+            (
+                {
+                    "columns": ["field1", "field2"],
+                    "order_by": (["field1"], True),
+                },
+                "SELECT field1, field2 FROM default.table_name ORDER BY field1 DESC",
+            ),
+            (
+                {
+                    "columns": ["field1", "field2", "field3"],
+                    "order_by": (["field1", "field2"], False),
+                },
+                "SELECT field1, field2, field3 FROM default.table_name "
+                "ORDER BY (field1, field2)",
+            ),
+            (
+                {
+                    "columns": ["field1", "field2", "field3"],
+                    "order_by": (["field1", "field2"], True),
+                },
+                "SELECT field1, field2, field3 FROM default.table_name "
+                "ORDER BY (field1, field2) DESC",
+            ),
+            (
+                {
+                    "columns": ["field1", "field2", "field3"],
+                    "order_by": (["field1", "field2", "field3"], True),
+                },
+                "SELECT field1, field2, field3 FROM default.table_name "
+                "ORDER BY (field1, field2, field3) DESC",
+            ),
+            (
+                {
+                    "columns": None,
+                    "filter": "label = 2",
+                },
+                "SELECT * FROM default.table_name WHERE label = 2",
+            ),
+            (
+                {
+                    "columns": ["field1", "field2"],
+                    "filter": "label = 2 AND text IS NOT NULL",
+                    "order_by": (["field1"], False),
+                },
+                "SELECT field1, field2 FROM default.table_name WHERE label = 2 AND "
+                "text IS NOT NULL ORDER BY field1",
+            ),
+        ],
+    )
+    def test_generate_query_full(
+        self, mock_init_client, datasource, query_params, expected_query
+    ):
+        mock_client = MagicMock()
+        mock_init_client.return_value = mock_client
+        mock_client.query.return_value = MagicMock()
+        datasource._columns = query_params.get("columns")
+        datasource._filter = query_params.get("filter")
+        datasource._order_by = query_params.get("order_by")
+        generated_query = datasource._generate_query()
+        assert expected_query == generated_query
+
+    @pytest.mark.parametrize("parallelism", [1, 2, 3, 4])
+    def test_get_read_tasks_ordered_table(self, datasource, parallelism):
+        batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field1"])
+        batch2 = pa.record_batch(
+            [pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field1"]
+        )
+        mock_stream = MagicMock()
+        mock_client = mock.MagicMock()
+        mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream
+        mock_stream.__iter__.return_value = [batch1, batch2]
+        datasource.MIN_ROWS_PER_READ_TASK = 4
+        datasource._init_client = MagicMock(return_value=mock_client)
+        datasource._get_estimate_count = MagicMock(return_value=16)
+        datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema))
+        read_tasks = datasource.get_read_tasks(parallelism)
+        expected_num_tasks = parallelism
+        assert len(read_tasks) == expected_num_tasks
+        total_rows = sum(batch.num_rows for batch in [batch1, batch2])
+        rows_per_task = total_rows // parallelism
+        extra_rows = total_rows % parallelism
+        for i, read_task in enumerate(read_tasks):
+            expected_rows = rows_per_task + (1 if i < extra_rows else 0)
+            assert read_task.metadata.num_rows == expected_rows
+
+    @pytest.mark.parametrize("parallelism", [1, 4])
+    def test_get_read_tasks_no_ordering(self, datasource, parallelism):
+        datasource._order_by = None
+        batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field2"])
+        batch2 = pa.record_batch(
+            [pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field2"]
+        )
+        mock_stream = MagicMock()
+        mock_client = mock.MagicMock()
+        mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream
+        mock_stream.__iter__.return_value = [batch1, batch2]
+        datasource.MIN_ROWS_PER_READ_TASK = 4
+        datasource._init_client = MagicMock(return_value=mock_client)
+        datasource._get_estimate_count = MagicMock(return_value=16)
+        datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema))
+        read_tasks = datasource.get_read_tasks(parallelism)
+        assert len(read_tasks) == 1
+        for i, read_task in enumerate(read_tasks):
+            assert read_task.metadata.num_rows == 16
+
+    def test_get_read_tasks_no_batches(self, datasource, mock_clickhouse_client):
+        mock_reader = mock.MagicMock()
+        mock_reader.__iter__.return_value = iter([])
+        datasource._init_client = MagicMock(return_value=mock_clickhouse_client)
+        datasource._get_estimate_count = MagicMock(return_value=0)
+        mock_block_accessor = mock.MagicMock()
+        datasource._get_sampled_estimates = MagicMock(return_value=(0, None))
+        datasource._get_sample_block = MagicMock(return_value=mock_block_accessor)
+        read_tasks = datasource.get_read_tasks(parallelism=2)
+        assert len(read_tasks) == 0
+
+    @mock.patch.object(ClickHouseDatasource, "_init_client")
+    @pytest.mark.parametrize(
+        "filter_str, expect_error, expected_error_substring",
+        [
+            ("label = 2 AND text IS NOT NULL", False, None),
+            ("some_col = 'my;string' AND another_col > 10", False, None),
+            ("AND label = 2", True, "Error: Simulated parse error"),
+            ("some_col =", True, "Error: Simulated parse error"),
+            ("col = 'someval", True, "Error: Simulated parse error"),
+            ("col = NULL", True, "Error: Simulated parse error"),
+            (
+                "col = 123; DROP TABLE foobar",
+                True,
+                "Invalid characters outside of string literals",
+            ),
+        ],
+    )
+    def test_filter_validation(
+        self, mock_init_client, filter_str, expect_error, expected_error_substring
+    ):
+        mock_client = MagicMock()
+        mock_init_client.return_value = mock_client
+        if expect_error:
+            if "Invalid characters" not in expected_error_substring:
+                mock_client.query.side_effect = Exception("Simulated parse error")
+            with pytest.raises(ValueError) as exc_info:
+                ClickHouseDatasource(
+                    table="default.table_name",
+                    dsn="clickhouse://user:password@localhost:8123/default",
+                    filter=filter_str,
+                )
+            assert expected_error_substring in str(exc_info.value), (
+                f"Expected substring '{expected_error_substring}' "
+                f"not found in: {exc_info.value}"
+            )
+        else:
+            mock_client.query.return_value = MagicMock()
+            ds = ClickHouseDatasource(
+                table="default.table_name",
+                dsn="clickhouse://user:password@localhost:8123/default",
+                filter=filter_str,
+            )
+            assert f"WHERE {filter_str}" in ds._query
+
+    @pytest.mark.parametrize("parallelism", [1, 4])
+    def test_get_read_tasks_with_filter(self, datasource, parallelism):
+        datasource._filter = "label = 2 AND text IS NOT NULL"
+        batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field2"])
+        batch2 = pa.record_batch(
+            [pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field2"]
+        )
+        mock_stream = MagicMock()
+        mock_client = mock.MagicMock()
+        mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream
+        mock_stream.__iter__.return_value = [batch1, batch2]
+        datasource.MIN_ROWS_PER_READ_TASK = 4
+        datasource._init_client = MagicMock(return_value=mock_client)
+        datasource._get_estimate_count = MagicMock(return_value=16)
+        datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema))
+        read_tasks = datasource.get_read_tasks(parallelism)
+        assert len(read_tasks) == 1
+        assert read_tasks[0].metadata.num_rows == 16
+
+    def test_filter_none(self):
+        table_name = "default.table_name"
+        dsn = "clickhouse://user:password@localhost:8123/default"
+        with mock.patch.object(ClickHouseDatasource, "_init_client") as mocked_init:
+            mock_client = MagicMock()
+            mocked_init.return_value = mock_client
+            ds = ClickHouseDatasource(table=table_name, dsn=dsn, filter=None)
+        assert "WHERE" not in ds._query
+        assert ds._filter is None
+
+
+@pytest.fixture
+def mock_clickhouse_sink_client():
+    client = MagicMock()
+    client.insert_arrow.return_value = QuerySummary({"written_rows": 3})
+    return client
+
+
+@pytest.fixture(autouse=True)
+def patch_global_get_client(mock_clickhouse_sink_client):
+    with patch(
+        "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
+    ):
+        yield
+
+
+@pytest.mark.usefixtures("ray_start_2_cpus_shared")
+class TestClickHouseDatasink:
+    @pytest.fixture
+    def datasink(self, mock_clickhouse_sink_client):
+        sink = ClickHouseDatasink(
+            table="default.test_table",
+            dsn="clickhouse+http://user:pass@localhost:8123/default",
+            mode=SinkMode.APPEND,
+            table_settings=ClickHouseTableSettings(engine="MergeTree()"),
+        )
+        return sink
+
+    @pytest.mark.parametrize(
+        "mode",
+        [
+            SinkMode.OVERWRITE,
+            SinkMode.APPEND,
+            SinkMode.CREATE,
+        ],
+    )
+    @pytest.mark.parametrize("table_exists", [True, False])
+    def test_on_write_start_modes(
+        self, datasink, mock_clickhouse_sink_client, mode, table_exists
+    ):
+        datasink._mode = mode
+        if (mode in [SinkMode.OVERWRITE, SinkMode.CREATE]) or (
+            mode == SinkMode.APPEND and not table_exists
+        ):
+            datasink._schema = pa.schema([("col1", pa.int32())])
+        with patch.object(
+            datasink, "_table_exists", return_value=table_exists
+        ) as mock_tbl_exists, patch.object(
+            datasink, "_get_existing_order_by", return_value="(prev_col)"
+        ) as mock_get_order:
+            if mode == SinkMode.CREATE and table_exists:
+                with pytest.raises(ValueError, match="already exists.*CREATE"):
+                    datasink.on_write_start()
+                mock_tbl_exists.assert_called_once()
+                mock_get_order.assert_not_called()
+                mock_clickhouse_sink_client.command.assert_not_called()
+            else:
+                datasink.on_write_start()
+                mock_tbl_exists.assert_called_once()
+                if mode == SinkMode.OVERWRITE:
+                    drop_cmd = "DROP TABLE IF EXISTS default.test_table"
+                    mock_clickhouse_sink_client.command.assert_any_call(drop_cmd)
+                    if table_exists:
+                        mock_get_order.assert_called_once()
+                    else:
+                        mock_get_order.assert_not_called()
+
+                elif mode == SinkMode.APPEND:
+                    if table_exists:
+                        mock_get_order.assert_called_once()
+                    else:
+                        mock_get_order.assert_not_called()
+                        create_cmds = [
+                            call_args[0][0]
+                            for call_args in mock_clickhouse_sink_client.command.call_args_list
+                            if "CREATE TABLE" in call_args[0][0]
+                        ]
+                        assert (
+                            len(create_cmds) == 1
+                        ), "Expected one CREATE TABLE for append + !exists."
+                elif mode == SinkMode.CREATE:
+                    if not table_exists:
+                        mock_get_order.assert_not_called()
+                        create_cmds = [
+                            call_args[0][0]
+                            for call_args in mock_clickhouse_sink_client.command.call_args_list
+                            if "CREATE TABLE" in call_args[0][0]
+                        ]
+                        assert (
+                            len(create_cmds) == 1
+                        ), "Expected one CREATE TABLE for create + !exists."
+
+    @pytest.mark.parametrize("mode", [SinkMode.OVERWRITE, SinkMode.APPEND])
+    @pytest.mark.parametrize("table_exists", [True, False])
+    @pytest.mark.parametrize("user_order_by", [None, "user_defined_col", "tuple()"])
+    def test_write_behavior(
+        self,
+        datasink,
+        mock_clickhouse_sink_client,
+        mode,
+        table_exists,
+        user_order_by,
+    ):
+        datasink._mode = mode
+        if user_order_by is not None:
+            datasink._table_settings.order_by = user_order_by
+        else:
+            datasink._table_settings.order_by = None
+        with patch.object(datasink, "_table_exists", return_value=table_exists), patch(
+            "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
+        ):
+            if not table_exists or mode == SinkMode.OVERWRITE:
+                datasink._schema = pa.schema([("col1", pa.int32())])
+            datasink.on_write_start()
+            rb = pa.record_batch([pa.array([1, 2, 3])], names=["col1"])
+            block_data = pa.Table.from_batches([rb])
+            ctx = TaskContext(1, "")
+            results = datasink.write([block_data], ctx=ctx)
+            assert results == [3]
+            mock_clickhouse_sink_client.insert_arrow.assert_called()
+
+    @pytest.mark.parametrize(
+        "schema, expected_order_by",
+        [
+            (pa.schema([]), "tuple()"),
+            (pa.schema([("ts", pa.timestamp("ns")), ("col2", pa.string())]), "ts"),
+            (pa.schema([("col1", pa.string()), ("val", pa.int64())]), "val"),
+            (pa.schema([("s1", pa.string()), ("s2", pa.large_string())]), "s1"),
+        ],
+    )
+    def test_pick_best_arrow_field_for_order_by(
+        self, datasink, mock_clickhouse_sink_client, schema, expected_order_by
+    ):
+        datasink._mode = SinkMode.OVERWRITE
+        datasink._table_settings.order_by = None
+        datasink._schema = schema
+        with patch.object(datasink, "_table_exists", return_value=False), patch(
+            "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
+        ):
+            datasink.on_write_start()
+            # Build an empty table: 0 rows
+            empty_table = pa.Table.from_batches([], schema=schema)
+            datasink.write([empty_table], ctx=None)
+            # Since we're skipping empty inserts now, we expect 0 calls:
+            mock_clickhouse_sink_client.insert_arrow.assert_not_called()
+
+    @pytest.mark.parametrize(
+        "ddl_str, expected_order_by",
+        [
+            (
+                "CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree() ORDER BY col1",
+                "col1",
+            ),
+            ("CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree()", None),
+            (
+                "CREATE TABLE default.test_table (col1 Int32) ORDER BY city ENGINE = MergeTree()",
+                "city",
+            ),
+            (
+                "CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree() PARTITION BY toYYYYMMDD(date_col)",
+                None,
+            ),
+        ],
+    )
+    def test_get_existing_order_by(
+        self, datasink, mock_clickhouse_sink_client, ddl_str, expected_order_by
+    ):
+        mock_clickhouse_sink_client.command.return_value = ddl_str
+        result = datasink._get_existing_order_by(mock_clickhouse_sink_client)
+        assert result == expected_order_by
+
+    @pytest.mark.parametrize(
+        "table_settings, schema, expected_engine, expected_order_by_part, expected_clauses",
+        [
+            (
+                ClickHouseTableSettings(),
+                pa.schema([("col1", pa.int32())]),
+                "MergeTree()",
+                "ORDER BY col1",
+                [],
+            ),
+            (
+                ClickHouseTableSettings(engine="ReplacingMergeTree()"),
+                pa.schema([("col1", pa.int32())]),
+                "ReplacingMergeTree()",
+                "ORDER BY col1",
+                [],
+            ),
+            (
+                ClickHouseTableSettings(order_by="user_col"),
+                pa.schema([("col1", pa.int32())]),
+                "MergeTree()",
+                "ORDER BY user_col",
+                [],
+            ),
+            (
+                ClickHouseTableSettings(partition_by="toYYYYMMDD(ts)"),
+                pa.schema([("ts", pa.timestamp("ns"))]),
+                "MergeTree()",
+                "ORDER BY ts",
+                ["PARTITION BY toYYYYMMDD(ts)"],
+            ),
+            (
+                ClickHouseTableSettings(primary_key="id"),
+                pa.schema([("id", pa.int64()), ("val", pa.string())]),
+                "MergeTree()",
+                "ORDER BY id",
+                ["PRIMARY KEY (id)"],
+            ),
+            (
+                ClickHouseTableSettings(settings="index_granularity=8192"),
+                pa.schema([("id", pa.int64())]),
+                "MergeTree()",
+                "ORDER BY id",
+                ["SETTINGS index_granularity=8192"],
+            ),
+            (
+                ClickHouseTableSettings(
+                    engine="SummingMergeTree()",
+                    order_by="col2",
+                    partition_by="toYYYYMMDD(ts)",
+                    primary_key="id",
+                    settings="index_granularity=8192",
+                ),
+                pa.schema(
+                    [
+                        ("id", pa.int64()),
+                        ("col2", pa.float64()),
+                        ("ts", pa.timestamp("ns")),
+                    ]
+                ),
+                "SummingMergeTree()",
+                "ORDER BY col2",
+                [
+                    "PARTITION BY toYYYYMMDD(ts)",
+                    "PRIMARY KEY (id)",
+                    "SETTINGS index_granularity=8192",
+                ],
+            ),
+        ],
+    )
+    def test_generate_create_table_sql(
+        self,
+        datasink,
+        mock_clickhouse_sink_client,
+        table_settings,
+        schema,
+        expected_engine,
+        expected_order_by_part,
+        expected_clauses,
+    ):
+        datasink._mode = SinkMode.OVERWRITE
+        datasink._table_settings = table_settings
+        datasink._schema = schema
+        with patch.object(datasink, "_table_exists", return_value=False), patch(
+            "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
+        ):
+            datasink.on_write_start()
+            arrays = []
+            for field in schema:
+                if pa.types.is_integer(field.type):
+                    arrays.append(pa.array([1, 2, 3], type=field.type))
+                elif pa.types.is_floating(field.type):
+                    arrays.append(pa.array([1.1, 2.2, 3.3], type=field.type))
+                elif pa.types.is_timestamp(field.type):
+                    arrays.append(pa.array([1, 2, 3], type=field.type))
+                else:
+                    arrays.append(pa.array(["a", "b", "c"], type=field.type))
+            block_data = pa.Table.from_arrays(arrays, names=[f.name for f in schema])
+            datasink.write([block_data], ctx=TaskContext(1, ""))
+            create_sql = None
+            for call_arg in mock_clickhouse_sink_client.command.call_args_list:
+                sql_arg = call_arg[0][0]
+                if "CREATE TABLE" in sql_arg:
+                    create_sql = sql_arg
+                    break
+            assert create_sql is not None, "No CREATE TABLE statement was generated!"
+            assert f"ENGINE = {expected_engine}" in create_sql
+            assert expected_order_by_part in create_sql
+            for clause in expected_clauses:
+                assert clause in create_sql
+
+    @pytest.mark.parametrize(
+        "provided_schema,block_fields,expected_create_columns",
+        [
+            (
+                pa.schema([("my_col", pa.float64()), ("ts", pa.timestamp("ns"))]),
+                [("my_col", pa.int32()), ("ts", pa.int64())],
+                ["`my_col` Float64", "`ts` DateTime64(3)"],
+            ),
+            (
+                pa.schema([("my_col", pa.float64()), ("col2", pa.string())]),
+                [("my_col", pa.int64()), ("col2", pa.large_string())],
+                [
+                    "`my_col` Float64",
+                    "`col2` String",
+                ],
+            ),
+            (
+                pa.schema([("id", pa.int32()), ("val", pa.string())]),
+                [("id", pa.int64()), ("val", pa.large_string())],
+                [
+                    "`id` Int32",
+                    "`val` String",
+                ],
+            ),
+            (
+                pa.schema([("f1", pa.int32()), ("f2", pa.float64())]),
+                [("f1", pa.int32()), ("f2", pa.int32())],
+                [
+                    "`f1` Int32",
+                    "`f2` Float64",
+                ],
+            ),
+        ],
+    )
+    def test_write_schema_override(
+        self,
+        datasink,
+        mock_clickhouse_sink_client,
+        provided_schema,
+        block_fields,
+        expected_create_columns,
+    ):
+        datasink._mode = SinkMode.CREATE
+        datasink._table_settings.order_by = None
+        with patch.object(datasink, "_table_exists", return_value=False), patch(
+            "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
+        ):
+            datasink._schema = provided_schema
+            datasink.on_write_start()
+            arrays = []
+            for name, typ in block_fields:
+                if pa.types.is_integer(typ):
+                    arrays.append(pa.array([1, 2, 3], type=typ))
+                elif pa.types.is_string(typ) or pa.types.is_large_string(typ):
+                    arrays.append(pa.array(["a", "b", "c"], type=typ))
+                elif pa.types.is_timestamp(typ):
+                    arrays.append(pa.array([1, 2, 3], type=typ))
+                else:
+                    arrays.append(pa.array([1.0, 2.0, 3.0], type=typ))
+            block_data = pa.Table.from_arrays(
+                arrays, names=[n for (n, _) in block_fields]
+            )
+            datasink.write([block_data], ctx=TaskContext(1, ""))
+            create_sql = None
+            for call_arg in mock_clickhouse_sink_client.command.call_args_list:
+                sql_arg = call_arg[0][0]
+                if "CREATE TABLE" in sql_arg:
+                    create_sql = sql_arg
+                    break
+            assert create_sql is not None, "Expected CREATE TABLE to be issued."
+            for expected_col_def in expected_create_columns:
+                assert expected_col_def in create_sql
+
+    @pytest.mark.parametrize(
+        "max_insert_block_rows,block_sizes,expected_insert_calls",
+        [
+            (2, [6], [3]),
+            (2, [6, 3], [3, 2]),
+            (None, [6, 3], [1, 1]),
+            (3, [3, 5, 2], [1, 2, 1]),
+        ],
+    )
+    def test_chunked_inserts(
+        self,
+        datasink,
+        mock_clickhouse_sink_client,
+        max_insert_block_rows,
+        block_sizes,
+        expected_insert_calls,
+    ):
+        datasink._mode = SinkMode.CREATE
+        datasink._schema = pa.schema([("col1", pa.int32())])
+        datasink._max_insert_block_rows = max_insert_block_rows
+        with patch.object(datasink, "_table_exists", return_value=False), patch(
+            "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
+        ):
+            datasink.on_write_start()
+            blocks = []
+            for size in block_sizes:
+                arr = pa.array(range(size), type=pa.int32())
+                block_table = pa.Table.from_arrays([arr], names=["col1"])
+                blocks.append(block_table)
+            datasink.write(blocks, ctx=TaskContext(1, ""))
+            insert_calls = [
+                call_args[0][1]
+                for call_args in mock_clickhouse_sink_client.insert_arrow.call_args_list
+            ]
+            actual_inserts = len(insert_calls)
+            assert actual_inserts == sum(expected_insert_calls), (
+                f"Expected total insert calls {sum(expected_insert_calls)}, "
+                f"got {actual_inserts}."
+            )
+            offset = 0
+            for block_idx, size in enumerate(block_sizes):
+                calls_for_block = expected_insert_calls[block_idx]
+                chunk_tables = insert_calls[offset : offset + calls_for_block]
+                offset += calls_for_block
+                total_rows = sum(tbl.num_rows for tbl in chunk_tables)
+                assert total_rows == size, (
+                    f"Block of size {size} was split incorrectly. "
+                    f"Sum of chunk sizes is {total_rows}."
+                )
+
+    @pytest.mark.parametrize(
+        "table_exists,mode,user_schema,block_fields,expected_error_regex",
+        [
+            (
+                False,
+                SinkMode.CREATE,
+                pa.schema([("id", pa.int32())]),
+                [("id", pa.int32()), ("extra_col", pa.int32())],
+                r"(ArrowInvalid|Could not convert|field names are not matching|columns not in target schema.*)",
+            ),
+            (
+                True,
+                SinkMode.OVERWRITE,
+                pa.schema([("id", pa.timestamp("ns"))]),
+                [("id", pa.int32())],
+                r"(ArrowInvalid|Could not convert|field names are not matching|columns not in target schema|Unsupported cast.*)",
+            ),
+        ],
+    )
+    def test_user_schema_block_mismatch(
+        self,
+        datasink,
+        mock_clickhouse_sink_client,
+        table_exists,
+        mode,
+        user_schema,
+        block_fields,
+        expected_error_regex,
+    ):
+        datasink._mode = mode
+        datasink._schema = user_schema
+        with patch.object(datasink, "_table_exists", return_value=table_exists), patch(
+            "clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
+        ):
+            try:
+                datasink.on_write_start()
+            except ValueError:
+                pass
+            arrays = []
+            for name, typ in block_fields:
+                arrays.append(pa.array([1, 2, 3], type=typ))
+            block_data = pa.Table.from_arrays(
+                arrays, names=[n for (n, _) in block_fields]
+            )
+            with pytest.raises(
+                (ValueError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError),
+                match=expected_error_regex,
+            ):
+                datasink.write([block_data], ctx=TaskContext(1, ""))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_csv.py

@@ -0,0 +1,213 @@
+import os
+
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from packaging.version import Version
+
+import ray
+from ray.data import Schema
+from ray.data._internal.util import rows_same
+from ray.data.block import BlockAccessor
+from ray.data.datasource.path_util import _unwrap_protocol
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+
+def df_to_csv(dataframe, path, **kwargs):
+    dataframe.to_csv(path, **kwargs)
+
+
+def test_csv_read(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    # Single file.
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path1 = os.path.join(tmp_path, "test1.csv")
+    df1.to_csv(path1, index=False)
+    ds = ray.data.read_csv(path1, partitioning=None)
+    dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
+    pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Test metadata ops.
+    assert ds.count() == 3
+    assert ds.input_files() == [_unwrap_protocol(path1)]
+    assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
+
+    # Two files, override_num_blocks=2.
+    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+    path2 = os.path.join(tmp_path, "test2.csv")
+    df2.to_csv(path2, index=False)
+    ds = ray.data.read_csv([path1, path2], override_num_blocks=2, partitioning=None)
+    dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
+    df = pd.concat([df1, df2], ignore_index=True)
+    pd.testing.assert_frame_equal(df.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Test metadata ops.
+    for entry in ds._execute().blocks:
+        assert (
+            # pyrefly: ignore[no-matching-overload]
+            BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
+            == entry.metadata.size_bytes
+        )
+
+    # Three files, override_num_blocks=2.
+    df3 = pd.DataFrame({"one": [7, 8, 9], "two": ["h", "i", "j"]})
+    path3 = os.path.join(tmp_path, "test3.csv")
+    df3.to_csv(path3, index=False)
+    ds = ray.data.read_csv(
+        [path1, path2, path3],
+        override_num_blocks=2,
+        partitioning=None,
+    )
+    df = pd.concat([df1, df2, df3], ignore_index=True)
+    dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
+    pd.testing.assert_frame_equal(df.astype(dsdf.dtypes.to_dict()), dsdf)
+
+
+def test_csv_write(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    input_df = pd.DataFrame({"id": [0]})
+    ds = ray.data.from_blocks([input_df])
+
+    ds.write_csv(tmp_path)
+
+    output_df = pd.concat(
+        [
+            pd.read_csv(os.path.join(tmp_path, filename))
+            for filename in os.listdir(tmp_path)
+        ]
+    )
+    assert rows_same(input_df, output_df)
+
+
+@pytest.mark.parametrize("override_num_blocks", [None, 2])
+def test_csv_roundtrip(
+    ray_start_regular_shared,
+    tmp_path,
+    override_num_blocks,
+    target_max_block_size_infinite_or_default,
+):
+    df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+
+    ds = ray.data.from_pandas([df], override_num_blocks=override_num_blocks)
+    ds.write_csv(tmp_path)
+
+    ds2 = ray.data.read_csv(tmp_path)
+    ds2df = ds2.to_pandas()
+    assert rows_same(ds2df, df)
+    for entry in ds2._execute().blocks:
+        # pyrefly: ignore[no-matching-overload]
+        assert (
+            BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
+            == entry.metadata.size_bytes
+        )
+
+
+def test_csv_read_invalid_format(ray_start_regular_shared, tmp_path):
+    df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+
+    # Setup: CSV and Parquet files in the same directory.
+    csv_path = os.path.join(tmp_path, "test.csv")
+    df.to_csv(csv_path, index=False)
+
+    table = pa.Table.from_pandas(df)
+    parquet_path = os.path.join(tmp_path, "test.parquet")
+    pq.write_table(table, parquet_path)
+
+    # Test 1: CSV parser should fail on Parquet file.
+    error_message = "Failed to read CSV file"
+    with pytest.raises(ValueError, match=error_message):
+        ray.data.read_csv(parquet_path).materialize()
+
+    # Test 2: CSV parser should fail when directory contains non-CSV files.
+    with pytest.raises(ValueError, match=error_message):
+        ray.data.read_csv(tmp_path).materialize()
+
+
+def test_csv_read_no_header(ray_start_regular_shared, tmp_path):
+    from pyarrow import csv
+
+    file_path = os.path.join(tmp_path, "test.csv")
+    df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    df.to_csv(file_path, index=False, header=False)
+    ds = ray.data.read_csv(
+        file_path,
+        read_options=csv.ReadOptions(column_names=["one", "two"]),
+    )
+    out_df = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(out_df.dtypes.to_dict()), out_df)
+
+
+def test_csv_read_with_column_type_specified(ray_start_regular_shared, tmp_path):
+    from pyarrow import csv
+
+    file_path = os.path.join(tmp_path, "test.csv")
+    df = pd.DataFrame({"one": [1, 2, 3e1], "two": ["a", "b", "c"]})
+    df.to_csv(file_path, index=False)
+
+    # Incorrect to parse scientific notation in int64 as PyArrow represents
+    # it as double.
+    with pytest.raises(ValueError):
+        ray.data.read_csv(
+            file_path,
+            convert_options=csv.ConvertOptions(
+                column_types={"one": "int64", "two": "string"}
+            ),
+        ).schema()
+
+    # Parsing scientific notation in double should work.
+    ds = ray.data.read_csv(
+        file_path,
+        convert_options=csv.ConvertOptions(
+            column_types={"one": "float64", "two": "string"}
+        ),
+    )
+    expected_df = pd.DataFrame({"one": [1.0, 2.0, 30.0], "two": ["a", "b", "c"]})
+    actual_df = ds.to_pandas()
+    pd.testing.assert_frame_equal(
+        expected_df.astype(actual_df.dtypes.to_dict()), actual_df
+    )
+
+
+@pytest.mark.skipif(
+    Version(pa.__version__) < Version("7.0.0"),
+    reason="invalid_row_handler was added in pyarrow 7.0.0",
+)
+def test_csv_invalid_file_handler(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    from pyarrow import csv
+
+    invalid_txt = "f1,f2\n2,3\nx\n4,5"
+    invalid_file = os.path.join(tmp_path, "invalid.csv")
+    with open(invalid_file, "wt") as f:
+        f.write(invalid_txt)
+
+    ray.data.read_csv(
+        invalid_file,
+        parse_options=csv.ParseOptions(
+            delimiter=",", invalid_row_handler=lambda i: "skip"
+        ),
+    )
+
+
+def test_read_example_data(ray_start_regular_shared, tmp_path):
+    ds = ray.data.read_csv("example://iris.csv")
+    assert ds.count() == 150
+    assert ds.take(1) == [
+        {
+            "sepal.length": 5.1,
+            "sepal.width": 3.5,
+            "petal.length": 1.4,
+            "petal.width": 0.2,
+            "variety": "Setosa",
+        }
+    ]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_daft.py

@@ -0,0 +1,59 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+
+@pytest.fixture(scope="module")
+def ray_start(request):
+    """Initialize Ray for Daft tests."""
+    import ray
+
+    try:
+        yield ray.init(
+            num_cpus=16,
+        )
+    finally:
+        ray.shutdown()
+
+
+def test_daft_round_trip(ray_start):
+    import daft
+
+    import ray
+
+    data = {
+        "int_col": list(range(128)),
+        "str_col": [str(i) for i in range(128)],
+        "nested_list_col": [[i] * 3 for i in range(128)],
+        "tensor_col": [np.array([[i] * 3] * 3) for i in range(128)],
+    }
+    df = daft.from_pydict(data)
+    ds = ray.data.from_daft(df)
+
+    # Ray stores data in Arrow format, so to_pandas() returns Arrow-backed
+    # dtypes (e.g. int64[pyarrow]) while Daft may return numpy dtypes.
+    # Compare values only, not dtypes.
+    pd.testing.assert_frame_equal(ds.to_pandas(), df.to_pandas(), check_dtype=False)
+
+    df2 = ds.to_daft()
+    df_pandas = df.to_pandas()
+    df2_pandas = df2.to_pandas()
+
+    for c in data.keys():
+        # NOTE: tensor behavior on round-trip is different because Ray Data provides
+        # Daft with more information about a column being a fixed-shape-tensor.
+        #
+        # Hence the Pandas representation of `df1` is "just" an object column, but
+        # `df2` knows that this is actually a numpy fixed shaped tensor column
+        if c == "tensor_col":
+            original = np.array(list(df_pandas[c]))
+            roundtripped = np.array(list(df2_pandas[c]))
+            np.testing.assert_array_equal(original, roundtripped)
+        else:
+            pd.testing.assert_series_equal(df_pandas[c], df2_pandas[c])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_databricks_credentials.py

@@ -0,0 +1,289 @@
+"""Unit tests for Databricks credential providers."""
+
+import os
+from unittest import mock
+
+import pytest
+
+from ray.data._internal.datasource.databricks_credentials import (
+    DatabricksCredentialProvider,
+    DatabricksTableCredentialConfig,
+    EnvironmentCredentialProvider,
+    StaticCredentialProvider,
+    UnityCatalogCredentialConfig,
+    resolve_credential_provider,
+)
+
+SAMPLE_TOKEN = "dapi_test_token_abc123"
+SAMPLE_HOST = "https://my-workspace.cloud.databricks.com"
+SAMPLE_URL = "https://uc-workspace.databricks.com"
+ALT_TOKEN = "dapi_alt_token_xyz789"
+ALT_HOST = "https://alt-workspace.databricks.com"
+
+
+class TestDatabricksCredentialProvider:
+    """Tests for the abstract DatabricksCredentialProvider base class."""
+
+    def test_cannot_instantiate_abstract_class(self):
+        """Verify DatabricksCredentialProvider cannot be instantiated directly."""
+        with pytest.raises(TypeError, match="Can't instantiate abstract class"):
+            DatabricksCredentialProvider()
+
+    def test_abstract_methods_defined(self):
+        """Verify all abstract methods are defined."""
+        abstract_methods = DatabricksCredentialProvider.__abstractmethods__
+        assert "get_token" in abstract_methods
+        assert "get_host" in abstract_methods
+        assert "invalidate" in abstract_methods
+
+
+class TestStaticCredentialProvider:
+    """Tests for StaticCredentialProvider."""
+
+    def test_init_with_valid_token_and_host(self):
+        """Test successful initialization with token and host."""
+        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
+        assert provider.get_token() == SAMPLE_TOKEN
+        assert provider.get_host() == SAMPLE_HOST
+
+    @pytest.mark.parametrize(
+        "token,host,expected_error",
+        [
+            ("", SAMPLE_HOST, "Token cannot be empty"),
+            (None, SAMPLE_HOST, "Token cannot be empty"),
+            (SAMPLE_TOKEN, "", "Host cannot be empty"),
+            (SAMPLE_TOKEN, None, "Host cannot be empty"),
+        ],
+    )
+    def test_init_with_invalid_inputs_raises_error(self, token, host, expected_error):
+        """Test that invalid token or host raises ValueError."""
+        with pytest.raises(ValueError, match=expected_error):
+            StaticCredentialProvider(token=token, host=host)
+
+    def test_invalidate_is_noop(self):
+        """Test that invalidate doesn't affect the static token."""
+        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
+        provider.invalidate()
+        assert provider.get_token() == SAMPLE_TOKEN
+        assert provider.get_host() == SAMPLE_HOST
+
+    def test_get_token_returns_same_value(self):
+        """Test that get_token always returns the same value."""
+        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
+        assert provider.get_token() == SAMPLE_TOKEN
+        assert provider.get_token() == SAMPLE_TOKEN
+
+
+class TestEnvironmentCredentialProvider:
+    """Tests for EnvironmentCredentialProvider."""
+
+    def test_get_token_from_env(self):
+        """Test get_token reads from environment variable."""
+        with mock.patch.dict(
+            os.environ,
+            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
+        ):
+            provider = EnvironmentCredentialProvider()
+            assert provider.get_token() == SAMPLE_TOKEN
+
+    def test_get_host_from_env(self):
+        """Test get_host reads from environment variable."""
+        with mock.patch.dict(
+            os.environ,
+            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
+        ):
+            provider = EnvironmentCredentialProvider()
+            assert provider.get_host() == SAMPLE_HOST
+
+    @pytest.mark.parametrize(
+        "env_vars,expected_error",
+        [
+            ({"DATABRICKS_HOST": SAMPLE_HOST}, "DATABRICKS_TOKEN.*not set"),
+            (
+                {"DATABRICKS_TOKEN": SAMPLE_TOKEN},
+                "set environment variable.*DATABRICKS_HOST",
+            ),
+        ],
+    )
+    def test_init_raises_when_env_var_not_set(self, env_vars, expected_error):
+        """Test __init__ raises ValueError when required env var is not set."""
+        with mock.patch.dict(os.environ, env_vars, clear=True):
+            with pytest.raises(ValueError, match=expected_error):
+                EnvironmentCredentialProvider()
+
+    def test_host_detected_from_databricks_runtime(self):
+        """Test host is detected from Databricks runtime when env var not set."""
+        detected_host = "detected-host.databricks.com"
+        with (
+            mock.patch.dict(os.environ, {"DATABRICKS_TOKEN": SAMPLE_TOKEN}, clear=True),
+            mock.patch.object(
+                EnvironmentCredentialProvider,
+                "_detect_databricks_host",
+                return_value=detected_host,
+            ),
+        ):
+            provider = EnvironmentCredentialProvider()
+            assert provider.get_host() == detected_host
+
+    def test_custom_env_var_names(self):
+        """Test using custom environment variable names."""
+        with mock.patch.dict(
+            os.environ, {"MY_TOKEN": SAMPLE_TOKEN, "MY_HOST": SAMPLE_HOST}
+        ):
+            provider = EnvironmentCredentialProvider(
+                token_env_var="MY_TOKEN", host_env_var="MY_HOST"
+            )
+            assert provider.get_token() == SAMPLE_TOKEN
+            assert provider.get_host() == SAMPLE_HOST
+
+    def test_invalidate_refreshes_token_from_env(self):
+        """Test that invalidate re-reads token from environment."""
+        refreshed_token = "dapi_refreshed_token_456"
+        with mock.patch.dict(
+            os.environ,
+            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
+        ):
+            provider = EnvironmentCredentialProvider()
+            assert provider.get_token() == SAMPLE_TOKEN
+
+            # Simulate external token refresh
+            os.environ["DATABRICKS_TOKEN"] = refreshed_token
+            provider.invalidate()
+            assert provider.get_token() == refreshed_token
+
+    def test_invalidate_keeps_token_if_env_unset(self):
+        """Test that invalidate keeps existing token if env var is unset."""
+        with mock.patch.dict(
+            os.environ,
+            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
+        ):
+            provider = EnvironmentCredentialProvider()
+
+            # Remove env var after initialization
+            del os.environ["DATABRICKS_TOKEN"]
+            provider.invalidate()
+            # Should keep the old token rather than failing
+            assert provider.get_token() == SAMPLE_TOKEN
+
+
+class TestDatabricksTableCredentialConfig:
+    """Tests for DatabricksTableCredentialConfig and resolve_credential_provider."""
+
+    def test_resolve_with_explicit_provider(self):
+        """Test that explicit credential_provider is returned as-is."""
+        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
+        config = DatabricksTableCredentialConfig(credential_provider=provider)
+        result = resolve_credential_provider(config)
+        assert result is provider
+
+    @pytest.mark.parametrize("credential_provider_arg", [None, "no_arg"])
+    def test_resolve_with_none_returns_environment_provider(
+        self, credential_provider_arg
+    ):
+        """Test that EnvironmentCredentialProvider is returned when none provided."""
+        with mock.patch.dict(
+            os.environ,
+            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
+        ):
+            if credential_provider_arg == "no_arg":
+                config = DatabricksTableCredentialConfig()
+            else:
+                config = DatabricksTableCredentialConfig(
+                    credential_provider=credential_provider_arg
+                )
+            result = resolve_credential_provider(config)
+            assert isinstance(result, EnvironmentCredentialProvider)
+
+
+class TestUnityCatalogCredentialConfig:
+    """Tests for UnityCatalogCredentialConfig and resolve_credential_provider."""
+
+    def test_resolve_with_explicit_provider(self):
+        """Test that explicit credential_provider is returned as-is."""
+        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
+        config = UnityCatalogCredentialConfig(credential_provider=provider)
+        result = resolve_credential_provider(config)
+        assert result is provider
+
+    def test_resolve_with_explicit_provider_ignores_url_and_token(self):
+        """Test that url/token are ignored when credential_provider is given."""
+        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
+        config = UnityCatalogCredentialConfig(
+            credential_provider=provider, url=ALT_HOST, token=ALT_TOKEN
+        )
+        result = resolve_credential_provider(config)
+        assert result is provider
+
+    def test_resolve_with_url_and_token(self):
+        """Test that url and token create a StaticCredentialProvider."""
+        config = UnityCatalogCredentialConfig(url=SAMPLE_URL, token=SAMPLE_TOKEN)
+        result = resolve_credential_provider(config)
+        assert isinstance(result, StaticCredentialProvider)
+        assert result.get_token() == SAMPLE_TOKEN
+        assert result.get_host() == SAMPLE_URL
+
+    @pytest.mark.parametrize(
+        "kwargs",
+        [
+            {},
+            {"url": SAMPLE_URL},
+            {"token": SAMPLE_TOKEN},
+        ],
+        ids=["no_args", "only_url", "only_token"],
+    )
+    def test_config_raises_with_incomplete_args(self, kwargs):
+        """Test that ValueError is raised when args are missing or incomplete."""
+        config = UnityCatalogCredentialConfig(**kwargs)
+        with pytest.raises(ValueError, match="Either 'credential_provider' or both"):
+            resolve_credential_provider(config)
+
+    @pytest.mark.parametrize(
+        "url,token",
+        [
+            ("", SAMPLE_TOKEN),
+            (SAMPLE_URL, ""),
+        ],
+        ids=["empty_url", "empty_token"],
+    )
+    def test_resolve_with_empty_string_raises(self, url, token):
+        """Test that empty strings for url or token raise ValueError."""
+        config = UnityCatalogCredentialConfig(url=url, token=token)
+        with pytest.raises(ValueError):
+            resolve_credential_provider(config)
+
+
+class TestCredentialProviderSerialization:
+    """Tests for credential provider serialization (needed for Ray workers)."""
+
+    @pytest.mark.parametrize(
+        "provider_type,expected_token,expected_host",
+        [
+            ("static", SAMPLE_TOKEN, SAMPLE_HOST),
+            ("environment", SAMPLE_TOKEN, SAMPLE_HOST),
+        ],
+    )
+    def test_provider_is_picklable(self, provider_type, expected_token, expected_host):
+        """Verify credential providers can be pickled and unpickled."""
+        import pickle
+
+        with mock.patch.dict(
+            os.environ,
+            {"DATABRICKS_TOKEN": expected_token, "DATABRICKS_HOST": expected_host},
+        ):
+            if provider_type == "static":
+                provider = StaticCredentialProvider(
+                    token=expected_token, host=expected_host
+                )
+            else:
+                provider = EnvironmentCredentialProvider()
+
+            pickled = pickle.dumps(provider)
+            unpickled = pickle.loads(pickled)
+            assert unpickled.get_token() == expected_token
+            assert unpickled.get_host() == expected_host
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_databricks_uc_datasource.py

@@ -0,0 +1,700 @@
+"""Tests for Databricks Unity Catalog datasource."""
+
+import json
+import os
+import re
+import tempfile
+import uuid
+from contextlib import contextmanager
+from dataclasses import dataclass
+from unittest import mock
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+import ray.cloudpickle as pickle
+from ray.data._internal.datasource.databricks_credentials import (
+    DatabricksCredentialProvider,
+    StaticCredentialProvider,
+)
+from ray.data._internal.datasource.databricks_uc_datasource import (
+    DatabricksUCDatasource,
+)
+from ray.data._internal.util import rows_same
+from ray.data.tests.datasource.databricks_test_utils import (
+    MockResponse,
+    RefreshableCredentialProvider,
+)
+from ray.tests.conftest import *  # noqa
+
+# =============================================================================
+# Dataclasses for mock objects
+# =============================================================================
+
+
+@dataclass
+class MockChunk:
+    """Mock chunk data for testing."""
+
+    index: int
+    row_count: int
+    byte_count: int
+    data: bytes
+
+
+# =============================================================================
+# Mock credential providers for testing
+# =============================================================================
+
+
+class TokenTrackingProvider(DatabricksCredentialProvider):
+    """A credential provider that returns incrementing tokens to track fetches."""
+
+    def __init__(self):
+        self.token_fetch_count = 0
+
+    def get_token(self) -> str:
+        self.token_fetch_count += 1
+        return f"token_{self.token_fetch_count}"
+
+    def get_host(self) -> str:
+        return "test_host"
+
+    def invalidate(self) -> None:
+        pass
+
+
+# =============================================================================
+# Pytest fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def databricks_env():
+    """Fixture that sets up Databricks environment variables."""
+    with mock.patch.dict(
+        os.environ,
+        {"DATABRICKS_HOST": "test_host", "DATABRICKS_TOKEN": "test_token"},
+    ):
+        yield
+
+
+@pytest.fixture
+def refreshable_credential_provider():
+    """Fixture that provides a refreshable credential provider."""
+    return RefreshableCredentialProvider(host="test_host")
+
+
+@pytest.fixture
+def token_tracking_provider():
+    """Fixture that provides a token tracking credential provider."""
+    return TokenTrackingProvider()
+
+
+@pytest.fixture
+def requests_mocker():
+    """Fixture that mocks requests.get and requests.post."""
+    with mock.patch("requests.get") as mock_get:
+        with mock.patch("requests.post") as mock_post:
+            yield {"get": mock_get, "post": mock_post}
+
+
+@pytest.fixture
+def test_data():
+    """Fixture that provides test DataFrame and configuration."""
+    return {
+        "expected_df": pd.DataFrame(
+            {
+                "c1": range(10000),
+                "c2": [f"str{i}" for i in range(10000)],
+            }
+        ),
+        "token": "test_token",
+        "warehouse_id": "test_warehouse_id",
+        "catalog": "catalog1",
+        "schema": "db1",
+        "query": "select * from table1",
+        "rows_per_chunk": 700,
+    }
+
+
+# =============================================================================
+# Helper functions
+# =============================================================================
+
+
+def create_mock_chunks(df: pd.DataFrame, rows_per_chunk: int) -> list[MockChunk]:
+    """Create mock chunks from a DataFrame."""
+    chunks = []
+    num_rows = len(df)
+    cur_pos = 0
+    index = 0
+
+    while cur_pos < num_rows:
+        chunk_rows = min(rows_per_chunk, num_rows - cur_pos)
+        chunk_df = df[cur_pos : cur_pos + chunk_rows]
+        chunk_pa_table = pa.Table.from_pandas(chunk_df)
+
+        sink = pa.BufferOutputStream()
+        with pa.ipc.new_stream(sink, chunk_pa_table.schema) as writer:
+            writer.write_table(chunk_pa_table)
+
+        chunks.append(
+            MockChunk(
+                index=index,
+                row_count=chunk_rows,
+                byte_count=len(sink.getvalue()),
+                data=sink.getvalue(),
+            )
+        )
+        index += 1
+        cur_pos += rows_per_chunk
+
+    return chunks
+
+
+# =============================================================================
+# Test classes
+# =============================================================================
+
+
+class TestDatabricksUCDatasourceIntegration:
+    """Integration tests for DatabricksUCDatasource."""
+
+    _MOCK_ENV_VAR = "RAY_DATABRICKS_UC_DATASOURCE_READ_FN_MOCK_TEST_SETUP_FN_PATH"
+
+    @contextmanager
+    def _setup_mock(self, test_data: dict, mock_chunks: list[MockChunk]):
+        """Set up mocks for integration tests."""
+        chunk_meta_json = [
+            {
+                "chunk_index": chunk.index,
+                "row_count": chunk.row_count,
+                "byte_count": chunk.byte_count,
+            }
+            for chunk in mock_chunks
+        ]
+        chunk_meta_json.reverse()
+        valid_statement_ids = set()
+
+        def request_post_mock(url, data=None, json=None, **kwargs):
+            import json as jsonlib
+
+            headers = kwargs["headers"]
+
+            if url == "https://test_shard/api/2.0/sql/statements/":
+                assert headers == {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {test_data['token']}",
+                }
+                assert jsonlib.loads(data) == {
+                    "statement": test_data["query"],
+                    "warehouse_id": test_data["warehouse_id"],
+                    "wait_timeout": "0s",
+                    "disposition": "EXTERNAL_LINKS",
+                    "format": "ARROW_STREAM",
+                    "catalog": test_data["catalog"],
+                    "schema": test_data["schema"],
+                }
+
+                statement_id = uuid.uuid4().hex
+                valid_statement_ids.add(statement_id)
+
+                return MockResponse(
+                    status_code=200,
+                    content=b"",
+                    _json_data={
+                        "statement_id": statement_id,
+                        "status": {"state": "PENDING"},
+                    },
+                )
+
+            assert False, "Invalid request."
+
+        def request_get_mock(url, params=None, **kwargs):
+            headers = kwargs["headers"]
+
+            if match := re.match(
+                r"^https://test_shard/api/2\.0/sql/statements/([^/]*)/$", url
+            ):
+                statement_id = match.group(1)
+                assert headers == {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {test_data['token']}",
+                }
+                assert statement_id in valid_statement_ids
+
+                return MockResponse(
+                    status_code=200,
+                    _json_data={
+                        "status": {"state": "SUCCEEDED"},
+                        "manifest": {
+                            "truncated": False,
+                            "chunks": chunk_meta_json,
+                        },
+                    },
+                )
+
+            if match := re.match(
+                r"^https://test_shard/api/2\.0/sql/"
+                r"statements/([^/]*)/result/chunks/([^/]*)$",
+                url,
+            ):
+                assert headers == {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {test_data['token']}",
+                }
+
+                chunk_index = match.group(2)
+                external_link = f"https://test_external_link/{chunk_index}"
+                return MockResponse(
+                    status_code=200,
+                    _json_data={"external_links": [{"external_link": external_link}]},
+                )
+
+            if match := re.match(r"^https://test_external_link/([^/]*)$", url):
+                assert headers is None
+                chunk_index = int(match.group(1))
+                return MockResponse(
+                    status_code=200,
+                    content=mock_chunks[chunk_index].data,
+                )
+
+            assert False, "Invalid request."
+
+        with (
+            mock.patch("requests.get", request_get_mock),
+            mock.patch("requests.post", request_post_mock),
+            mock.patch.dict(
+                os.environ,
+                {
+                    "DATABRICKS_HOST": "test_shard",
+                    "DATABRICKS_TOKEN": test_data["token"],
+                },
+            ),
+        ):
+            yield
+
+    @contextmanager
+    def _setup_integration_test(self, test_data: dict):
+        """Set up complete integration test environment with mocks and Ray."""
+        mock_chunks = create_mock_chunks(
+            test_data["expected_df"], test_data["rows_per_chunk"]
+        )
+
+        setup_mock_fn_path = os.path.join(tempfile.mkdtemp(), "setup_mock_fn.pkl")
+        with open(setup_mock_fn_path, "wb") as fp:
+            pickle.dump(lambda: self._setup_mock(test_data, mock_chunks), fp)
+
+        with (
+            self._setup_mock(test_data, mock_chunks),
+            mock.patch.dict(os.environ, {self._MOCK_ENV_VAR: setup_mock_fn_path}),
+        ):
+            ray.shutdown()
+            ray.init()
+            yield
+
+    def test_read_with_table_name(self, test_data):
+        """Test reading data using table name."""
+        with self._setup_integration_test(test_data):
+            result = ray.data.read_databricks_tables(
+                warehouse_id=test_data["warehouse_id"],
+                table="table1",
+                catalog=test_data["catalog"],
+                schema=test_data["schema"],
+                override_num_blocks=5,
+            ).to_pandas()
+
+            assert rows_same(result, test_data["expected_df"])
+
+    def test_read_with_sql_query(self, test_data):
+        """Test reading data using SQL query."""
+        with self._setup_integration_test(test_data):
+            result = ray.data.read_databricks_tables(
+                warehouse_id=test_data["warehouse_id"],
+                query=test_data["query"],
+                catalog=test_data["catalog"],
+                schema=test_data["schema"],
+                override_num_blocks=5,
+            ).to_pandas()
+
+            assert rows_same(result, test_data["expected_df"])
+
+    @pytest.mark.parametrize("num_blocks", [5, 100])
+    def test_read_with_different_parallelism(self, test_data, num_blocks):
+        """Test reading data with different parallelism settings."""
+        with self._setup_integration_test(test_data):
+            result = ray.data.read_databricks_tables(
+                warehouse_id=test_data["warehouse_id"],
+                query=test_data["query"],
+                catalog=test_data["catalog"],
+                schema=test_data["schema"],
+                override_num_blocks=num_blocks,
+            ).to_pandas()
+
+            assert rows_same(result, test_data["expected_df"])
+
+
+class TestDatabricksUCDatasourceCredentials:
+    """Tests for credential provider handling."""
+
+    def test_schema_name_does_not_shadow_datasource_fields(self, requests_mocker):
+        """Test that schema name is stored without using the `schema` attribute.
+        This is a regression test for https://github.com/ray-project/ray/issues/46481.
+        """
+        requests_mocker["post"].return_value = mock.Mock(
+            status_code=200,
+            raise_for_status=lambda: None,
+            json=lambda: {"statement_id": "test_stmt", "status": {"state": "PENDING"}},
+        )
+        requests_mocker["get"].return_value = mock.Mock(
+            status_code=200,
+            raise_for_status=lambda: None,
+            json=lambda: {
+                "status": {"state": "SUCCEEDED"},
+                "manifest": {"truncated": False, "chunks": []},
+            },
+        )
+
+        provider = StaticCredentialProvider(token="my_provider_token", host="test_host")
+        datasource = DatabricksUCDatasource(
+            warehouse_id="test_warehouse",
+            catalog="test_catalog",
+            schema="test_schema",
+            query="SELECT 1",
+            credential_provider=provider,
+        )
+
+        assert datasource.schema_name == "test_schema"
+        assert "schema" not in datasource.__dict__
+
+        call_kwargs = requests_mocker["post"].call_args[1]
+        payload = json.loads(call_kwargs["data"])
+        assert payload["schema"] == "test_schema"
+
+    def test_with_credential_provider(self, requests_mocker):
+        """Test DatabricksUCDatasource with credential_provider parameter."""
+        requests_mocker["post"].return_value = mock.Mock(
+            status_code=200,
+            raise_for_status=lambda: None,
+            json=lambda: {"statement_id": "test_stmt", "status": {"state": "PENDING"}},
+        )
+        requests_mocker["get"].return_value = mock.Mock(
+            status_code=200,
+            raise_for_status=lambda: None,
+            json=lambda: {
+                "status": {"state": "SUCCEEDED"},
+                "manifest": {"truncated": False},
+            },
+        )
+
+        provider = StaticCredentialProvider(token="my_provider_token", host="test_host")
+        _datasource = DatabricksUCDatasource(
+            warehouse_id="test_warehouse",
+            catalog="test_catalog",
+            schema="test_schema",
+            query="SELECT 1",
+            credential_provider=provider,
+        )
+
+        # Verify the token from provider was used in requests
+        call_kwargs = requests_mocker["post"].call_args[1]
+        assert "Authorization" in call_kwargs["headers"]
+        assert "Bearer my_provider_token" in call_kwargs["headers"]["Authorization"]
+
+    def test_fresh_token_per_request(self, requests_mocker, token_tracking_provider):
+        """Test that fresh tokens are fetched for each request during polling."""
+        tokens_used = []
+
+        def capture_post(url, *args, **kwargs):
+            tokens_used.append(kwargs["headers"]["Authorization"])
+            return mock.Mock(
+                status_code=200,
+                raise_for_status=lambda: None,
+                json=lambda: {
+                    "statement_id": "test_stmt",
+                    "status": {"state": "PENDING"},
+                },
+            )
+
+        poll_count = [0]
+
+        def capture_get(url, *args, **kwargs):
+            tokens_used.append(kwargs["headers"]["Authorization"])
+            poll_count[0] += 1
+            state = "PENDING" if poll_count[0] < 3 else "SUCCEEDED"
+            return mock.Mock(
+                status_code=200,
+                raise_for_status=lambda: None,
+                json=lambda: {
+                    "status": {"state": state},
+                    "manifest": {"truncated": False, "chunks": []},
+                },
+            )
+
+        requests_mocker["post"].side_effect = capture_post
+        requests_mocker["get"].side_effect = capture_get
+
+        DatabricksUCDatasource(
+            warehouse_id="test_warehouse",
+            catalog="test_catalog",
+            schema="test_schema",
+            query="SELECT 1",
+            credential_provider=token_tracking_provider,
+        )
+
+        # Verify fresh token was fetched for each request:
+        # 1 POST (statement creation) + 3 GETs (polling)
+        assert token_tracking_provider.token_fetch_count == 4
+        assert tokens_used == [
+            "Bearer token_1",  # POST
+            "Bearer token_2",  # GET poll 1
+            "Bearer token_3",  # GET poll 2
+            "Bearer token_4",  # GET poll 3
+        ]
+
+
+class TestDatabricksUCDatasource401Retry:
+    """Tests for 401 retry behavior."""
+
+    def test_401_during_initial_post(
+        self, requests_mocker, refreshable_credential_provider
+    ):
+        """Test that 401 during initial POST triggers credential invalidation and retry."""
+        post_call_count = [0]
+        post_headers_captured = []
+
+        def post_side_effect(url, *args, **kwargs):
+            post_call_count[0] += 1
+            headers = kwargs.get("headers", {})
+            post_headers_captured.append(headers.get("Authorization", ""))
+
+            # First POST returns 401
+            if post_call_count[0] == 1:
+                return mock.Mock(status_code=401)
+            # Retry succeeds
+            return mock.Mock(
+                status_code=200,
+                json=lambda: {
+                    "statement_id": "test_stmt",
+                    "status": {"state": "SUCCEEDED"},
+                    "manifest": {"truncated": False, "chunks": []},
+                },
+            )
+
+        requests_mocker["post"].side_effect = post_side_effect
+
+        DatabricksUCDatasource(
+            warehouse_id="test_warehouse",
+            catalog="test_catalog",
+            schema="test_schema",
+            query="SELECT 1",
+            credential_provider=refreshable_credential_provider,
+        )
+
+        # Verify retry occurred
+        assert (
+            post_call_count[0] == 2
+        ), "Expected POST to be called twice (initial + retry)"
+
+        # Verify invalidate was called
+        assert refreshable_credential_provider.invalidate_count == 1
+
+        # Verify first request used expired token, retry used refreshed token
+        assert "expired_token" in post_headers_captured[0]
+        assert "refreshed_token" in post_headers_captured[1]
+
+    def test_401_during_polling(self, requests_mocker, refreshable_credential_provider):
+        """Test that 401 during polling triggers credential invalidation and retry."""
+        poll_call_count = [0]
+        poll_headers_captured = []
+
+        requests_mocker["post"].return_value = mock.Mock(
+            status_code=200,
+            json=lambda: {
+                "statement_id": "test_stmt",
+                "status": {"state": "PENDING"},
+            },
+        )
+
+        def get_side_effect(url, *args, **kwargs):
+            poll_call_count[0] += 1
+            headers = kwargs.get("headers", {})
+            poll_headers_captured.append(headers.get("Authorization", ""))
+
+            # First poll returns 401 with expired token
+            if poll_call_count[0] == 1:
+                return mock.Mock(status_code=401)
+            # Retry succeeds
+            return mock.Mock(
+                status_code=200,
+                json=lambda: {
+                    "status": {"state": "SUCCEEDED"},
+                    "manifest": {"truncated": False, "chunks": []},
+                },
+            )
+
+        requests_mocker["get"].side_effect = get_side_effect
+
+        DatabricksUCDatasource(
+            warehouse_id="test_warehouse",
+            catalog="test_catalog",
+            schema="test_schema",
+            query="SELECT 1",
+            credential_provider=refreshable_credential_provider,
+        )
+
+        # Verify retry occurred
+        assert (
+            poll_call_count[0] == 2
+        ), "Expected GET to be called twice (initial + retry)"
+
+        # Verify invalidate was called once
+        assert refreshable_credential_provider.invalidate_count == 1
+
+        # Verify first request used expired token, retry used refreshed token
+        assert "expired_token" in poll_headers_captured[0]
+        assert "refreshed_token" in poll_headers_captured[1]
+
+    def test_401_during_chunk_fetch(
+        self, requests_mocker, refreshable_credential_provider
+    ):
+        """Test that 401 during chunk fetch triggers credential invalidation and retry."""
+        chunk_fetch_count = [0]
+        chunk_fetch_headers = []
+
+        # Create Arrow data for external URL response
+        table = pa.Table.from_pydict({"col1": [1, 2, 3]})
+        sink = pa.BufferOutputStream()
+        with pa.ipc.new_stream(sink, table.schema) as writer:
+            writer.write_table(table)
+        arrow_data = sink.getvalue().to_pybytes()
+
+        # POST for statement creation succeeds
+        requests_mocker["post"].return_value = mock.Mock(
+            status_code=200,
+            json=lambda: {
+                "statement_id": "test_stmt",
+                "status": {"state": "SUCCEEDED"},
+                "manifest": {
+                    "truncated": False,
+                    "chunks": [{"chunk_index": 0, "row_count": 10, "byte_count": 100}],
+                },
+            },
+        )
+
+        def get_side_effect(url, *args, **kwargs):
+            headers = kwargs.get("headers", {})
+
+            # External URL fetch (no auth headers)
+            if url.startswith("https://external/"):
+                return mock.Mock(status_code=200, content=arrow_data)
+
+            if "/result/chunks/" in url:
+                chunk_fetch_count[0] += 1
+                chunk_fetch_headers.append(headers.get("Authorization", ""))
+
+                # First chunk fetch returns 401
+                if chunk_fetch_count[0] == 1:
+                    return mock.Mock(status_code=401)
+                # Retry succeeds
+                return mock.Mock(
+                    status_code=200,
+                    json=lambda: {
+                        "external_links": [{"external_link": "https://external/data"}]
+                    },
+                )
+            else:
+                # Polling response (already succeeded in POST)
+                return mock.Mock(
+                    status_code=200,
+                    json=lambda: {
+                        "status": {"state": "SUCCEEDED"},
+                        "manifest": {
+                            "truncated": False,
+                            "chunks": [
+                                {"chunk_index": 0, "row_count": 10, "byte_count": 100}
+                            ],
+                        },
+                    },
+                )
+
+        requests_mocker["get"].side_effect = get_side_effect
+
+        # Create datasource
+        datasource = DatabricksUCDatasource(
+            warehouse_id="test_warehouse",
+            catalog="test_catalog",
+            schema="test_schema",
+            query="SELECT 1",
+            credential_provider=refreshable_credential_provider,
+        )
+
+        # Get read tasks and execute the read function to trigger chunk fetch
+        read_tasks = datasource.get_read_tasks(parallelism=1)
+        assert len(read_tasks) == 1
+
+        # Execute the read function - this triggers chunk fetch
+        read_fn = read_tasks[0].read_fn
+        results = list(read_fn())
+
+        # Verify chunk fetch retry occurred
+        assert (
+            chunk_fetch_count[0] == 2
+        ), "Expected chunk fetch to be called twice (initial + retry)"
+
+        # Verify invalidate was called during chunk fetch
+        assert refreshable_credential_provider.invalidate_count == 1
+
+        # Verify first chunk fetch used expired token, retry used refreshed token
+        assert "expired_token" in chunk_fetch_headers[0]
+        assert "refreshed_token" in chunk_fetch_headers[1]
+
+        # Verify we got results
+        assert len(results) == 1
+
+
+class TestDatabricksUCDatasourceEmptyResult:
+    """Tests for empty result handling."""
+
+    def test_empty_result_returns_zero_count(self, requests_mocker, databricks_env):
+        """Test that empty result returns zero count."""
+
+        def post_mock(url, *args, **kwargs):
+            return MockResponse(
+                status_code=200,
+                _json_data={
+                    "statement_id": "test_stmt",
+                    "status": {"state": "PENDING"},
+                },
+            )
+
+        def get_mock(url, *args, **kwargs):
+            return MockResponse(
+                status_code=200,
+                _json_data={
+                    "status": {"state": "SUCCEEDED"},
+                    "manifest": {"truncated": False},
+                },
+            )
+
+        requests_mocker["post"].side_effect = post_mock
+        requests_mocker["get"].side_effect = get_mock
+
+        ds = ray.data.read_databricks_tables(
+            warehouse_id="dummy_warehouse",
+            query="select * from dummy_table",
+            catalog="dummy_catalog",
+            schema="dummy_schema",
+            override_num_blocks=1,
+        )
+
+        assert ds.count() == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_datasink.py

@@ -0,0 +1,145 @@
+from dataclasses import dataclass
+from typing import Iterable, List
+
+import numpy
+import pytest
+
+import ray
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.block import Block
+from ray.data.datasource import Datasink
+from ray.data.datasource.datasink import DummyOutputDatasink, WriteResult
+
+
+def test_write_datasink(ray_start_regular_shared):
+    output = DummyOutputDatasink()
+    ds = ray.data.range(10, override_num_blocks=2)
+    ds.write_datasink(output)
+    assert output.num_ok == 1
+    assert output.num_failed == 0
+    assert ray.get(output.data_sink.get_rows_written.remote()) == 10
+
+    output.enabled = False
+    ds = ray.data.range(10, override_num_blocks=2)
+    with pytest.raises(ValueError):
+        ds.write_datasink(output, ray_remote_args={"max_retries": 0})
+    assert output.num_ok == 1
+    assert output.num_failed == 1
+    assert ray.get(output.data_sink.get_rows_written.remote()) == 10
+
+
+@pytest.mark.parametrize("min_rows_per_write", [25, 50])
+def test_min_rows_per_write(tmp_path, ray_start_regular_shared, min_rows_per_write):
+    class MockDatasink(Datasink[None]):
+        def __init__(self, min_rows_per_write):
+            self._min_rows_per_write = min_rows_per_write
+
+        def write(self, blocks: Iterable[Block], ctx: TaskContext) -> None:
+            assert sum(len(block) for block in blocks) == self._min_rows_per_write
+
+        @property
+        def min_rows_per_write(self):
+            return self._min_rows_per_write
+
+    ray.data.range(100, override_num_blocks=4).write_datasink(
+        MockDatasink(min_rows_per_write)
+    )
+
+
+def test_write_result(ray_start_regular_shared):
+    """Test the write_result argument in `on_write_complete`."""
+
+    @dataclass
+    class CustomWriteResult:
+
+        ids: List[int]
+
+    class CustomDatasink(Datasink[CustomWriteResult]):
+        def __init__(self) -> None:
+            self.ids = []
+            self.num_rows = 0
+            self.size_bytes = 0
+
+        def write(self, blocks: Iterable[Block], ctx: TaskContext):
+            ids = []
+            for b in blocks:
+                ids.extend(b["id"].to_pylist())
+            return CustomWriteResult(ids=ids)
+
+        def on_write_complete(self, write_result: WriteResult[CustomWriteResult]):
+            ids = []
+            for result in write_result.write_returns:
+                ids.extend(result.ids)
+            self.ids = sorted(ids)
+            self.num_rows = write_result.num_rows
+            self.size_bytes = write_result.size_bytes
+
+    num_items = 10
+    size_bytes_per_row = 500
+
+    def map_fn(row):
+        row["data"] = numpy.zeros(size_bytes_per_row, dtype=numpy.int8)
+        return row
+
+    ds = ray.data.range(num_items).map(map_fn)
+
+    datasink = CustomDatasink()
+    ds.write_datasink(datasink)
+
+    assert datasink.ids == list(range(num_items))
+    assert datasink.num_rows == num_items
+    assert datasink.size_bytes == pytest.approx(num_items * size_bytes_per_row, rel=0.1)
+
+
+class NodeLoggerOutputDatasink(Datasink[None]):
+    """A writable datasource that logs node IDs of write tasks, for testing."""
+
+    def __init__(self, node_id: str):
+
+        self.num_ok = 0
+        self.num_failed = 0
+        self.node_id = node_id
+        self.num_rows_written = 0
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+
+        node_id = ray.get_runtime_context().get_node_id()
+        assert node_id == self.node_id
+
+    def on_write_complete(self, write_result: WriteResult[None]):
+        self.num_ok += 1
+        self.num_rows_written += write_result.num_rows
+
+    def on_write_failed(self, error: Exception) -> None:
+        self.num_failed += 1
+
+
+def test_write_datasink_ray_remote_args(ray_start_cluster):
+    ray.shutdown()
+    cluster = ray_start_cluster
+    cluster.add_node(
+        resources={"foo": 100},
+        num_cpus=1,
+    )
+    bar_worker = cluster.add_node(resources={"bar": 100}, num_cpus=1)
+    bar_node_id = bar_worker.node_id
+
+    ray.init(cluster.address)
+
+    output = NodeLoggerOutputDatasink(bar_node_id)
+    ds = ray.data.range(100, override_num_blocks=10)
+    # Pin write tasks to node with "bar" resource.
+    ds.write_datasink(output, ray_remote_args={"resources": {"bar": 1}})
+    assert output.num_ok == 1
+    assert output.num_failed == 0
+    assert output.num_rows_written == 100
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_delta.py

@@ -0,0 +1,171 @@
+import os
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+from packaging.version import parse as parse_version
+
+import ray
+from ray.data import Schema
+from ray.data._internal.util import rows_same
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+# deltalake's write_deltalake requires pyarrow >= 15 for the Arrow C Stream interface.
+_pa_version = get_pyarrow_version()
+assert _pa_version is not None, "pyarrow must be installed to run these tests"
+
+pytestmark = pytest.mark.skipif(
+    _pa_version < parse_version("15.0.0"),
+    reason="deltalake write_deltalake requires pyarrow >= 15.0",
+)
+
+
+@pytest.mark.parametrize(
+    "batch_size",
+    [1, 100],
+)
+@pytest.mark.parametrize(
+    "write_mode",
+    ["append", "overwrite"],
+)
+def test_delta_read_basic(tmp_path, batch_size, write_mode):
+    from deltalake import write_deltalake
+
+    # Parse the data path.
+    path = os.path.join(tmp_path, "tmp_test_delta")
+
+    # Create a sample Delta Lake table
+    df = pd.DataFrame(
+        {"x": [42] * batch_size, "y": ["a"] * batch_size, "z": [3.14] * batch_size}
+    )
+    table = pa.Table.from_pandas(df)
+    if write_mode == "append":
+        write_deltalake(path, table, mode=write_mode)
+        write_deltalake(path, table, mode=write_mode)
+        expected = pd.concat([df, df], ignore_index=True)
+    elif write_mode == "overwrite":
+        write_deltalake(path, table, mode=write_mode)
+        expected = df
+    else:
+        raise ValueError(f"Unexpected write_mode: {write_mode}")
+
+    # Read the Delta Lake table
+    ds = ray.data.read_delta(path)
+
+    assert ds.schema() == Schema(
+        pa.schema(
+            {
+                "x": pa.int64(),
+                "y": pa.string(),
+                "z": pa.float64(),
+            }
+        )
+    )
+    assert rows_same(ds.to_pandas(), expected)
+
+
+@pytest.mark.parametrize(
+    "columns, expected_columns",
+    [
+        (["a", "c"], ["a", "c"]),
+        (["b"], ["b"]),
+        (["a", "b", "c"], ["a", "b", "c"]),
+    ],
+)
+def test_delta_read_column_selection(tmp_path, columns, expected_columns):
+    from deltalake import write_deltalake
+
+    path = os.path.join(tmp_path, "tmp_test_delta_cols")
+    df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"], "c": [1.0, 2.0, 3.0]})
+    write_deltalake(path, pa.Table.from_pandas(df))
+
+    ds = ray.data.read_delta(path, columns=columns)
+    expected = df[expected_columns]
+
+    assert ds.schema().names == expected_columns
+    assert rows_same(ds.to_pandas(), expected)
+
+
+@pytest.mark.parametrize(
+    "version, expected_data",
+    [
+        (0, {"x": [1, 2]}),
+        (1, {"x": [3, 4, 5]}),
+        (None, {"x": [3, 4, 5]}),
+    ],
+)
+def test_delta_read_version(tmp_path, version, expected_data):
+    from deltalake import write_deltalake
+
+    path = os.path.join(tmp_path, "tmp_test_delta_version")
+    write_deltalake(path, pa.table({"x": [1, 2]}))
+    write_deltalake(path, pa.table({"x": [3, 4, 5]}), mode="overwrite")
+
+    ds = ray.data.read_delta(path, version=version)
+    expected = pd.DataFrame(expected_data)
+
+    assert rows_same(ds.to_pandas(), expected)
+
+
+def test_delta_read_schema_evolution(tmp_path):
+    """Older files missing newer columns should be null-filled."""
+    from deltalake import write_deltalake
+
+    path = os.path.join(tmp_path, "tmp_test_delta_schema_evo")
+
+    write_deltalake(path, pa.table({"x": [1, 2]}))
+    write_deltalake(
+        path,
+        pa.table({"x": [3, 4], "y": ["a", "b"]}),
+        mode="append",
+        schema_mode="merge",  # pyrefly: ignore[unexpected-keyword]
+    )
+
+    ds = ray.data.read_delta(path)
+    expected = pd.DataFrame(
+        {"x": [1, 2, 3, 4], "y": [None, None, "a", "b"]},
+    )
+    # Match the Arrow-backed null sentinel produced by ``to_pandas()``.
+    expected["y"] = expected["y"].astype("string")
+
+    assert rows_same(ds.to_pandas(), expected)
+
+
+@pytest.mark.parametrize(
+    "storage_options",
+    [{}, None],
+)
+def test_delta_read_storage_options(tmp_path, storage_options):
+    """Verify that storage_options are forwarded to DeltaTable."""
+    from deltalake import write_deltalake
+
+    path = os.path.join(tmp_path, "tmp_test_delta_storage_opts")
+    df = pd.DataFrame({"x": [1, 2, 3]})
+    write_deltalake(path, pa.Table.from_pandas(df))
+
+    ds = ray.data.read_delta(path, storage_options=storage_options)
+    assert rows_same(ds.to_pandas(), df)
+
+
+def test_delta_read_empty_table(tmp_path):
+    from deltalake import write_deltalake
+
+    path = os.path.join(tmp_path, "tmp_test_delta_empty")
+    write_deltalake(path, pa.table({"x": pa.array([], type=pa.int64())}))
+
+    ds = ray.data.read_delta(path)
+    assert ds.count() == 0
+
+
+def test_delta_read_rejects_multiple_paths():
+    with pytest.raises(ValueError, match="Only a single Delta Lake table path"):
+        ray.data.read_delta(["path1", "path2"])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_delta_sharing.py

@@ -0,0 +1,294 @@
+import json
+import unittest
+from typing import TYPE_CHECKING, Optional
+from unittest import mock
+from unittest.mock import MagicMock, patch
+
+import pytest
+from delta_sharing.protocol import Table
+from delta_sharing.rest_client import DataSharingRestClient
+
+from ray.data._internal.datasource.delta_sharing_datasource import (
+    DeltaSharingDatasource,
+    _parse_delta_sharing_url,
+)
+from ray.data.block import BlockMetadata
+from ray.data.dataset import Dataset
+from ray.data.datasource.datasource import ReadTask
+from ray.data.read_api import read_delta_sharing_tables
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+
+class TestDeltaSharingDatasource(unittest.TestCase):
+    def setUp(self):
+        self.url = "path/to/profile#share.schema.table"
+        self.limit = 1000
+        self.version = 1
+        self.json_predicate_hints = '{"column":"value"}'
+        self.table = Table(name="table", share="share", schema="schema")
+
+        self.mock_rest_client = mock.create_autospec(DataSharingRestClient)
+        self.mock_response = mock.Mock()
+        self.mock_rest_client.list_files_in_table.return_value = self.mock_response
+
+        self.mock_response.add_files = [
+            {"url": "file1", "id": "1"},
+            {"url": "file2", "id": "2"},
+        ]
+        self.mock_response.metadata.schema_string = json.dumps(
+            {
+                "type": "struct",
+                "fields": [
+                    {
+                        "name": "column1",
+                        "type": "string",
+                        "nullable": True,
+                        "metadata": {},
+                    }
+                ],
+            }
+        )
+
+    @patch(
+        "ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource."
+        "setup_delta_sharing_connections"
+    )
+    def test_init(self, mock_setup_delta_sharing_connections):
+        mock_setup_delta_sharing_connections.return_value = (
+            self.table,
+            self.mock_rest_client,
+        )
+        datasource = DeltaSharingDatasource(
+            url=self.url,
+            json_predicate_hints=self.json_predicate_hints,
+            limit=self.limit,
+            version=self.version,
+            timestamp=None,
+        )
+
+        self.assertEqual(datasource._url, self.url)
+        self.assertEqual(datasource._json_predicate_hints, self.json_predicate_hints)
+        self.assertEqual(datasource._limit, self.limit)
+        self.assertEqual(datasource._version, self.version)
+        self.assertEqual(datasource._timestamp, None)
+
+    @patch(
+        "ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource."
+        "setup_delta_sharing_connections"
+    )
+    def test_get_read_tasks(self, mock_setup_delta_sharing_connections):
+        mock_setup_delta_sharing_connections.return_value = (
+            self.table,
+            self.mock_rest_client,
+        )
+        datasource = DeltaSharingDatasource(
+            url=self.url,
+            json_predicate_hints=self.json_predicate_hints,
+            limit=self.limit,
+            version=self.version,
+            timestamp=None,
+        )
+
+        read_tasks = datasource.get_read_tasks(parallelism=2)
+        self.assertEqual(len(read_tasks), 2)
+        self.assertTrue(all(isinstance(task, ReadTask) for task in read_tasks))
+
+        for task in read_tasks:
+            metadata = task.metadata
+            self.assertIsInstance(metadata, BlockMetadata)
+            self.assertEqual(len(metadata.input_files), 1)
+            self.assertTrue(metadata.input_files[0]["url"] in ["file1", "file2"])
+            self.assertEqual(metadata.num_rows, None)
+            self.assertEqual(metadata.size_bytes, None)
+            self.assertEqual(task.schema, None)
+            self.assertEqual(metadata.exec_stats, None)
+
+
+class TestParseDeltaSharingUrl(unittest.TestCase):
+    def test_valid_url(self):
+        url = "profile#share.schema.table"
+        expected_result = ("profile", "share", "schema", "table")
+        self.assertEqual(_parse_delta_sharing_url(url), expected_result)
+
+    def test_missing_hash(self):
+        url = "profile-share.schema.table"
+        with self.assertRaises(ValueError) as context:
+            _parse_delta_sharing_url(url)
+        self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
+
+    def test_missing_fragments(self):
+        url = "profile#share.schema"
+        with self.assertRaises(ValueError) as context:
+            _parse_delta_sharing_url(url)
+        self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
+
+    def test_empty_profile(self):
+        url = "#share.schema.table"
+        with self.assertRaises(ValueError) as context:
+            _parse_delta_sharing_url(url)
+        self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
+
+    def test_empty_share(self):
+        url = "profile#.schema.table"
+        with self.assertRaises(ValueError) as context:
+            _parse_delta_sharing_url(url)
+        self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
+
+    def test_empty_schema(self):
+        url = "profile#share..table"
+        with self.assertRaises(ValueError) as context:
+            _parse_delta_sharing_url(url)
+        self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
+
+    def test_empty_table(self):
+        url = "profile#share.schema."
+        with self.assertRaises(ValueError) as context:
+            _parse_delta_sharing_url(url)
+        self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
+
+
+class MockDeltaSharingDatasource:
+    def __init__(
+        self, url, json_predicate_hints=None, limit=None, version=None, timestamp=None
+    ):
+        self._url = url
+        self._json_predicate_hints = json_predicate_hints
+        self._limit = limit
+        self._version = version
+        self._timestamp = timestamp
+
+    def setup_delta_sharing_connections(self, url):
+        # Return mock objects for table and rest_client
+        table = MagicMock()
+        rest_client = MagicMock()
+
+        # Mock the rest_client's list_files_in_table method
+        rest_client.list_files_in_table.return_value = MagicMock(
+            add_files=[
+                {
+                    "url": "https://s3-bucket-name.s3.us-west-2.amazonaws.com/delta-exchange-test/table2/date%3D2021-04-28/part-00000-591723a8-6a27-4240-a90e-57426f4736d2.c000.snappy.parquet",  # noqa E501
+                    "id": "591723a8-6a27-4240-a90e-57426f4736d2",
+                    "size": 573,
+                    "partitionValues": {"date": "2021-04-28"},
+                    "stats": '{"numRecords":1,"minValues":{"eventTime":"2021-04-28T23:33:48.719Z"},"maxValues":{"eventTime":"2021-04-28T23:33:48.719Z"},"nullCount":{"eventTime":0}}',  # noqa E501
+                    "expirationTimestamp": 1652140800000,
+                }
+            ],
+            metadata=MagicMock(
+                schema_string='{"type":"struct","fields":[{"name":"eventTime","type":"timestamp","nullable":true,"metadata":{}},{"name":"date","type":"date","nullable":true,"metadata":{}}]}'  # noqa E501
+            ),
+        )
+        return table, rest_client
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ):
+        self._table, self._rest_client = self.setup_delta_sharing_connections(self._url)
+        response = self._rest_client.list_files_in_table(
+            self._table,
+            jsonPredicateHints=self._json_predicate_hints,
+            limitHint=self._limit,
+            version=self._version,
+            timestamp=self._timestamp,
+        )
+
+        read_tasks = []
+        for _ in range(parallelism):
+            read_task = MagicMock()
+            read_task.metadata = MagicMock(
+                num_rows=1,
+                schema=None,
+                input_files=[file["url"] for file in response.add_files],
+                size_bytes=573,
+                exec_stats=None,
+            )
+            read_task.data = MagicMock(
+                return_value=[
+                    {
+                        "eventTime": "2021-04-28T23:33:48.719Z",
+                        "date": "2021-04-28",
+                    }
+                ]
+            )
+            read_task.per_task_row_limit = per_task_row_limit
+            read_tasks.append(read_task)
+        return read_tasks
+
+
+@pytest.fixture
+def mock_delta_sharing_datasource(mocker):
+    mock_datasource = mocker.patch(
+        "ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource",
+        new=MockDeltaSharingDatasource,
+    )
+    return mock_datasource
+
+
+@pytest.fixture
+def mock_ray_data_read_datasource(mocker):
+    mock_read_datasource = mocker.patch("ray.data.read_datasource")
+    mock_read_datasource.return_value = MagicMock(spec=Dataset)
+    return mock_read_datasource
+
+
+@pytest.fixture
+def setup_profile_file(tmpdir):
+    profile_content = {
+        "shareCredentialsVersion": 1,
+        "endpoint": "https://sharing.delta.io/delta-sharing/",
+        "bearerToken": "<token>",
+        "expirationTime": "2021-11-12T00:12:29.0Z",
+    }
+    profile_file = tmpdir.join("profile.json")
+    profile_file.write(json.dumps(profile_content))
+    return str(profile_file)
+
+
+def test_read_delta_sharing_tables(
+    mock_delta_sharing_datasource, mock_ray_data_read_datasource, setup_profile_file
+):
+    url = f"{setup_profile_file}#share.schema.table"
+    limit = 100
+    version = 1
+    timestamp = "2021-01-01T00:00:00Z"
+    json_predicate_hints = '{"eventTime": "2021-04-28T23:33:48.719Z"}'
+    ray_remote_args = {"num_cpus": 2}
+    concurrency = 4
+    override_num_blocks = 2
+
+    # Call the function under test
+    result = read_delta_sharing_tables(
+        url=url,
+        limit=limit,
+        version=version,
+        timestamp=timestamp,
+        json_predicate_hints=json_predicate_hints,
+        ray_remote_args=ray_remote_args,
+        concurrency=concurrency,
+        override_num_blocks=override_num_blocks,
+    )
+
+    # Assert the result and interactions
+    assert isinstance(result, Dataset)
+    mock_ray_data_read_datasource.assert_called_once()
+    args, kwargs = mock_ray_data_read_datasource.call_args
+    datasource = kwargs["datasource"]
+    assert datasource._url == url
+    assert datasource._json_predicate_hints == json_predicate_hints
+    assert datasource._limit == limit
+    assert datasource._version == version
+    assert datasource._timestamp == timestamp
+    assert kwargs["ray_remote_args"] == ray_remote_args
+    assert kwargs["concurrency"] == concurrency
+    assert kwargs["override_num_blocks"] == override_num_blocks
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_file_based_datasource.py

@@ -0,0 +1,508 @@
+import os
+from typing import Any, Dict, Iterator, List
+from urllib.parse import urlparse
+
+import pyarrow
+import pytest
+from pytest_lazy_fixtures import lf as lazy_fixture
+
+import ray
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasource import ReadTask
+from ray.data.datasource.file_based_datasource import (
+    FileBasedDatasource,
+)
+from ray.data.datasource.partitioning import (
+    Partitioning,
+    PartitionStyle,
+    PathPartitionFilter,
+)
+
+
+class MockFileBasedDatasource(FileBasedDatasource):
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        builder = DelegatingBlockBuilder()
+        builder.add({"data": f.readall()})
+        yield builder.build()
+
+
+def execute_read_tasks(tasks: List[ReadTask]) -> List[Dict[str, Any]]:
+    """Execute the read tasks and return the resulting rows.
+
+    The motivation for this utility function is so that we can test datasources without
+    scheduling Ray tasks.
+    """
+    builder = DelegatingBlockBuilder()
+    for task in tasks:
+        for block in task():
+            builder.add_block(block)
+    block = builder.build()
+
+    block_accessor = BlockAccessor.for_block(block)
+    rows = list(block_accessor.iter_rows(public_row_format=True))
+
+    return rows
+
+
+def strip_scheme(uri):
+    """Remove scheme from a URI, if it exists."""
+    parsed = urlparse(uri)
+    if parsed.scheme:
+        return uri.split("://", 1)[1]  # remove scheme
+    return uri  # no scheme, return as-is
+
+
+@pytest.mark.parametrize(
+    "filesystem,dir_path,endpoint_url",
+    [
+        (None, lazy_fixture("local_path"), None),
+        (lazy_fixture("local_fs"), lazy_fixture("local_path"), None),
+        (lazy_fixture("s3_fs"), lazy_fixture("s3_path"), lazy_fixture("s3_server")),
+        (
+            lazy_fixture("s3_fs_with_space"),
+            lazy_fixture("s3_path_with_space"),
+            lazy_fixture("s3_server"),
+        ),
+        (
+            lazy_fixture("s3_fs_with_special_chars"),
+            lazy_fixture("s3_path_with_special_chars"),
+            lazy_fixture("s3_server"),
+        ),
+    ],
+)
+def test_read_single_file(ray_start_regular_shared, filesystem, dir_path, endpoint_url):
+    # `FileBasedDatasource` should read from the local filesystem if you don't specify
+    # one.
+    write_filesystem = filesystem
+    if write_filesystem is None:
+        write_filesystem = pyarrow.fs.LocalFileSystem()
+
+    file_uri = os.path.join(dir_path, "file.txt")
+
+    # PyArrow filesystems expect paths without schemes. `FileBasedDatasource` handles
+    # this internally, but we need to manually strip the scheme for the test setup.
+    write_path = strip_scheme(file_uri)
+    with write_filesystem.open_output_stream(write_path) as f:
+        f.write(b"spam")
+
+    datasource = MockFileBasedDatasource(file_uri, filesystem=filesystem)
+    tasks = datasource.get_read_tasks(1)
+    rows = execute_read_tasks(tasks)
+
+    assert rows == [{"data": b"spam"}]
+
+
+def test_read_single_directory(ray_start_regular_shared, tmp_path):
+    dir_path = tmp_path / "dir"
+    dir_path.mkdir()
+
+    p1 = dir_path / "a.txt"
+    p1.write_bytes(b"a")
+
+    p2 = dir_path / "b.txt"
+    p2.write_bytes(b"b")
+
+    datasource = MockFileBasedDatasource(dir_path)
+    rows = execute_read_tasks(datasource.get_read_tasks(1))
+
+    assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"b"}]
+
+
+def test_read_dir_and_file_mixed(ray_start_regular_shared, tmp_path):
+    dir_path = tmp_path / "dir"
+    dir_path.mkdir()
+
+    p1 = dir_path / "a.txt"
+    p1.write_bytes(b"a")
+
+    p2 = tmp_path / "c.txt"
+    p2.write_bytes(b"c")
+
+    datasource = MockFileBasedDatasource([str(dir_path), str(p2)])
+    rows = execute_read_tasks(datasource.get_read_tasks(1))
+
+    assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"c"}]
+
+
+def test_pathlib_paths(ray_start_regular_shared, tmp_path):
+    """Test that FileBasedDatasource accepts pathlib.Path objects."""
+    from pathlib import Path
+
+    path = Path(tmp_path) / "test_pathlib"
+    path.mkdir()
+
+    # Create pathlib.Path objects
+    file1 = path / "file1.txt"
+    file2 = path / "file2.txt"
+
+    file1.write_bytes(b"hello")
+    file2.write_bytes(b"world")
+
+    # Verify list of pathlib.Path works
+    datasource = MockFileBasedDatasource([file1, file2])
+    rows = execute_read_tasks(datasource.get_read_tasks(1))
+    assert sorted(rows, key=lambda r: r["data"]) == [
+        {"data": b"hello"},
+        {"data": b"world"},
+    ]
+
+    # Verify single pathlib.Path works
+    datasource = MockFileBasedDatasource(file1)
+    rows = execute_read_tasks(datasource.get_read_tasks(1))
+    assert rows == [{"data": b"hello"}]
+
+
+def test_single_file_infinite_target_max_block_size(
+    ray_start_regular_shared, target_max_block_size_infinite_or_default, tmp_path
+):
+    path = tmp_path / "file.txt"
+    path.write_bytes(b"spam")
+
+    datasource = MockFileBasedDatasource(path)
+    rows = execute_read_tasks(datasource.get_read_tasks(1))
+
+    assert rows == [{"data": b"spam"}]
+
+
+def test_partitioning_hive(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "country=us")
+    os.mkdir(path)
+    with open(os.path.join(path, "file.txt"), "wb") as file:
+        file.write(b"")
+
+    datasource = MockFileBasedDatasource(tmp_path, partitioning=Partitioning("hive"))
+
+    tasks = datasource.get_read_tasks(1)
+    rows = execute_read_tasks(tasks)
+
+    assert rows == [{"data": b"", "country": "us"}]
+
+
+def test_partition_filter_hive(ray_start_regular_shared, tmp_path):
+    for country in ["us", "jp"]:
+        path = os.path.join(tmp_path, f"country={country}")
+        os.mkdir(path)
+        with open(os.path.join(path, "file.txt"), "wb") as file:
+            file.write(b"")
+
+    filter = PathPartitionFilter.of(
+        style=PartitionStyle.HIVE,
+        filter_fn=lambda partitions: partitions["country"] == "us",
+    )
+    datasource = MockFileBasedDatasource(
+        tmp_path, partitioning=Partitioning("hive"), partition_filter=filter
+    )
+
+    tasks = datasource.get_read_tasks(1)
+    rows = execute_read_tasks(tasks)
+
+    assert rows == [{"data": b"", "country": "us"}]
+
+
+def test_partitioning_dir(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "us")
+    os.mkdir(path)
+    with open(os.path.join(path, "file.txt"), "wb") as file:
+        file.write(b"")
+
+    datasource = MockFileBasedDatasource(
+        tmp_path,
+        partitioning=Partitioning("dir", field_names=["country"], base_dir=tmp_path),
+    )
+
+    tasks = datasource.get_read_tasks(1)
+    rows = execute_read_tasks(tasks)
+
+    assert rows == [{"data": b"", "country": "us"}]
+
+
+def test_partition_filter_dir(ray_start_regular_shared, tmp_path):
+    for country in ["us", "jp"]:
+        path = os.path.join(tmp_path, country)
+        os.mkdir(path)
+        with open(os.path.join(path, "file.txt"), "wb") as file:
+            file.write(b"")
+
+    filter = PathPartitionFilter.of(
+        style=PartitionStyle.DIRECTORY,
+        base_dir=tmp_path,
+        field_names=["country"],
+        filter_fn=lambda partitions: partitions["country"] == "us",
+    )
+    partitioning = Partitioning("dir", field_names=["country"], base_dir=tmp_path)
+    datasource = MockFileBasedDatasource(
+        tmp_path, partitioning=partitioning, partition_filter=filter
+    )
+
+    tasks = datasource.get_read_tasks(1)
+    rows = execute_read_tasks(tasks)
+
+    assert rows == [{"data": b"", "country": "us"}]
+
+
+def test_partitioning_raises_on_mismatch(ray_start_regular_shared, tmp_path):
+    """Test when the partition key already exists in the data."""
+
+    class StubDatasource(FileBasedDatasource):
+        def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+            builder = DelegatingBlockBuilder()
+            builder.add({"country": f.readall()})
+            yield builder.build()
+
+    path = os.path.join(tmp_path, "country=us")
+    os.mkdir(path)
+    with open(os.path.join(path, "file.txt"), "wb") as file:
+        file.write(b"jp")
+
+    datasource = StubDatasource(tmp_path, partitioning=Partitioning("hive"))
+
+    # The data is `jp`, but the path contains `us`. Since the values are different,
+    # the datasource should raise a ValueError.
+    with pytest.raises(ValueError):
+        tasks = datasource.get_read_tasks(1)
+        execute_read_tasks(tasks)
+
+
+def test_ignore_missing_paths_true(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "file.txt")
+    with open(path, "wb") as file:
+        file.write(b"")
+
+    datasource = MockFileBasedDatasource(
+        [path, "missing.txt"], ignore_missing_paths=True
+    )
+
+    tasks = datasource.get_read_tasks(1)
+    rows = execute_read_tasks(tasks)
+
+    assert rows == [{"data": b""}]
+
+
+def test_ignore_missing_paths_false(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "file.txt")
+    with open(path, "wb") as file:
+        file.write(b"")
+
+    with pytest.raises(FileNotFoundError):
+        datasource = MockFileBasedDatasource(
+            [path, "missing.txt"], ignore_missing_paths=False
+        )
+        tasks = datasource.get_read_tasks(1)
+        execute_read_tasks(tasks)
+
+
+def test_local_paths(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "test.txt")
+    with open(path, "w"):
+        pass
+
+    datasource = MockFileBasedDatasource(path)
+    assert datasource.supports_distributed_reads
+
+    datasource = MockFileBasedDatasource(f"local://{path}")
+    assert not datasource.supports_distributed_reads
+
+
+def test_local_paths_with_client_raises_error(ray_start_cluster_enabled, tmp_path):
+    ray_start_cluster_enabled.add_node(num_cpus=1)
+    ray_start_cluster_enabled.head_node._ray_params.ray_client_server_port = "10004"
+    ray_start_cluster_enabled.head_node.start_ray_client_server()
+    ray.init("ray://localhost:10004")
+
+    path = os.path.join(tmp_path, "test.txt")
+    with open(path, "w"):
+        pass
+
+    with pytest.raises(ValueError):
+        MockFileBasedDatasource(f"local://{path}")
+
+
+def test_include_paths(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "test.txt")
+    with open(path, "w"):
+        pass
+
+    datasource = MockFileBasedDatasource(path, include_paths=True)
+    ds = ray.data.read_datasource(datasource)
+
+    paths = [row["path"] for row in ds.take_all()]
+    assert paths == [path]
+
+
+def test_file_extensions(ray_start_regular_shared, tmp_path):
+    csv_path = os.path.join(tmp_path, "file.csv")
+    with open(csv_path, "w") as file:
+        file.write("spam")
+
+    txt_path = os.path.join(tmp_path, "file.txt")
+    with open(txt_path, "w") as file:
+        file.write("ham")
+
+    datasource = MockFileBasedDatasource([csv_path, txt_path], file_extensions=None)
+    ds = ray.data.read_datasource(datasource)
+    assert sorted(ds.input_files()) == sorted([csv_path, txt_path])
+
+    datasource = MockFileBasedDatasource([csv_path, txt_path], file_extensions=["csv"])
+    ds = ray.data.read_datasource(datasource)
+    assert ds.input_files() == [csv_path]
+
+
+def test_file_extensions_no_match_raises(ray_start_regular_shared, tmp_path):
+    txt_path = tmp_path / "file.txt"
+    txt_path.write_bytes(b"ham")
+
+    with pytest.raises(
+        ValueError,
+        match="No input files found to read with the following file extensions",
+    ):
+        MockFileBasedDatasource([str(txt_path)], file_extensions=["csv"])
+
+
+def test_flaky_read_task_retries(ray_start_regular_shared, tmp_path):
+    """Test that flaky read tasks are retried for both the
+    default set of retried errors and a custom set of retried errors."""
+    csv_path = os.path.join(tmp_path, "file.csv")
+    with open(csv_path, "w") as file:
+        file.write("spam")
+
+    class Counter:
+        def __init__(self):
+            self.value = 0
+
+        def increment(self):
+            self.value += 1
+            return self.value
+
+    default_retried_error = ray.data.context.DEFAULT_RETRIED_IO_ERRORS[0]
+    custom_retried_error = "AWS Error ACCESS_DENIED"
+
+    class FlakyFileBasedDatasource(MockFileBasedDatasource):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+            CounterActor = ray.remote(Counter)
+            # This actor ref is shared across all read tasks.
+            self.counter = CounterActor.remote()
+
+        def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+            count = ray.get(self.counter.increment.remote())
+            if count == 1:
+                raise RuntimeError(default_retried_error)
+            elif count == 2:
+                raise RuntimeError(custom_retried_error)
+            else:
+                yield from super()._read_stream(f, path)
+
+    ray.data.DataContext.get_current().retried_io_errors.append(custom_retried_error)
+
+    datasource = FlakyFileBasedDatasource([csv_path])
+    ds = ray.data.read_datasource(datasource)
+    assert len(ds.take()) == 1
+
+
+@pytest.mark.parametrize(
+    "fs",
+    [pyarrow.fs.S3FileSystem(), pyarrow.fs.LocalFileSystem()],
+)
+@pytest.mark.parametrize(
+    "wrap_with_retries",
+    [True, False],
+)
+def test_s3_filesystem_serialization(fs, wrap_with_retries):
+    """Tests that the S3FileSystem can be serialized and deserialized with
+    the serialization workaround (_S3FileSystemWrapper).
+
+    Also checks that filesystems wrapped with RetryingPyFileSystem are
+    properly unwrapped.
+    """
+    import ray.cloudpickle as ray_pickle
+    from ray.data._internal.util import RetryingPyFileSystem
+    from ray.data.datasource.file_based_datasource import (
+        _unwrap_s3_serialization_workaround,
+        _wrap_s3_serialization_workaround,
+    )
+
+    orig_fs = fs
+
+    if wrap_with_retries:
+        fs = RetryingPyFileSystem.wrap(fs, retryable_errors=["DUMMY ERROR"])
+
+    wrapped_fs = _wrap_s3_serialization_workaround(fs)
+    unpickled_fs = ray_pickle.loads(ray_pickle.dumps(wrapped_fs))
+    unwrapped_fs = _unwrap_s3_serialization_workaround(unpickled_fs)
+
+    if wrap_with_retries:
+        assert isinstance(unwrapped_fs, RetryingPyFileSystem)
+        assert isinstance(unwrapped_fs.unwrap(), orig_fs.__class__)
+        assert unwrapped_fs.retryable_errors == ["DUMMY ERROR"]
+    else:
+        assert isinstance(unwrapped_fs, orig_fs.__class__)
+
+
+@pytest.mark.parametrize("shuffle", [True, False, "file"])
+def test_invalid_shuffle_arg_raises_error(ray_start_regular_shared, shuffle):
+    with pytest.raises(ValueError):
+        FileBasedDatasource("example://iris.csv", shuffle=shuffle)
+
+
+@pytest.mark.parametrize("shuffle", [None, "files"])
+def test_valid_shuffle_arg_does_not_raise_error(ray_start_regular_shared, shuffle):
+    FileBasedDatasource("example://iris.csv", shuffle=shuffle)
+
+
+def test_shuffle_files_changes_order(ray_start_regular_shared, tmp_path):
+    NUM_FILES = 10
+    NUM_RUNS = 5
+
+    for i in range(NUM_FILES):
+        (tmp_path / f"file_{i:02d}.txt").write_bytes(f"data_{i}".encode())
+
+    datasource = MockFileBasedDatasource(
+        str(tmp_path), shuffle="files", include_paths=True
+    )
+
+    output_paths_list = []
+    # Run NUM_RUNS times to verify shuffle produces different orderings
+    for _ in range(NUM_RUNS):
+        tasks = datasource.get_read_tasks(1)
+        rows = execute_read_tasks(tasks)
+        output_filenames = [os.path.basename(row["path"]) for row in rows]
+        output_paths_list.append(output_filenames)
+
+    expected_order = [f"file_{i:02d}.txt" for i in range(NUM_FILES)]
+
+    # Verify shuffle produces non-deterministic orderings across runs
+    unique_orderings = {tuple(paths) for paths in output_paths_list}
+    assert len(unique_orderings) >= 2
+
+    # Verify all files are present in each run
+    for output_paths in output_paths_list:
+        assert sorted(output_paths) == sorted(expected_order)
+
+
+def test_read_s3_file_error(shutdown_only, s3_path):
+    from ray.data.datasource.file_meta_provider import _handle_read_os_error
+
+    dummy_path = s3_path + "_dummy"
+    error_message = "Please check that file exists and has properly configured access."
+    with pytest.raises(OSError, match=error_message):
+        ray.data.read_parquet(dummy_path)
+    with pytest.raises(OSError, match=error_message):
+        ray.data.read_binary_files(dummy_path)
+    with pytest.raises(OSError, match=error_message):
+        ray.data.read_csv(dummy_path)
+    with pytest.raises(OSError, match=error_message):
+        ray.data.read_json(dummy_path)
+    with pytest.raises(OSError, match=error_message):
+        error = OSError(
+            f"Error creating dataset. Could not read schema from {dummy_path}: AWS "
+            "Error [code 15]: No response body.. Is this a 'parquet' file?"
+        )
+        _handle_read_os_error(error, dummy_path)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_file_datasink.py

@@ -0,0 +1,176 @@
+import os
+from typing import Any, Dict
+
+import pyarrow
+import pytest
+from pyarrow.fs import LocalFileSystem
+
+import ray
+from ray.data.block import BlockAccessor
+from ray.data.datasource import BlockBasedFileDatasink, RowBasedFileDatasink
+
+
+class FlakyOutputStream:
+    def __init__(self, stream: pyarrow.NativeFile, num_attempts: int):
+        self._stream = stream
+        self._num_attempts = num_attempts
+
+    def __enter__(self):
+        return self._stream.__enter__()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self._num_attempts < 2:
+            raise RuntimeError("AWS Error NETWORK_CONNECTION")
+
+        self._stream.__exit__(exc_type, exc_value, traceback)
+
+
+def test_flaky_block_based_open_output_stream(ray_start_regular_shared, tmp_path):
+    class FlakyCSVDatasink(BlockBasedFileDatasink):
+        def __init__(self, path: str):
+            super().__init__(path)
+            self._num_attempts = 0
+            self._filesystem = LocalFileSystem()
+
+        def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
+            stream = self._filesystem.open_output_stream(path)
+            flaky_stream = FlakyOutputStream(stream, self._num_attempts)
+            self._num_attempts += 1
+            return flaky_stream
+
+        def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+            block.to_pandas().to_csv(file)
+
+    ds = ray.data.range(100)
+
+    ds.write_datasink(FlakyCSVDatasink(tmp_path))
+
+    expected_values = list(range(100))
+    written_values = [row["id"] for row in ray.data.read_csv(tmp_path).take_all()]
+    assert sorted(written_values) == sorted(expected_values)
+
+
+def test_flaky_row_based_open_output_stream(ray_start_regular_shared, tmp_path):
+    class FlakyTextDatasink(RowBasedFileDatasink):
+        def __init__(self, path: str):
+            super().__init__(path)
+            self._num_attempts = 0
+            self._filesystem = LocalFileSystem()
+
+        def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
+            stream = self._filesystem.open_output_stream(path)
+            flaky_stream = FlakyOutputStream(stream, self._num_attempts)
+            self._num_attempts += 1
+            return flaky_stream
+
+        def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
+            file.write(f"{row['id']}".encode())
+
+    ds = ray.data.range(100)
+
+    ds.write_datasink(FlakyTextDatasink(tmp_path))
+
+    expected_values = [str(i) for i in range(100)]
+    written_values = [row["text"] for row in ray.data.read_text(tmp_path).take_all()]
+    assert sorted(written_values) == sorted(expected_values)
+
+
+def test_flaky_write_block_to_file(ray_start_regular_shared, tmp_path):
+    class FlakyCSVDatasink(BlockBasedFileDatasink):
+        def __init__(self, path: str):
+            super().__init__(path)
+            self._num_attempts = 0
+
+        def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+            if self._num_attempts < 2:
+                self._num_attempts += 1
+                raise RuntimeError("AWS Error INTERNAL_FAILURE")
+
+            block.to_pandas().to_csv(file)
+
+    ds = ray.data.range(100)
+
+    ds.write_datasink(FlakyCSVDatasink(tmp_path))
+
+    expected_values = list(range(100))
+    written_values = [row["id"] for row in ray.data.read_csv(tmp_path).take_all()]
+    assert sorted(written_values) == sorted(expected_values)
+
+
+def test_flaky_write_row_to_file(ray_start_regular_shared, tmp_path):
+    class FlakyTextDatasink(RowBasedFileDatasink):
+        def __init__(self, path: str):
+            super().__init__(path)
+            self._num_attempts = 0
+
+        def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
+            if self._num_attempts < 2:
+                self._num_attempts += 1
+                raise RuntimeError("AWS Error INTERNAL_FAILURE")
+
+            file.write(f"{row['id']}".encode())
+
+    ds = ray.data.range(100)
+
+    ds.write_datasink(FlakyTextDatasink(tmp_path))
+
+    expected_values = [str(i) for i in range(100)]
+    written_values = [row["text"] for row in ray.data.read_text(tmp_path).take_all()]
+    assert sorted(written_values) == sorted(expected_values)
+
+
+@pytest.mark.parametrize("num_rows", [0, 1])
+def test_write_preserves_user_directory(num_rows, tmp_path, ray_start_regular_shared):
+    class MockFileDatasink(BlockBasedFileDatasink):
+        def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+            file.write(b"")
+
+    ds = ray.data.range(num_rows)
+    path = os.path.join(tmp_path, "test")
+    os.mkdir(path)  # User-created directory
+
+    ds.write_datasink(MockFileDatasink(path=path))
+
+    assert os.path.isdir(path)
+
+
+def test_write_creates_dir(tmp_path, ray_start_regular_shared):
+    class MockFileDatasink(BlockBasedFileDatasink):
+        def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+            file.write(b"")
+
+    ds = ray.data.range(1)
+    path = os.path.join(tmp_path, "test")
+
+    ds.write_datasink(MockFileDatasink(path=path, try_create_dir=True))
+
+    assert os.path.isdir(path)
+
+
+@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
+def test_write_min_rows_per_file(tmp_path, ray_start_regular_shared, min_rows_per_file):
+    class MockFileDatasink(BlockBasedFileDatasink):
+        def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+            for _ in range(block.num_rows()):
+                file.write(b"row\n")
+
+    ds = ray.data.range(100, override_num_blocks=20)
+
+    ds.write_datasink(
+        MockFileDatasink(path=tmp_path, min_rows_per_file=min_rows_per_file)
+    )
+
+    num_rows_written_total = 0
+    for filename in os.listdir(tmp_path):
+        with open(os.path.join(tmp_path, filename), "r") as file:
+            num_rows_written = len(file.read().splitlines())
+            assert num_rows_written == min_rows_per_file
+            num_rows_written_total += num_rows_written
+
+    assert num_rows_written_total == 100
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_hudi.py

@@ -0,0 +1,161 @@
+import os
+import zipfile
+
+import pytest
+from packaging.version import parse as parse_version
+from pytest_lazy_fixtures import lf as lazy_fixture
+
+import ray
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.data.datasource.path_util import (
+    _resolve_paths_and_filesystem,
+    _unwrap_protocol,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+MIN_PYARROW_VERSION_FOR_HUDI = parse_version("11.0.0")
+PYARROW_VERSION = get_pyarrow_version()
+PYARROW_VERSION_MEETS_REQUIREMENT = (
+    PYARROW_VERSION and PYARROW_VERSION >= MIN_PYARROW_VERSION_FOR_HUDI
+)
+PYARROW_HUDI_TEST_SKIP_REASON = (
+    f"Hudi only supported if pyarrow >= {MIN_PYARROW_VERSION_FOR_HUDI}"
+)
+
+
+def _extract_testing_table(fixture_path: str, table_dir: str, target_dir: str) -> str:
+    with zipfile.ZipFile(fixture_path, "r") as zip_ref:
+        zip_ref.extractall(target_dir)
+    return os.path.join(target_dir, table_dir)
+
+
+def _get_hudi_table_path(fs, data_path, table_name, testing_dir="test_hudi") -> str:
+    setup_data_path = _unwrap_protocol(data_path)
+    target_testing_dir = os.path.join(setup_data_path, testing_dir)
+    fixture_path, _ = _resolve_paths_and_filesystem(
+        f"example://hudi-tables/{table_name}.zip", fs
+    )
+    return _extract_testing_table(fixture_path[0], table_name, target_testing_dir)
+
+
+@pytest.mark.skipif(
+    not PYARROW_VERSION_MEETS_REQUIREMENT,
+    reason=PYARROW_HUDI_TEST_SKIP_REASON,
+)
+@pytest.mark.parametrize(
+    "fs,data_path",
+    [
+        (None, lazy_fixture("local_path")),
+        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
+    ],
+)
+def test_hudi_snapshot_query_v6_trips_table(ray_start_regular_shared, fs, data_path):
+    table_path = _get_hudi_table_path(fs, data_path, "v6_trips_8i1u")
+
+    ds = ray.data.read_hudi(table_path, filters=[("city", "=", "san_francisco")])
+
+    assert ds.schema().names == [
+        "_hoodie_commit_time",
+        "_hoodie_commit_seqno",
+        "_hoodie_record_key",
+        "_hoodie_partition_path",
+        "_hoodie_file_name",
+        "ts",
+        "uuid",
+        "rider",
+        "driver",
+        "fare",
+        "city",
+    ]
+    assert ds.count() == 4
+    rows = (
+        ds.select_columns(["_hoodie_commit_time", "ts", "rider", "fare"])
+        .sort("fare")
+        .take_all()
+    )
+    first_commit = "20250715043008154"
+    second_commit = "20250715043011090"
+    assert rows == [
+        {
+            "_hoodie_commit_time": first_commit,
+            "ts": 1695159649087,
+            "rider": "rider-A",
+            "fare": 19.10,
+        },
+        {
+            "_hoodie_commit_time": second_commit,
+            "ts": 1695046462179,
+            "rider": "rider-D",
+            "fare": 25.0,
+        },
+        {
+            "_hoodie_commit_time": first_commit,
+            "ts": 1695091554788,
+            "rider": "rider-C",
+            "fare": 27.70,
+        },
+        {
+            "_hoodie_commit_time": first_commit,
+            "ts": 1695332066204,
+            "rider": "rider-E",
+            "fare": 93.50,
+        },
+    ]
+
+
+@pytest.mark.skipif(
+    not PYARROW_VERSION_MEETS_REQUIREMENT,
+    reason=PYARROW_HUDI_TEST_SKIP_REASON,
+)
+@pytest.mark.parametrize(
+    "fs,data_path",
+    [
+        (None, lazy_fixture("local_path")),
+        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
+    ],
+)
+def test_hudi_incremental_query_v6_trips_table(ray_start_regular_shared, fs, data_path):
+    table_path = _get_hudi_table_path(fs, data_path, "v6_trips_8i1u")
+
+    first_commit = "20250715043008154"
+    second_commit = "20250715043011090"
+    ds = ray.data.read_hudi(
+        table_path,
+        query_type="incremental",
+        hudi_options={
+            "hoodie.read.file_group.start_timestamp": first_commit,
+            "hoodie.read.file_group.end_timestamp": second_commit,
+        },
+    )
+
+    assert ds.schema().names == [
+        "_hoodie_commit_time",
+        "_hoodie_commit_seqno",
+        "_hoodie_record_key",
+        "_hoodie_partition_path",
+        "_hoodie_file_name",
+        "ts",
+        "uuid",
+        "rider",
+        "driver",
+        "fare",
+        "city",
+    ]
+    assert ds.count() == 1
+    rows = ds.select_columns(["_hoodie_commit_time", "ts", "rider", "fare"]).take_all()
+    assert rows == [
+        {
+            "_hoodie_commit_time": second_commit,
+            "ts": 1695046462179,
+            "rider": "rider-D",
+            "fare": 25.0,
+        },
+    ]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_huggingface.py

@@ -0,0 +1,464 @@
+from unittest.mock import MagicMock, patch
+
+import datasets
+import pyarrow
+import pytest
+import requests
+from packaging.version import Version
+
+import ray
+from ray.data.dataset import Dataset, MaterializedDataset
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def mock_hf_dataset():
+    """Create a mock HuggingFace dataset for testing."""
+    texts = [
+        "Climate change is a serious threat to our planet",
+        "We need to take action on global warming",
+        "Renewable energy is the future",
+        "Fossil fuels are destroying the environment",
+        "Solar power is becoming more affordable",
+        "Wind energy is growing rapidly",
+        "Electric vehicles are the way forward",
+        "Carbon emissions must be reduced",
+        "Green technology is advancing quickly",
+        "Sustainability is important for future generations",
+        "Climate science is well established",
+        "Ocean levels are rising due to warming",
+        "Extreme weather events are increasing",
+        "Biodiversity loss is accelerating",
+        "Deforestation contributes to climate change",
+        "Clean energy jobs are growing",
+        "Energy efficiency saves money",
+        "Public transportation reduces emissions",
+        "Plant-based diets help the environment",
+        "Recycling is essential for sustainability",
+    ]
+
+    # Create labels array with exactly the same length as texts
+    labels = [i % 2 for i in range(len(texts))]  # Alternating 0s and 1s
+
+    return datasets.Dataset.from_dict(
+        {
+            "text": texts,
+            "label": labels,
+        }
+    )
+
+
+@pytest.fixture
+def mock_hf_dataset_dict(mock_hf_dataset):
+    """Create a mock HuggingFace DatasetDict for testing."""
+    return datasets.DatasetDict({"train": mock_hf_dataset})
+
+
+@pytest.fixture
+def mock_hf_iterable_dataset():
+    """Create a mock HuggingFace IterableDataset for testing."""
+    texts = [
+        "Streaming climate tweet 1: The planet is warming",
+        "Streaming climate tweet 2: Renewable energy is key",
+        "Streaming climate tweet 3: We must act now",
+        "Streaming climate tweet 4: Solar panels everywhere",
+        "Streaming climate tweet 5: Wind turbines are beautiful",
+        "Streaming climate tweet 6: Electric cars are the future",
+        "Streaming climate tweet 7: Carbon neutral by 2050",
+        "Streaming climate tweet 8: Green energy revolution",
+        "Streaming climate tweet 9: Climate action needed",
+        "Streaming climate tweet 10: Sustainable development",
+        "Streaming climate tweet 11: Ocean conservation",
+        "Streaming climate tweet 12: Forest protection",
+        "Streaming climate tweet 13: Clean air matters",
+        "Streaming climate tweet 14: Water conservation",
+        "Streaming climate tweet 15: Biodiversity protection",
+    ]
+
+    labels = [1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1]
+
+    dataset = datasets.Dataset.from_dict(
+        {
+            "text": texts,
+            "label": labels,
+        }
+    )
+    iterable_dataset = dataset.to_iterable_dataset()
+    iterable_dataset.expected_count = len(texts)
+    return iterable_dataset
+
+
+@pytest.fixture
+def mock_parquet_urls():
+    """Fixture providing mock parquet URLs for testing."""
+    return [
+        "https://huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
+        "https://huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
+    ]
+
+
+@pytest.fixture
+def mock_resolved_urls():
+    """Fixture providing mock resolved URLs (after HTTP redirects) for testing."""
+    return [
+        "https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
+        "https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
+    ]
+
+
+@pytest.fixture
+def mock_ray_dataset(mock_hf_dataset):
+    """Fixture providing a mock Ray dataset that matches the mock HuggingFace dataset."""
+    return ray.data.from_items(
+        [
+            {"text": text, "label": label}
+            for text, label in zip(mock_hf_dataset["text"], mock_hf_dataset["label"])
+        ]
+    )
+
+
+@pytest.fixture
+def mock_successful_http_responses(mock_parquet_urls):
+    """Fixture providing mock successful HTTP responses for URL resolution."""
+    mock_responses = []
+    for url in mock_parquet_urls:
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.url = url
+        mock_responses.append(mock_response)
+    return mock_responses
+
+
+@pytest.fixture
+def mock_redirected_http_responses(mock_parquet_urls, mock_resolved_urls):
+    """Fixture providing mock HTTP responses that simulate redirects."""
+    mock_responses = []
+    for original_url, resolved_url in zip(mock_parquet_urls, mock_resolved_urls):
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.url = resolved_url
+        mock_responses.append(mock_response)
+    return mock_responses
+
+
+@pytest.fixture
+def mock_huggingface_datasource():
+    """Fixture providing the HuggingFaceDatasource class for mocking."""
+    from ray.data._internal.datasource.huggingface_datasource import (
+        HuggingFaceDatasource,
+    )
+
+    return HuggingFaceDatasource
+
+
+def verify_http_requests(mock_requests_head, expected_urls):
+    """Verify that HTTP requests were made correctly."""
+    assert mock_requests_head.call_count == len(expected_urls)
+
+    for i, url in enumerate(expected_urls):
+        call_args = mock_requests_head.call_args_list[i]
+        assert call_args[0][0] == url
+        assert call_args[1]["allow_redirects"] is True
+        assert call_args[1]["timeout"] == 5
+
+
+def verify_read_parquet_call(mock_read_parquet, expected_urls):
+    """Verify that read_parquet was called with correct parameters."""
+    mock_read_parquet.assert_called_once()
+    call_args = mock_read_parquet.call_args
+
+    # Check that the parquet URLs were passed
+    assert call_args[0][0] == expected_urls
+
+    # Check that the filesystem is HTTPFileSystem
+    assert "filesystem" in call_args[1]
+    assert "HTTPFileSystem" in str(type(call_args[1]["filesystem"]))
+
+    # Check that retry_exceptions includes FileNotFoundError and ClientResponseError
+    assert "ray_remote_args" in call_args[1]
+    assert FileNotFoundError in call_args[1]["ray_remote_args"]["retry_exceptions"]
+
+
+def verify_dataset_creation(ds, mock_hf_dataset):
+    """Verify that the dataset was created successfully."""
+    assert isinstance(ds, MaterializedDataset)
+    assert ds.count() == mock_hf_dataset.num_rows
+
+
+def setup_parquet_mocks(
+    mock_huggingface_datasource,
+    mock_parquet_urls,
+    mock_http_responses,
+    mock_ray_dataset,
+):
+    """Setup common mocking pattern for parquet-based tests."""
+    patches = []
+
+    # Mock the list_parquet_urls_from_dataset method
+    datasource_patch = patch.object(
+        mock_huggingface_datasource,
+        "list_parquet_urls_from_dataset",
+        return_value=mock_parquet_urls,
+    )
+    patches.append(datasource_patch)
+
+    # Mock the requests.head calls
+    requests_patch = patch("requests.head")
+    patches.append(requests_patch)
+
+    # Mock the read_parquet function
+    read_parquet_patch = patch("ray.data.read_api.read_parquet")
+    patches.append(read_parquet_patch)
+
+    # Start all patches
+    datasource_mock = datasource_patch.start()
+    requests_mock = requests_patch.start()
+    read_parquet_mock = read_parquet_patch.start()
+
+    # Configure mocks
+    requests_mock.side_effect = mock_http_responses
+    read_parquet_mock.return_value = mock_ray_dataset
+
+    return datasource_mock, requests_mock, read_parquet_mock, patches
+
+
+def hfds_assert_equals(hfds: datasets.Dataset, ds: Dataset):
+    hfds_table = hfds.data.table
+    ds_table = pyarrow.concat_tables([ray.get(tbl) for tbl in ds.to_arrow_refs()])
+
+    sorting = [(name, "descending") for name in hfds_table.column_names]
+    hfds_table = hfds_table.sort_by(sorting)
+    ds_table = ds_table.sort_by(sorting)
+
+    assert hfds_table.equals(ds_table)
+
+
+@pytest.mark.parametrize("num_par", [1, 4])
+def test_from_huggingface(mock_hf_dataset_dict, ray_start_regular_shared, num_par):
+    # Check that DatasetDict is not directly supported.
+    assert isinstance(mock_hf_dataset_dict, datasets.DatasetDict)
+    with pytest.raises(
+        DeprecationWarning,
+        match="You provided a Hugging Face DatasetDict",
+    ):
+        ray.data.from_huggingface(mock_hf_dataset_dict)
+
+    ray_datasets = {
+        "train": ray.data.from_huggingface(
+            mock_hf_dataset_dict["train"], override_num_blocks=num_par
+        ),
+    }
+
+    assert isinstance(ray_datasets["train"], ray.data.Dataset)
+    hfds_assert_equals(mock_hf_dataset_dict["train"], ray_datasets["train"])
+
+    # Test reading in a split Hugging Face dataset yields correct individual datasets
+    base_hf_dataset = mock_hf_dataset_dict["train"]
+    hf_dataset_split = base_hf_dataset.train_test_split(test_size=0.2)
+    ray_dataset_split_train = ray.data.from_huggingface(hf_dataset_split["train"])
+    assert ray_dataset_split_train.count() == hf_dataset_split["train"].num_rows
+
+
+@pytest.mark.skipif(
+    datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
+    reason="IterableDataset.iter() added in 2.8.0",
+)
+@pytest.mark.skipif(
+    Version(pyarrow.__version__) < Version("8.0.0"),
+    reason="pyarrow.Table.to_reader() added in 8.0.0",
+)
+# Note, pandas is excluded here because IterableDatasets do not support pandas format.
+@pytest.mark.parametrize(
+    "batch_format",
+    [None, "numpy", "arrow", "torch", "tensorflow", "jax"],
+)
+def test_from_huggingface_streaming(
+    mock_hf_iterable_dataset, batch_format, ray_start_regular_shared
+):
+    hfds = mock_hf_iterable_dataset.with_format(batch_format)
+    assert isinstance(hfds, datasets.IterableDataset)
+
+    ds = ray.data.from_huggingface(hfds)
+    expected_count = mock_hf_iterable_dataset.expected_count
+    assert ds.count() == expected_count
+
+
+@pytest.mark.skipif(
+    datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
+    reason="IterableDataset.iter() added in 2.8.0",
+)
+def test_from_huggingface_dynamic_generated(ray_start_regular_shared):
+    # https://github.com/ray-project/ray/issues/49529
+    # Mock the dynamic dataset loading
+    mock_dataset = datasets.Dataset.from_dict(
+        {
+            "text": [
+                "dynamic tweet 1",
+                "dynamic tweet 2",
+                "dynamic tweet 3",
+                "dynamic tweet 4",
+                "dynamic tweet 5",
+            ],
+            "label": [0, 1, 0, 1, 0],
+        }
+    )
+    mock_iterable = mock_dataset.to_iterable_dataset()
+
+    with patch("datasets.load_dataset", return_value=mock_iterable):
+        hfds = datasets.load_dataset(
+            "dataset-org/dream",
+            split="test",
+            streaming=True,
+            trust_remote_code=True,
+        )
+        ds = ray.data.from_huggingface(hfds)
+        ds.take(1)
+
+
+@pytest.mark.parametrize("override_num_blocks", [1, 2, 4, 8])
+def test_from_huggingface_override_num_blocks(
+    mock_hf_dataset, ray_start_regular_shared, override_num_blocks
+):
+    """Test that override_num_blocks works correctly with HuggingFace datasets."""
+    hf_train = mock_hf_dataset
+
+    ds_subset = ray.data.from_huggingface(
+        hf_train, override_num_blocks=override_num_blocks
+    )
+
+    assert isinstance(ds_subset, MaterializedDataset)
+
+    # Verify number of blocks
+    assert ds_subset.num_blocks() == override_num_blocks
+
+    # Verify data integrity
+    assert ds_subset.count() == hf_train.num_rows
+    hfds_assert_equals(hf_train, ds_subset)
+
+    # Test with a smaller subset to test edge cases
+    small_size = max(override_num_blocks * 3, 10)
+    hf_small = hf_train.select(range(min(small_size, hf_train.num_rows)))
+    ds_small = ray.data.from_huggingface(
+        hf_small, override_num_blocks=override_num_blocks
+    )
+
+    # Verify number of blocks
+    assert ds_small.num_blocks() == override_num_blocks
+
+    # Verify data integrity
+    assert ds_small.count() == hf_small.num_rows
+    hfds_assert_equals(hf_small, ds_small)
+
+
+def test_from_huggingface_with_parquet_files(
+    mock_hf_dataset,
+    ray_start_regular_shared,
+    mock_parquet_urls,
+    mock_ray_dataset,
+    mock_successful_http_responses,
+    mock_huggingface_datasource,
+):
+    """Test the distributed read path when parquet file URLs are available."""
+    datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
+        mock_huggingface_datasource,
+        mock_parquet_urls,
+        mock_successful_http_responses,
+        mock_ray_dataset,
+    )
+
+    try:
+        ds = ray.data.from_huggingface(mock_hf_dataset)
+
+        # Verify HTTP requests
+        verify_http_requests(requests_mock, mock_parquet_urls)
+
+        # Verify read_parquet call
+        verify_read_parquet_call(read_parquet_mock, mock_parquet_urls)
+
+        # Verify dataset creation
+        verify_dataset_creation(ds, mock_hf_dataset)
+
+    finally:
+        # Stop all patches
+        for patch_obj in patches:
+            patch_obj.stop()
+
+
+def test_from_huggingface_with_resolved_urls(
+    mock_hf_dataset,
+    ray_start_regular_shared,
+    mock_parquet_urls,
+    mock_resolved_urls,
+    mock_ray_dataset,
+    mock_redirected_http_responses,
+    mock_huggingface_datasource,
+):
+    """Test the URL resolution logic when HTTP redirects are encountered."""
+    datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
+        mock_huggingface_datasource,
+        mock_parquet_urls,
+        mock_redirected_http_responses,
+        mock_ray_dataset,
+    )
+
+    try:
+        ds = ray.data.from_huggingface(mock_hf_dataset)
+
+        # Verify HTTP requests
+        verify_http_requests(requests_mock, mock_parquet_urls)
+
+        # Verify read_parquet call with resolved URLs
+        verify_read_parquet_call(read_parquet_mock, mock_resolved_urls)
+
+        # Verify dataset creation
+        verify_dataset_creation(ds, mock_hf_dataset)
+
+    finally:
+        # Stop all patches
+        for patch_obj in patches:
+            patch_obj.stop()
+
+
+def test_from_huggingface_url_resolution_failures(
+    mock_hf_dataset,
+    ray_start_regular_shared,
+    mock_parquet_urls,
+    mock_ray_dataset,
+    mock_huggingface_datasource,
+):
+    """Test URL resolution failures fall back to single node read."""
+    # Convert the mock dataset to an IterableDataset so it uses the read_datasource fallback
+    mock_iterable_dataset = mock_hf_dataset.to_iterable_dataset()
+
+    with patch.object(
+        mock_huggingface_datasource,
+        "list_parquet_urls_from_dataset",
+        return_value=mock_parquet_urls,
+    ):
+        # Mock the requests.head calls to simulate failures
+        with patch("requests.head") as mock_requests_head:
+            # Configure mock to raise an exception for all URLs
+            mock_requests_head.side_effect = requests.RequestException(
+                "Connection failed"
+            )
+
+            # Mock the fallback path
+            with patch("ray.data.read_api.read_datasource") as mock_read_datasource:
+                mock_read_datasource.return_value = mock_ray_dataset
+
+                ds = ray.data.from_huggingface(mock_iterable_dataset)
+
+                # Verify that requests.head was called for each URL
+                assert mock_requests_head.call_count == len(mock_parquet_urls)
+
+                # Verify that the fallback read_datasource was called
+                mock_read_datasource.assert_called_once()
+
+                # Verify the dataset was created successfully via fallback
+                verify_dataset_creation(ds, mock_hf_dataset)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_image.py

@@ -0,0 +1,198 @@
+import os
+from typing import Dict
+
+import numpy as np
+import pytest
+from fsspec.implementations.local import LocalFileSystem
+from PIL import Image
+
+import ray
+from ray.data._internal.datasource.image_datasource import (
+    ImageDatasource,
+    ImageFileMetadataProvider,
+)
+from ray.data._internal.tensor_extensions.arrow import (
+    get_arrow_extension_fixed_shape_tensor_types,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+
+class TestReadImages:
+    def test_basic(self, ray_start_regular_shared):
+        # "simple" contains three 32x32 RGB images.
+        ds = ray.data.read_images("example://image-datasets/simple")
+        assert ds.schema().names == ["image"]
+        column_type = ds.schema().types[0]
+        assert isinstance(column_type, get_arrow_extension_fixed_shape_tensor_types())
+        assert all(record["image"].shape == (32, 32, 3) for record in ds.take())
+
+    @pytest.mark.parametrize("num_threads", [-1, 0, 1, 2, 4])
+    def test_multi_threading(self, ray_start_regular_shared, num_threads, monkeypatch):
+        monkeypatch.setattr(
+            ray.data._internal.datasource.image_datasource.ImageDatasource,
+            "_NUM_THREADS_PER_TASK",
+            num_threads,
+        )
+        ds = ray.data.read_images(
+            "example://image-datasets/simple",
+            override_num_blocks=1,
+            include_paths=True,
+        )
+        paths = [item["path"][-len("image1.jpg") :] for item in ds.take_all()]
+        if num_threads > 1:
+            # If there are more than 1 threads, the order is not guaranteed.
+            paths = sorted(paths)
+        expected_paths = ["image1.jpg", "image2.jpg", "image3.jpg"]
+        assert paths == expected_paths
+
+    def test_size(self, ray_start_regular_shared):
+        # "different-sizes" contains RGB images with different heights and widths.
+        ds = ray.data.read_images(
+            "example://image-datasets/different-sizes", size=(32, 32)
+        )
+        assert all(record["image"].shape == (32, 32, 3) for record in ds.take())
+
+    def test_different_sizes(self, ray_start_regular_shared):
+        ds = ray.data.read_images("example://image-datasets/different-sizes")
+        assert sorted(record["image"].shape for record in ds.take()) == [
+            (16, 16, 3),
+            (32, 32, 3),
+            (64, 64, 3),
+        ]
+
+    @pytest.mark.parametrize("size", [(-32, 32), (32, -32), (-32, -32)])
+    def test_invalid_size(self, ray_start_regular_shared, size):
+        with pytest.raises(ValueError):
+            ray.data.read_images("example://image-datasets/simple", size=size)
+
+    @pytest.mark.parametrize(
+        "mode, expected_shape", [("L", (32, 32)), ("RGB", (32, 32, 3))]
+    )
+    def test_mode(
+        self,
+        mode,
+        expected_shape,
+        ray_start_regular_shared,
+    ):
+        # "different-modes" contains 32x32 images with modes "CMYK", "L", and "RGB"
+        ds = ray.data.read_images("example://image-datasets/different-modes", mode=mode)
+        assert all([record["image"].shape == expected_shape for record in ds.take()])
+
+    def test_e2e_prediction(self, shutdown_only):
+        import torch
+        from torchvision import transforms
+        from torchvision.models import resnet18
+
+        ray.shutdown()
+        ray.init(num_cpus=2)
+
+        dataset = ray.data.read_images("example://image-datasets/simple")
+        transform = transforms.ToTensor()
+
+        def preprocess(batch: Dict[str, np.ndarray]):
+            return {"out": np.stack([transform(image) for image in batch["image"]])}
+
+        dataset = dataset.map_batches(preprocess, batch_format="numpy")
+
+        class Predictor:
+            def __init__(self):
+                self.model = resnet18(pretrained=True)
+
+            def __call__(self, batch: Dict[str, np.ndarray]):
+                with torch.inference_mode():
+                    torch_tensor = torch.as_tensor(batch["out"])
+                    return {"prediction": self.model(torch_tensor)}
+
+        predictions = dataset.map_batches(
+            Predictor, compute=ray.data.ActorPoolStrategy(min_size=1), batch_size=4096
+        )
+
+        for _ in predictions.iter_batches():
+            pass
+
+    @pytest.mark.parametrize(
+        "image_size,image_mode,expected_size,expected_ratio",
+        [(64, "RGB", 30000, 4), (32, "L", 3500, 0.5), (256, "RGBA", 750000, 85)],
+    )
+    def test_data_size_estimate(
+        self,
+        ray_start_regular_shared,
+        image_size,
+        image_mode,
+        expected_size,
+        expected_ratio,
+    ):
+        root = "example://image-datasets/different-sizes"
+        ds = ray.data.read_images(
+            root, size=(image_size, image_size), mode=image_mode, override_num_blocks=1
+        )
+
+        data_size = ds.size_bytes()
+        assert data_size >= 0, "estimated data size is out of expected bound"
+        data_size = ds.materialize().size_bytes()
+        assert data_size >= 0, "actual data size is out of expected bound"
+
+        datasource = ImageDatasource(
+            paths=[root],
+            size=(image_size, image_size),
+            mode=image_mode,
+            filesystem=LocalFileSystem(),
+            partitioning=None,
+            meta_provider=ImageFileMetadataProvider(),
+        )
+        assert (
+            datasource._encoding_ratio >= expected_ratio
+            and datasource._encoding_ratio <= expected_ratio * 1.5
+        ), "encoding ratio is out of expected bound"
+        data_size = datasource.estimate_inmemory_data_size()
+        assert data_size >= 0, "estimated data size is out of expected bound"
+
+    def test_dynamic_block_split(ray_start_regular_shared):
+        ctx = ray.data.context.DataContext.get_current()
+        target_max_block_size = ctx.target_max_block_size
+        # Reduce target max block size to trigger block splitting on small input.
+        # Otherwise we have to generate big input files, which is unnecessary.
+        ctx.target_max_block_size = 1
+        try:
+            root = "example://image-datasets/simple"
+            ds = ray.data.read_images(root, override_num_blocks=1)
+            assert ds._logical_plan.initial_num_blocks() == 1
+            ds = ds.materialize()
+            # Verify dynamic block splitting taking effect to generate more blocks.
+            assert ds._logical_plan.initial_num_blocks() == 3
+
+            # Test union of same datasets
+            union_ds = ds.union(ds, ds, ds).materialize()
+            assert union_ds._logical_plan.initial_num_blocks() == 12
+        finally:
+            ctx.target_max_block_size = target_max_block_size
+
+    def test_unidentified_image_error(ray_start_regular_shared, tmp_path):
+        path = str(tmp_path / "invalid.png")
+        with open(path, "wb") as file:
+            file.write(b"spam")  # Invalid bytes for a PNG file
+
+        with pytest.raises(ValueError):
+            ray.data.read_images(paths=file.name).materialize()
+
+
+class TestWriteImages:
+    def test_write_images(ray_start_regular_shared, tmp_path):
+        ds = ray.data.read_images("example://image-datasets/simple")
+        ds.write_images(
+            path=tmp_path,
+            column="image",
+        )
+
+        assert len(os.listdir(tmp_path)) == ds.count()
+
+        for filename in os.listdir(tmp_path):
+            path = os.path.join(tmp_path, filename)
+            Image.open(path)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_json.py

@@ -0,0 +1,470 @@
+import gzip
+import json
+import os
+
+import pandas as pd
+import pyarrow as pa
+import pyarrow.fs as fs
+import pyarrow.json as pajson
+import pytest
+
+import ray
+from ray.data import Schema
+from ray.data._internal.datasource.json_datasource import PandasJSONDatasource
+from ray.data._internal.pandas_block import PandasBlockBuilder
+from ray.data._internal.util import rows_same
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_based_datasource import (
+    FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+# Set the test timeout to 6 minutes
+pytestmark = pytest.mark.timeout(360)
+
+
+def test_json_read(
+    ray_start_regular_shared, target_max_block_size_infinite_or_default, tmp_path
+):
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path1 = os.path.join(tmp_path, "test1.json")
+    df1.to_json(path1, orient="records", lines=True)
+    ds = ray.data.read_json(path1)
+    dsdf = ds.to_pandas()
+    pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Metadata ops.
+    assert ds.count() == 3
+    assert ds.input_files() == [path1]
+    assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
+
+
+def test_zipped_json_read(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path1 = os.path.join(tmp_path, "test1.json.gz")
+    df1.to_json(path1, compression="gzip", orient="records", lines=True)
+    ds = ray.data.read_json(path1)
+    dsdf = ds.to_pandas()
+    pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Metadata ops.
+    assert ds.count() == 3
+    assert ds.input_files() == [path1]
+
+
+def test_read_json_fallback_from_pyarrow_failure(
+    ray_start_regular_shared, local_path, target_max_block_size_infinite_or_default
+):
+    # Try to read this with read_json() to trigger fallback logic
+    # to read bytes with json.load().
+    data = [{"one": [1]}, {"one": [1, 2]}]
+    path1 = os.path.join(local_path, "test1.json")
+    with open(path1, "w") as f:
+        json.dump(data, f)
+
+    # pyarrow.json cannot read JSONs containing arrays of different lengths.
+    from pyarrow import ArrowInvalid
+
+    with pytest.raises(ArrowInvalid):
+        pajson.read_json(path1)
+
+    # Ray Data successfully reads this in by
+    # falling back to json.load() when pyarrow fails.
+    ds = ray.data.read_json(path1)
+    assert ds.take_all() == data
+
+
+def test_json_read_with_read_options(
+    ray_start_regular_shared,
+    tmp_path,
+    target_max_block_size_infinite_or_default,
+):
+    # Arrow's JSON ReadOptions isn't serializable in pyarrow < 8.0.0, so this test
+    # covers our custom ReadOptions serializer.
+    # TODO(Clark): Remove this test and our custom serializer once we require
+    # pyarrow >= 8.0.0.
+
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path1 = os.path.join(tmp_path, "test1.json")
+    df1.to_json(path1, orient="records", lines=True)
+    ds = ray.data.read_json(
+        path1,
+        read_options=pajson.ReadOptions(use_threads=False, block_size=2**30),
+    )
+    dsdf = ds.to_pandas()
+    pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Test metadata ops.
+    assert ds.count() == 3
+    assert ds.input_files() == [path1]
+    assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
+
+
+def test_json_read_with_parse_options(
+    ray_start_regular_shared,
+    tmp_path,
+    target_max_block_size_infinite_or_default,
+):
+    # Arrow's JSON ParseOptions isn't serializable in pyarrow < 8.0.0, so this test
+    # covers our custom ParseOptions serializer, similar to ReadOptions in above test.
+    # TODO(chengsu): Remove this test and our custom serializer once we require
+    # pyarrow >= 8.0.0.
+
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path1 = os.path.join(tmp_path, "test1.json")
+    df1.to_json(path1, orient="records", lines=True)
+    ds = ray.data.read_json(
+        path1,
+        parse_options=pajson.ParseOptions(
+            explicit_schema=pa.schema([("two", pa.string())]),
+            unexpected_field_behavior="ignore",
+        ),
+    )
+    dsdf = ds.to_pandas()
+    assert len(dsdf.columns) == 1
+    pd.testing.assert_series_equal(df1["two"].astype(dsdf["two"].dtype), dsdf["two"])
+    # Test metadata ops.
+    assert ds.count() == 3
+    assert ds.input_files() == [path1]
+    assert ds.schema() == Schema(pa.schema([("two", pa.string())]))
+
+
+@pytest.mark.parametrize("override_num_blocks", [None, 1, 3])
+def test_jsonl_lists(
+    ray_start_regular_shared,
+    tmp_path,
+    override_num_blocks,
+    target_max_block_size_infinite_or_default,
+):
+    """Test JSONL with mixed types and schemas."""
+    data = [
+        ["ray", "rocks", "hello"],
+        ["oh", "no"],
+        ["rocking", "with", "ray"],
+    ]
+
+    path = os.path.join(tmp_path, "test.jsonl")
+    with open(path, "w") as f:
+        for record in data:
+            json.dump(record, f)
+            f.write("\n")
+
+    ds = ray.data.read_json(path, lines=True, override_num_blocks=override_num_blocks)
+    result = ds.take_all()
+
+    assert result[0] == {"0": "ray", "1": "rocks", "2": "hello"}
+    assert result[1] == {"0": "oh", "1": "no", "2": None}
+    assert result[2] == {"0": "rocking", "1": "with", "2": "ray"}
+
+
+def test_jsonl_mixed_types(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    """Test JSONL with mixed types and schemas."""
+    data = [
+        {"a": 1, "b": {"c": 2}},  # Nested dict
+        {"a": 1, "b": {"c": 3}},  # Nested dict
+        {"a": 1, "b": {"c": {"hello": "world"}}},  # Mixed Schema
+    ]
+
+    path = os.path.join(tmp_path, "test.jsonl")
+    with open(path, "w") as f:
+        for record in data:
+            json.dump(record, f)
+            f.write("\n")
+
+    ds = ray.data.read_json(path, lines=True)
+    result = ds.take_all()
+
+    assert result[0] == data[0]  # Dict stays as is
+    assert result[1] == data[1]
+    assert result[2] == data[2]
+
+
+def test_json_write(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    input_df = pd.DataFrame({"id": [0]})
+    ds = ray.data.from_blocks([input_df])
+
+    ds.write_json(tmp_path)
+
+    output_df = pd.concat(
+        [
+            pd.read_json(os.path.join(tmp_path, filename), lines=True)
+            for filename in os.listdir(tmp_path)
+        ]
+    )
+
+    assert rows_same(input_df, output_df)
+
+
+@pytest.mark.parametrize("override_num_blocks", [None, 2])
+def test_json_roundtrip(
+    ray_start_regular_shared,
+    tmp_path,
+    override_num_blocks,
+    target_max_block_size_infinite_or_default,
+):
+    df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+
+    ds = ray.data.from_pandas([df], override_num_blocks=override_num_blocks)
+    ds.write_json(tmp_path)
+
+    ds2 = ray.data.read_json(tmp_path)
+    ds2df = ds2.to_pandas()
+    assert rows_same(ds2df, df)
+    for entry in ds2._execute().blocks:
+        assert (
+            # pyrefly: ignore[no-matching-overload]
+            BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
+            == entry.metadata.size_bytes
+        )
+
+
+def test_json_read_small_file_unit_block_size(
+    ray_start_regular_shared,
+    tmp_path,
+    target_max_block_size_infinite_or_default,
+):
+    """Test reading a small JSON file with unit block_size."""
+
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path1 = os.path.join(tmp_path, "test1.json")
+    df1.to_json(path1, orient="records", lines=True)
+    ds = ray.data.read_json(path1, read_options=pajson.ReadOptions(block_size=1))
+    dsdf = ds.to_pandas()
+    pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Test metadata ops.
+    assert ds.count() == 3
+    assert ds.input_files() == [path1]
+    assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
+
+
+def test_json_read_file_larger_than_block_size(
+    ray_start_regular_shared,
+    tmp_path,
+    target_max_block_size_infinite_or_default,
+):
+    """Test reading a JSON file larger than the block size."""
+    block_size = 1024
+    num_chars = 2500
+    num_rows = 3
+    df2 = pd.DataFrame(
+        {
+            "one": ["a" * num_chars for _ in range(num_rows)],
+            "two": ["b" * num_chars for _ in range(num_rows)],
+        }
+    )
+    path2 = os.path.join(tmp_path, "test2.json")
+    df2.to_json(path2, orient="records", lines=True)
+    ds = ray.data.read_json(
+        path2, read_options=pajson.ReadOptions(block_size=block_size)
+    )
+    dsdf = ds.to_pandas()
+    pd.testing.assert_frame_equal(df2.astype(dsdf.dtypes.to_dict()), dsdf)
+    # Test metadata ops.
+    assert ds.count() == num_rows
+    assert ds.input_files() == [path2]
+    assert ds.schema() == Schema(
+        pa.schema([("one", pa.string()), ("two", pa.string())])
+    )
+
+
+def test_json_read_negative_block_size_fallback(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    """Test reading JSON with negative block_size triggers fallback to json.load()."""
+
+    df3 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path3 = os.path.join(tmp_path, "test3.json")
+    df3.to_json(path3, orient="records", lines=True)
+
+    # Negative Buffer Size, fails with arrow but succeeds in fallback to json.load()
+    ds = ray.data.read_json(path3, read_options=pajson.ReadOptions(block_size=-1))
+    dsdf = ds.to_pandas()
+    pd.testing.assert_frame_equal(df3.astype(dsdf.dtypes.to_dict()), dsdf)
+
+
+def test_json_read_zero_block_size_failure(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    """Test reading JSON with zero block_size fails in both arrow and fallback."""
+
+    df3 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    path3 = os.path.join(tmp_path, "test3.json")
+    df3.to_json(path3, orient="records", lines=True)
+
+    # Zero Buffer Size, fails with arrow and fails in fallback to json.load()
+    with pytest.raises(json.decoder.JSONDecodeError, match="Extra data"):
+        ds = ray.data.read_json(path3, read_options=pajson.ReadOptions(block_size=0))
+        dsdf = ds.to_pandas()
+        assert dsdf.equals(df3)
+
+
+@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
+def test_write_min_rows_per_file(
+    tmp_path,
+    ray_start_regular_shared,
+    min_rows_per_file,
+    target_max_block_size_infinite_or_default,
+):
+    ray.data.range(100, override_num_blocks=20).write_json(
+        tmp_path, min_rows_per_file=min_rows_per_file
+    )
+
+    for filename in os.listdir(tmp_path):
+        with open(os.path.join(tmp_path, filename), "r") as file:
+            num_rows_written = len(file.read().splitlines())
+            assert num_rows_written == min_rows_per_file
+
+
+def test_mixed_gzipped_json_files(
+    ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
+):
+    # Create a non-empty gzipped JSON file
+    non_empty_file_path = os.path.join(tmp_path, "non_empty.json.gz")
+    data = [{"col1": "value1", "col2": "value2", "col3": "value3"}]
+    with gzip.open(non_empty_file_path, "wt", encoding="utf-8") as f:
+        for record in data:
+            json.dump(record, f)
+            f.write("\n")
+
+    # Create an empty gzipped JSON file
+    empty_file_path = os.path.join(tmp_path, "empty.json.gz")
+    with gzip.open(empty_file_path, "wt", encoding="utf-8"):
+        pass  # Write nothing to create an empty file
+
+    # Attempt to read both files with Ray
+    ds = ray.data.read_json(
+        [non_empty_file_path, empty_file_path],
+        arrow_open_stream_args={"compression": "gzip"},
+    )
+
+    # The dataset should only contain data from the non-empty file
+    assert ds.count() == 1
+    # Iterate through each row in the dataset and compare with the expected data
+    for row in ds.iter_rows():
+        assert row == data[0], f"Row {row} does not match expected {data[0]}"
+
+    # Verify the data content using take
+    retrieved_data = ds.take(1)[0]
+    assert (
+        retrieved_data == data[0]
+    ), f"Retrieved data {retrieved_data} does not match expected {data[0]}."
+
+
+def test_json_with_http_path_parallelization(
+    ray_start_regular_shared, httpserver, target_max_block_size_infinite_or_default
+):
+    num_files = FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
+    urls = []
+    for i in range(num_files):
+        httpserver.expect_request(f"/file{i}.json").respond_with_json({"id": i})
+        urls.append(httpserver.url_for(f"/file{i}.json"))
+
+    ds = ray.data.read_json(urls)
+    actual_rows = ds.take_all()
+
+    expected_rows = [{"id": i} for i in range(num_files)]
+    assert sorted(actual_rows, key=lambda row: row["id"]) == sorted(
+        expected_rows, key=lambda row: row["id"]
+    )
+
+
+class TestPandasJSONDatasource:
+    @pytest.mark.parametrize(
+        "data",
+        [{"a": []}, {"a": [1]}, {"a": [1, 2, 3]}],
+        ids=["empty", "single", "multiple"],
+    )
+    @pytest.mark.parametrize(
+        "compression,filename",
+        [("gzip", "test.json.gz"), ("infer", "test.json")],  # infer = default
+    )
+    def test_read_stream(
+        self,
+        data,
+        tmp_path,
+        compression,
+        filename,
+        target_max_block_size_infinite_or_default,
+    ):
+        # Setup test file.
+        df = pd.DataFrame(data)
+        path = os.path.join(tmp_path, filename)
+        df.to_json(path, orient="records", lines=True, compression=compression)
+
+        # Setup datasource.
+        local_filesystem = fs.LocalFileSystem()
+        source = PandasJSONDatasource(
+            path, target_output_size_bytes=1, filesystem=local_filesystem
+        )
+
+        # Read stream.
+        block_builder = PandasBlockBuilder()
+        with source._open_input_source(local_filesystem, path) as f:
+            for block in source._read_stream(f, path):
+                block_builder.add_block(block)
+        block = block_builder.build()
+
+        # Verify.
+        assert rows_same(block, df)
+
+    def test_read_stream_with_target_output_size_bytes(
+        self, tmp_path, target_max_block_size_infinite_or_default
+    ):
+        # Setup test file. It contains 16 lines, each line is 8 MiB.
+        df = pd.DataFrame({"data": ["a" * 8 * 1024 * 1024] * 16})
+        path = os.path.join(tmp_path, "test.json")
+        df.to_json(path, orient="records", lines=True)
+
+        # Setup datasource. It should read 32 MiB (4 lines) per output.
+        local_filesystem = fs.LocalFileSystem()
+        source = PandasJSONDatasource(
+            path,
+            target_output_size_bytes=32 * 1024 * 1024,
+            filesystem=local_filesystem,
+        )
+
+        # Read stream.
+        block_builder = PandasBlockBuilder()
+        with source._open_input_source(local_filesystem, path) as f:
+            for block in source._read_stream(f, path):
+                assert len(block) == 4
+                block_builder.add_block(block)
+        block = block_builder.build()
+
+        # Verify.
+        assert rows_same(block, df)
+
+    def test_read_stream_with_advanced_file_pointer(
+        self, tmp_path, target_max_block_size_infinite_or_default
+    ):
+        # Setup test file.
+        df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+        path = os.path.join(tmp_path, "test.json")
+        df.to_json(path, orient="records", lines=True)
+
+        # Setup datasource.
+        local_filesystem = fs.LocalFileSystem()
+        source = PandasJSONDatasource(
+            path, target_output_size_bytes=1, filesystem=local_filesystem
+        )
+
+        # Simulate retrying a stream read on a file handle that was already consumed.
+        block_builder = PandasBlockBuilder()
+        with source._open_input_source(local_filesystem, path) as f:
+            f.read(1)
+            for block in source._read_stream(f, path):
+                block_builder.add_block(block)
+        block = block_builder.build()
+
+        # Verify.
+        assert rows_same(block, df)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_lance.py

@@ -0,0 +1,390 @@
+import os
+
+import lance
+import pyarrow as pa
+import pytest
+from packaging.version import Version
+from pytest_lazy_fixtures import lf as lazy_fixture
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.data import Schema
+from ray.data._internal.datasource.lance_datasink import (
+    _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
+    LanceDatasink,
+    _write_fragment,
+)
+from ray.data.datasource import SaveMode
+from ray.data.datasource.path_util import _unwrap_protocol
+
+# Skip tests for older pylance versions (<=0.3.19) due to incompatible lance API changes with pyarrow v9.0.0
+pytestmark = pytest.mark.skipif(
+    Version(lance.__version__) <= Version("0.3.19"),
+    reason=f"pylance {lance.__version__} <= 0.3.19; API incompatible",
+)
+
+
+@pytest.mark.parametrize(
+    "fs,data_path",
+    [
+        (None, lazy_fixture("local_path")),
+        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
+        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
+        (
+            lazy_fixture("s3_fs_with_space"),
+            lazy_fixture("s3_path_with_space"),
+        ),  # Path contains space.
+        (
+            lazy_fixture("s3_fs_with_anonymous_crendential"),
+            lazy_fixture("s3_path_with_anonymous_crendential"),
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "batch_size",
+    [None, 100],
+)
+def test_lance_read_basic(fs, data_path, batch_size):
+    df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
+    setup_data_path = _unwrap_protocol(data_path)
+    path = os.path.join(setup_data_path, "test.lance")
+    lance.write_dataset(df1, path)
+
+    ds_lance = lance.dataset(path)
+    assert ds_lance is not None
+    df2 = pa.table(
+        {
+            "one": [1, 2, 3, 4, 5, 6],
+            "three": [4, 5, 8, 9, 12, 13],
+            "four": ["u", "v", "w", "x", "y", "z"],
+        }
+    )
+    ds_lance.merge(df2, "one")
+
+    if batch_size is None:
+        ds = ray.data.read_lance(path)
+    else:
+        ds = ray.data.read_lance(path, scanner_options={"batch_size": batch_size})
+
+    # Test metadata-only ops.
+    assert ds.count() == 6
+    assert ds.schema() == Schema(
+        pa.schema(
+            {
+                "one": pa.int64(),
+                "two": pa.string(),
+                "three": pa.int64(),
+                "four": pa.string(),
+            }
+        )
+    )
+
+    # Test read.
+    values = [[s["one"], s["two"]] for s in ds.take_all()]
+    assert sorted(values) == [
+        [1, "a"],
+        [2, "b"],
+        [3, "c"],
+        [4, "e"],
+        [5, "f"],
+        [6, "g"],
+    ]
+
+    # Test column projection.
+    ds = ray.data.read_lance(path, columns=["one"])
+    values = [s["one"] for s in ds.take_all()]
+    assert sorted(values) == [1, 2, 3, 4, 5, 6]
+    assert ds.schema().names == ["one"]
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_read_with_scanner_fragments(data_path):
+    table = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
+    setup_data_path = _unwrap_protocol(data_path)
+    path = os.path.join(setup_data_path, "test.lance")
+    dataset = lance.write_dataset(table, path, max_rows_per_file=2)
+    assert dataset is not None
+
+    fragments = dataset.get_fragments()
+    ds = ray.data.read_lance(path, scanner_options={"fragments": fragments[:1]})
+    values = [[s["one"], s["two"]] for s in ds.take_all()]
+    assert values == [
+        [2, "b"],
+        [1, "a"],
+    ]
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_read_many_files(data_path):
+    setup_data_path = _unwrap_protocol(data_path)
+    path = os.path.join(setup_data_path, "test.lance")
+    num_rows = 1024
+    data = pa.table({"id": pa.array(range(num_rows))})
+    lance.write_dataset(data, path, max_rows_per_file=1)
+
+    def test_lance():
+        ds = ray.data.read_lance(path)
+        return ds.count() == num_rows
+
+    wait_for_condition(test_lance, timeout=10)
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_write(data_path):
+    schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
+
+    ray.data.range(10).map(
+        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
+    ).write_lance(data_path, schema=schema)
+
+    ds = lance.dataset(data_path)
+    assert ds is not None
+    ds.count_rows() == 10
+    assert ds.schema.names == schema.names
+    # The schema is platform-dependent, because numpy uses int32 on Windows.
+    # So we observe the schema that is written and use that.
+    schema = ds.schema
+
+    tbl = ds.to_table()
+    assert sorted(tbl["id"].to_pylist()) == list(range(10))
+    assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(10)}
+
+    ray.data.range(10).map(
+        lambda x: {"id": x["id"] + 10, "str": f"str-{x['id'] + 10}"}
+    ).write_lance(data_path, mode=SaveMode.APPEND)
+
+    ds = lance.dataset(data_path)
+    assert ds is not None
+    ds.count_rows() == 20
+    tbl = ds.to_table()
+    assert sorted(tbl["id"].to_pylist()) == list(range(20))
+    assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(20)}
+
+    ray.data.range(10).map(
+        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
+    ).write_lance(data_path, schema=schema, mode=SaveMode.OVERWRITE)
+
+    ds = lance.dataset(data_path)
+    assert ds is not None
+    ds.count_rows() == 10
+    assert ds.schema == schema
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_write_create_errors_if_exists(data_path):
+    table_path = os.path.join(data_path, "my_table")
+    ds = ray.data.range(10)
+
+    # First CREATE succeeds on an empty destination.
+    ds.write_lance(table_path, mode=SaveMode.CREATE)
+    assert lance.dataset(table_path).count_rows() == 10
+
+    # A second CREATE must error instead of silently overwriting.
+    with pytest.raises(ValueError, match="already exists"):
+        ray.data.range(5).write_lance(table_path, mode=SaveMode.CREATE)
+
+    # Existing data is untouched.
+    assert lance.dataset(table_path).count_rows() == 10
+
+    # CREATE is also the default mode, so it must guard too.
+    with pytest.raises(ValueError, match="already exists"):
+        ray.data.range(5).write_lance(table_path)
+    assert lance.dataset(table_path).count_rows() == 10
+
+    # OVERWRITE replaces the existing data.
+    ray.data.range(5).write_lance(table_path, mode=SaveMode.OVERWRITE)
+    assert lance.dataset(table_path).count_rows() == 5
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_write_append_errors_if_missing(data_path):
+    table_path = os.path.join(data_path, "missing_table")
+    # APPEND surfaces Lance's own "not found" error. We don't pin the message,
+    # since it can change across Lance versions.
+    expected_errors: tuple[type[Exception], ...] = (
+        ValueError,
+        OSError,
+        FileNotFoundError,
+    )
+    with pytest.raises(expected_errors):
+        ray.data.range(5).write_lance(table_path, mode=SaveMode.APPEND)
+    assert not os.path.exists(table_path)
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_write_min_rows_per_file(data_path):
+    schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
+
+    ray.data.range(10).map(
+        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
+    ).write_lance(data_path, schema=schema, min_rows_per_file=100)
+
+    ds = lance.dataset(data_path)
+    assert ds is not None
+    assert ds.count_rows() == 10
+    assert ds.schema == schema
+
+    assert len(ds.get_fragments()) == 1
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_write_max_rows_per_file(data_path):
+    schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
+
+    ray.data.range(10).map(
+        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
+    ).write_lance(data_path, schema=schema, max_rows_per_file=1)
+
+    ds = lance.dataset(data_path)
+    assert ds is not None
+    assert ds.count_rows() == 10
+    assert ds.schema == schema
+
+    assert len(ds.get_fragments()) == 10
+
+
+@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
+def test_lance_read_with_version(data_path):
+    # Write an initial dataset (version 1)
+    df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
+    setup_data_path = _unwrap_protocol(data_path)
+    path = os.path.join(setup_data_path, "test_version.lance")
+    lance.write_dataset(df1, path)
+
+    # Merge new data to create a later version (latest)
+    ds_lance = lance.dataset(path)
+    assert ds_lance is not None
+    # Get the initial version
+    initial_version = ds_lance.version
+
+    df2 = pa.table(
+        {
+            "one": [1, 2, 3, 4, 5, 6],
+            "three": [4, 5, 8, 9, 12, 13],
+            "four": ["u", "v", "w", "x", "y", "z"],
+        }
+    )
+    ds_lance.merge(df2, "one")
+
+    # Default read should return the latest (merged) dataset.
+    ds_latest = ray.data.read_lance(path)
+
+    assert ds_latest.count() == 6
+    # Latest dataset should contain merged columns
+    assert "three" in ds_latest.schema().names
+
+    # Read the initial version and ensure it contains the original columns
+    ds_prev = ray.data.read_lance(path, version=initial_version)
+    assert ds_prev.count() == 6
+    assert ds_prev.schema().names == ["one", "two"]
+
+    values_prev = [[s["one"], s["two"]] for s in ds_prev.take_all()]
+    assert sorted(values_prev) == [
+        [1, "a"],
+        [2, "b"],
+        [3, "c"],
+        [4, "e"],
+        [5, "f"],
+        [6, "g"],
+    ]
+
+
+@pytest.fixture
+def mock_lance_write(monkeypatch):
+    captured = {}
+
+    class _FakeLanceDatasink:
+        def __init__(self, path, **kwargs):
+            captured["path"] = path
+            captured["kwargs"] = kwargs
+
+    def _fake_write_datasink(self, datasink, **kwargs):
+        captured["datasink"] = datasink
+        captured["write_kwargs"] = kwargs
+
+    monkeypatch.setattr(ray.data.dataset, "LanceDatasink", _FakeLanceDatasink)
+    monkeypatch.setattr(ray.data.Dataset, "write_datasink", _fake_write_datasink)
+
+    return captured, _FakeLanceDatasink
+
+
+def test_write_lance_passes_namespace_args(mock_lance_write):
+    captured, fake_lance_datasink_cls = mock_lance_write
+    table_id = ["db", "table"]
+    namespace_impl = "dir"
+    namespace_properties = {"path": "/tmp/ns"}
+
+    ds = ray.data.range(1)
+    ds.write_lance(
+        "/tmp/lance-namespace-test",
+        table_id=table_id,
+        namespace_impl=namespace_impl,
+        namespace_properties=namespace_properties,
+    )
+
+    assert captured["path"] == "/tmp/lance-namespace-test"
+    assert captured["kwargs"]["table_id"] == table_id
+    assert captured["kwargs"]["namespace_impl"] == namespace_impl
+    assert captured["kwargs"]["namespace_properties"] == namespace_properties
+    assert isinstance(captured["datasink"], fake_lance_datasink_cls)
+
+
+@pytest.mark.parametrize("mode", [SaveMode.APPEND, SaveMode.OVERWRITE])
+def test_lance_namespace_write_rejects_non_create_mode(monkeypatch, mode):
+    class _FakeNamespace:
+        pass
+
+    monkeypatch.setattr(
+        "ray.data._internal.datasource.lance_datasink.get_or_create_namespace",
+        lambda namespace_impl, namespace_properties: _FakeNamespace(),
+    )
+
+    with pytest.raises(ValueError, match="Namespace writes currently only support"):
+        LanceDatasink(
+            uri="/tmp/lance-namespace-test",
+            mode=mode,
+            table_id=["db", "table"],
+            namespace_impl="dir",
+            namespace_properties={"path": "/tmp/ns"},
+        )
+
+
+@pytest.mark.parametrize(
+    "max_attempts,expected_blocks_consumed_before_write",
+    [(1, 1), (2, 3)],
+)
+def test_write_fragment_only_materializes_stream_when_retrying(
+    monkeypatch, max_attempts, expected_blocks_consumed_before_write
+):
+    import lance.fragment
+
+    consumed = {"count": 0}
+    blocks = [pa.table({"id": [i]}) for i in range(3)]
+
+    def block_stream():
+        for block in blocks:
+            consumed["count"] += 1
+            yield block
+
+    def fake_write_fragments(reader, uri, **kwargs):
+        assert consumed["count"] == expected_blocks_consumed_before_write
+        return []
+
+    monkeypatch.setattr(lance.fragment, "write_fragments", fake_write_fragments)
+
+    _write_fragment(
+        block_stream(),
+        "/tmp/lance-materialization-test",
+        retry_params={
+            "description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
+            "match": [],
+            "max_attempts": max_attempts,
+            "max_backoff_s": 0,
+        },
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_mcap.py

@@ -0,0 +1,394 @@
+import importlib.util
+import json
+import os
+
+import pytest
+
+import ray
+from ray.data.datasource.path_util import _unwrap_protocol
+from ray.data.tests.conftest import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+# Skip all tests if mcap is not available
+MCAP_AVAILABLE = importlib.util.find_spec("mcap") is not None
+pytestmark = pytest.mark.skipif(
+    not MCAP_AVAILABLE,
+    reason="mcap module not available. Install with: pip install mcap",
+)
+
+
+def create_test_mcap_file(file_path: str, messages: list) -> None:
+    """Create a test MCAP file with given messages."""
+    from mcap.writer import Writer
+
+    with open(file_path, "wb") as stream:
+        writer = Writer(stream)
+        writer.start(profile="", library="ray-test")
+
+        # Register schema
+        schema_id = writer.register_schema(
+            name="test_schema",
+            encoding="jsonschema",
+            data=json.dumps(
+                {
+                    "type": "object",
+                    "properties": {
+                        "value": {"type": "number"},
+                        "name": {"type": "string"},
+                    },
+                }
+            ).encode(),
+        )
+
+        # Register channels and write messages
+        channels = {}
+        for msg in messages:
+            topic = msg["topic"]
+            if topic not in channels:
+                channels[topic] = writer.register_channel(
+                    schema_id=schema_id,
+                    topic=topic,
+                    message_encoding="json",
+                )
+
+            writer.add_message(
+                channel_id=channels[topic],
+                log_time=msg["log_time"],
+                publish_time=msg.get("publish_time", msg["log_time"]),
+                data=json.dumps(msg["data"]).encode(),
+            )
+
+        writer.finish()
+
+
+@pytest.fixture
+def simple_mcap_file(tmp_path):
+    """Fixture providing a simple MCAP file with one message."""
+    path = os.path.join(tmp_path, "test.mcap")
+    messages = [
+        {
+            "topic": "/test",
+            "data": {"value": 1},
+            "log_time": 1000000000,
+        }
+    ]
+    create_test_mcap_file(path, messages)
+    return path
+
+
+@pytest.fixture
+def basic_mcap_file(tmp_path):
+    """Fixture providing a basic MCAP file with two different topics."""
+    path = os.path.join(tmp_path, "test.mcap")
+    messages = [
+        {
+            "topic": "/camera/image",
+            "data": {"frame_id": 1, "timestamp": 1000},
+            "log_time": 1000000000,
+        },
+        {
+            "topic": "/lidar/points",
+            "data": {"point_count": 1024, "timestamp": 2000},
+            "log_time": 2000000000,
+        },
+    ]
+    create_test_mcap_file(path, messages)
+    return path
+
+
+@pytest.fixture
+def multi_topic_mcap_file(tmp_path):
+    """Fixture providing an MCAP file with 9 messages across 3 topics."""
+    path = os.path.join(tmp_path, "multi_topic.mcap")
+    base_time = 1000000000
+    messages = []
+    for i in range(9):
+        topics = ["/topic_a", "/topic_b", "/topic_c"]
+        topic = topics[i % 3]
+        messages.append(
+            {
+                "topic": topic,
+                "data": {"seq": i, "topic": topic},
+                "log_time": base_time + i * 1000000,
+            }
+        )
+    create_test_mcap_file(path, messages)
+    return path
+
+
+@pytest.fixture
+def time_series_mcap_file(tmp_path):
+    """Fixture providing an MCAP file with 10 time-sequenced messages."""
+    path = os.path.join(tmp_path, "time_test.mcap")
+    base_time = 1000000000
+    messages = [
+        {
+            "topic": "/test_topic",
+            "data": {"seq": i},
+            "log_time": base_time + i * 1000000,
+        }
+        for i in range(10)
+    ]
+    create_test_mcap_file(path, messages)
+    return path, base_time
+
+
+def test_read_mcap_basic(ray_start_regular_shared, basic_mcap_file):
+    """Test basic MCAP file reading."""
+    ds = ray.data.read_mcap(basic_mcap_file)
+
+    # Test metadata operations
+    assert ds.count() == 2
+    assert ds.input_files() == [_unwrap_protocol(basic_mcap_file)]
+
+    # Verify basic fields are present
+    rows = ds.take_all()
+    for row in rows:
+        assert "data" in row
+        assert "topic" in row
+        assert "log_time" in row
+        assert "publish_time" in row
+
+
+def test_read_mcap_multiple_files(ray_start_regular_shared, tmp_path):
+    """Test reading multiple MCAP files."""
+    paths = []
+    for i in range(2):
+        path = os.path.join(tmp_path, f"test_{i}.mcap")
+        messages = [
+            {
+                "topic": f"/test_{i}",
+                "data": {"file_id": i},
+                "log_time": 1000000000 + i * 1000000,
+            }
+        ]
+        create_test_mcap_file(path, messages)
+        paths.append(path)
+
+    ds = ray.data.read_mcap(paths)
+    assert ds.count() == 2
+    assert set(ds.input_files()) == {_unwrap_protocol(p) for p in paths}
+
+    rows = ds.take_all()
+    file_ids = {row["data"]["file_id"] for row in rows}
+    assert file_ids == {0, 1}
+
+
+def test_read_mcap_directory(ray_start_regular_shared, tmp_path):
+    """Test reading MCAP files from a directory."""
+    # Create MCAP files in directory
+    for i in range(2):
+        path = os.path.join(tmp_path, f"data_{i}.mcap")
+        messages = [
+            {
+                "topic": f"/dir_test_{i}",
+                "data": {"index": i},
+                "log_time": 1000000000 + i * 1000000,
+            }
+        ]
+        create_test_mcap_file(path, messages)
+
+    ds = ray.data.read_mcap(tmp_path)
+    assert ds.count() == 2
+
+
+def test_read_mcap_topic_filtering(ray_start_regular_shared, multi_topic_mcap_file):
+    """Test filtering by topics."""
+    # Test topic filtering
+    topics = {"/topic_a", "/topic_b"}
+    ds = ray.data.read_mcap(multi_topic_mcap_file, topics=topics)
+
+    rows = ds.take_all()
+    actual_topics = {row["topic"] for row in rows}
+    assert actual_topics.issubset(topics)
+    assert len(rows) == 6  # 2/3 of messages
+
+
+def test_read_mcap_time_range_filtering(
+    ray_start_regular_shared, time_series_mcap_file
+):
+    """Test filtering by time range."""
+    path, base_time = time_series_mcap_file
+
+    # Filter to first 5 messages
+    time_range = (base_time, base_time + 5000000)
+    ds = ray.data.read_mcap(path, time_range=time_range)
+
+    rows = ds.take_all()
+    assert len(rows) <= 5
+    for row in rows:
+        assert base_time <= row["log_time"] <= base_time + 5000000
+
+
+def test_read_mcap_message_type_filtering(ray_start_regular_shared, simple_mcap_file):
+    """Test filtering by message types."""
+    # Filter with existing schema
+    ds = ray.data.read_mcap(simple_mcap_file, message_types={"test_schema"})
+    assert ds.count() == 1
+
+    # Filter with non-existent schema
+    ds = ray.data.read_mcap(simple_mcap_file, message_types={"nonexistent"})
+    assert ds.count() == 0
+
+
+@pytest.mark.parametrize("include_metadata", [True, False])
+def test_read_mcap_include_metadata(
+    ray_start_regular_shared, simple_mcap_file, include_metadata
+):
+    """Test include_metadata option."""
+    ds = ray.data.read_mcap(simple_mcap_file, include_metadata=include_metadata)
+    rows = ds.take_all()
+
+    if include_metadata:
+        assert "schema_name" in rows[0]
+        assert "channel_id" in rows[0]
+    else:
+        assert "schema_name" not in rows[0]
+        assert "channel_id" not in rows[0]
+
+
+def test_read_mcap_include_paths(ray_start_regular_shared, simple_mcap_file):
+    """Test include_paths option."""
+    ds = ray.data.read_mcap(simple_mcap_file, include_paths=True)
+    rows = ds.take_all()
+
+    for row in rows:
+        assert "path" in row
+        assert simple_mcap_file in row["path"]
+
+
+def test_read_mcap_invalid_time_range(ray_start_regular_shared, simple_mcap_file):
+    """Test validation of time range parameters."""
+    # Start time >= end time
+    with pytest.raises(ValueError, match="start_time must be less than end_time"):
+        ray.data.read_mcap(simple_mcap_file, time_range=(2000, 1000))
+
+    # Negative times
+    with pytest.raises(ValueError, match="time values must be non-negative"):
+        ray.data.read_mcap(simple_mcap_file, time_range=(-1000, 2000))
+
+
+def test_read_mcap_missing_dependency(ray_start_regular_shared, simple_mcap_file):
+    """Test graceful failure when mcap library is missing."""
+    from unittest.mock import patch
+
+    with patch.dict("sys.modules", {"mcap": None}):
+        with pytest.raises(ImportError, match="MCAPDatasource.*depends on 'mcap'"):
+            ray.data.read_mcap(simple_mcap_file)
+
+
+def test_read_mcap_nonexistent_file(ray_start_regular_shared):
+    """Test handling of nonexistent files."""
+    with pytest.raises(Exception):  # FileNotFoundError or similar
+        ds = ray.data.read_mcap("/nonexistent/file.mcap")
+        ds.materialize()  # Force execution
+
+
+@pytest.mark.parametrize("override_num_blocks", [1, 2])
+def test_read_mcap_override_num_blocks(
+    ray_start_regular_shared, tmp_path, override_num_blocks
+):
+    """Test override_num_blocks parameter."""
+    path = os.path.join(tmp_path, "blocks_test.mcap")
+    messages = [
+        {
+            "topic": "/test",
+            "data": {"seq": i},
+            "log_time": 1000000000 + i * 1000000,
+        }
+        for i in range(3)
+    ]
+    create_test_mcap_file(path, messages)
+
+    ds = ray.data.read_mcap(path, override_num_blocks=override_num_blocks)
+
+    # Should still read all the data
+    assert ds.count() == 3
+    rows = ds.take_all()
+    assert len(rows) == 3
+
+
+def test_read_mcap_file_extensions(ray_start_regular_shared, tmp_path):
+    """Test file extension filtering."""
+    # Create MCAP file
+    mcap_path = os.path.join(tmp_path, "data.mcap")
+    messages = [
+        {
+            "topic": "/test",
+            "data": {"test": "mcap_data"},
+            "log_time": 1000000000,
+        }
+    ]
+    create_test_mcap_file(mcap_path, messages)
+
+    # Create non-MCAP file
+    other_path = os.path.join(tmp_path, "data.txt")
+    with open(other_path, "w") as f:
+        f.write("not mcap data")
+
+    # Should only read .mcap files by default
+    ds = ray.data.read_mcap(tmp_path)
+    assert ds.count() == 1
+    rows = ds.take_all()
+    assert rows[0]["data"]["test"] == "mcap_data"
+
+
+@pytest.mark.parametrize("ignore_missing_paths", [True, False])
+def test_read_mcap_ignore_missing_paths(
+    ray_start_regular_shared, simple_mcap_file, ignore_missing_paths
+):
+    """Test ignore_missing_paths parameter."""
+    paths = [simple_mcap_file, "/nonexistent/missing.mcap"]
+
+    if ignore_missing_paths:
+        ds = ray.data.read_mcap(paths, ignore_missing_paths=ignore_missing_paths)
+        assert ds.count() == 1
+        assert ds.input_files() == [_unwrap_protocol(simple_mcap_file)]
+    else:
+        with pytest.raises(Exception):  # FileNotFoundError or similar
+            ds = ray.data.read_mcap(paths, ignore_missing_paths=ignore_missing_paths)
+            ds.materialize()
+
+
+def test_read_mcap_json_decoding(ray_start_regular_shared, tmp_path):
+    """Test that JSON-encoded messages are properly decoded."""
+    path = os.path.join(tmp_path, "json_test.mcap")
+
+    # Test data with nested JSON structure
+    test_data = {
+        "sensor_data": {
+            "temperature": 23.5,
+            "humidity": 45.0,
+            "readings": [1, 2, 3, 4, 5],
+        },
+        "metadata": {"device_id": "sensor_001", "location": "room_a"},
+    }
+
+    messages = [
+        {
+            "topic": "/sensor/data",
+            "data": test_data,
+            "log_time": 1000000000,
+        }
+    ]
+
+    create_test_mcap_file(path, messages)
+    assert os.path.exists(path), f"Test MCAP file was not created at {path}"
+
+    ds = ray.data.read_mcap(path)
+    rows = ds.take_all()
+
+    assert len(rows) == 1, f"Expected 1 row, got {len(rows)}"
+    row = rows[0]
+
+    # Verify the data field is properly decoded as a Python dict, not bytes
+    assert isinstance(row["data"], dict), f"Expected dict, got {type(row['data'])}"
+    assert row["data"]["sensor_data"]["temperature"] == 23.5
+    assert row["data"]["metadata"]["device_id"] == "sensor_001"
+    assert row["data"]["sensor_data"]["readings"] == [1, 2, 3, 4, 5]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_mongo.py

@@ -0,0 +1,301 @@
+import shutil
+import subprocess
+import tempfile
+import time
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+# To run tests locally, make sure you install mongodb-org and have mongod
+# available on your PATH. Started directly since mongodb-org has no SysV init
+# script. See https://hub.docker.com/_/mongo
+
+
+@pytest.fixture
+def start_mongo():
+    import pymongo
+    import pymongo.errors
+
+    dbpath = tempfile.mkdtemp(prefix="mongod_test_")
+    proc = subprocess.Popen(
+        ["mongod", "--dbpath", dbpath, "--bind_ip", "127.0.0.1"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    # Wait for mongod to accept connections.
+    mongo_url = "mongodb://localhost:27017"
+    for _ in range(30):
+        if proc.poll() is not None:
+            raise RuntimeError(
+                f"mongod exited unexpectedly (returncode={proc.returncode})"
+            )
+        try:
+            client = pymongo.MongoClient(mongo_url, serverSelectionTimeoutMS=1000)
+            client.admin.command("ping")
+            break
+        except pymongo.errors.PyMongoError:
+            time.sleep(0.5)
+    else:
+        proc.kill()
+        raise RuntimeError("mongod failed to start")
+
+    # Make sure a clean slate for each test by dropping
+    # previously created ones (if any).
+    for db in client.list_database_names():
+        # Keep the MongoDB default databases.
+        if db not in ("admin", "local", "config"):
+            client.drop_database(db)
+    yield client, mongo_url
+
+    proc.terminate()
+    proc.wait(timeout=10)
+    shutil.rmtree(dbpath)
+
+
+def test_read_write_mongo(ray_start_regular_shared, start_mongo):
+    from pymongo.errors import ServerSelectionTimeoutError
+    from pymongoarrow.api import Schema
+
+    client, mongo_url = start_mongo
+    foo_db = "foo-db"
+    foo_collection = "foo-collection"
+    foo = client[foo_db][foo_collection]
+    foo.delete_many({})
+
+    # Read nonexistent URI.
+    with pytest.raises(ServerSelectionTimeoutError):
+        ds = ray.data.read_mongo(
+            uri="nonexistent-uri",
+            database=foo_db,
+            collection=foo_collection,
+        )
+    # Read nonexistent database.
+    with pytest.raises(ValueError):
+        ds = ray.data.read_mongo(
+            uri=mongo_url,
+            database="nonexistent-db",
+            collection=foo_collection,
+        )
+    # Read nonexistent collection.
+    with pytest.raises(ValueError):
+        ds = ray.data.read_mongo(
+            uri=mongo_url,
+            database=foo_db,
+            collection="nonexistent-collection",
+        )
+
+    # Inject 5 test docs.
+    docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5)]
+    df = pd.DataFrame(docs).astype({"int_field": "int32"})
+    foo.insert_many(docs)
+
+    # Read a non-empty database, with schema specified.
+    schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        schema=schema,
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [3, 2]
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["float_field", "int_field"]
+    assert ds_schema.types == [pa.float64(), pa.int32()]
+    result = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with schema inference, which will read all columns (including the auto
+    # generated internal column "_id").
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [3, 2]
+    assert ds.count() == 5
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    # We are not testing the datatype of _id here, because it varies per platform
+    assert ds.schema().types[1:] == [
+        pa.float64(),
+        pa.int32(),
+    ]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read a subset of the collection.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [2, 1]
+    assert ds.count() == 3
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())
+
+    # Read with auto-tuned parallelism.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+    )
+
+    assert ds.count() == 5
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    # We are not testing the datatype of _id here, because it varies per platform
+    assert ds.schema().types[1:] == [
+        pa.float64(),
+        pa.int32(),
+    ]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with a parallelism larger than number of rows.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=1000,
+    )
+
+    assert ds.count() == 5
+    assert ds.schema().names == ["_id", "float_field", "int_field"]
+    # We are not testing the datatype of _id here, because it varies per platform
+    assert ds.schema().types[1:] == [
+        pa.float64(),
+        pa.int32(),
+    ]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Add a column and then write back to MongoDB.
+    # Inject 2 more test docs.
+    new_docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5, 7)]
+    new_df = pd.DataFrame(new_docs).astype({"int_field": "int32"})
+    ds2 = ray.data.from_pandas(new_df)
+    ds2.write_mongo(uri=mongo_url, database=foo_db, collection=foo_collection)
+
+    # Read again to verify the content.
+    expected_ds = ds.drop_columns(["_id"]).union(ds2)
+    ds3 = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+    )
+    ds3.drop_columns(["_id"]).to_pandas().equals(expected_ds.to_pandas())
+
+    # Destination database doesn't exist.
+    with pytest.raises(ValueError):
+        ray.data.range(10).write_mongo(
+            uri=mongo_url, database="nonexistent-db", collection=foo_collection
+        )
+    # Destination collection doesn't exist.
+    with pytest.raises(ValueError):
+        ray.data.range(10).write_mongo(
+            uri=mongo_url, database=foo_db, collection="nonexistent-collection"
+        )
+
+
+def test_mongo_datasource(ray_start_regular_shared, start_mongo):
+    from pymongoarrow.api import Schema
+
+    client, mongo_url = start_mongo
+    foo_db = "foo-db"
+    foo_collection = "foo-collection"
+    foo = client[foo_db][foo_collection]
+    foo.delete_many({})
+
+    # Inject 5 test docs.
+    docs = [{"float_field": 2.0 * key, "int_field": key} for key in range(5)]
+    df = pd.DataFrame(docs).astype({"int_field": "int32"})
+    foo.insert_many(docs)
+
+    # Read non-empty datasource with a specified schema.
+    schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        schema=schema,
+        override_num_blocks=2,
+    ).materialize()
+    assert ds._block_num_rows() == [3, 2]
+    assert ds.num_blocks() == 2
+    assert ds.count() == 5
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["float_field", "int_field"]
+    assert ds_schema.types == [pa.float64(), pa.int32()]
+    result = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with schema inference, which will read all columns (including the auto
+    # generated internal column "_id").
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=2,
+    ).materialize()
+    assert ds._block_num_rows() == [3, 2]
+    assert ds.num_blocks() == 2
+    assert ds.count() == 5
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["_id", "float_field", "int_field"]
+    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with auto-tuned parallelism.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+    ).materialize()
+    assert ds.num_blocks() == 2
+    assert ds.count() == 5
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["_id", "float_field", "int_field"]
+    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read with a parallelism larger than number of rows.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        override_num_blocks=1000,
+    )
+    assert ds.schema(fetch_if_missing=False) is None
+    result = ds.drop_columns(["_id"]).to_pandas()
+    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
+
+    # Read a subset of the collection.
+    ds = ray.data.read_mongo(
+        uri=mongo_url,
+        database=foo_db,
+        collection=foo_collection,
+        pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
+        override_num_blocks=2,
+    )
+    assert ds._block_num_rows() == [2, 1]
+    ds_schema = ds.schema()
+    assert ds_schema.names == ["_id", "float_field", "int_field"]
+    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
+    df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_numpy.py

@@ -0,0 +1,123 @@
+import os
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data.context import DataContext
+from ray.data.dataset import Schema
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.util import extract_values
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.mark.parametrize("from_ref", [False, True])
+def test_from_numpy(ray_start_regular_shared, from_ref):
+    arr1 = np.expand_dims(np.arange(0, 4), axis=1)
+    arr2 = np.expand_dims(np.arange(4, 8), axis=1)
+    arrs = [arr1, arr2]
+    if from_ref:
+        ds = ray.data.from_numpy_refs([ray.put(arr) for arr in arrs])
+    else:
+        ds = ray.data.from_numpy(arrs)
+    values = np.stack(extract_values("data", ds.take(8)))
+    np.testing.assert_array_equal(values, np.concatenate((arr1, arr2)))
+    # Check that conversion task is included in stats.
+    assert "FromNumpy" in ds.stats()
+
+    # Test from single NumPy ndarray.
+    if from_ref:
+        ds = ray.data.from_numpy_refs(ray.put(arr1))
+    else:
+        ds = ray.data.from_numpy(arr1)
+    values = np.stack(extract_values("data", ds.take(4)))
+    np.testing.assert_array_equal(values, arr1)
+    # Check that conversion task is included in stats.
+    assert "FromNumpy" in ds.stats()
+
+
+def test_from_numpy_variable_shaped(ray_start_regular_shared):
+    arr = np.array([np.ones((2, 2)), np.ones((3, 3))], dtype=object)
+    ds = ray.data.from_numpy(arr)
+    values = np.array(extract_values("data", ds.take(2)), dtype=object)
+
+    def recursive_to_list(a):
+        if not isinstance(a, (list, np.ndarray)):
+            return a
+        return [recursive_to_list(e) for e in a]
+
+    # Convert to a nested Python list in order to circumvent failed comparisons on
+    # ndarray raggedness.
+    np.testing.assert_equal(recursive_to_list(values), recursive_to_list(arr))
+
+
+def test_to_numpy_refs(ray_start_regular_shared):
+    # Tensor Dataset
+    ds = ray.data.range_tensor(10, override_num_blocks=2)
+    arr = np.concatenate(extract_values("data", ray.get(ds.to_numpy_refs())))
+    np.testing.assert_equal(arr, np.expand_dims(np.arange(0, 10), 1))
+
+    # Table Dataset
+    ds = ray.data.range(10)
+    arr = np.concatenate([t["id"] for t in ray.get(ds.to_numpy_refs())])
+    np.testing.assert_equal(arr, np.arange(0, 10))
+
+    # Test multi-column Arrow dataset.
+    ds = ray.data.from_arrow(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    arrs = ray.get(ds.to_numpy_refs())
+    np.testing.assert_equal(
+        arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}]
+    )
+
+    # Test multi-column Pandas dataset.
+    ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    arrs = ray.get(ds.to_numpy_refs())
+    np.testing.assert_equal(
+        arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}]
+    )
+
+
+def test_numpy_roundtrip(ray_start_regular_shared, tmp_path):
+    tensor_type = DataContext.get_current().arrow_fixed_shape_tensor_format.to_type()
+
+    ds = ray.data.range_tensor(10, override_num_blocks=2)
+    ds.write_numpy(tmp_path, column="data")
+    ds = ray.data.read_numpy(tmp_path)
+    assert ds.count() == 10
+    assert ds.schema() == Schema(pa.schema([("data", tensor_type((1,), pa.int64()))]))
+    assert sorted(ds.take_all(), key=lambda row: row["data"]) == [
+        {"data": np.array([i])} for i in range(10)
+    ]
+
+
+def test_numpy_read_x(ray_start_regular_shared, tmp_path):
+    tensor_type = DataContext.get_current().arrow_fixed_shape_tensor_format.to_type()
+
+    path = os.path.join(tmp_path, "test_np_dir")
+    os.mkdir(path)
+    np.save(os.path.join(path, "test.npy"), np.expand_dims(np.arange(0, 10), 1))
+    ds = ray.data.read_numpy(path, override_num_blocks=1)
+    assert ds.count() == 10
+    assert ds.schema() == Schema(pa.schema([("data", tensor_type((1,), pa.int64()))]))
+    np.testing.assert_equal(
+        extract_values("data", ds.take(2)), [np.array([0]), np.array([1])]
+    )
+
+
+def test_numpy_write(ray_start_regular_shared, tmp_path):
+    ds = ray.data.range_tensor(1)
+
+    ds.write_numpy(tmp_path, column="data")
+
+    actual_array = np.concatenate(
+        [np.load(os.path.join(tmp_path, filename)) for filename in os.listdir(tmp_path)]
+    )
+    assert actual_array == np.array((0,))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_pandas.py

@@ -0,0 +1,223 @@
+from typing import Iterator
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data._internal.tensor_extensions.arrow import (
+    ArrowTensorArray,
+    get_arrow_extension_fixed_shape_tensor_types,
+)
+from ray.data.block import Block
+from ray.data.extensions import TensorDtype
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+from ray.types import ObjectRef
+
+
+def _get_first_block(bundles: Iterator[RefBundle]) -> ObjectRef[Block]:
+    return next(bundles).block_refs[0]
+
+
+@pytest.mark.parametrize("enable_pandas_block", [False, True])
+def test_from_pandas(ray_start_regular_shared, enable_pandas_block):
+    ctx = ray.data.context.DataContext.get_current()
+    old_enable_pandas_block = ctx.enable_pandas_block
+    ctx.enable_pandas_block = enable_pandas_block
+    try:
+        df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+        df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+        ds = ray.data.from_pandas([df1, df2])
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(6)]
+        rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+
+        # test from single pandas dataframe
+        ds = ray.data.from_pandas(df1)
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(3)]
+        rows = [(r.one, r.two) for _, r in df1.iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+    finally:
+        ctx.enable_pandas_block = old_enable_pandas_block
+
+
+@pytest.mark.parametrize("num_inputs", [1, 2])
+def test_from_pandas_override_num_blocks(num_inputs, ray_start_regular_shared):
+    df = pd.DataFrame({"number": [0]})
+
+    ds = ray.data.from_pandas([df] * num_inputs, override_num_blocks=2)
+
+    assert ds.materialize().num_blocks() == 2
+
+
+@pytest.mark.parametrize("enable_pandas_block", [False, True])
+def test_from_pandas_refs(ray_start_regular_shared, enable_pandas_block):
+    ctx = ray.data.context.DataContext.get_current()
+    old_enable_pandas_block = ctx.enable_pandas_block
+    ctx.enable_pandas_block = enable_pandas_block
+    try:
+        df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+        df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+        ds = ray.data.from_pandas_refs([ray.put(df1), ray.put(df2)])
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(6)]
+        rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+
+        # test from single pandas dataframe ref
+        ds = ray.data.from_pandas_refs(ray.put(df1))
+        block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
+        assert (
+            isinstance(block, pd.DataFrame)
+            if enable_pandas_block
+            else isinstance(block, pa.Table)
+        )
+        values = [(r["one"], r["two"]) for r in ds.take(3)]
+        rows = [(r.one, r.two) for _, r in df1.iterrows()]
+        assert values == rows
+        # Check that metadata fetch is included in stats.
+        assert "FromPandas" in ds.stats()
+    finally:
+        ctx.enable_pandas_block = old_enable_pandas_block
+
+
+def test_to_pandas(ray_start_regular_shared):
+    n = 5
+    df = pd.DataFrame({"id": list(range(n))})
+    ds = ray.data.range(n)
+    dfds = ds.to_pandas()
+    pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
+
+    # Test limit.
+    with pytest.raises(ValueError):
+        dfds = ds.to_pandas(limit=3)
+
+    # Test limit greater than number of rows.
+    dfds = ds.to_pandas(limit=6)
+    pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
+
+
+def test_to_pandas_different_block_types(ray_start_regular_shared):
+    # Test for https://github.com/ray-project/ray/issues/48575.
+    df = pd.DataFrame({"a": [0]})
+    ds1 = ray.data.from_pandas(df)
+
+    table = pa.Table.from_pandas(df)
+    ds2 = ray.data.from_arrow(table)
+
+    actual_df = ds1.union(ds2).to_pandas()
+
+    expected_df = pd.DataFrame({"a": [0, 0]}).astype(actual_df.dtypes.to_dict())
+    pd.testing.assert_frame_equal(actual_df, expected_df)
+
+
+def test_to_pandas_refs(ray_start_regular_shared):
+    n = 5
+    df = pd.DataFrame({"id": list(range(n))})
+    ds = ray.data.range(n)
+    dfds = pd.concat(ray.get(ds.to_pandas_refs()), ignore_index=True)
+    pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
+
+
+def test_pandas_roundtrip(ray_start_regular_shared, tmp_path):
+    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
+    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
+    ds = ray.data.from_pandas([df1, df2], override_num_blocks=2)
+    dfds = ds.to_pandas()
+    expected = pd.concat([df1, df2], ignore_index=True)
+    pd.testing.assert_frame_equal(expected.astype(dfds.dtypes.to_dict()), dfds)
+
+
+def test_to_pandas_tensor_column_cast_pandas(ray_start_regular_shared):
+    # Check that tensor column casting occurs when converting a Dataset to a Pandas
+    # DataFrame.
+    data = np.arange(12).reshape((3, 2, 2))
+    ctx = ray.data.context.DataContext.get_current()
+    original = ctx.enable_tensor_extension_casting
+    try:
+        ctx.enable_tensor_extension_casting = True
+        in_df = pd.DataFrame({"a": [data]})
+        ds = ray.data.from_pandas(in_df)
+        dtypes = ds.schema().base_schema.types
+        assert len(dtypes) == 1
+        # Tensor column should be automatically cast to Tensor extension.
+        assert isinstance(dtypes[0], TensorDtype)
+        # Original df should not be changed.
+        assert not isinstance(in_df.dtypes[0], TensorDtype)
+        out_df = ds.to_pandas()
+        # Column should be cast back to object dtype when returning back to user.
+        assert out_df["a"].dtype.type is np.object_
+        expected_df = pd.DataFrame({"a": [data]})
+        pd.testing.assert_frame_equal(out_df, expected_df)
+    finally:
+        ctx.enable_tensor_extension_casting = original
+
+
+def test_to_pandas_tensor_column_cast_arrow(ray_start_regular_shared):
+    # Check that tensor column casting occurs when converting a Dataset to a Pandas
+    # DataFrame.
+    data = np.arange(12).reshape((3, 2, 2))
+    ctx = ray.data.context.DataContext.get_current()
+    original = ctx.enable_tensor_extension_casting
+    try:
+        ctx.enable_tensor_extension_casting = True
+        in_table = pa.table({"a": ArrowTensorArray.from_numpy(data)})
+        ds = ray.data.from_arrow(in_table)
+        dtype = ds.schema().base_schema.field(0).type
+        assert isinstance(dtype, get_arrow_extension_fixed_shape_tensor_types())
+        out_df = ds.to_pandas()
+        assert out_df["a"].dtype.type is np.object_
+        expected_df = pd.DataFrame({"a": list(data)})
+        pd.testing.assert_frame_equal(out_df, expected_df)
+    finally:
+        ctx.enable_tensor_extension_casting = original
+
+
+def test_read_pandas_data_array_column(ray_start_regular_shared):
+    df = pd.DataFrame(
+        {
+            "one": [1, 2, 3],
+            "array": [
+                np.array([1, 1, 1]),
+                np.array([2, 2, 2]),
+                np.array([3, 3, 3]),
+            ],
+        }
+    )
+    ds = ray.data.from_pandas(df)
+    row = ds.take(1)[0]
+    assert row["one"] == 1
+    assert all(row["array"] == [1, 1, 1])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_raydp.py

@@ -0,0 +1,71 @@
+import pandas
+import pytest
+import raydp
+
+import ray
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.test_util import _check_usage_record
+
+
+# RayDP tests require Ray Java. Make sure ray jar is built before running this test.
+@pytest.fixture(scope="function")
+def spark(request):
+    ray.init(num_cpus=2, include_dashboard=False)
+    spark_session = raydp.init_spark("test", 1, 1, "500M")
+
+    def stop_all():
+        raydp.stop_spark()
+        ray.shutdown()
+
+    request.addfinalizer(stop_all)
+    return spark_session
+
+
+def test_raydp_roundtrip(spark):
+    spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])
+    rows = [(r.one, r.two) for r in spark_df.take(3)]
+    ds = ray.data.from_spark(spark_df)
+    values = [(r["one"], r["two"]) for r in ds.take(6)]
+    assert values == rows
+    df = ds.to_spark(spark)
+    rows_2 = [(r.one, r.two) for r in df.take(3)]
+    assert values == rows_2
+
+
+def test_raydp_to_spark(spark):
+    n = 5
+    ds = ray.data.range(n)
+    values = [r["id"] for r in ds.take(5)]
+    df = ds.to_spark(spark)
+    rows = [r.id for r in df.take(5)]
+    assert values == rows
+
+
+def test_from_spark_e2e(spark):
+    spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])
+
+    rows = [(r.one, r.two) for r in spark_df.take(3)]
+    ds = ray.data.from_spark(spark_df)
+    assert len(ds.take_all()) == len(rows)
+    values = [(r["one"], r["two"]) for r in ds.take(6)]
+    assert values == rows
+
+    # Check that metadata fetch is included in stats.
+    assert "FromArrow" in ds.stats()
+    # Underlying implementation uses `FromArrow` operator
+    assert ds._logical_plan.dag.name == "FromArrow"
+    _check_usage_record(["FromArrow"])
+
+
+def test_to_pandas(spark):
+    df = spark.range(100)
+    ds = ray.data.from_spark(df)
+    pdf = ds.to_pandas()
+    pdf2 = df.toPandas().astype(pdf.dtypes.to_dict())
+    pandas.testing.assert_frame_equal(pdf, pdf2)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_read_parquet_v2.py

@@ -0,0 +1,214 @@
+"""Integration-ish tests for ``read_parquet()`` on the DataSourceV2 path.
+
+These tests exercise planning-time behavior: schema inference,
+``ListFiles → ReadFiles`` attachment to the logical plan, and
+unsupported-option gating. They call ``ray.data.read_parquet`` which
+triggers Ray auto-init, so they live alongside the other datasource
+integration tests rather than under ``tests/unit/``.
+"""
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+
+import ray
+from ray.data._internal.datasource_v2.partitioners.round_robin_partitioner import (
+    RoundRobinPartitioner,
+)
+from ray.data._internal.datasource_v2.scanners.parquet_scanner import ParquetScanner
+from ray.data._internal.logical.operators import ListFiles, ReadFiles
+from ray.data.context import DataContext
+
+
+def _write(path, table):
+    pq.write_table(table, str(path))
+
+
+@pytest.fixture
+def restore_ctx():
+    ctx = DataContext.get_current()
+    original = ctx.use_datasource_v2
+    try:
+        yield ctx
+    finally:
+        ctx.use_datasource_v2 = original
+
+
+def test_v2_flag_default():
+    # The default is driven by ``DEFAULT_USE_DATASOURCE_V2``. Asserting
+    # either direction here would be brittle, so just check that the
+    # default is a bool.
+    ctx = DataContext()
+    assert isinstance(ctx.use_datasource_v2, bool)
+
+
+def test_read_parquet_builds_list_files_read_files_chain(tmp_path, restore_ctx):
+    f = tmp_path / "data.parquet"
+    _write(f, pa.table({"a": [1, 2, 3], "b": ["x", "y", "z"]}))
+
+    restore_ctx.use_datasource_v2 = True
+    ds = ray.data.read_parquet(str(tmp_path))
+
+    assert isinstance(ds._logical_plan.dag, ReadFiles)
+    assert isinstance(ds._logical_plan.dag.input_dependencies[0], ListFiles)
+    schema = ds.schema()
+    assert schema is not None
+    assert "a" in schema.names
+    assert "b" in schema.names
+
+
+def test_read_parquet_v2_hive_partitioned(tmp_path, restore_ctx):
+    for p in ["a", "b"]:
+        d = tmp_path / f"color={p}"
+        d.mkdir()
+        _write(d / "data.parquet", pa.table({"x": [1, 2]}))
+
+    restore_ctx.use_datasource_v2 = True
+    ds = ray.data.read_parquet(str(tmp_path))
+    schema = ds.schema()
+    assert "x" in schema.names
+    assert "color" in schema.names
+
+
+def test_read_parquet_v2_include_paths(tmp_path, restore_ctx):
+    _write(tmp_path / "data.parquet", pa.table({"a": [1]}))
+
+    restore_ctx.use_datasource_v2 = True
+    ds = ray.data.read_parquet(str(tmp_path), include_paths=True)
+    schema = ds.schema()
+    assert "path" in schema.names
+
+
+def test_read_parquet_v2_include_row_hash(tmp_path, restore_ctx):
+    _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
+
+    restore_ctx.use_datasource_v2 = True
+    ds = ray.data.read_parquet(str(tmp_path), include_row_hash=True)
+    schema = ds.schema()
+    assert schema is not None
+    assert "row_hash" in schema.names
+    assert schema.types[schema.names.index("row_hash")] == pa.uint64()
+
+
+def test_read_parquet_v2_columns_applies_select_columns(tmp_path, restore_ctx):
+    from ray.data._internal.logical.operators.map_operator import Project
+
+    _write(tmp_path / "data.parquet", pa.table({"a": [1], "b": [2]}))
+
+    restore_ctx.use_datasource_v2 = True
+    with pytest.warns(DeprecationWarning, match="`columns=` on `read_parquet`"):
+        ds = ray.data.read_parquet(str(tmp_path), columns=["a"])
+
+    # ``columns=`` is applied via ``ds.select_columns([...])``, which
+    # wraps the ReadFiles op in a Project node.
+    dag = ds._logical_plan.dag
+    assert isinstance(dag, Project)
+    assert [expr.name for expr in dag.exprs] == ["a"]
+    assert isinstance(dag.input_dependencies[0], ReadFiles)
+
+
+def test_read_parquet_v2_columns_with_include_paths_preserves_path(
+    tmp_path, restore_ctx
+):
+    from ray.data._internal.logical.operators.map_operator import Project
+
+    _write(tmp_path / "data.parquet", pa.table({"a": [1], "b": [2]}))
+
+    restore_ctx.use_datasource_v2 = True
+    with pytest.warns(DeprecationWarning, match="`columns=` on `read_parquet`"):
+        ds = ray.data.read_parquet(str(tmp_path), columns=["a"], include_paths=True)
+
+    dag = ds._logical_plan.dag
+    assert isinstance(dag, Project)
+    # V1 ``columns=[...]`` retained ``"path"`` implicitly when
+    # ``include_paths=True``; the V2 path appends it to keep that
+    # behavior.
+    assert [expr.name for expr in dag.exprs] == ["a", "path"]
+
+
+def test_read_parquet_v2_override_num_blocks_drives_partitioner(tmp_path, restore_ctx):
+    _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
+
+    restore_ctx.use_datasource_v2 = True
+    original = restore_ctx.read_op_min_num_blocks
+    ds = ray.data.read_parquet(str(tmp_path), override_num_blocks=7)
+
+    # The override should drive the ListFiles partitioner's bucket count
+    # for this read only — the global DataContext must not be mutated.
+    list_files_op = ds._logical_plan.dag.input_dependencies[0]
+    assert isinstance(list_files_op, ListFiles)
+    assert isinstance(list_files_op.file_partitioner, RoundRobinPartitioner)
+    assert list_files_op.file_partitioner.num_buckets == 7
+    assert restore_ctx.read_op_min_num_blocks == original
+
+
+def test_read_parquet_v2_filter_raises(tmp_path, restore_ctx):
+    import pyarrow.dataset as pds
+
+    _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
+
+    restore_ctx.use_datasource_v2 = True
+    with pytest.raises(ValueError, match="`filter=` on `read_parquet`"):
+        ray.data.read_parquet(str(tmp_path), filter=pds.field("a") > 1)
+
+
+def test_read_parquet_v2_dataset_kwargs_rejects_partitioning(tmp_path, restore_ctx):
+    _write(tmp_path / "data.parquet", pa.table({"a": [1]}))
+
+    restore_ctx.use_datasource_v2 = True
+    with pytest.warns(DeprecationWarning, match="`dataset_kwargs`"):
+        with pytest.raises(
+            ValueError, match="'partitioning' parameter isn't supported"
+        ):
+            ray.data.read_parquet(
+                str(tmp_path), dataset_kwargs={"partitioning": "hive"}
+            )
+
+
+def test_read_parquet_v2_dataset_kwargs_rejects_filters(tmp_path, restore_ctx):
+    _write(tmp_path / "data.parquet", pa.table({"a": [1]}))
+
+    restore_ctx.use_datasource_v2 = True
+    with pytest.warns(DeprecationWarning, match="`dataset_kwargs`"):
+        with pytest.raises(ValueError, match="Row filtering via 'filters'"):
+            ray.data.read_parquet(
+                str(tmp_path), dataset_kwargs={"filters": [("a", ">", 0)]}
+            )
+
+
+def test_read_parquet_v2_dataset_kwargs_threads_through_to_scanner(
+    tmp_path, restore_ctx
+):
+    _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
+
+    restore_ctx.use_datasource_v2 = True
+    with pytest.warns(DeprecationWarning, match="`dataset_kwargs`"):
+        ds = ray.data.read_parquet(
+            str(tmp_path),
+            dataset_kwargs={
+                "coerce_int96_timestamp_unit": "ms",
+                "read_dictionary": ["a"],
+            },
+        )
+
+    # ``read_dictionary`` is renamed to ``dictionary_columns`` to match
+    # ``pds.ParquetFileFormat``; ``coerce_int96_timestamp_unit`` passes
+    # through unchanged.
+    read_files_op = ds._logical_plan.dag
+    assert isinstance(read_files_op, ReadFiles)
+    assert isinstance(read_files_op.scanner, ParquetScanner)
+    assert read_files_op.scanner.parquet_format_kwargs == {
+        "coerce_int96_timestamp_unit": "ms",
+        "dictionary_columns": ["a"],
+    }
+
+
+def test_read_parquet_v2_empty_dir_raises(tmp_path, restore_ctx):
+    restore_ctx.use_datasource_v2 = True
+    with pytest.raises(ValueError, match="no files found"):
+        ray.data.read_parquet(str(tmp_path))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-xvs"]))

python/ray/data/tests/datasource/test_snowflake.py

@@ -0,0 +1,120 @@
+import base64
+import os
+import random
+import string
+from typing import Any, Dict, List, Tuple
+
+import pytest
+from snowflake.connector import connect
+
+import ray
+from ray.tests.conftest import *  # noqa
+
+# Note: Snowflake secrets are only used in postmerge authenticated tests.
+
+
+@pytest.fixture
+def connection_parameters():
+    private_key_b64 = os.getenv("SNOWFLAKE_PRIVATE_KEY")
+    private_key_bytes = base64.b64decode(private_key_b64)
+    parameters = {
+        "user": os.getenv("SNOWFLAKE_USER"),
+        "account": os.getenv("SNOWFLAKE_ACCOUNT"),
+        "database": os.getenv("SNOWFLAKE_DATABASE"),
+        "schema": os.getenv("SNOWFLAKE_SCHEMA"),
+        "warehouse": os.getenv("SNOWFLAKE_WAREHOUSE"),
+        "private_key": private_key_bytes,
+    }
+
+    yield parameters
+
+
+@pytest.fixture
+def temp_table(connection_parameters):
+    table_name = "".join([random.choice(string.ascii_uppercase) for _ in range(8)])
+
+    yield table_name
+
+    with connect(**connection_parameters) as connection, connection.cursor() as cursor:
+        cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
+        connection.commit()
+
+
+@pytest.mark.needs_credentials
+def test_read(ray_start_regular_shared, connection_parameters):
+    # This query fetches a small dataset with a variety of column types.
+    query = "SELECT * FROM SNOWFLAKE_SAMPLE_DATA.TPCDS_SF100TCL.CALL_CENTER"
+
+    # Read the data and check contents.
+    dataset = ray.data.read_snowflake(query, connection_parameters)
+    actual_column_names = dataset.schema().names
+    actual_rows = [tuple(row.values()) for row in dataset.take_all()]
+    expected_column_names, expected_rows = execute(query, connection_parameters)
+
+    assert actual_column_names == expected_column_names
+    assert sorted(actual_rows) == sorted(expected_rows)
+
+
+@pytest.mark.needs_credentials
+def test_write(ray_start_regular_shared, temp_table, connection_parameters):
+    expected_column_names = ["title", "year", "score"]
+    expected_rows = [
+        ("Monty Python and the Holy Grail", 1975, 8.2),
+        ("And Now for Something Completely Different", 1971, 7.5),
+    ]
+
+    # Create the table first
+    create_table_sql = f"""
+    CREATE TABLE IF NOT EXISTS {temp_table} (
+        "title" VARCHAR(255),
+        "year" INTEGER,
+        "score" FLOAT
+    )
+    """
+    execute(create_table_sql, connection_parameters)
+
+    items = [dict(zip(expected_column_names, row)) for row in expected_rows]
+    dataset = ray.data.from_items(items)
+
+    dataset.write_snowflake(temp_table, connection_parameters)
+    actual_column_names, actual_rows = execute(
+        f"SELECT * FROM {temp_table}", connection_parameters
+    )
+
+    assert actual_column_names == expected_column_names
+    assert sorted(actual_rows) == sorted(expected_rows)
+
+
+@pytest.mark.needs_credentials
+def execute(
+    query: str, connection_parameters: Dict[str, str]
+) -> Tuple[List[str], List[Tuple[Any]]]:
+    """Execute a query on Snowflake and return the resulting data.
+
+    Args:
+        query: The SQL query to execute.
+        connection_parameters: Connection params for snowflake.
+
+    Returns:
+        A two-tuple containing the column names and rows.
+    """
+    with connect(**connection_parameters) as connection, connection.cursor() as cursor:
+        cursor.execute(query)
+        column_names = [column_metadata.name for column_metadata in cursor.description]
+        rows = cursor.fetchall()
+
+    # TODO(mowen): Figure out how to actually handle the Decimal objects, we don't
+    # want a divergenece in behavior here.
+    # The Snowflake Python Connector represents numbers as `Decimal` objects.
+    # rows = [
+    #     tuple(float(value) if isinstance(value, Decimal) else value for value in row)
+    #     for row in rows
+    # ]
+
+    return column_names, rows
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_sql.py

@@ -0,0 +1,176 @@
+import sqlite3
+import tempfile
+from typing import Generator
+
+import pytest
+
+import ray
+from ray.tests.conftest import *  # noqa  # noqa
+
+
+@pytest.fixture(name="temp_database")
+def temp_database_fixture() -> Generator[str, None, None]:
+    with tempfile.NamedTemporaryFile(suffix=".db") as file:
+        yield file.name
+
+
+def test_read_sql(temp_database: str):
+    connection = sqlite3.connect(temp_database)
+    connection.execute("CREATE TABLE movie(title, year, score)")
+    expected_values = [
+        ("Monty Python and the Holy Grail", 1975, 8.2),
+        ("And Now for Something Completely Different", 1971, 7.5),
+    ]
+    connection.executemany("INSERT INTO movie VALUES (?, ?, ?)", expected_values)
+    connection.commit()
+    connection.close()
+
+    dataset = ray.data.read_sql(
+        "SELECT * FROM movie",
+        lambda: sqlite3.connect(temp_database),
+    )
+    actual_values = [tuple(record.values()) for record in dataset.take_all()]
+
+    assert sorted(actual_values) == sorted(expected_values)
+
+
+@pytest.mark.parametrize(
+    "sql, sql_params",
+    [
+        ("SELECT * FROM movie WHERE year >= ?", (1975,)),
+        ("SELECT * FROM movie WHERE year >= ?", [1975]),
+        ("SELECT * FROM movie WHERE year >= :year", {"year": 1975}),
+    ],
+)
+def test_read_sql_with_params(temp_database: str, sql: str, sql_params):
+    connection = sqlite3.connect(temp_database)
+    connection.execute("CREATE TABLE movie(title, year, score)")
+    expected_values = [
+        ("Monty Python and the Holy Grail", 1975, 8.2),
+        ("And Now for Something Completely Different", 1971, 7.5),
+        ("Monty Python's Life of Brian", 1979, 8.0),
+    ]
+    connection.executemany("INSERT INTO movie VALUES (?, ?, ?)", expected_values)
+    connection.commit()
+    connection.close()
+
+    dataset = ray.data.read_sql(
+        sql,
+        lambda: sqlite3.connect(temp_database),
+        sql_params=sql_params,
+    )
+    actual_values = [tuple(record.values()) for record in dataset.take_all()]
+
+    assert sorted(actual_values) == sorted(
+        [row for row in expected_values if row[1] >= 1975]
+    )
+
+
+def test_read_sql_with_parallelism_fallback(temp_database: str):
+    connection = sqlite3.connect(temp_database)
+    connection.execute("CREATE TABLE grade(name, id, score)")
+    base_tuple = ("xiaoming", 1, 8.2)
+    # Generate 200 elements
+    expected_values = [
+        (f"{base_tuple[0]}{i}", i, base_tuple[2] + i + 1) for i in range(500)
+    ]
+    connection.executemany("INSERT INTO grade VALUES (?, ?, ?)", expected_values)
+    connection.commit()
+    connection.close()
+
+    num_blocks = 2
+    dataset = ray.data.read_sql(
+        "SELECT * FROM grade",
+        lambda: sqlite3.connect(temp_database),
+        override_num_blocks=num_blocks,
+        shard_hash_fn="unicode",
+        shard_keys=["id"],
+    )
+    dataset = dataset.materialize()
+    assert dataset.num_blocks() == num_blocks
+
+    actual_values = [tuple(record.values()) for record in dataset.take_all()]
+    assert sorted(actual_values) == sorted(expected_values)
+
+
+# for mysql test
+@pytest.mark.skip(reason="skip this test because mysql env is not ready")
+def test_read_sql_with_parallelism_mysql(temp_database: str):
+    # connect mysql
+    import pymysql
+
+    connection = pymysql.connect(
+        host="10.10.xx.xx", user="root", password="22222", database="test"
+    )
+    cursor = connection.cursor()
+
+    cursor.execute(
+        "CREATE TABLE IF NOT EXISTS grade (name VARCHAR(255), id INT, score FLOAT)"
+    )
+
+    base_tuple = ("xiaoming", 1, 8.2)
+    expected_values = [
+        (f"{base_tuple[0]}{i}", i, base_tuple[2] + i + 1) for i in range(200)
+    ]
+
+    cursor.executemany(
+        "INSERT INTO grade (name, id, score) VALUES (%s, %s, %s)", expected_values
+    )
+    connection.commit()
+
+    cursor.close()
+    connection.close()
+
+    dataset = ray.data.read_sql(
+        "SELECT * FROM grade",
+        lambda: pymysql.connect(host="xxxxx", user="xx", password="xx", database="xx"),
+        parallelism=4,
+        shard_keys=["id"],
+    )
+    actual_values = [tuple(record.values()) for record in dataset.take_all()]
+
+    assert sorted(actual_values) == sorted(expected_values)
+    assert dataset.materialize().num_blocks() == 4
+
+
+def test_write_sql(temp_database: str):
+    connection = sqlite3.connect(temp_database)
+    connection.cursor().execute("CREATE TABLE test(string, number)")
+    dataset = ray.data.from_items(
+        [{"string": "spam", "number": 0}, {"string": "ham", "number": 1}]
+    )
+
+    dataset.write_sql(
+        "INSERT INTO test VALUES(?, ?)", lambda: sqlite3.connect(temp_database)
+    )
+
+    result = connection.cursor().execute("SELECT * FROM test ORDER BY number")
+    assert result.fetchall() == [("spam", 0), ("ham", 1)]
+
+
+@pytest.mark.parametrize("num_blocks", (1, 20))
+def test_write_sql_many_rows(num_blocks: int, temp_database: str):
+    connection = sqlite3.connect(temp_database)
+    connection.cursor().execute("CREATE TABLE test(id)")
+    dataset = ray.data.range(1000).repartition(num_blocks)
+
+    dataset.write_sql(
+        "INSERT INTO test VALUES(?)", lambda: sqlite3.connect(temp_database)
+    )
+
+    result = connection.cursor().execute("SELECT * FROM test ORDER BY id")
+    assert result.fetchall() == [(i,) for i in range(1000)]
+
+
+def test_write_sql_nonexistant_table(temp_database: str):
+    dataset = ray.data.range(1)
+    with pytest.raises(sqlite3.OperationalError):
+        dataset.write_sql(
+            "INSERT INTO test VALUES(?)", lambda: sqlite3.connect(temp_database)
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_tensorflow_datasets.py

@@ -0,0 +1,38 @@
+import sys
+
+import pytest
+
+import ray
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.test_util import _check_usage_record
+from ray.data.tests.util import extract_values
+from ray.tests.conftest import *  # noqa
+
+
+def test_from_tf_e2e(ray_start_regular_shared_2_cpus):
+    import tensorflow as tf
+    import tensorflow_datasets as tfds
+
+    tf_dataset = tfds.load("mnist", split=["train"], as_supervised=True)[0]
+    tf_dataset = tf_dataset.take(8)  # Use subset to make test run faster.
+
+    ray_dataset = ray.data.from_tf(tf_dataset)
+
+    actual_data = extract_values("item", ray_dataset.take_all())
+    expected_data = list(tf_dataset)
+    assert len(actual_data) == len(expected_data)
+    for (expected_features, expected_label), (actual_features, actual_label) in zip(
+        expected_data, actual_data
+    ):
+        tf.debugging.assert_equal(expected_features, actual_features)
+        tf.debugging.assert_equal(expected_label, actual_label)
+
+    # Check that metadata fetch is included in stats.
+    assert "FromItems" in ray_dataset.stats()
+    # Underlying implementation uses `FromItems` operator
+    assert ray_dataset._logical_plan.dag.name == "FromItems"
+    _check_usage_record(["FromItems"])
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_text.py

@@ -0,0 +1,100 @@
+import os
+
+import pytest
+from fsspec.implementations.http import HTTPFileSystem
+
+import ray
+from ray.data._internal.execution.interfaces.ref_bundle import (
+    _ref_bundles_iterator_to_block_refs_list,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.mock_http_server import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+
+def _to_lines(rows):
+    return [row["text"] for row in rows]
+
+
+def test_empty_text_files(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "test_text")
+    os.mkdir(path)
+    # 2 empty files.
+    _ = open(os.path.join(path, "file1.txt"), "w")
+    _ = open(os.path.join(path, "file2.txt"), "w")
+    ds = ray.data.read_text(path)
+    assert ds.count() == 0
+    ds = ray.data.read_text(path, drop_empty_lines=False)
+    assert ds.count() == 0
+
+
+def test_read_text(ray_start_regular_shared, tmp_path):
+    path = os.path.join(tmp_path, "test_text")
+    os.mkdir(path)
+    with open(os.path.join(path, "file1.txt"), "w") as f:
+        f.write("hello\n")
+        f.write("world")
+    with open(os.path.join(path, "file2.txt"), "w") as f:
+        f.write("goodbye")
+    with open(os.path.join(path, "file3.txt"), "w") as f:
+        f.write("ray\n")
+    ds = ray.data.read_text(path)
+    assert sorted(_to_lines(ds.take())) == ["goodbye", "hello", "ray", "world"]
+    ds = ray.data.read_text(path, drop_empty_lines=False)
+    assert ds.count() == 4
+
+
+def test_read_text_remote_args(ray_start_cluster, tmp_path):
+    cluster = ray_start_cluster
+    cluster.add_node(
+        resources={"foo": 100},
+        num_cpus=1,
+        _system_config={"max_direct_call_object_size": 0},
+    )
+    cluster.add_node(resources={"bar": 100}, num_cpus=1)
+
+    ray.shutdown()
+    ray.init(cluster.address)
+
+    @ray.remote
+    def get_node_id():
+        return ray.get_runtime_context().get_node_id()
+
+    bar_node_id = ray.get(get_node_id.options(resources={"bar": 1}).remote())
+
+    path = os.path.join(tmp_path, "test_text")
+    os.mkdir(path)
+    with open(os.path.join(path, "file1.txt"), "w") as f:
+        f.write("hello\n")
+        f.write("world")
+    with open(os.path.join(path, "file2.txt"), "w") as f:
+        f.write("goodbye")
+
+    ds = ray.data.read_text(
+        path, override_num_blocks=2, ray_remote_args={"resources": {"bar": 1}}
+    )
+
+    block_refs = _ref_bundles_iterator_to_block_refs_list(
+        ds.iter_internal_ref_bundles()
+    )
+    ray.wait(block_refs, num_returns=len(block_refs), fetch_local=False)
+    location_data = ray.experimental.get_object_locations(block_refs)
+    locations = []
+    for block in block_refs:
+        locations.extend(location_data[block]["node_ids"])
+    assert set(locations) == {bar_node_id}, locations
+    assert sorted(_to_lines(ds.take())) == ["goodbye", "hello", "world"]
+
+
+def test_fsspec_http_file_system(ray_start_regular_shared, http_server, http_file):
+    ds = ray.data.read_text(http_file, filesystem=HTTPFileSystem())
+    assert ds.count() > 0
+    # Test auto-resolve of HTTP file system when it is not provided.
+    ds = ray.data.read_text(http_file)
+    assert ds.count() > 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_tf.py

@@ -0,0 +1,412 @@
+import sys
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import ray
+from ray import train
+from ray.data.preprocessors import Concatenator
+from ray.train import ScalingConfig
+
+if sys.version_info <= (3, 12):
+    # Skip this test for Python 3.12+ due to tensorflow incompatibility
+    import tensorflow as tf
+
+    # if tf version is > 2.16, errors cannot be imported as functions
+    # parse version with packaging
+    from packaging import version
+
+    from ray.train.tensorflow import TensorflowTrainer
+
+    if version.parse(tf.__version__) >= version.parse("2.16"):
+        mse = tf.keras.losses.MeanSquaredError()
+        mae = tf.keras.losses.MeanAbsoluteError()
+    else:
+        mse = tf.keras.losses.mean_squared_error
+        mae = tf.keras.losses.mean_absolute_error
+
+
+class TestToTF:
+    def test_autosharding_is_disabled(self):
+        ds = ray.data.from_items([{"spam": 0, "ham": 0}])
+
+        dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
+
+        actual_auto_shard_policy = (
+            dataset.options().experimental_distribute.auto_shard_policy
+        )
+        expected_auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
+        assert actual_auto_shard_policy is expected_auto_shard_policy
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_type(self, include_additional_columns):
+        ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns="spam", label_columns="ham", additional_columns="weight"
+            )
+            feature_spec, label_spec, additional_spec = dataset.element_spec
+        else:
+            dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
+            feature_spec, label_spec = dataset.element_spec
+
+        assert isinstance(feature_spec, tf.TypeSpec)
+        assert isinstance(label_spec, tf.TypeSpec)
+        if include_additional_columns:
+            assert isinstance(additional_spec, tf.TypeSpec)
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_user_provided(self, include_additional_columns):
+        ds = ray.data.from_items([{"spam": 0, "ham": 0, "eggs": 0, "weight": 0}])
+
+        if include_additional_columns:
+            dataset1 = ds.to_tf(
+                feature_columns=["spam", "ham"],
+                label_columns="eggs",
+                additional_columns="weight",
+            )
+            feature_spec, label_spec, additional_spec = dataset1.element_spec
+            dataset2 = ds.to_tf(
+                feature_columns=["spam", "ham"],
+                label_columns="eggs",
+                additional_columns="weight",
+                feature_type_spec=feature_spec,
+                label_type_spec=label_spec,
+                additional_type_spec=additional_spec,
+            )
+            (
+                feature_output_spec,
+                label_output_spec,
+                additional_output_spec,
+            ) = dataset2.element_spec
+        else:
+            dataset1 = ds.to_tf(feature_columns=["spam", "ham"], label_columns="eggs")
+            feature_spec, label_spec = dataset1.element_spec
+            dataset2 = ds.to_tf(
+                feature_columns=["spam", "ham"],
+                label_columns="eggs",
+                feature_type_spec=feature_spec,
+                label_type_spec=label_spec,
+            )
+            feature_output_spec, label_output_spec = dataset2.element_spec
+
+        assert isinstance(label_output_spec, tf.TypeSpec)
+        assert isinstance(feature_output_spec, dict)
+        assert feature_output_spec.keys() == {"spam", "ham"}
+        assert all(
+            isinstance(value, tf.TypeSpec) for value in feature_output_spec.values()
+        )
+        if include_additional_columns:
+            assert isinstance(additional_output_spec, tf.TypeSpec)
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_type_with_multiple_columns(self, include_additional_columns):
+        ds = ray.data.from_items(
+            [{"spam": 0, "ham": 0, "eggs": 0, "weight1": 0, "weight2": 0}]
+        )
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns=["spam", "ham"],
+                label_columns="eggs",
+                additional_columns=["weight1", "weight2"],
+            )
+            (
+                feature_output_signature,
+                _,
+                additional_output_signature,
+            ) = dataset.element_spec
+        else:
+            dataset = ds.to_tf(feature_columns=["spam", "ham"], label_columns="eggs")
+            feature_output_signature, _ = dataset.element_spec
+
+        assert isinstance(feature_output_signature, dict)
+        assert feature_output_signature.keys() == {"spam", "ham"}
+        assert all(
+            isinstance(value, tf.TypeSpec)
+            for value in feature_output_signature.values()
+        )
+
+        if include_additional_columns:
+            assert isinstance(additional_output_signature, dict)
+            assert additional_output_signature.keys() == {"weight1", "weight2"}
+            assert all(
+                isinstance(value, tf.TypeSpec)
+                for value in additional_output_signature.values()
+            )
+
+        df = pd.DataFrame(
+            {
+                "feature1": [0, 1, 2],
+                "feature2": [3, 4, 5],
+                "label": [0, 1, 1],
+                "weight1": [0, 0.1, 0.2],
+                "weight2": [0.3, 0.4, 0.5],
+            }
+        )
+        ds = ray.data.from_pandas(df)
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns=["feature1", "feature2"],
+                label_columns="label",
+                additional_columns=["weight1", "weight2"],
+                batch_size=3,
+            )
+            (
+                feature_output_signature,
+                _,
+                additional_output_signature,
+            ) = dataset.element_spec
+            assert isinstance(additional_output_signature, dict)
+            assert additional_output_signature.keys() == {"weight1", "weight2"}
+            assert all(
+                isinstance(value, tf.TypeSpec)
+                for value in additional_output_signature.values()
+            )
+        else:
+            dataset = ds.to_tf(
+                feature_columns=["feature1", "feature2"],
+                label_columns="label",
+                batch_size=3,
+            )
+            feature_output_signature, _ = dataset.element_spec
+
+        assert isinstance(feature_output_signature, dict)
+        assert feature_output_signature.keys() == {"feature1", "feature2"}
+        assert all(
+            isinstance(value, tf.TypeSpec)
+            for value in feature_output_signature.values()
+        )
+
+        if include_additional_columns:
+            features, labels, additional_metadata = next(iter(dataset))
+            assert (
+                additional_metadata["weight1"].numpy() == df["weight1"].values
+            ).all()
+            assert (
+                additional_metadata["weight2"].numpy() == df["weight2"].values
+            ).all()
+        else:
+            features, labels = next(iter(dataset))
+        assert (labels.numpy() == df["label"].values).all()
+        assert (features["feature1"].numpy() == df["feature1"].values).all()
+        assert (features["feature2"].numpy() == df["feature2"].values).all()
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_name(self, include_additional_columns):
+        ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns="spam", label_columns="ham", additional_columns="weight"
+            )
+            feature_spec, label_spec, additional_spec = dataset.element_spec
+        else:
+            dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
+            feature_spec, label_spec = dataset.element_spec
+
+        assert feature_spec.name == "spam"
+        assert label_spec.name == "ham"
+        if include_additional_columns:
+            assert additional_spec.name == "weight"
+
+    @pytest.mark.parametrize(
+        "data, expected_dtype",
+        # Skip this test for Python 3.12+ due to tensorflow incompatibility
+        [
+            (0, tf.int64),
+            (0.0, tf.double),
+            (False, tf.bool),
+            ("eggs", tf.string),
+            ([1.0, 2.0], tf.float64),
+            (np.zeros([2, 2], dtype=np.float32), tf.float32),
+        ]
+        if sys.version_info <= (3, 12)
+        else [],
+    )
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_dtype(self, data, expected_dtype, include_additional_columns):
+        ds = ray.data.from_items([{"spam": data, "ham": data, "weight": data}])
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns="spam",
+                label_columns="ham",
+                additional_columns="weight",
+            )
+            feature_spec, label_spec, additional_spec = dataset.element_spec
+        else:
+            dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
+            feature_spec, label_spec = dataset.element_spec
+
+        assert feature_spec.dtype == expected_dtype
+        assert label_spec.dtype == expected_dtype
+        if include_additional_columns:
+            assert additional_spec.dtype == expected_dtype
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_shape(self, include_additional_columns):
+        ds = ray.data.from_items(8 * [{"spam": 0, "ham": 0, "weight": 0}])
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns="spam",
+                label_columns="ham",
+                additional_columns="weight",
+                batch_size=4,
+            )
+            feature_spec, label_spec, additional_spec = dataset.element_spec
+            assert tuple(additional_spec.shape) == (None,)
+        else:
+            dataset = ds.to_tf(
+                feature_columns="spam", label_columns="ham", batch_size=4
+            )
+            feature_spec, label_spec = dataset.element_spec
+
+        assert tuple(feature_spec.shape) == (None,)
+        assert tuple(label_spec.shape) == (None,)
+
+        if include_additional_columns:
+            features, labels, additional_metadata = next(iter(dataset))
+            assert tuple(additional_metadata.shape) == (4,)
+        else:
+            features, labels = next(iter(dataset))
+        assert tuple(features.shape) == (4,)
+        assert tuple(labels.shape) == (4,)
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_shape_with_tensors(self, include_additional_columns):
+        ds = ray.data.from_items(
+            8
+            * [
+                {
+                    "spam": np.zeros([3, 32, 32]),
+                    "ham": 0,
+                    "weight": np.zeros([3, 32, 32]),
+                }
+            ]
+        )
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns="spam",
+                label_columns="ham",
+                additional_columns="weight",
+                batch_size=4,
+            )
+            feature_spec, _, additional_spec = dataset.element_spec
+            assert tuple(additional_spec.shape) == (None, 3, 32, 32)
+        else:
+            dataset = ds.to_tf(
+                feature_columns="spam", label_columns="ham", batch_size=4
+            )
+            feature_spec, _ = dataset.element_spec
+
+        assert tuple(feature_spec.shape) == (None, 3, 32, 32)
+
+        if include_additional_columns:
+            features, labels, additional_metadata = next(iter(dataset))
+            assert tuple(additional_metadata.shape) == (4, 3, 32, 32)
+        else:
+            features, labels = next(iter(dataset))
+        assert tuple(features.shape) == (4, 3, 32, 32)
+        assert tuple(labels.shape) == (4,)
+
+    @pytest.mark.parametrize("batch_size", [1, 2])
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_element_spec_shape_with_ragged_tensors(
+        self, batch_size, include_additional_columns
+    ):
+        df = pd.DataFrame(
+            {
+                "spam": [np.zeros([32, 32, 3]), np.zeros([64, 64, 3])],
+                "ham": [0, 0],
+                "weight": [np.zeros([32, 32, 3]), np.zeros([64, 64, 3])],
+            }
+        )
+        ds = ray.data.from_pandas(df)
+
+        if include_additional_columns:
+            dataset = ds.to_tf(
+                feature_columns="spam",
+                label_columns="ham",
+                additional_columns="weight",
+                batch_size=batch_size,
+            )
+            feature_spec, _, additional_spec = dataset.element_spec
+            assert tuple(additional_spec.shape) == (None, None, None, None)
+        else:
+            dataset = ds.to_tf(
+                feature_columns="spam", label_columns="ham", batch_size=batch_size
+            )
+            feature_spec, _ = dataset.element_spec
+
+        assert tuple(feature_spec.shape) == (None, None, None, None)
+
+        if include_additional_columns:
+            features, labels, additional_metadata = next(iter(dataset))
+            assert tuple(additional_metadata.shape) == (batch_size, None, None, None)
+        else:
+            features, labels = next(iter(dataset))
+        assert tuple(features.shape) == (batch_size, None, None, None)
+        assert tuple(labels.shape) == (batch_size,)
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_training(self, include_additional_columns):
+        def build_model() -> tf.keras.Model:
+            return tf.keras.Sequential([tf.keras.layers.Dense(1)])
+
+        def train_func():
+            strategy = tf.distribute.MultiWorkerMirroredStrategy()
+            with strategy.scope():
+                multi_worker_model = build_model()
+                multi_worker_model.compile(
+                    optimizer=tf.keras.optimizers.SGD(),
+                    loss=mae,
+                    metrics=[mse],
+                )
+
+            if include_additional_columns:
+                dataset = train.get_dataset_shard("train").to_tf(
+                    "X", "Y", additional_columns="W", batch_size=4
+                )
+            else:
+                dataset = train.get_dataset_shard("train").to_tf("X", "Y", batch_size=4)
+            multi_worker_model.fit(dataset)
+
+        dataset = ray.data.from_items(8 * [{"X0": 0, "X1": 0, "Y": 0, "W": 0}])
+        concatenator = Concatenator(columns=["X0", "X1"], output_column_name="X")
+        dataset = concatenator.transform(dataset)
+
+        trainer = TensorflowTrainer(
+            train_loop_per_worker=train_func,
+            scaling_config=ScalingConfig(num_workers=2),
+            datasets={"train": dataset},
+        )
+        trainer.fit()
+
+    @pytest.mark.parametrize("include_additional_columns", [False, True])
+    def test_invalid_column_raises_error(self, include_additional_columns):
+        ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
+        with pytest.raises(ValueError):
+            if include_additional_columns:
+                ds.to_tf(
+                    feature_columns="spam",
+                    label_columns="ham",
+                    additional_columns="baz",
+                )
+            else:
+                ds.to_tf(feature_columns="foo", label_columns="bar")
+
+
+if __name__ == "__main__":
+    import sys
+
+    if sys.version_info >= (3, 12):
+        # Skip this test for Python 3.12+ due to to incompatibility tensorflow
+        sys.exit(0)
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_tfrecords.py

@@ -0,0 +1,822 @@
+import json
+import os
+import sys
+from typing import TYPE_CHECKING
+from unittest.mock import MagicMock
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+from pandas.api.types import is_float_dtype, is_int64_dtype, is_object_dtype
+
+import ray
+from ray.data.dataset import Dataset
+from ray.tests.conftest import *  # noqa: F401,F403
+
+if TYPE_CHECKING:
+    from tensorflow_metadata.proto.v0 import schema_pb2
+
+if sys.version_info <= (3, 12):
+    # Skip this test for Python 3.12+ due to to incompatibility tensorflow
+    import tensorflow as tf
+
+
+def _is_object_like(dtype):
+    """Match the pre-Arrow-dtype semantics of ``is_object_dtype``: pandas used
+    object dtype for lists, bytes, and strings; ArrowBlockAccessor.to_pandas()
+    now preserves these as ``pd.ArrowDtype`` via a ``types_mapper``."""
+    if is_object_dtype(dtype):
+        return True
+    if isinstance(dtype, pd.ArrowDtype):
+        pa_type = dtype.pyarrow_dtype
+        return (
+            pa.types.is_list(pa_type)
+            or pa.types.is_large_list(pa_type)
+            or pa.types.is_binary(pa_type)
+            or pa.types.is_large_binary(pa_type)
+            or pa.types.is_string(pa_type)
+            or pa.types.is_large_string(pa_type)
+        )
+    return False
+
+
+def _is_int64_like(dtype):
+    if is_int64_dtype(dtype):
+        return True
+    if isinstance(dtype, pd.ArrowDtype):
+        return dtype.pyarrow_dtype == pa.int64()
+    return False
+
+
+def _is_float_like(dtype):
+    if is_float_dtype(dtype):
+        return True
+    if isinstance(dtype, pd.ArrowDtype):
+        return pa.types.is_floating(dtype.pyarrow_dtype)
+    return False
+
+
+def tf_records_partial():
+    """Underlying data corresponds to `data_partial` fixture."""
+    import tensorflow as tf
+
+    return [
+        # Record one (corresponding to row one).
+        tf.train.Example(
+            features=tf.train.Features(
+                feature={
+                    "int_item": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[1])
+                    ),
+                    "int_list": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[2, 2, 3])
+                    ),
+                    "int_partial": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[])
+                    ),
+                    "float_item": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[1.0])
+                    ),
+                    "float_list": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[2.0, 3.0, 4.0])
+                    ),
+                    "float_partial": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[1.0])
+                    ),
+                    "bytes_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"abc"])
+                    ),
+                    "bytes_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"def", b"1234"])
+                    ),
+                    "bytes_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                    "string_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"uvw"])
+                    ),
+                    "string_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"xyz", b"999"])
+                    ),
+                    "string_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                }
+            )
+        ),
+        # Record two (corresponding to row two).
+        tf.train.Example(
+            features=tf.train.Features(
+                feature={
+                    "int_item": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[2])
+                    ),
+                    "int_list": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[3, 3, 4])
+                    ),
+                    "int_partial": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[9, 2])
+                    ),
+                    "float_item": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[2.0])
+                    ),
+                    "float_list": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[5.0, 6.0, 7.0])
+                    ),
+                    "float_partial": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[])
+                    ),
+                    "bytes_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"ghi"])
+                    ),
+                    "bytes_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"jkl", b"5678"])
+                    ),
+                    "bytes_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"hello"])
+                    ),
+                    "string_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"mno"])
+                    ),
+                    "string_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"pqr", b"111"])
+                    ),
+                    "string_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"world"])
+                    ),
+                }
+            )
+        ),
+    ]
+
+
+def data_partial(with_tf_schema):
+    """TFRecords generated from this corresponds to `tf_records_partial`."""
+    return [
+        # Row one.
+        {
+            "int_item": [1] if with_tf_schema else 1,
+            "int_list": [2, 2, 3],
+            "int_partial": [],
+            "float_item": [1.0] if with_tf_schema else 1.0,
+            "float_list": [2.0, 3.0, 4.0],
+            "float_partial": [1.0] if with_tf_schema else 1.0,
+            "bytes_item": [b"abc"] if with_tf_schema else b"abc",
+            "bytes_list": [b"def", b"1234"],
+            "bytes_partial": [] if with_tf_schema else None,
+            "string_item": ["uvw"] if with_tf_schema else "uvw",
+            "string_list": ["xyz", "999"],
+            "string_partial": [] if with_tf_schema else None,
+        },
+        # Row two.
+        {
+            "int_item": [2] if with_tf_schema else 2,
+            "int_list": [3, 3, 4],
+            "int_partial": [9, 2],
+            "float_item": [2.0] if with_tf_schema else 2.0,
+            "float_list": [5.0, 6.0, 7.0],
+            "float_partial": [] if with_tf_schema else None,
+            "bytes_item": [b"ghi"] if with_tf_schema else b"ghi",
+            "bytes_list": [b"jkl", b"5678"],
+            "bytes_partial": [b"hello"] if with_tf_schema else b"hello",
+            "string_item": ["mno"] if with_tf_schema else "mno",
+            "string_list": ["pqr", "111"],
+            "string_partial": ["world"] if with_tf_schema else "world",
+        },
+    ]
+
+
+def tf_records_empty():
+    """Underlying data corresponds to `data_empty` fixture."""
+    import tensorflow as tf
+
+    return [
+        # Record one (corresponding to row one).
+        tf.train.Example(
+            features=tf.train.Features(
+                feature={
+                    "int_item": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[1])
+                    ),
+                    "int_list": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[2, 2, 3])
+                    ),
+                    "int_partial": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[])
+                    ),
+                    "int_empty": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[])
+                    ),
+                    "float_item": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[1.0])
+                    ),
+                    "float_list": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[2.0, 3.0, 4.0])
+                    ),
+                    "float_partial": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[1.0])
+                    ),
+                    "float_empty": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[])
+                    ),
+                    "bytes_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"abc"])
+                    ),
+                    "bytes_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"def", b"1234"])
+                    ),
+                    "bytes_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                    "bytes_empty": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                    "string_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"uvw"])
+                    ),
+                    "string_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"xyz", b"999"])
+                    ),
+                    "string_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                    "string_empty": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                }
+            )
+        ),
+        # Record two (corresponding to row two).
+        tf.train.Example(
+            features=tf.train.Features(
+                feature={
+                    "int_item": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[2])
+                    ),
+                    "int_list": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[3, 3, 4])
+                    ),
+                    "int_partial": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[9, 2])
+                    ),
+                    "int_empty": tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[])
+                    ),
+                    "float_item": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[2.0])
+                    ),
+                    "float_list": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[5.0, 6.0, 7.0])
+                    ),
+                    "float_partial": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[])
+                    ),
+                    "float_empty": tf.train.Feature(
+                        float_list=tf.train.FloatList(value=[])
+                    ),
+                    "bytes_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"ghi"])
+                    ),
+                    "bytes_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"jkl", b"5678"])
+                    ),
+                    "bytes_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"hello"])
+                    ),
+                    "bytes_empty": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                    "string_item": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"mno"])
+                    ),
+                    "string_list": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"pqr", b"111"])
+                    ),
+                    "string_partial": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[b"world"])
+                    ),
+                    "string_empty": tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[])
+                    ),
+                }
+            )
+        ),
+    ]
+
+
+def data_empty(with_tf_schema):
+    """TFRecords generated from this corresponds to
+    the `tf_records_empty` fixture."""
+    return [
+        # Row one.
+        {
+            "int_item": [1] if with_tf_schema else 1,
+            "int_list": [2, 2, 3],
+            "int_partial": [],
+            "int_empty": [],
+            "float_item": [1.0] if with_tf_schema else 1.0,
+            "float_list": [2.0, 3.0, 4.0],
+            "float_partial": [1.0] if with_tf_schema else 1.0,
+            "float_empty": [],
+            "bytes_item": [b"abc"] if with_tf_schema else b"abc",
+            "bytes_list": [b"def", b"1234"],
+            "bytes_partial": [],
+            "bytes_empty": [],
+            "string_item": ["uvw"] if with_tf_schema else "uvw",
+            "string_list": ["xyz", "999"],
+            "string_partial": [] if with_tf_schema else None,
+            "string_empty": [],
+        },
+        # Row two.
+        {
+            "int_item": [2] if with_tf_schema else 2,
+            "int_list": [3, 3, 4],
+            "int_partial": [9, 2],
+            "int_empty": [],
+            "float_item": [2.0] if with_tf_schema else 2.0,
+            "float_list": [5.0, 6.0, 7.0],
+            "float_partial": [],
+            "float_empty": [],
+            "bytes_item": [b"ghi"] if with_tf_schema else b"ghi",
+            "bytes_list": [b"jkl", b"5678"],
+            "bytes_partial": [b"hello"] if with_tf_schema else b"hello",
+            "bytes_empty": [],
+            "string_item": ["mno"] if with_tf_schema else "mno",
+            "string_list": ["pqr", "111"],
+            "string_partial": ["world"] if with_tf_schema else "world",
+            "string_empty": [],
+        },
+    ]
+
+
+def _features_to_schema(features: "tf.train.Features") -> "schema_pb2.Schema":
+    from tensorflow_metadata.proto.v0 import schema_pb2
+
+    tf_schema = schema_pb2.Schema()
+    for feature_name, feature_msg in features.feature.items():
+        schema_feature = tf_schema.feature.add()
+        schema_feature.name = feature_name
+        if feature_msg.HasField("bytes_list"):
+            schema_feature.type = schema_pb2.FeatureType.BYTES
+        elif feature_msg.HasField("float_list"):
+            schema_feature.type = schema_pb2.FeatureType.FLOAT
+        elif feature_msg.HasField("int64_list"):
+            schema_feature.type = schema_pb2.FeatureType.INT
+    return tf_schema
+
+
+def _ds_eq_streaming(ds_expected, ds_actual) -> bool:
+    # Casting the strings to bytes for comparing string features
+    def _str2bytes(d):
+        for k, v in d.items():
+            if "string" in k:
+                if isinstance(v, list):
+                    d[k] = [vv.encode() for vv in v]
+                elif isinstance(v, str):
+                    d[k] = v.encode()
+        return d
+
+    ds_expected = ds_expected.map(_str2bytes)
+    assert ds_expected.take() == ds_actual.take()
+
+
+@pytest.mark.parametrize(
+    "with_tf_schema,compression",
+    [
+        (True, None),
+        (False, None),
+    ],
+)
+def test_read_tfrecords(
+    with_tf_schema,
+    compression,
+    ray_start_regular_shared_2_cpus,
+    tmp_path,
+):
+    import pandas as pd
+    import tensorflow as tf
+
+    example = tf_records_empty()[0]
+
+    tf_schema = None
+    if with_tf_schema:
+        tf_schema = _features_to_schema(example.features)
+
+    path = os.path.join(tmp_path, "data.tfrecords")
+    with tf.io.TFRecordWriter(
+        path=path, options=tf.io.TFRecordOptions(compression_type=compression)
+    ) as writer:
+        writer.write(example.SerializeToString())
+
+    arrow_open_stream_args = None
+    if compression:
+        arrow_open_stream_args = {"compression": compression}
+
+    ds = ray.data.read_tfrecords(
+        path,
+        tf_schema=tf_schema,
+        arrow_open_stream_args=arrow_open_stream_args,
+    )
+
+    df = ds.to_pandas()
+    # Protobuf serializes features in a non-deterministic order.
+    if with_tf_schema:
+        assert _is_object_like(dict(df.dtypes)["int_item"])
+    else:
+        assert _is_int64_like(dict(df.dtypes)["int_item"])
+    assert _is_object_like(dict(df.dtypes)["int_list"])
+    assert _is_object_like(dict(df.dtypes)["int_partial"])
+    assert _is_object_like(dict(df.dtypes)["int_empty"])
+
+    if with_tf_schema:
+        assert _is_object_like(dict(df.dtypes)["float_item"])
+        assert _is_object_like(dict(df.dtypes)["float_partial"])
+    else:
+        assert _is_float_like(dict(df.dtypes)["float_item"])
+        assert _is_float_like(dict(df.dtypes)["float_partial"])
+    assert _is_object_like(dict(df.dtypes)["float_list"])
+    assert _is_object_like(dict(df.dtypes)["float_empty"])
+
+    assert _is_object_like(dict(df.dtypes)["bytes_item"])
+    assert _is_object_like(dict(df.dtypes)["bytes_partial"])
+    assert _is_object_like(dict(df.dtypes)["bytes_list"])
+    assert _is_object_like(dict(df.dtypes)["bytes_empty"])
+
+    assert _is_object_like(dict(df.dtypes)["string_item"])
+    assert _is_object_like(dict(df.dtypes)["string_partial"])
+    assert _is_object_like(dict(df.dtypes)["string_list"])
+    assert _is_object_like(dict(df.dtypes)["string_empty"])
+
+    # If the schema is specified, we should not perform the
+    # automatic unwrapping of single-element lists.
+    if with_tf_schema:
+        assert isinstance(df["int_item"], pd.Series)
+        assert df["int_item"].tolist() == [[1]]
+    else:
+        assert list(df["int_item"]) == [1]
+    assert np.array_equal(df["int_list"][0], np.array([2, 2, 3]))
+    assert np.array_equal(df["int_partial"][0], np.array([], dtype=np.int64))
+    assert np.array_equal(df["int_empty"][0], np.array([], dtype=np.int64))
+
+    if with_tf_schema:
+        assert isinstance(df["float_item"], pd.Series)
+        assert df["float_item"].tolist() == [[1.0]]
+        assert df["float_partial"].tolist() == [[1.0]]
+    else:
+        assert list(df["float_item"]) == [1.0]
+        assert list(df["float_partial"]) == [1.0]
+    assert np.array_equal(df["float_list"][0], np.array([2.0, 3.0, 4.0]))
+    assert np.array_equal(df["float_empty"][0], np.array([], dtype=np.float32))
+
+    if with_tf_schema:
+        assert isinstance(df["bytes_item"], pd.Series)
+        assert df["bytes_item"].tolist() == [[b"abc"]]
+        assert isinstance(df["string_item"], pd.Series)
+        assert df["string_item"].tolist() == [[b"uvw"]]  # strings are read as bytes
+    else:
+        assert list(df["bytes_item"]) == [b"abc"]
+        assert list(df["string_item"]) == [b"uvw"]
+    assert np.array_equal(df["bytes_list"][0], np.array([b"def", b"1234"]))
+    assert np.array_equal(df["bytes_partial"][0], np.array([], dtype=np.bytes_))
+    assert np.array_equal(df["bytes_empty"][0], np.array([], dtype=np.bytes_))
+
+    assert np.array_equal(df["string_list"][0], np.array([b"xyz", b"999"]))
+    assert np.array_equal(df["string_partial"][0], np.array([], dtype=np.bytes_))
+    assert np.array_equal(df["string_empty"][0], np.array([], dtype=np.bytes_))
+
+
+@pytest.fixture
+def mock_ray_data_read_tfrecords(mocker):
+    mock_read_tfrecords = mocker.patch("ray.data.read_tfrecords")
+    mock_read_tfrecords.return_value = MagicMock(spec=Dataset)
+    return mock_read_tfrecords
+
+
+@pytest.mark.parametrize("num_cpus", [1, 2, 4])
+def test_read_tfrecords_ray_remote_args(
+    ray_start_regular_shared_2_cpus,
+    mock_ray_data_read_tfrecords,
+    tmp_path,
+    num_cpus,
+):
+    import tensorflow as tf
+
+    example = tf_records_empty()[0]
+    path = os.path.join(tmp_path, "data.tfrecords")
+    with tf.io.TFRecordWriter(path=path) as writer:
+        writer.write(example.SerializeToString())
+    ray_remote_args = {"num_cpus": num_cpus}
+    ds = ray.data.read_tfrecords(
+        paths=[path],
+        ray_remote_args=ray_remote_args,
+    )
+    assert isinstance(ds, Dataset)
+    mock_ray_data_read_tfrecords.assert_called_once()
+    args, kwargs = mock_ray_data_read_tfrecords.call_args
+    assert kwargs["paths"] == [path]
+    assert kwargs["ray_remote_args"] == ray_remote_args
+
+
+@pytest.mark.parametrize("with_tf_schema", (True, False))
+def test_write_tfrecords(
+    with_tf_schema,
+    ray_start_regular_shared_2_cpus,
+    tmp_path,
+):
+    """Test that write_tfrecords writes TFRecords correctly.
+
+    Test this by writing a Dataset to a TFRecord (function under test),
+    reading it back out into a tf.train.Example,
+    and checking that the result is analogous to the original Dataset.
+    """
+
+    import tensorflow as tf
+
+    # The dataset we will write to a .tfrecords file.
+    ds = ray.data.from_items(
+        data_partial(with_tf_schema),
+        # Here, we specify `override_num_blocks=1` to ensure that all rows end up in
+        # the same block, which is required for type inference involving partially
+        # missing columns.
+        override_num_blocks=1,
+    )
+
+    # The corresponding tf.train.Example that we would expect to read
+    # from this dataset.
+    expected_records = tf_records_partial()
+
+    tf_schema = None
+    if with_tf_schema:
+        features = expected_records[0].features
+        tf_schema = _features_to_schema(features)
+
+    # Perform the test.
+    # Write the dataset to a .tfrecords file.
+    ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
+
+    # Read the Examples back out from the .tfrecords file.
+    # This follows the offical TFRecords tutorial:
+    # https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file_2
+
+    filenames = sorted(os.listdir(tmp_path))
+    filepaths = [os.path.join(tmp_path, filename) for filename in filenames]
+    raw_dataset = tf.data.TFRecordDataset(filepaths)
+
+    tfrecords = []
+    for raw_record in raw_dataset:
+        example = tf.train.Example()
+        example.ParseFromString(raw_record.numpy())
+        tfrecords.append(example)
+
+    assert tfrecords == expected_records
+
+
+@pytest.mark.parametrize("with_tf_schema", (True, False))
+def test_write_tfrecords_empty_features(
+    with_tf_schema,
+    ray_start_regular_shared_2_cpus,
+    tmp_path,
+):
+    """Test that write_tfrecords writes TFRecords with completely empty features
+    correctly (i.e. the case where type inference from partially filled features
+    is not possible). We expect this to succeed when passing an explicit `tf_schema`
+    param, and otherwise will raise a `ValueError`.
+
+    Test this by writing a Dataset to a TFRecord (function under test),
+    reading it back out into a tf.train.Example,
+    and checking that the result is analogous to the original Dataset.
+    """
+
+    import tensorflow as tf
+
+    # The dataset we will write to a .tfrecords file.
+    ds = ray.data.from_items(data_empty(with_tf_schema))
+
+    # The corresponding tf.train.Example that we would expect to read
+    # from this dataset.
+    expected_records = tf_records_empty()
+
+    if not with_tf_schema:
+        with pytest.raises(ValueError):
+            # Type inference from fully empty columns should fail if
+            # no schema is specified.
+            ds.write_tfrecords(tmp_path)
+    else:
+        features = expected_records[0].features
+        tf_schema = _features_to_schema(features)
+
+        # Perform the test.
+        # Write the dataset to a .tfrecords file.
+        ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
+
+        # Read the Examples back out from the .tfrecords file.
+        # This follows the offical TFRecords tutorial:
+        # https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file_2
+
+        filenames = sorted(os.listdir(tmp_path))
+        filepaths = [os.path.join(tmp_path, filename) for filename in filenames]
+        raw_dataset = tf.data.TFRecordDataset(filepaths)
+
+        tfrecords = []
+        for raw_record in raw_dataset:
+            example = tf.train.Example()
+            example.ParseFromString(raw_record.numpy())
+            tfrecords.append(example)
+
+        assert tfrecords == expected_records
+
+
+@pytest.mark.parametrize("with_tf_schema", (True, False))
+def test_readback_tfrecords(
+    ray_start_regular_shared_2_cpus,
+    tmp_path,
+    with_tf_schema,
+):
+    """
+    Test reading back TFRecords written using datasets.
+    The dataset we read back should be the same that we wrote.
+    """
+
+    # The dataset we will write to a .tfrecords file.
+    # Here and in the read_tfrecords call below, we specify `override_num_blocks=1`
+    # to ensure that all rows end up in the same block, which is required
+    # for type inference involving partially missing columns.
+    ds = ray.data.from_items(data_partial(with_tf_schema), override_num_blocks=1)
+    expected_records = tf_records_partial()
+
+    tf_schema = None
+    if with_tf_schema:
+        features = expected_records[0].features
+        tf_schema = _features_to_schema(features)
+
+    # Write the TFRecords.
+    ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
+    # Read the TFRecords.
+    readback_ds = ray.data.read_tfrecords(
+        tmp_path, tf_schema=tf_schema, override_num_blocks=1
+    )
+    _ds_eq_streaming(ds, readback_ds)
+
+
+@pytest.mark.parametrize("with_tf_schema", (True, False))
+def test_readback_tfrecords_empty_features(
+    ray_start_regular_shared_2_cpus,
+    tmp_path,
+    with_tf_schema,
+):
+    """
+    Test reading back TFRecords written using datasets.
+    The dataset we read back should be the same that we wrote.
+    """
+
+    # The dataset we will write to a .tfrecords file.
+    ds = ray.data.from_items(data_empty(with_tf_schema))
+    if not with_tf_schema:
+        with pytest.raises(ValueError):
+            # With no schema specified, this should fail because
+            # type inference on completely empty columns is ambiguous.
+            ds.write_tfrecords(tmp_path)
+    else:
+        ds = ray.data.from_items(data_empty(with_tf_schema), override_num_blocks=1)
+        expected_records = tf_records_empty()
+
+        features = expected_records[0].features
+        tf_schema = _features_to_schema(features)
+
+        # Write the TFRecords.
+        ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
+
+        # Read the TFRecords.
+        readback_ds = ray.data.read_tfrecords(
+            tmp_path,
+            tf_schema=tf_schema,
+            override_num_blocks=1,
+        )
+        _ds_eq_streaming(ds, readback_ds)
+
+
+def test_write_tfrecords_tensor(
+    ray_start_regular_shared_2_cpus, tmp_path, tensor_format_context
+):
+    """Test that write_tfrecords handles tensor data by serializing
+    tensors to bytes via tf.io.serialize_tensor, preserving shape and dtype."""
+    import tensorflow as tf
+
+    ds = ray.data.range_tensor(3, shape=(2, 2))
+
+    ds.write_tfrecords(tmp_path)
+
+    # Read back the raw TFRecord examples and deserialize tensors.
+    filenames = sorted(os.listdir(tmp_path))
+    filepaths = [os.path.join(tmp_path, filename) for filename in filenames]
+    raw_dataset = tf.data.TFRecordDataset(filepaths)
+
+    results = []
+    for raw_record in raw_dataset:
+        example = tf.train.Example()
+        example.ParseFromString(raw_record.numpy())
+        serialized = example.features.feature["data"].bytes_list.value[0]
+        tensor = tf.io.parse_tensor(serialized, out_type=tf.int64)
+        results.append(tensor.numpy())
+
+    assert len(results) == 3
+    for i, result in enumerate(results):
+        assert result.shape == (2, 2)
+        expected = np.full((2, 2), i)
+        np.testing.assert_array_equal(result, expected)
+
+
+def test_write_invalid_tfrecords(ray_start_regular_shared_2_cpus, tmp_path):
+    """
+    If we try to write a dataset with invalid TFRecord datatypes,
+    ValueError should be raised.
+    """
+
+    ds = ray.data.from_items([{"item": None}])
+
+    with pytest.raises(ValueError):
+        ds.write_tfrecords(tmp_path)
+
+
+def test_read_invalid_tfrecords(ray_start_regular_shared_2_cpus, tmp_path):
+    file_path = os.path.join(tmp_path, "file.json")
+    with open(file_path, "w") as file:
+        json.dump({"number": 0, "string": "foo"}, file)
+
+    # Expect RuntimeError raised when reading JSON as TFRecord file.
+    with pytest.raises(RuntimeError, match="Failed to read TFRecord file"):
+        ray.data.read_tfrecords(file_path).schema()
+
+
+def test_read_with_invalid_schema(
+    ray_start_regular_shared_2_cpus,
+    tmp_path,
+):
+    from tensorflow_metadata.proto.v0 import schema_pb2
+
+    # The dataset we will write to a .tfrecords file.
+    ds = ray.data.from_items(data_partial(True), override_num_blocks=1)
+    expected_records = tf_records_partial()
+
+    # Build fake schema proto with missing/incorrect field name
+    tf_schema_wrong_name = schema_pb2.Schema()
+    schema_feature = tf_schema_wrong_name.feature.add()
+    schema_feature.name = "wrong_name"
+    schema_feature.type = schema_pb2.FeatureType.INT
+
+    # Build a fake schema proto with incorrect type
+    tf_schema_wrong_type = _features_to_schema(expected_records[0].features)
+    for schema_feature in tf_schema_wrong_type.feature:
+        if schema_feature.name == "bytes_item":
+            schema_feature.type = schema_pb2.FeatureType.INT
+            break
+
+    # Writing with incorrect schema should raise a `ValueError`
+    with pytest.raises(ValueError) as e:
+        ds.write_tfrecords(tmp_path, tf_schema=tf_schema_wrong_name)
+    assert "Found extra unexpected feature" in str(e.value.args[0])
+
+    with pytest.raises(ValueError) as e:
+        ds.write_tfrecords(tmp_path, tf_schema=tf_schema_wrong_type)
+    assert str(e.value.args[0]) == (
+        "Schema field type mismatch during write: "
+        "specified type is int, but underlying type is bytes"
+    )
+
+    # Complete a valid write, then try reading with incorrect schema,
+    # which should raise a `ValueError`.
+    ds.write_tfrecords(tmp_path)
+    with pytest.raises(ValueError) as e:
+        ray.data.read_tfrecords(tmp_path, tf_schema=tf_schema_wrong_name).materialize()
+    assert "Found extra unexpected feature" in str(e.value.args[0])
+
+    with pytest.raises(ValueError) as e:
+        ray.data.read_tfrecords(tmp_path, tf_schema=tf_schema_wrong_type).materialize()
+    assert str(e.value.args[0]) == (
+        "Schema field type mismatch during read: "
+        "specified type is int, but underlying type is bytes"
+    )
+
+
+@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
+def test_write_min_rows_per_file(
+    tmp_path, ray_start_regular_shared_2_cpus, min_rows_per_file
+):
+    ray.data.range(100, override_num_blocks=20).write_tfrecords(
+        tmp_path, min_rows_per_file=min_rows_per_file
+    )
+
+    for filename in os.listdir(tmp_path):
+        dataset = tf.data.TFRecordDataset(os.path.join(tmp_path, filename))
+        assert len(list(dataset)) == min_rows_per_file
+
+
+if __name__ == "__main__":
+    import sys
+
+    if sys.version_info >= (3, 12):
+        # Skip this test for Python 3.12+ due to to incompatibility tensorflow
+        sys.exit(0)
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_torch.py

@@ -0,0 +1,93 @@
+import pytest
+import torch
+
+import ray
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.util import extract_values
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.mark.parametrize("local_read", [True, False])
+def test_from_torch_map_style_dataset(ray_start_10_cpus_shared, local_read):
+    class StubDataset(torch.utils.data.Dataset):
+        def __len__(self):
+            return 1
+
+        def __getitem__(self, index):
+            return index
+
+    torch_dataset = StubDataset()
+
+    ray_dataset = ray.data.from_torch(torch_dataset, local_read=local_read)
+
+    actual_data = ray_dataset.take_all()
+    assert actual_data == [{"item": 0}]
+
+
+def test_from_torch_iterable_style_dataset(ray_start_10_cpus_shared):
+    class StubIterableDataset(torch.utils.data.IterableDataset):
+        def __len__(self):
+            return 1
+
+        def __iter__(self):
+            return iter([0])
+
+    iter_torch_dataset = StubIterableDataset()
+
+    ray_dataset = ray.data.from_torch(iter_torch_dataset)
+
+    actual_data = ray_dataset.take_all()
+    assert actual_data == [{"item": 0}]
+
+
+@pytest.mark.parametrize("local_read", [True, False])
+def test_from_torch_boundary_conditions(ray_start_10_cpus_shared, local_read):
+    """
+    Tests that from_torch respects __len__ for map-style datasets
+    """
+    from torch.utils.data import Dataset
+
+    class BoundaryTestMapDataset(Dataset):
+        """A map-style dataset where __len__ is less than the underlying data size."""
+
+        def __init__(self, data, length):
+            super().__init__()
+            self._data = data
+            self._length = length
+            assert self._length <= len(
+                self._data
+            ), "Length must be <= data size to properly test boundary conditions"
+
+        def __len__(self):
+            return self._length
+
+        def __getitem__(self, index):
+            if not (0 <= index < self._length):
+                # Note: don't use IndexError because we want to fail clearly if
+                # Ray Data tries to access beyond __len__ - 1
+                raise RuntimeError(
+                    f"Index {index} out of bounds for dataset with length {self._length}"
+                )
+            return self._data[index]
+
+    source_data = list(range(10))
+    dataset_len = 8  # Intentionally less than len(source_data)
+
+    # --- Test MapDataset ---
+    map_ds = BoundaryTestMapDataset(source_data, dataset_len)
+    # Expected data only includes elements up to dataset_len - 1
+    expected_items = source_data[:dataset_len]
+
+    ray_ds_map = ray.data.from_torch(map_ds, local_read=local_read)
+    actual_items_map = extract_values("item", list(ray_ds_map.take_all()))
+
+    # This assertion verifies that ray_ds_map didn't try to access index 8 or 9,
+    # which would have raised an IndexError in BoundaryTestMapDataset.__getitem__
+    assert actual_items_map == expected_items
+    assert len(actual_items_map) == dataset_len
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_turbopuffer_datasink.py

@@ -0,0 +1,703 @@
+"""Tests for TurbopufferDatasink.
+
+Organized by critical paths:
+1. Constructor validation
+2. Client initialization
+3. Arrow table preparation
+4. Single-namespace batching
+5. Transform to Turbopuffer format
+6. Retry logic
+7. End-to-end write orchestration
+8. Streaming behavior
+9. Multi-namespace writes
+10. Serialization
+"""
+
+import pickle
+import sys
+import time
+import uuid
+from typing import List
+from unittest.mock import MagicMock, patch
+
+import pyarrow as pa
+import pytest
+from packaging.version import parse as parse_version
+
+from ray.data._internal.datasource.turbopuffer_datasink import TurbopufferDatasink
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+
+# Skip all tests if PyArrow version is less than 19.0
+pytestmark = pytest.mark.skipif(
+    get_pyarrow_version() < parse_version("19.0.0"),
+    reason="TurbopufferDatasink tests require PyArrow >= 19.0",
+)
+
+# =============================================================================
+# Fixtures
+# =============================================================================
+
+
+@pytest.fixture(autouse=True)
+def mock_turbopuffer_module(monkeypatch):
+    """Provide a fake turbopuffer module so imports in the datasink succeed."""
+    fake_module = MagicMock()
+    fake_module.Turbopuffer = MagicMock()
+    with patch.dict(sys.modules, {"turbopuffer": fake_module}):
+        yield fake_module
+
+
+@pytest.fixture
+def sink():
+    """Default sink with minimal required arguments."""
+    return TurbopufferDatasink(
+        namespace="default_ns",
+        region="gcp-us-central1",
+        api_key="test-api-key",
+    )
+
+
+@pytest.fixture
+def mock_client():
+    """Mock Turbopuffer client with namespace support."""
+    client = MagicMock()
+    client.namespace.return_value = MagicMock()
+    return client
+
+
+@pytest.fixture
+def sample_table():
+    """Standard table with id and vector columns."""
+    return pa.table(
+        {
+            "id": [1, 2, 3],
+            "vector": [[0.1], [0.2], [0.3]],
+        }
+    )
+
+
+def make_sink(**kwargs) -> TurbopufferDatasink:
+    """Helper to construct a sink with minimal required arguments."""
+    params = {
+        "namespace": "default_ns",
+        "region": "gcp-us-central1",
+        "api_key": "test-api-key",
+    }
+    params.update(kwargs)
+    return TurbopufferDatasink(**params)
+
+
+# =============================================================================
+# 1. Constructor validation
+# =============================================================================
+
+
+class TestConstructorValidation:
+    """Tests for constructor argument validation."""
+
+    def test_requires_namespace_or_namespace_column(self):
+        """Must provide exactly one of namespace / namespace_column."""
+        with pytest.raises(ValueError, match="Either.*must be provided"):
+            TurbopufferDatasink(
+                region="gcp-us-central1",
+                api_key="k",
+            )
+
+    def test_rejects_both_namespace_and_namespace_column(self):
+        """Cannot provide both namespace and namespace_column."""
+        with pytest.raises(ValueError, match="exactly one"):
+            TurbopufferDatasink(
+                namespace="ns",
+                namespace_column="ns_col",
+                region="gcp-us-central1",
+                api_key="k",
+            )
+
+    def test_namespace_column_cannot_be_id_or_vector(self):
+        """namespace_column must not collide with id_column or vector_column."""
+        with pytest.raises(ValueError, match="namespace_column.*must not be the same"):
+            make_sink(namespace=None, namespace_column="id")
+
+        with pytest.raises(ValueError, match="namespace_column.*must not be the same"):
+            make_sink(namespace=None, namespace_column="vector")
+
+    def test_api_key_from_env(self, monkeypatch):
+        """API key can come from environment variable."""
+        monkeypatch.delenv("TURBOPUFFER_API_KEY", raising=False)
+
+        # No api_key and no env var -> error
+        with pytest.raises(ValueError):
+            TurbopufferDatasink(namespace="ns", region="gcp-us-central1")
+
+        # With env var, init should succeed
+        monkeypatch.setenv("TURBOPUFFER_API_KEY", "env-api-key")
+        sink = TurbopufferDatasink(namespace="ns", region="gcp-us-central1")
+        assert sink.api_key == "env-api-key"
+
+    def test_rejects_same_id_and_vector_column(self):
+        """id_column and vector_column must be distinct."""
+        with pytest.raises(ValueError, match="id_column and vector_column"):
+            make_sink(id_column="doc_id", vector_column="doc_id")
+
+    def test_accepts_region_only(self):
+        """Constructor succeeds with region and no base_url."""
+        sink = make_sink(region="gcp-us-central1")
+        assert sink.region == "gcp-us-central1"
+        assert sink.base_url is None
+
+    def test_accepts_base_url_only(self):
+        """Constructor succeeds with base_url and no region."""
+        sink = make_sink(
+            region=None,
+            base_url="https://gcp-us-central1.turbopuffer.com",
+        )
+        assert sink.base_url == "https://gcp-us-central1.turbopuffer.com"
+        assert sink.region is None
+
+    def test_rejects_both_region_and_base_url(self):
+        """Cannot provide both region and base_url."""
+        with pytest.raises(ValueError, match="exactly one of 'region' or 'base_url'"):
+            make_sink(
+                region="gcp-us-central1",
+                base_url="https://gcp-us-central1.turbopuffer.com",
+            )
+
+    def test_rejects_neither_region_nor_base_url(self):
+        """Must provide at least one of region or base_url."""
+        with pytest.raises(ValueError, match="Either 'region' or 'base_url'"):
+            TurbopufferDatasink(
+                namespace="ns",
+                api_key="k",
+            )
+
+
+# =============================================================================
+# 2. Client initialization
+# =============================================================================
+
+
+class TestClientInitialization:
+    """Tests for Turbopuffer client lazy initialization."""
+
+    def test_lazy_initialization(self, sink, mock_turbopuffer_module):
+        """Client is created lazily and cached."""
+        client1 = sink._get_client()
+        client2 = sink._get_client()
+
+        assert client1 is client2
+        mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
+            api_key="test-api-key",
+            region="gcp-us-central1",
+        )
+
+    def test_uses_explicit_region(self, mock_turbopuffer_module):
+        """Client uses the configured region."""
+        sink = make_sink(region="custom-region")
+        sink._get_client()
+
+        mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
+            api_key="test-api-key",
+            region="custom-region",
+        )
+
+    def test_uses_base_url(self, mock_turbopuffer_module):
+        """Client uses base_url when region is not provided."""
+        sink = make_sink(
+            region=None,
+            base_url="https://gcp-us-central1.turbopuffer.com",
+        )
+        sink._get_client()
+
+        mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
+            api_key="test-api-key",
+            base_url="https://gcp-us-central1.turbopuffer.com",
+        )
+
+    def test_base_url_does_not_pass_region(self, mock_turbopuffer_module):
+        """When base_url is used, region is not passed to the client."""
+        sink = make_sink(
+            region=None,
+            base_url="https://custom.turbopuffer.com",
+        )
+        sink._get_client()
+
+        call_kwargs = mock_turbopuffer_module.Turbopuffer.call_args[1]
+        assert "region" not in call_kwargs
+        assert call_kwargs["base_url"] == "https://custom.turbopuffer.com"
+
+    def test_region_does_not_pass_base_url(self, mock_turbopuffer_module):
+        """When region is used, base_url is not passed to the client."""
+        sink = make_sink(region="gcp-us-central1")
+        sink._get_client()
+
+        call_kwargs = mock_turbopuffer_module.Turbopuffer.call_args[1]
+        assert "base_url" not in call_kwargs
+        assert call_kwargs["region"] == "gcp-us-central1"
+
+
+# =============================================================================
+# 3. Arrow table preparation
+# =============================================================================
+
+
+class TestArrowTablePreparation:
+    """Tests for _prepare_arrow_table."""
+
+    def test_renames_columns_and_filters_null_ids(self):
+        """Custom columns are renamed and null IDs filtered."""
+        table = pa.table(
+            {
+                "doc_id": [1, 2, None],
+                "emb": [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
+            }
+        )
+        sink = make_sink(id_column="doc_id", vector_column="emb")
+
+        prepared = sink._prepare_arrow_table(table)
+
+        # Null ID row filtered, columns renamed to id/vector
+        expected = pa.table(
+            {
+                "id": [1, 2],
+                "vector": [[0.1, 0.2], [0.3, 0.4]],
+            }
+        )
+        assert prepared.equals(expected)
+
+    def test_missing_id_column_raises(self):
+        """Missing custom ID column raises ValueError."""
+        table = pa.table({"other": [1, 2, 3]})
+        sink = make_sink(id_column="doc_id")
+
+        with pytest.raises(ValueError):
+            sink._prepare_arrow_table(table)
+
+    def test_missing_vector_column_raises(self):
+        """Missing vector column raises ValueError."""
+        table = pa.table({"id": [1, 2, 3]})
+        sink = make_sink(vector_column="embedding")
+
+        with pytest.raises(ValueError, match="Vector column 'embedding' not found"):
+            sink._prepare_arrow_table(table)
+
+    @pytest.mark.parametrize(
+        "existing_col,custom_col,expected_match",
+        [
+            ("id", "doc_id", "already has.*'id' column"),
+            ("vector", "emb", "already has.*'vector' column"),
+        ],
+        ids=["id_conflict", "vector_conflict"],
+    )
+    def test_conflicting_column_names_raise(
+        self, existing_col, custom_col, expected_match
+    ):
+        """Raise if table already has target column name."""
+        if existing_col == "id":
+            table = pa.table(
+                {"id": [1, 2], "doc_id": [10, 20], "vector": [[0.1], [0.2]]}
+            )
+            sink = make_sink(id_column="doc_id")
+        else:
+            table = pa.table(
+                {"id": [1, 2], "vector": [[0.1], [0.2]], "emb": [[0.3], [0.4]]}
+            )
+            sink = make_sink(vector_column="emb")
+
+        with pytest.raises(ValueError, match=expected_match):
+            sink._prepare_arrow_table(table)
+
+
+# =============================================================================
+# 4. Single-namespace batching
+# =============================================================================
+
+
+class TestSingleNamespaceBatching:
+    """Tests for write batching behavior."""
+
+    def test_batches_by_batch_size(self, mock_client):
+        """Large tables are split into batches."""
+        num_rows = 25
+        table = pa.table(
+            {
+                "id": list(range(num_rows)),
+                "vector": [[float(i)] for i in range(num_rows)],
+            }
+        )
+        sink = make_sink(batch_size=10)
+        batch_sizes: List[int] = []
+
+        def track_batch(ns, batch, namespace_name=None):
+            # batch is a RecordBatch, get its row count
+            batch_sizes.append(batch.num_rows)
+
+        with patch.object(sink, "_get_client", return_value=mock_client):
+            with patch.object(sink, "_write_batch_with_retry", side_effect=track_batch):
+                sink.write([table], ctx=None)
+
+        assert batch_sizes == [10, 10, 5]
+
+    def test_skips_empty_blocks(self, sink):
+        """Empty blocks don't trigger namespace writes."""
+        empty_table = pa.table({"id": [], "vector": []})
+
+        with patch.object(sink, "_get_client") as mock_get_client:
+            with patch.object(sink, "_write_batch_with_retry") as mock_write:
+                mock_get_client.return_value = MagicMock()
+                sink.write([empty_table], ctx=None)
+
+        mock_write.assert_not_called()
+
+
+# =============================================================================
+# 5. Transform to Turbopuffer format
+# =============================================================================
+
+
+class TestTransformToTurbopufferFormat:
+    """Tests for _transform_to_turbopuffer_format."""
+
+    def test_requires_id_column(self, sink):
+        """Table must have 'id' column."""
+        table = pa.table({"col": [1, 2, 3]})
+
+        with pytest.raises(ValueError):
+            sink._transform_to_turbopuffer_format(table)
+
+    def test_converts_uuid_bytes_to_native_uuid(self, sink):
+        """16-byte binary IDs become native uuid.UUID objects.
+
+        Per Turbopuffer performance docs, native UUIDs (16 bytes) are more
+        efficient than string UUIDs (36 bytes).
+        """
+        u = uuid.uuid4()
+
+        # ID column must be binary(16) for UUID conversion
+        table = pa.table(
+            {
+                "id": pa.array([u.bytes], type=pa.binary(16)),
+                "vector": [[0.1, 0.2]],
+            }
+        )
+
+        columns = sink._transform_to_turbopuffer_format(table)
+
+        expected = {
+            "id": [u],  # Native uuid.UUID, not bytes
+            "vector": [[0.1, 0.2]],
+        }
+        assert columns == expected
+        assert isinstance(columns["id"][0], uuid.UUID)
+
+
+# =============================================================================
+# 6. Retry logic
+# =============================================================================
+
+
+class TestRetryLogic:
+    """Tests for _write_batch_with_retry."""
+
+    @pytest.fixture
+    def sample_batch(self):
+        """A simple batch for retry tests."""
+        return pa.table({"id": [1], "vector": [[0.1]]})
+
+    def test_success_first_try(self, sink, sample_batch):
+        """Successful write on first attempt."""
+        namespace = MagicMock()
+
+        sink._write_batch_with_retry(namespace, sample_batch)
+
+        namespace.write.assert_called_once_with(
+            upsert_columns={"id": [1], "vector": [[0.1]]},
+            schema=None,
+            distance_metric="cosine_distance",
+        )
+
+    def test_retries_then_succeeds(self, sink, sample_batch, monkeypatch):
+        """Transient failures are retried."""
+        monkeypatch.setattr(time, "sleep", lambda _: None)
+        namespace = MagicMock()
+        attempts = {"count": 0}
+
+        def flaky_write(*args, **kwargs):
+            attempts["count"] += 1
+            if attempts["count"] < 3:
+                raise RuntimeError("temporary error")
+
+        namespace.write.side_effect = flaky_write
+
+        sink._write_batch_with_retry(namespace, sample_batch)
+
+        assert attempts["count"] == 3
+
+    def test_exhausts_retries_and_raises(self, sink, sample_batch, monkeypatch):
+        """Persistent failures exhaust retries and raise."""
+        monkeypatch.setattr(time, "sleep", lambda _: None)
+        namespace = MagicMock()
+        namespace.write.side_effect = RuntimeError("persistent error")
+
+        with pytest.raises(RuntimeError, match="persistent error"):
+            sink._write_batch_with_retry(namespace, sample_batch)
+
+        assert namespace.write.call_count == 5  # max_attempts=5
+
+    @pytest.mark.parametrize(
+        "schema,distance_metric",
+        [
+            ({"field": "value"}, "cosine_distance"),
+            (None, "euclidean_squared"),
+            ({"type": "string"}, "euclidean_squared"),
+        ],
+        ids=["with_schema", "alt_metric", "both"],
+    )
+    def test_configurable_options(self, schema, distance_metric):
+        """Schema and distance_metric are passed to write."""
+        sink = make_sink(schema=schema, distance_metric=distance_metric)
+        namespace = MagicMock()
+        batch = pa.table({"id": [1], "vector": [[0.1]]})
+
+        sink._write_batch_with_retry(namespace, batch)
+
+        namespace.write.assert_called_once_with(
+            upsert_columns={"id": [1], "vector": [[0.1]]},
+            schema=schema,
+            distance_metric=distance_metric,
+        )
+
+
+# =============================================================================
+# 7. End-to-end write orchestration
+# =============================================================================
+
+
+class TestWriteOrchestration:
+    """Tests for top-level write() method."""
+
+    def test_write_multiple_blocks(self, sink):
+        """Multiple blocks are processed and written."""
+        blocks = [
+            pa.table({"id": [1, 2], "vector": [[1.0], [2.0]]}),
+            pa.table({"id": [3], "vector": [[3.0]]}),
+        ]
+        write_calls = []
+
+        def track_write(ns, batch, namespace_name=None):
+            write_calls.append(batch.num_rows)
+
+        with patch.object(sink, "_get_client") as mock_get_client:
+            mock_client = MagicMock()
+            mock_get_client.return_value = mock_client
+            with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
+                sink.write(blocks, ctx=None)
+
+        # Two blocks written
+        assert len(write_calls) == 2
+        assert write_calls == [2, 1]
+
+        # Namespace accessed with correct name
+        mock_client.namespace.assert_called_with("default_ns")
+
+
+# =============================================================================
+# 8. Streaming behavior (memory efficiency)
+# =============================================================================
+
+
+class TestStreamingBehavior:
+    """Tests for memory-efficient streaming writes."""
+
+    def test_processes_blocks_independently(self, sink):
+        """Each block is processed and written separately."""
+        blocks = [pa.table({"id": [i], "vector": [[float(i)]]}) for i in range(5)]
+        write_counts = []
+
+        def track_write(ns, batch, namespace_name=None):
+            write_counts.append(batch.num_rows)
+
+        with patch.object(sink, "_get_client", return_value=MagicMock()):
+            with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
+                sink.write(blocks, ctx=None)
+
+        # 5 blocks → 5 writes of 1 row each
+        assert len(write_counts) == 5
+        assert all(c == 1 for c in write_counts)
+
+
+# =============================================================================
+# 9. Multi-namespace writes
+# =============================================================================
+
+
+class TestMultiNamespaceWrites:
+    """Tests for namespace_column-driven multi-namespace writes."""
+
+    def test_routes_rows_to_correct_namespaces(self):
+        """Rows are grouped by namespace_column and written to the right ns."""
+        sink = make_sink(namespace=None, namespace_column="tenant")
+        table = pa.table(
+            {
+                "tenant": ["ns_a", "ns_b", "ns_a", "ns_b"],
+                "id": [1, 2, 3, 4],
+                "vector": [[0.1], [0.2], [0.3], [0.4]],
+            }
+        )
+
+        writes = {}  # namespace_name -> list of row counts
+
+        def track_write(ns, batch, namespace_name=None):
+            writes.setdefault(namespace_name, []).append(batch.num_rows)
+
+        mock_client = MagicMock()
+        mock_client.namespace.return_value = MagicMock()
+        with patch.object(sink, "_get_client", return_value=mock_client):
+            with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
+                sink.write([table], ctx=None)
+
+        assert "ns_a" in writes
+        assert "ns_b" in writes
+        assert sum(writes["ns_a"]) == 2
+        assert sum(writes["ns_b"]) == 2
+
+    def test_drops_namespace_column_before_writing(self):
+        """The namespace column is not included in the written data."""
+        sink = make_sink(namespace=None, namespace_column="tenant")
+        table = pa.table(
+            {
+                "tenant": ["ns_a"],
+                "id": [1],
+                "vector": [[0.1]],
+            }
+        )
+
+        written_batches = []
+
+        def capture_batch(ns, batch, namespace_name=None):
+            written_batches.append(batch)
+
+        mock_client = MagicMock()
+        mock_client.namespace.return_value = MagicMock()
+        with patch.object(sink, "_get_client", return_value=mock_client):
+            with patch.object(
+                sink, "_write_batch_with_retry", side_effect=capture_batch
+            ):
+                sink.write([table], ctx=None)
+
+        assert len(written_batches) == 1
+        assert "tenant" not in written_batches[0].column_names
+        assert "id" in written_batches[0].column_names
+
+    def test_missing_namespace_column_raises(self):
+        """Missing namespace column in data raises ValueError."""
+        sink = make_sink(namespace=None, namespace_column="tenant")
+        table = pa.table(
+            {
+                "id": [1],
+                "vector": [[0.1]],
+            }
+        )
+
+        mock_client = MagicMock()
+        with patch.object(sink, "_get_client", return_value=mock_client):
+            with pytest.raises(ValueError, match="Namespace column.*not found"):
+                sink.write([table], ctx=None)
+
+    def test_null_namespace_values_raise(self):
+        """Null values in namespace column raise ValueError."""
+        sink = make_sink(namespace=None, namespace_column="tenant")
+        table = pa.table(
+            {
+                "tenant": ["ns_a", None],
+                "id": [1, 2],
+                "vector": [[0.1], [0.2]],
+            }
+        )
+
+        mock_client = MagicMock()
+        with patch.object(sink, "_get_client", return_value=mock_client):
+            with pytest.raises(ValueError, match="contains null values"):
+                sink.write([table], ctx=None)
+
+    def test_skips_empty_blocks_in_multi_namespace(self):
+        """Empty blocks are skipped in multi-namespace mode."""
+        sink = make_sink(namespace=None, namespace_column="tenant")
+        empty_table = pa.table(
+            {
+                "tenant": pa.array([], type=pa.string()),
+                "id": pa.array([], type=pa.int64()),
+                "vector": pa.array([], type=pa.list_(pa.float64())),
+            }
+        )
+
+        mock_client = MagicMock()
+        with patch.object(sink, "_get_client", return_value=mock_client):
+            with patch.object(sink, "_write_batch_with_retry") as mock_write:
+                sink.write([empty_table], ctx=None)
+
+        mock_write.assert_not_called()
+
+
+# =============================================================================
+# 10. Serialization behavior
+# =============================================================================
+
+
+class TestSerialization:
+    """Tests for pickle serialization support."""
+
+    def test_preserves_configuration(self, sink, mock_turbopuffer_module):
+        """Configuration is preserved after pickle round-trip."""
+        pickled = pickle.dumps(sink)
+        unpickled = pickle.loads(pickled)
+
+        assert unpickled.namespace == sink.namespace
+        assert unpickled.namespace_column == sink.namespace_column
+        assert unpickled.api_key == sink.api_key
+        assert unpickled.region == sink.region
+        assert unpickled.base_url == sink.base_url
+        assert unpickled.batch_size == sink.batch_size
+        assert unpickled._client is None
+
+        # Lazy initialization works after unpickling
+        client = unpickled._get_client()
+        assert client is not None
+        mock_turbopuffer_module.Turbopuffer.assert_called()
+
+    def test_preserves_namespace_column_configuration(self, mock_turbopuffer_module):
+        """namespace_column configuration survives pickle round-trip."""
+        sink = make_sink(namespace=None, namespace_column="tenant")
+        pickled = pickle.dumps(sink)
+        unpickled = pickle.loads(pickled)
+
+        assert unpickled.namespace is None
+        assert unpickled.namespace_column == "tenant"
+        assert unpickled._client is None
+
+    def test_preserves_base_url_configuration(self, mock_turbopuffer_module):
+        """base_url configuration survives pickle round-trip."""
+        sink = make_sink(
+            region=None,
+            base_url="https://gcp-us-central1.turbopuffer.com",
+        )
+        pickled = pickle.dumps(sink)
+        unpickled = pickle.loads(pickled)
+
+        assert unpickled.region is None
+        assert unpickled.base_url == "https://gcp-us-central1.turbopuffer.com"
+        assert unpickled._client is None
+
+        # Lazy initialization works and uses base_url
+        unpickled._get_client()
+        mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
+            api_key="test-api-key",
+            base_url="https://gcp-us-central1.turbopuffer.com",
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_video.py

@@ -0,0 +1,35 @@
+import numpy as np
+import pyarrow as pa
+import pytest
+
+import ray
+
+
+def test_read_videos():
+    uri = "s3://anonymous@ray-example-data/basketball.mp4"
+    ds = ray.data.read_videos(uri, include_timestamps=True)
+
+    assert ds.count() == 333
+    assert ds.schema().names == ["frame", "frame_index", "frame_timestamp"]
+
+    frame_indices = ds.select_columns(["frame_index"]).take_all()
+    assert sorted(frame_indices, key=lambda item: item["frame_index"]) == [
+        {"frame_index": i} for i in range(333)
+    ]
+
+    frame_timestamps = ds.select_columns(["frame_timestamp"]).take_all()
+    for t in frame_timestamps:
+        assert isinstance(t["frame_timestamp"], np.ndarray)
+        assert t["frame_timestamp"].shape[0] == 2
+
+    frame_type, frame_index_type, _ = ds.schema().types
+
+    assert frame_type.shape == (720, 1280, 3)
+    assert frame_type.value_type == pa.uint8()
+    assert frame_index_type == pa.int64()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_webdataset.py

@@ -0,0 +1,335 @@
+# Copyright NVIDIA Corporation 2023
+# SPDX-License-Identifier: Apache-2.0
+
+import glob
+import io
+import os
+import pickle
+import tarfile
+
+import pytest
+import webdataset as wds
+
+import ray
+from ray.tests.conftest import *  # noqa
+
+
+class TarWriter:
+    def __init__(self, path):
+        self.path = path
+        self.tar = tarfile.open(path, "w")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.tar.close()
+
+    def write(self, name, data):
+        f = self.tar.tarinfo()
+        f.name = name
+        f.size = len(data)
+        self.tar.addfile(f, io.BytesIO(data))
+
+
+def test_webdataset_read(ray_start_2_cpus, tmp_path):
+    path = os.path.join(tmp_path, "bar_000000.tar")
+    with TarWriter(path) as tf:
+        for i in range(100):
+            tf.write(f"{i}.a", str(i).encode("utf-8"))
+            tf.write(f"{i}.b", str(i**2).encode("utf-8"))
+    assert os.path.exists(path)
+    assert len(glob.glob(f"{tmp_path}/*.tar")) == 1
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
+    samples = ds.take(100)
+    assert len(samples) == 100
+    for i, sample in enumerate(samples):
+        assert isinstance(sample, dict), sample
+        assert sample["__key__"] == str(i)
+        assert sample["a"].decode("utf-8") == str(i)
+        assert sample["b"].decode("utf-8") == str(i**2)
+
+
+@pytest.fixture
+def allow_unsafe_deserialization(monkeypatch):
+    monkeypatch.setenv("RAY_DATA_WEBDATASET_ALLOW_UNSAFE_DESERIALIZATION", "1")
+
+
+def test_webdataset_expand_json(
+    ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
+):
+    import numpy as np
+    import torch
+
+    image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+    gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
+    dstruct = dict(a=[1, 2], b=dict(c=2), d="hello")
+    ttensor = torch.tensor([1, 2, 3]).numpy()
+
+    sample = {
+        "__key__": "foo",
+        "jpg": image,
+        "gray.png": gray,
+        "mp": dstruct,
+        "json": dstruct,
+        "pt": ttensor,
+        "und": b"undecoded",
+        "custom": b"nothing",
+    }
+
+    # write the encoded data using the default encoder
+    data = [sample]
+    ds = ray.data.from_items(data).repartition(1)
+    ds.write_webdataset(path=tmp_path, try_create_dir=True)
+    ds = ray.data.read_webdataset(
+        paths=[str(tmp_path)], override_num_blocks=1, expand_json=True
+    )
+    record = ds.take(1)
+    assert [1, 2] == record[0]["a"]
+
+
+def test_webdataset_suffixes(ray_start_2_cpus, tmp_path):
+    path = os.path.join(tmp_path, "bar_000000.tar")
+    with TarWriter(path) as tf:
+        for i in range(100):
+            tf.write(f"{i}.txt", str(i).encode("utf-8"))
+            tf.write(f"{i}.test.txt", str(i**2).encode("utf-8"))
+            tf.write(f"{i}.cls", str(i**2).encode("utf-8"))
+            tf.write(f"{i}.test.cls2", str(i**2).encode("utf-8"))
+    assert os.path.exists(path)
+    assert len(glob.glob(f"{tmp_path}/*.tar")) == 1
+
+    # test simple suffixes
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["txt", "cls"])
+    samples = ds.take(100)
+    assert len(samples) == 100
+    for i, sample in enumerate(samples):
+        assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls"}
+
+    # test fnmatch patterns for suffixes
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["*.txt", "*.cls"])
+    samples = ds.take(100)
+    assert len(samples) == 100
+    for i, sample in enumerate(samples):
+        assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls", "test.txt"}
+
+    # test selection function
+    def select(name):
+        return name.endswith("txt")
+
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=select)
+    samples = ds.take(100)
+    assert len(samples) == 100
+    for i, sample in enumerate(samples):
+        assert set(sample.keys()) == {"__url__", "__key__", "txt", "test.txt"}
+
+    # test filerename
+    def renamer(name):
+        result = name.replace("txt", "text")
+        print("***", name, result)
+        return result
+
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)], filerename=renamer)
+    samples = ds.take(100)
+    assert len(samples) == 100
+    for i, sample in enumerate(samples):
+        assert set(sample.keys()) == {
+            "__url__",
+            "__key__",
+            "text",
+            "cls",
+            "test.text",
+            "test.cls2",
+        }
+
+
+def test_webdataset_write(ray_start_2_cpus, tmp_path):
+    print(ray.available_resources())
+    data = [dict(__key__=str(i), a=str(i), b=str(i**2)) for i in range(100)]
+    ds = ray.data.from_items(data).repartition(1)
+    ds.write_webdataset(path=tmp_path, try_create_dir=True)
+    paths = glob.glob(f"{tmp_path}/*.tar")
+    assert len(paths) == 1
+    with open(paths[0], "rb") as stream:
+        tf = tarfile.open(fileobj=stream)
+        for i in range(100):
+            assert tf.extractfile(f"{i}.a").read().decode("utf-8") == str(i)
+            assert tf.extractfile(f"{i}.b").read().decode("utf-8") == str(i**2)
+
+
+def custom_decoder(sample):
+    for key, value in sample.items():
+        if key == "png":
+            # check that images have already been decoded
+            assert not isinstance(value, bytes)
+        elif key.endswith("custom"):
+            sample[key] = "custom-value"
+    return sample
+
+
+def test_webdataset_coding(ray_start_2_cpus, tmp_path, allow_unsafe_deserialization):
+    import numpy as np
+    import PIL.Image
+    import torch
+
+    image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+    gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
+    dstruct = dict(a=[1], b=dict(c=2), d="hello")
+    ttensor = torch.tensor([1, 2, 3]).numpy()
+
+    sample = {
+        "__key__": "foo",
+        "jpg": image,
+        "gray.png": gray,
+        "mp": dstruct,
+        "json": dstruct,
+        "pt": ttensor,
+        "und": b"undecoded",
+        "custom": b"nothing",
+    }
+
+    # write the encoded data using the default encoder
+    data = [sample]
+    ds = ray.data.from_items(data).repartition(1)
+    ds.write_webdataset(path=tmp_path, try_create_dir=True)
+
+    # read the encoded data using the default decoder
+    paths = glob.glob(f"{tmp_path}/*.tar")
+    assert len(paths) == 1
+    path = paths[0]
+    assert os.path.exists(path)
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
+    samples = ds.take(1)
+    assert len(samples) == 1
+    for sample in samples:
+        assert isinstance(sample, dict), sample
+        assert sample["__key__"] == "foo"
+        assert isinstance(sample["jpg"], np.ndarray)
+        assert sample["jpg"].shape == (100, 100, 3)
+        assert isinstance(sample["gray.png"], np.ndarray)
+        assert sample["gray.png"].shape == (100, 100)
+        assert isinstance(sample["mp"], dict)
+        assert sample["mp"]["a"] == [1]
+        assert sample["mp"]["b"]["c"] == 2
+        assert isinstance(sample["json"], dict)
+        assert sample["json"]["a"] == [1]
+        assert isinstance(sample["pt"], np.ndarray)
+        assert sample["pt"].tolist() == [1, 2, 3]
+
+    # test the format argument to the default decoder and multiple decoders
+    ds = ray.data.read_webdataset(
+        paths=[str(tmp_path)], decoder=["PIL", custom_decoder]
+    )
+    samples = ds.take(1)
+    assert len(samples) == 1
+    for sample in samples:
+        assert isinstance(sample, dict), sample
+        assert sample["__key__"] == "foo"
+        assert isinstance(sample["jpg"], PIL.Image.Image)
+        assert isinstance(sample["gray.png"], PIL.Image.Image)
+        assert isinstance(sample["und"], bytes)
+        assert sample["und"] == b"undecoded"
+        assert sample["custom"] == "custom-value"
+
+
+def test_webdataset_decoding(ray_start_2_cpus, tmp_path):
+    import numpy as np
+    import torch
+
+    image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+    gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
+    dstruct = dict(a=np.nan, b=dict(c=2), d="hello", e={"img_filename": "for_test.jpg"})
+    ttensor = torch.tensor([1, 2, 3]).numpy()
+
+    sample = {
+        "__key__": "foo",
+        "jpg": image,
+        "gray.png": gray,
+        "mp": dstruct,
+        "json": dstruct,
+        "pt": ttensor,
+        "und": b"undecoded",
+        "custom": b"nothing",
+    }
+
+    # write the encoded data using the default encoder
+    data = [sample]
+    ds = ray.data.from_items(data).repartition(1)
+    ds.write_webdataset(path=tmp_path, try_create_dir=True)
+
+    ds = ray.data.read_webdataset(
+        paths=[str(tmp_path)],
+        override_num_blocks=1,
+        decoder=None,
+    )
+    samples = ds.take(1)
+    import json
+
+    meta_json = json.loads(samples[0]["json"].decode("utf-8"))
+    assert meta_json["e"]["img_filename"] == "for_test.jpg"
+
+
+@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
+def test_write_min_rows_per_file(tmp_path, ray_start_2_cpus, min_rows_per_file):
+    ray.data.from_items(
+        [{"id": str(i)} for i in range(100)], override_num_blocks=20
+    ).write_webdataset(tmp_path, min_rows_per_file=min_rows_per_file)
+
+    for filename in os.listdir(tmp_path):
+        dataset = wds.WebDataset(os.path.join(tmp_path, filename))
+        assert len(list(dataset)) == min_rows_per_file
+
+
+@pytest.mark.parametrize(
+    "filename",
+    ["000000.pkl", "000000.pickle", "000000.pt", "000000.pth"],
+)
+def test_default_decoder_rejects_unsafe_extensions(
+    ray_start_2_cpus, tmp_path, filename
+):
+    path = os.path.join(tmp_path, "unsafe.tar")
+    with TarWriter(path) as tf:
+        tf.write(filename, b"fake-payload")
+
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
+    with pytest.raises(Exception, match="Refusing to"):
+        ds.take_all()
+
+
+def test_default_decoder_allows_unsafe_with_env_var(
+    ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
+):
+    path = os.path.join(tmp_path, "trusted.tar")
+    with TarWriter(path) as tf:
+        tf.write("000000.pkl", pickle.dumps({"key": "value"}))
+
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
+    rows = ds.take_all()
+
+    assert len(rows) == 1
+    assert rows[0]["pkl"] == {"key": "value"}
+
+
+def test_custom_decoder_bypasses_unsafe_guard(ray_start_2_cpus, tmp_path):
+    path = os.path.join(tmp_path, "custom.tar")
+    with TarWriter(path) as tf:
+        tf.write("000000.pkl", pickle.dumps({"key": "value"}))
+
+    def safe_pkl_decoder(sample):
+        sample = dict(sample)
+        for key, value in sample.items():
+            if key == "pkl":
+                sample[key] = pickle.loads(value)
+        return sample
+
+    ds = ray.data.read_webdataset(paths=[str(tmp_path)], decoder=safe_pkl_decoder)
+    rows = ds.take_all()
+
+    assert len(rows) == 1
+    assert rows[0]["pkl"] == {"key": "value"}
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/datasource/test_zarrv2.py

@@ -0,0 +1,937 @@
+import logging
+import os
+from pathlib import Path
+from typing import Any
+
+import fsspec
+import numpy as np
+import pandas as pd
+import pyarrow.fs
+import pytest
+import zarr
+from pytest_lazy_fixtures import lf as lazy_fixture
+
+import ray
+from ray.data._internal.datasource import zarrv2_datasource
+from ray.data.block import BlockAccessor
+from ray.data.tests.conftest import *  # noqa: F401, F403
+
+
+def _execute_read_tasks(tasks) -> pd.DataFrame:
+    frames = [
+        BlockAccessor.for_block(block).to_pandas() for task in tasks for block in task()
+    ]
+    return pd.concat(frames, ignore_index=True)
+
+
+def _reconstruct_array(df: pd.DataFrame, array_name: str) -> np.ndarray:
+    """Concatenate all chunks of one array from a long-form result frame."""
+    sub = df[df["array"] == array_name].sort_values(
+        "chunk_index", key=lambda col: col.map(tuple)
+    )
+    return np.concatenate(list(sub["chunk"]), axis=0)
+
+
+def _write_real_zarr_store(
+    store_path: Path,
+    arrays: dict,  # {name: (data, chunks)}
+) -> Path:
+    """Write a real Zarr v2 store from numpy arrays and consolidate metadata."""
+    root = zarr.open_group(str(store_path), mode="w")
+    for name, (data, chunks) in arrays.items():
+        root.create_dataset(name, data=data, chunks=chunks, dtype=data.dtype)
+    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
+    return store_path
+
+
+@pytest.fixture
+def zarrv2_group_store(tmp_path) -> Path:
+    """Two arrays at the store root, both 2-D and 1-D, axis-0-aligned (shape[0]==5)."""
+    return _write_real_zarr_store(
+        tmp_path / "group.zarr",
+        {
+            "images": (np.arange(20, dtype="<i4").reshape(5, 4), (2, 4)),
+            "nested": (np.arange(5, dtype="|u1"), (2,)),
+        },
+    )
+
+
+@pytest.fixture
+def zarrv2_root_store(tmp_path) -> Path:
+    """Single-array store with the array sitting directly at the store root."""
+    store_path = tmp_path / "root.zarr"
+    arr = zarr.open(
+        str(store_path),
+        mode="w",
+        shape=(5, 4),
+        chunks=(2, 4),
+        dtype="<i4",
+    )
+    arr[:] = np.arange(20, dtype="<i4").reshape(5, 4)
+    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
+    return store_path
+
+
+@pytest.fixture
+def local_fsspec_fs():
+    """fsspec local filesystem (for parametrized cross-fs read tests)."""
+    return fsspec.filesystem("file")
+
+
+@pytest.fixture
+def heterogeneous_zarrv2_store(tmp_path) -> Path:
+    """A store mixing different ranks, shape[0]s, dtypes, and native chunk sizes.
+
+    Mirrors the UMI-style real-world layout where ``data/*`` arrays share an
+    axis-0 timestep count but differ in everything else, and ``meta/*``
+    arrays live in a separate axis-0 universe entirely. The chunk-per-row
+    datasource handles all of these in one read; nothing has to align.
+    """
+    store_path = tmp_path / "heterogeneous.zarr"
+    root = zarr.open_group(str(store_path), mode="w")
+    # 4-D image tensor with tiny axis-0 chunks (1 image per chunk).
+    root.create_dataset(
+        "data/camera0_rgb",
+        data=np.arange(20 * 2 * 2 * 3, dtype="|u1").reshape(20, 2, 2, 3),
+        chunks=(1, 2, 2, 3),
+    )
+    # 2-D pose array, same shape[0]=20, much larger axis-0 chunks (10).
+    root.create_dataset(
+        "data/robot0_eef_pos",
+        data=np.arange(20 * 3, dtype="<f4").reshape(20, 3),
+        chunks=(10, 3),
+    )
+    # Episode-boundary metadata: separate axis-0 universe.
+    root.create_dataset(
+        "meta/episode_ends",
+        data=np.array([5, 12, 20], dtype="<i8"),
+        chunks=(3,),
+    )
+    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
+    return store_path
+
+
+@pytest.fixture
+def unconsolidated_zarrv2_store(tmp_path) -> Path:
+    """Two arrays at the store root, no ``.zmetadata``.
+
+    Exercises the no-``.zmetadata`` code paths (per-array ``.zarray``
+    discovery and full-store walk) — the common shape of real-world stores
+    behind plain HTTPS or other listing-less filesystems.
+    """
+    store_path = tmp_path / "unconsolidated.zarr"
+    root = zarr.open_group(str(store_path), mode="w")
+    root.create_dataset(
+        "images", data=np.arange(20, dtype="<i4").reshape(5, 4), chunks=(2, 4)
+    )
+    root.create_dataset("nested", data=np.arange(5, dtype="|u1"), chunks=(2,))
+    return store_path
+
+
+@pytest.fixture
+def aligned_zarrv2_store(tmp_path) -> Path:
+    """Three arrays sharing ``shape[0]=8``, different ranks and native chunks.
+
+    Models the UMI-style case where data arrays co-stride on the timestep
+    axis but differ in everything else.
+    """
+    store_path = tmp_path / "aligned.zarr"
+    root = zarr.open_group(str(store_path), mode="w")
+    root.create_dataset(
+        "img",
+        data=np.arange(8 * 4 * 4 * 3, dtype="|u1").reshape(8, 4, 4, 3),
+        chunks=(2, 4, 4, 3),
+    )
+    root.create_dataset(
+        "state",
+        data=np.arange(8 * 3, dtype="<f4").reshape(8, 3),
+        chunks=(4, 3),  # different native axis-0 chunks than img
+    )
+    root.create_dataset(
+        "label",
+        data=np.arange(8, dtype="<i8"),
+        chunks=(8,),
+    )
+    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
+    return store_path
+
+
+@pytest.fixture
+def zarr_zip_store(tmp_path) -> Path:
+    """A small Zarr store packed into a ``.zip`` for URL-detection tests."""
+    src = tmp_path / "src.zarr"
+    _write_real_zarr_store(
+        src,
+        {
+            "data": (np.arange(12, dtype="<i4").reshape(6, 2), (3, 2)),
+        },
+    )
+    zip_path = tmp_path / "store.zarr.zip"
+    import shutil
+
+    shutil.make_archive(
+        base_name=str(tmp_path / "store.zarr"),
+        format="zip",
+        root_dir=str(src),
+    )
+    assert zip_path.exists()
+    return zip_path
+
+
+# ---------------------------------------------------------------------------
+# Metadata discovery
+# ---------------------------------------------------------------------------
+
+
+def test_normalizes_requested_root_array_path(zarrv2_root_store):
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(zarrv2_root_store),
+        array_paths=[""],
+    )
+    assert list(datasource._metadata_by_path) == [""]
+
+
+def test_normalizes_requested_array_paths(zarrv2_group_store):
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(zarrv2_group_store),
+        array_paths=["images/", "nested"],
+    )
+    assert list(datasource._metadata_by_path) == ["images", "nested"]
+
+
+def test_rejects_missing_array_paths(zarrv2_group_store):
+    with pytest.raises(
+        ValueError,
+        match=r"Array\(s\) not found: 'missing'\. Available: 'images', 'nested'",
+    ):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(zarrv2_group_store),
+            array_paths=["missing"],
+        )
+
+
+def test_loads_per_array_zarray_without_zmetadata(unconsolidated_zarrv2_store):
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(unconsolidated_zarrv2_store),
+        array_paths=["images", "nested"],
+    )
+    assert set(datasource._metadata_by_path) == {"images", "nested"}
+
+
+def test_full_scan_discovers_arrays_without_zmetadata(unconsolidated_zarrv2_store):
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(unconsolidated_zarrv2_store),
+        allow_full_metadata_scan=True,
+    )
+    assert set(datasource._metadata_by_path) == {"images", "nested"}
+
+
+def test_requires_array_paths_or_full_scan_when_unconsolidated(
+    unconsolidated_zarrv2_store,
+):
+    with pytest.raises(
+        ValueError,
+        match=(
+            r"No array_paths were provided and this Zarr store does not "
+            r"contain \.zmetadata"
+        ),
+    ):
+        zarrv2_datasource.ZarrV2Datasource(str(unconsolidated_zarrv2_store))
+
+
+def test_array_paths_missing_zarray_file_raises_value_error(
+    unconsolidated_zarrv2_store,
+):
+    with pytest.raises(
+        ValueError,
+        match=r"Array path 'missing' not found",
+    ):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(unconsolidated_zarrv2_store),
+            array_paths=["missing"],
+        )
+
+
+def test_local_scheme_pins_reads_to_driver_node(zarrv2_group_store):
+    """``local://`` stores can't be read distributed; plain/cloud paths can."""
+    local = zarrv2_datasource.ZarrV2Datasource("local://" + str(zarrv2_group_store))
+    assert local.supports_distributed_reads is False
+
+    plain = zarrv2_datasource.ZarrV2Datasource(str(zarrv2_group_store))
+    assert plain.supports_distributed_reads is True
+
+
+def test_consolidation_detected_via_open_consolidated(
+    zarrv2_group_store, unconsolidated_zarrv2_store
+):
+    """``_consolidated`` reflects whether ``.zmetadata`` actually opened."""
+    consolidated = zarrv2_datasource.ZarrV2Datasource(
+        str(zarrv2_group_store), array_paths=["images"]
+    )
+    assert consolidated._consolidated is True
+
+    unconsolidated = zarrv2_datasource.ZarrV2Datasource(
+        str(unconsolidated_zarrv2_store), array_paths=["images"]
+    )
+    assert unconsolidated._consolidated is False
+
+
+def test_array_paths_rejects_group_path(tmp_path):
+    """Requesting a group path (not an array) on an unconsolidated store errors."""
+    store_path = tmp_path / "withgroup.zarr"
+    root = zarr.open_group(str(store_path), mode="w")
+    grp = root.create_group("grp")
+    grp.create_dataset("inner", data=np.arange(4, dtype="<i4"), chunks=(2,))
+    # Not consolidated -> the per-array ``.zarray`` lookup path.
+    with pytest.raises(ValueError, match="is a group, not an array"):
+        zarrv2_datasource.ZarrV2Datasource(str(store_path), array_paths=["grp"])
+
+
+def test_root_array_rejects_non_root_array_paths(zarrv2_root_store):
+    """A single root-level array rejects array_paths that aren't the root ''."""
+    with pytest.raises(ValueError, match="single root-level array"):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(zarrv2_root_store), array_paths=["missing"]
+        )
+
+
+# ---------------------------------------------------------------------------
+# chunk_shapes validation
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "chunk_shapes, match",
+    [
+        ("invalid", "positive integers"),
+        ({"images": 1}, "positive integers"),
+        ({"does_not_exist": [2]}, "Unknown array path"),
+    ],
+)
+def test_rejects_invalid_chunk_shapes(zarrv2_group_store, chunk_shapes, match):
+    with pytest.raises(ValueError, match=match):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(zarrv2_group_store), chunk_shapes=chunk_shapes
+        )
+
+
+@pytest.mark.parametrize(
+    "chunk_shapes,array_paths,expected",
+    [
+        # No chunk_shapes: every array reads at its native chunk size.
+        # 4-D image with tiny chunks coexists with 2-D pose with big chunks —
+        # nothing is forced into a shared min/max.
+        (
+            None,
+            None,
+            {
+                "data/camera0_rgb": (1, 2, 2, 3),
+                "data/robot0_eef_pos": (10, 3),
+                "meta/episode_ends": (3,),
+            },
+        ),
+        # ``[5]`` prefix overrides axis 0 across arrays of all ranks at once.
+        (
+            [5],
+            None,
+            {
+                "data/camera0_rgb": (5, 2, 2, 3),
+                "data/robot0_eef_pos": (5, 3),
+                "meta/episode_ends": (5,),
+            },
+        ),
+        # Length-2 prefix overrides axes 0+1; needs every selected array to
+        # have rank >= 2, so we filter out ``meta/episode_ends`` (rank 1).
+        (
+            [5, 1],
+            ["data/camera0_rgb", "data/robot0_eef_pos"],
+            {
+                "data/camera0_rgb": (5, 1, 2, 3),
+                "data/robot0_eef_pos": (5, 1),
+            },
+        ),
+        # Per-array overrides may retile only some arrays while others keep
+        # their native chunks.
+        (
+            {
+                "data/camera0_rgb": [5],
+                "data/robot0_eef_pos": [7],
+            },
+            None,
+            {
+                "data/camera0_rgb": (5, 2, 2, 3),
+                "data/robot0_eef_pos": (7, 3),
+                "meta/episode_ends": (3,),
+            },
+        ),
+    ],
+)
+def test_chunk_shapes_resolution_across_mixed_rank(
+    heterogeneous_zarrv2_store, chunk_shapes, array_paths, expected
+):
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(heterogeneous_zarrv2_store),
+        chunk_shapes=chunk_shapes,
+        array_paths=array_paths,
+    )
+    assert datasource._array_chunks == expected
+
+
+# ---------------------------------------------------------------------------
+# align_axis_0 (wide-form mode)
+# ---------------------------------------------------------------------------
+
+
+def test_align_axis_0_emits_wide_rows(ray_start_regular_shared, aligned_zarrv2_store):
+    """Wide-row schema: ``t_start``, ``t_stop``, one column per selected array."""
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(aligned_zarrv2_store),
+        align_axis_0=True,
+        chunk_shapes=[4],
+    )
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
+    assert set(df.columns) == {"t_start", "t_stop", "img", "state", "label"}
+    # shape[0]=8, chunk_shapes=[4] -> 2 rows.
+    assert len(df) == 2
+    # Reconstruct each array by concatenating slices in order.
+    img_recon = np.concatenate(list(df["img"]), axis=0)
+    assert img_recon.shape == (8, 4, 4, 3)
+    state_recon = np.concatenate(list(df["state"]), axis=0)
+    assert state_recon.shape == (8, 3)
+    label_recon = np.concatenate(list(df["label"]), axis=0)
+    assert label_recon.shape == (8,)
+    # t_start/t_stop are correct.
+    starts = sorted(df["t_start"].tolist())
+    stops = sorted(df["t_stop"].tolist())
+    assert starts == [0, 4]
+    assert stops == [4, 8]
+
+
+def test_align_axis_0_column_set(ray_start_regular_shared, aligned_zarrv2_store):
+    """array_paths selects which arrays are read; aligned mode emits one column
+    per selected array (plus t_start/t_stop)."""
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(aligned_zarrv2_store),
+        array_paths=["img", "state"],
+        align_axis_0=True,
+        chunk_shapes=[4],
+    )
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
+    assert set(df.columns) == {"t_start", "t_stop", "img", "state"}
+
+
+def test_align_axis_0_rejects_misaligned_shape0(heterogeneous_zarrv2_store):
+    """Misalignment raises with the per-array shape[0] breakdown."""
+    with pytest.raises(
+        ValueError,
+        match=r"All selected arrays must share shape\[0\]",
+    ):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(heterogeneous_zarrv2_store),
+            align_axis_0=True,
+            chunk_shapes=[5],
+        )
+
+
+def test_align_axis_0_rejects_divergent_axis_0_chunks(aligned_zarrv2_store):
+    """If aligned arrays end up with different axis-0 chunks, error clearly.
+
+    The native chunks differ (img=2, state=4, label=8) — without a
+    ``chunk_shapes`` re-tile they all stay at native, and the validator
+    catches the mismatch.
+    """
+    with pytest.raises(
+        ValueError, match="Aligned arrays must share the same axis-0 chunk size"
+    ):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(aligned_zarrv2_store),
+            align_axis_0=True,
+        )
+
+
+# ---------------------------------------------------------------------------
+# overlap (aligned-mode lookahead)
+# ---------------------------------------------------------------------------
+
+
+def test_overlap_extends_chunk_data(ray_start_regular_shared, aligned_zarrv2_store):
+    """``overlap=N`` makes each row's per-array slice cover ``N`` extra timesteps.
+
+    Aligned store has shape[0]=8, ``chunk_shapes=[4]`` -> rows own [0,4) and [4,8).
+    With ``overlap=2``, row 0's data covers [0,6) and row 1's data covers [4,8)
+    (clipped at the store end since 4+4+2 > 8).
+    """
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(aligned_zarrv2_store),
+        align_axis_0=True,
+        chunk_shapes=[4],
+        overlap=2,
+    )
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
+    # Ownership unchanged: 2 rows of width 4 each.
+    assert sorted(zip(df["t_start"], df["t_stop"])) == [(0, 4), (4, 8)]
+    # Data extents: row 0 has 6 timesteps, row 1 has 4 (clipped at shape[0]=8).
+    rows = sorted(df.to_dict("records"), key=lambda r: r["t_start"])
+    assert rows[0]["img"].shape[0] == 6  # 4 owned + 2 overlap
+    assert rows[0]["state"].shape[0] == 6
+    assert rows[1]["img"].shape[0] == 4  # 4 owned + 0 overlap (clipped)
+    assert rows[1]["state"].shape[0] == 4
+
+
+def test_overlap_requires_align_axis_0(aligned_zarrv2_store):
+    """``overlap`` in long-form (no ``align_axis_0``) is a clear error."""
+    with pytest.raises(ValueError, match="overlap requires align_axis_0=True"):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(aligned_zarrv2_store),
+            overlap=2,
+        )
+
+
+def test_overlap_rejects_negative_and_non_int(aligned_zarrv2_store):
+    bad_values: list[Any] = [-1, 1.5, "two"]
+
+    for bad in bad_values:
+        with pytest.raises(ValueError, match="overlap must be a non-negative integer"):
+            zarrv2_datasource.ZarrV2Datasource(
+                str(aligned_zarrv2_store),
+                align_axis_0=True,
+                chunk_shapes=[4],
+                overlap=bad,
+            )
+
+
+def test_chunk_shapes_rejected_when_longer_than_smallest_array(
+    heterogeneous_zarrv2_store,
+):
+    """A shared ``chunk_shapes`` override longer than a target rank is an error."""
+    with pytest.raises(
+        ValueError,
+        match=r"chunk_shapes override for array .* has 2 axes but array of shape .* has rank 1",
+    ):
+        zarrv2_datasource.ZarrV2Datasource(
+            str(heterogeneous_zarrv2_store),
+            chunk_shapes=[2, 2],  # OK for 2-D and 4-D, fails for 1-D episode_ends
+        )
+
+
+# ---------------------------------------------------------------------------
+# Filesystem handling
+# ---------------------------------------------------------------------------
+
+
+def test_accepts_pyarrow_fs_filesystem(zarrv2_group_store):
+    """A pyarrow.fs.FileSystem passed in is wrapped into fsspec internally."""
+    datasource = zarrv2_datasource.ZarrV2Datasource(
+        str(zarrv2_group_store),
+        filesystem=pyarrow.fs.LocalFileSystem(),
+    )
+    from fsspec.spec import AbstractFileSystem
+
+    assert isinstance(datasource._fs, AbstractFileSystem)
+    assert set(datasource._metadata_by_path) == {"images", "nested"}
+
+
+def test_rejects_unsupported_filesystem_type():
+    """Filesystem that's neither pyarrow.fs nor fsspec raises ``TypeError``."""
+    with pytest.raises(
+        TypeError,
+        match=r"filesystem must be pyarrow\.fs\.FileSystem or",
+    ):
+        zarrv2_datasource.ZarrV2Datasource(
+            "/tmp/some.zarr",
+            filesystem="not-a-filesystem",
+        )
+
+
+# ---------------------------------------------------------------------------
+# .zarr.zip URL support
+# ---------------------------------------------------------------------------
+
+
+def test_reads_zarr_zip_local_path(ray_start_regular_shared, zarr_zip_store):
+    """A local ``.zarr.zip`` path auto-wires fsspec's ZipFileSystem."""
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store))
+    # The store has one array "data" of shape (6, 2) chunks (3, 2) -> 2 chunks.
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
+    assert len(df) == 2
+    assert set(df["array"]) == {"data"}
+    recon = _reconstruct_array(df, "data")
+    np.testing.assert_array_equal(recon, np.arange(12, dtype="<i4").reshape(6, 2))
+
+
+# ---------------------------------------------------------------------------
+# Read task generation and execution (end-to-end)
+# ---------------------------------------------------------------------------
+
+
+def test_get_read_tasks_batches_chunks_by_parallelism(tmp_path):
+    """5 chunks split across parallelism=3 produces batches [2, 2, 1]."""
+    store_path = tmp_path / "store.zarr"
+    _write_real_zarr_store(
+        store_path,
+        {"images": (np.arange(5 * 4, dtype="<i4").reshape(5, 4), (1, 4))},
+    )
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
+
+    read_tasks = datasource.get_read_tasks(parallelism=3)
+
+    assert len(read_tasks) == 3
+    assert [task.metadata.num_rows for task in read_tasks] == [2, 2, 1]
+    assert all(task.metadata.input_files == (str(store_path),) for task in read_tasks)
+
+
+def test_long_form_chunk_index_order_matches_grid(ray_start_regular_shared, tmp_path):
+    """Lazy grid-range tasks emit chunk_index in the same row-major order as a
+    full grid enumeration (regression guard for the lazy-unravel refactor)."""
+    from itertools import product
+
+    store_path = tmp_path / "order.zarr"
+    # shape (6, 4), chunks (2, 2) -> grid (3, 2) = 6 chunks.
+    _write_real_zarr_store(
+        store_path, {"a": (np.arange(6 * 4, dtype="<i4").reshape(6, 4), (2, 2))}
+    )
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
+    # parallelism=2 -> two flat-index ranges; concatenated they must be in order.
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
+    got = [tuple(int(x) for x in ci) for ci in df["chunk_index"]]
+    assert got == list(product(range(3), range(2)))
+
+
+def test_per_task_row_limit_caps_chunks_read(
+    ray_start_regular_shared, tmp_path, monkeypatch
+):
+    """per_task_row_limit bounds how many chunks a task actually reads, so a
+    downstream ``limit`` doesn't pull the whole batch's I/O."""
+    store_path = tmp_path / "limit.zarr"
+    _write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (1,))})
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
+
+    reads = []
+    real_read_chunk = zarrv2_datasource._read_chunk
+
+    def _spy(*args, **kwargs):
+        reads.append(1)
+        return real_read_chunk(*args, **kwargs)
+
+    monkeypatch.setattr(zarrv2_datasource, "_read_chunk", _spy)
+
+    # parallelism=1 -> one task batching all 10 chunks; cap it at 3.
+    tasks = datasource.get_read_tasks(parallelism=1, per_task_row_limit=3)
+    blocks = [block for task in tasks for block in task()]
+
+    total_rows = sum(BlockAccessor.for_block(b).num_rows() for b in blocks)
+    assert total_rows == 3
+    # The fix: only 3 chunks were actually read (not all 10, then truncated).
+    assert len(reads) == 3
+
+
+def test_read_chunk_retries_transient_io(monkeypatch):
+    """_read_chunk retries reads whose error matches retry_match (Ray Data's
+    DataContext.retried_io_errors), then succeeds."""
+    monkeypatch.setattr("time.sleep", lambda *_: None)  # no backoff in the test
+
+    class _FlakyArray:
+        attempts = 0
+
+        def __getitem__(self, _idx):
+            type(self).attempts += 1
+            if self.attempts < 3:
+                raise OSError("Connection reset by peer")
+            return np.arange(4, dtype="<i4")
+
+    class _Root:
+        def __getitem__(self, _name):
+            return _FlakyArray()
+
+    out = zarrv2_datasource._read_chunk(
+        _Root(),  # pyrefly: ignore[bad-argument-type]
+        "x",
+        ((0, 4),),
+        retry_match=["Connection reset"],
+    )
+    np.testing.assert_array_equal(out, np.arange(4, dtype="<i4"))
+    assert _FlakyArray.attempts == 3  # failed twice, then succeeded
+
+
+def test_long_form_schema_and_materialization(ray_start_regular_shared, tmp_path):
+    """End-to-end: long-form rows are emitted with the expected columns and data."""
+    store_path = tmp_path / "aligned.zarr"
+    images_src = np.arange(20, dtype="<i4").reshape(5, 4)
+    labels_src = np.arange(5, dtype="|u1")
+    _write_real_zarr_store(
+        store_path,
+        {
+            "images": (images_src, (2, 4)),
+            "labels": (labels_src, (2,)),
+        },
+    )
+
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
+
+    # Schema is the long-form quad.
+    assert list(df.columns) == ["array", "chunk_index", "chunk_slices", "chunk"]
+    # 3 chunks for images (5/2), 3 chunks for labels (5/2) = 6 rows total.
+    assert len(df) == 6
+    assert set(df["array"]) == {"images", "labels"}
+
+    np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
+    np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)
+
+    # ``chunk_slices`` matches the actual chunk shape and indexes back to
+    # the source array: arr[start:stop, ...] equals the chunk.
+    for _, row in df.iterrows():
+        slices = row["chunk_slices"]
+        chunk = row["chunk"]
+        assert len(slices) == chunk.ndim
+        for axis, (start, stop) in enumerate(slices):
+            assert stop - start == chunk.shape[axis]
+        if row["array"] == "images":
+            np.testing.assert_array_equal(
+                chunk,
+                images_src[slices[0][0] : slices[0][1], slices[1][0] : slices[1][1]],
+            )
+
+
+def test_chunk_shapes_override_changes_grid(ray_start_regular_shared, tmp_path):
+    """User-supplied chunk_shapes controls the chunk grid and row count."""
+    store_path = tmp_path / "tile.zarr"
+    src = np.arange(10, dtype="<i4")
+    _write_real_zarr_store(store_path, {"data": (src, (2,))})  # native: 5 chunks
+
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path), chunk_shapes=[5])
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
+    assert sorted(chunk.shape[0] for chunk in df["chunk"]) == [5, 5]
+
+
+def test_heterogeneous_store_emits_one_row_per_chunk(
+    ray_start_regular_shared, heterogeneous_zarrv2_store
+):
+    """Mixed-rank/shape/dtype arrays each contribute their chunk count to the output."""
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(heterogeneous_zarrv2_store))
+    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
+
+    # Expected chunk counts:
+    #   data/camera0_rgb       shape=(20,2,2,3) chunks=(1,2,2,3) → 20 chunks
+    #   data/robot0_eef_pos    shape=(20,3)     chunks=(10,3)    → 2 chunks
+    #   meta/episode_ends      shape=(3,)       chunks=(3,)      → 1 chunk
+    counts = df.groupby("array").size().to_dict()
+    assert counts == {
+        "data/camera0_rgb": 20,
+        "data/robot0_eef_pos": 2,
+        "meta/episode_ends": 1,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Estimator
+# ---------------------------------------------------------------------------
+
+
+def test_estimate_inmemory_data_size(tmp_path):
+    """Estimate = sum over arrays of numel * dtype.itemsize."""
+    store_path = tmp_path / "est.zarr"
+    _write_real_zarr_store(
+        store_path,
+        {
+            "a": (np.zeros((5, 4), dtype="<i4"), (2, 4)),
+            "b": (np.zeros(5, dtype="|u1"), (2,)),
+        },
+    )
+    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
+    # 5*4*4 (a) + 5*1 (b) = 80 + 5 = 85
+    assert datasource.estimate_inmemory_data_size() == 85
+
+
+# ---------------------------------------------------------------------------
+# Cross-filesystem end-to-end (Ray Data convention)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "fs",
+    [
+        None,
+        lazy_fixture("local_fs"),  # pyarrow.fs (gets wrapped to fsspec internally)
+        lazy_fixture("local_fsspec_fs"),  # native fsspec
+    ],
+)
+def test_read_zarr_basic_across_filesystems(ray_start_regular_shared, fs, local_path):
+    """Round-trip a real Zarr store through read_zarr for each filesystem flavor.
+
+    Mirrors the parametrized read-path coverage other Ray Data datasources use
+    (lance, parquet, json, hudi, …) — exercises None / pyarrow.fs / fsspec
+    input shapes against the same store written to a local path.
+    """
+    store_path = os.path.join(local_path, "data.zarr")
+    images_src = np.arange(20, dtype="<i4").reshape(5, 4)
+    labels_src = np.arange(5, dtype="|u1")
+    _write_real_zarr_store(
+        Path(store_path),
+        {
+            "images": (images_src, (2, 4)),
+            "labels": (labels_src, (2,)),
+        },
+    )
+
+    ds = ray.data.read_zarr(store_path, filesystem=fs)
+
+    # 3 chunks each for images and labels (5/2 → ceil = 3) → 6 rows total.
+    assert ds.count() == 6
+    df = pd.DataFrame(ds.take_all())
+    np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
+    np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)
+
+
+def test_rejects_zarr_v3(tmp_path, monkeypatch):
+    """read_zarr targets zarr-python 2.x; an incompatible v3 install must raise a
+    clear, actionable error at construction, not a cryptic ImportError mid-read."""
+    monkeypatch.setattr(zarr, "__version__", "3.0.1")
+    with pytest.raises(ImportError, match=r"zarr-python 2\.x"):
+        zarrv2_datasource.ZarrV2Datasource(str(tmp_path))
+
+
+def test_explicit_filesystem_strips_uri_scheme(ray_start_regular_shared, tmp_path):
+    """An explicit ``filesystem=`` plus a scheme-prefixed path must strip the
+    scheme so the store path is backend-relative. Regression: pyarrow
+    filesystems can't resolve a ``file://`` / ``gs://`` prefix in the path."""
+    store_path = tmp_path / "scheme.zarr"
+    _write_real_zarr_store(store_path, {"data": (np.arange(6, dtype="<i4"), (2,))})
+
+    ds = zarrv2_datasource.ZarrV2Datasource(
+        f"file://{store_path}", filesystem=pyarrow.fs.LocalFileSystem()
+    )
+    assert ds._store_path == str(store_path)
+    df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
+    assert len(df) == 3
+
+
+def test_get_read_tasks_parallelism_zero(tmp_path):
+    """parallelism=0 must not divide by zero; fall back to a single task."""
+    store_path = tmp_path / "p0.zarr"
+    _write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (2,))})
+    ds = zarrv2_datasource.ZarrV2Datasource(str(store_path))
+    tasks = ds.get_read_tasks(parallelism=0)
+    assert len(tasks) >= 1
+
+
+def test_align_axis_0_rejects_scalar_array(tmp_path):
+    """align_axis_0=True with a 0-D (scalar) array must raise a clear error
+    rather than an IndexError when reading the (empty) axis-0 chunk size."""
+    store_path = tmp_path / "scalar.zarr"
+    root = zarr.open_group(str(store_path), mode="w")
+    root.create_dataset("vec", data=np.arange(8, dtype="<i4"), chunks=(4,))
+    root.create_dataset("scalar", data=np.array(42, dtype="<i4"))  # 0-D
+    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
+
+    with pytest.raises(ValueError, match=r"0-D \(scalar\)"):
+        zarrv2_datasource.ZarrV2Datasource(str(store_path), align_axis_0=True)
+
+
+def test_reads_zarr_zip_with_explicit_zip_filesystem(
+    ray_start_regular_shared, zarr_zip_store
+):
+    """A .zip path read through an explicitly-passed fsspec ZipFileSystem must
+    resolve the store at the archive root (store path ``""``), not treat the
+    ``.zip`` name as an entry inside the archive."""
+    zip_fs = fsspec.filesystem("zip", fo=str(zarr_zip_store))
+    ds = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store), filesystem=zip_fs)
+    assert ds._store_path == ""
+    df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
+    assert len(df) == 2
+
+
+def test_align_axis_0_columns_unify_across_blocks(
+    ray_start_regular_shared, aligned_zarrv2_store
+):
+    """Wide-form gives each array its own column, so blocks combine cleanly
+    across the dataset even with trailing edge chunks of differing shape -- the
+    batch-safe schema for row-aligned arrays."""
+    from ray.data._internal.arrow_ops.transform_pyarrow import unify_schemas
+    from ray.data.block import BlockAccessor
+
+    ds = zarrv2_datasource.ZarrV2Datasource(
+        str(aligned_zarrv2_store), align_axis_0=True, chunk_shapes=[3]
+    )
+    blocks = [block for task in ds.get_read_tasks(parallelism=64) for block in task()]
+    assert len(blocks) > 1  # actually exercise cross-block unification
+    schemas = [BlockAccessor.for_block(b).to_arrow().schema for b in blocks]
+    unified = unify_schemas(schemas)  # must not raise
+    assert {"t_start", "t_stop", "img", "state", "label"}.issubset(set(unified.names))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))
+
+
+# ---------------------------------------------------------------------------
+# Custom codec registration in Ray workers
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def fresh_ray():
+    """A clean Ray for a test that needs its own ``ray.init`` (e.g. a custom
+    ``runtime_env``). Unlike ``shutdown_only`` (teardown only), it also shuts
+    down any pre-existing cluster, so isolation doesn't depend on test order.
+    """
+    if ray.is_initialized():
+        ray.shutdown()
+    yield
+    if ray.is_initialized():
+        ray.shutdown()
+
+
+def test_custom_codec_succeeds_with_worker_setup_hook(fresh_ray, tmp_path):
+    """Test that we successfully register a custom codec.
+
+    numcodecs' registry is process-local.
+    """
+    import numcodecs
+
+    def _register_codec():
+        import numcodecs
+        import numpy as np
+
+        class _RayZarrTestCodec(numcodecs.abc.Codec):
+            codec_id = "ray_zarr_test_codec"
+
+            def encode(self, buf):
+                return bytes(buf)
+
+            def decode(self, buf, out=None):
+                arr = np.frombuffer(buf, dtype=np.uint8)
+                if out is not None:
+                    out[:] = arr.view(out.dtype)
+                    return out
+                return arr.copy()
+
+        numcodecs.register_codec(_RayZarrTestCodec)
+
+    # Register driver-side so we can write the store.
+    _register_codec()
+
+    store_path = tmp_path / "codec_test.zarr"
+    arr = zarr.open(
+        str(store_path),
+        mode="w",
+        shape=(8,),
+        chunks=(4,),
+        dtype="u1",
+        compressor=numcodecs.get_codec({"id": "ray_zarr_test_codec"}),
+    )
+    arr[:] = np.arange(8, dtype="u1")
+    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
+
+    ray.init(
+        num_cpus=1,
+        logging_level=logging.ERROR,
+        log_to_driver=False,
+        runtime_env={"worker_process_setup_hook": _register_codec},
+    )
+    ds = ray.data.read_zarr(str(store_path))
+    rows = sorted(ds.take_all(), key=lambda r: tuple(r["chunk_index"]))
+    recon = np.concatenate([r["chunk"] for r in rows])
+    np.testing.assert_array_equal(recon, np.arange(8, dtype="u1"))