chore: import upstream snapshot with attribution

python/ray/data/tests/unit/README.md

@@ -0,0 +1,26 @@
+# Unit Tests
+
+This directory contains unit tests that do not depend on distributed infrastructure or external dependencies.
+
+## Requirements
+
+Unit tests in this directory must be:
+- **Fast**: Execute in milliseconds, not seconds
+- **Isolated**: No dependencies on Ray runtime, external services, or file I/O
+- **Deterministic**: No randomness or time-based behavior
+
+## Restrictions
+
+Tests should NOT:
+- Initialize or use the Ray distributed runtime
+- Use `time.sleep()` or other time-based delays
+- Depend on external services, databases, or file systems
+- Make network calls
+
+## Enforcement
+
+The `conftest.py` in this directory enforces these restrictions by preventing:
+- `ray.init()` from being called
+- `time.sleep()` from being used
+
+If a test requires any of these, it should be moved to the main test directory instead.

python/ray/data/tests/unit/conftest.py

@@ -0,0 +1,24 @@
+import time
+
+import pytest
+
+import ray
+
+
+@pytest.fixture(autouse=True)
+def disallow_ray_init(monkeypatch):
+    def raise_on_init(*args, **kwargs):
+        raise RuntimeError("Unit tests should not depend on Ray being initialized.")
+
+    monkeypatch.setattr(ray, "init", raise_on_init)
+
+
+@pytest.fixture(autouse=True)
+def disallow_time_sleep(monkeypatch):
+    def raise_on_sleep(seconds):
+        raise RuntimeError(
+            f"Unit tests should not use time.sleep({seconds}). "
+            "Unit tests should be fast and deterministic."
+        )
+
+    monkeypatch.setattr(time, "sleep", raise_on_sleep)

python/ray/data/tests/unit/datasource_v2/test_file_chunker.py

@@ -0,0 +1,192 @@
+"""Unit tests for ``FileChunker`` implementations in DataSourceV2."""
+from pathlib import Path
+from typing import cast
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+
+from ray.data._internal.datasource_v2.chunkers.file_chunker import (
+    ChunkMetadata,
+    LineDelimitedFileChunker,
+    LineDelimitedFileChunkMetadata,
+    ParquetFileChunker,
+    ParquetFileChunkMetadata,
+    WholeFileChunker,
+    create_chunk_metadata,
+)
+
+
+def _write_parquet_with_row_groups(
+    path: str, num_row_groups: int, rows_per_group: int = 10
+) -> int:
+    """Write a Parquet file with exactly ``num_row_groups`` row groups.
+
+    Returns the on-disk file size in bytes.
+    """
+    n = num_row_groups * rows_per_group
+    table = pa.table({"a": list(range(n))})
+    pq.write_table(table, path, row_group_size=rows_per_group)
+    return Path(path).stat().st_size
+
+
+class TestCreateChunkMetadata:
+    def test_validates_missing_keys(self):
+        with pytest.raises(ValueError, match="Missing required keys"):
+            create_chunk_metadata(ParquetFileChunkMetadata, row_group_start=0)
+
+    def test_validates_unexpected_keys(self):
+        with pytest.raises(ValueError, match="Unexpected keys"):
+            create_chunk_metadata(
+                ParquetFileChunkMetadata,
+                row_group_start=0,
+                row_group_end=1,
+                extra_field="boom",
+            )
+
+    def test_returns_dict_with_keys(self):
+        md = create_chunk_metadata(
+            ParquetFileChunkMetadata, row_group_start=2, row_group_end=5
+        )
+        assert md == {"row_group_start": 2, "row_group_end": 5}
+
+
+class TestWholeFileChunker:
+    def test_yields_single_none_chunk(self):
+        chunker = WholeFileChunker()
+        chunks = list(chunker.generate_chunk_metadatas("foo.bin", 12345))
+        assert chunks == [(None, 12345)]
+
+    def test_does_not_read_metadata(self):
+        assert WholeFileChunker.reads_file_metadata is False
+
+
+class TestLineDelimitedFileChunker:
+    def test_chunks_uncompressed_file(self):
+        chunker = LineDelimitedFileChunker()
+        # 600MB file at 256MB chunks -> 3 chunks (256, 256, 88).
+        chunks = list(chunker.generate_chunk_metadatas("data.jsonl", 600 * 1024 * 1024))
+        assert len(chunks) == 3
+        for i, (md, size) in enumerate(chunks):
+            assert md is not None
+            md = cast(LineDelimitedFileChunkMetadata, md)
+            assert md["chunk_byte_start_idx"] == i * 256 * 1024 * 1024
+            assert size == md["chunk_byte_end_idx"] - md["chunk_byte_start_idx"]
+        # Final chunk should clip to file_size.
+        last_md = cast(LineDelimitedFileChunkMetadata, chunks[-1][0])
+        assert last_md["chunk_byte_end_idx"] == 600 * 1024 * 1024
+
+    def test_compressed_file_yields_whole(self):
+        chunker = LineDelimitedFileChunker()
+        chunks = list(chunker.generate_chunk_metadatas("data.jsonl.gz", 1024))
+        assert chunks == [(None, 1024)]
+
+    def test_does_not_read_metadata(self):
+        assert LineDelimitedFileChunker.reads_file_metadata is False
+
+
+class TestParquetFileChunker:
+    def test_reads_file_metadata_flag(self):
+        assert ParquetFileChunker.reads_file_metadata is True
+
+    @pytest.mark.parametrize(
+        "target_chunk_size, expected_ranges",
+        [
+            # target < row-group size → K=1, one chunk per row group.
+            (1, [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)]),
+            # target ≫ file size → all row groups bundled into one chunk.
+            (10 * 1024**3, [(0, 5)]),
+        ],
+        ids=["target_below_rg_size", "target_above_file_size"],
+    )
+    def test_row_group_bundling_by_target(
+        self, tmp_path, target_chunk_size, expected_ranges
+    ):
+        p = str(tmp_path / "d.parquet")
+        size = _write_parquet_with_row_groups(p, num_row_groups=5)
+        chunker = ParquetFileChunker(target_chunk_size=target_chunk_size)
+        chunks = list(chunker.generate_chunk_metadatas(p, size))
+        ranges = [(m["row_group_start"], m["row_group_end"]) for m, _ in chunks]
+        assert ranges == expected_ranges
+
+    def test_chunk_size_equals_summed_row_group_bytes(self, tmp_path):
+        p = str(tmp_path / "d.parquet")
+        size = _write_parquet_with_row_groups(p, num_row_groups=4)
+        chunker = ParquetFileChunker(target_chunk_size=1)
+        chunks = list(chunker.generate_chunk_metadatas(p, size))
+        md = pq.read_metadata(p)
+        expected_total = sum(
+            md.row_group(i).column(c).total_compressed_size
+            for i in range(md.num_row_groups)
+            for c in range(md.row_group(i).num_columns)
+        )
+        assert sum(sz for _, sz in chunks) == expected_total
+
+    def test_ranges_are_contiguous_and_cover_all_row_groups(self, tmp_path):
+        """Bundled ranges partition [0, num_row_groups) with no gaps/overlap."""
+        p = str(tmp_path / "d.parquet")
+        # Many small row groups + a moderate target → some bundling.
+        size = _write_parquet_with_row_groups(p, num_row_groups=12, rows_per_group=5)
+        n = pq.read_metadata(p).num_row_groups
+        chunker = ParquetFileChunker(target_chunk_size=2_000)
+        chunks = list(chunker.generate_chunk_metadatas(p, size))
+        ranges = [(m["row_group_start"], m["row_group_end"]) for m, _ in chunks]
+        # Contiguous, starts at 0, ends at n, each non-empty.
+        assert ranges[0][0] == 0
+        assert ranges[-1][1] == n
+        for (_, prev_end), (next_start, _) in zip(ranges, ranges[1:]):
+            assert prev_end == next_start
+        assert all(start < end for start, end in ranges)
+
+    def test_corrupt_footer_falls_back_to_whole_file(self, tmp_path):
+        p = str(tmp_path / "bad.parquet")
+        Path(p).write_bytes(b"this is not a parquet file")
+        chunker = ParquetFileChunker(target_chunk_size=1)
+        chunks = list(chunker.generate_chunk_metadatas(p, 26))
+        assert chunks == [(None, 26)]
+
+    @pytest.mark.parametrize(
+        "ctor_arg, ctx_chunk_size, ctx_min_block, expected",
+        [
+            # ctor arg wins over the context knobs.
+            (2048, 1024, 7777, 2048),
+            # An explicit 0 is honored (resolved with ``is not None``, not
+            # ``or``), so it isn't silently treated as "unset" and overridden.
+            (0, 1024, 7777, 0),
+            # ctx chunk-size knob used when there's no ctor arg.
+            (None, 1024, 7777, 1024),
+            # Falls back to target_min_block_size when the chunk-size knob is unset.
+            (None, None, 7777, 7777),
+        ],
+        ids=["ctor_arg", "ctor_arg_zero", "ctx_knob", "fallback_min_block"],
+    )
+    def test_target_chunk_size_resolution(
+        self, restore_data_context, ctor_arg, ctx_chunk_size, ctx_min_block, expected
+    ):
+        from ray.data.context import DataContext
+
+        ctx = DataContext.get_current()
+        ctx.parquet_chunker_target_chunk_size = ctx_chunk_size
+        ctx.target_min_block_size = ctx_min_block
+        chunker = ParquetFileChunker(target_chunk_size=ctor_arg)
+        assert chunker._target_chunk_size == expected
+
+
+def test_chunk_metadata_subclasses_are_typeddicts():
+    # Ensures the subclasses don't accidentally inherit unrelated keys.
+    pmd: ChunkMetadata = create_chunk_metadata(
+        ParquetFileChunkMetadata, row_group_start=0, row_group_end=1
+    )
+    lmd: ChunkMetadata = create_chunk_metadata(
+        LineDelimitedFileChunkMetadata,
+        chunk_byte_start_idx=0,
+        chunk_byte_end_idx=10,
+    )
+    assert set(pmd.keys()) == {"row_group_start", "row_group_end"}
+    assert set(lmd.keys()) == {"chunk_byte_start_idx", "chunk_byte_end_idx"}
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/datasource_v2/test_file_indexer.py

@@ -0,0 +1,352 @@
+import os
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from pyarrow.fs import LocalFileSystem
+
+from ray.data._internal.datasource_v2.chunkers.file_chunker import (
+    FileChunker,
+    LineDelimitedFileChunker,
+    ParquetFileChunker,
+    WholeFileChunker,
+)
+from ray.data._internal.datasource_v2.listing.file_indexer import (
+    FileInfo,
+    NonSamplingFileIndexer,
+)
+from ray.data._internal.datasource_v2.listing.file_pruners import FileExtensionPruner
+
+
+class _CountingChunker(FileChunker):
+    """Fake chunker that emits a fixed number of chunks per file and reports it
+    reads metadata, so the indexer takes the threaded fan-out path."""
+
+    reads_file_metadata = True
+
+    def __init__(self, chunks_per_file: int):
+        self._chunks_per_file = chunks_per_file
+
+    def generate_chunk_metadatas(self, path, file_size, filesystem=None):
+        for i in range(self._chunks_per_file):
+            yield {"chunk_index": i}, file_size
+
+
+def _list_all(indexer, paths, **kwargs):
+    """Run list_files and flatten all manifests into (path, size) pairs."""
+    pa_paths = pa.array(paths)
+    fs = LocalFileSystem()
+    manifests = list(indexer.list_files(pa_paths, filesystem=fs, **kwargs))
+    results = []
+    for m in manifests:
+        for p, s in zip(m.paths, m.file_sizes):
+            results.append((str(p), int(s)))
+    return sorted(results)
+
+
+@pytest.fixture(params=[1, 2], ids=["sequential", "threaded"])
+def indexer(request):
+    """Yield a NonSamplingFileIndexer using sequential or threaded listing."""
+    return NonSamplingFileIndexer(ignore_missing_paths=False, num_workers=request.param)
+
+
+class TestListFiles:
+    def test_single_file(self, tmp_path, indexer):
+        f = tmp_path / "data.csv"
+        f.write_bytes(b"x" * 42)
+
+        results = _list_all(indexer, [str(f)])
+        assert results == [(str(f), 42)]
+
+    def test_directory(self, tmp_path, indexer):
+        for name in ["a.csv", "b.csv", "c.csv"]:
+            (tmp_path / name).write_bytes(b"x" * 10)
+
+        results = _list_all(indexer, [str(tmp_path)])
+        assert len(results) == 3
+        assert all(size == 10 for _, size in results)
+
+    def test_nested_directories(self, tmp_path, indexer):
+        (tmp_path / "top.csv").write_bytes(b"x" * 100)
+        (tmp_path / "sub").mkdir()
+        (tmp_path / "sub" / "nested.csv").write_bytes(b"x" * 100)
+        (tmp_path / "sub" / "deep").mkdir()
+        (tmp_path / "sub" / "deep" / "leaf.csv").write_bytes(b"x" * 100)
+
+        results = _list_all(indexer, [str(tmp_path)])
+        assert len(results) == 3
+        basenames = sorted(os.path.basename(p) for p, _ in results)
+        assert basenames == ["leaf.csv", "nested.csv", "top.csv"]
+
+    def test_multiple_paths(self, tmp_path, indexer):
+        f1 = tmp_path / "one.csv"
+        f2 = tmp_path / "two.csv"
+        f1.write_bytes(b"x" * 10)
+        f2.write_bytes(b"x" * 20)
+
+        results = _list_all(indexer, [str(f1), str(f2)])
+        assert sorted(results) == [(str(f1), 10), (str(f2), 20)]
+
+    @pytest.mark.parametrize(
+        "filename",
+        [".hidden", "_metadata", "_SUCCESS", ".gitignore"],
+        ids=["dot-prefix", "underscore-prefix", "underscore-upper", "dotfile"],
+    )
+    def test_excludes_hidden_and_metadata_files(self, tmp_path, indexer, filename):
+        (tmp_path / filename).write_bytes(b"x" * 100)
+        (tmp_path / "visible.csv").write_bytes(b"x" * 100)
+
+        results = _list_all(indexer, [str(tmp_path)])
+        assert len(results) == 1
+        assert os.path.basename(results[0][0]) == "visible.csv"
+
+    @pytest.mark.parametrize(
+        "filename",
+        ["_metadata", "_my_file.csv", ".hidden_data"],
+        ids=["underscore-metadata", "underscore-csv", "dot-hidden"],
+    )
+    def test_includes_excluded_prefix_files_in_subdirectories(
+        self, tmp_path, indexer, filename
+    ):
+        """Files whose names start with _ or . should only be excluded when
+        they appear at the top level of the listed directory, not when they
+        appear inside a subdirectory. The relative path from the root is
+        e.g. "subdir/_metadata" which starts with "s", not "_"."""
+        sub = tmp_path / "subdir"
+        sub.mkdir()
+        (sub / filename).write_bytes(b"x" * 50)
+        (sub / "normal.csv").write_bytes(b"x" * 50)
+
+        results = _list_all(indexer, [str(tmp_path)])
+        basenames = sorted(os.path.basename(p) for p, _ in results)
+        assert filename in basenames
+        assert "normal.csv" in basenames
+
+    def test_skips_zero_size_files(self, tmp_path, indexer):
+        (tmp_path / "empty.csv").write_bytes(b"")
+        (tmp_path / "real.csv").write_bytes(b"x" * 50)
+
+        results = _list_all(indexer, [str(tmp_path)])
+        assert len(results) == 1
+        assert os.path.basename(results[0][0]) == "real.csv"
+
+    def test_empty_directory(self, tmp_path, indexer):
+        os.makedirs(tmp_path / "empty_dir", exist_ok=True)
+        results = _list_all(indexer, [str(tmp_path / "empty_dir")])
+        assert results == []
+
+
+class TestPruners:
+    @pytest.fixture
+    def indexer(self):
+        return NonSamplingFileIndexer(ignore_missing_paths=False, num_workers=1)
+
+    @pytest.mark.parametrize(
+        "extensions, expected_basenames",
+        [
+            (["csv"], ["a.csv"]),
+            (["json"], ["b.json"]),
+            (["csv", "json"], ["a.csv", "b.json"]),
+            (["parquet"], []),
+        ],
+        ids=["csv-only", "json-only", "csv-and-json", "no-match"],
+    )
+    def test_extension_pruner(self, tmp_path, indexer, extensions, expected_basenames):
+        (tmp_path / "a.csv").write_bytes(b"x" * 100)
+        (tmp_path / "b.json").write_bytes(b"x" * 100)
+        (tmp_path / "c.txt").write_bytes(b"x" * 100)
+
+        pruner = FileExtensionPruner(extensions)
+        results = _list_all(indexer, [str(tmp_path)], pruners=[pruner])
+        basenames = sorted(os.path.basename(p) for p, _ in results)
+        assert basenames == sorted(expected_basenames)
+
+    def test_multiple_pruners_intersect(self, tmp_path, indexer):
+        """Multiple pruners are AND'd — a file must pass all of them."""
+        (tmp_path / "a.csv").write_bytes(b"x" * 100)
+        (tmp_path / "b.json").write_bytes(b"x" * 100)
+
+        pruner_csv = FileExtensionPruner(["csv", "json"])
+        pruner_json = FileExtensionPruner(["json"])
+        results = _list_all(indexer, [str(tmp_path)], pruners=[pruner_csv, pruner_json])
+        basenames = [os.path.basename(p) for p, _ in results]
+        assert basenames == ["b.json"]
+
+
+class TestMissingPaths:
+    def test_raises_on_missing_path(self, tmp_path):
+        indexer = NonSamplingFileIndexer(ignore_missing_paths=False)
+        missing = str(tmp_path / "nonexistent")
+
+        with pytest.raises(FileNotFoundError):
+            _list_all(indexer, [missing])
+
+    def test_ignores_missing_path(self, tmp_path):
+        indexer = NonSamplingFileIndexer(ignore_missing_paths=True)
+        missing = str(tmp_path / "nonexistent")
+
+        results = _list_all(indexer, [missing])
+        assert results == []
+
+    def test_mixed_existing_and_missing(self, tmp_path):
+        indexer = NonSamplingFileIndexer(ignore_missing_paths=True)
+        real = tmp_path / "real.csv"
+        real.write_bytes(b"x" * 10)
+        missing = str(tmp_path / "gone")
+
+        results = _list_all(indexer, [str(real), missing])
+        assert results == [(str(real), 10)]
+
+
+class TestManifestBatching:
+    def test_splits_into_multiple_manifests(self, tmp_path):
+        indexer = NonSamplingFileIndexer(
+            ignore_missing_paths=False, max_paths_per_output=3
+        )
+
+        for i in range(7):
+            (tmp_path / f"file_{i}.csv").write_bytes(b"x" * 100)
+
+        pa_paths = pa.array([str(tmp_path)])
+        fs = LocalFileSystem()
+        manifests = list(indexer.list_files(pa_paths, filesystem=fs))
+
+        assert len(manifests) == 3  # ceil(7/3)
+        assert len(manifests[0]) == 3
+        assert len(manifests[1]) == 3
+        assert len(manifests[2]) == 1
+
+        total_files = sum(len(m) for m in manifests)
+        assert total_files == 7
+
+
+class TestFileChunkerIntegration:
+    """Cover ``NonSamplingFileIndexer`` interaction with a ``FileChunker``."""
+
+    def test_default_uses_whole_file_chunker(self):
+        indexer = NonSamplingFileIndexer(ignore_missing_paths=False)
+        assert isinstance(indexer.file_chunker, WholeFileChunker)
+
+    def test_explicit_chunker_is_exposed(self):
+        chunker = ParquetFileChunker(target_chunk_size=1024)
+        indexer = NonSamplingFileIndexer(
+            ignore_missing_paths=False, file_chunker=chunker
+        )
+        assert indexer.file_chunker is chunker
+
+    def test_whole_file_chunker_yields_none_chunk_metadata(self, tmp_path):
+        (tmp_path / "a.csv").write_bytes(b"x" * 100)
+        indexer = NonSamplingFileIndexer(ignore_missing_paths=False, num_workers=1)
+        fs = LocalFileSystem()
+        manifests = list(indexer.list_files(pa.array([str(tmp_path)]), filesystem=fs))
+        assert len(manifests) == 1
+        manifest = manifests[0]
+        assert len(manifest) == 1
+        # ``WholeFileChunker`` emits one ``None`` chunk per file.
+        assert list(manifest.file_chunk_metadatas) == [None]
+        assert list(manifest.file_sizes) == [100]
+
+    def test_parquet_chunker_splits_file_on_row_group_boundaries(self, tmp_path):
+        # A real Parquet file with 8 row groups; the chunker reads the footer
+        # at listing time and (target < row-group size) emits one chunk per
+        # row group with an explicit half-open range.
+        table = pa.table({"a": list(range(80))})
+        pq.write_table(table, str(tmp_path / "big.parquet"), row_group_size=10)
+        chunker = ParquetFileChunker(target_chunk_size=1)
+        indexer = NonSamplingFileIndexer(
+            ignore_missing_paths=False,
+            num_workers=1,
+            file_chunker=chunker,
+        )
+        fs = LocalFileSystem()
+        manifests = list(indexer.list_files(pa.array([str(tmp_path)]), filesystem=fs))
+        rows = []
+        for m in manifests:
+            for path, size, md in zip(m.paths, m.file_sizes, m.file_chunk_metadatas):
+                rows.append((str(path), int(size), md))
+
+        # 8 row groups → 8 chunks, contiguous half-open ranges.
+        assert len(rows) == 8
+        ranges = [(md["row_group_start"], md["row_group_end"]) for _, _, md in rows]
+        assert ranges == [(i, i + 1) for i in range(8)]
+
+    def test_parquet_chunker_parallel_footer_reads(self, tmp_path):
+        # With num_workers > 1 the chunker's footer reads fan across the
+        # thread pool (reads_file_metadata=True). Verify correctness is
+        # unaffected: every file's row groups are represented exactly once.
+        for f in range(4):
+            table = pa.table({"a": list(range(30))})
+            pq.write_table(table, str(tmp_path / f"f{f}.parquet"), row_group_size=10)
+        chunker = ParquetFileChunker(target_chunk_size=1)
+        indexer = NonSamplingFileIndexer(
+            ignore_missing_paths=False,
+            num_workers=4,
+            file_chunker=chunker,
+        )
+        fs = LocalFileSystem()
+        manifests = list(indexer.list_files(pa.array([str(tmp_path)]), filesystem=fs))
+        per_path_ranges = {}
+        for m in manifests:
+            for path, _, md in zip(m.paths, m.file_sizes, m.file_chunk_metadatas):
+                per_path_ranges.setdefault(str(path), []).append(
+                    (md["row_group_start"], md["row_group_end"])
+                )
+        # 4 files × 3 row groups each.
+        assert len(per_path_ranges) == 4
+        for ranges in per_path_ranges.values():
+            assert sorted(ranges) == [(0, 1), (1, 2), (2, 3)]
+
+    def test_line_delimited_chunker_byte_ranges(self, tmp_path):
+        (tmp_path / "a.jsonl").write_bytes(b"x" * 10_000)
+        chunker = LineDelimitedFileChunker()
+        # Force smaller chunks via a private override so the unit test
+        # doesn't need a 256 MB file on disk.
+        chunker._CHUNK_BYTE_SIZE = 1024
+        indexer = NonSamplingFileIndexer(
+            ignore_missing_paths=False,
+            num_workers=1,
+            file_chunker=chunker,
+        )
+        fs = LocalFileSystem()
+        manifests = list(indexer.list_files(pa.array([str(tmp_path)]), filesystem=fs))
+        rows = []
+        for m in manifests:
+            for path, size, md in zip(m.paths, m.file_sizes, m.file_chunk_metadatas):
+                rows.append((str(path), int(size), md))
+        assert len(rows) == 10
+        # Byte ranges must tile the file exactly.
+        assert rows[0][2]["chunk_byte_start_idx"] == 0
+        assert rows[-1][2]["chunk_byte_end_idx"] == 10_000
+
+    def test_parallel_chunking_preserves_discovery_order(self):
+        """Regression: with num_workers>1 and preserve_order, a multi-chunk-per-
+        file chunker's records stay grouped per file in discovery order.
+
+        ``make_async_gen`` only preserves order for a 1:1 map, so the indexer
+        must emit one record list per file (not yield chunk rows individually) --
+        otherwise its round-robin merge interleaves chunks across the files
+        processed concurrently, e.g. f0,f1,f2,f3,f0,f1,... instead of f0,f0,f0,...
+        """
+        chunker = _CountingChunker(chunks_per_file=3)
+        indexer = NonSamplingFileIndexer(
+            ignore_missing_paths=False, num_workers=4, file_chunker=chunker
+        )
+        file_infos = [FileInfo(path=f"f{i}.parquet", size=100 + i) for i in range(8)]
+        records = list(
+            indexer._generate_chunk_records(
+                file_infos, filesystem=None, preserve_order=True
+            )
+        )
+        # 8 files x 3 chunks each, contiguous per file in discovery order.
+        assert [path for path, _, _ in records] == [
+            f"f{i}.parquet" for i in range(8) for _ in range(3)
+        ]
+        assert [md["chunk_index"] for _, _, md in records] == [
+            i for _ in range(8) for i in range(3)
+        ]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/datasource_v2/test_list_files_op.py

@@ -0,0 +1,78 @@
+"""Unit tests for :class:`ListFiles` logical op.
+
+Full physical-planning tests live in the CI parquet regression suite
+(they need Ray initialized for ``ray.put`` on the listing input
+bundles). Here we exercise just the logical op shape and the shuffle
+factory semantics.
+"""
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+
+from ray.data._internal.datasource_v2.listing.file_indexer import (
+    NonSamplingFileIndexer,
+)
+from ray.data._internal.datasource_v2.listing.file_manifest import (
+    FILE_SIZE_COLUMN_NAME,
+    PATH_COLUMN_NAME,
+)
+from ray.data._internal.logical.operators import ListFiles
+from ray.data.datasource.file_based_datasource import FileShuffleConfig
+
+
+def _mk_indexer():
+    return NonSamplingFileIndexer(ignore_missing_paths=False)
+
+
+def _mk_list_files(tmp_path, num_files: int = 3, shuffle_seed=None):
+    for i in range(num_files):
+        pq.write_table(pa.table({"x": [i]}), str(tmp_path / f"f{i}.parquet"))
+    paths = [str(tmp_path / f"f{i}.parquet") for i in range(num_files)]
+
+    def _shuffle_factory():
+        if shuffle_seed is None:
+            return None
+        return FileShuffleConfig(seed=shuffle_seed)
+
+    import pyarrow.fs as pafs
+
+    return ListFiles(
+        paths=paths,
+        file_indexer=_mk_indexer(),
+        filesystem=pafs.LocalFileSystem(),
+        source_paths=paths,
+        shuffle_config_factory=_shuffle_factory,
+    )
+
+
+def test_list_files_infers_manifest_schema(tmp_path):
+    op = _mk_list_files(tmp_path, num_files=1)
+    schema = op.infer_schema()
+    assert schema.names == [PATH_COLUMN_NAME, FILE_SIZE_COLUMN_NAME]
+    assert schema.field(PATH_COLUMN_NAME).type == pa.string()
+    assert schema.field(FILE_SIZE_COLUMN_NAME).type == pa.int64()
+
+
+def test_list_files_has_no_input_dependencies(tmp_path):
+    op = _mk_list_files(tmp_path, num_files=1)
+    assert op.input_dependencies == []
+    assert op.num_outputs is None
+    assert op.output_data() is None
+
+
+def test_shuffle_config_factory_none_when_unconfigured(tmp_path):
+    op = _mk_list_files(tmp_path, num_files=1, shuffle_seed=None)
+    assert op.shuffle_config_factory() is None
+
+
+def test_shuffle_config_factory_returns_config_when_seeded(tmp_path):
+    op = _mk_list_files(tmp_path, num_files=1, shuffle_seed=42)
+    config = op.shuffle_config_factory()
+    assert isinstance(config, FileShuffleConfig)
+    assert config.seed == 42
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-xvs"]))

python/ray/data/tests/unit/expressions/test_arithmetic.py

@@ -0,0 +1,526 @@
+"""Tests for arithmetic expression operations.
+
+This module tests:
+- Basic arithmetic: ADD, SUB, MUL, DIV, FLOORDIV
+- Reverse arithmetic: radd, rsub, rmul, rtruediv, rfloordiv
+- Rounding helpers: ceil, floor, round, trunc
+- Logarithmic helpers: ln, log10, log2, exp
+- Trigonometric helpers: sin, cos, tan, asin, acos, atan
+- Arithmetic helpers: negate, sign, power, abs
+"""
+
+import math
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+from pkg_resources import parse_version
+
+from ray.data._internal.planner.plan_expression.expression_evaluator import eval_expr
+from ray.data.expressions import BinaryExpr, Operation, UDFExpr, col, lit
+from ray.data.tests.conftest import get_pyarrow_version
+
+pytestmark = pytest.mark.skipif(
+    get_pyarrow_version() < parse_version("20.0.0"),
+    reason="Expression unit tests require PyArrow >= 20.0.0",
+)
+
+# ──────────────────────────────────────
+# Basic Arithmetic Operations
+# ──────────────────────────────────────
+
+
+class TestBasicArithmetic:
+    """Tests for basic arithmetic operations (+, -, *, /, //)."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for arithmetic tests."""
+        return pd.DataFrame(
+            {
+                "a": [10, 20, 30, 40],
+                "b": [2, 4, 5, 8],
+                "c": [1.5, 2.5, 3.5, 4.5],
+            }
+        )
+
+    # ── Addition ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_name,expected_values",
+        [
+            (col("a") + 5, "add_literal", [15, 25, 35, 45]),
+            (col("a") + col("b"), "add_cols", [12, 24, 35, 48]),
+            (col("a") + lit(10), "add_lit", [20, 30, 40, 50]),
+        ],
+        ids=["col_plus_int", "col_plus_col", "col_plus_lit"],
+    )
+    def test_addition(self, sample_data, expr, expected_name, expected_values):
+        """Test addition operations."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.ADD
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values, name=None),
+            check_names=False,
+        )
+
+    def test_reverse_addition(self, sample_data):
+        """Test reverse addition (literal + expr)."""
+        expr = 5 + col("a")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.ADD
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([15, 25, 35, 45])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_string_concat_invalid_input_type(self):
+        """Reject non-string-like inputs in string concatenation."""
+        table = pa.table({"name": ["a", "b"], "age": [1, 2]})
+        expr = col("name") + col("age")
+        with pytest.raises(TypeError, match="string-like pyarrow.*int64"):
+            eval_expr(expr, table)
+
+    # ── Subtraction ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("a") - 5, [5, 15, 25, 35]),
+            (col("a") - col("b"), [8, 16, 25, 32]),
+        ],
+        ids=["col_minus_int", "col_minus_col"],
+    )
+    def test_subtraction(self, sample_data, expr, expected_values):
+        """Test subtraction operations."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.SUB
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values, name=None),
+            check_names=False,
+        )
+
+    def test_reverse_subtraction(self, sample_data):
+        """Test reverse subtraction (literal - expr)."""
+        expr = 100 - col("a")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.SUB
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([90, 80, 70, 60])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Multiplication ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("a") * 2, [20, 40, 60, 80]),
+            (col("a") * col("b"), [20, 80, 150, 320]),
+        ],
+        ids=["col_times_int", "col_times_col"],
+    )
+    def test_multiplication(self, sample_data, expr, expected_values):
+        """Test multiplication operations."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.MUL
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values, name=None),
+            check_names=False,
+        )
+
+    def test_reverse_multiplication(self, sample_data):
+        """Test reverse multiplication (literal * expr)."""
+        expr = 3 * col("b")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.MUL
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([6, 12, 15, 24])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Division ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("a") / 2, [5.0, 10.0, 15.0, 20.0]),
+            (col("a") / col("b"), [5.0, 5.0, 6.0, 5.0]),
+        ],
+        ids=["col_div_int", "col_div_col"],
+    )
+    def test_division(self, sample_data, expr, expected_values):
+        """Test division operations."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.DIV
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values, name=None),
+            check_names=False,
+        )
+
+    def test_reverse_division(self, sample_data):
+        """Test reverse division (literal / expr)."""
+        expr = 100 / col("a")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.DIV
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([10.0, 5.0, 100 / 30, 2.5])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Floor Division ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("a") // 3, [3, 6, 10, 13]),
+            (col("a") // col("b"), [5, 5, 6, 5]),
+        ],
+        ids=["col_floordiv_int", "col_floordiv_col"],
+    )
+    def test_floor_division(self, sample_data, expr, expected_values):
+        """Test floor division operations."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.FLOORDIV
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values, name=None),
+            check_names=False,
+        )
+
+    def test_reverse_floor_division(self, sample_data):
+        """Test reverse floor division (literal // expr)."""
+        expr = 100 // col("a")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.FLOORDIV
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([10, 5, 3, 2])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Modulo ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("a") % 3, [1, 2, 0, 1]),
+            (col("a") % col("c"), [1.0, 0.0, 2.0, 4.0]),
+            (10 % col("b"), [0, 2, 0, 2]),
+        ],
+        ids=["col_mod_int", "col_mod_fp", "col_rmod_int"],
+    )
+    def test_modulo(self, sample_data, expr, expected_values):
+        """Test modulo operations."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.MOD
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values, name=None),
+            check_names=False,
+        )
+
+
+# ──────────────────────────────────────
+# Complex Arithmetic Expressions
+# ──────────────────────────────────────
+
+
+class TestComplexArithmetic:
+    """Tests for complex arithmetic expressions with multiple operations."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for complex arithmetic tests."""
+        return pd.DataFrame(
+            {
+                "x": [1.0, 2.0, 3.0, 4.0],
+                "y": [4.0, 3.0, 2.0, 1.0],
+                "z": [2.0, 2.0, 2.0, 2.0],
+            }
+        )
+
+    def test_chained_operations(self, sample_data):
+        """Test chained arithmetic operations."""
+        expr = (col("x") + col("y")) * col("z")
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([10.0, 10.0, 10.0, 10.0])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_nested_operations(self, sample_data):
+        """Test nested arithmetic operations."""
+        expr = ((col("x") * 2) + (col("y") / 2)) - 1
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([3.0, 4.5, 6.0, 7.5])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_order_of_operations(self, sample_data):
+        """Test that order of operations is respected."""
+        # Should compute x + (y * z) due to operator precedence
+        expr = col("x") + col("y") * col("z")
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([9.0, 8.0, 7.0, 6.0])  # 1+8, 2+6, 3+4, 4+2
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# Rounding Operations
+# ──────────────────────────────────────
+
+
+class TestRoundingOperations:
+    """Tests for rounding helper methods."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with decimal values for rounding tests."""
+        return pd.DataFrame(
+            {
+                "value": [1.2, 2.5, 3.7, -1.3, -2.5, -3.8],
+            }
+        )
+
+    @pytest.mark.parametrize(
+        "method,expected_values",
+        [
+            ("ceil", [2, 3, 4, -1, -2, -3]),
+            ("floor", [1, 2, 3, -2, -3, -4]),
+            ("trunc", [1, 2, 3, -1, -2, -3]),
+        ],
+        ids=["ceil", "floor", "trunc"],
+    )
+    def test_rounding_methods(self, sample_data, method, expected_values):
+        """Test rounding methods (ceil, floor, trunc)."""
+        expr = getattr(col("value"), method)()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        # Convert to list for comparison since PyArrow might return different types
+        result_list = result.tolist()
+        assert result_list == expected_values
+
+    def test_round_method(self, sample_data):
+        """Test round method (may differ due to banker's rounding)."""
+        expr = col("value").round()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        # PyArrow uses banker's rounding (round half to even)
+        # Just verify it runs and returns numeric values
+        assert len(result) == len(sample_data)
+
+
+# ──────────────────────────────────────
+# Logarithmic Operations
+# ──────────────────────────────────────
+
+
+class TestLogarithmicOperations:
+    """Tests for logarithmic helper methods."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with positive values for logarithmic tests."""
+        return pd.DataFrame(
+            {
+                "value": [1.0, math.e, 10.0, 100.0],
+            }
+        )
+
+    def test_ln(self, sample_data):
+        """Test natural logarithm."""
+        expr = col("value").ln()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        expected = [0.0, 1.0, math.log(10), math.log(100)]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_log10(self, sample_data):
+        """Test base-10 logarithm."""
+        expr = col("value").log10()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        expected = [0.0, math.log10(math.e), 1.0, 2.0]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_log2(self):
+        """Test base-2 logarithm."""
+        data = pd.DataFrame({"value": [1.0, 2.0, 4.0, 8.0]})
+        expr = col("value").log2()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, data)
+        expected = [0.0, 1.0, 2.0, 3.0]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_exp(self):
+        """Test exponential function."""
+        data = pd.DataFrame({"value": [0.0, 1.0, 2.0]})
+        expr = col("value").exp()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, data)
+        expected = [1.0, math.e, math.e**2]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+
+# ──────────────────────────────────────
+# Trigonometric Operations
+# ──────────────────────────────────────
+
+
+class TestTrigonometricOperations:
+    """Tests for trigonometric helper methods."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with angles in radians for trig tests."""
+        return pd.DataFrame(
+            {
+                "angle": [0.0, math.pi / 6, math.pi / 4, math.pi / 3, math.pi / 2],
+            }
+        )
+
+    def test_sin(self, sample_data):
+        """Test sine function."""
+        expr = col("angle").sin()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        expected = [0.0, 0.5, math.sqrt(2) / 2, math.sqrt(3) / 2, 1.0]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_cos(self, sample_data):
+        """Test cosine function."""
+        expr = col("angle").cos()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        expected = [1.0, math.sqrt(3) / 2, math.sqrt(2) / 2, 0.5, 0.0]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_tan(self):
+        """Test tangent function."""
+        data = pd.DataFrame({"angle": [0.0, math.pi / 4]})
+        expr = col("angle").tan()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, data)
+        expected = [0.0, 1.0]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_asin(self):
+        """Test arcsine function."""
+        data = pd.DataFrame({"value": [0.0, 0.5, 1.0]})
+        expr = col("value").asin()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, data)
+        expected = [0.0, math.pi / 6, math.pi / 2]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_acos(self):
+        """Test arccosine function."""
+        data = pd.DataFrame({"value": [1.0, 0.5, 0.0]})
+        expr = col("value").acos()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, data)
+        expected = [0.0, math.pi / 3, math.pi / 2]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+    def test_atan(self):
+        """Test arctangent function."""
+        data = pd.DataFrame({"value": [0.0, 1.0]})
+        expr = col("value").atan()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, data)
+        expected = [0.0, math.pi / 4]
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+
+# ──────────────────────────────────────
+# Arithmetic Helper Operations
+# ──────────────────────────────────────
+
+
+class TestArithmeticHelpers:
+    """Tests for arithmetic helper methods (negate, sign, power, abs)."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for arithmetic helper tests."""
+        return pd.DataFrame(
+            {
+                "value": [5, -3, 0, 10, -7],
+            }
+        )
+
+    def test_negate(self, sample_data):
+        """Test negate method."""
+        expr = col("value").negate()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        expected = [-5, 3, 0, -10, 7]
+        assert result.tolist() == expected
+
+    def test_sign(self, sample_data):
+        """Test sign method."""
+        expr = col("value").sign()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        expected = [1, -1, 0, 1, -1]
+        assert result.tolist() == expected
+
+    def test_abs(self, sample_data):
+        """Test abs method."""
+        expr = col("value").abs()
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, sample_data)
+        expected = [5, 3, 0, 10, 7]
+        assert result.tolist() == expected
+
+    @pytest.mark.parametrize(
+        "base_values,exponent,expected",
+        [
+            ([2, 3, 4], 2, [4, 9, 16]),
+            ([2, 3, 4], 3, [8, 27, 64]),
+            ([4, 9, 16], 0.5, [2.0, 3.0, 4.0]),
+        ],
+        ids=["square", "cube", "sqrt"],
+    )
+    def test_power(self, base_values, exponent, expected):
+        """Test power method with various exponents."""
+        data = pd.DataFrame({"value": base_values})
+        expr = col("value").power(exponent)
+        assert isinstance(expr, UDFExpr)
+        result = eval_expr(expr, data)
+        for r, e in zip(result.tolist(), expected):
+            assert abs(r - e) < 1e-10
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/expressions/test_boolean.py

@@ -0,0 +1,300 @@
+"""Tests for boolean/logical expression operations.
+
+This module tests:
+- Logical operators: AND (&), OR (|), NOT (~)
+- Boolean expression combinations
+- Complex nested boolean expressions
+"""
+
+import pandas as pd
+import pytest
+
+from ray.data._internal.planner.plan_expression.expression_evaluator import eval_expr
+from ray.data.expressions import BinaryExpr, Operation, UnaryExpr, col, lit
+
+# ──────────────────────────────────────
+# Logical AND Operations
+# ──────────────────────────────────────
+
+
+class TestLogicalAnd:
+    """Tests for logical AND (&) operations."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for logical AND tests."""
+        return pd.DataFrame(
+            {
+                "is_active": [True, True, False, False],
+                "is_verified": [True, False, True, False],
+                "age": [25, 17, 30, 15],
+            }
+        )
+
+    def test_and_two_booleans(self, sample_data):
+        """Test AND of two boolean columns."""
+        expr = col("is_active") & col("is_verified")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.AND
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_and_two_comparisons(self, sample_data):
+        """Test AND of two comparison expressions."""
+        expr = (col("is_active") == lit(True)) & (col("age") >= 18)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_and_chained(self, sample_data):
+        """Test chained AND operations."""
+        expr = (col("is_active")) & (col("is_verified")) & (col("age") >= 18)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# Logical OR Operations
+# ──────────────────────────────────────
+
+
+class TestLogicalOr:
+    """Tests for logical OR (|) operations."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for logical OR tests."""
+        return pd.DataFrame(
+            {
+                "is_admin": [True, False, False, False],
+                "is_moderator": [False, True, False, False],
+                "age": [25, 17, 30, 15],
+            }
+        )
+
+    def test_or_two_booleans(self, sample_data):
+        """Test OR of two boolean columns."""
+        expr = col("is_admin") | col("is_moderator")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.OR
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, True, False, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_or_two_comparisons(self, sample_data):
+        """Test OR of two comparison expressions."""
+        expr = (col("is_admin") == lit(True)) | (col("age") >= 18)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_or_chained(self, sample_data):
+        """Test chained OR operations."""
+        expr = (col("is_admin")) | (col("is_moderator")) | (col("age") >= 21)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# Logical NOT Operations
+# ──────────────────────────────────────
+
+
+class TestLogicalNot:
+    """Tests for logical NOT (~) operations."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for logical NOT tests."""
+        return pd.DataFrame(
+            {
+                "is_active": [True, False, True, False],
+                "is_banned": [False, False, True, True],
+                "age": [25, 17, 30, 15],
+            }
+        )
+
+    def test_not_boolean_column(self, sample_data):
+        """Test NOT of a boolean column."""
+        expr = ~col("is_active")
+        assert isinstance(expr, UnaryExpr)
+        assert expr.op == Operation.NOT
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_not_comparison(self, sample_data):
+        """Test NOT of a comparison expression."""
+        expr = ~(col("age") >= 18)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_double_negation(self, sample_data):
+        """Test double negation (~~)."""
+        expr = ~~col("is_active")
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# Complex Boolean Combinations
+# ──────────────────────────────────────
+
+
+class TestComplexBooleanExpressions:
+    """Tests for complex boolean expression combinations."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for complex boolean tests."""
+        return pd.DataFrame(
+            {
+                "age": [17, 21, 25, 30, 65],
+                "is_student": [True, True, False, False, False],
+                "is_member": [False, True, True, False, True],
+                "country": ["USA", "UK", "USA", "Canada", "USA"],
+            }
+        )
+
+    def test_and_or_combination(self, sample_data):
+        """Test combination of AND and OR."""
+        # (age >= 21) AND (is_student OR is_member)
+        expr = (col("age") >= 21) & (col("is_student") | col("is_member"))
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_not_with_and_or(self, sample_data):
+        """Test NOT combined with AND and OR."""
+        # NOT(age < 18) AND (is_member OR is_student)
+        expr = ~(col("age") < 18) & (col("is_member") | col("is_student"))
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_demorgans_law_and(self, sample_data):
+        """Test De Morgan's law: ~(A & B) == (~A) | (~B)."""
+        # ~(is_student & is_member)
+        expr1 = ~(col("is_student") & col("is_member"))
+        # (~is_student) | (~is_member)
+        expr2 = (~col("is_student")) | (~col("is_member"))
+
+        result1 = eval_expr(expr1, sample_data)
+        result2 = eval_expr(expr2, sample_data)
+
+        pd.testing.assert_series_equal(
+            result1.reset_index(drop=True),
+            result2.reset_index(drop=True),
+            check_names=False,
+        )
+
+    def test_demorgans_law_or(self, sample_data):
+        """Test De Morgan's law: ~(A | B) == (~A) & (~B)."""
+        # ~(is_student | is_member)
+        expr1 = ~(col("is_student") | col("is_member"))
+        # (~is_student) & (~is_member)
+        expr2 = (~col("is_student")) & (~col("is_member"))
+
+        result1 = eval_expr(expr1, sample_data)
+        result2 = eval_expr(expr2, sample_data)
+
+        pd.testing.assert_series_equal(
+            result1.reset_index(drop=True),
+            result2.reset_index(drop=True),
+            check_names=False,
+        )
+
+    def test_deeply_nested_boolean(self, sample_data):
+        """Test deeply nested boolean expression."""
+        # ((age >= 21) & (country == "USA")) | ((is_student) & (is_member))
+        expr = ((col("age") >= 21) & (col("country") == "USA")) | (
+            (col("is_student")) & (col("is_member"))
+        )
+        result = eval_expr(expr, sample_data)
+        # Row 0: (17>=21 & USA) | (True & False) = False | False = False
+        # Row 1: (21>=21 & UK) | (True & True) = False | True = True
+        # Row 2: (25>=21 & USA) | (False & True) = True | False = True
+        # Row 3: (30>=21 & Canada) | (False & False) = False | False = False
+        # Row 4: (65>=21 & USA) | (False & True) = True | False = True
+        expected = pd.Series([False, True, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# Boolean Expression Structural Equality
+# ──────────────────────────────────────
+
+
+class TestBooleanStructuralEquality:
+    """Tests for structural equality of boolean expressions."""
+
+    def test_and_structural_equality(self):
+        """Test structural equality for AND expressions."""
+        expr1 = col("a") & col("b")
+        expr2 = col("a") & col("b")
+        expr3 = col("b") & col("a")  # Order matters
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+    def test_or_structural_equality(self):
+        """Test structural equality for OR expressions."""
+        expr1 = col("a") | col("b")
+        expr2 = col("a") | col("b")
+        expr3 = col("a") | col("c")
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+    def test_not_structural_equality(self):
+        """Test structural equality for NOT expressions."""
+        expr1 = ~col("a")
+        expr2 = ~col("a")
+        expr3 = ~col("b")
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+    def test_complex_boolean_structural_equality(self):
+        """Test structural equality for complex boolean expressions."""
+        expr1 = (col("a") > 10) & ((col("b") < 5) | ~col("c"))
+        expr2 = (col("a") > 10) & ((col("b") < 5) | ~col("c"))
+        expr3 = (col("a") > 10) & ((col("b") < 6) | ~col("c"))
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/expressions/test_comparison.py

@@ -0,0 +1,372 @@
+"""Tests for comparison expression operations.
+
+This module tests:
+- Comparison operators: GT (>), LT (<), GE (>=), LE (<=), EQ (==), NE (!=)
+- Comparison with columns and literals
+- Reverse comparisons (literal compared to column)
+"""
+
+import pandas as pd
+import pytest
+
+from ray.data._internal.planner.plan_expression.expression_evaluator import eval_expr
+from ray.data.expressions import BinaryExpr, Operation, col, lit
+
+# ──────────────────────────────────────
+# Basic Comparison Operations
+# ──────────────────────────────────────
+
+
+class TestComparisonOperators:
+    """Tests for comparison operators (>, <, >=, <=, ==, !=)."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for comparison tests."""
+        return pd.DataFrame(
+            {
+                "age": [18, 21, 25, 30, 16],
+                "score": [50, 75, 100, 60, 85],
+                "status": ["active", "inactive", "active", "pending", "active"],
+            }
+        )
+
+    # ── Greater Than ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("age") > 21, [False, False, True, True, False]),
+            (col("age") > col("score") / 10, [True, True, True, True, True]),
+        ],
+        ids=["col_gt_literal", "col_gt_col_expr"],
+    )
+    def test_greater_than(self, sample_data, expr, expected_values):
+        """Test greater than (>) comparisons."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.GT
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values),
+            check_names=False,
+        )
+
+    def test_greater_than_reverse(self, sample_data):
+        """Test reverse greater than (literal > col)."""
+        expr = 22 > col("age")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.LT  # Reverse: 22 > age becomes age < 22
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, True, False, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Less Than ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("age") < 21, [True, False, False, False, True]),
+            (col("score") < 70, [True, False, False, True, False]),
+        ],
+        ids=["col_lt_literal", "score_lt_70"],
+    )
+    def test_less_than(self, sample_data, expr, expected_values):
+        """Test less than (<) comparisons."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.LT
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values),
+            check_names=False,
+        )
+
+    def test_less_than_reverse(self, sample_data):
+        """Test reverse less than (literal < col)."""
+        expr = 20 < col("age")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.GT  # Reverse: 20 < age becomes age > 20
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Greater Than or Equal ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("age") >= 21, [False, True, True, True, False]),
+            (col("score") >= 75, [False, True, True, False, True]),
+        ],
+        ids=["col_ge_21", "score_ge_75"],
+    )
+    def test_greater_equal(self, sample_data, expr, expected_values):
+        """Test greater than or equal (>=) comparisons."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.GE
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values),
+            check_names=False,
+        )
+
+    def test_greater_equal_reverse(self, sample_data):
+        """Test reverse greater equal (literal >= col)."""
+        expr = 21 >= col("age")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.LE  # Reverse
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, True, False, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Less Than or Equal ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("age") <= 21, [True, True, False, False, True]),
+            (col("score") <= 60, [True, False, False, True, False]),
+        ],
+        ids=["col_le_21", "score_le_60"],
+    )
+    def test_less_equal(self, sample_data, expr, expected_values):
+        """Test less than or equal (<=) comparisons."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.LE
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values),
+            check_names=False,
+        )
+
+    def test_less_equal_reverse(self, sample_data):
+        """Test reverse less equal (literal <= col)."""
+        expr = 25 <= col("age")
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.GE  # Reverse
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, False, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    # ── Equality ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("age") == 21, [False, True, False, False, False]),
+            (col("status") == "active", [True, False, True, False, True]),
+            (col("score") == lit(100), [False, False, True, False, False]),
+        ],
+        ids=["age_eq_21", "status_eq_active", "score_eq_100"],
+    )
+    def test_equality(self, sample_data, expr, expected_values):
+        """Test equality (==) comparisons."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.EQ
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values),
+            check_names=False,
+        )
+
+    # ── Not Equal ──
+
+    @pytest.mark.parametrize(
+        "expr,expected_values",
+        [
+            (col("age") != 21, [True, False, True, True, True]),
+            (col("status") != "active", [False, True, False, True, False]),
+        ],
+        ids=["age_ne_21", "status_ne_active"],
+    )
+    def test_not_equal(self, sample_data, expr, expected_values):
+        """Test not equal (!=) comparisons."""
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.NE
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values),
+            check_names=False,
+        )
+
+
+# ──────────────────────────────────────
+# Column vs Column Comparisons
+# ──────────────────────────────────────
+
+
+class TestColumnToColumnComparison:
+    """Tests for comparing columns against other columns."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with comparable columns."""
+        return pd.DataFrame(
+            {
+                "value_a": [10, 20, 30, 40],
+                "value_b": [15, 20, 25, 45],
+                "threshold": [12, 18, 35, 35],
+            }
+        )
+
+    @pytest.mark.parametrize(
+        "expr_fn,expected_values",
+        [
+            (lambda: col("value_a") > col("value_b"), [False, False, True, False]),
+            (lambda: col("value_a") < col("threshold"), [True, False, True, False]),
+            (lambda: col("value_a") == col("value_b"), [False, True, False, False]),
+            (lambda: col("value_a") >= col("threshold"), [False, True, False, True]),
+        ],
+        ids=["a_gt_b", "a_lt_threshold", "a_eq_b", "a_ge_threshold"],
+    )
+    def test_column_to_column_comparisons(self, sample_data, expr_fn, expected_values):
+        """Test various column-to-column comparisons."""
+        expr = expr_fn()
+        result = eval_expr(expr, sample_data)
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True),
+            pd.Series(expected_values),
+            check_names=False,
+        )
+
+
+# ──────────────────────────────────────
+# Comparison with Expressions
+# ──────────────────────────────────────
+
+
+class TestComparisonWithExpressions:
+    """Tests for comparing expressions against other expressions."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for expression comparison tests."""
+        return pd.DataFrame(
+            {
+                "price": [100, 200, 150],
+                "discount": [10, 50, 30],
+                "min_price": [80, 160, 130],
+            }
+        )
+
+    def test_compare_computed_values(self, sample_data):
+        """Test comparing computed expression results."""
+        # (price - discount) > min_price
+        expr = (col("price") - col("discount")) > col("min_price")
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False])  # 90>80, 150>160, 120>130
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_compare_scaled_values(self, sample_data):
+        """Test comparing scaled column values."""
+        # price * 0.9 >= min_price (check if 10% discount still meets minimum)
+        expr = col("price") * 0.9 >= col("min_price")
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, True, True])  # 90>=80, 180>=160, 135>=130
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# String Comparisons
+# ──────────────────────────────────────
+
+
+class TestStringComparison:
+    """Tests for string equality and inequality."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with string columns."""
+        return pd.DataFrame(
+            {
+                "name": ["Alice", "Bob", "Charlie", "Alice"],
+                "city": ["NYC", "LA", "NYC", "SF"],
+            }
+        )
+
+    def test_string_equality(self, sample_data):
+        """Test string equality comparison."""
+        expr = col("name") == "Alice"
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_string_inequality(self, sample_data):
+        """Test string inequality comparison."""
+        expr = col("city") != "NYC"
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# Boolean Comparisons
+# ──────────────────────────────────────
+
+
+class TestBooleanComparison:
+    """Tests for boolean value comparisons."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with boolean columns."""
+        return pd.DataFrame(
+            {
+                "is_active": [True, False, True, False],
+                "is_verified": [True, True, False, False],
+            }
+        )
+
+    def test_boolean_equality_true(self, sample_data):
+        """Test boolean equality with True."""
+        expr = col("is_active") == lit(True)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_boolean_equality_false(self, sample_data):
+        """Test boolean equality with False."""
+        expr = col("is_active") == lit(False)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_boolean_column_to_column(self, sample_data):
+        """Test comparing two boolean columns."""
+        expr = col("is_active") == col("is_verified")
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/expressions/test_conversion.py

@@ -0,0 +1,628 @@
+"""Tests for expression conversion to PyArrow and Iceberg.
+
+This module tests:
+- Conversion to PyArrow compute expressions (to_pyarrow)
+- Conversion to PyIceberg expressions (IcebergExpressionVisitor)
+- Unsupported expression handling
+"""
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.dataset as ds
+import pytest
+from packaging.version import parse as version_parse
+from pyiceberg.expressions import (
+    And,
+    EqualTo,
+    GreaterThan,
+    GreaterThanOrEqual,
+    In,
+    IsNull,
+    LessThan,
+    LessThanOrEqual,
+    Not,
+    NotEqualTo,
+    NotIn,
+    NotNull,
+    Or,
+    Reference,
+    literal,
+)
+
+from ray.data._internal.datasource.iceberg_datasource import _IcebergExpressionVisitor
+from ray.data.datatype import DataType
+from ray.data.expressions import (
+    BinaryExpr,
+    Operation,
+    UDFExpr,
+    col,
+    download,
+    lit,
+    star,
+)
+
+# ──────────────────────────────────────
+# PyArrow Conversion Tests
+# ──────────────────────────────────────
+
+
+class TestToPyArrow:
+    """Test conversion of Ray Data expressions to PyArrow compute expressions."""
+
+    @pytest.fixture
+    def test_table(self):
+        """Sample PyArrow table for testing expressions."""
+        return pa.table(
+            {
+                "age": [15, 25, 45, 70],
+                "x": [1, 2, 3, 4],
+                "price": [10.0, 20.0, 30.0, 40.0],
+                "quantity": [2, 3, 1, 5],
+                "tax": [1.0, 2.0, 3.0, 4.0],
+                "status": ["active", "pending", "inactive", "active"],
+                "value": [1, None, 3, None],
+                "active": [True, False, True, False],
+            }
+        )
+
+    # ── Basic Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_pyarrow_expr",
+        [
+            (col("age"), lambda: pc.field("age")),
+            (lit(42), lambda: pc.scalar(42)),
+            (lit("hello"), lambda: pc.scalar("hello")),
+        ],
+        ids=["col", "int_lit", "str_lit"],
+    )
+    def test_basic_expressions(self, test_table, ray_expr, equivalent_pyarrow_expr):
+        """Test conversion of basic expressions."""
+        converted = ray_expr.to_pyarrow()
+        expected = equivalent_pyarrow_expr()
+        assert converted.equals(expected)
+
+    # ── Arithmetic Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_pyarrow_expr",
+        [
+            (
+                col("x") + 5,
+                lambda: pc.add(pc.field("x"), pc.scalar(5)),
+            ),
+            (
+                col("x") - 3,
+                lambda: pc.subtract(pc.field("x"), pc.scalar(3)),
+            ),
+            (
+                col("x") * 2,
+                lambda: pc.multiply(pc.field("x"), pc.scalar(2)),
+            ),
+            (
+                col("x") / 2,
+                lambda: pc.divide(pc.field("x"), pc.scalar(2)),
+            ),
+        ],
+        ids=["add", "sub", "mul", "div"],
+    )
+    def test_arithmetic_expressions(
+        self, test_table, ray_expr, equivalent_pyarrow_expr
+    ):
+        """Test conversion of arithmetic expressions."""
+        converted = ray_expr.to_pyarrow()
+        expected = equivalent_pyarrow_expr()
+        assert converted.equals(expected)
+
+    # ── Comparison Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_pyarrow_expr,description",
+        [
+            (
+                col("age") > 18,
+                lambda: pc.greater(pc.field("age"), pc.scalar(18)),
+                "greater than",
+            ),
+            (
+                col("age") < 65,
+                lambda: pc.less(pc.field("age"), pc.scalar(65)),
+                "less than",
+            ),
+            (
+                col("age") >= 21,
+                lambda: pc.greater_equal(pc.field("age"), pc.scalar(21)),
+                "greater equal",
+            ),
+            (
+                col("age") <= 30,
+                lambda: pc.less_equal(pc.field("age"), pc.scalar(30)),
+                "less equal",
+            ),
+            (
+                col("status") == "active",
+                lambda: pc.equal(pc.field("status"), pc.scalar("active")),
+                "equality",
+            ),
+            (
+                col("status") != "deleted",
+                lambda: pc.not_equal(pc.field("status"), pc.scalar("deleted")),
+                "not equal",
+            ),
+        ],
+        ids=["gt", "lt", "ge", "le", "eq", "ne"],
+    )
+    def test_comparison_expressions(
+        self, test_table, ray_expr, equivalent_pyarrow_expr, description
+    ):
+        """Test conversion of comparison expressions."""
+        converted = ray_expr.to_pyarrow()
+        expected = equivalent_pyarrow_expr()
+
+        # Verify they produce the same results on sample data
+        dataset = ds.dataset(test_table)
+        try:
+            result_converted = dataset.scanner(filter=converted).to_table()
+            result_expected = dataset.scanner(filter=expected).to_table()
+            assert result_converted.equals(result_expected)
+        except (TypeError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError):
+            pass
+
+    # ── Boolean Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_pyarrow_expr,description",
+        [
+            (
+                (col("age") > 18) & (col("age") < 65),
+                lambda: pc.and_kleene(
+                    pc.greater(pc.field("age"), pc.scalar(18)),
+                    pc.less(pc.field("age"), pc.scalar(65)),
+                ),
+                "logical AND",
+            ),
+            (
+                (col("status") == "active") | (col("status") == "pending"),
+                lambda: pc.or_kleene(
+                    pc.equal(pc.field("status"), pc.scalar("active")),
+                    pc.equal(pc.field("status"), pc.scalar("pending")),
+                ),
+                "logical OR",
+            ),
+            (
+                ~col("active"),
+                lambda: pc.invert(pc.field("active")),
+                "logical NOT",
+            ),
+        ],
+        ids=["and", "or", "not"],
+    )
+    def test_boolean_expressions(
+        self, test_table, ray_expr, equivalent_pyarrow_expr, description
+    ):
+        """Test conversion of boolean expressions."""
+        converted = ray_expr.to_pyarrow()
+        expected = equivalent_pyarrow_expr()
+
+        dataset = ds.dataset(test_table)
+        try:
+            result_converted = dataset.scanner(filter=converted).to_table()
+            result_expected = dataset.scanner(filter=expected).to_table()
+            assert result_converted.equals(result_expected)
+        except (TypeError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError):
+            pass
+
+    # ── Predicate Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_pyarrow_expr,description",
+        [
+            (
+                col("value").is_null(),
+                lambda: pc.is_null(pc.field("value")),
+                "is_null check",
+            ),
+            (
+                col("value").is_not_null(),
+                lambda: pc.is_valid(pc.field("value")),
+                "is_not_null check",
+            ),
+            (
+                col("status").is_in(["active", "pending"]),
+                lambda: pc.is_in(pc.field("status"), pa.array(["active", "pending"])),
+                "is_in with list",
+            ),
+        ],
+        ids=["is_null", "is_not_null", "is_in"],
+    )
+    def test_predicate_expressions(
+        self, test_table, ray_expr, equivalent_pyarrow_expr, description
+    ):
+        """Test conversion of predicate expressions."""
+        converted = ray_expr.to_pyarrow()
+        expected = equivalent_pyarrow_expr()
+
+        dataset = ds.dataset(test_table)
+        try:
+            result_converted = dataset.scanner(filter=converted).to_table()
+            result_expected = dataset.scanner(filter=expected).to_table()
+            assert result_converted.equals(result_expected)
+        except (TypeError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError):
+            pass
+
+    # ── Nested Expressions ──
+
+    def test_nested_arithmetic(self, test_table):
+        """Test nested arithmetic expressions."""
+        ray_expr = (col("price") * col("quantity")) + col("tax")
+        converted = ray_expr.to_pyarrow()
+        assert isinstance(converted, pc.Expression)
+
+    # ── Alias Expressions ──
+
+    def test_alias_expressions(self, test_table):
+        """Test that alias expressions unwrap to inner expression."""
+        ray_expr = (col("x") + 5).alias("result")
+        converted = ray_expr.to_pyarrow()
+        assert isinstance(converted, pc.Expression)
+
+    # ── PyArrow Compute UDF Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_pyarrow_expr",
+        [
+            pytest.param(
+                col("name").str.match_regex("foo.*bar"),
+                lambda: pc.match_substring_regex(pc.field("name"), "foo.*bar"),
+                id="match_regex",
+            ),
+            pytest.param(
+                col("name").str.starts_with("foo"),
+                lambda: pc.starts_with(pc.field("name"), "foo"),
+                id="starts_with",
+            ),
+            pytest.param(
+                col("name").str.ends_with("bar"),
+                lambda: pc.ends_with(pc.field("name"), "bar"),
+                id="ends_with",
+            ),
+            pytest.param(
+                col("name").str.contains("baz"),
+                lambda: pc.match_substring(pc.field("name"), "baz"),
+                id="contains",
+            ),
+            pytest.param(
+                col("name").str.upper(),
+                lambda: pc.utf8_upper(pc.field("name")),
+                id="upper",
+            ),
+            pytest.param(
+                col("x").ceil(),
+                lambda: pc.ceil(pc.field("x")),
+                id="ceil",
+            ),
+            pytest.param(
+                col("x").abs(),
+                lambda: pc.abs_checked(pc.field("x")),
+                id="abs",
+            ),
+        ],
+    )
+    def test_pyarrow_compute_udf_expressions(
+        self, test_table, ray_expr, equivalent_pyarrow_expr
+    ):
+        """Test that PyArrow-compute-backed UDFs convert to PyArrow expressions."""
+        converted = ray_expr.to_pyarrow()
+        expected = equivalent_pyarrow_expr()
+        assert converted.equals(expected)
+
+    @pytest.mark.skipif(
+        version_parse(pa.__version__) < version_parse("19.0.0"),
+        reason="Requires PyArrow >= 19 for string compute UDF pushdown",
+    )
+    def test_negated_pyarrow_compute_udf(self, test_table):
+        """Test that negated PyArrow compute UDF expressions convert correctly."""
+        ray_expr = ~col("status").str.match_regex("act.*")
+        converted = ray_expr.to_pyarrow()
+        assert isinstance(converted, pc.Expression)
+
+        dataset = ds.dataset(test_table)
+        result = dataset.to_table(filter=converted)
+        assert all(
+            not bool(pc.match_substring_regex(s, "act.*"))
+            for s in result.column("status").to_pylist()
+        )
+
+    def test_pyarrow_compute_udf_as_dataset_filter(self, test_table):
+        """Test that converted UDF expressions work as dataset scan filters."""
+        ray_expr = col("status").str.match_regex("^active$")
+        pa_expr = ray_expr.to_pyarrow()
+
+        dataset = ds.dataset(test_table)
+        result = dataset.to_table(filter=pa_expr)
+        assert all(s == "active" for s in result.column("status").to_pylist())
+
+    # ── Unsupported Expressions ──
+
+    def test_user_udf_expression_raises(self):
+        """Test that user-defined UDF expressions raise TypeError."""
+
+        def dummy_fn(x):
+            return x
+
+        udf_expr = UDFExpr(
+            fn=dummy_fn,
+            args=[col("x")],
+            kwargs={},
+            data_type=DataType(int),
+        )
+
+        with pytest.raises(TypeError, match="UDF expressions cannot be converted"):
+            udf_expr.to_pyarrow()
+
+    def test_download_expression_raises(self):
+        """Test that download expressions raise TypeError."""
+        with pytest.raises(TypeError, match="Download expressions cannot be converted"):
+            download("uri").to_pyarrow()
+
+    def test_star_expression_raises(self):
+        """Test that star expressions raise TypeError."""
+        with pytest.raises(TypeError, match="Star expressions cannot be converted"):
+            star().to_pyarrow()
+
+
+# ──────────────────────────────────────
+# Iceberg Conversion Tests
+# ──────────────────────────────────────
+
+
+class TestIcebergExpressionVisitor:
+    """Test conversion of Ray Data expressions to PyIceberg expressions."""
+
+    # ── Basic Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_iceberg_expr,description",
+        [
+            (col("age"), lambda: Reference("age"), "column reference"),
+            (lit(42), lambda: literal(42), "integer literal"),
+            (lit("active"), lambda: literal("active"), "string literal"),
+        ],
+        ids=["col", "int_lit", "str_lit"],
+    )
+    def test_basic_expressions(self, ray_expr, equivalent_iceberg_expr, description):
+        """Test conversion of basic expressions."""
+        visitor = _IcebergExpressionVisitor()
+        converted = visitor.visit(ray_expr)
+        expected = equivalent_iceberg_expr()
+        assert converted == expected
+
+    # ── Comparison Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_iceberg_expr,description",
+        [
+            (
+                col("age") > 18,
+                lambda: GreaterThan(Reference("age"), literal(18)),
+                "greater than",
+            ),
+            (
+                col("age") >= 21,
+                lambda: GreaterThanOrEqual(Reference("age"), literal(21)),
+                "greater than or equal",
+            ),
+            (
+                col("age") < 65,
+                lambda: LessThan(Reference("age"), literal(65)),
+                "less than",
+            ),
+            (
+                col("age") <= 100,
+                lambda: LessThanOrEqual(Reference("age"), literal(100)),
+                "less than or equal",
+            ),
+            (
+                col("status") == "active",
+                lambda: EqualTo(Reference("status"), literal("active")),
+                "equality",
+            ),
+            (
+                col("status") != "inactive",
+                lambda: NotEqualTo(Reference("status"), literal("inactive")),
+                "not equal",
+            ),
+        ],
+        ids=["gt", "ge", "lt", "le", "eq", "ne"],
+    )
+    def test_comparison_expressions(
+        self, ray_expr, equivalent_iceberg_expr, description
+    ):
+        """Test conversion of comparison expressions."""
+        visitor = _IcebergExpressionVisitor()
+        converted = visitor.visit(ray_expr)
+        expected = equivalent_iceberg_expr()
+        assert converted == expected
+
+    # ── Boolean Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_iceberg_expr,description",
+        [
+            (
+                (col("age") > 18) & (col("age") < 65),
+                lambda: And(
+                    GreaterThan(Reference("age"), literal(18)),
+                    LessThan(Reference("age"), literal(65)),
+                ),
+                "logical AND",
+            ),
+            (
+                (col("is_member") == lit(True)) | (col("is_premium") == lit(True)),
+                lambda: Or(
+                    EqualTo(Reference("is_member"), literal(True)),
+                    EqualTo(Reference("is_premium"), literal(True)),
+                ),
+                "logical OR",
+            ),
+            (
+                ~(col("deleted") == lit(True)),
+                lambda: Not(EqualTo(Reference("deleted"), literal(True))),
+                "logical NOT",
+            ),
+        ],
+        ids=["and", "or", "not"],
+    )
+    def test_boolean_expressions(self, ray_expr, equivalent_iceberg_expr, description):
+        """Test conversion of boolean expressions."""
+        visitor = _IcebergExpressionVisitor()
+        converted = visitor.visit(ray_expr)
+        expected = equivalent_iceberg_expr()
+        assert converted == expected
+
+    # ── Predicate Expressions ──
+
+    @pytest.mark.parametrize(
+        "ray_expr,equivalent_iceberg_expr,description",
+        [
+            (
+                col("value").is_null(),
+                lambda: IsNull(Reference("value")),
+                "is_null check",
+            ),
+            (
+                col("name").is_not_null(),
+                lambda: NotNull(Reference("name")),
+                "is_not_null check",
+            ),
+            (
+                col("status").is_in(["active", "pending"]),
+                lambda: In(Reference("status"), ["active", "pending"]),
+                "is_in with list",
+            ),
+            (
+                col("status").not_in(["inactive", "deleted"]),
+                lambda: NotIn(Reference("status"), ["inactive", "deleted"]),
+                "not_in with list",
+            ),
+        ],
+        ids=["is_null", "is_not_null", "is_in", "not_in"],
+    )
+    def test_predicate_expressions(
+        self, ray_expr, equivalent_iceberg_expr, description
+    ):
+        """Test conversion of predicate expressions."""
+        visitor = _IcebergExpressionVisitor()
+        converted = visitor.visit(ray_expr)
+        expected = equivalent_iceberg_expr()
+        assert converted == expected
+
+    # ── Complex Nested Expressions ──
+
+    def test_complex_nested_boolean(self):
+        """Test complex nested boolean expression."""
+        ray_expr = (
+            (col("age") >= 21)
+            & (col("country") == "USA")
+            & col("verified").is_not_null()
+        )
+        expected = And(
+            And(
+                GreaterThanOrEqual(Reference("age"), literal(21)),
+                EqualTo(Reference("country"), literal("USA")),
+            ),
+            NotNull(Reference("verified")),
+        )
+
+        visitor = _IcebergExpressionVisitor()
+        converted = visitor.visit(ray_expr)
+        assert converted == expected
+
+    def test_aliased_expression(self):
+        """Test that alias expressions unwrap to inner expression."""
+        ray_expr = (col("age") > 18).alias("is_adult")
+        expected = GreaterThan(Reference("age"), literal(18))
+
+        visitor = _IcebergExpressionVisitor()
+        converted = visitor.visit(ray_expr)
+        assert converted == expected
+
+    # ── Unsupported Arithmetic ──
+
+    @pytest.mark.parametrize(
+        "ray_expr",
+        [
+            col("price") + 10,
+            col("quantity") * 2,
+            col("total") - col("discount"),
+            col("revenue") / col("count"),
+            col("items") // 5,
+        ],
+        ids=["add", "mul", "sub", "div", "floordiv"],
+    )
+    def test_arithmetic_raises(self, ray_expr):
+        """Test that arithmetic operations raise appropriate errors."""
+        visitor = _IcebergExpressionVisitor()
+        with pytest.raises(
+            ValueError, match="Unsupported binary operation for Iceberg"
+        ):
+            visitor.visit(ray_expr)
+
+    # ── Unsupported Expressions ──
+
+    def test_udf_expression_raises(self):
+        """Test that UDF expressions raise TypeError."""
+
+        def dummy_fn(x):
+            return x
+
+        udf_expr = UDFExpr(
+            fn=dummy_fn,
+            args=[col("x")],
+            kwargs={},
+            data_type=DataType(int),
+        )
+
+        visitor = _IcebergExpressionVisitor()
+        with pytest.raises(
+            TypeError, match="UDF expressions cannot be converted to Iceberg"
+        ):
+            visitor.visit(udf_expr)
+
+    def test_download_expression_raises(self):
+        """Test that download expressions raise TypeError."""
+        visitor = _IcebergExpressionVisitor()
+        with pytest.raises(
+            TypeError, match="Download expressions cannot be converted to Iceberg"
+        ):
+            visitor.visit(download("uri"))
+
+    def test_star_expression_raises(self):
+        """Test that star expressions raise TypeError."""
+        visitor = _IcebergExpressionVisitor()
+        with pytest.raises(
+            TypeError, match="Star expressions cannot be converted to Iceberg"
+        ):
+            visitor.visit(star())
+
+    def test_is_in_requires_literal_list(self):
+        """Test that IN/NOT_IN operations require literal lists."""
+        visitor = _IcebergExpressionVisitor()
+
+        # This should work - literal list
+        expr = col("status").is_in(["active", "pending"])
+        result = visitor.visit(expr)
+        assert isinstance(result, In)
+
+        # This should fail - column reference on right side
+        with pytest.raises(
+            ValueError, match="IN operation requires right operand to be a literal list"
+        ):
+            invalid_expr = BinaryExpr(Operation.IN, col("a"), col("b"))
+            visitor.visit(invalid_expr)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/expressions/test_core.py

@@ -0,0 +1,537 @@
+"""Tests for core expression types and basic functionality.
+
+This module tests:
+- ColumnExpr, LiteralExpr, BinaryExpr, UnaryExpr, AliasExpr, StarExpr
+- Structural equality for all expression types
+- Expression tree repr (string representation)
+- UDFExpr structural equality
+"""
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
+
+from ray.data._internal.planner.plan_expression.expression_visitors import (
+    _InlineExprReprVisitor,
+)
+from ray.data.datatype import DataType
+from ray.data.expressions import (
+    BinaryExpr,
+    ColumnExpr,
+    Expr,
+    LiteralExpr,
+    Operation,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+    col,
+    download,
+    lit,
+    star,
+    udf,
+)
+
+# ──────────────────────────────────────
+# Column Expression Tests
+# ──────────────────────────────────────
+
+
+class TestColumnExpr:
+    """Tests for ColumnExpr functionality."""
+
+    def test_column_creation(self):
+        """Test that col() creates a ColumnExpr with correct name."""
+        expr = col("age")
+        assert isinstance(expr, ColumnExpr)
+        assert expr.name == "age"
+
+    def test_column_name_property(self):
+        """Test that name property returns the column name."""
+        expr = col("my_column")
+        assert expr.name == "my_column"
+
+    @pytest.mark.parametrize(
+        "name1,name2,expected",
+        [
+            ("a", "a", True),
+            ("a", "b", False),
+            ("column_name", "column_name", True),
+            ("COL", "col", False),  # Case sensitive
+        ],
+        ids=["same_name", "different_name", "long_name", "case_sensitive"],
+    )
+    def test_column_structural_equality(self, name1, name2, expected):
+        """Test structural equality for column expressions."""
+        assert col(name1).structurally_equals(col(name2)) is expected
+
+
+# ──────────────────────────────────────
+# Literal Expression Tests
+# ──────────────────────────────────────
+
+
+class TestLiteralExpr:
+    """Tests for LiteralExpr functionality."""
+
+    @pytest.mark.parametrize(
+        "value",
+        [42, 3.14, "hello", True, False, None, [1, 2, 3]],
+        ids=["int", "float", "string", "bool_true", "bool_false", "none", "list"],
+    )
+    def test_literal_creation(self, value):
+        """Test that lit() creates a LiteralExpr with correct value."""
+        expr = lit(value)
+        assert isinstance(expr, LiteralExpr)
+        assert expr.value == value
+
+    @pytest.mark.parametrize(
+        "val1,val2,expected",
+        [
+            (1, 1, True),
+            (1, 2, False),
+            ("x", "y", False),
+            ("x", "x", True),
+            (1, 1.0, False),  # Different types
+            (True, True, True),
+            (True, False, False),
+            ([1, 2], [1, 2], True),
+            ([1, 2], [1, 3], False),
+        ],
+        ids=[
+            "same_int",
+            "different_int",
+            "different_str",
+            "same_str",
+            "int_vs_float",
+            "same_bool",
+            "different_bool",
+            "same_list",
+            "different_list",
+        ],
+    )
+    def test_literal_structural_equality(self, val1, val2, expected):
+        """Test structural equality for literal expressions."""
+        assert lit(val1).structurally_equals(lit(val2)) is expected
+
+
+# ──────────────────────────────────────
+# Binary Expression Tests
+# ──────────────────────────────────────
+
+
+class TestBinaryExpr:
+    """Tests for BinaryExpr structure (not operation semantics)."""
+
+    def test_binary_expression_structure(self):
+        """Test that binary expressions have correct structure."""
+        expr = col("a") + lit(1)
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.ADD
+        assert isinstance(expr.left, ColumnExpr)
+        assert isinstance(expr.right, LiteralExpr)
+
+    @pytest.mark.parametrize(
+        "expr1,expr2,expected",
+        [
+            (col("a") + 1, col("a") + 1, True),
+            (col("a") + 1, col("a") + 2, False),  # Different literal
+            (col("a") + 1, col("b") + 1, False),  # Different column
+            (col("a") + 1, col("a") - 1, False),  # Different operator
+            # Nested binary expressions
+            ((col("a") * 2) + (col("b") / 3), (col("a") * 2) + (col("b") / 3), True),
+            ((col("a") * 2) + (col("b") / 3), (col("a") * 2) - (col("b") / 3), False),
+            ((col("a") * 2) + (col("b") / 3), (col("c") * 2) + (col("b") / 3), False),
+            ((col("a") * 2) + (col("b") / 3), (col("a") * 2) + (col("b") / 4), False),
+            # Commutative operations are not structurally equal
+            (col("a") + col("b"), col("b") + col("a"), False),
+            (lit(1) * col("c"), col("c") * lit(1), False),
+        ],
+        ids=[
+            "same_simple",
+            "different_literal",
+            "different_column",
+            "different_operator",
+            "same_nested",
+            "nested_diff_op",
+            "nested_diff_col",
+            "nested_diff_lit",
+            "commutative_add",
+            "commutative_mul",
+        ],
+    )
+    def test_binary_structural_equality(self, expr1, expr2, expected):
+        """Test structural equality for binary expressions."""
+        assert expr1.structurally_equals(expr2) is expected
+        # Test symmetry
+        assert expr2.structurally_equals(expr1) is expected
+
+
+# ──────────────────────────────────────
+# Unary Expression Tests
+# ──────────────────────────────────────
+
+
+class TestUnaryExpr:
+    """Tests for UnaryExpr structure."""
+
+    @pytest.mark.parametrize(
+        "expr,expected_op",
+        [
+            (col("age").is_null(), Operation.IS_NULL),
+            (col("name").is_not_null(), Operation.IS_NOT_NULL),
+            (~col("active"), Operation.NOT),
+        ],
+        ids=["is_null", "is_not_null", "not"],
+    )
+    def test_unary_expression_structure(self, expr, expected_op):
+        """Test that unary expressions have correct structure."""
+        assert isinstance(expr, UnaryExpr)
+        assert expr.op == expected_op
+        assert isinstance(expr.operand, Expr)
+
+    def test_unary_structural_equality(self):
+        """Test structural equality for unary expressions."""
+        # Same expressions should be equal
+        assert col("age").is_null().structurally_equals(col("age").is_null())
+        assert (
+            col("active").is_not_null().structurally_equals(col("active").is_not_null())
+        )
+        assert (~col("flag")).structurally_equals(~col("flag"))
+
+        # Different operations should not be equal
+        assert not col("age").is_null().structurally_equals(col("age").is_not_null())
+
+        # Different operands should not be equal
+        assert not col("age").is_null().structurally_equals(col("name").is_null())
+
+
+# ──────────────────────────────────────
+# Alias Expression Tests
+# ──────────────────────────────────────
+
+
+class TestAliasExpr:
+    """Tests for AliasExpr functionality."""
+
+    @pytest.mark.parametrize(
+        "expr,alias_name,expected_alias",
+        [
+            (col("price"), "product_price", "product_price"),
+            (lit(42), "answer", "answer"),
+            (col("a") + col("b"), "sum", "sum"),
+            ((col("price") * col("qty")) + lit(5), "total_with_fee", "total_with_fee"),
+            (col("age") >= lit(18), "is_adult", "is_adult"),
+        ],
+        ids=["col_alias", "lit_alias", "binary_alias", "complex_alias", "comparison"],
+    )
+    def test_alias_functionality(self, expr, alias_name, expected_alias):
+        """Test alias creation and properties."""
+        aliased_expr = expr.alias(alias_name)
+        assert aliased_expr.name == expected_alias
+        assert aliased_expr.expr.structurally_equals(expr)
+        # Data type should be preserved
+        assert aliased_expr.data_type == expr.data_type
+
+    @pytest.mark.parametrize(
+        "expr1,expr2,expected",
+        [
+            (col("a").alias("b"), col("a").alias("b"), True),
+            (col("a").alias("b"), col("a").alias("c"), False),  # Different alias
+            (col("a").alias("b"), col("b").alias("b"), False),  # Different column
+            ((col("a") + 1).alias("result"), (col("a") + 1).alias("result"), True),
+            (
+                (col("a") + 1).alias("result"),
+                (col("a") + 2).alias("result"),
+                False,
+            ),  # Different expr
+            (col("a").alias("b"), col("a"), False),  # Alias vs non-alias
+        ],
+        ids=[
+            "same_alias",
+            "different_alias_name",
+            "different_column",
+            "same_complex",
+            "different_expr",
+            "alias_vs_non_alias",
+        ],
+    )
+    def test_alias_structural_equality(self, expr1, expr2, expected):
+        """Test structural equality for alias expressions."""
+        assert expr1.structurally_equals(expr2) is expected
+
+    def test_alias_structural_equality_respects_rename_flag(self):
+        expr = col("a")
+        aliased = expr.alias("b")
+        renamed = expr._rename("b")
+
+        assert aliased.structurally_equals(aliased)
+        assert renamed.structurally_equals(renamed)
+        assert not aliased.structurally_equals(renamed)
+        assert not aliased.structurally_equals(expr.alias("c"))
+
+    def test_alias_evaluation_equivalence(self):
+        """Test that alias evaluation produces same result as original."""
+        import pandas as pd
+
+        from ray.data._internal.planner.plan_expression.expression_evaluator import (
+            eval_expr,
+        )
+
+        test_data = pd.DataFrame({"price": [10, 20], "qty": [2, 3]})
+        expr = col("price") * col("qty")
+        aliased = expr.alias("total")
+
+        original_result = eval_expr(expr, test_data)
+        aliased_result = eval_expr(aliased, test_data)
+
+        assert original_result.equals(aliased_result)
+
+
+# ──────────────────────────────────────
+# Star Expression Tests
+# ──────────────────────────────────────
+
+
+class TestStarExpr:
+    """Tests for StarExpr functionality."""
+
+    def test_star_creation(self):
+        """Test that star() creates a StarExpr."""
+        expr = star()
+        assert isinstance(expr, StarExpr)
+
+    def test_star_structural_equality(self):
+        """Test structural equality for star expressions."""
+        assert star().structurally_equals(star())
+        assert not star().structurally_equals(col("a"))
+
+
+# ──────────────────────────────────────
+# UDF Expression Tests
+# ──────────────────────────────────────
+
+
+class TestUDFExpr:
+    """Tests for UDFExpr structural equality."""
+
+    def test_regular_function_udf_structural_equality(self):
+        """Test that regular function UDFs compare fn correctly."""
+
+        @udf(return_dtype=DataType.int32())
+        def add_one(x: pa.Array) -> pa.Array:
+            return pc.add(x, 1)
+
+        @udf(return_dtype=DataType.int32())
+        def add_two(x: pa.Array) -> pa.Array:
+            return pc.add(x, 2)
+
+        expr1 = add_one(col("value"))
+        expr2 = add_one(col("value"))
+        expr3 = add_two(col("value"))
+
+        # Same function should be equal
+        assert expr1.structurally_equals(expr2)
+
+        # Different functions should not be equal
+        assert not expr1.structurally_equals(expr3)
+
+    def test_callable_class_udf_structural_equality(self):
+        """Test that callable class UDFs with same spec are structurally equal."""
+
+        @udf(return_dtype=DataType.int32())
+        class AddOffset:
+            def __init__(self, offset):
+                self.offset = offset
+
+            def __call__(self, x: pa.Array) -> pa.Array:
+                return pc.add(x, self.offset)
+
+        # Create the same callable class instance
+        add_five = AddOffset(5)
+
+        # Each call creates a new _placeholder function internally,
+        # but the callable_class_spec should be the same
+        expr1 = add_five(col("value"))
+        expr2 = add_five(col("value"))
+
+        # These should be structurally equal
+        assert expr1.structurally_equals(expr2)
+        assert expr2.structurally_equals(expr1)
+
+        # Different constructor args should not be equal
+        add_ten = AddOffset(10)
+        expr3 = add_ten(col("value"))
+        assert not expr1.structurally_equals(expr3)
+
+        # Different column args should not be equal
+        expr4 = add_five(col("other"))
+        assert not expr1.structurally_equals(expr4)
+
+    def test_callable_class_vs_regular_function_udf(self):
+        """Test that callable class UDFs are not equal to regular function UDFs."""
+
+        @udf(return_dtype=DataType.int32())
+        class AddOne:
+            def __call__(self, x: pa.Array) -> pa.Array:
+                return pc.add(x, 1)
+
+        @udf(return_dtype=DataType.int32())
+        def add_one(x: pa.Array) -> pa.Array:
+            return pc.add(x, 1)
+
+        class_expr = AddOne()(col("value"))
+        func_expr = add_one(col("value"))
+
+        # Different types of UDFs should not be equal
+        assert not class_expr.structurally_equals(func_expr)
+        assert not func_expr.structurally_equals(class_expr)
+
+
+# ──────────────────────────────────────
+# Cross-type Equality Tests
+# ──────────────────────────────────────
+
+
+class TestCrossTypeEquality:
+    """Test that different expression types are not structurally equal."""
+
+    @pytest.mark.parametrize(
+        "expr1,expr2",
+        [
+            (col("a"), lit("a")),
+            (col("a"), col("a") + 0),
+            (lit(1), lit(1) + 0),
+            (col("a"), col("a").alias("a")),
+            (col("a"), star()),
+        ],
+        ids=[
+            "col_vs_lit",
+            "col_vs_binary",
+            "lit_vs_binary",
+            "col_vs_alias",
+            "col_vs_star",
+        ],
+    )
+    def test_different_types_not_equal(self, expr1, expr2):
+        """Test that different expression types are not structurally equal."""
+        assert not expr1.structurally_equals(expr2)
+        assert not expr2.structurally_equals(expr1)
+
+    def test_operator_eq_is_not_structural_eq(self):
+        """Confirms that == builds an expression, while structurally_equals compares."""
+        # `==` returns a BinaryExpr, not a boolean
+        op_eq_expr = col("a") == col("a")
+        assert isinstance(op_eq_expr, Expr)
+        assert not isinstance(op_eq_expr, bool)
+
+        # `structurally_equals` returns a boolean
+        struct_eq_result = col("a").structurally_equals(col("a"))
+        assert isinstance(struct_eq_result, bool)
+        assert struct_eq_result is True
+
+
+# ──────────────────────────────────────
+# Expression Repr Tests
+# ──────────────────────────────────────
+
+
+def _build_complex_expr():
+    """Build a convoluted expression that exercises all visitor code paths."""
+
+    def custom_udf(x, y):
+        return x + y
+
+    # Create UDF expression
+    udf_expr = UDFExpr(
+        fn=custom_udf,
+        args=[col("value"), lit(10)],
+        kwargs={"z": col("multiplier")},
+        data_type=DataType(int),
+    )
+
+    # Build the mega-complex expression
+    inner_expr = (
+        ((col("age") + lit(10)) * col("rate") / lit(2.5) >= lit(100))
+        & (
+            col("name").is_not_null()
+            | (col("status").is_in(["active", "pending"]) & col("verified"))
+        )
+        & ((col("count") - lit(5)) // lit(2) <= col("limit"))
+        & ~(col("deleted").is_null() | (col("score") != lit(0)))
+        & (download("uri") < star())
+        & (udf_expr.alias("udf_result") > lit(50))
+    ).alias("complex_filter")
+
+    return ~inner_expr
+
+
+class TestExpressionRepr:
+    """Test expression string representations."""
+
+    def test_tree_repr(self):
+        """Test tree representation of expressions."""
+        expr = _build_complex_expr()
+        expected = """NOT
+    └── operand: ALIAS('complex_filter')
+        └── AND
+            ├── left: AND
+            │   ├── left: AND
+            │   │   ├── left: AND
+            │   │   │   ├── left: AND
+            │   │   │   │   ├── left: GE
+            │   │   │   │   │   ├── left: DIV
+            │   │   │   │   │   │   ├── left: MUL
+            │   │   │   │   │   │   │   ├── left: ADD
+            │   │   │   │   │   │   │   │   ├── left: COL('age')
+            │   │   │   │   │   │   │   │   └── right: LIT(10)
+            │   │   │   │   │   │   │   └── right: COL('rate')
+            │   │   │   │   │   │   └── right: LIT(2.5)
+            │   │   │   │   │   └── right: LIT(100)
+            │   │   │   │   └── right: OR
+            │   │   │   │       ├── left: IS_NOT_NULL
+            │   │   │   │       │   └── operand: COL('name')
+            │   │   │   │       └── right: AND
+            │   │   │   │           ├── left: IN
+            │   │   │   │           │   ├── left: COL('status')
+            │   │   │   │           │   └── right: LIT(['active', 'pending'])
+            │   │   │   │           └── right: COL('verified')
+            │   │   │   └── right: LE
+            │   │   │       ├── left: FLOORDIV
+            │   │   │       │   ├── left: SUB
+            │   │   │       │   │   ├── left: COL('count')
+            │   │   │       │   │   └── right: LIT(5)
+            │   │   │       │   └── right: LIT(2)
+            │   │   │       └── right: COL('limit')
+            │   │   └── right: NOT
+            │   │       └── operand: OR
+            │   │           ├── left: IS_NULL
+            │   │           │   └── operand: COL('deleted')
+            │   │           └── right: NE
+            │   │               ├── left: COL('score')
+            │   │               └── right: LIT(0)
+            │   └── right: LT
+            │       ├── left: DOWNLOAD('uri')
+            │       └── right: COL(*)
+            └── right: GT
+                ├── left: ALIAS('udf_result')
+                │   └── UDF(custom_udf)
+                │       ├── arg[0]: COL('value')
+                │       ├── arg[1]: LIT(10)
+                │       └── kwarg['z']: COL('multiplier')
+                └── right: LIT(50)"""
+        assert repr(expr) == expected
+
+    def test_inline_repr_prefix(self):
+        """Test that inline representation starts correctly."""
+        expr = _build_complex_expr()
+        visitor = _InlineExprReprVisitor()
+        inline_repr = visitor.visit(expr)
+        expected_prefix = "~((((((((col('age') + 10) * col('rate')) / 2.5) >= 100) & (col('name').is_not_null() | ((col('status')"
+        assert inline_repr.startswith(expected_prefix)
+        assert inline_repr.endswith(".alias('complex_filter')")
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/expressions/test_namespace_list.py

@@ -0,0 +1,23 @@
+"""Unit tests for list namespace expressions.
+
+These tests verify expression construction logic without requiring Ray.
+"""
+
+import pytest
+
+from ray.data.expressions import col
+
+
+class TestListNamespaceErrors:
+    """Tests for proper error handling in list namespace."""
+
+    def test_list_invalid_index_type(self):
+        """Test list bracket notation rejects invalid types."""
+        with pytest.raises(TypeError, match="List indices must be integers or slices"):
+            col("items").list["invalid"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/expressions/test_predicate.py

@@ -0,0 +1,352 @@
+"""Tests for predicate expression operations.
+
+This module tests:
+- Null predicates: is_null(), is_not_null()
+- Membership predicates: is_in(), not_in()
+"""
+
+import pandas as pd
+import pytest
+
+from ray.data._internal.planner.plan_expression.expression_evaluator import eval_expr
+from ray.data.expressions import BinaryExpr, Operation, UnaryExpr, col, lit
+
+# ──────────────────────────────────────
+# Null Predicate Operations
+# ──────────────────────────────────────
+
+
+class TestIsNull:
+    """Tests for is_null() predicate."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with null values for null predicate tests."""
+        return pd.DataFrame(
+            {
+                "value": [1.0, None, 3.0, None, 5.0],
+                "name": ["Alice", None, "Charlie", "Diana", None],
+            }
+        )
+
+    def test_is_null_numeric(self, sample_data):
+        """Test is_null on numeric column."""
+        expr = col("value").is_null()
+        assert isinstance(expr, UnaryExpr)
+        assert expr.op == Operation.IS_NULL
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, False, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_null_string(self, sample_data):
+        """Test is_null on string column."""
+        expr = col("name").is_null()
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, False, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_null_structural_equality(self):
+        """Test structural equality for is_null expressions."""
+        expr1 = col("value").is_null()
+        expr2 = col("value").is_null()
+        expr3 = col("other").is_null()
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+
+class TestIsNotNull:
+    """Tests for is_not_null() predicate."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with null values."""
+        return pd.DataFrame(
+            {
+                "value": [1.0, None, 3.0, None, 5.0],
+                "name": ["Alice", None, "Charlie", "Diana", None],
+            }
+        )
+
+    def test_is_not_null_numeric(self, sample_data):
+        """Test is_not_null on numeric column."""
+        expr = col("value").is_not_null()
+        assert isinstance(expr, UnaryExpr)
+        assert expr.op == Operation.IS_NOT_NULL
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_not_null_string(self, sample_data):
+        """Test is_not_null on string column."""
+        expr = col("name").is_not_null()
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_not_null_structural_equality(self):
+        """Test structural equality for is_not_null expressions."""
+        expr1 = col("value").is_not_null()
+        expr2 = col("value").is_not_null()
+        expr3 = col("other").is_not_null()
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+
+class TestNullPredicateCombinations:
+    """Tests for null predicates combined with other operations."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with null values and other columns."""
+        return pd.DataFrame(
+            {
+                "value": [10.0, None, 30.0, None, 50.0],
+                "threshold": [5.0, 20.0, 25.0, 10.0, 40.0],
+            }
+        )
+
+    def test_null_aware_comparison(self, sample_data):
+        """Test null-aware comparison (is_not_null AND comparison)."""
+        # Filter: value is not null AND value > threshold
+        expr = col("value").is_not_null() & (col("value") > col("threshold"))
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_null_or_condition(self, sample_data):
+        """Test is_null combined with OR."""
+        # value is null OR value > 40
+        expr = col("value").is_null() | (col("value") > 40)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, True, False, True, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+# ──────────────────────────────────────
+# Membership Predicate Operations
+# ──────────────────────────────────────
+
+
+class TestIsIn:
+    """Tests for is_in() predicate."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for membership tests."""
+        return pd.DataFrame(
+            {
+                "status": ["active", "inactive", "pending", "active", "deleted"],
+                "category": ["A", "B", "C", "A", "D"],
+                "value": [1, 2, 3, 4, 5],
+            }
+        )
+
+    def test_is_in_string_list(self, sample_data):
+        """Test is_in with string list."""
+        expr = col("status").is_in(["active", "pending"])
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.IN
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_in_single_value_list(self, sample_data):
+        """Test is_in with single-value list."""
+        expr = col("status").is_in(["active"])
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_in_numeric_list(self, sample_data):
+        """Test is_in with numeric list."""
+        expr = col("value").is_in([1, 3, 5])
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_in_empty_list(self, sample_data):
+        """Test is_in with empty list (should return all False)."""
+        expr = col("status").is_in([])
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([False, False, False, False, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_in_with_literal_expr(self, sample_data):
+        """Test is_in with LiteralExpr containing list."""
+        values_expr = lit(["A", "C"])
+        expr = col("category").is_in(values_expr)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_in_structural_equality(self):
+        """Test structural equality for is_in expressions."""
+        expr1 = col("status").is_in(["active", "pending"])
+        expr2 = col("status").is_in(["active", "pending"])
+        expr3 = col("status").is_in(["active"])
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+
+class TestNotIn:
+    """Tests for not_in() predicate."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for membership tests."""
+        return pd.DataFrame(
+            {
+                "status": ["active", "inactive", "pending", "active", "deleted"],
+                "value": [1, 2, 3, 4, 5],
+            }
+        )
+
+    def test_not_in_string_list(self, sample_data):
+        """Test not_in with string list."""
+        expr = col("status").not_in(["inactive", "deleted"])
+        assert isinstance(expr, BinaryExpr)
+        assert expr.op == Operation.NOT_IN
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_not_in_numeric_list(self, sample_data):
+        """Test not_in with numeric list."""
+        expr = col("value").not_in([2, 4])
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, False, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_not_in_empty_list(self, sample_data):
+        """Test not_in with empty list (should return all True)."""
+        expr = col("status").not_in([])
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, True, True, True, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_not_in_structural_equality(self):
+        """Test structural equality for not_in expressions."""
+        expr1 = col("status").not_in(["deleted"])
+        expr2 = col("status").not_in(["deleted"])
+        expr3 = col("status").not_in(["deleted", "inactive"])
+
+        assert expr1.structurally_equals(expr2)
+        assert not expr1.structurally_equals(expr3)
+
+
+class TestMembershipWithNulls:
+    """Tests for membership predicates with null values."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data with null values for membership tests."""
+        return pd.DataFrame(
+            {
+                "status": ["active", None, "pending", None, "deleted"],
+                "value": [1, None, 3, None, 5],
+            }
+        )
+
+    def test_is_in_with_nulls_in_data(self, sample_data):
+        """Test is_in when data contains nulls."""
+        expr = col("status").is_in(["active", "pending"])
+        result = eval_expr(expr, sample_data)
+        # Nulls should return False (null is not in any list)
+        expected = pd.Series([True, False, True, False, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_not_in_with_nulls_in_data(self, sample_data):
+        """Test not_in when data contains nulls."""
+        expr = col("status").not_in(["active"])
+        result = eval_expr(expr, sample_data)
+        # Nulls should return True (null is not in the exclusion list)
+        expected = pd.Series([False, True, True, True, True])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+class TestMembershipCombinations:
+    """Tests for membership predicates combined with other operations."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for combination tests."""
+        return pd.DataFrame(
+            {
+                "status": ["active", "inactive", "pending", "active", "deleted"],
+                "priority": ["high", "low", "high", "medium", "low"],
+                "value": [100, 50, 75, 200, 25],
+            }
+        )
+
+    def test_is_in_and_comparison(self, sample_data):
+        """Test is_in combined with comparison."""
+        # status in ["active", "pending"] AND value > 50
+        expr = col("status").is_in(["active", "pending"]) & (col("value") > 50)
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_multiple_is_in(self, sample_data):
+        """Test multiple is_in predicates."""
+        # status in ["active"] AND priority in ["high", "medium"]
+        expr = col("status").is_in(["active"]) & col("priority").is_in(
+            ["high", "medium"]
+        )
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, False, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+    def test_is_in_or_not_in(self, sample_data):
+        """Test is_in combined with not_in."""
+        # status in ["active"] OR priority not_in ["low"]
+        expr = col("status").is_in(["active"]) | col("priority").not_in(["low"])
+        result = eval_expr(expr, sample_data)
+        expected = pd.Series([True, False, True, True, False])
+        pd.testing.assert_series_equal(
+            result.reset_index(drop=True), expected, check_names=False
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/expressions/test_schema_resolution.py

@@ -0,0 +1,318 @@
+"""Tests for schema-aware expression resolution.
+
+Covers ``Expr.get_type``, ``Expr.nullable``, ``Expr.to_field``, and the
+``exprlist_to_fields`` helper. These are the building blocks Phase 1
+operators (``Project``, ``Aggregate``, ``Join``, etc.) use to compute
+``infer_schema()`` without executing the plan.
+"""
+
+import pyarrow as pa
+import pytest
+
+from ray.data.datatype import DataType
+from ray.data.expressions import (
+    DownloadExpr,
+    MonotonicallyIncreasingIdExpr,
+    RandomExpr,
+    UDFExpr,
+    UUIDExpr,
+    col,
+    expand_star_exprs,
+    exprlist_to_fields,
+    lit,
+    star,
+    udf,
+)
+from ray.data.tests.util import assert_exprs_equal
+
+
+@pytest.fixture
+def schema():
+    return pa.schema(
+        [
+            pa.field("a", pa.int32(), nullable=True),
+            pa.field("b", pa.float32(), nullable=False),
+            pa.field("name", pa.string(), nullable=True),
+            pa.field("flag", pa.bool_(), nullable=True),
+        ]
+    )
+
+
+class TestColumnExpr:
+    def test_to_field_returns_input_field_verbatim(self):
+        expected = pa.field("x", pa.int64(), nullable=False, metadata={b"k": b"v"})
+        in_schema = pa.schema([expected])
+        assert col("x").to_field(in_schema) == expected
+
+    def test_to_field_int_column(self, schema):
+        assert col("a").to_field(schema) == pa.field("a", pa.int32(), nullable=True)
+
+    def test_to_field_non_nullable_column(self, schema):
+        assert col("b").to_field(schema) == pa.field("b", pa.float32(), nullable=False)
+
+    def test_to_field_missing(self, schema):
+        assert col("missing").to_field(schema) is None
+
+
+class TestLiteralExpr:
+    @pytest.mark.parametrize(
+        "value,expected",
+        [
+            (5, pa.field("v", pa.int64(), nullable=False)),
+            (5.0, pa.field("v", pa.float64(), nullable=False)),
+            ("hello", pa.field("v", pa.string(), nullable=False)),
+            (True, pa.field("v", pa.bool_(), nullable=False)),
+            (None, pa.field("v", pa.null(), nullable=True)),
+        ],
+    )
+    def test_to_field(self, schema, value, expected):
+        # ``LiteralExpr`` has no name, so ``to_field`` returns ``None``;
+        # we wrap in an alias to get a named field.
+        assert lit(value).alias("v").to_field(schema) == expected
+
+
+class TestBinaryExpr:
+    @pytest.mark.parametrize(
+        "expr_factory,expected",
+        [
+            # int32 + int64 -> int64 (PyArrow promotion); a is nullable.
+            (lambda: col("a") + lit(5), pa.field("out", pa.int64(), nullable=True)),
+            # int32 + float32 -> float; a is nullable, so output is nullable.
+            (lambda: col("a") + col("b"), pa.field("out", pa.float32(), nullable=True)),
+            # float32 * float64 -> double; both operands are non-nullable
+            # (b is non-null, literal is not None) -> output is non-nullable.
+            (
+                lambda: col("b") * lit(2.0),
+                pa.field("out", pa.float64(), nullable=False),
+            ),
+            # comparisons -> bool
+            (lambda: col("a") > lit(0), pa.field("out", pa.bool_(), nullable=True)),
+            (lambda: col("a") == lit(1), pa.field("out", pa.bool_(), nullable=True)),
+            # logical -> bool
+            (
+                lambda: col("flag") & col("flag"),
+                pa.field("out", pa.bool_(), nullable=True),
+            ),
+            (
+                lambda: col("flag") | col("flag"),
+                pa.field("out", pa.bool_(), nullable=True),
+            ),
+            # in/not_in -> bool
+            (
+                lambda: col("a").is_in([1, 2, 3]),
+                pa.field("out", pa.bool_(), nullable=True),
+            ),
+            (
+                lambda: col("a").not_in([1, 2, 3]),
+                pa.field("out", pa.bool_(), nullable=True),
+            ),
+            # string concat
+            (
+                lambda: col("name") + lit("!"),
+                pa.field("out", pa.string(), nullable=True),
+            ),
+        ],
+    )
+    def test_to_field(self, schema, expr_factory, expected):
+        assert expr_factory().alias("out").to_field(schema) == expected
+
+    def test_unresolvable_returns_none(self, schema):
+        assert (col("missing") + lit(1)).alias("x").to_field(schema) is None
+
+
+class TestUnaryExpr:
+    def test_is_null(self, schema):
+        expected = pa.field("isnull", pa.bool_(), nullable=False)
+        assert col("a").is_null().alias("isnull").to_field(schema) == expected
+
+    def test_is_not_null(self, schema):
+        expected = pa.field("isnotnull", pa.bool_(), nullable=False)
+        assert col("a").is_not_null().alias("isnotnull").to_field(schema) == expected
+
+    def test_not_bool(self, schema):
+        expected = pa.field("neg", pa.bool_(), nullable=True)
+        assert (~col("flag")).alias("neg").to_field(schema) == expected
+
+
+class TestAliasExpr:
+    def test_to_field_renames(self, schema):
+        # Wraps a column field; alias swaps the name, preserves type/nullability.
+        assert col("a").alias("renamed").to_field(schema) == pa.field(
+            "renamed", pa.int32(), nullable=True
+        )
+
+    def test_to_field_around_binary(self, schema):
+        assert (col("a") + col("b")).alias("sum").to_field(schema) == pa.field(
+            "sum", pa.float32(), nullable=True
+        )
+
+
+class TestSelfContainedExprs:
+    def test_udf_uses_return_dtype(self, schema):
+        @udf(return_dtype=DataType.float64())  # pyrefly: ignore[missing-attribute]
+        def double(x):
+            return x
+
+        assert double(col("a")).alias("d").to_field(  # pyrefly: ignore[not-callable]
+            schema
+        ) == pa.field("d", pa.float64(), nullable=True)
+
+    def test_download_is_binary(self, schema):
+        assert DownloadExpr("uri").alias("bytes").to_field(schema) == pa.field(
+            "bytes", pa.binary(), nullable=True
+        )
+
+    def test_monotonically_increasing_id(self, schema):
+        assert MonotonicallyIncreasingIdExpr().alias("id").to_field(schema) == pa.field(
+            "id", pa.int64(), nullable=False
+        )
+
+    def test_random(self, schema):
+        assert RandomExpr().alias("r").to_field(schema) == pa.field(
+            "r", pa.float64(), nullable=False
+        )
+
+    def test_uuid(self, schema):
+        assert UUIDExpr().alias("u").to_field(schema) == pa.field(
+            "u", pa.string(), nullable=False
+        )
+
+
+class TestStarExpr:
+    def test_to_field_returns_none(self, schema):
+        # ``StarExpr`` represents many columns; ``exprlist_to_fields``
+        # expands it inline rather than calling ``to_field`` on it.
+        assert star().to_field(schema) is None
+
+
+class TestExprlistToFields:
+    def test_simple_columns(self, schema):
+        expected = pa.schema(
+            [
+                pa.field("a", pa.int32(), nullable=True),
+                pa.field("b", pa.float32(), nullable=False),
+            ]
+        )
+        result = pa.schema(exprlist_to_fields([col("a"), col("b")], schema))
+        assert result == expected
+
+    def test_star_expansion(self, schema):
+        # Star expands to all input fields verbatim.
+        result = pa.schema(exprlist_to_fields([star()], schema))
+        assert result == schema
+
+    def test_star_with_rename(self, schema):
+        result = pa.schema(
+            exprlist_to_fields([star(), col("a")._rename("renamed_a")], schema)
+        )
+        # Renaming "a" -> "renamed_a" substitutes the renamed field at
+        # "a"'s position (matching runtime ``eval_projection``).
+        expected = pa.schema(
+            [
+                pa.field("renamed_a", pa.int32(), nullable=True),
+                pa.field("b", pa.float32(), nullable=False),
+                pa.field("name", pa.string(), nullable=True),
+                pa.field("flag", pa.bool_(), nullable=True),
+            ]
+        )
+        assert result == expected
+
+    def test_star_with_rename_missing_source_returns_none(self, schema):
+        # Renaming an absent column must fail resolution (matching the
+        # runtime, which raises "column not found"), not silently append.
+        assert exprlist_to_fields([star(), col("missing")._rename("x")], schema) is None
+
+    def test_star_with_with_column(self, schema):
+        # with_column-style: [star(), expr.alias(name)] preserves all input
+        # columns and appends the new computed column.
+        result = pa.schema(
+            exprlist_to_fields([star(), (col("a") + col("b")).alias("sum")], schema)
+        )
+        expected = pa.schema(
+            list(schema) + [pa.field("sum", pa.float32(), nullable=True)]
+        )
+        assert result == expected
+
+    def test_returns_none_on_unresolvable(self, schema):
+        assert exprlist_to_fields([col("missing")], schema) is None
+
+    def test_with_column_overrides_existing_column(self, schema):
+        # ``with_column("a", new_expr)`` builds ``[star(), new_expr.alias("a")]``.
+        # The override should replace the existing "a" in place (last-wins,
+        # matching runtime ``eval_projection``'s upsert behavior), not
+        # produce a duplicate.
+        result = pa.schema(
+            exprlist_to_fields([star(), (col("a") + lit(10)).alias("a")], schema)
+        )
+        expected = pa.schema(
+            [
+                pa.field("a", pa.int64(), nullable=True),  # new type from a + 10
+                pa.field("b", pa.float32(), nullable=False),
+                pa.field("name", pa.string(), nullable=True),
+                pa.field("flag", pa.bool_(), nullable=True),
+            ]
+        )
+        assert result == expected
+
+    def test_returns_none_on_udf_without_return_dtype(self, schema):
+        # Construct a synthetic UDFExpr with object data_type to simulate
+        # the "untyped UDF" case.
+        e = UDFExpr(
+            fn=lambda x: x,
+            args=[col("a")],
+            kwargs={},
+            data_type=DataType(object),
+        )
+        assert exprlist_to_fields([e.alias("out")], schema) is None
+
+
+class TestExpandStarExprs:
+    """Tests for ``expand_star_exprs`` (eager expansion in ``Project``)."""
+
+    def test_passthrough_without_star(self, schema):
+        # No StarExpr -> input list returned unchanged.
+        exprs = [col("a"), (col("a") + col("b")).alias("sum")]
+        assert expand_star_exprs(exprs, schema) is exprs
+
+    def test_passthrough_when_schema_is_none(self):
+        exprs = [star(), col("a")]
+        assert expand_star_exprs(exprs, None) is exprs
+
+    def test_simple_star(self, schema):
+        # ``[star()]`` -> one ``col()`` per input column.
+        result = expand_star_exprs([star()], schema)
+        assert_exprs_equal(result, [col("a"), col("b"), col("name"), col("flag")])
+
+    def test_star_with_with_column(self, schema):
+        # ``with_column``-style: ``[star(), expr.alias("new")]`` expands
+        # to ``[col(a), col(b), col(name), col(flag), expr.alias("new")]``.
+        new_expr = (col("a") + col("b")).alias("new")
+        result = expand_star_exprs([star(), new_expr], schema)
+        assert_exprs_equal(
+            result, [col("a"), col("b"), col("name"), col("flag"), new_expr]
+        )
+
+    def test_star_with_rename(self, schema):
+        # ``rename_columns({"a": "renamed_a"})``: the rename substitutes for
+        # its source column *in place* (matching runtime ``eval_projection``
+        # / ``exprlist_to_fields``), so ``a`` becomes ``renamed_a`` at
+        # position 0 rather than moving to the end.
+        rename = col("a")._rename("renamed_a")
+        result = expand_star_exprs([star(), rename], schema)
+        assert_exprs_equal(result, [rename, col("b"), col("name"), col("flag")])
+
+    def test_star_with_rename_source_missing(self, schema):
+        # A rename whose source column isn't in the input schema stays in
+        # its trailing position so it still errors ("column not found") at
+        # runtime instead of being silently dropped.
+        rename = col("missing")._rename("renamed")
+        result = expand_star_exprs([star(), rename], schema)
+        assert_exprs_equal(
+            result, [col("a"), col("b"), col("name"), col("flag"), rename]
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-xvs"]))

python/ray/data/tests/unit/test_arrow_block.py

@@ -0,0 +1,276 @@
+import sys
+from typing import Union
+
+import numpy as np
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.arrow_block import (
+    ArrowBlockAccessor,
+    ArrowBlockBuilder,
+    ArrowBlockColumnAccessor,
+    _get_max_chunk_size,
+)
+from ray.data._internal.arrow_ops.transform_pyarrow import combine_chunked_array, concat
+from ray.data._internal.tensor_extensions.arrow import (
+    ArrowTensorArray,
+)
+
+
+def simple_array():
+    return pa.array([1, 2, None, 6], type=pa.int64())
+
+
+def simple_chunked_array():
+    return pa.chunked_array([pa.array([1, 2]), pa.array([None, 6])])
+
+
+def _wrap_as_pa_scalar(v, dtype: pa.DataType):
+    return pa.scalar(v, type=dtype)
+
+
+@pytest.mark.parametrize("arr", [simple_array(), simple_chunked_array()])
+@pytest.mark.parametrize("as_py", [True, False])
+class TestArrowBlockColumnAccessor:
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 3),
+            (False, 4),
+        ],
+    )
+    def test_count(self, arr, ignore_nulls, as_py, expected):
+        accessor = ArrowBlockColumnAccessor(arr)
+        result = accessor.count(ignore_nulls=ignore_nulls, as_py=as_py)
+
+        if not as_py:
+            expected = _wrap_as_pa_scalar(expected, dtype=pa.int64())
+
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 9),
+            (False, None),
+        ],
+    )
+    def test_sum(self, arr, ignore_nulls, as_py, expected):
+        accessor = ArrowBlockColumnAccessor(arr)
+        result = accessor.sum(ignore_nulls=ignore_nulls, as_py=as_py)
+
+        if not as_py:
+            expected = _wrap_as_pa_scalar(expected, dtype=pa.int64())
+
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 1),
+            (False, None),
+        ],
+    )
+    def test_min(self, arr, ignore_nulls, as_py, expected):
+        accessor = ArrowBlockColumnAccessor(arr)
+        result = accessor.min(ignore_nulls=ignore_nulls, as_py=as_py)
+
+        if not as_py:
+            expected = _wrap_as_pa_scalar(expected, dtype=pa.int64())
+
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 6),
+            (False, None),
+        ],
+    )
+    def test_max(self, arr, ignore_nulls, as_py, expected):
+        accessor = ArrowBlockColumnAccessor(arr)
+        result = accessor.max(ignore_nulls=ignore_nulls, as_py=as_py)
+
+        if not as_py:
+            expected = _wrap_as_pa_scalar(expected, dtype=pa.int64())
+
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "ignore_nulls, expected",
+        [
+            (True, 3),
+            (False, None),
+        ],
+    )
+    def test_mean(self, arr, ignore_nulls, as_py, expected):
+        accessor = ArrowBlockColumnAccessor(arr)
+        result = accessor.mean(ignore_nulls=ignore_nulls, as_py=as_py)
+
+        if not as_py:
+            expected = _wrap_as_pa_scalar(expected, dtype=pa.float64())
+
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "provided_mean, expected",
+        [
+            (3.0, 14.0),
+            (None, 14.0),
+        ],
+    )
+    def test_sum_of_squared_diffs_from_mean(self, arr, provided_mean, as_py, expected):
+        accessor = ArrowBlockColumnAccessor(arr)
+        result = accessor.sum_of_squared_diffs_from_mean(
+            ignore_nulls=True, mean=provided_mean, as_py=as_py
+        )
+
+        if not as_py:
+            expected = _wrap_as_pa_scalar(expected, dtype=pa.float64())
+
+        assert result == expected
+
+    def test_to_pylist(self, arr, as_py):
+        accessor = ArrowBlockColumnAccessor(arr)
+        assert accessor.to_pylist() == arr.to_pylist()
+
+
+@pytest.mark.parametrize(
+    "input_,expected_output",
+    [
+        # Empty chunked array
+        (pa.chunked_array([], type=pa.int8()), pa.array([], type=pa.int8())),
+        # Fixed-shape tensors
+        (
+            pa.chunked_array(
+                [
+                    ArrowTensorArray.from_numpy(np.arange(3).reshape(3, 1)),
+                    ArrowTensorArray.from_numpy(np.arange(3).reshape(3, 1)),
+                ]
+            ),
+            ArrowTensorArray.from_numpy(
+                np.concatenate(
+                    [
+                        np.arange(3).reshape(3, 1),
+                        np.arange(3).reshape(3, 1),
+                    ]
+                )
+            ),
+        ),
+        # Ragged (variable-shaped) tensors
+        (
+            pa.chunked_array(
+                [
+                    ArrowTensorArray.from_numpy(np.arange(3).reshape(3, 1)),
+                    ArrowTensorArray.from_numpy(np.arange(5).reshape(5, 1)),
+                ]
+            ),
+            ArrowTensorArray.from_numpy(
+                np.concatenate(
+                    [
+                        np.arange(3).reshape(3, 1),
+                        np.arange(5).reshape(5, 1),
+                    ]
+                )
+            ),
+        ),
+        # Small (< 2 GiB) arrays
+        (
+            pa.chunked_array(
+                [
+                    pa.array([1, 2, 3], type=pa.int16()),
+                    pa.array([4, 5, 6], type=pa.int16()),
+                ]
+            ),
+            pa.array([1, 2, 3, 4, 5, 6], type=pa.int16()),
+        ),
+    ],
+)
+def test_combine_chunked_array_small(
+    input_, expected_output: Union[pa.Array, pa.ChunkedArray]
+):
+    result = combine_chunked_array(input_)
+
+    assert expected_output.equals(result)
+
+
+@pytest.mark.parametrize(
+    "input_block, fill_column_name, fill_value, expected_output_block",
+    [
+        (
+            pa.Table.from_pydict({"a": [0, 1]}),
+            "b",
+            2,
+            pa.Table.from_pydict({"a": [0, 1], "b": [2, 2]}),
+        ),
+        (
+            pa.Table.from_pydict({"a": [0, 1]}),
+            "b",
+            pa.scalar(2),
+            pa.Table.from_pydict({"a": [0, 1], "b": [2, 2]}),
+        ),
+    ],
+)
+def test_fill_column(input_block, fill_column_name, fill_value, expected_output_block):
+    block_accessor = ArrowBlockAccessor.for_block(input_block)
+
+    actual_output_block = block_accessor.fill_column(fill_column_name, fill_value)
+
+    assert actual_output_block.equals(expected_output_block)
+
+
+def test_add_blocks_with_different_column_names():
+    builder = ArrowBlockBuilder()
+
+    builder.add_block(pa.Table.from_pydict({"col1": ["spam"]}))
+    builder.add_block(pa.Table.from_pydict({"col2": ["foo"]}))
+    block = builder.build()
+
+    expected_table = pa.Table.from_pydict(
+        {"col1": ["spam", None], "col2": [None, "foo"]}
+    )
+    assert block.equals(expected_table)
+
+
+@pytest.mark.parametrize(
+    "table_data,max_chunk_size_bytes,expected",
+    [
+        ({"a": []}, 100, None),
+        ({"a": list(range(100))}, 7, 1),
+        ({"a": list(range(100))}, 10, 1),
+        ({"a": list(range(100))}, 25, 3),
+        ({"a": list(range(100))}, 50, 6),
+        ({"a": list(range(100))}, 100, 12),
+    ],
+)
+def test_arrow_block_max_chunk_size(table_data, max_chunk_size_bytes, expected):
+    table = pa.table(table_data)
+    assert _get_max_chunk_size(table, max_chunk_size_bytes) == expected
+
+
+def test_arrow_block_concat():
+    table1 = pa.table(
+        {
+            "a": [1, 2, 3],
+            "s": [{"x": 1} for _ in range(3)],
+        }
+    )
+    table2 = pa.table(
+        {
+            "b": [4, 5, 6],
+        }
+    )
+    concatenated = concat([table1, table2])
+    assert set(concatenated.column_names) == {"a", "s", "b"}
+    expected = pa.table(
+        {
+            "a": [1, 2, 3, None, None, None],
+            "s": [{"x": 1} for _ in range(3)] + [None] * 3,
+            "b": [None, None, None, 4, 5, 6],
+        }
+    )
+    assert concatenated.select(["a", "s", "b"]) == expected
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_arrow_serialization.py

@@ -0,0 +1,205 @@
+import sys
+from unittest import mock
+
+import numpy as np
+import pyarrow as pa
+import pytest
+from packaging.version import parse as parse_version
+
+from ray._private.arrow_serialization import (
+    PicklableArrayPayload,
+    _align_bit_offset,
+    _bytes_for_bits,
+    _copy_bitpacked_buffer_if_needed,
+    _copy_buffer_if_needed,
+    _copy_normal_buffer_if_needed,
+    _copy_offsets_buffer_if_needed,
+)
+
+
+@pytest.mark.parametrize(
+    "n,expected",
+    [(0, 0)] + [(i, 8) for i in range(1, 9)] + [(i, 16) for i in range(9, 17)],
+)
+def test_bytes_for_bits_manual(n, expected):
+    assert _bytes_for_bits(n) == expected
+
+
+def test_bytes_for_bits_auto():
+    M = 128
+    expected = [((n - 1) // 8 + 1) * 8 for n in range(M)]
+    for n, e in enumerate(expected):
+        assert _bytes_for_bits(n) == e, n
+
+
+def test_align_bit_offset_auto():
+    M = 10
+    n = M * (2**8 - 1)
+    # Represent an integer as a Pyarrow buffer of bytes.
+    bytes_ = n.to_bytes(M, sys.byteorder)
+    buf = pa.py_buffer(bytes_)
+    for slice_len in range(1, M):
+        for bit_offset in range(1, n - slice_len * 8):
+            byte_length = _bytes_for_bits(bit_offset + slice_len * 8) // 8
+            # Shift the buffer to eliminate the offset.
+            out_buf = _align_bit_offset(buf, bit_offset, byte_length)
+            # Check that shifted buffer is equivalent to our base int shifted by the
+            # same number of bits.
+            assert int.from_bytes(out_buf.to_pybytes(), sys.byteorder) == (
+                n >> bit_offset
+            )
+
+
+@mock.patch("ray._private.arrow_serialization._copy_normal_buffer_if_needed")
+@mock.patch("ray._private.arrow_serialization._copy_bitpacked_buffer_if_needed")
+def test_copy_buffer_if_needed(mock_bitpacked, mock_normal):
+    # Test that type-based buffer copy dispatch works as expected.
+    bytes_ = b"abcd"
+    buf = pa.py_buffer(bytes_)
+    offset = 1
+    length = 2
+
+    # Normal (non-boolean) buffer copy path.
+    type_ = pa.int32()
+    _copy_buffer_if_needed(buf, type_, offset, length)
+    expected_byte_width = 4
+    mock_normal.assert_called_once_with(buf, expected_byte_width, offset, length)
+    mock_normal.reset_mock()
+
+    type_ = pa.int64()
+    _copy_buffer_if_needed(buf, type_, offset, length)
+    expected_byte_width = 8
+    mock_normal.assert_called_once_with(buf, expected_byte_width, offset, length)
+    mock_normal.reset_mock()
+
+    # Boolean buffer copy path.
+    type_ = pa.bool_()
+    _copy_buffer_if_needed(buf, type_, offset, length)
+    mock_bitpacked.assert_called_once_with(buf, offset, length)
+    mock_bitpacked.reset_mock()
+
+
+def test_copy_normal_buffer_if_needed():
+    bytes_ = b"abcd"
+    buf = pa.py_buffer(bytes_)
+    byte_width = 1
+    uncopied_buf = _copy_normal_buffer_if_needed(buf, byte_width, 0, len(bytes_))
+    assert uncopied_buf.address == buf.address
+    assert uncopied_buf.size == len(bytes_)
+    for offset in range(1, len(bytes_) - 1):
+        for length in range(1, len(bytes_) - offset):
+            copied_buf = _copy_normal_buffer_if_needed(buf, byte_width, offset, length)
+            assert copied_buf.address != buf.address
+            assert copied_buf.size == length
+
+
+def test_copy_bitpacked_buffer_if_needed():
+    M = 20
+    n = M * 8
+    # Represent an integer as a pyarrow buffer of bytes.
+    bytes_ = (n * 8).to_bytes(M, sys.byteorder)
+    buf = pa.py_buffer(bytes_)
+    for offset in range(0, n - 1):
+        for length in range(1, n - offset):
+            copied_buf = _copy_bitpacked_buffer_if_needed(buf, offset, length)
+            if offset > 0:
+                assert copied_buf.address != buf.address
+            else:
+                assert copied_buf.address == buf.address
+            # Buffer needs to include bits remaining in byte after adjusting for bit
+            # offset..
+            assert copied_buf.size == ((length + (offset % 8) - 1) // 8) + 1
+
+
+@pytest.mark.parametrize(
+    "arr_type,expected_offset_type",
+    [
+        (pa.list_(pa.int64()), pa.int32()),
+        (pa.string(), pa.int32()),
+        (pa.binary(), pa.int32()),
+        (pa.large_list(pa.int64()), pa.int64()),
+        (pa.large_string(), pa.int64()),
+        (pa.large_binary(), pa.int64()),
+    ],
+)
+def test_copy_offsets_buffer_if_needed(arr_type, expected_offset_type):
+    offset_arr = pa.array([0, 1, 3, 6, 10, 15, 21], type=expected_offset_type)
+    buf = offset_arr.buffers()[1]
+    offset = 2
+    length = 3
+    offset_buf, data_offset, data_length = _copy_offsets_buffer_if_needed(
+        buf, arr_type, offset, length
+    )
+    assert data_offset == 3
+    assert data_length == 12
+    truncated_offset_arr = pa.Array.from_buffers(
+        expected_offset_type, length, [None, offset_buf]
+    )
+    expected_offset_arr = pa.array([0, 3, 7], type=expected_offset_type)
+    assert truncated_offset_arr.equals(expected_offset_arr)
+
+
+@pytest.mark.skipif(
+    parse_version(str(pa.__version__)) < parse_version("10.0.0"),
+    reason="FixedShapeTensorArray is not supported in PyArrow < 10.0.0",
+)
+def test_fixed_shape_tensor_array_serialization():
+    a = pa.FixedShapeTensorArray.from_numpy_ndarray(  # pyrefly: ignore[missing-attribute]
+        np.arange(4 * 2 * 3).reshape(4, 2, 3)
+    )
+    payload = PicklableArrayPayload.from_array(a)
+    a2 = payload.to_array()
+    assert a == a2
+
+
+class _VariableShapeTensorType(pa.ExtensionType):
+    def __init__(
+        self,
+        value_type: pa.DataType,
+        ndim: int,
+    ) -> None:
+        self.value_type = value_type
+        self.ndim = ndim
+        super().__init__(
+            pa.struct(
+                [
+                    pa.field("data", pa.list_(value_type)),
+                    pa.field("shape", pa.list_(pa.int32(), ndim)),
+                ]
+            ),
+            "variable_shape_tensor",
+        )
+
+    def __arrow_ext_serialize__(self) -> bytes:
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(cls, storage_type: pa.DataType, serialized: bytes):
+        ndim = storage_type[1].type.list_size
+        value_type = storage_type[0].type.value_type
+        return cls(value_type, ndim)
+
+
+def test_variable_shape_tensor_serialization():
+    t = _VariableShapeTensorType(pa.float32(), 2)
+    values = [
+        {
+            "data": np.arange(2 * 3, dtype=np.float32).tolist(),
+            "shape": [2, 3],
+        },
+        {
+            "data": np.arange(4 * 5, dtype=np.float32).tolist(),
+            "shape": [4, 5],
+        },
+    ]
+    storage = pa.array(values, type=t.storage_type)
+    ar = pa.ExtensionArray.from_storage(t, storage)
+    payload = PicklableArrayPayload.from_array(ar)
+    ar2 = payload.to_array()
+    assert ar == ar2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_arrow_type_conversion.py

@@ -0,0 +1,205 @@
+import gc
+from dataclasses import dataclass, field
+
+import numpy as np
+import pyarrow as pa
+import pytest
+from packaging.version import parse as parse_version
+
+from ray.data import DataContext
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.tensor_extensions.arrow import (
+    ArrowConversionError,
+    ArrowTensorArray,
+    _convert_to_pyarrow_native_array,
+    _infer_pyarrow_type,
+    convert_to_pyarrow_array,
+)
+from ray.data._internal.tensor_extensions.utils import create_ragged_ndarray
+from ray.data._internal.util import MiB
+from ray.data._internal.utils.arrow_utils import get_pyarrow_version
+from ray.tests.conftest import *  # noqa
+
+import psutil
+
+
+@dataclass
+class UserObj:
+    i: int = field()
+
+
+@pytest.mark.parametrize(
+    "input",
+    [
+        # Python native lists
+        [
+            [1, 2],
+            [3, 4],
+        ],
+        # Python native tuples
+        [
+            (1, 2),
+            (3, 4),
+        ],
+        # Lists as PA scalars
+        [
+            pa.scalar([1, 2]),
+            pa.scalar([3, 4]),
+        ],
+    ],
+)
+def test_arrow_native_list_conversion(input, disable_fallback_to_object_extension):
+    """Test asserts that nested lists are represented as native Arrow lists
+    upon serialization into Arrow format (and are NOT converted to numpy
+    tensor using extension)"""
+
+    if isinstance(input[0], pa.Scalar) and get_pyarrow_version() <= parse_version(
+        "13.0.0"
+    ):
+        pytest.skip(
+            "Pyarrow < 13.0 not able to properly infer native types from its own Scalars"
+        )
+
+    pa_arr = convert_to_pyarrow_array(input, "a")
+
+    # Should be able to natively convert back to Pyarrow array,
+    # not using any extensions
+    assert pa_arr.type == pa.list_(pa.int64()), pa_arr.type
+    assert pa.array(input) == pa_arr, pa_arr
+
+
+@pytest.mark.parametrize("arg_type", ["list", "ndarray"])
+@pytest.mark.parametrize(
+    "numpy_precision, expected_arrow_timestamp_type",
+    [
+        ("ms", pa.timestamp("ms")),
+        ("us", pa.timestamp("us")),
+        ("ns", pa.timestamp("ns")),
+        # The coarsest resolution Arrow supports is seconds.
+        ("Y", pa.timestamp("s")),
+        ("M", pa.timestamp("s")),
+        ("D", pa.timestamp("s")),
+        ("h", pa.timestamp("s")),
+        ("m", pa.timestamp("s")),
+        ("s", pa.timestamp("s")),
+        # The finest resolution Arrow supports is nanoseconds.
+        ("ps", pa.timestamp("ns")),
+        ("fs", pa.timestamp("ns")),
+        ("as", pa.timestamp("ns")),
+    ],
+)
+def test_convert_datetime_array(
+    numpy_precision: str,
+    expected_arrow_timestamp_type: pa.TimestampType,
+    arg_type: str,
+    restore_data_context,
+):
+    DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
+
+    ndarray = np.ones(1, dtype=f"datetime64[{numpy_precision}]")
+
+    if arg_type == "ndarray":
+        column_values = ndarray
+    elif arg_type == "list":
+        column_values = [ndarray]
+    else:
+        pytest.fail(f"Unknown type: {arg_type}")
+
+    # Step 1: Convert to PA array
+    converted = convert_to_pyarrow_array(column_values, "")
+
+    if arg_type == "ndarray":
+        expected = pa.array(
+            column_values.astype(f"datetime64[{expected_arrow_timestamp_type.unit}]")
+        )
+    elif arg_type == "list":
+        expected = ArrowTensorArray.from_numpy(
+            [
+                column_values[0].astype(
+                    f"datetime64[{expected_arrow_timestamp_type.unit}]"
+                )
+            ]
+        )
+    else:
+        pytest.fail(f"Unknown type: {arg_type}")
+
+    assert expected.type == converted.type
+    assert expected == converted
+
+
+@pytest.mark.parametrize("dtype", ["int64", "float64", "datetime64[ns]"])
+def test_infer_type_does_not_leak_memory(dtype):
+    # Test for https://github.com/apache/arrow/issues/45493.
+    ndarray = np.zeros(923040, dtype=dtype)  # A ~7 MiB column
+
+    process = psutil.Process()
+    gc.collect()
+    pa.default_memory_pool().release_unused()
+    before = process.memory_info().rss
+
+    # Call the function several times. If there's a memory leak, this loop will leak
+    # as much as 1 GiB of memory with 8 repetitions. 8 was chosen arbitrarily.
+    num_repetitions = 8
+    for _ in range(num_repetitions):
+        _infer_pyarrow_type(ndarray)
+
+    gc.collect()
+    pa.default_memory_pool().release_unused()
+    after = process.memory_info().rss
+
+    margin_of_error = 64 * MiB
+    assert after - before < margin_of_error, memory_string(after - before)
+
+
+def test_pa_infer_type_failing_to_infer():
+    # Represent a single column that will be using `ArrowPythonObjectExtension` type
+    # to ser/de native Python objects into bytes
+    column_vals = create_ragged_ndarray(
+        [
+            "hi",
+            1,
+            None,
+            [[[[]]]],
+            {"a": [[{"b": 2, "c": UserObj(i=123)}]]},
+            UserObj(i=456),
+        ]
+    )
+
+    inferred_dtype = _infer_pyarrow_type(column_vals)
+
+    # Arrow (17.0) seems to fallback to assume the dtype of the first element
+    assert pa.string().equals(inferred_dtype)
+
+
+def test_convert_to_pyarrow_array_object_ext_type_fallback():
+    column_values = create_ragged_ndarray(
+        [
+            "hi",
+            1,
+            None,
+            [[[[]]]],
+            {"a": [[{"b": 2, "c": UserObj(i=123)}]]},
+            UserObj(i=456),
+        ]
+    )
+    column_name = "py_object_column"
+
+    # First, assert that straightforward conversion into Arrow native types fails
+    with pytest.raises(ArrowConversionError) as exc_info:
+        _convert_to_pyarrow_native_array(column_values, column_name)
+
+    assert (
+        str(exc_info.value)
+        == "Error converting column 'py_object_column' (target type: string) to Arrow: ['hi' 1 None list([[[[]]]]) {'a': [[{'b': 2, 'c': UserObj(i=123)}]]}\n UserObj(i=456)]"  # noqa: E501
+    )
+
+    # Subsequently, assert that fallback to `ArrowObjectExtensionType` succeeds
+    pa_array = convert_to_pyarrow_array(column_values, column_name)
+
+    assert pa_array.to_pylist() == column_values.tolist()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/data/tests/unit/test_auto_batch_size.py

@@ -0,0 +1,95 @@
+"""Tests for batch_size="auto" in BatchMapTransformFn and map_batches."""
+import numpy as np
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_transformer import (
+    BatchMapTransformFn,
+    MapTransformer,
+    _compute_auto_batch_size,
+)
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.planner.plan_udf_map_op import (
+    _generate_transform_fn_for_map_batches,
+)
+from ray.data.block import BatchFormat
+from ray.data.context import DEFAULT_TARGET_MAX_BLOCK_SIZE
+
+
+def test_compute_auto_batch_size_basic():
+    """Batch size is target_bytes / bytes_per_row from the first block."""
+    # 10 int64 rows = 80 bytes (8 bytes/row). Target of 800 -> batch_size = 100.
+    block = pa.table({"x": pa.array(np.arange(10, dtype=np.int64))})
+    batch_size, _ = _compute_auto_batch_size(iter([block]), target_batch_size_bytes=800)
+    assert batch_size == 100
+
+
+def test_compute_auto_batch_size_clamped_to_one():
+    """When the target is smaller than one row, batch size clamps to 1."""
+    # 10 int64 rows = 80 bytes (8 bytes/row). Target of 1 byte < 1 row -> clamp to 1.
+    block = pa.table({"x": pa.array(np.arange(10, dtype=np.int64))})
+    batch_size, _ = _compute_auto_batch_size(iter([block]), target_batch_size_bytes=1)
+    assert batch_size == 1
+
+
+@pytest.mark.parametrize(
+    "blocks",
+    [
+        pytest.param([], id="empty_iterator"),
+        pytest.param([pa.table({"x": pa.array([], type=pa.int64())})], id="zero_rows"),
+    ],
+)
+def test_compute_auto_batch_size_returns_none(blocks):
+    """Empty and zero-row blocks return None (whole block becomes one batch)."""
+    batch_size, _ = _compute_auto_batch_size(iter(blocks))
+    assert batch_size is None
+
+
+def test_compute_auto_batch_size_iterator_includes_peeked_block():
+    """The returned iterator contains all input blocks, including the peeked block."""
+    blocks = [
+        pa.table({"x": pa.array(np.arange(10, dtype=np.int64))}),
+        pa.table({"x": pa.array(np.arange(10, 20, dtype=np.int64))}),
+    ]
+    _, returned_blocks = _compute_auto_batch_size(iter(blocks))
+    returned = list(returned_blocks)
+    assert len(returned) == 2
+    assert returned[0].equals(blocks[0])
+    assert returned[1].equals(blocks[1])
+
+
+def test_auto_batches_respect_target_size():
+    """With 'auto', rows are grouped to approximate the target byte size."""
+    # 1000 int64 rows = 8000 bytes (8 bytes/row). Target of 80 -> batch_size = 10.
+    block = pa.table({"x": pa.array(range(1000), type=pa.int64())})
+
+    received_sizes = []
+
+    def identity(batch):
+        received_sizes.append(len(batch))
+        return batch
+
+    transformer = MapTransformer(
+        [
+            BatchMapTransformFn(
+                _generate_transform_fn_for_map_batches(identity),
+                batch_size="auto",
+                batch_format=BatchFormat.ARROW,
+                output_block_size_option=OutputBlockSizeOption.of(
+                    target_max_block_size=DEFAULT_TARGET_MAX_BLOCK_SIZE
+                ),
+                target_batch_size_bytes=80,
+            )
+        ]
+    )
+    ctx = TaskContext(task_idx=0, op_name="test")
+    list(transformer.apply_transform(iter([block]), ctx))
+
+    assert sum(received_sizes) == 1000
+    assert all(s == 10 for s in received_sizes[:-1])
+
+
+if __name__ == "__main__":
+
+    pytest.main([__file__, "-v"])

python/ray/data/tests/unit/test_average_calculator.py

@@ -0,0 +1,67 @@
+from unittest.mock import patch
+
+import pytest
+
+from ray.data._internal.average_calculator import TimeWindowAverageCalculator
+
+
+@pytest.fixture
+def current_time():
+    class MutableInt:
+        def __init__(self, value: int = 0):
+            self.value = value
+
+        def __repr__(self):
+            return f"MutableInt({self.value})"
+
+        def increment(self):
+            self.value += 1
+
+        def get_value(self) -> int:
+            return self.value
+
+    _current_time = MutableInt()
+
+    def time():
+        return _current_time.get_value()
+
+    with patch("time.time", time):
+        yield _current_time
+
+
+def test_calcuate_time_window_average(current_time):
+    """Test TimeWindowAverageCalculator."""
+    window_s = 10
+    values_to_report = [i + 1 for i in range(20)]
+
+    calculator = TimeWindowAverageCalculator(window_s)
+    assert calculator.get_average() is None
+
+    for value in values_to_report:
+        # Report values, test `get_average`.
+        # and proceed the time by 1 second each time.
+        calculator.report(value)
+        avg = calculator.get_average()
+        values_in_window = values_to_report[
+            max(current_time.get_value() - 10, 0) : current_time.get_value() + 1
+        ]
+        expected = sum(values_in_window) / len(values_in_window)
+        assert avg == expected, current_time.get_value()
+        current_time.increment()
+
+    for _ in range(10):
+        # Keep proceeding the time, and test `get_average`.
+        avg = calculator.get_average()
+        values_in_window = values_to_report[max(current_time.get_value() - 10, 0) : 20]
+        expected = sum(values_in_window) / len(values_in_window)
+        assert avg == expected, current_time.get_value()
+        current_time.increment()
+
+    # Now no values in the time window, `get_average` should return None.
+    assert calculator.get_average() is None
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/data/tests/unit/test_block.py

@@ -0,0 +1,164 @@
+from datetime import datetime
+
+import numpy as np
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+from ray.data.block import BlockAccessor, BlockColumnAccessor
+
+
+def test_find_partitions_single_column_ascending():
+    table = pa.table({"value": [1, 2, 2, 3, 5]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(0,), (2,), (4,), (6,)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"value": []}  # <0
+    assert partitions[1].to_pydict() == {"value": [1]}  # [0,2)
+    assert partitions[2].to_pydict() == {"value": [2, 2, 3]}  # [2,4)
+    assert partitions[3].to_pydict() == {"value": [5]}  # [4,6)
+    assert partitions[4].to_pydict() == {"value": []}  # >=6
+
+
+def test_find_partitions_single_column_descending():
+    table = pa.table({"value": [5, 3, 2, 2, 1]})
+    sort_key = SortKey(key=["value"], descending=[True])
+    boundaries = [(6,), (3,), (2,), (0,)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"value": []}  # >=6
+    assert partitions[1].to_pydict() == {"value": [5, 3]}  # [3, 6)
+    assert partitions[2].to_pydict() == {"value": [2, 2]}  # [2, 3)
+    assert partitions[3].to_pydict() == {"value": [1]}  # [0, 2)
+    assert partitions[4].to_pydict() == {"value": []}  # <0
+
+
+def test_find_partitions_multi_column_ascending_first():
+    table = pa.table({"col1": [1, 1, 1, 1, 1, 2, 2], "col2": [4, 3, 2.5, 2, 1, 2, 1]})
+    sort_key = SortKey(key=["col1", "col2"], descending=[False, True])
+    boundaries = [(1, 3), (1, 2), (2, 2), (2, 0)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"col1": [1], "col2": [4]}
+    assert partitions[1].to_pydict() == {"col1": [1, 1], "col2": [3, 2.5]}
+    assert partitions[2].to_pydict() == {"col1": [1, 1], "col2": [2, 1]}
+    assert partitions[3].to_pydict() == {"col1": [2, 2], "col2": [2, 1]}
+    assert partitions[4].to_pydict() == {"col1": [], "col2": []}
+
+
+def test_find_partitions_multi_column_descending_first():
+    table = pa.table({"col1": [2, 2, 1, 1, 1, 1, 1], "col2": [1, 2, 1, 2, 3, 4, 5]})
+    sort_key = SortKey(key=["col1", "col2"], descending=[True, False])
+    boundaries = [(2, 0), (2, 2), (1, 2), (1, 6)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 5
+    assert partitions[0].to_pydict() == {"col1": [], "col2": []}
+    assert partitions[1].to_pydict() == {"col1": [2, 2], "col2": [1, 2]}
+    assert partitions[2].to_pydict() == {"col1": [1, 1], "col2": [1, 2]}
+    assert partitions[3].to_pydict() == {"col1": [1, 1, 1], "col2": [3, 4, 5]}
+    assert partitions[4].to_pydict() == {"col1": [], "col2": []}
+
+
+@pytest.mark.parametrize("null", [None, np.nan])
+def test_find_partitions_table_with_nulls(null):
+    table = pa.table({"value": [1, 2, 3, null, null]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(2,), (4,)]
+
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+
+    assert len(partitions) == 3
+    assert partitions[0].to_pydict() == {"value": [1]}  # <2
+    assert partitions[1].to_pydict() == {"value": [2, 3]}  # [2, 4)
+
+    if null is None:
+        assert partitions[2].to_pydict() == {"value": [null, null]}  # >=4
+    else:
+        # NOTE: NaNs couldn't be compared directly
+        result = partitions[2].to_pydict()
+        assert len(result["value"]) == 2 and all([np.isnan(v) for v in result["value"]])
+
+
+@pytest.mark.parametrize("dtype", ["str", "datetime"])
+def test_find_partitions_object_table_with_nulls(dtype):
+    if dtype == "str":
+        col = pa.array(["a", "b", "c", None, None])
+    elif dtype == "datetime":
+        col = pa.array(
+            [
+                datetime.fromordinal(1),
+                datetime.fromordinal(2),
+                datetime.fromordinal(3),
+                None,
+                None,
+            ]
+        )
+    else:
+        raise ValueError(f"Unexpected dtype={dtype}")
+
+    table = pa.table({"value": col})
+
+    sort_key = SortKey(key=["value"], descending=[False])
+
+    ndarray = BlockColumnAccessor.for_column(col).to_numpy(zero_copy_only=False)
+
+    # Compose boundaries
+    boundaries = [(ndarray[1],), (None,)]
+
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+
+    assert len(partitions) == 3
+    assert (
+        partitions[0]["value"].to_pylist() == col[0:1].to_pylist()
+    )  # < second element
+    assert partitions[1]["value"].to_pylist() == col[1:3].to_pylist()  # < None
+    assert partitions[2]["value"].to_pylist() == col[3:].to_pylist()  # remaining
+
+
+@pytest.mark.parametrize("null", [None, np.nan])
+def test_find_partitions_null_boundary(null):
+    table = pa.table({"value": [1, 2, 3]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(2,), (null,)]
+
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+
+    assert len(partitions) == 3
+    assert partitions[0].to_pydict() == {"value": [1]}  # <2
+    assert partitions[1].to_pydict() == {"value": [2, 3]}  # < null (nulls go last)
+    assert partitions[2].to_pydict() == {"value": []}
+
+
+def test_find_partitions_duplicates():
+    table = pa.table({"value": [2, 2, 2, 2, 2]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    boundaries = [(1,), (2,), (3,)]
+    partitions = BlockAccessor.for_block(table)._find_partitions_sorted(
+        boundaries, sort_key
+    )
+    assert len(partitions) == 4
+    assert partitions[0].to_pydict() == {"value": []}  # <1
+    assert partitions[1].to_pydict() == {"value": []}  # [1,2)
+    assert partitions[2].to_pydict() == {"value": [2, 2, 2, 2, 2]}  # [2,3)
+    assert partitions[3].to_pydict() == {"value": []}  # >=3
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_block_boundaries.py

@@ -0,0 +1,70 @@
+import numpy as np
+
+from ray.data.block import _get_group_boundaries_sorted_numpy
+
+
+def test_groupby_map_groups_get_block_boundaries():
+    """Test for cases with Nan or None"""
+    indices = _get_group_boundaries_sorted_numpy(
+        [
+            np.array([1, 1, 2, 2, 3, 3]),
+            np.array([1, 1, 2, 2, 3, 4]),
+        ]
+    )
+
+    assert list(indices) == [0, 2, 4, 5, 6]
+
+    indices = _get_group_boundaries_sorted_numpy(
+        [
+            np.array([1, 1, 2, 2, 3, 3]),
+            np.array(["a", "b", "a", "a", "b", "b"]),
+        ]
+    )
+
+    assert list(indices) == [0, 1, 2, 4, 6]
+
+    indices = _get_group_boundaries_sorted_numpy([np.array([1, 1, 2, 2, 3, 3])])
+
+    assert list(indices) == [0, 2, 4, 6]
+
+
+def test_groupby_map_groups_get_block_boundaries_with_nan():
+    """Test for cases with Nan or None. Since the arrays are sorted
+    in the groupby, they are located at the end. Also, nans/None are
+    treated as the same group.
+    """
+
+    indices = _get_group_boundaries_sorted_numpy(
+        [
+            np.array([1, 1, 2, 2, 3, np.nan, np.nan]),
+            np.array([1, 1, 2, 2, 3, 4, np.nan]),
+        ]
+    )
+
+    assert list(indices) == [0, 2, 4, 5, 6, 7]
+
+    indices = _get_group_boundaries_sorted_numpy(
+        [
+            np.array([1, 1, 2, 2, 3, 3, np.nan]),
+            np.array(["a", "b", "a", "a", "b", "b", None]),
+        ]
+    )
+
+    assert list(indices) == [0, 1, 2, 4, 6, 7]
+
+    indices = _get_group_boundaries_sorted_numpy(
+        [
+            np.array([1, 1, 2, 2, 3, 3, 4]),
+            np.array(["a", "b", "a", "a", "b", "b", None]),
+        ]
+    )
+
+    assert list(indices) == [0, 1, 2, 4, 6, 7]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_block_metadata_schema_cache.py

@@ -0,0 +1,150 @@
+"""Tests for the `BlockMetadataWithSchema` pickle round-trip path and its
+schema-deserialization cache (`_read_arrow_schema_cached`).
+
+The cache is on the StreamingExecutor scheduler thread's hot path; we want to
+verify:
+  1. Round-tripping preserves schema equality and field metadata.
+  2. Repeated unpickles of identical bytes reuse the cached `pa.Schema`.
+  3. Distinct schema bytes produce distinct cached entries.
+  4. Pandas / None schemas still work (no caching path taken).
+"""
+
+import pickle
+
+import pyarrow as pa
+import pytest
+
+from ray.data.block import (
+    BlockMetadata,
+    BlockMetadataWithSchema,
+    _read_arrow_schema_cached,
+)
+
+
+def _wide_arrow_schema(num_cols: int = 50) -> pa.Schema:
+    fields = []
+    for i in range(num_cols):
+        if i % 3 == 0:
+            fields.append(pa.field(f"scalar_{i}", pa.float32()))
+        elif i % 3 == 1:
+            fields.append(pa.field(f"vec64_{i}", pa.list_(pa.float32(), 64)))
+        else:
+            fields.append(pa.field(f"vec32_{i}", pa.list_(pa.float32(), 32)))
+    return pa.schema(fields)
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    _read_arrow_schema_cached.cache_clear()
+    yield
+    _read_arrow_schema_cached.cache_clear()
+
+
+def _make_bm(schema: "pa.Schema | None") -> BlockMetadataWithSchema:
+    md = BlockMetadata(
+        num_rows=10,
+        size_bytes=1024,
+        exec_stats=None,
+        input_files=None,
+        task_exec_stats=None,
+    )
+    return BlockMetadataWithSchema.from_metadata(md, schema=schema)
+
+
+def test_round_trip_preserves_schema():
+    schema = _wide_arrow_schema(20)
+    bm = _make_bm(schema)
+    restored = pickle.loads(pickle.dumps(bm))
+    assert restored.schema.equals(schema)
+    assert restored.num_rows == 10
+    assert restored.size_bytes == 1024
+
+
+def test_cache_hit_on_repeated_pickle_loads():
+    schema = _wide_arrow_schema(20)
+    payload = pickle.dumps(_make_bm(schema))
+
+    info_before = _read_arrow_schema_cached.cache_info()
+
+    restored = [pickle.loads(payload) for _ in range(50)]
+
+    info_after = _read_arrow_schema_cached.cache_info()
+    # Exactly one miss for the first load, the rest are hits.
+    assert info_after.misses - info_before.misses == 1
+    assert info_after.hits - info_before.hits == 49
+
+    # All decoded schemas are identical and reference-equal to the cached one.
+    first = restored[0].schema
+    for r in restored[1:]:
+        assert r.schema is first
+
+
+def test_distinct_schemas_distinct_cache_entries():
+    s1 = _wide_arrow_schema(10)
+    s2 = _wide_arrow_schema(20)
+    p1 = pickle.dumps(_make_bm(s1))
+    p2 = pickle.dumps(_make_bm(s2))
+
+    info_before = _read_arrow_schema_cached.cache_info()
+    a = pickle.loads(p1)
+    b = pickle.loads(p2)
+    c = pickle.loads(p1)
+    d = pickle.loads(p2)
+    info_after = _read_arrow_schema_cached.cache_info()
+
+    assert info_after.misses - info_before.misses == 2
+    assert info_after.hits - info_before.hits == 2
+    assert a.schema is c.schema
+    assert b.schema is d.schema
+    assert a.schema is not b.schema
+    assert a.schema.equals(s1)
+    assert b.schema.equals(s2)
+
+
+def test_none_schema_unaffected_by_cache():
+    bm = _make_bm(None)
+    info_before = _read_arrow_schema_cached.cache_info()
+    restored = pickle.loads(pickle.dumps(bm))
+    info_after = _read_arrow_schema_cached.cache_info()
+    assert restored.schema is None
+    # No cache traffic at all.
+    assert info_after.misses == info_before.misses
+    assert info_after.hits == info_before.hits
+
+
+def test_bytearray_schema_payload_is_decoded():
+    """If state["schema"] arrives as bytearray (e.g. from an alternative
+    serialization path), __setstate__ must still decode it to a pa.Schema —
+    not store it raw — and must reuse the LRU-cached entry keyed by the
+    equivalent bytes payload."""
+    schema = _wide_arrow_schema(20)
+    bm = _make_bm(schema)
+    state = bm.__getstate__()
+    assert isinstance(state["schema"], bytes)
+
+    # Prime the cache via the normal bytes path.
+    bytes_restored = pickle.loads(pickle.dumps(bm))
+    assert bytes_restored.schema.equals(schema)
+    info_after_bytes = _read_arrow_schema_cached.cache_info()
+
+    # Now feed __setstate__ a bytearray with the same contents.
+    bytearray_state = dict(state)
+    bytearray_state["schema"] = bytearray(state["schema"])
+    bytearray_restored = BlockMetadataWithSchema.from_metadata(bm.metadata, schema=None)
+    bytearray_restored.__setstate__(bytearray_state)
+
+    # It must be decoded to a real pa.Schema, not stored raw.
+    assert isinstance(bytearray_restored.schema, pa.Schema)
+    assert bytearray_restored.schema.equals(schema)
+
+    # And it must hit the same cache entry as the bytes path (no extra miss).
+    info_after_bytearray = _read_arrow_schema_cached.cache_info()
+    assert info_after_bytearray.misses == info_after_bytes.misses
+    assert info_after_bytearray.hits == info_after_bytes.hits + 1
+    assert bytearray_restored.schema is bytes_restored.schema
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-vv", __file__]))

python/ray/data/tests/unit/test_block_ref_counter.py

@@ -0,0 +1,156 @@
+import threading
+from typing import Callable, Dict
+
+import pytest
+
+import ray
+from ray.data._internal.execution.block_ref_counter import BlockRefCounter
+
+
+def _ref(uid: int) -> "ray.ObjectRef":
+    """Real ObjectRef with a deterministic distinct 28-byte ID, no Ray cluster needed."""
+    return ray.ObjectRef(uid.to_bytes(28, "big"))
+
+
+class FakeAddObjectOutOfScopeCallback:
+    """Test double for CoreWorker.add_object_out_of_scope_callback.
+
+    Records each registered callback keyed by the block's object-ID bytes so a
+    test can fire it explicitly. Set should_registration_fail=True to simulate an object
+    that is already out of scope at registration time.
+    """
+
+    def __init__(self, should_registration_fail: bool = False):
+        self._should_fail = should_registration_fail
+        self._callbacks: Dict[bytes, Callable[[bytes], None]] = {}
+
+    def __call__(
+        self, object_ref: "ray.ObjectRef", callback: Callable[[bytes], None]
+    ) -> bool:
+        if not self._should_fail:
+            self._callbacks[object_ref.binary()] = callback
+        return not self._should_fail
+
+    def fire(self, object_ref: "ray.ObjectRef") -> None:
+        id_binary = object_ref.binary()
+        self._callbacks[id_binary](id_binary)
+
+
+class TestBlockRefCounterAccounting:
+    def test_single_block_produced_and_released(self):
+        add_cb = FakeAddObjectOutOfScopeCallback()
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+        ref = _ref(1)
+
+        counter.on_block_produced(ref, 1, "op_a")
+        assert counter.get_object_store_memory_usage("op_a") == 1
+
+        add_cb.fire(ref)
+        assert counter.get_object_store_memory_usage("op_a") == 0
+
+    def test_multiple_blocks_same_producer(self):
+        add_cb = FakeAddObjectOutOfScopeCallback()
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+        ref1, ref2 = _ref(1), _ref(2)
+
+        counter.on_block_produced(ref1, 1, "op_a")
+        counter.on_block_produced(ref2, 1, "op_a")
+        assert counter.get_object_store_memory_usage("op_a") == 2
+
+        add_cb.fire(ref1)
+        assert counter.get_object_store_memory_usage("op_a") == 1
+        add_cb.fire(ref2)
+        assert counter.get_object_store_memory_usage("op_a") == 0
+
+    def test_multiple_producers_isolated(self):
+        add_cb = FakeAddObjectOutOfScopeCallback()
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+        ref1, ref2 = _ref(1), _ref(2)
+
+        counter.on_block_produced(ref1, 1, "op_a")
+        counter.on_block_produced(ref2, 1, "op_b")
+        assert counter.get_object_store_memory_usage("op_a") == 1
+        assert counter.get_object_store_memory_usage("op_b") == 1
+
+        add_cb.fire(ref1)
+        assert counter.get_object_store_memory_usage("op_a") == 0
+        assert counter.get_object_store_memory_usage("op_b") == 1
+
+    def test_duplicate_registration_is_noop(self):
+        """on_block_produced is idempotent: a duplicate ref is silently ignored.
+
+        This matters when an AllToAllOperator forwards an input ref unchanged;
+        the ref was already registered by the upstream producer."""
+        add_cb = FakeAddObjectOutOfScopeCallback()
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+        ref = _ref(1)
+
+        counter.on_block_produced(ref, 1, "op_a")
+        counter.on_block_produced(ref, 1, "op_a")
+        assert counter.get_object_store_memory_usage("op_a") == 1
+
+        add_cb.fire(ref)
+        assert counter.get_object_store_memory_usage("op_a") == 0
+
+    def test_register_when_object_already_out_of_scope(self):
+        """If registration reports the object is already gone, the increment is
+        undone immediately so the producer nets to zero."""
+        add_cb = FakeAddObjectOutOfScopeCallback(should_registration_fail=True)
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+
+        counter.on_block_produced(_ref(1), 1, "op_a")
+        assert counter.get_object_store_memory_usage("op_a") == 0
+
+
+class TestBlockRefCounterClear:
+    def test_clear_resets_usage(self):
+        add_cb = FakeAddObjectOutOfScopeCallback()
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+        counter.on_block_produced(_ref(1), 1, "op_a")
+        assert counter.get_object_store_memory_usage("op_a") == 1
+
+        counter.clear()
+        assert counter.get_object_store_memory_usage("op_a") == 0
+
+    def test_stale_callback_after_clear_is_noop(self):
+        """A stale callback firing after clear() must not touch accounting
+        recorded after the reset."""
+        add_cb = FakeAddObjectOutOfScopeCallback()
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+        stale_ref = _ref(1)
+        counter.on_block_produced(stale_ref, 1, "op_a")
+
+        counter.clear()
+
+        counter.on_block_produced(_ref(2), 1, "op_a")
+        assert counter.get_object_store_memory_usage("op_a") == 1
+
+        add_cb.fire(stale_ref)
+        assert counter.get_object_store_memory_usage("op_a") == 1
+
+
+class TestBlockRefCounterThreadSafety:
+    def test_concurrent_callbacks_dont_corrupt_state(self):
+        """Many threads firing callbacks at once must not corrupt the count."""
+        add_cb = FakeAddObjectOutOfScopeCallback()
+        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
+        producer_id = "op_concurrent"
+        n = 50
+        refs = [_ref(i) for i in range(n)]
+        for ref in refs:
+            counter.on_block_produced(ref, 1, producer_id)
+        assert counter.get_object_store_memory_usage(producer_id) == n
+
+        threads = [threading.Thread(target=add_cb.fire, args=(ref,)) for ref in refs]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        assert counter.get_object_store_memory_usage(producer_id) == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_bundler.py

@@ -0,0 +1,608 @@
+import uuid
+from typing import Any, List
+
+import pandas as pd
+import pytest
+
+import ray
+from ray.data._internal.execution.bundle_queue import (
+    EstimateSize,
+    ExactMultipleSize,
+    RebundleQueue,
+)
+from ray.data._internal.execution.interfaces.ref_bundle import BlockEntry, RefBundle
+from ray.data.block import BlockAccessor
+
+
+def _make_ref_bundles_for_unit_test(raw_bundles: List[List[List[Any]]]) -> tuple:
+    """Create RefBundles with fake object refs for unit testing (no Ray required).
+
+    Args:
+        raw_bundles: A list of bundles, where each bundle is a list of blocks,
+            and each block is a list of values.
+
+    Returns:
+        A tuple of (list of RefBundles, block_data_map) where block_data_map
+        maps fake object refs to their actual DataFrame data.
+    """
+    output_bundles = []
+    block_data_map = {}
+
+    for raw_bundle in raw_bundles:
+        blocks = []
+        schema = None
+        for raw_block in raw_bundle:
+            block = pd.DataFrame({"id": raw_block})
+            # Use UUID to generate unique fake object refs
+            block_ref = ray.ObjectRef(uuid.uuid4().hex[:28].encode())
+            block_data_map[block_ref] = block
+
+            blocks.append(
+                BlockEntry(block_ref, BlockAccessor.for_block(block).get_metadata())
+            )
+            schema = BlockAccessor.for_block(block).schema()
+
+        output_bundle = RefBundle(blocks=blocks, owns_blocks=True, schema=schema)
+        output_bundles.append(output_bundle)
+
+    return output_bundles, block_data_map
+
+
+def _get_bundle_values(bundle: RefBundle, block_data_map: dict) -> List[List[Any]]:
+    """Extract values from a bundle using block_data_map (no ray.get needed)."""
+    output = []
+    for block_ref in bundle.block_refs:
+        output.append(list(block_data_map[block_ref]["id"]))
+    return output
+
+
+@pytest.mark.parametrize(
+    "target,in_bundles,expected_row_counts",
+    [
+        (
+            # Target of 2 rows per bundle
+            2,
+            [[[1]], [[2]], [[3]], [[4]]],
+            [2, 2],  # Expected output: 2 bundles of 2 rows each
+        ),
+        (
+            # Target of 3 rows with uneven inputs
+            3,
+            [[[1, 2]], [[3, 4, 5]], [[6]]],
+            [3, 3],  # Expected: [1,2,3] and [4,5,6]
+        ),
+        (
+            # Target of 4 rows with leftover
+            4,
+            [[[1, 2]], [[3, 4]], [[5, 6, 7]]],
+            [4, 3],  # Expected: [1,2,3,4] and [5,6,7]
+        ),
+        (
+            # Larger target with various input sizes
+            5,
+            [[[1, 2, 3]], [[4, 5, 6, 7]], [[8, 9]], [[10, 11, 12]]],
+            [5, 5, 2],  # Expected: [1-5], [6-10], [11-12]
+        ),
+        (
+            # Test with empty blocks
+            3,
+            [[[1]], [[]], [[2, 3]], [[]], [[4, 5]]],
+            [3, 2],  # Expected: [1,2,3] and [4,5]
+        ),
+        (
+            # Test with last block smaller than target num rows per block
+            100,
+            [[[1]], [[2]], [[3]], [[4]], [[5]]],
+            [5],
+        ),
+    ],
+)
+def test_streaming_repartition_ref_bundler(target, in_bundles, expected_row_counts):
+    """Test RebundleQueue with various input patterns (unit test)."""
+
+    bundler = RebundleQueue(ExactMultipleSize(target))
+    bundles, block_data_map = _make_ref_bundles_for_unit_test(in_bundles)
+    out_bundles = []
+
+    for bundle in bundles:
+        bundler.add(bundle)
+
+    # NOTE: The check for num bundles/blocks is harder to reason about since we rebundle bundles together
+    og_total_size_bytes = bundler.estimate_size_bytes()
+    og_total_num_rows = bundler.num_rows()
+    assert sum(bundle.size_bytes() for bundle in bundles) == og_total_size_bytes
+    assert sum(bundle.num_rows() for bundle in bundles) == og_total_num_rows
+
+    all_original_bundles = []
+    while bundler.has_next():
+        out_bundle, original_bundles = bundler.get_next_with_original()
+        out_bundles.append(out_bundle)
+        all_original_bundles.extend(original_bundles)
+
+    bundler.finalize()
+
+    while bundler.has_next():
+        out_bundle, original_bundles = bundler.get_next_with_original()
+        out_bundles.append(out_bundle)
+        all_original_bundles.extend(original_bundles)
+
+    # Verify number of output bundles
+    assert len(out_bundles) == len(
+        expected_row_counts
+    ), f"Expected {len(expected_row_counts)} bundles, got {len(out_bundles)}"
+
+    # Verify row counts for each bundle
+    for i, (out_bundle, expected_count) in enumerate(
+        zip(out_bundles, expected_row_counts)
+    ):
+        assert (
+            out_bundle.num_rows() == expected_count
+        ), f"Bundle {i}: expected {expected_count} rows, got {out_bundle.num_rows()}"
+
+    # Verify all bundles have been ingested
+    assert bundler.num_blocks() == 0
+    assert bundler.num_rows() == 0
+    assert len(bundler) == 0
+    assert bundler.estimate_size_bytes() == 0
+
+    # Verify all output bundles except the last are exact multiples of target
+    for i, out_bundle in enumerate(out_bundles[:-1]):
+        assert (
+            out_bundle.num_rows() % target == 0
+        ), f"Bundle {i} has {out_bundle.num_rows()} rows, not a multiple of {target}"
+
+    # Verify data integrity - all input data is preserved in order (bundler slicing is in order)
+    total_input_rows = sum(sum(len(block) for block in bundle) for bundle in in_bundles)
+    total_output_rows = sum(bundle.num_rows() for bundle in out_bundles)
+    assert total_output_rows == total_input_rows
+
+    # Verify block content - extract all values from output bundles
+    output_values = []
+    total_num_rows = 0
+    total_size_bytes = 0
+    for bundle in out_bundles:
+        for entry, block_slice in zip(bundle.blocks, bundle.slices):
+            block_ref = entry.ref
+            # Look up the actual block data from our map (no ray.get needed)
+            data = block_data_map[block_ref]["id"]
+            if block_slice is not None:
+                # We apply the slice here manually because this is just for testing bundler
+                # and the block slicing is happened in map operator for streaming repartition
+                data = data[block_slice.start_offset : block_slice.end_offset]
+            output_values.extend(data)
+        total_num_rows += bundle.num_rows()
+        total_size_bytes += bundle.size_bytes()
+
+    assert og_total_size_bytes == total_size_bytes
+    assert og_total_num_rows == total_num_rows
+
+    # Expected values are all input values flattened in order
+    expected_values = [
+        value for bundle in in_bundles for block in bundle for value in block
+    ]
+    assert (
+        output_values == expected_values
+    ), f"Output values {output_values} don't match expected {expected_values}"
+
+    # Verify get_next_with_original tracks all non-empty original bundles
+    # (empty bundles are accumulated separately and have no originals)
+    non_empty_bundles = [b for b in bundles if b.num_rows() > 0]
+    assert len(all_original_bundles) == len(non_empty_bundles)
+    for orig, expected in zip(all_original_bundles, non_empty_bundles):
+        assert orig is expected
+
+
+def test_peek_next():
+    """Test that peek_next returns the next bundle without removing it."""
+    bundler = RebundleQueue(ExactMultipleSize(2))
+    bundles, _ = _make_ref_bundles_for_unit_test([[[1]], [[2]], [[3]]])
+
+    # Peek on empty queue returns None
+    assert bundler.peek_next() is None
+
+    # Add bundles until we have a ready bundle
+    bundler.add(bundles[0])
+    assert bundler.peek_next() is None  # Not enough rows yet
+
+    bundler.add(bundles[1])
+    assert bundler.has_next()
+
+    # Peek should return the bundle without removing it
+    peeked = bundler.peek_next()
+    assert peeked is not None
+    assert peeked.num_rows() == 2
+
+    # Peek again should return the same bundle
+    peeked2 = bundler.peek_next()
+    assert peeked2 is peeked
+
+    # Metrics should be unchanged after peek
+    initial_rows = bundler.num_rows()
+    initial_len = len(bundler)
+    bundler.peek_next()
+    assert bundler.num_rows() == initial_rows
+    assert len(bundler) == initial_len
+
+    # get_next should return the same bundle
+    got = bundler.get_next()
+    assert got.num_rows() == peeked.num_rows()
+
+
+def test_clear():
+    """Test that clear resets the bundler to empty state."""
+    bundler = RebundleQueue(ExactMultipleSize(2))
+    bundles, _ = _make_ref_bundles_for_unit_test([[[1]], [[2]], [[3]], [[4]]])
+
+    # Add some bundles
+    for bundle in bundles:
+        bundler.add(bundle)
+
+    # Verify bundler has content
+    assert bundler.has_next()
+    assert bundler.num_rows() > 0
+    assert len(bundler) > 0
+    assert bundler.estimate_size_bytes() > 0
+
+    # Clear the bundler
+    bundler.clear()
+
+    # Verify bundler is empty
+    assert not bundler.has_next()
+    assert bundler.num_rows() == 0
+    assert len(bundler) == 0
+    assert bundler.num_blocks() == 0
+    assert bundler.estimate_size_bytes() == 0
+    assert bundler.peek_next() is None
+
+    # Verify we can add bundles again after clear
+    new_bundles, _ = _make_ref_bundles_for_unit_test([[[10]], [[20]]])
+    for bundle in new_bundles:
+        bundler.add(bundle)
+
+    assert bundler.has_next()
+    out = bundler.get_next()
+    assert out.num_rows() == 2
+
+
+def test_add_updates_metrics():
+    """Test that add correctly updates queue metrics."""
+    bundler = RebundleQueue(ExactMultipleSize(10))  # High target so nothing gets built
+    bundles, _ = _make_ref_bundles_for_unit_test([[[1, 2]], [[3, 4, 5]]])
+
+    # Initially empty
+    assert bundler.num_rows() == 0
+    assert bundler.num_blocks() == 0
+    assert bundler.estimate_size_bytes() == 0
+
+    # Add first bundle
+    bundler.add(bundles[0])
+    assert bundler.num_rows() == 2
+    assert bundler.num_blocks() == 1
+    assert bundler.estimate_size_bytes() == bundles[0].size_bytes()
+
+    # Add second bundle
+    bundler.add(bundles[1])
+    assert bundler.num_rows() == 5
+    assert bundler.num_blocks() == 2
+    expected_bytes = bundles[0].size_bytes() + bundles[1].size_bytes()
+    assert bundler.estimate_size_bytes() == expected_bytes
+
+
+# =============================================================================
+# Tests for EstimateSize strategy
+# =============================================================================
+
+
+@pytest.mark.parametrize(
+    "target,in_bundles,expected_bundles",
+    [
+        (
+            # Unit target, should leave unchanged.
+            1,
+            [
+                # Input bundles
+                [[1]],
+                [[2]],
+                [[3, 4]],
+                [[5]],
+            ],
+            [
+                # Output bundles
+                [[1]],
+                [[2]],
+                [[3, 4]],
+                [[5]],
+            ],
+        ),
+        (
+            # No target, should leave unchanged.
+            None,
+            [
+                # Input bundles
+                [[1]],
+                [[2]],
+                [[3, 4]],
+                [[5]],
+            ],
+            [
+                # Output bundles
+                [[1]],
+                [[2]],
+                [[3, 4]],
+                [[5]],
+            ],
+        ),
+        (
+            # Proper handling of empty blocks
+            2,
+            [
+                # Input bundles
+                [[1]],
+                [[]],
+                [[]],
+                [[2, 3]],
+                [[]],
+                [[]],
+            ],
+            [
+                # Output bundles
+                [[1], [], [], [2, 3]],
+                [[], []],
+            ],
+        ),
+        (
+            # Test bundling, finalizing, passing, leftovers, etc.
+            2,
+            [
+                # Input bundles
+                [[1], [2]],
+                [[3, 4, 5]],
+                [[6], [7]],
+                [[8]],
+                [[9, 10], [11]],
+            ],
+            [[[1], [2]], [[3, 4, 5]], [[6], [7]], [[8], [9, 10], [11]]],
+        ),
+        (
+            # Test bundling, finalizing, passing, leftovers, etc.
+            3,
+            [
+                # Input bundles
+                [[1]],
+                [[2, 3]],
+                [[4, 5, 6, 7]],
+                [[8, 9], [10, 11]],
+            ],
+            [
+                # Output bundles
+                [[1], [2, 3]],
+                [[4, 5, 6, 7]],
+                [[8, 9], [10, 11]],
+            ],
+        ),
+    ],
+)
+def test_estimate_size_bundler_basic(target, in_bundles, expected_bundles):
+    """Test RebundleQueue with EstimateSize strategy creates expected output bundles."""
+    bundler = RebundleQueue(EstimateSize(target))
+    bundles, block_data_map = _make_ref_bundles_for_unit_test(in_bundles)
+    out_bundles = []
+    for bundle in bundles:
+        bundler.add(bundle)
+        while bundler.has_next():
+            out_bundle = _get_bundle_values(bundler.get_next(), block_data_map)
+            out_bundles.append(out_bundle)
+
+    bundler.finalize()
+
+    if bundler.has_next():
+        out_bundle = _get_bundle_values(bundler.get_next(), block_data_map)
+        out_bundles.append(out_bundle)
+
+    # Assert expected output
+    assert out_bundles == expected_bundles
+    # Assert that all bundles have been ingested
+    assert bundler.num_blocks() == 0
+
+    for bundle, expected in zip(out_bundles, expected_bundles):
+        assert bundle == expected
+
+
+@pytest.mark.parametrize(
+    "target,n,num_bundles,num_out_bundles,out_bundle_size",
+    [
+        (5, 20, 20, 4, 5),
+        (5, 24, 10, 4, 6),
+        (8, 16, 4, 2, 8),
+    ],
+)
+def test_estimate_size_bundler_uniform(
+    target, n, num_bundles, num_out_bundles, out_bundle_size
+):
+    """Test RebundleQueue with EstimateSize creates expected number of bundles."""
+    import numpy as np
+
+    bundler = RebundleQueue(EstimateSize(target))
+    data = np.arange(n)
+    pre_bundles = [arr.tolist() for arr in np.array_split(data, num_bundles)]
+    # Convert to expected format: each bundle has one block
+    raw_bundles = [[list(arr)] for arr in pre_bundles]
+    bundles, block_data_map = _make_ref_bundles_for_unit_test(raw_bundles)
+
+    out_bundles = []
+    for bundle in bundles:
+        bundler.add(bundle)
+        while bundler.has_next():
+            out_bundle = bundler.get_next()
+            out_bundles.append(out_bundle)
+    bundler.finalize()
+    if bundler.has_next():
+        out_bundle = bundler.get_next()
+        out_bundles.append(out_bundle)
+
+    assert len(out_bundles) == num_out_bundles
+    for out_bundle in out_bundles:
+        assert out_bundle.num_rows() == out_bundle_size
+
+    flat_out = [
+        i
+        for bundle in out_bundles
+        for block_ref in bundle.block_refs
+        for i in list(block_data_map[block_ref]["id"])
+    ]
+    assert flat_out == list(range(n))
+
+
+def test_estimate_size_peek_next():
+    """Test peek_next with EstimateSize strategy."""
+    bundler = RebundleQueue(EstimateSize(2))
+    bundles, _ = _make_ref_bundles_for_unit_test([[[1]], [[2]], [[3]]])
+
+    # Peek on empty queue returns None
+    assert bundler.peek_next() is None
+
+    # Add bundles until we have a ready bundle
+    bundler.add(bundles[0])
+    assert bundler.peek_next() is None  # Not enough rows yet
+
+    bundler.add(bundles[1])
+    assert bundler.has_next()
+
+    # Peek should return the bundle without removing it
+    peeked = bundler.peek_next()
+    assert peeked is not None
+    assert peeked.num_rows() == 2
+
+    # Peek again should return the same bundle
+    peeked2 = bundler.peek_next()
+    assert peeked2 is peeked
+
+    # get_next should return the same bundle
+    got = bundler.get_next()
+    assert got.num_rows() == peeked.num_rows()
+
+
+def test_estimate_size_clear():
+    """Test clear with EstimateSize strategy."""
+    bundler = RebundleQueue(EstimateSize(2))
+    bundles, _ = _make_ref_bundles_for_unit_test([[[1]], [[2]], [[3]], [[4]]])
+
+    # Add some bundles
+    for bundle in bundles:
+        bundler.add(bundle)
+
+    # Verify bundler has content
+    assert bundler.has_next()
+    assert bundler.num_rows() > 0
+
+    # Clear the bundler
+    bundler.clear()
+
+    # Verify bundler is empty
+    assert not bundler.has_next()
+    assert bundler.num_rows() == 0
+    assert len(bundler) == 0
+    assert bundler.num_blocks() == 0
+
+    # Verify we can add bundles again after clear
+    new_bundles, _ = _make_ref_bundles_for_unit_test([[[10]], [[20]]])
+    for bundle in new_bundles:
+        bundler.add(bundle)
+
+    assert bundler.has_next()
+    out = bundler.get_next()
+    assert out.num_rows() == 2
+
+
+def test_estimate_size_add_updates_metrics():
+    """Test add updates metrics with EstimateSize strategy."""
+    bundler = RebundleQueue(EstimateSize(10))  # High target so nothing gets built
+    bundles, _ = _make_ref_bundles_for_unit_test([[[1, 2]], [[3, 4, 5]]])
+
+    # Initially empty
+    assert bundler.num_rows() == 0
+    assert bundler.num_blocks() == 0
+    assert bundler.estimate_size_bytes() == 0
+
+    # Add first bundle
+    bundler.add(bundles[0])
+    assert bundler.num_rows() == 2
+    assert bundler.num_blocks() == 1
+    assert bundler.estimate_size_bytes() == bundles[0].size_bytes()
+
+    # Add second bundle
+    bundler.add(bundles[1])
+    assert bundler.num_rows() == 5
+    assert bundler.num_blocks() == 2
+    expected_bytes = bundles[0].size_bytes() + bundles[1].size_bytes()
+    assert bundler.estimate_size_bytes() == expected_bytes
+
+
+def test_empty_bundle_merges_with_previous_pending():
+    """Test that empty bundles merge into the last pending bundle
+    rather than accumulating separately."""
+    bundler = RebundleQueue(EstimateSize(3))
+    bundles, _ = _make_ref_bundles_for_unit_test([[[1, 2]], [[]], [[]], [[3]]])
+
+    for bundle in bundles:
+        bundler.add(bundle)
+
+    # The two empty bundles merged into the [1,2] pending bundle,
+    # then [3] arrived pushing total to 3 → ready bundle built.
+    assert bundler.has_next()
+    out = bundler.get_next()
+    assert out.num_rows() == 3
+    # Should have 4 blocks: [1,2], [], [], [3]
+    assert len(out.block_refs) == 4
+
+    assert bundler.num_rows() == 0
+    assert len(bundler) == 0
+
+
+def test_empty_bundle_no_previous_pending():
+    """Test that empty bundles with no previous pending just go to pending
+    and merge with subsequent empties."""
+    bundler = RebundleQueue(EstimateSize(2))
+    bundles, _ = _make_ref_bundles_for_unit_test([[[]], [[]], [[1, 2]]])
+
+    for bundle in bundles:
+        bundler.add(bundle)
+
+    # Two empties merged together in pending, then [1,2] triggers ready.
+    assert bundler.has_next()
+    out = bundler.get_next()
+    assert out.num_rows() == 2
+    # 3 blocks: [], [], [1,2]
+    assert len(out.block_refs) == 3
+
+    assert bundler.num_rows() == 0
+    assert len(bundler) == 0
+
+
+def test_empty_bundle_only():
+    """Test that a queue receiving only empty bundles flushes at finalize."""
+    bundler = RebundleQueue(EstimateSize(2))
+    bundles, _ = _make_ref_bundles_for_unit_test([[[]], [[]]])
+
+    for bundle in bundles:
+        bundler.add(bundle)
+
+    # Empties sit in pending (merged together), can't trigger a ready bundle
+    assert not bundler.has_next()
+
+    bundler.finalize()
+    assert bundler.has_next()
+
+    out = bundler.get_next()
+    assert out.num_rows() == 0
+    assert len(out.block_refs) == 2
+
+    assert not bundler.has_next()
+    assert bundler.num_rows() == 0
+    assert len(bundler) == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_data_batch_conversion.py

@@ -0,0 +1,361 @@
+import warnings
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.tensor_extensions.arrow import ArrowTensorArray
+from ray.data._internal.tensor_extensions.pandas import TensorArray, TensorDtype
+from ray.data.constants import TENSOR_COLUMN_NAME
+from ray.data.util.data_batch_conversion import (
+    BatchFormat,
+    _cast_ndarray_columns_to_tensor_extension,
+    _cast_tensor_columns_to_ndarrays,
+    _convert_batch_type_to_numpy,
+    _convert_batch_type_to_pandas,
+    _convert_pandas_to_batch_type,
+)
+from ray.data.util.torch_utils import convert_ndarray_to_torch_tensor
+
+
+def test_pandas_pandas():
+    input_data = pd.DataFrame({"x": [1, 2, 3]})
+    expected_output = input_data
+    actual_output = _convert_batch_type_to_pandas(input_data)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    actual_output = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.PANDAS
+    )
+    pd.testing.assert_frame_equal(actual_output, input_data)
+
+
+def test_numpy_to_numpy():
+    input_data = {"x": np.arange(12).reshape(3, 4)}
+    expected_output = input_data
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    assert expected_output == actual_output
+
+    input_data = {
+        "column_1": np.arange(12).reshape(3, 4),
+        "column_2": np.arange(12).reshape(3, 4),
+    }
+    expected_output = {
+        "column_1": np.arange(12).reshape(3, 4),
+        "column_2": np.arange(12).reshape(3, 4),
+    }
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    assert input_data.keys() == expected_output.keys()
+    np.testing.assert_array_equal(input_data["column_1"], expected_output["column_1"])
+    np.testing.assert_array_equal(input_data["column_2"], expected_output["column_2"])
+
+    input_data = np.arange(12).reshape(3, 4)
+    expected_output = input_data
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    np.testing.assert_array_equal(expected_output, actual_output)
+
+
+def test_arrow_to_numpy():
+    input_data = pa.table({"column_1": [1, 2, 3, 4]})
+    expected_output = {"column_1": np.array([1, 2, 3, 4])}
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    assert expected_output.keys() == actual_output.keys()
+    np.testing.assert_array_equal(
+        expected_output["column_1"], actual_output["column_1"]
+    )
+
+    input_data = pa.table(
+        {
+            TENSOR_COLUMN_NAME: ArrowTensorArray.from_numpy(
+                np.arange(12).reshape(3, 2, 2)
+            )
+        }
+    )
+    expected_output = np.arange(12).reshape(3, 2, 2)
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    np.testing.assert_array_equal(expected_output, actual_output)
+
+    input_data = pa.table(
+        {
+            "column_1": [1, 2, 3, 4],
+            "column_2": [1, -1, 1, -1],
+        }
+    )
+    expected_output = {
+        "column_1": np.array([1, 2, 3, 4]),
+        "column_2": np.array([1, -1, 1, -1]),
+    }
+
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    assert expected_output.keys() == actual_output.keys()
+    np.testing.assert_array_equal(
+        expected_output["column_1"], actual_output["column_1"]
+    )
+    np.testing.assert_array_equal(
+        expected_output["column_2"], actual_output["column_2"]
+    )
+
+
+def test_pd_dataframe_to_numpy():
+    input_data = pd.DataFrame({"column_1": [1, 2, 3, 4]})
+    expected_output = np.array([1, 2, 3, 4])
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    np.testing.assert_array_equal(expected_output, actual_output)
+
+    input_data = pd.DataFrame(
+        {TENSOR_COLUMN_NAME: TensorArray(np.arange(12).reshape(3, 4))}
+    )
+    expected_output = np.arange(12).reshape(3, 4)
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    np.testing.assert_array_equal(expected_output, actual_output)
+
+    input_data = pd.DataFrame({"column_1": [1, 2, 3, 4], "column_2": [1, -1, 1, -1]})
+    expected_output = {
+        "column_1": np.array([1, 2, 3, 4]),
+        "column_2": np.array([1, -1, 1, -1]),
+    }
+    actual_output = _convert_batch_type_to_numpy(input_data)
+    assert expected_output.keys() == actual_output.keys()
+    np.testing.assert_array_equal(
+        expected_output["column_1"], actual_output["column_1"]
+    )
+    np.testing.assert_array_equal(
+        expected_output["column_2"], actual_output["column_2"]
+    )
+
+
+@pytest.mark.parametrize("use_tensor_extension_for_input", [True, False])
+@pytest.mark.parametrize("cast_tensor_columns", [True, False])
+def test_pandas_multi_dim_pandas(cast_tensor_columns, use_tensor_extension_for_input):
+    input_tensor = np.arange(12).reshape((3, 2, 2))
+    input_data = pd.DataFrame(
+        {
+            "x": TensorArray(input_tensor)
+            if use_tensor_extension_for_input
+            else list(input_tensor)
+        }
+    )
+    expected_output = pd.DataFrame(
+        {
+            "x": (
+                list(input_tensor)
+                if cast_tensor_columns or not use_tensor_extension_for_input
+                else TensorArray(input_tensor)
+            )
+        }
+    )
+    actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    actual_output = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.PANDAS, cast_tensor_columns=cast_tensor_columns
+    )
+    pd.testing.assert_frame_equal(actual_output, input_data)
+
+
+def test_no_pandas_future_warning():
+    """Tests that Pandas in-place FutureWarning is
+    suppressed during tensor extension casting."""
+
+    input_tensor = np.arange(12).reshape((3, 2, 2))
+    input_data = pd.DataFrame({"x": TensorArray(input_tensor)})
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", category=FutureWarning)
+        data_no_tensor_array = _cast_tensor_columns_to_ndarrays(input_data)
+        _cast_ndarray_columns_to_tensor_extension(data_no_tensor_array)
+
+
+@pytest.mark.parametrize("cast_tensor_columns", [True, False])
+def test_numpy_pandas(cast_tensor_columns):
+    input_data = np.array([1, 2, 3])
+    expected_output = pd.DataFrame({TENSOR_COLUMN_NAME: input_data})
+    actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    output_array = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
+    )
+
+    np.testing.assert_equal(output_array, input_data)
+
+
+@pytest.mark.parametrize("cast_tensor_columns", [True, False])
+def test_numpy_multi_dim_pandas(cast_tensor_columns):
+    input_data = np.arange(12).reshape((3, 2, 2))
+    expected_output = pd.DataFrame({TENSOR_COLUMN_NAME: list(input_data)})
+    actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    output_array = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
+    )
+    np.testing.assert_array_equal(np.array(list(output_array)), input_data)
+
+
+def test_numpy_object_pandas():
+    input_data = np.array([[1, 2, 3], [1]], dtype=object)
+    expected_output = pd.DataFrame({TENSOR_COLUMN_NAME: input_data})
+    actual_output = _convert_batch_type_to_pandas(input_data)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    np.testing.assert_array_equal(
+        _convert_pandas_to_batch_type(actual_output, type=BatchFormat.NUMPY), input_data
+    )
+
+
+@pytest.mark.parametrize("writable", [False, True])
+def test_numpy_to_tensor_warning(writable):
+    input_data = np.array([[1, 2, 3]], dtype=int)
+    input_data.setflags(write=writable)
+
+    with pytest.warns(None) as record:
+        tensor = convert_ndarray_to_torch_tensor(input_data)
+    assert not record.list, [w.message for w in record.list]
+    assert tensor is not None
+
+
+def test_dict_fail():
+    input_data = {"x": "y"}
+    with pytest.raises(ValueError):
+        _convert_batch_type_to_pandas(input_data)
+
+
+@pytest.mark.parametrize("cast_tensor_columns", [True, False])
+def test_dict_pandas(cast_tensor_columns):
+    input_data = {"x": np.array([1, 2, 3])}
+    expected_output = pd.DataFrame({"x": input_data["x"]})
+    actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    output_array = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
+    )
+
+    np.testing.assert_array_equal(output_array, input_data["x"])
+
+
+@pytest.mark.parametrize("cast_tensor_columns", [True, False])
+def test_dict_multi_dim_to_pandas(cast_tensor_columns):
+    tensor = np.arange(12).reshape((3, 2, 2))
+    input_data = {"x": tensor}
+    expected_output = pd.DataFrame({"x": list(tensor)})
+    actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    output_array = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
+    )
+    np.testing.assert_array_equal(np.array(list(output_array)), input_data["x"])
+
+
+@pytest.mark.parametrize("cast_tensor_columns", [True, False])
+def test_dict_pandas_multi_column(cast_tensor_columns):
+    array_dict = {"x": np.array([1, 2, 3]), "y": np.array([4, 5, 6])}
+    expected_output = pd.DataFrame(array_dict)
+    actual_output = _convert_batch_type_to_pandas(array_dict, cast_tensor_columns)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    output_dict = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.NUMPY, cast_tensor_columns=cast_tensor_columns
+    )
+
+    for k, v in output_dict.items():
+        np.testing.assert_array_equal(v, array_dict[k])
+
+
+def test_arrow_pandas():
+    df = pd.DataFrame({"x": [1, 2, 3]})
+    input_data = pa.Table.from_pandas(df)
+    expected_output = df
+    actual_output = _convert_batch_type_to_pandas(input_data)
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    assert _convert_pandas_to_batch_type(actual_output, type=BatchFormat.ARROW).equals(
+        input_data
+    )
+
+
+@pytest.mark.parametrize("cast_tensor_columns", [True, False])
+def test_arrow_tensor_pandas(cast_tensor_columns):
+    np_array = np.arange(12).reshape((3, 2, 2))
+    input_data = pa.Table.from_arrays(
+        [ArrowTensorArray.from_numpy(np_array)], names=["x"]
+    )
+    actual_output = _convert_batch_type_to_pandas(input_data, cast_tensor_columns)
+    expected_output = pd.DataFrame({"x": list(np_array)})
+    expected_output = pd.DataFrame(
+        {"x": (list(np_array) if cast_tensor_columns else TensorArray(np_array))}
+    )
+    pd.testing.assert_frame_equal(expected_output, actual_output)
+
+    arrow_output = _convert_pandas_to_batch_type(
+        actual_output, type=BatchFormat.ARROW, cast_tensor_columns=cast_tensor_columns
+    )
+    assert arrow_output.equals(input_data)
+
+
+def _make_object_column(arrays):
+    """Build a 1-D object-dtype ndarray whose elements are the given ndarrays.
+
+    ``np.array([...], dtype=object)`` would build a 2-D array, so we fill an
+    empty object array element-by-element to keep each ndarray as a cell value.
+    """
+    out = np.empty(len(arrays), dtype=object)
+    for i, arr in enumerate(arrays):
+        out[i] = arr
+    return out
+
+
+def test_cast_ndarray_columns_duplicate_names():
+    """
+    Casting ndarray to tensor columns must handle duplicate column names,
+    keeping each column's data intact.
+    """
+    col_a = [np.array([1, 2, 3]), np.array([4, 5, 6]), np.array([7, 8, 9])]
+    col_b = [np.array([10, 20]), np.array([30, 40]), np.array([50, 60])]
+    df = pd.DataFrame(
+        {"_a": _make_object_column(col_a), "_b": _make_object_column(col_b)}
+    )
+    df.columns = ["x", "x"]
+
+    actual = _cast_ndarray_columns_to_tensor_extension(df)
+
+    # Both physical columns must be cast to the tensor extension type.
+    assert [isinstance(dt, TensorDtype) for _, dt in actual.dtypes.items()] == [
+        True,
+        True,
+    ]
+    # Values must be preserved per physical column (no write-back broadcasting
+    # one column's data across both duplicate labels).
+    np.testing.assert_array_equal(np.array(list(actual.iloc[:, 0])), col_a)
+    np.testing.assert_array_equal(np.array(list(actual.iloc[:, 1])), col_b)
+
+
+def test_cast_tensor_columns_duplicate_names():
+    """
+    Casting tensor columns back to ndarrays must handle duplicate names,
+    keeping each column's data intact.
+    """
+    col_a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    col_b = np.array([[10, 20], [30, 40], [50, 60]])
+    df = pd.DataFrame({"_a": TensorArray(col_a), "_b": TensorArray(col_b)})
+    df.columns = ["x", "x"]
+
+    actual = _cast_tensor_columns_to_ndarrays(df)
+
+    # Neither physical column should remain a tensor extension column.
+    assert [isinstance(dt, TensorDtype) for _, dt in actual.dtypes.items()] == [
+        False,
+        False,
+    ]
+    # Each physical column must retain its own original data.
+    np.testing.assert_array_equal(np.array(list(actual.iloc[:, 0])), col_a)
+    np.testing.assert_array_equal(np.array(list(actual.iloc[:, 1])), col_b)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/data/tests/unit/test_dataset_repr.py

@@ -0,0 +1,61 @@
+import pyarrow as pa
+
+from ray.data._internal.dataset_repr import _build_dataset_ascii_repr_from_rows
+from ray.data.dataset import Schema
+
+
+def test_dataset_repr_from_rows_not_materialized():
+    schema = Schema(pa.schema([("a", pa.int64()), ("b", pa.string())]))
+    text = _build_dataset_ascii_repr_from_rows(
+        schema=schema,
+        num_rows=5,
+        dataset_name="test_ds",
+        is_materialized=False,
+        head_rows=[],
+        tail_rows=[],
+    )
+    assert text == (
+        "name: test_ds\n"
+        "shape: (5, 2)\n"
+        "╭───────┬────────╮\n"
+        "│ a     ┆ b      │\n"
+        "│ ---   ┆ ---    │\n"
+        "│ int64 ┆ string │\n"
+        "╰───────┴────────╯\n"
+        "(Dataset isn't materialized)"
+    )
+
+
+def test_dataset_repr_from_rows_gap():
+    schema = Schema(pa.schema([("id", pa.int64())]))
+    text = _build_dataset_ascii_repr_from_rows(
+        schema=schema,
+        num_rows=12,
+        dataset_name=None,
+        is_materialized=True,
+        head_rows=[["0"], ["1"]],
+        tail_rows=[["10"], ["11"]],
+    )
+    assert text == (
+        "shape: (12, 1)\n"
+        "╭───────╮\n"
+        "│ id    │\n"
+        "│ ---   │\n"
+        "│ int64 │\n"
+        "╞═══════╡\n"
+        "│ 0     │\n"
+        "│ 1     │\n"
+        "│ …     │\n"
+        "│ 10    │\n"
+        "│ 11    │\n"
+        "╰───────╯\n"
+        "(Showing 4 of 12 rows)"
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_datatype.py

@@ -0,0 +1,882 @@
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+from packaging import version
+
+from ray.data.datatype import DataType, TypeCategory
+
+# Skip all tests if PyArrow version is less than 19.0
+pytestmark = pytest.mark.skipif(
+    version.parse(pa.__version__) < version.parse("19.0.0"),
+    reason="DataType tests require PyArrow >= 19.0",
+)
+
+
+class TestDataTypeFactoryMethods:
+    """Test the generated factory methods."""
+
+    @pytest.mark.parametrize(
+        "method_name,pa_type,description",
+        [
+            ("int8", pa.int8(), "8-bit signed integer"),
+            ("int16", pa.int16(), "16-bit signed integer"),
+            ("int32", pa.int32(), "32-bit signed integer"),
+            ("int64", pa.int64(), "64-bit signed integer"),
+            ("uint8", pa.uint8(), "8-bit unsigned integer"),
+            ("uint16", pa.uint16(), "16-bit unsigned integer"),
+            ("uint32", pa.uint32(), "32-bit unsigned integer"),
+            ("uint64", pa.uint64(), "64-bit unsigned integer"),
+            ("float32", pa.float32(), "32-bit floating point number"),
+            ("float64", pa.float64(), "64-bit floating point number"),
+            ("string", pa.string(), "variable-length string"),
+            ("bool", pa.bool_(), "boolean value"),
+            ("binary", pa.binary(), "variable-length binary data"),
+        ],
+    )
+    def test_factory_method_creates_correct_type(
+        self, method_name, pa_type, description
+    ):
+        """Test that factory methods create DataType with correct PyArrow type."""
+        factory_method = getattr(DataType, method_name)
+        result = factory_method()
+
+        assert isinstance(result, DataType)
+        assert result.is_arrow_type()
+        assert result._physical_dtype == pa_type
+
+    @pytest.mark.parametrize(
+        "method_name",
+        [
+            "int8",
+            "int16",
+            "int32",
+            "int64",
+            "uint8",
+            "uint16",
+            "uint32",
+            "uint64",
+            "float32",
+            "float64",
+            "string",
+            "bool",
+            "binary",
+        ],
+    )
+    def test_factory_method_has_proper_docstring(self, method_name):
+        """Test that generated factory methods have proper docstrings."""
+        factory_method = getattr(DataType, method_name)
+        doc = factory_method.__doc__
+
+        assert "Create a DataType representing" in doc
+        assert "Returns:" in doc
+        assert f"DataType with PyArrow {method_name} type" in doc
+
+
+class TestDataTypeValidation:
+    """Test DataType validation and initialization."""
+
+    @pytest.mark.parametrize(
+        "valid_type",
+        [
+            pa.int64(),
+            pa.string(),
+            pa.timestamp("s"),
+            np.dtype("int32"),
+            np.dtype("float64"),
+            int,
+            str,
+            float,
+        ],
+    )
+    def test_post_init_accepts_valid_types(self, valid_type):
+        """Test that __post_init__ accepts valid type objects."""
+        # Should not raise
+        dt = DataType(valid_type)
+        assert dt._physical_dtype == valid_type
+
+    @pytest.mark.parametrize(
+        "invalid_type",
+        [
+            "string",
+            123,
+            [1, 2, 3],
+            {"key": "value"},
+            object(),
+        ],
+    )
+    def test_post_init_rejects_invalid_types(self, invalid_type):
+        """Test that __post_init__ rejects invalid type objects."""
+        with pytest.raises(
+            TypeError,
+            match="DataType supports only PyArrow DataType, NumPy dtype, or Python type",
+        ):
+            DataType(invalid_type)
+
+
+class TestDataTypeCheckers:
+    """Test type checking methods."""
+
+    @pytest.mark.parametrize(
+        "datatype,is_arrow,is_numpy,is_python",
+        [
+            (DataType.from_arrow(pa.int64()), True, False, False),
+            (DataType.from_arrow(pa.string()), True, False, False),
+            (DataType.from_numpy(np.dtype("int32")), False, True, False),
+            (DataType.from_numpy(np.dtype("float64")), False, True, False),
+            (DataType(int), False, False, True),
+            (DataType(str), False, False, True),
+        ],
+    )
+    def test_type_checkers(self, datatype, is_arrow, is_numpy, is_python):
+        """Test is_arrow_type, is_numpy_type, and is_python_type methods."""
+        assert datatype.is_arrow_type() == is_arrow
+        assert datatype.is_numpy_type() == is_numpy
+        assert datatype.is_python_type() == is_python
+
+
+class TestDataTypeFactories:
+    """Test factory methods from external systems."""
+
+    @pytest.mark.parametrize(
+        "pa_type",
+        [
+            pa.int32(),
+            pa.string(),
+            pa.timestamp("s"),
+            pa.list_(pa.int32()),
+            pa.decimal128(10, 2),
+        ],
+    )
+    def test_from_arrow(self, pa_type):
+        """Test from_arrow factory method."""
+        dt = DataType.from_arrow(pa_type)
+
+        assert isinstance(dt, DataType)
+        assert dt.is_arrow_type()
+        assert dt._physical_dtype == pa_type
+
+    @pytest.mark.parametrize(
+        "numpy_input,expected_dtype",
+        [
+            (np.dtype("int32"), np.dtype("int32")),
+            (np.dtype("float64"), np.dtype("float64")),
+            ("int64", np.dtype("int64")),
+            ("float32", np.dtype("float32")),
+        ],
+    )
+    def test_from_numpy(self, numpy_input, expected_dtype):
+        """Test from_numpy factory method."""
+        dt = DataType.from_numpy(numpy_input)
+
+        assert isinstance(dt, DataType)
+        assert dt.is_numpy_type()
+        assert dt._physical_dtype == expected_dtype
+
+    @pytest.mark.parametrize(
+        "cudf_dtype,expected_dtype",
+        [
+            ("int32", np.dtype("int32")),
+            ("uint64", np.dtype("uint64")),
+            ("float32", np.dtype("float32")),
+            ("float64", np.dtype("float64")),
+            ("bool", np.dtype("bool")),
+            ("datetime64[ns]", np.dtype("datetime64[ns]")),
+        ],
+    )
+    def test_from_cudf_primitives(self, cudf_dtype, expected_dtype):
+        """Test from_cudf handles primitive cuDF dtypes."""
+        cudf = pytest.importorskip("cudf")
+
+        dt = DataType.from_cudf(cudf.dtype(cudf_dtype))
+
+        assert dt.is_numpy_type()
+        assert dt._physical_dtype == expected_dtype
+
+    def test_from_cudf_extension_dtype(self):
+        """Test from_cudf handles cuDF dtypes with Arrow representations."""
+        pytest.importorskip("cudf")
+        from cudf.core import dtypes
+
+        dt = DataType.from_cudf(dtypes.ListDtype("int32"))
+
+        assert dt.is_arrow_type()
+        assert dt.to_arrow_dtype() == pa.list_(pa.int32())
+
+
+class TestDataTypeConversions:
+    """Test type conversion methods."""
+
+    def test_to_arrow_dtype_arrow_passthrough(self):
+        """Test that Arrow types return themselves."""
+        dt = DataType.from_arrow(pa.int64())
+        result = dt.to_arrow_dtype()
+        assert result == pa.int64()
+
+    def test_to_arrow_dtype_numpy_conversion(self):
+        """Test conversion from NumPy to Arrow types."""
+        dt = DataType.from_numpy(np.dtype("int32"))
+        result = dt.to_arrow_dtype()
+        assert result == pa.int32()
+
+    def test_to_arrow_dtype_python_conversion(self):
+        """Test conversion from Python to Arrow types."""
+        dt = DataType(int)
+        result = dt.to_arrow_dtype([1])
+        # Python int should map to int64 in Arrow
+        assert result == pa.int64()
+
+    @pytest.mark.parametrize(
+        "source_dt,expected_result",
+        [
+            # NumPy types should return themselves
+            (DataType.from_numpy(np.dtype("int32")), np.dtype("int32")),
+            (DataType.from_numpy(np.dtype("float64")), np.dtype("float64")),
+            # Python types should fall back to object
+            (DataType(str), np.dtype("object")),
+            (DataType(list), np.dtype("object")),
+        ],
+    )
+    def test_to_numpy_dtype(self, source_dt, expected_result):
+        """Test to_numpy_dtype conversion."""
+        result = source_dt.to_numpy_dtype()
+        assert result == expected_result
+
+    def test_to_numpy_dtype_arrow_basic_types(self):
+        """Test Arrow to NumPy conversion for types that should work."""
+        # Test basic types that should convert properly
+        test_cases = [
+            (pa.int32(), np.dtype("int32")),
+            (pa.float64(), np.dtype("float64")),
+            (pa.bool_(), np.dtype("bool")),
+        ]
+
+        for pa_type, expected_np_dtype in test_cases:
+            dt = DataType.from_arrow(pa_type)
+            result = dt.to_numpy_dtype()
+            # Some Arrow types may not convert exactly as expected,
+            # so let's just verify the result is a valid numpy dtype
+            assert isinstance(result, np.dtype)
+
+    def test_to_numpy_dtype_complex_arrow_fallback(self):
+        """Test that complex Arrow types fall back to object dtype."""
+        complex_dt = DataType.from_arrow(pa.list_(pa.int32()))
+        result = complex_dt.to_numpy_dtype()
+        assert result == np.dtype("object")
+
+    def test_to_cudf_type(self):
+        """Test conversion to cuDF-compatible dtypes."""
+        cudf = pytest.importorskip("cudf")
+        from cudf.core import dtypes as cudf_dtypes
+
+        assert DataType.from_numpy("int32").to_cudf_type() == cudf.dtype("int32")
+        assert DataType.int32().to_cudf_type() == cudf.dtype("int32")
+        assert DataType.from_numpy("uint64").to_cudf_type() == cudf.dtype("uint64")
+        assert DataType.from_numpy("float64").to_cudf_type() == cudf.dtype("float64")
+        assert DataType.from_numpy("bool").to_cudf_type() == cudf.dtype("bool")
+        assert DataType.string().to_cudf_type() == "str"
+        assert DataType.from_arrow(pa.decimal32(7, 2)).to_cudf_type() == (
+            cudf_dtypes.Decimal32Dtype(7, 2)
+        )
+        assert DataType.from_arrow(pa.decimal64(18, 2)).to_cudf_type() == (
+            cudf_dtypes.Decimal64Dtype(18, 2)
+        )
+        assert DataType.from_arrow(pa.decimal128(38, 2)).to_cudf_type() == (
+            cudf_dtypes.Decimal128Dtype(38, 2)
+        )
+
+        with pytest.raises(TypeError):
+            DataType.from_arrow(pa.decimal256(40, 2)).to_cudf_type()
+
+    def test_cudf_methods_require_cudf(self, monkeypatch):
+        """Test cuDF conversion methods fail clearly without cuDF installed."""
+        import builtins
+
+        real_import = builtins.__import__
+
+        def import_without_cudf(name, *args, **kwargs):
+            if name == "cudf" or name.startswith("cudf."):
+                raise ImportError("No module named 'cudf'")
+            return real_import(name, *args, **kwargs)
+
+        monkeypatch.setattr(builtins, "__import__", import_without_cudf)
+
+        with pytest.raises(ImportError, match="requires cuDF"):
+            DataType.from_cudf("int32")
+
+        with pytest.raises(ImportError, match="requires cuDF"):
+            DataType.from_numpy("int32").to_cudf_type()
+
+    @pytest.mark.parametrize("python_type", [int, str, float, bool, list])
+    def test_to_python_type_success(self, python_type):
+        """Test to_python_type returns the original Python type."""
+        dt = DataType(python_type)
+        result = dt.to_python_type()
+        assert result == python_type
+
+    @pytest.mark.parametrize(
+        "non_python_dt",
+        [
+            DataType.from_arrow(pa.int64()),
+            DataType.from_numpy(np.dtype("float32")),
+        ],
+    )
+    def test_to_python_type_failure(self, non_python_dt):
+        """Test to_python_type raises ValueError for non-Python types."""
+        with pytest.raises(ValueError, match="is not backed by a Python type"):
+            non_python_dt.to_python_type()
+
+
+class TestDataTypeFromDtype:
+    """Test conversion from dtype descriptors."""
+
+    def test_from_dtype_dispatch(self):
+        assert DataType.from_dtype(pa.int32()) == DataType.from_arrow(pa.int32())
+        assert DataType.from_dtype(np.dtype("int32")) == DataType.from_numpy("int32")
+        assert DataType.from_dtype("float32") == DataType.from_numpy("float32")
+        assert DataType.from_dtype(np.int64) == DataType.from_numpy("int64")
+        assert DataType.from_dtype(int) == DataType(int)
+        assert DataType.from_dtype(pd.StringDtype()) == DataType.from_arrow(pa.string())
+
+    def test_from_dtype_pandas_extension_dtypes(self):
+        test_cases = [
+            (pd.ArrowDtype(pa.int32()), DataType.from_arrow(pa.int32())),
+            (pd.Int8Dtype(), DataType.from_numpy("int8")),
+            (pd.UInt64Dtype(), DataType.from_numpy("uint64")),
+            (pd.Float32Dtype(), DataType.from_numpy("float32")),
+            (pd.BooleanDtype(), DataType.from_numpy("bool")),
+            (pd.StringDtype(), DataType.from_arrow(pa.string())),
+        ]
+        for pandas_dtype, expected_dtype in test_cases:
+            assert DataType.from_dtype(pandas_dtype) == expected_dtype
+
+    def test_from_dtype_passthrough_and_unknown(self):
+        dt = DataType.from_numpy("int64")
+        assert DataType.from_dtype(dt) is dt
+        assert DataType.from_dtype(None) is None
+        assert DataType.from_dtype(object()) is None
+
+    def test_from_dtype_unsupported_strings_return_none(self, monkeypatch):
+        def raise_import_error(cls, dtype):
+            raise ImportError("No module named 'cudf'")
+
+        monkeypatch.setattr(DataType, "from_cudf", classmethod(raise_import_error))
+
+        assert DataType.from_dtype("string") is None
+        assert DataType.from_dtype("category") is None
+
+    def test_from_dtype_ignores_cudf_runtime_failure(self, monkeypatch):
+        def raise_runtime_error(cls, dtype):
+            raise RuntimeError("CUDA runtime unavailable")
+
+        monkeypatch.setattr(DataType, "from_cudf", classmethod(raise_runtime_error))
+
+        assert DataType.from_dtype(object()) is None
+
+
+class TestDataTypeInference:
+    """Test type inference from values."""
+
+    @pytest.mark.parametrize(
+        "numpy_value,expected_dtype",
+        [
+            (np.array([1, 2, 3], dtype="int32"), np.dtype("int32")),
+            (np.array([1.0, 2.0], dtype="float64"), np.dtype("float64")),
+            (np.int64(42), np.dtype("int64")),
+            (np.float32(3.14), np.dtype("float32")),
+        ],
+    )
+    def test_infer_dtype_numpy_values(self, numpy_value, expected_dtype):
+        """Test inference of NumPy arrays and scalars."""
+        dt = DataType.infer_dtype(numpy_value)
+
+        assert dt.is_numpy_type()
+        assert dt._physical_dtype == expected_dtype
+
+    @pytest.mark.parametrize(
+        "python_value",
+        [
+            42,  # int
+            3.14,  # float
+            "hello",  # str
+            True,  # bool
+            [1, 2, 3],  # list
+        ],
+    )
+    def test_infer_dtype_python_values_arrow_success(self, python_value):
+        """Test inference of Python values that Arrow can handle."""
+        dt = DataType.infer_dtype(python_value)
+
+        # Should infer to Arrow type for basic Python values
+        assert dt.is_arrow_type()
+
+
+class TestDataTypeStringRepresentation:
+    """Test string representation methods."""
+
+    @pytest.mark.parametrize(
+        "datatype,expected_repr",
+        [
+            (DataType.from_arrow(pa.int64()), "DataType(arrow:int64)"),
+            (DataType.from_arrow(pa.string()), "DataType(arrow:string)"),
+            (DataType.from_numpy(np.dtype("float32")), "DataType(numpy:float32)"),
+            (DataType.from_numpy(np.dtype("int64")), "DataType(numpy:int64)"),
+            (DataType(str), "DataType(python:str)"),
+            (DataType(int), "DataType(python:int)"),
+        ],
+    )
+    def test_repr(self, datatype, expected_repr):
+        """Test __repr__ method for different type categories."""
+        assert repr(datatype) == expected_repr
+
+
+class TestDataTypeEqualityAndHashing:
+    """Test equality and hashing behavior."""
+
+    @pytest.mark.parametrize(
+        "dt1,dt2,should_be_equal",
+        [
+            # Same types should be equal
+            (DataType.from_arrow(pa.int64()), DataType.from_arrow(pa.int64()), True),
+            (
+                DataType.from_numpy(np.dtype("float32")),
+                DataType.from_numpy(np.dtype("float32")),
+                True,
+            ),
+            (DataType(str), DataType(str), True),
+            # Different Arrow types should not be equal
+            (DataType.from_arrow(pa.int64()), DataType.from_arrow(pa.int32()), False),
+            # Same conceptual type but different systems should not be equal
+            (
+                DataType.from_arrow(pa.int64()),
+                DataType.from_numpy(np.dtype("int64")),
+                False,
+            ),
+        ],
+    )
+    def test_equality(self, dt1, dt2, should_be_equal):
+        """Test __eq__ method."""
+        if should_be_equal:
+            assert dt1 == dt2
+            assert hash(dt1) == hash(dt2)
+        else:
+            assert dt1 != dt2
+
+    def test_numpy_vs_python_inequality(self):
+        """Test that numpy int64 and python int are not equal."""
+        numpy_dt = DataType.from_numpy(np.dtype("int64"))
+        python_dt = DataType(int)
+
+        # These represent the same conceptual type but with different systems
+        # so they should not be equal
+
+        # First verify they have different internal types
+        assert type(numpy_dt._physical_dtype) is not type(python_dt._physical_dtype)
+        assert numpy_dt._physical_dtype is not python_dt._physical_dtype
+
+        # Test the type checkers return different results
+        assert numpy_dt.is_numpy_type() and not python_dt.is_numpy_type()
+        assert python_dt.is_python_type() and not numpy_dt.is_python_type()
+
+        # They should not be equal
+        assert numpy_dt != python_dt
+
+    @pytest.mark.parametrize(
+        "non_datatype_value",
+        [
+            "not_a_datatype",
+            42,
+            [1, 2, 3],
+            {"key": "value"},
+            None,
+        ],
+    )
+    def test_inequality_with_non_datatype(self, non_datatype_value):
+        """Test that DataType is not equal to non-DataType objects."""
+        dt = DataType.from_arrow(pa.int64())
+        assert dt != non_datatype_value
+
+    def test_hashability(self):
+        """Test that DataType objects can be used in sets and as dict keys."""
+        dt1 = DataType.from_arrow(pa.int64())
+        dt2 = DataType.from_arrow(pa.int64())  # Same as dt1
+        dt3 = DataType.from_arrow(pa.int32())  # Different
+
+        # Test in set
+        dt_set = {dt1, dt2, dt3}
+        assert len(dt_set) == 2  # dt1 and dt2 are the same
+
+        # Test as dict keys
+        dt_dict = {dt1: "first", dt3: "second"}
+        assert dt_dict[dt2] == "first"  # dt2 should match dt1
+
+
+class TestIsOf:
+    """Test is_of method with TypeCategory."""
+
+    @pytest.mark.parametrize(
+        "dt,category,expected",
+        [
+            (DataType.list(DataType.int64()), TypeCategory.LIST, True),
+            (DataType.large_list(DataType.int64()), TypeCategory.LIST, True),
+            (DataType.large_list(DataType.int64()), TypeCategory.LARGE_LIST, True),
+            (DataType.fixed_size_list(DataType.int32(), 3), TypeCategory.LIST, True),
+            (DataType.struct([("x", DataType.int64())]), TypeCategory.STRUCT, True),
+            (DataType.map(DataType.string(), DataType.int64()), TypeCategory.MAP, True),
+            (DataType.tensor((3, 4), DataType.float32()), TypeCategory.TENSOR, True),
+            (DataType.temporal("timestamp"), TypeCategory.TEMPORAL, True),
+            (DataType.temporal("date32"), TypeCategory.TEMPORAL, True),
+            # Negatives
+            (DataType.int64(), TypeCategory.LIST, False),
+            (DataType.list(DataType.int64()), TypeCategory.LARGE_LIST, False),
+            (
+                DataType.fixed_size_list(DataType.int32(), 3),
+                TypeCategory.LARGE_LIST,
+                False,
+            ),
+            (DataType.list(DataType.int64()), TypeCategory.STRUCT, False),
+            (DataType.struct([("x", DataType.int64())]), TypeCategory.MAP, False),
+        ],
+    )
+    def test_is_of(self, dt, category, expected):
+        assert dt.is_of(category) == expected
+        # Test with string representation too
+        assert dt.is_of(category.value) == expected
+
+    def test_is_of_invalid_category(self):
+        dt = DataType.int64()
+        assert dt.is_of("invalid_category") is False
+
+
+class TestNestedTypeFactories:
+    """Test factory methods for nested types (list, struct, map, etc.)."""
+
+    @pytest.mark.parametrize(
+        "factory_call,expected_arrow_type",
+        [
+            (lambda: DataType.list(DataType.int64()), pa.list_(pa.int64())),
+            (lambda: DataType.list(DataType.string()), pa.list_(pa.string())),
+            (
+                lambda: DataType.large_list(DataType.float32()),
+                pa.large_list(pa.float32()),
+            ),
+            (
+                lambda: DataType.fixed_size_list(DataType.int32(), 5),
+                pa.list_(pa.int32(), 5),
+            ),
+        ],
+    )
+    def test_list_type_factories(self, factory_call, expected_arrow_type):
+        """Test list-type factory methods create correct Arrow types."""
+        dt = factory_call()
+        assert dt.is_arrow_type()
+        assert dt._physical_dtype == expected_arrow_type
+
+    @pytest.mark.parametrize(
+        "fields,expected_arrow_type",
+        [
+            (
+                [("x", DataType.int64()), ("y", DataType.float64())],
+                pa.struct([("x", pa.int64()), ("y", pa.float64())]),
+            ),
+            (
+                [("name", DataType.string()), ("age", DataType.int32())],
+                pa.struct([("name", pa.string()), ("age", pa.int32())]),
+            ),
+        ],
+    )
+    def test_struct_factory(self, fields, expected_arrow_type):
+        """Test struct factory method creates correct Arrow types."""
+        dt = DataType.struct(fields)
+        assert dt.is_arrow_type()
+        assert dt._physical_dtype == expected_arrow_type
+
+    @pytest.mark.parametrize(
+        "key_type,value_type,expected_arrow_type",
+        [
+            (DataType.string(), DataType.int64(), pa.map_(pa.string(), pa.int64())),
+            (DataType.int32(), DataType.float32(), pa.map_(pa.int32(), pa.float32())),
+        ],
+    )
+    def test_map_factory(self, key_type, value_type, expected_arrow_type):
+        """Test map factory method creates correct Arrow types."""
+        dt = DataType.map(key_type, value_type)
+        assert dt.is_arrow_type()
+        assert dt._physical_dtype == expected_arrow_type
+
+    @pytest.mark.parametrize(
+        "temporal_type,unit,tz,expected_type",
+        [
+            ("timestamp", "s", None, pa.timestamp("s")),
+            ("timestamp", "us", "UTC", pa.timestamp("us", tz="UTC")),
+            ("date32", None, None, pa.date32()),
+            ("date64", None, None, pa.date64()),
+            ("time32", "s", None, pa.time32("s")),
+            ("time64", "us", None, pa.time64("us")),
+            ("duration", "ms", None, pa.duration("ms")),
+        ],
+    )
+    def test_temporal_factory(self, temporal_type, unit, tz, expected_type):
+        """Test temporal factory method creates correct Arrow types."""
+        if tz is not None:
+            dt = DataType.temporal(temporal_type, unit=unit, tz=tz)
+        elif unit is not None:
+            dt = DataType.temporal(temporal_type, unit=unit)
+        else:
+            dt = DataType.temporal(temporal_type)
+
+        assert dt.is_arrow_type()
+        assert dt._physical_dtype == expected_type
+
+    @pytest.mark.parametrize(
+        "temporal_type,unit,error_msg",
+        [
+            ("time32", "us", "time32 unit must be 's' or 'ms'"),
+            ("time64", "ms", "time64 unit must be 'us' or 'ns'"),
+            ("invalid", None, "Invalid temporal_type"),
+        ],
+    )
+    def test_temporal_factory_validation(self, temporal_type, unit, error_msg):
+        """Test temporal factory validates inputs correctly."""
+        with pytest.raises(ValueError, match=error_msg):
+            DataType.temporal(temporal_type, unit=unit)
+
+
+class TestTypePredicates:
+    """Test type predicate methods (is_list_type, is_struct_type, etc.)."""
+
+    @pytest.mark.parametrize(
+        "datatype,expected",
+        [
+            (
+                DataType.from_arrow(pa.null()),
+                {
+                    "null": True,
+                    "boolean": False,
+                    "integer": False,
+                    "floating": False,
+                    "uint64": False,
+                },
+            ),
+            (
+                DataType.bool(),
+                {
+                    "null": False,
+                    "boolean": True,
+                    "integer": False,
+                    "floating": False,
+                    "uint64": False,
+                },
+            ),
+            (
+                DataType.int64(),
+                {
+                    "null": False,
+                    "boolean": False,
+                    "integer": True,
+                    "floating": False,
+                    "uint64": False,
+                },
+            ),
+            (
+                DataType.from_numpy("uint64"),
+                {
+                    "null": False,
+                    "boolean": False,
+                    "integer": True,
+                    "floating": False,
+                    "uint64": True,
+                },
+            ),
+            (
+                DataType.float32(),
+                {
+                    "null": False,
+                    "boolean": False,
+                    "integer": False,
+                    "floating": True,
+                    "uint64": False,
+                },
+            ),
+            (
+                DataType(bool),
+                {
+                    "null": False,
+                    "boolean": True,
+                    "integer": False,
+                    "floating": False,
+                    "uint64": False,
+                },
+            ),
+        ],
+    )
+    def test_scalar_dtype_predicates(self, datatype, expected):
+        """Test scalar dtype predicate helpers."""
+        assert datatype.is_null_type() == expected["null"]
+        assert datatype.is_boolean_type() == expected["boolean"]
+        assert datatype.is_integer_type() == expected["integer"]
+        assert datatype.is_floating_type() == expected["floating"]
+        assert datatype.is_uint64_type() == expected["uint64"]
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            # List types
+            (DataType.list(DataType.int64()), True),
+            (DataType.large_list(DataType.string()), True),
+            (DataType.fixed_size_list(DataType.float32(), 3), True),
+            # Tensor types (should return False)
+            (DataType.tensor(shape=(3, 4), dtype=DataType.float32()), False),
+            (DataType.variable_shaped_tensor(dtype=DataType.float64(), ndim=2), False),
+            # Non-list types
+            (DataType.int64(), False),
+            (DataType.string(), False),
+            (DataType.struct([("x", DataType.int32())]), False),
+        ],
+    )
+    def test_is_list_type(self, datatype, expected_result):
+        """Test is_list_type predicate."""
+        assert datatype.is_list_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            (DataType.tensor(shape=(3, 4), dtype=DataType.float32()), True),
+            (DataType.variable_shaped_tensor(dtype=DataType.float64(), ndim=2), True),
+        ],
+    )
+    def test_is_tensor_type(self, datatype, expected_result):
+        """Test is_tensor_type predicate."""
+        assert datatype.is_tensor_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            (DataType.struct([("x", DataType.int64())]), True),
+            (
+                DataType.struct([("a", DataType.string()), ("b", DataType.float32())]),
+                True,
+            ),
+            (DataType.list(DataType.int64()), False),
+            (DataType.int64(), False),
+        ],
+    )
+    def test_is_struct_type(self, datatype, expected_result):
+        """Test is_struct_type predicate."""
+        assert datatype.is_struct_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            (DataType.map(DataType.string(), DataType.int64()), True),
+            (DataType.map(DataType.int32(), DataType.float32()), True),
+            (DataType.list(DataType.int64()), False),
+            (DataType.int64(), False),
+        ],
+    )
+    def test_is_map_type(self, datatype, expected_result):
+        """Test is_map_type predicate."""
+        assert datatype.is_map_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            # Nested types
+            (DataType.list(DataType.int64()), True),
+            (DataType.struct([("x", DataType.int32())]), True),
+            (DataType.map(DataType.string(), DataType.int64()), True),
+            # Non-nested types
+            (DataType.int64(), False),
+            (DataType.string(), False),
+            (DataType.float32(), False),
+        ],
+    )
+    def test_is_nested_type(self, datatype, expected_result):
+        """Test is_nested_type predicate."""
+        assert datatype.is_nested_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            # Numerical Arrow types
+            (DataType.int64(), True),
+            (DataType.int32(), True),
+            (DataType.float32(), True),
+            (DataType.float64(), True),
+            # Numerical NumPy types
+            (DataType.from_numpy(np.dtype("int32")), True),
+            (DataType.from_numpy(np.dtype("float64")), True),
+            # Numerical Python types
+            (DataType(int), True),
+            (DataType(float), True),
+            # Non-numerical types
+            (DataType.string(), False),
+            (DataType.binary(), False),
+            (DataType(str), False),
+        ],
+    )
+    def test_is_numerical_type(self, datatype, expected_result):
+        """Test is_numerical_type predicate."""
+        assert datatype.is_numerical_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            # String Arrow types
+            (DataType.string(), True),
+            (DataType.from_arrow(pa.large_string()), True),
+            # String NumPy types
+            (DataType.from_numpy(np.dtype("U10")), True),
+            # String Python types
+            (DataType(str), True),
+            # Non-string types
+            (DataType.int64(), False),
+            (DataType.binary(), False),
+        ],
+    )
+    def test_is_string_type(self, datatype, expected_result):
+        """Test is_string_type predicate."""
+        assert datatype.is_string_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            # Binary Arrow types
+            (DataType.binary(), True),
+            (DataType.from_arrow(pa.large_binary()), True),
+            (DataType.from_arrow(pa.binary(10)), True),  # fixed_size_binary
+            # Binary Python types
+            (DataType(bytes), True),
+            (DataType(bytearray), True),
+            # Non-binary types
+            (DataType.string(), False),
+            (DataType.int64(), False),
+        ],
+    )
+    def test_is_binary_type(self, datatype, expected_result):
+        """Test is_binary_type predicate."""
+        assert datatype.is_binary_type() == expected_result
+
+    @pytest.mark.parametrize(
+        "datatype,expected_result",
+        [
+            # Temporal Arrow types
+            (DataType.temporal("timestamp", unit="s"), True),
+            (DataType.temporal("date32"), True),
+            (DataType.temporal("time64", unit="us"), True),
+            (DataType.temporal("duration", unit="ms"), True),
+            # Temporal NumPy types
+            (DataType.from_numpy(np.dtype("datetime64[D]")), True),
+            (DataType.from_numpy(np.dtype("timedelta64[s]")), True),
+            # Non-temporal types
+            (DataType.int64(), False),
+            (DataType.string(), False),
+        ],
+    )
+    def test_is_temporal_type(self, datatype, expected_result):
+        """Test is_temporal_type predicate."""
+        assert datatype.is_temporal_type() == expected_result
+
+
+if __name__ == "__main__":
+    pytest.main(["-v", __file__])

python/ray/data/tests/unit/test_deduping_schema.py

@@ -0,0 +1,289 @@
+import logging
+from types import SimpleNamespace
+from typing import Optional
+
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data._internal.execution.streaming_executor_state import (
+    OpState,
+    _build_schemas_mismatch_warning,
+    dedupe_schemas_with_validation,
+)
+from ray.data._internal.pandas_block import PandasBlockSchema
+from ray.data.block import Schema, _is_empty_schema
+
+
+@pytest.mark.parametrize(
+    "incoming_schema",
+    [
+        pa.schema([pa.field("uuid", pa.string())]),  # NOTE: diff from old_schema
+        pa.schema([]),  # Empty Schema
+        PandasBlockSchema(names=["col1"], types=[int]),
+        PandasBlockSchema(names=[], types=[]),
+        None,  # Null Schema
+    ],
+)
+@pytest.mark.parametrize(
+    "old_schema",
+    [
+        pa.schema([pa.field("id", pa.int64())]),
+        pa.schema([]),  # Empty Schema
+        PandasBlockSchema(names=["col2"], types=[int]),
+        PandasBlockSchema(names=[], types=[]),
+        None,  # Null Schema
+    ],
+)
+def test_dedupe_schema_handle_empty(
+    old_schema: Optional["Schema"],
+    incoming_schema: Optional["Schema"],
+):
+
+    incoming_bundle = RefBundle([], owns_blocks=False, schema=incoming_schema)
+    out_bundle, diverged = dedupe_schemas_with_validation(
+        old_schema, incoming_bundle, enforce_schemas=False
+    )
+
+    if _is_empty_schema(old_schema):
+        # old_schema is invalid
+        assert not diverged, (old_schema, incoming_schema)
+        assert out_bundle.schema == incoming_schema, (old_schema, incoming_schema)
+    else:
+        # old_schema is valid
+        assert diverged, (old_schema, incoming_schema)
+        assert incoming_schema != old_schema, (old_schema, incoming_schema)
+        assert old_schema == out_bundle.schema, (old_schema, incoming_schema)
+
+
+@pytest.mark.parametrize("enforce_schemas", [False, True])
+@pytest.mark.parametrize(
+    "incoming_schema", [pa.schema([pa.field("uuid", pa.string())])]
+)
+@pytest.mark.parametrize("old_schema", [pa.schema([pa.field("id", pa.int64())])])
+def test_dedupe_schema_divergence(
+    enforce_schemas: bool,
+    old_schema: Optional["Schema"],
+    incoming_schema: Optional["Schema"],
+):
+
+    incoming_bundle = RefBundle([], owns_blocks=False, schema=incoming_schema)
+    out_bundle, diverged = dedupe_schemas_with_validation(
+        old_schema, incoming_bundle, enforce_schemas=enforce_schemas
+    )
+
+    assert diverged
+
+    if enforce_schemas:
+        assert out_bundle.schema == pa.schema(list(old_schema) + list(incoming_schema))
+    else:
+        assert out_bundle.schema == old_schema
+
+
+@pytest.mark.parametrize(
+    "incoming_schema",
+    [
+        pa.schema([]),  # Empty Schema
+        PandasBlockSchema(names=[], types=[]),
+        None,  # Null Schema
+    ],
+)
+def test_build_mismatch_warning_empty(incoming_schema: Optional["Schema"]):
+    old_schema = pa.schema(
+        [
+            pa.field("foo", pa.int32()),
+            pa.field("bar", pa.string()),
+            pa.field("baz", pa.bool_()),
+        ]
+    )
+
+    msg = _build_schemas_mismatch_warning(old_schema, incoming_schema)
+
+    assert (
+        """Operator produced a RefBundle with an empty/unknown schema. (3 total):
+    foo: int32
+    bar: string
+    baz: bool
+This may lead to unexpected behavior."""
+        == msg
+    )
+
+
+@pytest.mark.parametrize("truncation_length", [20, 4, 3, 2])
+def test_build_mismatch_warning_truncation(truncation_length: int):
+    old_schema = pa.schema(
+        [
+            pa.field("foo", pa.int32()),
+            pa.field("bar", pa.string()),
+            pa.field("baz", pa.bool_()),
+            pa.field("quux", pa.uint16()),
+            pa.field("corge", pa.int64()),
+            pa.field("grault", pa.int32()),
+        ]
+    )
+    incoming_schema = pa.schema(
+        [
+            pa.field("foo", pa.int64()),
+            pa.field("baz", pa.bool_()),
+            pa.field("qux", pa.float64()),
+            pa.field("quux", pa.uint32()),
+            pa.field("corge", pa.int32()),
+            pa.field("grault", pa.int64()),
+        ]
+    )
+
+    msg = _build_schemas_mismatch_warning(
+        old_schema, incoming_schema, truncation_length=truncation_length
+    )
+
+    if truncation_length >= 4:
+        assert (
+            """Operator produced a RefBundle with a different schema than the previous one.
+Fields exclusive to the incoming schema (1 total):
+    qux: double
+Fields exclusive to the old schema (1 total):
+    bar: string
+Fields that have different types across the old and the incoming schemas (4 total):
+    foo: int32 => int64
+    quux: uint16 => uint32
+    corge: int64 => int32
+    grault: int32 => int64
+This may lead to unexpected behavior."""
+            == msg
+        )
+    elif truncation_length == 3:
+        assert (
+            """Operator produced a RefBundle with a different schema than the previous one.
+Fields exclusive to the incoming schema (1 total):
+    qux: double
+Fields exclusive to the old schema (1 total):
+    bar: string
+Fields that have different types across the old and the incoming schemas (4 total):
+    foo: int32 => int64
+    quux: uint16 => uint32
+    corge: int64 => int32
+    ... and 1 more
+This may lead to unexpected behavior."""
+            == msg
+        )
+    else:
+        assert (
+            """Operator produced a RefBundle with a different schema than the previous one.
+Fields exclusive to the incoming schema (1 total):
+    qux: double
+Fields exclusive to the old schema (1 total):
+    bar: string
+Fields that have different types across the old and the incoming schemas (4 total):
+    foo: int32 => int64
+    quux: uint16 => uint32
+    ... and 2 more
+This may lead to unexpected behavior."""
+            == msg
+        )
+
+
+def test_build_mismatch_warning_disordered():
+    old_schema = pa.schema(
+        [
+            pa.field("foo", pa.int32()),
+            pa.field("bar", pa.string()),
+            pa.field("baz", pa.bool_()),
+            pa.field("qux", pa.float64()),
+        ]
+    )
+    incoming_schema = pa.schema(
+        [
+            pa.field("foo", pa.int32()),
+            pa.field("baz", pa.bool_()),
+            pa.field("bar", pa.string()),
+            pa.field("qux", pa.float64()),
+        ]
+    )
+
+    msg = _build_schemas_mismatch_warning(old_schema, incoming_schema)
+
+    assert (
+        """Operator produced a RefBundle with a different schema than the previous one.
+Some fields are ordered differently across the old and the incoming schemas.
+This may lead to unexpected behavior."""
+        == msg
+    )
+
+
+class _DummyOp:
+    def __init__(self, enforce_schemas=True):
+        self.name = "dummy"
+        self.data_context = SimpleNamespace(enforce_schemas=enforce_schemas)
+        self.input_dependencies = []
+        self.output_dependencies = []
+        self.metrics = SimpleNamespace(
+            num_alive_actors=0,
+            num_restarting_actors=0,
+            num_pending_actors=0,
+            num_external_inqueue_blocks=0,
+            num_external_inqueue_bytes=0,
+            num_external_outqueue_blocks=0,
+            num_external_outqueue_bytes=0,
+        )
+
+    def num_output_splits(self):
+        return 1
+
+    def get_actor_info(self):
+        return SimpleNamespace(
+            running=0,
+            restarting=0,
+            pending=0,
+            active=0,
+            idle=0,
+            pool_utilization=0,
+            tasks_in_flight=0,
+        )
+
+
+@pytest.mark.parametrize("enforce_schemas", [False, True])
+def test_add_output_emits_warning(enforce_schemas, caplog, propagate_logs):
+    op = _DummyOp(enforce_schemas=enforce_schemas)
+    state = OpState(op, [])
+
+    old_schema = pa.schema(
+        [
+            pa.field("foo", pa.int32()),
+            pa.field("bar", pa.bool_()),
+        ]
+    )
+    new_schema = pa.schema(
+        [
+            pa.field("foo", pa.int32()),
+            pa.field("baz", pa.uint32()),
+        ]
+    )
+
+    # First output seeds schema (no warning)
+    state.add_output(RefBundle([], owns_blocks=False, schema=old_schema))
+
+    assert not caplog.messages
+
+    # Second output diverges, should warn once
+    with caplog.at_level(logging.WARNING):
+        state.add_output(RefBundle([], owns_blocks=False, schema=new_schema))
+
+    assert not (enforce_schemas ^ len(caplog.messages))
+
+    if enforce_schemas:
+        msg = "\n".join(caplog.messages)
+        assert (
+            msg
+            == """Operator produced a RefBundle with a different schema than the previous one.
+Fields exclusive to the incoming schema (1 total):
+    baz: uint32
+Fields exclusive to the old schema (1 total):
+    bar: bool
+This may lead to unexpected behavior."""
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_distribution_tracker.py

@@ -0,0 +1,167 @@
+import numpy as np
+import pytest
+
+try:
+    import datasketches
+except ImportError:
+    datasketches = None
+
+from ray.data._internal.execution.interfaces.distribution_tracker import (
+    DistributionTracker,
+)
+
+
+def test_empty_tracker_has_zero_moments_and_no_extremes():
+    tracker = DistributionTracker()
+
+    assert tracker.num_samples == 0
+    assert tracker.mean == 0.0
+    assert tracker.variance == 0.0
+    assert tracker.min is None
+    assert tracker.max is None
+
+
+def test_moments_match_numpy_after_adding_samples():
+    tracker = DistributionTracker()
+    samples = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
+    for s in samples:
+        tracker.add_sample(s)
+
+    assert tracker.num_samples == len(samples)
+    assert pytest.approx(tracker.mean) == np.mean(samples)
+    assert pytest.approx(tracker.variance) == np.var(samples, ddof=1)
+    assert pytest.approx(tracker.stddev) == np.std(samples, ddof=1)
+
+
+def test_extremes_track_min_and_max():
+    tracker = DistributionTracker()
+    samples = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
+    for s in samples:
+        tracker.add_sample(s)
+
+    assert tracker.min == 2.0
+    assert tracker.max == 9.0
+
+
+def test_as_dict_contains_all_fields():
+    tracker = DistributionTracker()
+    tracker.add_sample(1.0)
+
+    d = tracker.as_dict()
+    assert set(d.keys()) == {
+        "num_samples",
+        "mean",
+        "variance",
+        "min",
+        "max",
+        "p25",
+        "p50",
+        "p75",
+        "p90",
+        "p95",
+        "p99",
+    }
+
+
+@pytest.mark.skipif(datasketches is None, reason="datasketches not installed")
+def test_percentiles_approximate_expected_quantiles():
+    tracker = DistributionTracker()
+    for i in range(1, 101):
+        tracker.add_sample(float(i))
+
+    assert tracker.p50 is not None and 45 <= tracker.p50 <= 55
+    assert tracker.p90 is not None and 85 <= tracker.p90 <= 95
+    assert tracker.p99 is not None and 95 <= tracker.p99 <= 100
+
+
+def _build(samples):
+    tracker = DistributionTracker()
+    for s in samples:
+        tracker.add_sample(s)
+    return tracker
+
+
+def test_merge_moments_match_numpy_on_concatenation():
+    a = _build([2.0, 4.0, 4.0, 4.0])
+    b = _build([5.0, 5.0, 7.0, 9.0])
+
+    a.merge(b)
+    combined = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]
+
+    assert a.num_samples == len(combined)
+    assert pytest.approx(a.mean) == np.mean(combined)
+    assert pytest.approx(a.variance) == np.var(combined, ddof=1)
+    assert a.min == min(combined)
+    assert a.max == max(combined)
+
+
+def test_merge_is_commutative():
+    samples_a = [2.0, 4.0, 4.0, 4.0]
+    samples_b = [5.0, 5.0, 7.0, 9.0]
+
+    ab = _build(samples_a)
+    ab.merge(_build(samples_b))
+    ba = _build(samples_b)
+    ba.merge(_build(samples_a))
+
+    assert ab.num_samples == ba.num_samples
+    assert pytest.approx(ab.mean) == ba.mean
+    assert pytest.approx(ab.variance) == ba.variance
+    assert ab.min == ba.min
+    assert ab.max == ba.max
+
+
+def test_merge_with_empty_other_is_noop():
+    tracker = _build([2.0, 4.0, 6.0])
+
+    tracker.merge(DistributionTracker())
+
+    assert tracker.num_samples == 3
+    assert pytest.approx(tracker.mean) == np.mean([2.0, 4.0, 6.0])
+    assert tracker.min == 2.0
+    assert tracker.max == 6.0
+
+
+def test_merge_self_is_noop():
+    tracker = _build([2.0, 4.0, 6.0])
+
+    tracker.merge(tracker)
+
+    assert tracker.num_samples == 3
+    assert pytest.approx(tracker.mean) == 4.0
+
+
+@pytest.mark.skipif(datasketches is None, reason="datasketches not installed")
+def test_cloudpickle_roundtrip_preserves_sketch():
+    # ``kll_doubles_sketch`` is C++-backed and not natively picklable —
+    # without DistributionTracker's serialize/deserialize hooks, any
+    # Ray Data path that cloudpickles a Dataset (it carries Timers,
+    # which carry DistributionTrackers) fails with
+    # ``TypeError: cannot pickle 'kll_doubles_sketch' object``.
+    import pickle
+
+    import cloudpickle
+
+    tracker = DistributionTracker()
+    for i in range(1, 101):
+        tracker.add_sample(float(i))
+
+    for dumps, loads in [
+        (pickle.dumps, pickle.loads),
+        (cloudpickle.dumps, cloudpickle.loads),
+    ]:
+        restored = loads(dumps(tracker))
+        # Welford moments are exact across the round-trip.
+        assert restored.num_samples == tracker.num_samples
+        assert restored.mean == tracker.mean
+        assert restored.min == tracker.min
+        assert restored.max == tracker.max
+        # The deserialized sketch must still answer quantile queries.
+        assert restored.p50 is not None
+        assert restored.p50 == tracker.p50
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_dynamic_work_queue.py

@@ -0,0 +1,331 @@
+import pytest
+
+from ray.data._internal.dynamic_work_queue import parallel_process_work_stealing
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_flat(num_workers):
+    """Flat (non-recursive) processing: every seed item produces exactly one
+    result with no dynamically-added work."""
+
+    def process(item, add_work, add_result):
+        add_result(item * 10)
+
+    results = list(
+        parallel_process_work_stealing(
+            [1, 2, 3, 4, 5],
+            process,
+            num_workers=num_workers,
+        )
+    )
+
+    assert sorted(results) == [10, 20, 30, 40, 50]
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_flat_ordered(num_workers):
+    """Flat processing with preserve_order sorts by order_key."""
+
+    def process(item, add_work, add_result):
+        add_result(item * 10)
+
+    results = list(
+        parallel_process_work_stealing(
+            [3, 1, 5, 2, 4],
+            process,
+            num_workers=num_workers,
+            preserve_order=True,
+            order_key=lambda x: x,
+        )
+    )
+
+    assert results == [10, 20, 30, 40, 50]
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_recursive(num_workers):
+    """Recursive work generation: processing one item spawns sub-items,
+    mimicking a tree traversal."""
+
+    tree = {
+        "a": ["a/1", "a/2"],
+        "b": ["b/1"],
+    }
+
+    def process(item, add_work, add_result):
+        if item in tree:
+            for child in tree[item]:
+                add_work(child)
+        else:
+            add_result(item)
+
+    results = list(
+        parallel_process_work_stealing(
+            ["a", "b"],
+            process,
+            num_workers=num_workers,
+        )
+    )
+
+    assert sorted(results) == ["a/1", "a/2", "b/1"]
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_recursive_ordered(num_workers):
+    """Recursive work generation with preserve_order sorts all results
+    globally by order_key."""
+
+    tree = {
+        "a": ["a/2", "a/1"],
+        "b": ["b/1"],
+    }
+
+    def process(item, add_work, add_result):
+        if item in tree:
+            for child in tree[item]:
+                add_work(child)
+        else:
+            add_result(item)
+
+    results = list(
+        parallel_process_work_stealing(
+            ["a", "b"],
+            process,
+            num_workers=num_workers,
+            preserve_order=True,
+            order_key=lambda x: x,
+        )
+    )
+
+    assert results == ["a/1", "a/2", "b/1"]
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_deep_chain(num_workers):
+    """A single seed item spawns a long chain of work (depth > breadth)."""
+
+    def process(item, add_work, add_result):
+        if item > 0:
+            add_work(item - 1)
+        add_result(item)
+
+    results = list(
+        parallel_process_work_stealing(
+            [5],
+            process,
+            num_workers=num_workers,
+        )
+    )
+
+    assert sorted(results) == [0, 1, 2, 3, 4, 5]
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_deep_chain_ordered(num_workers):
+    """Deep chain with preserve_order sorts results by order_key."""
+
+    def process(item, add_work, add_result):
+        if item > 0:
+            add_work(item - 1)
+        add_result(item)
+
+    results = list(
+        parallel_process_work_stealing(
+            [5],
+            process,
+            num_workers=num_workers,
+            preserve_order=True,
+            order_key=lambda x: x,
+        )
+    )
+
+    assert results == [0, 1, 2, 3, 4, 5]
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_error_propagation(num_workers):
+    """Exceptions raised in process_fn are propagated to the caller."""
+
+    def process(item, add_work, add_result):
+        if item == 3:
+            raise ValueError("boom")
+        add_result(item)
+
+    with pytest.raises(ValueError, match="boom"):
+        list(
+            parallel_process_work_stealing(
+                [1, 2, 3, 4],
+                process,
+                num_workers=num_workers,
+            )
+        )
+
+
+def test_parallel_process_work_stealing_empty_seeds():
+    """No seed items produces an empty generator."""
+    results = list(
+        parallel_process_work_stealing([], lambda item, aw, ar: None, num_workers=1)
+    )
+    assert results == []
+
+
+def test_parallel_process_work_stealing_no_results():
+    """Seed items that produce no results yield an empty sequence."""
+
+    def process(item, add_work, add_result):
+        pass  # intentionally produce nothing
+
+    results = list(parallel_process_work_stealing([1, 2, 3], process, num_workers=2))
+    assert results == []
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_preserve_order_sorts_globally(num_workers):
+    """With preserve_order=True, all results are sorted globally by
+    order_key regardless of which seed item produced them."""
+
+    def process(item, add_work, add_result):
+        for suffix in ["c", "a", "b"]:
+            add_result(f"{item}-{suffix}")
+
+    results = list(
+        parallel_process_work_stealing(
+            ["X", "Y"],
+            process,
+            num_workers=num_workers,
+            preserve_order=True,
+            order_key=lambda x: x,
+        )
+    )
+
+    assert results == ["X-a", "X-b", "X-c", "Y-a", "Y-b", "Y-c"]
+
+
+def test_parallel_process_work_stealing_preserve_order_requires_order_key():
+    """preserve_order=True without order_key raises ValueError."""
+    with pytest.raises(ValueError, match="order_key is required"):
+        list(
+            parallel_process_work_stealing(
+                [1],
+                lambda item, aw, ar: ar(item),
+                preserve_order=True,
+            )
+        )
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_early_stop(num_workers):
+    """Generator can be stopped early (via break) without hanging."""
+
+    def process(item, add_work, add_result):
+        add_result(item)
+        # Spawn more work to ensure the queue isn't empty at break time.
+        if item < 100:
+            add_work(item + 1)
+
+    gen = parallel_process_work_stealing(
+        [0],
+        process,
+        num_workers=num_workers,
+        preserve_order=False,
+    )
+
+    collected = []
+    for result in gen:
+        collected.append(result)
+        if len(collected) >= 5:
+            break
+
+    assert len(collected) == 5
+
+
+def test_parallel_process_work_stealing_invalid_num_workers():
+    """num_workers < 1 raises ValueError."""
+    with pytest.raises(ValueError, match="num_workers must be at least 1"):
+        list(
+            parallel_process_work_stealing(
+                [1], lambda item, aw, ar: ar(item), num_workers=0
+            )
+        )
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_early_stop_no_thread_leak(num_workers):
+    """Threads spawned by the generator must not leak after early termination."""
+    import threading
+    import time
+
+    baseline = threading.active_count()
+
+    def process(item, add_work, add_result):
+        add_result(item)
+        if item < 1000:
+            add_work(item + 1)
+
+    gen = parallel_process_work_stealing(
+        [0],
+        process,
+        num_workers=num_workers,
+        preserve_order=False,
+    )
+
+    collected = []
+    for result in gen:
+        collected.append(result)
+        if len(collected) >= 3:
+            break
+
+    assert len(collected) == 3
+
+    time.sleep(3)
+    assert threading.active_count() <= baseline + 1
+
+
+def test_parallel_process_work_stealing_error_clears_exception():
+    """_WorkerError.exception should be cleared after re-raising to avoid
+    traceback reference cycles that delay GC of worker-frame locals."""
+    import gc
+    import weakref
+
+    class BigObject:
+        pass
+
+    refs: "list[weakref.ref]" = []
+
+    def process(item, add_work, add_result):
+        obj = BigObject()
+        refs.append(weakref.ref(obj))
+        raise ValueError("test error")
+
+    with pytest.raises(ValueError, match="test error"):
+        list(parallel_process_work_stealing([1], process, num_workers=1))
+
+    gc.collect()
+    assert refs[0]() is None, "BigObject leaked via traceback reference cycle"
+
+
+@pytest.mark.parametrize("num_workers", [1, 4])
+def test_parallel_process_work_stealing_error_in_dynamically_added_work(num_workers):
+    """Errors raised in dynamically-added (non-seed) work items must still
+    propagate to the caller."""
+
+    def process(item, add_work, add_result):
+        if item == "child":
+            raise RuntimeError("child error")
+        add_work("child")
+        add_result(item)
+
+    with pytest.raises(RuntimeError, match="child error"):
+        list(
+            parallel_process_work_stealing(
+                ["root"],
+                process,
+                num_workers=num_workers,
+            )
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_expression_evaluator.py

@@ -0,0 +1,429 @@
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from pkg_resources import parse_version
+
+from ray.data._internal.logical.operators import CSE_TEMP_COLUMN_PREFIX
+from ray.data._internal.planner.plan_expression.expression_evaluator import (
+    ExpressionEvaluator,
+    eval_projection,
+)
+from ray.data.expressions import col, star
+from ray.data.tests.conftest import get_pyarrow_version
+
+
+@pytest.fixture(scope="module")
+def sample_data(tmpdir_factory):
+    """Fixture to create and yield sample Parquet data, and clean up afterwards."""
+    # Sample data for testing purposes
+    data = {
+        "age": [
+            25,
+            32,
+            45,
+            29,
+            40,
+            np.nan,
+        ],  # List of ages, including a None value for testing
+        "city": [
+            "New York",
+            "San Francisco",
+            "Los Angeles",
+            "Los Angeles",
+            "San Francisco",
+            "San Jose",
+        ],
+        "is_student": [False, True, False, False, True, None],  # Including a None value
+    }
+
+    # Define the schema explicitly
+    schema = pa.schema(
+        [("age", pa.float64()), ("city", pa.string()), ("is_student", pa.bool_())]
+    )
+
+    # Create a PyArrow table from the sample data
+    table = pa.table(data, schema=schema)
+
+    # Use tmpdir_factory to create a temporary directory
+    temp_dir = tmpdir_factory.mktemp("data")
+    parquet_file = temp_dir.join("sample_data.parquet")
+
+    # Write the table to a Parquet file in the temporary directory
+    pq.write_table(table, str(parquet_file))
+
+    # Yield the path to the Parquet file for testing
+    yield str(parquet_file), schema
+
+
+expressions_and_expected_data = [
+    # Parameterized test cases with expressions and their expected results
+    # Comparison Ops
+    (
+        "40 > age",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    (
+        "40 >= age",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "30 < age",
+        [
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "age >= 30",
+        [
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "age == 40",
+        [{"age": 40, "city": "San Francisco", "is_student": True}],
+    ),
+    (
+        "is_student != True",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    (
+        "age < 0",
+        [],
+    ),
+    (
+        "is_student == True",
+        [
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "is_student == False",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    # Op 'in'
+    (
+        "city in ['Los Angeles', 'New York']",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    (
+        "city in ['Los Angeles']",
+        [
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    (
+        "city in ['New York', 'San Francisco']",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "age in []",
+        [],
+    ),
+    (
+        "age in [25, 32, 45, 29, 40]",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "age in [25, 32, None]",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    # Op 'not in'
+    (
+        "is_student not in [None]",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "is_student not in [True, None]",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    # Logical Ops 'and'
+    (
+        "age > 30 and is_student == True",
+        [
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "city == 'Los Angeles' and age < 40",
+        [
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    (
+        "age < 40 and city in ['New York', 'Los Angeles']",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    # Logical Ops 'or'
+    (
+        "age < 30 or is_student == True",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "city == 'New York' or city == 'San Francisco'",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "age < 30 or age > 40",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    # Logical Ops combination 'and' and 'or'
+    (
+        "(age < 30 or age > 40) and is_student != True",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+        ],
+    ),
+    # Op 'is_null'
+    (
+        "is_null(is_student)",
+        [
+            {"age": None, "city": "San Jose", "is_student": None},
+        ],
+    ),
+    (
+        "age < 40 or is_null(is_student)",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+            {"age": None, "city": "San Jose", "is_student": None},
+        ],
+    ),
+    # Op 'is_nan'
+    (
+        "is_nan(age)",
+        [
+            {"age": None, "city": "San Jose", "is_student": None},
+        ],
+    ),
+    (
+        "city in ['San Jose', 'Los Angeles'] and is_nan(age)",
+        [
+            {"age": None, "city": "San Jose", "is_student": None},
+        ],
+    ),
+    # Op 'is_valid'
+    (
+        "is_valid(is_student)",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+    (
+        "is_valid(is_student) and is_valid(age)",
+        [
+            {"age": 25, "city": "New York", "is_student": False},
+            {"age": 32, "city": "San Francisco", "is_student": True},
+            {"age": 45, "city": "Los Angeles", "is_student": False},
+            {"age": 29, "city": "Los Angeles", "is_student": False},
+            {"age": 40, "city": "San Francisco", "is_student": True},
+        ],
+    ),
+]
+
+
+@pytest.mark.skipif(
+    get_pyarrow_version() < parse_version("20.0.0"),
+    reason="test_filter requires PyArrow >= 20.0.0",
+)
+@pytest.mark.parametrize("expression, expected_data", expressions_and_expected_data)
+def test_filter(sample_data, expression, expected_data):
+    """Test the filter functionality of the ExpressionEvaluator."""
+
+    # Instantiate the ExpressionEvaluator with valid column names
+    sample_data_path, _ = sample_data
+    filters = ExpressionEvaluator.get_filters(expression=expression)
+
+    # Read the table from the Parquet file with the applied filters
+    filtered_table = pq.read_table(sample_data_path, filters=filters)
+
+    # Convert the filtered table back to a list of dictionaries for comparison
+    result = filtered_table.to_pandas().to_dict(orient="records")
+
+    def convert_nan_to_none(data):
+        return [
+            {k: (None if pd.isna(v) else v) for k, v in record.items()}
+            for record in data
+        ]
+
+    # Convert NaN to None for comparison
+    result_converted = convert_nan_to_none(result)
+
+    assert result_converted == expected_data
+
+
+def test_filter_equal_negative_number():
+    df = pd.DataFrame.from_dict(
+        {"A": [-1, -1, 1, 2, -1, 3, 4, 5], "B": [-1, -1, 1, 2, -1, 3, 4, 5]}
+    )
+    expression = ExpressionEvaluator.get_filters(expression="A == -1")
+    result = pa.table(df).filter(expression)
+    result_df = result.to_pandas().to_dict(orient="records")
+    expected = df[df["A"] == -1].to_dict(orient="records")
+    assert result_df == expected
+
+
+@pytest.mark.skipif(
+    get_pyarrow_version() < parse_version("20.0.0"),
+    reason="test_filter requires PyArrow >= 20.0.0",
+)
+def test_filter_bad_expression(sample_data):
+    with pytest.raises(ValueError, match="Invalid syntax in the expression"):
+        ExpressionEvaluator.get_filters(expression="bad filter")
+
+    filters = ExpressionEvaluator.get_filters(expression="hi > 3")
+
+    sample_data_path, _ = sample_data
+    with pytest.raises(pa.ArrowInvalid):
+        pq.read_table(sample_data_path, filters=filters)
+
+
+def test_eval_projection_star_rename_missing_source_raises():
+    """A rename targeting a column not present in the block must raise rather
+    than be silently dropped during star expansion."""
+    block = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})
+    projection = [star(), col("nonexistent")._rename("x")]
+
+    with pytest.raises(KeyError):
+        eval_projection(projection, block)
+
+
+def test_eval_projection_with_common_sub_exprs_arrow():
+    block = pa.table({"a": [1, 2, 3]})
+    common = (col("a") + 1).alias(f"{CSE_TEMP_COLUMN_PREFIX}test_0")
+    projection = [
+        (
+            col(f"{CSE_TEMP_COLUMN_PREFIX}test_0")
+            + col(f"{CSE_TEMP_COLUMN_PREFIX}test_0")
+        ).alias("y")
+    ]
+
+    out = eval_projection(
+        projection,
+        block,
+        common_sub_exprs=[common],
+    )
+
+    assert out.column_names == ["y"]
+    assert out["y"].to_pylist() == [4, 6, 8]
+
+
+def test_eval_projection_cse_temp_columns_do_not_leak_with_star():
+    block = pa.table({"a": [1, 2, 3]})
+    common = (col("a") + 1).alias(f"{CSE_TEMP_COLUMN_PREFIX}test_0")
+
+    out = eval_projection(
+        [star(), col(f"{CSE_TEMP_COLUMN_PREFIX}test_0").alias("y")],
+        block,
+        common_sub_exprs=[common],
+    )
+
+    assert out.column_names == ["a", "y"]
+    assert out["y"].to_pylist() == [2, 3, 4]
+
+
+def test_eval_projection_preserves_reserved_prefix_without_cse():
+    block = pa.table({f"{CSE_TEMP_COLUMN_PREFIX}user": [1, 2]})
+    out = eval_projection([star()], block)
+    assert out.column_names == [f"{CSE_TEMP_COLUMN_PREFIX}user"]
+
+
+def test_eval_projection_with_common_sub_exprs_pandas():
+    block = pd.DataFrame({"a": [1, 2, 3]})
+    common = (col("a") + 1).alias(f"{CSE_TEMP_COLUMN_PREFIX}test_0")
+    projection = [
+        (
+            col(f"{CSE_TEMP_COLUMN_PREFIX}test_0")
+            + col(f"{CSE_TEMP_COLUMN_PREFIX}test_0")
+        ).alias("y")
+    ]
+
+    out = eval_projection(
+        projection,
+        block,
+        common_sub_exprs=[common],
+    )
+
+    assert out.columns.tolist() == ["y"]
+    assert out["y"].tolist() == [4, 6, 8]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_expression_visitor.py

@@ -0,0 +1,126 @@
+from collections import Counter
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
+
+from ray.data._internal.planner.plan_expression.expression_visitors import (
+    _ColumnReferenceCollector,
+    _StructuralFingerprintOccurrenceCollector,
+    _StructuralFingerprintVisitor,
+)
+from ray.data.datatype import DataType
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    PyArrowComputeUDFExpr,
+    UDFExpr,
+    col,
+    lit,
+    monotonically_increasing_id,
+    random,
+    udf,
+    uuid,
+)
+
+
+@udf(return_dtype=DataType.int64())
+def add_one(x: pa.Array) -> pa.Array:
+    return pc.add(x, 1)
+
+
+def _fingerprint(expr):
+    return _StructuralFingerprintVisitor().visit(expr)
+
+
+def test_structural_fingerprint_matches_structural_equality():
+    expr = (add_one(col("a")) + lit(1)).alias("result")
+    equivalent_expr = (add_one(col("a")) + lit(1)).alias("result")
+    different_child_expr = (add_one(col("b")) + lit(1)).alias("result")
+    different_alias_expr = (add_one(col("a")) + lit(1)).alias("other")
+
+    assert expr.structurally_equals(equivalent_expr)
+    assert _fingerprint(expr) == _fingerprint(equivalent_expr)
+    assert _fingerprint(expr) != _fingerprint(different_child_expr)
+    assert _fingerprint(expr) != _fingerprint(different_alias_expr)
+
+
+def test_structural_fingerprint_handles_pyarrow_compute_udfs():
+    expr = col("a").abs()
+    equivalent_expr = col("a").abs()
+    different_child_expr = col("b").abs()
+
+    assert isinstance(expr, PyArrowComputeUDFExpr)
+    assert expr.structurally_equals(equivalent_expr)
+    assert _fingerprint(expr) == _fingerprint(equivalent_expr)
+    assert _fingerprint(expr) != _fingerprint(different_child_expr)
+
+
+def test_occurrence_collector_records_bottom_up_keys_and_depths():
+    expr = (add_one(col("a")) + add_one(col("a"))).alias("result")
+
+    collector = _StructuralFingerprintOccurrenceCollector()
+    root_key = collector.visit(expr)
+    occurrences = collector.get_occurrences()
+
+    assert root_key == _fingerprint(expr)
+    assert [type(occurrence.expr) for occurrence in occurrences] == [
+        ColumnExpr,
+        UDFExpr,
+        ColumnExpr,
+        UDFExpr,
+        BinaryExpr,
+        AliasExpr,
+    ]
+    assert [occurrence.depth for occurrence in occurrences] == [3, 2, 3, 2, 1, 0]
+    assert all(
+        occurrence.key == _fingerprint(occurrence.expr) for occurrence in occurrences
+    )
+
+    udf_occurrences = [
+        occurrence for occurrence in occurrences if isinstance(occurrence.expr, UDFExpr)
+    ]
+    assert len(udf_occurrences) == 2
+    assert udf_occurrences[0].key == udf_occurrences[1].key
+    assert udf_occurrences[0].expr.structurally_equals(udf_occurrences[1].expr)
+
+
+@pytest.mark.parametrize(
+    "expr,expected",
+    [
+        # idempotent leaves and composites
+        (col("a"), True),
+        (lit(1), True),
+        (col("a") + lit(1), True),
+        (add_one(col("a")), True),
+        (add_one(col("a")).alias("y"), True),
+        # non-idempotent leaves
+        (random(), False),
+        (random(seed=42), False),
+        (uuid(), False),
+        (monotonically_increasing_id(), False),
+        # non-idempotency propagates through composites
+        (random() + lit(1), False),
+        ((col("a") + uuid()).alias("x"), False),
+        (add_one(monotonically_increasing_id()), False),
+    ],
+)
+def test_is_idempotent(expr, expected):
+    assert expr.is_idempotent() is expected
+
+
+def test_column_reference_collector_counts_multiplicity():
+    collector = _ColumnReferenceCollector()
+    collector.visit(col("x") + col("x") + col("y"))
+
+    # get_counts() counts repeats within a single expression...
+    assert collector.get_counts() == Counter({"x": 2, "y": 1})
+    # ...while get_column_refs() stays ordered and de-duplicated.
+    assert collector.get_column_refs() == ["x", "y"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_fifo_bundle_queue.py

@@ -0,0 +1,192 @@
+from typing import Any
+from uuid import uuid4
+
+import pandas as pd
+import pytest
+
+import ray
+from ray.data._internal.execution.bundle_queue import FIFOBundleQueue
+from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
+from ray.data.block import BlockAccessor
+
+
+def _create_bundle(data: Any) -> RefBundle:
+    """Create a RefBundle with a single row with the given data using artificial refs."""
+    block = pd.DataFrame({"data": [data]})
+    # Create artificial object ref without calling ray.put()
+    block_ref = ray.ObjectRef(uuid4().hex[:28].encode())
+    metadata = BlockAccessor.for_block(block).get_metadata()
+    schema = BlockAccessor.for_block(block).schema()
+    return RefBundle(
+        [BlockEntry(block_ref, metadata)], owns_blocks=False, schema=schema
+    )
+
+
+def test_fifo_queue_add_and_length():
+    """Test adding bundles and checking length."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+
+    queue.add(bundle1)
+    assert len(queue) == 1
+
+    queue.add(bundle2)
+    assert len(queue) == 2
+
+
+def test_fifo_queue_get_next_fifo_order():
+    """Test that bundles are returned in FIFO order."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+    bundle3 = _create_bundle("data111")
+
+    queue.add(bundle1)
+    queue.add(bundle2)
+    queue.add(bundle3)
+
+    assert queue.get_next() is bundle1
+    assert queue.get_next() is bundle2
+    assert queue.get_next() is bundle3
+
+
+def test_fifo_queue_init_with_bundles():
+    """Test initializing queue with a list of bundles."""
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+
+    queue = FIFOBundleQueue(bundles=[bundle1, bundle2])
+
+    assert len(queue) == 2
+    assert queue.get_next() is bundle1
+    assert queue.get_next() is bundle2
+
+
+def test_fifo_queue_peek_next():
+    """Test peeking at the next bundle without removing it."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+
+    queue.add(bundle1)
+    queue.add(bundle2)
+
+    # Peek should return bundle1 without removing
+    assert queue.peek_next() is bundle1
+    assert len(queue) == 2
+
+    # Peek again should return the same bundle
+    assert queue.peek_next() is bundle1
+
+
+def test_fifo_queue_peek_next_empty():
+    """Test peeking when queue is empty."""
+    queue = FIFOBundleQueue()
+    assert queue.peek_next() is None
+
+
+def test_fifo_queue_has_next():
+    """Test has_next correctly reflects queue state."""
+    queue = FIFOBundleQueue()
+    assert not queue.has_next()
+
+    bundle1 = _create_bundle("data1")
+    queue.add(bundle1)
+    assert queue.has_next()
+
+    queue.get_next()
+    assert not queue.has_next()
+
+
+def test_fifo_queue_get_next_empty_raises():
+    """Test that get_next raises when queue is empty."""
+    queue = FIFOBundleQueue()
+
+    with pytest.raises(ValueError, match="Popping from empty"):
+        queue.get_next()
+
+
+def test_fifo_queue_clear():
+    """Test clearing the queue resets everything."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+
+    queue.add(bundle1)
+    queue.add(bundle2)
+
+    queue.clear()
+
+    assert len(queue) == 0
+    assert queue.estimate_size_bytes() == 0
+    assert queue.num_blocks() == 0
+    assert not queue.has_next()
+
+
+def test_fifo_queue_metrics():
+    """Test that metrics are tracked correctly."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+
+    queue.add(bundle1)
+    assert queue.estimate_size_bytes() == bundle1.size_bytes()
+    assert queue.num_blocks() == 1
+
+    queue.add(bundle2)
+    assert queue.estimate_size_bytes() == bundle1.size_bytes() + bundle2.size_bytes()
+    assert queue.num_blocks() == 2
+
+    queue.get_next()
+    assert queue.estimate_size_bytes() == bundle2.size_bytes()
+    assert queue.num_blocks() == 1
+
+
+def test_fifo_queue_iter():
+    """Test iterating over the queue."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+    bundle3 = _create_bundle("data111")
+
+    queue.add(bundle1)
+    queue.add(bundle2)
+    queue.add(bundle3)
+
+    # Iterate without consuming
+    bundles = list(queue)
+    assert bundles == [bundle1, bundle2, bundle3]
+    assert len(queue) == 3  # Queue unchanged
+
+
+def test_fifo_queue_to_list():
+    """Test converting queue to list."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+    bundle2 = _create_bundle("data11")
+
+    queue.add(bundle1)
+    queue.add(bundle2)
+
+    bundles = queue.to_list()
+    assert bundles == [bundle1, bundle2]
+    assert len(queue) == 2  # Queue unchanged
+
+
+def test_fifo_queue_finalize_is_noop():
+    """Test that finalize does nothing (it's a no-op for FIFO queue)."""
+    queue = FIFOBundleQueue()
+    bundle1 = _create_bundle("data1")
+
+    queue.add(bundle1)
+    queue.finalize()  # Should not raise or change anything
+
+    assert len(queue) == 1
+    assert queue.get_next() is bundle1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_filename_provider.py

@@ -0,0 +1,67 @@
+import pytest
+
+from ray.data.datasource.filename_provider import FilenameProvider
+
+
+@pytest.fixture(params=["csv", None])
+def filename_provider(request):
+    yield FilenameProvider(dataset_uuid="", file_format=request.param)
+
+
+def test_default_filename_for_task_includes_task_index(filename_provider):
+    filename_0 = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+    filename_1 = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=1
+    )
+    assert filename_0 != filename_1
+    assert "000000" in filename_0
+    assert "000001" in filename_1
+
+
+def test_default_get_filename_for_task_is_deterministic(filename_provider):
+    """Test the new get_filename_for_task() method is deterministic."""
+
+    first_filename = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+    second_filename = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+
+    assert first_filename == second_filename
+
+
+def test_default_row_filenames_derived_from_task_are_unique(filename_provider):
+    """Row filenames derived from task filename with block/row index are unique."""
+    task_filename = filename_provider.get_filename_for_task(
+        write_uuid="spam", task_index=0
+    )
+    filenames = []
+    for block_index in range(2):
+        for row_index in range(4):
+            if "." in task_filename:
+                base, ext = task_filename.rsplit(".", 1)
+                filenames.append(f"{base}_{block_index:06}_{row_index:06}.{ext}")
+            else:
+                filenames.append(f"{task_filename}_{block_index:06}_{row_index:06}")
+    assert len(set(filenames)) == len(filenames)
+
+
+def test_default_get_filename_for_task_is_unique(filename_provider):
+    """Test the new get_filename_for_task() method generates unique filenames."""
+    filenames = [
+        filename_provider.get_filename_for_task(
+            write_uuid="spam",
+            task_index=task_index,
+        )
+        for task_index in range(4)
+    ]
+    assert len(set(filenames)) == len(filenames)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_jax_util.py

@@ -0,0 +1,322 @@
+import sys
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+
+from ray.data.util.jax_util import jax_sync_generator
+
+pytest.importorskip("jax")
+
+
+def test_jax_sync_generator_empty_batch(ray_start_regular_shared):
+    def empty_batch_iterable():
+        yield {"a": np.array([1, 2, 3])}
+        yield {}  # Empty dict batch
+        yield {"a": np.array([4, 5, 6])}
+
+    gen = jax_sync_generator(empty_batch_iterable(), drop_last=True, batch_size=3)
+    results = list(gen)
+
+    assert len(results) == 2
+    assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
+    assert np.array_equal(results[1]["a"], np.array([4, 5, 6]))
+
+
+def test_jax_sync_generator_empty_column(ray_start_regular_shared):
+    def empty_column_iterable():
+        yield {"a": np.array([1, 2, 3])}
+        yield {"a": np.array([])}  # Dict with empty column
+        yield {"a": np.array([4, 5, 6])}
+
+    gen = jax_sync_generator(empty_column_iterable(), drop_last=True, batch_size=3)
+    results = list(gen)
+
+    assert len(results) == 2
+    assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
+    assert np.array_equal(results[1]["a"], np.array([4, 5, 6]))
+
+
+def test_jax_sync_generator_no_sync(ray_start_regular_shared):
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+        yield {"a": np.array([4, 5, 6])}
+
+    # Should work fine in single process even with sync=False
+    gen = jax_sync_generator(
+        batches(), drop_last=True, batch_size=3, synchronize_batches=False
+    )
+    results = list(gen)
+    assert len(results) == 2
+
+
+def test_jax_sync_generator_padding(ray_start_regular_shared):
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+        yield {"a": np.array([4, 5])}
+
+    # Should pad the second batch to size 3 with value -1
+    gen = jax_sync_generator(
+        batches(),
+        drop_last=False,
+        batch_size=3,
+        paddings=-1,
+        synchronize_batches=False,
+    )
+    results = list(gen)
+
+    assert len(results) == 2
+    assert len(results[0]["a"]) == 3
+    assert len(results[1]["a"]) == 3
+    assert np.array_equal(results[1]["a"], np.array([4, 5, -1]))
+
+
+def test_jax_sync_generator_drop_last(ray_start_regular_shared):
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+        yield {"a": np.array([4, 5])}
+
+    # Should drop the second batch because it's not size 3 and padding=None
+    # Note: in single host, it doesn't drop unless we use _iter_batches(drop_last=True)
+    # But jax_sync_generator with drop_last=True will raise error if sizes don't match min.
+    # Actually, jax_sync_generator logic for single host just passes through if not sync.
+
+    # Let's test single host with divisibility check failure
+    gen = jax_sync_generator(
+        batches(),
+        drop_last=False,
+        batch_size=3,
+        paddings=None,
+        synchronize_batches=False,
+    )
+    results = list(gen)
+    assert len(results) == 2  # Both yielded because 2 is divisible by 1 local device
+
+    # Let's force num_local_devices = 4 for testing error handling
+    from unittest.mock import patch
+
+    with patch("jax.local_device_count", return_value=4):
+        gen = jax_sync_generator(
+            batches(),
+            drop_last=True,
+            batch_size=3,
+            paddings=None,
+            synchronize_batches=False,
+        )
+        with pytest.raises(
+            ValueError,
+            match="evenly divisible by the number of local JAX devices",
+        ):
+            list(gen)
+
+
+def test_jax_sync_generator_multi_host_uneven_batches_with_padding(
+    ray_start_regular_shared,
+):
+
+    from unittest.mock import patch
+
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+        # Host 0 ends here, Host 1 has more
+
+    # Mock jax environment: 2 hosts, 1 device per host
+    with patch("jax.process_count", return_value=2), patch(
+        "jax.local_device_count", return_value=1
+    ), patch(
+        "ray.data.util.jax_util._convert_batch",
+        side_effect=lambda x, sharding, **kwargs: x,
+    ):
+
+        def mock_process_allgather(arr):
+            # Simulate Host 1 having more data
+            # local_infos for host 0: [1, 3, 0, 0, ...] (from batches())
+            # arr is a JAX array because jax_sync_generator does jnp.array(local_infos)
+            # Convert to numpy for easy manipulation
+            h0_infos = np.array(arr)
+
+            h1_infos = np.zeros_like(h0_infos)
+            h1_infos[0] = 1  # batch 1 exists
+            h1_infos[1] = 3  # batch 1 size
+            h1_infos[2] = 1  # batch 2 exists
+            h1_infos[3] = 3  # batch 2 size
+
+            import jax.numpy as jnp
+
+            return jnp.array([h0_infos, h1_infos])
+
+        with patch(
+            "jax.experimental.multihost_utils.process_allgather",
+            side_effect=mock_process_allgather,
+        ):
+            # Host 0 uses jax_sync_generator
+            gen = jax_sync_generator(
+                batches(),
+                drop_last=False,
+                batch_size=3,
+                paddings=-1,
+                synchronize_batches=True,
+            )
+            # Should yield 2 batches: one real, one dummy
+            results = list(gen)
+            assert len(results) == 2
+            assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
+            assert np.array_equal(results[1]["a"], np.array([-1, -1, -1]))
+
+
+def test_jax_sync_generator_multi_host_uneven_batches_drop_last(
+    ray_start_regular_shared,
+):
+
+    from unittest.mock import patch
+
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+        yield {"a": np.array([4, 5, 6])}
+
+    with patch("jax.process_count", return_value=2), patch(
+        "jax.local_device_count", return_value=1
+    ), patch(
+        "ray.data.util.jax_util._convert_batch",
+        side_effect=lambda x, sharding, **kwargs: x,
+    ):
+
+        def mock_process_allgather(arr):
+            # Host 0 has 2 batches, Host 1 has 1 batch
+            h0_infos = np.array(arr)
+            h1_infos = np.zeros_like(h0_infos)
+            h1_infos[0] = 1
+            h1_infos[1] = 3
+            import jax.numpy as jnp
+
+            return jnp.array([h0_infos, h1_infos])
+
+        with patch(
+            "jax.experimental.multihost_utils.process_allgather",
+            side_effect=mock_process_allgather,
+        ):
+            gen = jax_sync_generator(
+                batches(),
+                drop_last=True,
+                batch_size=3,
+                synchronize_batches=True,
+            )
+            # Should yield only 1 batch and stop
+            results = list(gen)
+            assert len(results) == 1
+            assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
+
+
+def test_jax_sync_generator_multi_host_uneven_batch_sizes_fail(
+    ray_start_regular_shared,
+):
+
+    from unittest.mock import patch
+
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+
+    with patch("jax.process_count", return_value=2), patch(
+        "jax.local_device_count", return_value=1
+    ), patch(
+        "ray.data.util.jax_util._convert_batch",
+        side_effect=lambda x, sharding, **kwargs: x,
+    ):
+
+        def mock_process_allgather(arr):
+            # Host 0 batch size 3, Host 1 batch size 2
+            h0_infos = np.array(arr)
+            h1_infos = h0_infos.copy()
+            h1_infos[1] = 2
+            import jax.numpy as jnp
+
+            return jnp.array([h0_infos, h1_infos])
+
+        with patch(
+            "jax.experimental.multihost_utils.process_allgather",
+            side_effect=mock_process_allgather,
+        ):
+            gen = jax_sync_generator(
+                batches(),
+                drop_last=False,
+                batch_size=3,
+                synchronize_batches=True,
+            )
+            with pytest.raises(
+                ValueError, match="Uneven batch sizes detected across JAX workers"
+            ):
+                list(gen)
+
+
+def test_jax_sync_generator_multi_host_uneven_num_batches_fail(
+    ray_start_regular_shared,
+):
+
+    from unittest.mock import patch
+
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+        yield {"a": np.array([4, 5, 6])}
+
+    with patch("jax.process_count", return_value=2), patch(
+        "jax.local_device_count", return_value=1
+    ), patch(
+        "ray.data.util.jax_util._convert_batch",
+        side_effect=lambda x, sharding, **kwargs: x,
+    ):
+
+        def mock_process_allgather(arr):
+            # Host 0 has 2 batches, Host 1 has 1 batch, no padding
+            h0_infos = np.array(arr)
+            h1_infos = np.zeros_like(h0_infos)
+            h1_infos[0] = 1
+            h1_infos[1] = 3
+            import jax.numpy as jnp
+
+            return jnp.array([h0_infos, h1_infos])
+
+        with patch(
+            "jax.experimental.multihost_utils.process_allgather",
+            side_effect=mock_process_allgather,
+        ):
+            gen = jax_sync_generator(
+                batches(),
+                drop_last=False,
+                batch_size=3,
+                synchronize_batches=True,
+            )
+            with pytest.raises(
+                ValueError,
+                match="Uneven number of batches detected across JAX workers",
+            ):
+                list(gen)
+
+
+def test_jax_sync_generator_with_dtypes(ray_start_regular_shared):
+    def batches():
+        yield {"a": np.array([1, 2, 3])}
+
+    import jax.numpy as jnp
+
+    dtypes = {"a": jnp.float16}
+
+    # Mock _convert_batch to capture dtypes
+    mock_convert = MagicMock(side_effect=lambda x, sharding, dtypes=None: x)
+
+    with patch("ray.data.util.jax_util._convert_batch", mock_convert):
+        gen = jax_sync_generator(
+            batches(),
+            drop_last=False,
+            batch_size=3,
+            dtypes=dtypes,
+            synchronize_batches=False,
+        )
+        list(gen)
+
+    # Verify that dtypes was passed to _convert_batch
+    mock_convert.assert_called_once()
+    assert mock_convert.call_args[1]["dtypes"] == dtypes
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_limit_operator.py

@@ -0,0 +1,51 @@
+import pytest
+
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_operator import _per_block_limit_fn
+
+
+@pytest.mark.parametrize(
+    "blocks_data,per_block_limit,expected_output",
+    [
+        # Test case 1: Single block, limit less than block size
+        ([[1, 2, 3, 4, 5]], 3, [[1, 2, 3]]),
+        # Test case 2: Single block, limit equal to block size
+        ([[1, 2, 3]], 3, [[1, 2, 3]]),
+        # Test case 3: Single block, limit greater than block size
+        ([[1, 2]], 5, [[1, 2]]),
+        # Test case 4: Multiple blocks, limit spans across blocks
+        ([[1, 2], [3, 4], [5, 6]], 3, [[1, 2], [3]]),
+        # Test case 5: Multiple blocks, limit exactly at block boundary
+        ([[1, 2], [3, 4]], 2, [[1, 2]]),
+        # Test case 6: Empty blocks
+        ([], 5, []),
+        # Test case 7: Zero limit
+        ([[1, 2, 3]], 0, []),
+    ],
+)
+def test_per_block_limit_fn(blocks_data, per_block_limit, expected_output):
+    """Test the _per_block_limit_fn function with various inputs."""
+    import pandas as pd
+
+    # Convert test data to pandas blocks
+    blocks = [pd.DataFrame({"value": data}) for data in blocks_data]
+
+    # Create a mock TaskContext
+    ctx = TaskContext(op_name="test", task_idx=0, target_max_block_size_override=None)
+
+    # Call the function
+    result_blocks = list(_per_block_limit_fn(blocks, ctx, per_block_limit))
+
+    # Convert result back to lists for comparison
+    result_data = []
+    for block in result_blocks:
+        block_data = block["value"].tolist()
+        result_data.append(block_data)
+
+    assert result_data == expected_output
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_logical_plan.py

@@ -0,0 +1,74 @@
+from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan
+from ray.data.context import DataContext
+
+
+class DummyLogicalOperator(LogicalOperator):
+    def __init__(self, input_dependencies, name=None):
+        object.__setattr__(self, "_input_dependencies", input_dependencies)
+        if name is not None:
+            object.__setattr__(self, "_name", name)
+
+
+def test_sources_singleton():
+    ctx = DataContext.get_current()
+
+    source = DummyLogicalOperator([], name="source")
+    assert LogicalPlan(source, ctx).sources() == [source]
+
+
+def test_sources_chain():
+    ctx = DataContext.get_current()
+
+    source = DummyLogicalOperator([], name="source")
+    sink = DummyLogicalOperator([source], name="sink")
+    assert LogicalPlan(sink, ctx).sources() == [source]
+
+
+def test_sources_multiple_sources():
+    ctx = DataContext.get_current()
+
+    source1 = DummyLogicalOperator([], name="source1")
+    source2 = DummyLogicalOperator([], name="source2")
+    sink = DummyLogicalOperator([source1, source2], name="sink")
+    assert LogicalPlan(sink, ctx).sources() == [source1, source2]
+
+
+def test_logical_operator_defaults_name_to_class_name():
+    op = DummyLogicalOperator([])
+    assert op.name == "DummyLogicalOperator"
+    assert op.dag_str == "DummyLogicalOperator[DummyLogicalOperator]"
+
+
+def test_logical_operator_does_not_track_output_dependencies():
+    source = DummyLogicalOperator([], name="source")
+    sink = DummyLogicalOperator([source], name="sink")
+
+    # Logical operators should not maintain reverse output-dependency state.
+    assert not hasattr(source, "_output_dependencies")
+    assert not hasattr(sink, "_output_dependencies")
+
+    transformed = sink._apply_transform(lambda op: op)
+    assert transformed is sink
+    assert not hasattr(source, "_output_dependencies")
+    assert not hasattr(sink, "_output_dependencies")
+
+
+def test_logical_operator_transform_supports_custom_subclasses():
+    source = DummyLogicalOperator([], name="source")
+    replacement = DummyLogicalOperator([], name="replacement")
+    sink = DummyLogicalOperator([source], name="sink")
+
+    transformed = sink._apply_transform(lambda op: replacement if op is source else op)
+
+    assert transformed is not sink
+    assert transformed.name == "sink"
+    assert transformed.input_dependencies == [replacement]
+    assert sink.input_dependencies == [source]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_object_extension.py

@@ -0,0 +1,82 @@
+import types
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.object_extensions.arrow import (
+    ArrowPythonObjectArray,
+    ArrowPythonObjectType,
+)
+from ray.data._internal.object_extensions.pandas import PythonObjectArray
+
+
+def test_object_array_validation():
+    # Test unknown input type raises TypeError.
+    with pytest.raises(TypeError):
+        PythonObjectArray(object())
+
+    PythonObjectArray(np.array([object(), object()]))
+    PythonObjectArray([object(), object()])
+
+
+def test_arrow_scalar_object_array_roundtrip():
+    arr = np.array(
+        ["test", 20, False, {"some": "value"}, None, np.zeros((10, 10))], dtype=object
+    )
+    ata = ArrowPythonObjectArray.from_objects(arr)
+    assert isinstance(ata.type, ArrowPythonObjectType)
+    assert isinstance(ata, ArrowPythonObjectArray)
+    assert len(ata) == len(arr)
+    out = ata.to_numpy()
+    np.testing.assert_array_equal(out[:-1], arr[:-1])
+    assert np.all(out[-1] == arr[-1])
+
+
+def test_arrow_python_object_array_slice():
+    arr = np.array(["test", 20, "test2", 40, "test3", 60], dtype=object)
+    ata = ArrowPythonObjectArray.from_objects(arr)
+    assert list(ata[1:3].to_pandas()) == [20, "test2"]
+    assert ata[2:4].to_pylist() == ["test2", 40]
+
+
+def test_arrow_pandas_roundtrip():
+    obj = types.SimpleNamespace(a=1, b="test")
+    t1 = pa.table({"a": ArrowPythonObjectArray.from_objects([obj, obj]), "b": [0, 1]})
+    t2 = pa.Table.from_pandas(t1.to_pandas())
+    assert t1.equals(t2)
+
+
+def test_pandas_python_object_isna():
+    arr = np.array([1, np.nan, 3, 4, 5, np.nan, 7, 8, 9], dtype=object)
+    ta = PythonObjectArray(arr)
+    np.testing.assert_array_equal(ta.isna(), pd.isna(arr))
+
+
+def test_pandas_python_object_take():
+    arr = np.array([1, 2, 3, 4, 5], dtype=object)
+    ta = PythonObjectArray(arr)
+    indices = [1, 2, 3]
+    np.testing.assert_array_equal(ta.take(indices).to_numpy(), arr[indices])
+    indices = [1, 2, -1]
+    np.testing.assert_array_equal(
+        ta.take(indices, allow_fill=True, fill_value=100).to_numpy(),
+        np.array([2, 3, 100]),
+    )
+
+
+def test_pandas_python_object_concat():
+    arr1 = np.array([1, 2, 3, 4, 5], dtype=object)
+    arr2 = np.array([6, 7, 8, 9, 10], dtype=object)
+    ta1 = PythonObjectArray(arr1)
+    ta2 = PythonObjectArray(arr2)
+    concat_arr = PythonObjectArray._concat_same_type([ta1, ta2])
+    assert len(concat_arr) == arr1.shape[0] + arr2.shape[0]
+    np.testing.assert_array_equal(concat_arr.to_numpy(), np.concatenate([arr1, arr2]))
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))

python/ray/data/tests/unit/test_parquet_predicate_split.py

@@ -0,0 +1,348 @@
+"""Unit tests for Parquet predicate splitting optimization.
+
+This module tests the _split_predicate_by_columns function which optimizes
+predicate pushdown for partitioned Parquet datasets by splitting predicates
+into data-column, partition-column, and residual parts.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Set
+
+import pytest
+
+from ray.data._internal.datasource.parquet_datasource import (
+    _split_predicate_by_columns,
+)
+from ray.data.expressions import Expr, col
+
+
+@dataclass
+class PredicateSplitTestCase:
+    """Test case for predicate splitting."""
+
+    predicate: Expr
+    partition_cols: Set[str]
+    expected_data_predicate: Optional[Expr]
+    expected_partition_predicate: Optional[Expr]
+    description: str
+    expected_residual_predicate: Optional[Expr] = None
+
+
+# fmt: off
+TEST_CASES = [
+    # ====================================================================
+    # Pure data predicates - should push down everything
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=col("data1") > 5,
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data1") > 5,
+        expected_partition_predicate=None,
+        description="Simple data predicate",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & (col("data2") == "x"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) & (col("data2") == "x"),
+        expected_partition_predicate=None,
+        description="AND of data predicates",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) | (col("data2") == "x"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) | (col("data2") == "x"),
+        expected_partition_predicate=None,
+        description="OR of data predicates",
+    ),
+    PredicateSplitTestCase(
+        predicate=~(col("data1") > 5),
+        partition_cols={"partition_col"},
+        expected_data_predicate=~(col("data1") > 5),
+        expected_partition_predicate=None,
+        description="NOT of data predicate",
+    ),
+    PredicateSplitTestCase(
+        predicate=col("data1").is_null(),
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data1").is_null(),
+        expected_partition_predicate=None,
+        description="IS_NULL of data column",
+    ),
+    PredicateSplitTestCase(
+        predicate=col("data1").is_not_null(),
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data1").is_not_null(),
+        expected_partition_predicate=None,
+        description="IS_NOT_NULL of data column",
+    ),
+    PredicateSplitTestCase(
+        predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("data3") == "test"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("data3") == "test"),
+        expected_partition_predicate=None,
+        description="Complex nested data predicates",
+    ),
+    # ====================================================================
+    # Pure partition predicates - should enable pruning
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=col("partition_col") == "US",
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=col("partition_col") == "US",
+        description="Simple partition predicate",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
+        partition_cols={"partition1", "partition2"},
+        expected_data_predicate=None,
+        expected_partition_predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
+        description="AND of partition predicates",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("partition1") == "US") | (col("partition2") == "2020"),
+        partition_cols={"partition1", "partition2"},
+        expected_data_predicate=None,
+        expected_partition_predicate=(col("partition1") == "US") | (col("partition2") == "2020"),
+        description="OR of partition predicates",
+    ),
+    # ====================================================================
+    # Mixed predicates with AND - should split both parts
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & (col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data1") > 5,
+        expected_partition_predicate=col("partition_col") == "US",
+        description="Simple AND with data and partition",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("partition_col") == "US") & (col("data1") > 5),
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data1") > 5,
+        expected_partition_predicate=col("partition_col") == "US",
+        description="Simple AND with partition and data (reversed order)",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & (col("data2") < 10) & (col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
+        expected_partition_predicate=col("partition_col") == "US",
+        description="Multiple data predicates AND partition",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & (col("partition1") == "US") & (col("data2") < 10) & (col("partition2") == "2020"),
+        partition_cols={"partition1", "partition2"},
+        expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
+        expected_partition_predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
+        description="Interleaved data and partition predicates",
+    ),
+    PredicateSplitTestCase(
+        predicate=((col("data1") > 5) & (col("data2") < 10)) & (col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
+        expected_partition_predicate=col("partition_col") == "US",
+        description="Nested AND of data predicates with partition",
+    ),
+    # ====================================================================
+    # Mixed predicates with OR - CANNOT split safely
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) | (col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=None,
+        expected_residual_predicate=(col("data1") > 5) | (col("partition_col") == "US"),
+        description="OR with data and partition (unsafe to split)",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("partition_col") == "US") | (col("data1") > 5),
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=None,
+        expected_residual_predicate=(col("partition_col") == "US") | (col("data1") > 5),
+        description="OR with partition and data (unsafe to split)",
+    ),
+    PredicateSplitTestCase(
+        predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=None,
+        expected_residual_predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("partition_col") == "US"),
+        description="OR with complex data predicate and partition",
+    ),
+    # ====================================================================
+    # Mixed predicates with NOT - CANNOT split safely
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=~(col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=~(col("partition_col") == "US"),
+        description="NOT of partition predicate",
+    ),
+    PredicateSplitTestCase(
+        predicate=~((col("data1") > 5) & (col("partition_col") == "US")),
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=None,
+        expected_residual_predicate=~((col("data1") > 5) & (col("partition_col") == "US")),
+        description="NOT of mixed AND (becomes OR via De Morgan, unsafe)",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & ~(col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data1") > 5,
+        expected_partition_predicate=~(col("partition_col") == "US"),
+        description="AND with NOT of partition predicate (can extract both parts)",
+    ),
+    # ====================================================================
+    # Complex nested scenarios
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=((col("data1") > 5) & (col("data2") < 10)) & ((col("data3") == "test") & (col("partition_col") == "US")),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) & (col("data2") < 10) & (col("data3") == "test"),
+        expected_partition_predicate=col("partition_col") == "US",
+        description="Deeply nested ANDs with mixed columns",
+    ),
+    PredicateSplitTestCase(
+        predicate=((col("data1") > 5) | (col("data2") < 10)) & (col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) | (col("data2") < 10),
+        expected_partition_predicate=col("partition_col") == "US",
+        description="AND of complex data predicate (with OR) and partition",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & ((col("data2") < 10) & (col("partition_col") == "US")),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
+        expected_partition_predicate=col("partition_col") == "US",
+        description="Left-nested data with right-nested mixed",
+    ),
+    PredicateSplitTestCase(
+        predicate=((col("data1") > 5) & (col("partition1") == "US")) & ((col("data2") < 10) & (col("partition2") == "2020")),
+        partition_cols={"partition1", "partition2"},
+        expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
+        expected_partition_predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
+        description="Both sides have mixed predicates",
+    ),
+    PredicateSplitTestCase(
+        predicate=((col("data1") > 5) | (col("partition_col") == "US")) & (col("data2") < 10),
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data2") < 10,
+        expected_partition_predicate=None,
+        expected_residual_predicate=(col("data1") > 5) | (col("partition_col") == "US"),
+        description="AND with left side having OR with partition (residual carries the unsplittable OR)",
+    ),
+    # ====================================================================
+    # Edge cases
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & (col("data1") < 10),
+        partition_cols={"partition_col"},
+        expected_data_predicate=(col("data1") > 5) & (col("data1") < 10),
+        expected_partition_predicate=None,
+        description="Same column referenced multiple times (data)",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("partition_col") == "US") & (col("partition_col") != "UK"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=(col("partition_col") == "US") & (col("partition_col") != "UK"),
+        description="Same partition column referenced multiple times",
+    ),
+    PredicateSplitTestCase(
+        predicate=col("data1").is_null() & (col("partition_col") == "US"),
+        partition_cols={"partition_col"},
+        expected_data_predicate=col("data1").is_null(),
+        expected_partition_predicate=col("partition_col") == "US",
+        description="IS_NULL data predicate with partition",
+    ),
+    PredicateSplitTestCase(
+        predicate=col("partition_col").is_null(),
+        partition_cols={"partition_col"},
+        expected_data_predicate=None,
+        expected_partition_predicate=col("partition_col").is_null(),
+        description="IS_NULL on partition column",
+    ),
+    # ====================================================================
+    # No partition columns in dataset
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=col("data1") > 5,
+        partition_cols=set(),
+        expected_data_predicate=col("data1") > 5,
+        expected_partition_predicate=None,
+        description="No partition columns in dataset",
+    ),
+    PredicateSplitTestCase(
+        predicate=(col("data1") > 5) & (col("data2") < 10),
+        partition_cols=set(),
+        expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
+        expected_partition_predicate=None,
+        description="Complex predicate with no partition columns",
+    ),
+    # ====================================================================
+    # All columns are partition columns
+    # ====================================================================
+    PredicateSplitTestCase(
+        predicate=col("partition1") > 5,
+        partition_cols={"partition1", "partition2"},
+        expected_data_predicate=None,
+        expected_partition_predicate=col("partition1") > 5,
+        description="All columns are partition columns",
+    ),
+]
+
+def _assert_predicate_matches(
+    actual: Optional[Expr], expected: Optional[Expr], pred_type: str, description: str
+):
+    """Helper to assert predicate matches expected value."""
+    if expected is None:
+        assert actual is None, (
+            f"{description}: Expected no {pred_type} predicate (None), got {actual}"
+        )
+    else:
+        assert actual is not None, f"{description}: Expected {pred_type} predicate, got None"
+
+
+@pytest.mark.parametrize("test_case", TEST_CASES, ids=lambda tc: tc.description)
+def test_split_predicate_by_columns(test_case: PredicateSplitTestCase):
+    """Test predicate splitting for various scenarios.
+
+    This test covers:
+    - Pure data predicates (should extract data part only)
+    - Pure partition predicates (should extract partition part only)
+    - Mixed predicates with AND (should split both parts)
+    - Mixed predicates with OR (kept as residual)
+    - Mixed predicates with NOT (varies by case)
+    - Complex nested scenarios with residual carry-over
+    - Edge cases
+    """
+    result = _split_predicate_by_columns(test_case.predicate, test_case.partition_cols)
+
+    _assert_predicate_matches(
+        result.data_predicate,
+        test_case.expected_data_predicate,
+        "data",
+        test_case.description,
+    )
+    _assert_predicate_matches(
+        result.partition_predicate,
+        test_case.expected_partition_predicate,
+        "partition",
+        test_case.description,
+    )
+    _assert_predicate_matches(
+        result.residual_predicate,
+        test_case.expected_residual_predicate,
+        "residual",
+        test_case.description,
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_path_util.py

@@ -0,0 +1,163 @@
+from unittest import mock
+
+import pyarrow
+import pyarrow.fs
+import pytest
+from fsspec.implementations.http import HTTPFileSystem
+from pyarrow.fs import FSSpecHandler, PyFileSystem
+
+from ray.data._internal.util import RetryingPyFileSystem
+from ray.data.datasource.path_util import (
+    _has_file_extension,
+    _is_filesystem_compatible_with_scheme,
+    _is_local_windows_path,
+    _resolve_paths_and_filesystem,
+)
+
+
+@pytest.mark.parametrize(
+    "path, extensions, has_extension",
+    [
+        ("foo.csv", ["csv"], True),
+        ("foo.csv", ["json", "csv"], True),
+        ("foo.csv", ["json", "jsonl"], False),
+        ("foo.csv", [".csv"], True),
+        ("foo.parquet.crc", ["parquet"], False),
+        ("foo.parquet.crc", ["crc"], True),
+        ("s3://bucket/foo.parquet?versionId=abc123", ["parquet"], True),
+        ("bucket/foo.parquet?versionId=abc123", ["parquet"], True),
+        ("s3://bucket/foo.parquet?versionId=abc123", ["csv"], False),
+        ("s3://bucket/data#v2/file.parquet", ["parquet"], True),
+        ("C:\\data\\test.parquet", ["parquet"], True),
+        ("C:\\data\\parquet", ["parquet"], False),
+        ("foo.csv", None, True),
+    ],
+)
+def test_has_file_extension(path, extensions, has_extension):
+    assert _has_file_extension(path, extensions) == has_extension
+
+
+@pytest.mark.parametrize(
+    "filesystem", [None, PyFileSystem(FSSpecHandler(HTTPFileSystem()))]
+)
+def test_resolve_http_paths(filesystem):
+    resolved_paths, resolve_filesystem = _resolve_paths_and_filesystem(
+        "https://google.com", filesystem
+    )
+    # `_resolve_paths_and_filesystem` shouldn't remove the protocol/scheme from the
+    # path for HTTP paths.
+    assert resolved_paths == ["https://google.com"]
+    assert isinstance(resolve_filesystem, pyarrow.fs.PyFileSystem)
+    assert isinstance(resolve_filesystem.handler, pyarrow.fs.FSSpecHandler)
+    assert isinstance(resolve_filesystem.handler.fs, HTTPFileSystem)
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "c:/some/where",
+        "c:\\some\\where",
+        "c:\\some\\where/mixed",
+    ],
+)
+def test_windows_path(path):
+    with mock.patch("sys.platform", "win32"):
+        assert _is_local_windows_path(path)
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "some/file",
+        "some/file;semicolon",
+        "some/file?questionmark",
+        "some/file#hash",
+        "some/file;all?of the#above",
+    ],
+)
+def test_weird_local_paths(path):
+    resolved_paths, _ = _resolve_paths_and_filesystem(path)
+    assert resolved_paths[0] == path
+
+
+class TestIsFilesystemCompatibleWithScheme:
+    """Tests for _is_filesystem_compatible_with_scheme with real filesystem implementations."""
+
+    def test_native_local_filesystem(self):
+        """Native PyArrow LocalFileSystem should be compatible with empty scheme."""
+        fs = pyarrow.fs.LocalFileSystem()
+        assert _is_filesystem_compatible_with_scheme(fs, "") is True
+        assert _is_filesystem_compatible_with_scheme(fs, "s3") is False
+
+    def test_native_s3_filesystem(self):
+        """Native PyArrow S3FileSystem should be compatible with s3 scheme."""
+        fs = pyarrow.fs.S3FileSystem(anonymous=True, region="us-east-1")
+        assert _is_filesystem_compatible_with_scheme(fs, "s3") is True
+        assert _is_filesystem_compatible_with_scheme(fs, "") is False
+
+    def test_fsspec_s3_filesystem_with_s3_scheme(self):
+        """fsspec S3FileSystem wrapped in PyFileSystem should be compatible with s3 scheme."""
+        s3fs = pytest.importorskip("s3fs")
+        wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "s3") is True
+
+    def test_fsspec_s3_filesystem_with_bare_paths(self):
+        """fsspec S3FileSystem wrapped in PyFileSystem should be compatible with bare paths."""
+        s3fs = pytest.importorskip("s3fs")
+        wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "") is True
+
+    def test_fsspec_s3_filesystem_with_s3a_scheme(self):
+        """Real s3fs has protocol=('s3', 's3a'), so it should match s3a scheme."""
+        s3fs = pytest.importorskip("s3fs")
+        wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "s3a") is True
+
+    def test_fsspec_gcs_filesystem_with_gs_scheme(self):
+        """fsspec GCS filesystem should be compatible with gs scheme."""
+        gcsfs = pytest.importorskip("gcsfs")
+        wrapped_fs = PyFileSystem(FSSpecHandler(gcsfs.GCSFileSystem(token="anon")))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "gs") is True
+
+    def test_fsspec_gcs_filesystem_with_gcs_scheme(self):
+        """Real gcsfs has protocol=('gcs', 'gs'), so it should match gcs scheme."""
+        gcsfs = pytest.importorskip("gcsfs")
+        wrapped_fs = PyFileSystem(FSSpecHandler(gcsfs.GCSFileSystem(token="anon")))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "gcs") is True
+
+    def test_fsspec_http_filesystem(self):
+        """fsspec HTTPFileSystem wrapped in PyFileSystem should be compatible with http."""
+        wrapped_fs = PyFileSystem(FSSpecHandler(HTTPFileSystem()))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "http") is True
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "https") is True
+
+    def test_fsspec_s3_not_compatible_with_gs(self):
+        """fsspec S3FileSystem should NOT be compatible with gs scheme."""
+        s3fs = pytest.importorskip("s3fs")
+        wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "gs") is False
+
+    def test_fsspec_s3_not_compatible_with_http(self):
+        """fsspec S3FileSystem should NOT be compatible with http/https schemes."""
+        s3fs = pytest.importorskip("s3fs")
+        wrapped_fs = PyFileSystem(FSSpecHandler(s3fs.S3FileSystem(anon=True)))
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "http") is False
+        assert _is_filesystem_compatible_with_scheme(wrapped_fs, "https") is False
+
+    def test_retrying_wrapper_around_native_s3(self):
+        """RetryingPyFileSystem wrapping a native S3FileSystem should be compatible with s3."""
+        s3_fs = pyarrow.fs.S3FileSystem(anonymous=True)
+        retrying_fs = RetryingPyFileSystem.wrap(s3_fs, retryable_errors=["AWS Error"])
+        assert _is_filesystem_compatible_with_scheme(retrying_fs, "s3") is True
+        assert _is_filesystem_compatible_with_scheme(retrying_fs, "gs") is False
+
+    def test_unknown_scheme_trusts_filesystem(self):
+        """Unknown schemes should always return True (trust user's filesystem)."""
+        fs = pyarrow.fs.LocalFileSystem()
+        assert _is_filesystem_compatible_with_scheme(fs, "custom") is True
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_reordering_bundle_queue.py

@@ -0,0 +1,318 @@
+from typing import Any
+from uuid import uuid4
+
+import pandas as pd
+import pytest
+
+import ray
+from ray.data._internal.execution.bundle_queue import ReorderingBundleQueue
+from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
+from ray.data.block import BlockAccessor
+
+
+def _create_bundle(data: Any) -> RefBundle:
+    """Create a RefBundle with a single row with the given data using artificial refs."""
+    block = pd.DataFrame({"data": [data]})
+    # Create artificial object ref without calling ray.put()
+    block_ref = ray.ObjectRef(uuid4().hex[:28].encode())
+    metadata = BlockAccessor.for_block(block).get_metadata()
+    schema = BlockAccessor.for_block(block).schema()
+    return RefBundle(
+        [BlockEntry(block_ref, metadata)], owns_blocks=False, schema=schema
+    )
+
+
+def test_ordered_queue_add_and_get_in_order():
+    """Test adding and getting bundles in sequential order."""
+    queue = ReorderingBundleQueue()
+    bundle0 = _create_bundle("data1")
+    bundle1 = _create_bundle("data11")
+
+    queue.add(bundle0, key=0)
+    queue.add(bundle1, key=1)
+
+    assert len(queue) == 2
+    assert queue.has_next()
+
+    # Can only get from key 0 until it's finalized
+    assert queue.get_next() is bundle0
+    # Nothing else to dequeue for 1
+    assert not queue.has_next()
+    queue.finalize(key=0)
+
+    # Now can get from key 1
+    assert queue.get_next() is bundle1
+    assert not queue.has_next()
+    queue.finalize(key=1)
+
+    assert len(queue) == 0
+    assert not queue.has_next()
+
+
+def test_ordered_queue_add_out_of_order():
+    """Test that bundles added out of order are returned in key order."""
+    queue = ReorderingBundleQueue()
+    bundle0 = _create_bundle("data1")
+    bundle1 = _create_bundle("data11")
+    bundle2 = _create_bundle("data111")
+
+    # Add in reverse order
+    queue.add(bundle2, key=2)
+    queue.add(bundle0, key=0)
+    queue.add(bundle1, key=1)
+
+    assert len(queue) == 3
+
+    # Should still get in key order
+    assert queue.get_next() is bundle0
+    queue.finalize(key=0)
+
+    assert queue.get_next() is bundle1
+    queue.finalize(key=1)
+
+    assert queue.get_next() is bundle2
+    queue.finalize(key=2)
+
+
+def test_ordered_queue_multiple_bundles_per_key():
+    """Test adding multiple bundles for the same key."""
+    queue = ReorderingBundleQueue()
+    bundle1a = _create_bundle("data1a")
+    bundle1b = _create_bundle("data1b")
+    bundle2 = _create_bundle("data2")
+
+    queue.add(bundle1a, key=0)
+    queue.add(bundle1b, key=0)
+    queue.add(bundle2, key=1)
+
+    assert len(queue) == 3
+
+    # Get both bundles from key 0
+    assert queue.get_next() is bundle1a
+    assert queue.get_next() is bundle1b
+    queue.finalize(key=0)
+
+    # Now get from key 1
+    assert queue.get_next() is bundle2
+    queue.finalize(key=1)
+
+
+def test_ordered_queue_finalize_before_all_consumed():
+    """Test finalizing a key before all its bundles are consumed."""
+    queue = ReorderingBundleQueue()
+    bundle1a = _create_bundle("data1a")
+    bundle1b = _create_bundle("data1b")
+    bundle2 = _create_bundle("data2")
+
+    queue.add(bundle1a, key=0)
+    queue.add(bundle1b, key=0)
+    queue.add(bundle2, key=1)
+
+    # Finalize key 0 before consuming all bundles
+    queue.finalize(key=0)
+
+    # Should still be able to get all bundles from key 0
+    assert queue.get_next() is bundle1a
+    assert queue.get_next() is bundle1b
+
+    # After consuming all, automatically moves to key 1
+    assert queue.get_next() is bundle2
+
+
+def test_ordered_queue_has_next_blocked_by_earlier_key():
+    """Test that has_next returns False when current key has no bundles."""
+    queue = ReorderingBundleQueue()
+    bundle1 = _create_bundle("data11")
+
+    # Add bundle for key 1, but nothing for key 0
+    queue.add(bundle1, key=1)
+
+    # has_next should return False because key 0 (current) has no bundles
+    assert not queue.has_next()
+    assert len(queue) == 1
+
+    # Finalize key 0 (even though it's empty) to move to key 1
+    queue.finalize(key=0)
+
+    # Now has_next should return True
+    assert queue.has_next()
+    assert queue.get_next() is bundle1
+
+
+def test_ordered_queue_peek_next():
+    """Test peeking at the next bundle without removing it."""
+    queue = ReorderingBundleQueue()
+    bundle0 = _create_bundle("data1")
+    bundle1 = _create_bundle("data11")
+
+    queue.add(bundle0, key=0)
+    queue.add(bundle1, key=1)
+
+    # Peek should return bundle0 without removing
+    assert queue.peek_next() is bundle0
+    assert len(queue) == 2
+
+    # Peek again should return the same bundle
+    assert queue.peek_next() is bundle0
+
+
+def test_ordered_queue_peek_next_empty():
+    """Test peeking when current key has no bundles."""
+    queue = ReorderingBundleQueue()
+    bundle1 = _create_bundle("data11")
+
+    queue.add(bundle1, key=1)
+
+    # Current key 0 is empty
+    assert queue.peek_next() is None
+
+
+def test_ordered_queue_out_of_order():
+    """Tests that ordered queue works correctly under following conditions"""
+
+    queue = ReorderingBundleQueue()
+    bundle0 = _create_bundle("data0")
+    bundle1 = _create_bundle("data1")
+
+    # First, add bundle for key=1
+    queue.add(bundle1, key=1)
+    queue.finalize(key=1)
+    # No bundles can be retrieved yet as we're missing bundles for key=0
+    assert not queue.has_next()
+
+    # Next, add bundle for key=0
+    queue.add(bundle0, key=0)
+    assert queue.get_next() is bundle0
+    queue.finalize(key=0)
+
+    # Now able to retrieve bundle for key=1
+    assert queue.get_next() is bundle1
+
+    # `has_next` should return bundle0 without removing
+    assert not queue.has_next()
+    assert len(queue) == 0
+
+
+@pytest.mark.parametrize("target_op", ["get", "peek"])
+def test_ordered_queue_getting_stuck(target_op):
+    bundle2 = _create_bundle("data2")
+
+    queue = ReorderingBundleQueue()
+
+    # Task 2 produces output and completes (finalizes key)
+    queue.add(bundle2, key=2)
+    queue.finalize(key=2)
+
+    # _current_key = 0
+    # _completed_keys = {2}
+
+    # Task 1 completes with NO output (empty result)
+    queue.finalize(key=1)
+
+    # _current_key = 0
+    # _completed_keys = {1, 2}
+
+    # Task 0 completes with NO output (empty result)
+    #
+    # Previously this will trigger moving to the next key, since
+    # _current_key == 0 AND _inner[0] is empty
+    #   → _move_to_next_key()
+    #       → _current_key = 1
+    queue.finalize(key=0)
+
+    # Current state:
+    #   _current_key = 1
+    #   _completed_keys = {0, 1, 2}
+    #   _inner = {0: [], 1: [], 2: [bundle_2]}
+
+    # Previously
+    #   - `has_next` would return False, (_inner[_current_key] is empty)
+    #   - `get_next` will never be invoked (b/c `has_next` returns false)
+    #   - `finalize(key=1)` has already been invoked, no pointer advancement will happen
+    #
+    # This results in the last bundle getting stuck in the queue
+    if target_op == "get":
+        assert queue.get_next() is bundle2
+    elif target_op == "peek":
+        assert queue.peek_next() is bundle2
+        assert queue.get_next() is bundle2
+    else:
+        pytest.fail(f"unsupported {target_op}")
+
+    assert len(queue) == 0
+
+
+def test_ordered_queue_get_next_empty_raises():
+    """Test that get_next raises when current key is empty."""
+    queue = ReorderingBundleQueue()
+
+    with pytest.raises(ValueError, match="Cannot pop from empty queue"):
+        queue.get_next()
+
+
+def test_ordered_queue_clear():
+    """Test clearing the queue resets everything."""
+    queue = ReorderingBundleQueue()
+    bundle0 = _create_bundle("data1")
+    bundle1 = _create_bundle("data11")
+
+    queue.add(bundle0, key=0)
+    queue.add(bundle1, key=1)
+    queue.finalize(key=0)
+    queue.get_next()  # Consume bundle0, moves to key 1
+
+    queue.clear()
+
+    assert len(queue) == 0
+    assert queue.estimate_size_bytes() == 0
+    assert queue.num_blocks() == 0
+    assert not queue.has_next()
+
+
+def test_ordered_queue_metrics():
+    """Test that metrics are tracked correctly."""
+    queue = ReorderingBundleQueue()
+    bundle0 = _create_bundle("data1")
+    bundle1 = _create_bundle("data11")
+
+    queue.add(bundle0, key=0)
+    assert queue.estimate_size_bytes() == bundle0.size_bytes()
+    assert queue.num_blocks() == 1
+
+    queue.add(bundle1, key=1)
+    assert queue.estimate_size_bytes() == bundle0.size_bytes() + bundle1.size_bytes()
+    assert queue.num_blocks() == 2
+
+    queue.get_next()
+    queue.finalize(key=0)
+    assert queue.estimate_size_bytes() == bundle1.size_bytes()
+    assert queue.num_blocks() == 1
+
+
+def test_ordered_queue_finalize_out_of_order():
+    """Test that keys can be finalized out of order."""
+    queue = ReorderingBundleQueue()
+    bundle0 = _create_bundle("data1")
+    bundle1 = _create_bundle("data11")
+    bundle2 = _create_bundle("data111")
+
+    queue.add(bundle0, key=0)
+    queue.add(bundle1, key=1)
+    queue.add(bundle2, key=2)
+
+    # Finalize key 2 first, then 1, then 0
+    queue.finalize(key=2)
+    queue.finalize(key=1)
+
+    # Should still need to consume key 0 first
+    assert queue.get_next() is bundle0
+    queue.finalize(key=0)
+
+    assert queue.get_next() is bundle1
+    assert queue.get_next() is bundle2
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_resource_manager.py

@@ -0,0 +1,241 @@
+import pytest
+
+import ray
+from ray.data._internal.execution.block_ref_counter import BlockRefCounter
+from ray.data._internal.execution.interfaces import (
+    BlockEntry,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionOptions,
+    ExecutionResources,
+)
+from ray.data._internal.execution.operators.union_operator import UnionOperator
+from ray.data._internal.execution.resource_manager import (
+    ResourceManager,
+)
+from ray.data._internal.execution.streaming_executor_state import (
+    build_streaming_topology,
+)
+from ray.data.block import BlockMetadata
+from ray.data.context import DataContext
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.conftest import noop_counter
+
+
+def test_physical_operator_tracks_output_dependencies():
+    input_op = PhysicalOperator("input", [], DataContext.get_current())
+    downstream_op = PhysicalOperator(
+        "downstream", [input_op], DataContext.get_current()
+    )
+
+    assert input_op.output_dependencies == [downstream_op]
+
+
+def test_physical_apply_transform_rewires_all_input_output_dependencies():
+    ctx = DataContext.get_current()
+    left_input = PhysicalOperator("left_input", [], ctx)
+    right_input = PhysicalOperator("right_input", [], ctx)
+    root = PhysicalOperator("root", [left_input, right_input], ctx)
+    left_replacement = PhysicalOperator("left_replacement", [], ctx)
+
+    transformed_root = root._apply_transform(
+        lambda op: left_replacement if op is left_input else op
+    )
+
+    assert transformed_root is not root
+    assert transformed_root.id != root.id
+    assert transformed_root.metrics is not root.metrics
+    assert transformed_root.input_dependencies == [left_replacement, right_input]
+    assert transformed_root in left_replacement.output_dependencies
+    assert transformed_root in right_input.output_dependencies
+    assert root not in left_input.output_dependencies
+    assert root not in right_input.output_dependencies
+
+
+def test_physical_apply_transform_rewires_when_current_node_is_replaced():
+    ctx = DataContext.get_current()
+    left_input = PhysicalOperator("left_input", [], ctx)
+    right_input = PhysicalOperator("right_input", [], ctx)
+    root = PhysicalOperator("root", [left_input, right_input], ctx)
+
+    transformed_root = root._apply_transform(
+        lambda op: PhysicalOperator("replacement", [left_input], ctx)
+        if op is root
+        else op
+    )
+
+    assert transformed_root is not root
+    assert transformed_root in left_input.output_dependencies
+    assert root not in left_input.output_dependencies
+    assert root not in right_input.output_dependencies
+    assert transformed_root not in right_input.output_dependencies
+
+
+def test_physical_apply_transform_deep_chain_no_stale_downstream_refs():
+    ctx = DataContext.get_current()
+    leaf = PhysicalOperator("leaf", [], ctx)
+    mid = PhysicalOperator("mid", [leaf], ctx)
+    root = PhysicalOperator("root", [mid], ctx)
+
+    def transform(op: PhysicalOperator) -> PhysicalOperator:
+        if op is leaf:
+            return PhysicalOperator("leaf_replacement", [], ctx)
+        if op.name == "root":
+            return PhysicalOperator("root_replacement", op.input_dependencies, ctx)
+        return op
+
+    transformed_root = root._apply_transform(transform)
+    transformed_mid = transformed_root.input_dependencies[0]
+    transformed_leaf = transformed_mid.input_dependencies[0]
+
+    assert transformed_root.name == "root_replacement"
+    assert transformed_mid is not mid
+    assert transformed_leaf.name == "leaf_replacement"
+    assert root not in transformed_mid.output_dependencies
+    assert transformed_mid.output_dependencies == [transformed_root]
+
+
+def test_physical_apply_transform_rejects_in_place_input_mutation():
+    ctx = DataContext.get_current()
+    old_input = PhysicalOperator("old_input", [], ctx)
+    new_input = PhysicalOperator("new_input", [], ctx)
+    root = PhysicalOperator("root", [old_input], ctx)
+
+    def transform(op: PhysicalOperator) -> PhysicalOperator:
+        if op is root:
+            op._input_dependencies = [new_input]
+            return op
+        return op
+
+    with pytest.raises(
+        AssertionError,
+        match="In-place input mutation is not supported; return a new node instead.",
+    ):
+        root._apply_transform(transform)
+
+
+def test_does_not_double_count_usage_from_union():
+    """Regression test for https://github.com/ray-project/ray/pull/61040."""
+    # Create a mock topology:
+    #
+    #   input1 ───┐
+    #             ├─▶ union_op
+    #   input2 ───┘
+    input1 = PhysicalOperator("op1", [], DataContext.get_current())
+    input2 = PhysicalOperator("op2", [], DataContext.get_current())
+    union_op = UnionOperator(DataContext.get_current(), input1, input2)
+    topology = build_streaming_topology(union_op, ExecutionOptions(), noop_counter())
+
+    # Create a resource manager.
+    total_resources = ExecutionResources(cpu=0, object_store_memory=2)
+    resource_manager = ResourceManager(
+        topology,
+        ExecutionOptions(),
+        lambda: total_resources,
+        DataContext.get_current(),
+        BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
+    )
+
+    # Create two 1-byte `RefBundle`s.
+    block_ref1 = ray.ObjectRef(b"1" * 28)
+    block_ref2 = ray.ObjectRef(b"2" * 28)
+    block_metadata = BlockMetadata(
+        num_rows=1, size_bytes=1, input_files=None, exec_stats=None
+    )
+    bundle1 = RefBundle(
+        [BlockEntry(block_ref1, block_metadata)], owns_blocks=True, schema=None
+    )
+    bundle2 = RefBundle(
+        [BlockEntry(block_ref2, block_metadata)], owns_blocks=True, schema=None
+    )
+
+    # Add two 1-byte `RefBundle` to the union operator.
+    topology[union_op].add_output(bundle1)
+    topology[union_op].add_output(bundle2)
+    resource_manager.update_usages()
+
+    # The total object store memory usage should be 2. If the resource manager double-
+    # counts the usage from the union operator, the total object store memory usage can
+    # be greater than 2.
+    total_object_store_memory = sum(
+        [
+            resource_manager.get_op_usage(
+                op, include_ineligible_downstream=True
+            ).object_store_memory
+            for op in topology.keys()
+        ]
+    )
+    assert total_object_store_memory == 2, total_object_store_memory
+
+
+def test_per_input_inqueue_attribution_for_union():
+    """Test that per-input attribution correctly charges each upstream operator
+    only for the blocks it produced in the union's internal input queue.
+
+    When preserve_order=True, the union operator buffers blocks per-input.
+    The resource manager should attribute each input buffer's memory only to
+    the corresponding upstream operator, not to all upstream operators.
+    """
+    # Create a mock topology:
+    #
+    #   input1 ───┐
+    #             ├─▶ union_op
+    #   input2 ───┘
+    input1 = PhysicalOperator("op1", [], DataContext.get_current())
+    input2 = PhysicalOperator("op2", [], DataContext.get_current())
+    union_op = UnionOperator(DataContext.get_current(), input1, input2)
+
+    options = ExecutionOptions()
+    options.preserve_order = True
+    topology = build_streaming_topology(union_op, options, noop_counter())
+
+    # Create a resource manager.
+    total_resources = ExecutionResources(cpu=0, object_store_memory=200)
+    resource_manager = ResourceManager(
+        topology,
+        options,
+        lambda: total_resources,
+        DataContext.get_current(),
+        BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
+    )
+
+    # Create two 10-byte RefBundles with distinct block refs (simulates real execution
+    # where each block from a source has its own ObjectRef).
+    block_ref1 = ray.ObjectRef(b"1" * 28)
+    block_ref2 = ray.ObjectRef(b"2" * 28)
+    block_metadata = BlockMetadata(
+        num_rows=1, size_bytes=10, input_files=None, exec_stats=None
+    )
+    bundle1 = RefBundle(
+        [BlockEntry(block_ref1, block_metadata)], owns_blocks=True, schema=None
+    )
+    bundle2 = RefBundle(
+        [BlockEntry(block_ref2, block_metadata)], owns_blocks=True, schema=None
+    )
+
+    # Add blocks only to input2's buffer inside the union operator.
+    # With preserve_order=True, _add_input_inner routes to _input_buffers[input_index].
+    union_op.add_input(bundle1, input_index=1)
+    union_op.add_input(bundle2, input_index=1)
+
+    resource_manager.update_usages()
+
+    # input2 should be charged for its blocks in the union's input buffer (20 bytes).
+    input2_usage = resource_manager.get_op_usage(
+        input2, include_ineligible_downstream=True
+    ).object_store_memory
+    # input1 should NOT be charged for input2's blocks (0 bytes from union inqueue).
+    input1_usage = resource_manager.get_op_usage(
+        input1, include_ineligible_downstream=True
+    ).object_store_memory
+
+    assert input1_usage == 0
+    assert input2_usage == 20
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_ruleset.py

@@ -0,0 +1,209 @@
+from typing import List, Type
+
+import pytest
+
+from ray.data._internal.logical.interfaces.optimizer import Rule
+from ray.data._internal.logical.ruleset import Ruleset
+
+
+def test_add_rule():
+    class A(Rule):
+        pass
+
+    ruleset = Ruleset([A])
+    assert list(ruleset) == [A]
+
+
+def test_add_rule_with_dependencies():
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [A]
+
+    ruleset = Ruleset([A])
+    ruleset.add(B)
+    assert list(ruleset) == [A, B]
+
+
+def test_add_rule_with_dependents():
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        @classmethod
+        def dependents(cls) -> List[Type[Rule]]:
+            return [A]
+
+    ruleset = Ruleset([A])
+    ruleset.add(B)
+    assert list(ruleset) == [B, A]
+
+
+def test_add_rule_with_multiple_dependencies():
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        pass
+
+    class C(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [A, B]
+
+    ruleset = Ruleset([A, B])
+    ruleset.add(C)
+
+    rules = list(ruleset)
+    assert set(rules) == {A, B, C}
+    assert rules.index(A) < rules.index(B)
+    assert rules.index(B) < rules.index(C)
+
+
+def test_add_rule_with_multiple_dependents():
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        pass
+
+    class C(Rule):
+        @classmethod
+        def dependents(cls) -> List[Type[Rule]]:
+            return [A, B]
+
+    ruleset = Ruleset([A, B])
+    ruleset.add(C)
+
+    rules = list(ruleset)
+    assert set(rules) == {A, B, C}
+    assert rules[0] == C
+
+
+def test_add_rule_with_missing_dependencies():
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [A]
+
+    ruleset = Ruleset()
+    ruleset.add(B)
+    assert list(ruleset) == [B]
+
+
+def test_add_rule_with_missing_dependents():
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        @classmethod
+        def dependents(cls) -> List[Type[Rule]]:
+            return [A]
+
+    ruleset = Ruleset()
+    ruleset.add(B)
+    assert list(ruleset) == [B]
+
+
+def test_add_rule_with_cycle_raises_error():
+    class A(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [B]
+
+    class B(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [A]
+
+    ruleset = Ruleset([A])
+    with pytest.raises(ValueError):
+        ruleset.add(B)
+
+
+def test_edge_declared_from_both_ends_is_deduped():
+    # A must precede B, declared redundantly from both ends. The edge should be
+    # recorded once, not double-counted.
+    class A(Rule):
+        @classmethod
+        def dependents(cls) -> List[Type[Rule]]:
+            return [B]
+
+    class B(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [A]
+
+    ruleset = Ruleset([A, B])
+
+    nodes, indegree = ruleset._build_graph()
+    node_a = next(n for n in nodes if n.rule is A)
+    node_b = next(n for n in nodes if n.rule is B)
+    assert indegree[id(node_b)] == 1  # not 2
+    assert [n.rule for n in node_a.dependents] == [B]  # not [B, B]
+    assert list(ruleset) == [A, B]
+
+
+def test_disjoint_cycle_with_independent_root_raises_error():
+    # An acyclic root (A) alongside a disjoint cycle (B <-> C) must still be
+    # detected as a cycle -- the presence of a root must not mask it.
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [C]
+
+    class C(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [B]
+
+    ruleset = Ruleset([A, B])
+    with pytest.raises(ValueError):
+        ruleset.add(C)
+
+
+def test_remove_rule():
+    class A(Rule):
+        pass
+
+    ruleset = Ruleset([A])
+    ruleset.remove(A)
+    assert list(ruleset) == []
+
+
+def test_remove_rule_not_in_ruleset():
+    class A(Rule):
+        pass
+
+    ruleset = Ruleset([])
+    with pytest.raises(ValueError):
+        ruleset.remove(A)
+
+
+def test_remove_rule_with_dependants():
+    class A(Rule):
+        pass
+
+    class B(Rule):
+        @classmethod
+        def dependencies(cls) -> List[Type[Rule]]:
+            return [A]
+
+    ruleset = Ruleset([A, B])
+    ruleset.remove(A)
+    assert list(ruleset) == [B]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/unit/test_throughput_solver.py

@@ -0,0 +1,150 @@
+import pytest
+
+from ray.data._internal.cluster_autoscaler.throughput_solver import (
+    allocate_resources,
+    compute_optimal_throughput,
+)
+from ray.data._internal.execution.interfaces import ExecutionResources
+
+
+class TestComputeOptimalThroughput:
+    def test_one_op_cpu_bound(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0},
+            resource_requirements={"A": ExecutionResources(cpu=1)},
+            resource_limits=ExecutionResources.for_limits(cpu=2),
+            concurrency_limits={"A": float("inf")},
+        )
+        assert result == pytest.approx(2.0)
+
+    def test_one_op_gpu_bound(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0},
+            resource_requirements={"A": ExecutionResources(gpu=1)},
+            resource_limits=ExecutionResources.for_limits(gpu=2),
+            concurrency_limits={"A": float("inf")},
+        )
+        assert result == pytest.approx(2.0)
+
+    def test_one_op_memory_bound(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0},
+            resource_requirements={"A": ExecutionResources(memory=1e6)},
+            resource_limits=ExecutionResources.for_limits(memory=2e6),
+            concurrency_limits={"A": float("inf")},
+        )
+        assert result == pytest.approx(2.0)
+
+    def test_one_op_concurrency_bound(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0},
+            resource_requirements={"A": ExecutionResources(cpu=1)},
+            resource_limits=ExecutionResources.for_limits(),
+            concurrency_limits={"A": 2},
+        )
+        assert result == pytest.approx(2.0)
+
+    def test_two_ops_equal_rates(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0, "B": 1.0},
+            resource_requirements={
+                "A": ExecutionResources(cpu=1),
+                "B": ExecutionResources(cpu=1),
+            },
+            resource_limits=ExecutionResources.for_limits(cpu=2),
+            concurrency_limits={"A": float("inf"), "B": float("inf")},
+        )
+        assert result == pytest.approx(1.0)
+
+    def test_two_ops_different_rates(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0, "B": 2.0},
+            resource_requirements={
+                "A": ExecutionResources(cpu=1),
+                "B": ExecutionResources(cpu=1),
+            },
+            resource_limits=ExecutionResources.for_limits(cpu=3),
+            concurrency_limits={"A": float("inf"), "B": float("inf")},
+        )
+        assert result == pytest.approx(2.0)
+
+    def test_two_ops_different_resource_requirements(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0, "B": 1.0},
+            resource_requirements={
+                "A": ExecutionResources(cpu=1),
+                "B": ExecutionResources(cpu=2),
+            },
+            resource_limits=ExecutionResources.for_limits(cpu=3),
+            concurrency_limits={"A": float("inf"), "B": float("inf")},
+        )
+        assert result == pytest.approx(1.0)
+
+    def test_zero_resource_requirement(self):
+        result = compute_optimal_throughput(
+            rates={"A": 1.0},
+            resource_requirements={"A": ExecutionResources.zero()},
+            resource_limits=ExecutionResources.for_limits(cpu=1),
+            concurrency_limits={"A": float("inf")},
+        )
+        assert result == float("inf")
+
+
+class TestAllocateResources:
+    def test_empty_rates(self):
+        result = allocate_resources(
+            0.0,
+            rates={},
+            resource_requirements={},
+        )
+        assert result == {}
+
+    def test_zero_throughput(self):
+        result = allocate_resources(
+            0.0,
+            rates={"A": 1.0},
+            resource_requirements={
+                "A": ExecutionResources(cpu=1),
+            },
+        )
+        assert result == {
+            "A": ExecutionResources.zero(),
+        }
+
+    def test_one_op(self):
+        result = allocate_resources(
+            1.0,
+            rates={"A": 1.0},
+            resource_requirements={"A": ExecutionResources(cpu=1)},
+        )
+        assert result["A"] == ExecutionResources(cpu=1)
+
+    def test_two_ops_different_rates(self):
+        result = allocate_resources(
+            2.0,
+            rates={"A": 1.0, "B": 2.0},
+            resource_requirements={
+                "A": ExecutionResources(cpu=1),
+                "B": ExecutionResources(cpu=1),
+            },
+        )
+        assert result["A"] == ExecutionResources(cpu=2)
+        assert result["B"] == ExecutionResources(cpu=1)
+
+    def test_two_ops_different_resource_requirements(self):
+        result = allocate_resources(
+            1.0,
+            rates={"A": 1.0, "B": 1.0},
+            resource_requirements={
+                "A": ExecutionResources(cpu=1),
+                "B": ExecutionResources(cpu=2),
+            },
+        )
+        assert result["A"] == ExecutionResources(cpu=1)
+        assert result["B"] == ExecutionResources(cpu=2)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))