chore: import upstream snapshot with attribution

python/ray/data/tests/block_batching/test_block_batching.py

@@ -0,0 +1,96 @@
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+from ray.data._internal.arrow_block import ArrowBlockAccessor
+from ray.data._internal.block_batching.block_batching import batch_blocks
+
+
+def block_generator(num_rows: int, num_blocks: int):
+    for i in range(num_blocks):
+        yield pa.table({"foo": list(range(i * num_rows, (i + 1) * num_rows))})
+
+
+class TestBatchBlocks:
+    """Tests for batch_blocks function."""
+
+    @pytest.mark.parametrize("batch_format", ["pandas", "numpy", "pyarrow"])
+    def test_basic(self, batch_format):
+        """Test that batch_blocks yields all data in the requested format."""
+        blocks = block_generator(num_rows=3, num_blocks=2)
+        batches = list(batch_blocks(blocks, batch_format=batch_format))
+
+        assert len(batches) == 2
+
+        if batch_format == "pandas":
+            assert isinstance(batches[0], pd.DataFrame)
+            assert isinstance(batches[1], pd.DataFrame)
+            pd.testing.assert_frame_equal(
+                batches[0],
+                ArrowBlockAccessor(pa.table({"foo": [0, 1, 2]})).to_pandas(),
+            )
+            pd.testing.assert_frame_equal(
+                batches[1],
+                ArrowBlockAccessor(pa.table({"foo": [3, 4, 5]})).to_pandas(),
+            )
+        elif batch_format == "numpy":
+            assert isinstance(batches[0], dict)
+            assert isinstance(batches[1], dict)
+            np.testing.assert_array_equal(batches[0]["foo"], np.array([0, 1, 2]))
+            np.testing.assert_array_equal(batches[1]["foo"], np.array([3, 4, 5]))
+        elif batch_format == "pyarrow":
+            assert batches == [
+                pa.table({"foo": [0, 1, 2]}),
+                pa.table({"foo": [3, 4, 5]}),
+            ]
+        else:
+            pytest.fail(f"Unsupported batch format {batch_format}")
+
+    @pytest.mark.parametrize(
+        "batch_size,drop_last,expected_values",
+        [
+            # 6 rows, batch_size=2: yields 3 full batches
+            (2, False, [[0, 1], [2, 3], [4, 5]]),
+            # 6 rows, batch_size=4: yields 1 full + 1 partial batch
+            (4, False, [[0, 1, 2, 3], [4, 5]]),
+            # 6 rows, batch_size=4, drop_last: drops partial batch
+            (4, True, [[0, 1, 2, 3]]),
+            # 6 rows, batch_size=10, drop_last: no batches (all dropped)
+            (10, True, []),
+        ],
+    )
+    def test_batch_size(self, batch_size, drop_last, expected_values):
+        """Test batch_size and drop_last parameters."""
+        blocks = block_generator(num_rows=3, num_blocks=2)
+        batches = list(
+            batch_blocks(
+                blocks,
+                batch_size=batch_size,
+                drop_last=drop_last,
+                batch_format="numpy",
+            )
+        )
+
+        assert len(batches) == len(expected_values)
+        for batch, expected in zip(batches, expected_values):
+            np.testing.assert_array_equal(batch["foo"], np.array(expected))
+
+    def test_collate_fn(self):
+        """Test that collate_fn transforms batches."""
+
+        def double_values(batch):
+            return {"foo": [x * 2 for x in batch["foo"].tolist()]}
+
+        blocks = block_generator(num_rows=3, num_blocks=2)
+        batches = list(batch_blocks(blocks, collate_fn=double_values))
+
+        assert len(batches) == 2
+        assert batches[0]["foo"] == [0, 2, 4]
+        assert batches[1]["foo"] == [6, 8, 10]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/block_batching/test_interfaces.py

@@ -0,0 +1,170 @@
+"""Tests for ray.data._internal.block_batching.interfaces."""
+
+import pytest
+
+from ray.data._internal.block_batching.interfaces import (
+    Batch,
+    BatchMetadata,
+    BatchStageTimings,
+    BlockStageTimings,
+)
+from ray.data._internal.stats import IterationStage, TimeSpan
+
+
+class TestAccumulateBlockTimings:
+    """Tests for BatchStageTimings.accumulate_block_timings().
+
+    accumulate_block_timings appends each block's spans to the batch's lists
+    (no merging) so that overlap attribution can sum non-overlapping spans
+    without double-counting.
+    """
+
+    def test_single_block(self):
+        """Accumulating a single block appends its span."""
+        dst = BatchStageTimings()
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=1.0, end_s=2.0),
+                data_transfer=TimeSpan(start_s=2.0, end_s=3.0),
+            )
+        )
+        assert len(dst.production_wait) == 1
+        assert dst.production_wait[0].start_s == 1.0
+        assert dst.production_wait[0].end_s == 2.0
+
+    def test_multiple_blocks_kept_separate(self):
+        """Multiple blocks produce a list of separate spans (no merge)."""
+        dst = BatchStageTimings()
+
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=1.0, end_s=2.0),
+                data_transfer=TimeSpan(start_s=2.0, end_s=3.0),
+            )
+        )
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=3.0, end_s=4.0),
+                data_transfer=TimeSpan(start_s=4.0, end_s=5.0),
+            )
+        )
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=5.0, end_s=6.0),
+                data_transfer=TimeSpan(start_s=6.0, end_s=7.0),
+            )
+        )
+
+        assert len(dst.production_wait) == 3
+        assert [s.start_s for s in dst.production_wait] == [1.0, 3.0, 5.0]
+        assert [s.end_s for s in dst.production_wait] == [2.0, 4.0, 6.0]
+
+    def test_overlapping_blocks_kept_separate(self):
+        """Overlapping windows are NOT merged — kept as separate spans."""
+        dst = BatchStageTimings()
+
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=1.0, end_s=5.0),
+                data_transfer=TimeSpan(start_s=5.0, end_s=6.0),
+            )
+        )
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=3.0, end_s=7.0),
+                data_transfer=TimeSpan(start_s=7.0, end_s=8.0),
+            )
+        )
+
+        assert len(dst.production_wait) == 2
+        assert dst.production_wait[0].start_s == 1.0
+        assert dst.production_wait[1].end_s == 7.0
+
+    def test_into_empty_destination(self):
+        """Accumulating into an empty BatchStageTimings appends the span."""
+        dst = BatchStageTimings()
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=10.0, end_s=20.0),
+                data_transfer=TimeSpan(start_s=20.0, end_s=30.0),
+            )
+        )
+        assert len(dst.production_wait) == 1
+        assert dst.production_wait[0].start_s == 10.0
+
+    def test_data_transfer_multiple_blocks(self):
+        """data_transfer spans are kept separate across multiple blocks."""
+        dst = BatchStageTimings()
+
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=0.0, end_s=1.0),
+                data_transfer=TimeSpan(start_s=1.0, end_s=2.0),
+            )
+        )
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=2.0, end_s=3.0),
+                data_transfer=TimeSpan(start_s=3.0, end_s=4.0),
+            )
+        )
+
+        assert len(dst.data_transfer) == 2
+        assert [s.start_s for s in dst.data_transfer] == [1.0, 3.0]
+
+    def test_both_stages_independent(self):
+        """production_wait and data_transfer lists are independent."""
+        dst = BatchStageTimings()
+
+        # Block 1: prod [1,2], xfer [2,3]
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=1.0, end_s=2.0),
+                data_transfer=TimeSpan(start_s=2.0, end_s=3.0),
+            )
+        )
+        # Block 2: prod [5,6], xfer [6,7]
+        dst.accumulate_block_timings(
+            BlockStageTimings(
+                production_wait=TimeSpan(start_s=5.0, end_s=6.0),
+                data_transfer=TimeSpan(start_s=6.0, end_s=7.0),
+            )
+        )
+
+        assert len(dst.production_wait) == 2
+        assert len(dst.data_transfer) == 2
+        assert [s.start_s for s in dst.production_wait] == [1.0, 5.0]
+        assert [s.start_s for s in dst.data_transfer] == [2.0, 6.0]
+
+
+class TestStageTimingsFields:
+    """Tests that BatchStageTimings fields are accessible via stages()."""
+
+    def test_batch_carries_timings_through_pipeline(self):
+        """A Batch's metadata.stage_timings carries all stage windows."""
+        timings = BatchStageTimings()
+        timings.production_wait.append(TimeSpan(start_s=1.0, end_s=2.0))
+        timings.batching = TimeSpan(start_s=2.0, end_s=3.0)
+        timings.format = TimeSpan(start_s=3.0, end_s=4.0)
+        timings.collate = TimeSpan(start_s=4.0, end_s=5.0)
+        timings.finalize = TimeSpan(start_s=5.0, end_s=6.0)
+
+        batch = Batch(
+            BatchMetadata(batch_idx=0, num_rows=50, stage_timings=timings), None
+        )
+
+        # Verify all stages are accessible via stages() iterator
+        stage_dict = dict(batch.metadata.stage_timings.stages())
+        assert len(stage_dict) == 6
+        assert stage_dict[IterationStage.PRODUCTION_WAIT][0].start_s == 1.0
+        assert stage_dict[IterationStage.BATCHING][0].end_s == 3.0
+        assert stage_dict[IterationStage.FORMAT][0].start_s == 3.0
+        assert stage_dict[IterationStage.COLLATE][0].end_s == 5.0
+        assert stage_dict[IterationStage.FINALIZE][0].start_s == 5.0
+        assert batch.metadata.num_rows == 50
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/block_batching/test_iter_batches.py

@@ -0,0 +1,874 @@
+import queue
+import threading
+import time
+from typing import Iterator, List, Optional
+from unittest.mock import patch
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.block_batching.interfaces import (
+    Batch,
+    BatchMetadata,
+    BatchStageTimings,
+    BlockPrefetcher,
+)
+from ray.data._internal.block_batching.iter_batches import (
+    BatchIterator,
+    prefetch_batches_locally,
+    restore_original_order,
+)
+from ray.data._internal.block_batching.util import (
+    WaitBlockPrefetcher,
+)
+from ray.data._internal.execution.interfaces.ref_bundle import BlockEntry, RefBundle
+from ray.data._internal.stats import DatasetStats, TimeSpan
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.types import ObjectRef
+
+# Sleep duration injected into each scenario's bottleneck stage. Picked to be
+# large enough to dominate scheduling/measurement noise but small enough to
+# keep the test fast (5 batches × 0.3s ≈ 1.5s per scenario).
+SLEEP_S = 0.3
+
+
+def ref_bundle_generator(num_rows: int, num_blocks: int) -> Iterator[RefBundle]:
+    for i in range(num_blocks):
+        block = pa.table({"foo": [i] * num_rows})
+        metadata = BlockMetadata(
+            num_rows=num_rows,
+            size_bytes=0,
+            input_files=[],
+            exec_stats=None,
+        )
+        schema = block.schema
+        yield RefBundle(
+            blocks=(BlockEntry(ray.put(block), metadata),),
+            owns_blocks=True,
+            schema=schema,
+        )
+
+
+@pytest.mark.parametrize("num_batches_to_prefetch", [1, 2])
+@pytest.mark.parametrize("batch_size", [None, 1, 4])
+def test_prefetch_batches_locally(
+    ray_start_regular_shared, num_batches_to_prefetch, batch_size
+):
+    class DummyPrefetcher(BlockPrefetcher):
+        def __init__(self):
+            self.windows = []
+
+        def prefetch_blocks(self, block_refs: List[ObjectRef[Block]]):
+            if batch_size is None:
+                assert len(block_refs) == num_batches_to_prefetch
+            else:
+                assert (
+                    sum(len(ray.get(block_ref)) for block_ref in block_refs)
+                    >= batch_size * num_batches_to_prefetch
+                )
+            self.windows.append(block_refs)
+
+    num_blocks = 10
+    num_rows = 2
+    prefetcher = DummyPrefetcher()
+    ref_bundles = list(ref_bundle_generator(num_blocks=num_blocks, num_rows=num_rows))
+    prefetch_block_iter = prefetch_batches_locally(
+        iter(ref_bundles),
+        prefetcher=prefetcher,
+        num_batches_to_prefetch=num_batches_to_prefetch,
+        batch_size=batch_size,
+    )
+
+    block_count = 0
+    prefetched_blocks = []
+    previous_num_windows = 1
+
+    for block in prefetch_block_iter:
+        prefetched_blocks.append(block)
+        block_count += 1
+        remaining_rows = (num_blocks - block_count) * num_rows
+        if batch_size is None and block_count < num_blocks - num_batches_to_prefetch:
+            # Test that we are actually prefetching in advance if this is not the last
+            # block.
+            assert len(prefetcher.windows) == previous_num_windows + 1
+            previous_num_windows = len(prefetcher.windows)
+        elif (
+            batch_size is not None
+            and remaining_rows > batch_size * num_batches_to_prefetch
+        ):
+            # Test that we are actually prefetching in advance if this is not the last
+            # batch.
+            assert len(prefetcher.windows) == previous_num_windows + 1
+            previous_num_windows = len(prefetcher.windows)
+
+    # Test that original blocks are unchanged.
+    expected_blocks = []
+    for ref_bundle in ref_bundles:
+        expected_blocks.extend(ref_bundle.block_refs)
+    assert prefetched_blocks == expected_blocks
+
+
+def test_restore_from_original_order():
+    base_iterator = [
+        Batch(BatchMetadata(batch_idx=1), None),
+        Batch(BatchMetadata(batch_idx=0), None),
+        Batch(BatchMetadata(batch_idx=3), None),
+        Batch(BatchMetadata(batch_idx=2), None),
+    ]
+
+    ordered = list(restore_original_order(iter(base_iterator)))
+    idx = [batch.metadata.batch_idx for batch in ordered]
+    assert idx == [0, 1, 2, 3]
+
+
+def test_attribute_blocked_time_overlap_attribution():
+    stats = DatasetStats(metadata={}, parent=None)
+    batch_iterator = BatchIterator(iter([]), stats=stats)
+    timings = BatchStageTimings()
+    timings.production_wait.append(TimeSpan(start_s=10.0, end_s=20.0))
+    timings.batching = TimeSpan(start_s=20.0, end_s=30.0)
+    timings.format = TimeSpan(start_s=30.0, end_s=40.0)
+    timings.finalize = TimeSpan(start_s=50.0, end_s=60.0)
+    batch = Batch(BatchMetadata(batch_idx=0, num_rows=8, stage_timings=timings), None)
+
+    batch_iterator._attribute_blocked_time(
+        batch, blocked_start_s=15.0, blocked_end_s=35.0
+    )
+
+    assert stats.iter_blocked_production_wait_s.get() == pytest.approx(5.0)
+    assert stats.iter_blocked_batching_s.get() == pytest.approx(10.0)
+    assert stats.iter_blocked_format_s.get() == pytest.approx(5.0)
+    assert stats.iter_blocked_collate_s.get() == 0
+    assert stats.iter_blocked_finalize_s.get() == 0
+    assert stats.iter_batches_total == 1
+    assert stats.iter_rows_total == 8
+
+
+def _make_span(start: Optional[float], end: Optional[float]) -> Optional[TimeSpan]:
+    """Create a TimeSpan, or None if the stage didn't run."""
+    if start is None or end is None:
+        return None
+    return TimeSpan(start_s=start, end_s=end)
+
+
+def _make_batch_with_timings(
+    production_wait_start: Optional[float] = None,
+    production_wait_end: Optional[float] = None,
+    data_transfer_start: Optional[float] = None,
+    data_transfer_end: Optional[float] = None,
+    batching_start: Optional[float] = None,
+    batching_end: Optional[float] = None,
+    format_start: Optional[float] = None,
+    format_end: Optional[float] = None,
+    collate_start: Optional[float] = None,
+    collate_end: Optional[float] = None,
+    finalize_start: Optional[float] = None,
+    finalize_end: Optional[float] = None,
+    num_rows: int = 0,
+):
+    """Helper to construct a Batch with specific stage timing windows."""
+    timings = BatchStageTimings()
+    pw = _make_span(production_wait_start, production_wait_end)
+    if pw is not None:
+        timings.production_wait.append(pw)
+    dt = _make_span(data_transfer_start, data_transfer_end)
+    if dt is not None:
+        timings.data_transfer.append(dt)
+    timings.batching = _make_span(batching_start, batching_end)
+    timings.format = _make_span(format_start, format_end)
+    timings.collate = _make_span(collate_start, collate_end)
+    timings.finalize = _make_span(finalize_start, finalize_end)
+    return Batch(
+        BatchMetadata(batch_idx=0, num_rows=num_rows, stage_timings=timings), None
+    )
+
+
+def _make_test_iterator(stats):
+    """Create a BatchIterator wired to the given stats without a real pipeline."""
+    it = BatchIterator.__new__(BatchIterator)
+    it._stats = stats
+    return it
+
+
+class TestAttributeBlockedTimeEdgeCases:
+    """Edge case tests for overlap-based blocked attribution."""
+
+    def test_zero_overlap_stage_finished_before_blocked(self):
+        """Fetch [0, 1.5] finished before training blocked at t=2 → 0 attribution."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(
+            production_wait_start=0.0, production_wait_end=1.5
+        )
+        it._attribute_blocked_time(batch, blocked_start_s=2.0, blocked_end_s=3.0)
+        assert stats.iter_blocked_production_wait_s.get() == 0.0
+
+    def test_zero_overlap_blocked_before_stage(self):
+        """Training blocked [0, 1], stage ran [2, 3] → 0 attribution."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(format_start=2.0, format_end=3.0)
+        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=1.0)
+        assert stats.iter_blocked_format_s.get() == 0.0
+
+    def test_partial_overlap(self):
+        """Fetch [0, 2], blocked [1, 3] → overlap = min(2,3)-max(0,1) = 1.0."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(
+            production_wait_start=0.0, production_wait_end=2.0
+        )
+        it._attribute_blocked_time(batch, blocked_start_s=1.0, blocked_end_s=3.0)
+        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(1.0)
+
+    def test_full_overlap_stage_inside_blocked(self):
+        """Stage [1, 2] entirely inside blocked [0, 3] → full 1.0 credit."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(batching_start=1.0, batching_end=2.0)
+        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
+        assert stats.iter_blocked_batching_s.get() == pytest.approx(1.0)
+
+    def test_no_collate_fn_zero_attribution(self):
+        """collate stage has start_s=0 → skipped, 0 attribution."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(format_start=1.0, format_end=2.0)
+        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
+        assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
+        assert stats.iter_blocked_collate_s.get() == 0.0
+
+    def test_no_finalize_fn_zero_attribution(self):
+        """finalize stage has start_s=0 → skipped, 0 attribution."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(collate_start=1.0, collate_end=2.0)
+        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
+        assert stats.iter_blocked_collate_s.get() == pytest.approx(1.0)
+        assert stats.iter_blocked_finalize_s.get() == 0.0
+
+    def test_prefetch_hides_fetch_from_training(self):
+        """Effective prefetch: fetch done before training blocks → 0 fetch attribution."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(
+            production_wait_start=0.0,
+            production_wait_end=1.5,
+            collate_start=2.3,
+            collate_end=2.6,
+        )
+        # Training only starts blocking at t=2 (prefetch worked)
+        it._attribute_blocked_time(batch, blocked_start_s=2.0, blocked_end_s=2.6)
+        assert stats.iter_blocked_production_wait_s.get() == 0.0
+        assert stats.iter_blocked_collate_s.get() == pytest.approx(0.3)
+
+    def test_accumulation_across_batches(self):
+        """Two batches each contribute to fetch — values accumulate."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        # Batch 1: fetch [0,1], blocked [0,2] → overlap 1.0
+        b1 = _make_batch_with_timings(
+            production_wait_start=0.0, production_wait_end=1.0, num_rows=10
+        )
+        it._attribute_blocked_time(b1, blocked_start_s=0.0, blocked_end_s=2.0)
+        # Batch 2: fetch [5,6], blocked [5,7] → overlap 1.0
+        b2 = _make_batch_with_timings(
+            production_wait_start=5.0, production_wait_end=6.0, num_rows=20
+        )
+        it._attribute_blocked_time(b2, blocked_start_s=5.0, blocked_end_s=7.0)
+
+        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(2.0)
+        assert stats.iter_batches_total == 2
+        assert stats.iter_rows_total == 30
+
+    def test_overlap_invariant_sum_leq_total(self):
+        """sum(iter_blocked_*) <= iter_total_blocked_s holds for non-overlapping stages."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        stats.iter_total_blocked_s.add(5.0)
+        batch = _make_batch_with_timings(
+            production_wait_start=0.0,
+            production_wait_end=1.0,
+            batching_start=1.0,
+            batching_end=2.0,
+            format_start=2.0,
+            format_end=3.0,
+            num_rows=5,
+        )
+        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)
+
+        total = stats.iter_total_blocked_s.get()
+        sum_stages = (
+            stats.iter_blocked_production_wait_s.get()
+            + stats.iter_blocked_batching_s.get()
+            + stats.iter_blocked_format_s.get()
+            + stats.iter_blocked_collate_s.get()
+            + stats.iter_blocked_finalize_s.get()
+        )
+        assert sum_stages <= total + 1e-9
+
+    def test_blocked_inside_stage(self):
+        """Stage [0, 10] fully contains blocked [3, 5] → overlap = 2.0."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(
+            production_wait_start=0.0, production_wait_end=10.0
+        )
+        it._attribute_blocked_time(batch, blocked_start_s=3.0, blocked_end_s=5.0)
+        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(2.0)
+
+    def test_all_stages_simultaneous_overlap(self):
+        """Multiple stages overlap with blocked window simultaneously."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        batch = _make_batch_with_timings(
+            production_wait_start=0.0,
+            production_wait_end=1.0,
+            batching_start=1.0,
+            batching_end=2.0,
+            format_start=2.0,
+            format_end=3.0,
+            collate_start=3.0,
+            collate_end=4.0,
+            finalize_start=4.0,
+            finalize_end=5.0,
+            num_rows=100,
+        )
+        # Blocked window covers all stages
+        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)
+        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(1.0)
+        assert stats.iter_blocked_batching_s.get() == pytest.approx(1.0)
+        assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
+        assert stats.iter_blocked_collate_s.get() == pytest.approx(1.0)
+        assert stats.iter_blocked_finalize_s.get() == pytest.approx(1.0)
+        assert stats.iter_batches_total == 1
+        assert stats.iter_rows_total == 100
+
+    def test_overlapping_spans_not_double_counted(self):
+        """Two overlapping production_wait spans: union, not sum."""
+        stats = DatasetStats(metadata={}, parent=None)
+        it = _make_test_iterator(stats)
+        # Block 1: prod [0, 100], Block 2: prod [50, 150] — overlap [50, 100]
+        # Blocked [0, 200] covers both
+        batch = _make_batch_with_timings(
+            production_wait_start=0.0,
+            production_wait_end=100.0,
+            num_rows=10,
+        )
+        # Add a second production_wait span (multi-block batch)
+        batch.metadata.stage_timings.production_wait.append(
+            TimeSpan(start_s=50.0, end_s=150.0)
+        )
+        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=200.0)
+        # Union of [0,100] and [50,150] = [0,150] = 150, NOT 100+100=200
+        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(150.0)
+
+
+def test_attribute_blocked_time_all_stages_full_overlap():
+    """All stages with realistic timing, full overlap with blocked window."""
+    stats = DatasetStats(metadata={}, parent=None)
+    it = _make_test_iterator(stats)
+    stats.iter_total_blocked_s.add(5.0)
+
+    batch = _make_batch_with_timings(
+        production_wait_start=0.0,
+        production_wait_end=0.5,
+        batching_start=0.5,
+        batching_end=1.0,
+        format_start=1.0,
+        format_end=2.0,
+        collate_start=2.0,
+        collate_end=2.5,
+        finalize_start=2.5,
+        finalize_end=3.0,
+        num_rows=256,
+    )
+
+    # Blocked window covers all stages
+    it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)
+
+    # Each stage gets its full duration
+    assert stats.iter_blocked_production_wait_s.get() == pytest.approx(0.5)
+    assert stats.iter_blocked_batching_s.get() == pytest.approx(0.5)
+    assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
+    assert stats.iter_blocked_collate_s.get() == pytest.approx(0.5)
+    assert stats.iter_blocked_finalize_s.get() == pytest.approx(0.5)
+    assert stats.iter_batches_total == 1
+    assert stats.iter_rows_total == 256
+
+    # Invariant: sum = 3.0 <= total_blocked = 5.0
+    sum_stages = (
+        stats.iter_blocked_production_wait_s.get()
+        + stats.iter_blocked_batching_s.get()
+        + stats.iter_blocked_format_s.get()
+        + stats.iter_blocked_collate_s.get()
+        + stats.iter_blocked_finalize_s.get()
+    )
+    assert sum_stages == pytest.approx(3.0)
+    assert sum_stages <= stats.iter_total_blocked_s.get() + 1e-9
+
+
+def test_finalize_fn_uses_single_thread(ray_start_regular_shared):
+    """Tests that finalize_fn is not run with multiple threads."""
+    ref_bundles_iter = ref_bundle_generator(num_blocks=20, num_rows=2)
+
+    q = queue.Queue()
+    semaphore = threading.Semaphore(value=1)
+
+    def finalize_enforce_single_thread(batch):
+        already_acquired = not semaphore.acquire(blocking=False)
+        if already_acquired:
+            e = AssertionError("finalize_fn is being run concurrently.")
+            q.put(e, block=True)
+        semaphore.release()
+        return batch
+
+    # Test that finalize_fn is called in a single thread,
+    # even if prefetch_batches is set.
+    output_batches = BatchIterator(
+        ref_bundles_iter,
+        collate_fn=lambda batch: batch,
+        finalize_fn=finalize_enforce_single_thread,
+        prefetch_batches=4,
+    )
+
+    # Force execution of the iterator.
+    # This step should not raise an exception.
+    list(output_batches)
+
+    try:
+        e = q.get(block=False, timeout=0.1)
+        raise e
+    except queue.Empty:
+        pass
+
+
+# Test for 3 cases
+# 1. Batch size is less than block size
+# 2. Batch size is more than block size
+# 3. Block size is not divisble by batch size
+@pytest.mark.parametrize("batch_size", [1, 4, 3])
+@pytest.mark.parametrize("drop_last", [True, False])
+@pytest.mark.parametrize("prefetch_batches", [0, 1])
+def test_iter_batches_e2e(
+    ray_start_regular_shared, batch_size, drop_last, prefetch_batches
+):
+    def collate_fn(batch: pd.DataFrame):
+        return batch + 1
+
+    ref_bundles_iter = ref_bundle_generator(num_blocks=4, num_rows=2)
+
+    output_batches = BatchIterator(
+        ref_bundles_iter,
+        batch_size=batch_size,
+        prefetch_batches=prefetch_batches,
+        batch_format="pandas",
+        collate_fn=collate_fn,
+        drop_last=drop_last,
+        preserve_order=True,
+    )
+
+    output_batches = list(output_batches)
+
+    assert len(output_batches) > 0
+    for df in output_batches:
+        # Check batch formatting.
+        assert isinstance(df, pd.DataFrame)
+        # Check batch size.
+        if batch_size == 3 and not drop_last:
+            assert len(df) in {2, 3}
+        else:
+            assert len(df) == batch_size
+
+    concat_df = pd.concat(output_batches)
+    # Test that collate_fn is applied.
+    assert concat_df["foo"].iloc[0] == 1
+    # Make sure order is preserved.
+    for i in range(len(concat_df) - 1):
+        assert concat_df["foo"].iloc[i + 1] >= concat_df["foo"].iloc[i]
+
+
+def test_iter_batches_counts_rows_at_pipeline_exit(ray_start_regular_shared):
+    stats = DatasetStats(metadata={}, parent=None)
+    ref_bundles_iter = ref_bundle_generator(num_blocks=4, num_rows=2)
+
+    output_batches = list(
+        BatchIterator(
+            ref_bundles_iter,
+            stats=stats,
+            batch_size=3,
+            prefetch_batches=0,
+            batch_format="pandas",
+            drop_last=True,
+        )
+    )
+
+    assert len(output_batches) == 2
+    assert [len(batch) for batch in output_batches] == [3, 3]
+    assert stats.iter_batches_total == 2
+    assert stats.iter_rows_total == 6
+
+
+def test_iter_batches_e2e_async(ray_start_regular_shared):
+    """We add time.sleep in 3 places:
+    1. In the base generator to simulate streaming executor blocking on next results.
+    2. In the collate_fn to simulate expensive slicing/formatting/collation
+    3. In the user thread to simulate training.
+    """
+
+    def collate_fn(batch):
+        time.sleep(2)
+        return batch
+
+    ref_bundles = ref_bundle_generator(num_blocks=20, num_rows=2)
+    start_time = time.time()
+    output_batches = BatchIterator(
+        ref_bundles,
+        batch_size=None,
+        collate_fn=collate_fn,
+        prefetch_batches=4,
+    )
+    batches = []
+    for batch in output_batches:
+        time.sleep(1.5)
+        batches.append(batch)
+    end_time = time.time()
+
+    # 20 batches, 1.5 second sleep. Should be less than 45 seconds, even with some
+    # overhead.
+    # If there was no overlap, then we would expect this to take at least 20*2.5 = 50
+    assert end_time - start_time < 45, end_time - start_time
+
+    assert len(batches) == 20
+    assert all(len(batch) == 2 for batch in batches)
+
+
+@pytest.mark.parametrize("preserve_order", [True, False])
+def test_iter_batches_preserve_order_flag(
+    ray_start_regular_shared, preserve_order, restore_data_context
+):
+    """When `execution_options.preserve_order` is True, batches must come
+    out in input order even with a multi-worker format threadpool. When
+    False, ordering is not guaranteed (but the full set of batches must
+    still be produced)."""
+    # Variable per-batch collate cost makes worker-completion order
+    # arbitrary so the reorder path actually does work when enabled.
+    def collate_fn(batch):
+        idx = int(batch["foo"][0])
+        time.sleep(0.05 * (idx % 4))
+        return batch
+
+    num_blocks = 16
+    ref_bundles = ref_bundle_generator(num_blocks=num_blocks, num_rows=1)
+    output_batches = list(
+        BatchIterator(
+            ref_bundles,
+            batch_size=1,
+            collate_fn=collate_fn,
+            batch_format="pandas",
+            prefetch_batches=4,
+            preserve_order=preserve_order,
+        )
+    )
+
+    indices = [int(df["foo"].iloc[0]) for df in output_batches]
+    assert sorted(indices) == list(range(num_blocks))
+    if preserve_order:
+        assert indices == list(range(num_blocks)), indices
+
+
+def test_finalize_fn_runs_after_restore_original_order(ray_start_regular_shared):
+    """When preserve_order=True, finalize_fn must run after the reorder
+    buffer so that the buffer holds CPU batches rather than finalize_fn
+    outputs (e.g., GPU tensors). Asserts finalize_fn sees batches in
+    monotonically increasing order even when the format/collate threadpool
+    completes them out of order."""
+
+    def collate_fn(batch):
+        # Variable per-batch cost so worker-completion order is arbitrary.
+        idx = int(batch["foo"].iloc[0])
+        time.sleep(0.05 * (idx % 4))
+        return batch
+
+    seen_by_finalize = []
+    seen_lock = threading.Lock()
+
+    def finalize_fn(batch):
+        idx = int(batch["foo"].iloc[0])
+        with seen_lock:
+            seen_by_finalize.append(idx)
+        return batch
+
+    num_blocks = 16
+    ref_bundles = ref_bundle_generator(num_blocks=num_blocks, num_rows=1)
+    list(
+        BatchIterator(
+            ref_bundles,
+            batch_size=1,
+            collate_fn=collate_fn,
+            finalize_fn=finalize_fn,
+            batch_format="pandas",
+            prefetch_batches=4,
+            preserve_order=True,
+        )
+    )
+
+    assert seen_by_finalize == list(range(num_blocks)), seen_by_finalize
+
+
+def _ref_bundles_with_size(
+    num_blocks: int, num_rows: int, size_bytes_per_block: int
+) -> Iterator[RefBundle]:
+    """Create ref bundles with explicit size_bytes for testing."""
+    for i in range(num_blocks):
+        block = pa.table({"foo": [i] * num_rows})
+        metadata = BlockMetadata(
+            num_rows=num_rows,
+            size_bytes=size_bytes_per_block,
+            input_files=[],
+            exec_stats=None,
+        )
+        schema = block.schema
+        yield RefBundle(
+            blocks=(BlockEntry(ray.put(block), metadata),),
+            owns_blocks=True,
+            schema=schema,
+        )
+
+
+@pytest.mark.parametrize(
+    "num_batches_to_prefetch,expected_bytes_sequence",
+    [
+        # No prefetching: all 5 blocks report 0 prefetched bytes
+        (0, [0, 0, 0, 0, 0]),
+        # prefetch 2 blocks: with 5 blocks of 100 bytes each
+        # After yield block 0: window has 1,2 -> 200 (added block 2)
+        # After yield block 1: window has 2,3 -> 200 (added block 3)
+        # After yield block 2: window has 3,4 -> 200 (added block 4)
+        # After yield block 3: window has 4 -> 100 (no more to add)
+        # After yield block 4: window empty -> 0
+        (2, [200, 200, 200, 100, 0]),
+    ],
+)
+def test_prefetch_bytes_tracking(
+    ray_start_regular_shared, num_batches_to_prefetch, expected_bytes_sequence
+):
+    """Test iter_prefetched_bytes is set correctly during prefetching.
+
+    Tests prefetch_batches_locally directly to verify exact values,
+    bypassing async BatchIterator which has non-deterministic timing.
+    """
+    stats = DatasetStats(metadata={}, parent=None)
+
+    # Create 5 ref bundles, each with size_bytes=100
+    num_blocks = 5
+    ref_bundles = list(
+        _ref_bundles_with_size(num_blocks, num_rows=2, size_bytes_per_block=100)
+    )
+
+    prefetcher = WaitBlockPrefetcher()
+    block_iter = prefetch_batches_locally(
+        iter(ref_bundles),
+        prefetcher=prefetcher,
+        num_batches_to_prefetch=num_batches_to_prefetch,
+        batch_size=None,
+        stats=stats,
+    )
+
+    # Track iter_prefetched_bytes after each block is yielded
+    recorded_bytes = []
+    for _ in block_iter:
+        recorded_bytes.append(stats.iter_prefetched_bytes)
+
+    assert recorded_bytes == expected_bytes_sequence, f"Got {recorded_bytes}"
+
+
+@pytest.mark.parametrize("prefetch_batches", [0, 2])
+def test_prefetch_bytes_callback(ray_start_regular_shared, prefetch_batches):
+    """Test prefetch_bytes_callback is invoked correctly by BatchIterator."""
+    reported_bytes = []
+
+    def prefetch_callback(num_bytes: int):
+        reported_bytes.append(num_bytes)
+
+    stats = DatasetStats(metadata={}, parent=None)
+
+    # Create 5 ref bundles
+    num_blocks = 5
+    ref_bundles = list(
+        _ref_bundles_with_size(num_blocks, num_rows=2, size_bytes_per_block=100)
+    )
+
+    output_batches = BatchIterator(
+        iter(ref_bundles),
+        stats=stats,
+        batch_size=None,
+        prefetch_batches=prefetch_batches,
+        prefetch_bytes_callback=prefetch_callback,
+    )
+
+    # Consume all batches
+    batches = list(output_batches)
+
+    assert len(batches) == 5
+
+    # Callback is called 5 times (per batch) + 1 time at epoch end
+    assert len(reported_bytes) == 6, f"Expected 6, got {len(reported_bytes)}"
+
+    # All values should be non-negative
+    assert all(b >= 0 for b in reported_bytes), f"Negative: {reported_bytes}"
+
+    # Last value should be 0 (after_epoch_end)
+    assert reported_bytes[-1] == 0, f"Last should be 0: {reported_bytes}"
+
+
+@pytest.mark.parametrize(
+    "scenario,bound_stage",
+    [
+        ("production", "iter_blocked_production_wait_s"),
+        ("data_transfer", "iter_blocked_data_transfer_s"),
+        ("batching", "iter_blocked_batching_s"),
+        ("collate", "iter_blocked_collate_s"),
+        ("format", "iter_blocked_format_s"),
+        ("finalize", "iter_blocked_finalize_s"),
+    ],
+)
+def test_e2e_blocked_attribution_by_scenario(
+    ray_start_regular_shared, scenario, bound_stage
+):
+    """E2e: when a specific stage is the bottleneck, its blocked metric
+    should be the largest among all stages, and at least SLEEP_S."""
+    from ray.data._internal.stats import _StatsManager
+
+    iter_kwargs = {"batch_size": 10, "prefetch_batches": 0}
+    patches = []
+
+    if scenario == "production":
+        # Slow upstream map → production_wait dominates.
+        def slow_map(batch):
+            time.sleep(SLEEP_S)
+            return batch
+
+        ds = ray.data.range(50, override_num_blocks=5).map(slow_map)
+
+    elif scenario == "data_transfer":
+        # Patch ray.get ONLY in util.resolve_block_refs (not globally) so the
+        # streaming executor's own ray.get calls aren't slowed (which would
+        # inflate production_wait). We replace util_mod.ray with a proxy that
+        # has a slow `get` but delegates everything else to the real ray.
+        from ray.data._internal.block_batching import util as util_mod
+
+        orig_get = ray.get
+
+        class _SlowGetRayProxy:
+            """Proxy that sleeps on `.get` but delegates everything else."""
+
+            def __getattr__(self, name):
+                return getattr(ray, name)
+
+            @staticmethod
+            def get(ref):
+                time.sleep(SLEEP_S)
+                return orig_get(ref)
+
+        patches.append(patch.object(util_mod, "ray", _SlowGetRayProxy()))
+        ds = ray.data.range(50, override_num_blocks=5)
+
+    elif scenario == "batching":
+        # Patch Batcher.next_batch to inject slow batching.
+        from ray.data._internal.batcher import Batcher
+
+        orig_next_batch = Batcher.next_batch
+
+        def slow_next_batch(self):
+            time.sleep(SLEEP_S)
+            return orig_next_batch(self)
+
+        patches.append(patch.object(Batcher, "next_batch", slow_next_batch))
+        ds = ray.data.range(50, override_num_blocks=5)
+
+    elif scenario == "collate":
+        # Pass _collate_fn via _iter_batches (private signature accepts it;
+        # public iter_batches does not).
+        def slow_collate(batch):
+            time.sleep(SLEEP_S)
+            return batch
+
+        iter_kwargs["_collate_fn"] = slow_collate
+        ds = ray.data.range(50, override_num_blocks=5)
+
+    elif scenario == "format":
+        # Patch BlockAccessor.to_batch_format — it's called INSIDE
+        # _format_batch's _maybe_time context, so the sleep is captured by
+        # the format timing span.
+        orig_to_batch_format = BlockAccessor.to_batch_format
+
+        def slow_to_batch_format(self, batch_format):
+            time.sleep(SLEEP_S)
+            return orig_to_batch_format(self, batch_format)
+
+        patches.append(
+            patch.object(BlockAccessor, "to_batch_format", slow_to_batch_format)
+        )
+        ds = ray.data.range(50, override_num_blocks=5)
+
+    elif scenario == "finalize":
+        # Pass _finalize_fn via _iter_batches (private signature accepts it).
+        def slow_finalize(data):
+            time.sleep(SLEEP_S)
+            return data
+
+        iter_kwargs["_finalize_fn"] = slow_finalize
+        ds = ray.data.range(50, override_num_blocks=5)
+
+    it = ds.iterator()
+    captured = []
+    orig = _StatsManager.update_iteration_metrics
+
+    def spy(stats, dataset_tag):
+        captured.append(stats)
+        return orig(stats, dataset_tag)
+
+    patches.append(patch.object(_StatsManager, "update_iteration_metrics", spy))
+
+    import contextlib
+
+    with contextlib.ExitStack() as stack:
+        for p in patches:
+            stack.enter_context(p)
+        # Use _iter_batches (private) so we can pass _collate_fn / _finalize_fn
+        # which the public iter_batches signature does not expose.
+        for _ in it._iter_batches(**iter_kwargs):
+            pass
+
+    stats = captured[-1]
+    all_stages = [
+        stats.iter_blocked_production_wait_s.get(),
+        stats.iter_blocked_data_transfer_s.get(),
+        stats.iter_blocked_batching_s.get(),
+        stats.iter_blocked_format_s.get(),
+        stats.iter_blocked_collate_s.get(),
+        stats.iter_blocked_finalize_s.get(),
+    ]
+    bound_value = getattr(stats, bound_stage).get()
+    # The bottleneck stage should be at least the sleep time we injected,
+    # proving the timing capture is actually recording the stall.
+    assert bound_value >= SLEEP_S, (
+        f"{scenario}-bound: {bound_stage}={bound_value} < SLEEP_S={SLEEP_S}; "
+        "timing capture missed the injected stall"
+    )
+    # The bottleneck stage should be strictly greater than all others.
+    for v in all_stages:
+        if v == bound_value:
+            continue
+        assert (
+            bound_value > v
+        ), f"{scenario}-bound: {bound_stage}={bound_value} not > {v}"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/block_batching/test_util.py

@@ -0,0 +1,660 @@
+import logging
+import random
+import sys
+import threading
+import time
+from collections import Counter
+from os import urandom
+from typing import Callable, Iterator
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.block_batching.interfaces import (
+    Batch,
+    BatchMetadata,
+    ResolvedBlock,
+)
+from ray.data._internal.block_batching.util import (
+    _calculate_ref_hits,
+    blocks_to_batches,
+    collate,
+    finalize_batches,
+    format_batches,
+    iter_threaded,
+    resolve_block_refs,
+)
+from ray.data._internal.stats import DatasetStats
+from ray.data._internal.util import make_async_gen
+
+logger = logging.getLogger(__file__)
+
+
+def block_generator(num_rows: int, num_blocks: int):
+    for _ in range(num_blocks):
+        yield pa.table({"foo": [1] * num_rows})
+
+
+def test_resolve_block_refs(ray_start_regular_shared):
+    block_refs = [ray.put(0), ray.put(1), ray.put(2)]
+
+    resolved_iter = resolve_block_refs(iter(block_refs))
+    resolved = list(resolved_iter)
+    assert all(isinstance(b, ResolvedBlock) for b in resolved)
+    assert [b.block for b in resolved] == [0, 1, 2]
+
+
+def test_resolve_block_refs_accumulates_data_transfer_timer(
+    ray_start_regular_shared,
+):
+    """resolve_block_refs accumulates ray.get() time into iter_get_s and
+    captures a per-block data_transfer TimeSpan."""
+    block_refs = [ray.put(i) for i in range(3)]
+
+    stats = DatasetStats(metadata={}, parent=None)
+    resolved = list(resolve_block_refs(iter(block_refs), stats=stats))
+
+    assert len(resolved) == 3
+
+    # data_transfer TimeSpan captured per block.
+    for r in resolved:
+        assert r.stage_timings is not None
+        assert r.stage_timings.data_transfer is not None
+        assert r.stage_timings.data_transfer.duration >= 0.0
+
+
+def test_resolve_block_refs_captures_production_wait_span(
+    ray_start_regular_shared,
+):
+    """resolve_block_refs captures a per-block production_wait TimeSpan
+    around ``next(block_ref_iter)`` (manual capture, no Timer accumulation)."""
+    block_refs = [ray.put(i) for i in range(3)]
+
+    stats = DatasetStats(metadata={}, parent=None)
+    resolved = list(resolve_block_refs(iter(block_refs), stats=stats))
+
+    assert len(resolved) == 3
+    for r in resolved:
+        assert r.stage_timings is not None
+        assert r.stage_timings.production_wait is not None
+        assert r.stage_timings.production_wait.duration >= 0.0
+
+
+@pytest.mark.parametrize("block_size", [1, 10])
+@pytest.mark.parametrize("drop_last", [True, False])
+def test_blocks_to_batches(block_size, drop_last):
+    num_blocks = 5
+    block_iter = block_generator(num_rows=block_size, num_blocks=num_blocks)
+    # Wrap raw blocks in ResolvedBlock (stage_timings=None) as blocks_to_batches now expects
+    wrapped_blocks = (ResolvedBlock(block=b) for b in block_iter)
+
+    batch_size = 3
+    batch_iter = list(
+        blocks_to_batches(wrapped_blocks, batch_size=batch_size, drop_last=drop_last)
+    )
+
+    if drop_last:
+        for batch in batch_iter:
+            assert len(batch.data) == batch_size
+    else:
+        full_batches = 0
+        leftover_batches = 0
+
+        dataset_size = block_size * num_blocks
+        for batch in batch_iter:
+            if len(batch.data) == batch_size:
+                full_batches += 1
+            if len(batch.data) == (dataset_size % batch_size):
+                leftover_batches += 1
+
+        assert leftover_batches == 1
+        assert full_batches == (dataset_size // batch_size)
+
+    assert [batch.metadata.batch_idx for batch in batch_iter] == list(
+        range(len(batch_iter))
+    )
+
+
+@pytest.mark.parametrize("batch_format", ["pandas", "numpy", "pyarrow"])
+def test_format_batches(batch_format):
+    block_iter = block_generator(num_rows=2, num_blocks=2)
+    batch_iter = (
+        Batch(BatchMetadata(batch_idx=i), block) for i, block in enumerate(block_iter)
+    )
+    batch_iter = list(format_batches(batch_iter, batch_format=batch_format))
+
+    for batch in batch_iter:
+        if batch_format == "pandas":
+            assert isinstance(batch.data, pd.DataFrame)
+        elif batch_format == "arrow":
+            assert isinstance(batch.data, pa.Table)
+        elif batch_format == "numpy":
+            assert isinstance(batch.data, dict)
+            assert isinstance(batch.data["foo"], np.ndarray)
+
+    assert [batch.metadata.batch_idx for batch in batch_iter] == list(
+        range(len(batch_iter))
+    )
+
+
+def test_collate():
+    def collate_fn(batch):
+        return pa.table({"bar": [1] * 2})
+
+    batches = [
+        Batch(BatchMetadata(batch_idx=i), data)
+        for i, data in enumerate(block_generator(num_rows=2, num_blocks=2))
+    ]
+    batch_iter = collate(batches, collate_fn=collate_fn)
+
+    for i, batch in enumerate(batch_iter):
+        assert batch.metadata.batch_idx == i
+        assert batch.data == pa.table({"bar": [1] * 2})
+
+
+def test_finalize():
+    def finalize_fn(batch):
+        return pa.table({"bar": [1] * 2})
+
+    batches = [
+        Batch(BatchMetadata(batch_idx=i), data)
+        for i, data in enumerate(block_generator(num_rows=2, num_blocks=2))
+    ]
+    batch_iter = finalize_batches(batches, finalize_fn=finalize_fn)
+
+    for i, batch in enumerate(batch_iter):
+        assert batch.metadata.batch_idx == i
+        assert batch.data == pa.table({"bar": [1] * 2})
+
+
+@pytest.mark.parametrize("preserve_ordering", [True, False])
+@pytest.mark.parametrize("buffer_size", [0, 1, 2])
+def test_make_async_gen_fail(buffer_size: int, preserve_ordering):
+    """Tests that any errors raised in async threads are propagated to the main
+    thread."""
+
+    def gen(base_iterator):
+        raise ValueError("Fail")
+
+    iterator = make_async_gen(
+        base_iterator=iter([1]),
+        fn=gen,
+        preserve_ordering=preserve_ordering,
+        buffer_size=buffer_size,
+    )
+
+    with pytest.raises(ValueError) as e:
+        for _ in iterator:
+            pass
+
+    assert e.match("Fail")
+
+
+@pytest.mark.parametrize("preserve_ordering", [True, False])
+def test_make_async_gen_varying_seq_length_stress_test(preserve_ordering):
+    """This test executes make_async_gen against a function generating variable
+    length sequences to stress test its concurrency control.
+    """
+
+    num_workers = 4
+
+    c = 0
+
+    # Roll the dice 100 times
+    for i in range(100):
+        # Fetch 8b seed from urandom
+        seed = int.from_bytes(urandom(8), byteorder=sys.byteorder)
+        r = random.Random(seed)
+
+        print(f">>> Seed: {seed}")
+
+        # NOTE: Number of seqs >> number of workers
+        #       to saturate the input queue
+        num_seqs = num_workers * 10
+
+        lens = list(range(num_seqs))
+
+        r.shuffle(lens)
+
+        source = [range(len_) for len_ in lens]
+
+        print("===" * 8)
+        print(source)
+        print("===" * 8)
+
+        def flatten(list_iter):
+            for l in list_iter:
+                print(f">>> Flattening: {l}")
+                yield from l
+
+        it = make_async_gen(
+            iter(source),
+            flatten,
+            preserve_ordering=preserve_ordering,
+            num_workers=4,
+            buffer_size=1,
+        )
+
+        total = 0
+
+        for i in it:
+            total += i
+
+        assert total == 9880
+        c += 1
+
+    assert c == 100
+
+
+@pytest.mark.parametrize("preserve_ordering", [True, False])
+def test_make_async_gen_non_reentrant(preserve_ordering):
+    """This test is asserting that make_async_gen iterating over the
+    sequence as a whole and not re-entering provided transformation,
+    as this might have substantial performance impact in extreme case
+    of re-entering for every element of the sequence
+    """
+
+    logs = []
+    finished = False
+
+    def _transform_inner(it):
+        nonlocal finished
+
+        assert not finished
+
+        logs.append(">>> Entering Inner")
+
+        for i in it:
+            logs.append(f">>> Inner: {i}")
+            yield i
+
+        logs.append(">>> Leaving Inner")
+
+        # Once this transform finishes
+        finished = True
+
+    def _transform_b(it):
+        logs.append(">>> Entering Outer")
+
+        for i in _transform_inner(it):
+            logs.append(f">>> Outer: {i}")
+            yield i
+
+        logs.append(">>> Leaving Outer")
+
+    for _ in make_async_gen(
+        iter(range(3)),
+        _transform_b,
+        preserve_ordering=preserve_ordering,
+    ):
+        pass
+
+    assert [
+        ">>> Entering Outer",
+        ">>> Entering Inner",
+        ">>> Inner: 0",
+        ">>> Outer: 0",
+        ">>> Inner: 1",
+        ">>> Outer: 1",
+        ">>> Inner: 2",
+        ">>> Outer: 2",
+        ">>> Leaving Inner",
+        ">>> Leaving Outer",
+    ] == logs
+
+
+@pytest.mark.parametrize("preserve_ordering", [True, False])
+@pytest.mark.parametrize(
+    "buffer_size, expected_gen_time",
+    [
+        (0, 5.5),  # 5 x 1s + 0.5s buffer
+        (1, 7.5),  # 3 x 1s + 2 x 2s (limited buffer delay) + 0.5s buffer
+        (2, 5.5),  # 5 x 1s + 0.5s buffer
+    ],
+)
+def test_make_async_gen_x(buffer_size: int, expected_gen_time, preserve_ordering):
+    """Tests that make_async_gen overlaps compute."""
+
+    num_items = 5
+
+    def gen(base_iterator):
+        gen_start = time.perf_counter()
+
+        for i in base_iterator:
+            time.sleep(1)
+            yield i
+            print(f">>> ({time.time()}) Generating {i}")
+
+        gen_finish = time.perf_counter()
+
+        # 0.5s buffer
+        assert gen_finish - gen_start < expected_gen_time
+
+    def sleepy_udf(item):
+        time.sleep(2)
+        return item
+
+    iterator = make_async_gen(
+        base_iterator=iter(range(num_items)),
+        fn=gen,
+        preserve_ordering=preserve_ordering,
+        num_workers=1,
+        buffer_size=buffer_size,
+    )
+
+    outputs = []
+
+    iter_start = time.perf_counter()
+    for item in iterator:
+        print(f">>> ({time.time()}) Iterating over {item}")
+        print(item)
+        outputs.append(sleepy_udf(item))
+    iter_finish = time.perf_counter()
+
+    dur_s = iter_finish - iter_start
+
+    print(f">>> Took {dur_s}")
+
+    # 1s to yield first element
+    # 10s to iterate t/h all 5
+    # 0.5s extra buffer
+    assert dur_s < num_items * 2 + 1.5
+
+    # Assert ordering is preserved
+    assert outputs == list(range(num_items))
+
+
+@pytest.mark.parametrize("preserve_ordering", [True, False])
+@pytest.mark.parametrize("buffer_size", [0, 1, 2])
+def test_make_async_gen_multiple_threads(buffer_size: int, preserve_ordering):
+    """Tests that using multiple threads can overlap compute even more."""
+
+    num_items = 5
+
+    gen_sleep = 2
+    iter_sleep = 3
+
+    def gen(base_iterator):
+        for i in base_iterator:
+            time.sleep(gen_sleep)
+            yield i
+
+    def sleep_udf(item):
+        time.sleep(iter_sleep)
+        return item
+
+    # All 5 items should be fetched concurrently.
+    iterator = make_async_gen(
+        base_iterator=iter(range(num_items)),
+        fn=gen,
+        preserve_ordering=preserve_ordering,
+        num_workers=5,
+        buffer_size=buffer_size,
+    )
+
+    start_time = time.time()
+
+    # Only sleep for first item.
+    elements = [sleep_udf(next(iterator))] + list(iterator)
+
+    # All subsequent items should already be prefetched and should be ready.
+    end_time = time.time()
+
+    # Assert ordering is preserved
+    if preserve_ordering:
+        assert elements == list(range(num_items))
+
+    # - 2 second for every worker to handle their single element
+    # - 3 seconds for overlapping one
+    # - 0.5 seconds buffer
+    assert end_time - start_time < gen_sleep + iter_sleep + 0.5
+
+
+@pytest.mark.parametrize("preserve_ordering", [True, False])
+@pytest.mark.parametrize("buffer_size", [0, 1, 2])
+def test_make_async_gen_multiple_threads_unfinished(
+    buffer_size: int, preserve_ordering
+):
+    """Tests that using multiple threads can overlap compute even more.
+    Do not finish iteration with break in the middle.
+    """
+
+    num_items = 5
+
+    def gen(base_iterator):
+        for i in base_iterator:
+            time.sleep(4)
+            yield i
+
+    def sleep_udf(item):
+        time.sleep(5)
+        return item
+
+    # All 5 items should be fetched concurrently.
+    iterator = make_async_gen(
+        base_iterator=iter(range(num_items)),
+        fn=gen,
+        preserve_ordering=preserve_ordering,
+        num_workers=5,
+        buffer_size=buffer_size,
+    )
+
+    start_time = time.time()
+
+    # Only sleep for first item.
+    sleep_udf(next(iterator))
+
+    # All subsequent items should already be prefetched and should be ready.
+    for i, _ in enumerate(iterator):
+        if i > 2:
+            break
+    end_time = time.time()
+
+    # 4 second for first item, 5 seconds for udf, 0.5 seconds buffer
+    assert end_time - start_time < 9.5
+
+
+def test_calculate_ref_hits(ray_start_regular_shared):
+    refs = [ray.put(0), ray.put(1)]
+    hits, misses, unknowns = _calculate_ref_hits(refs)
+    # With ctx.enable_get_object_locations_for_metrics set to False
+    # by default, `_calculate_ref_hits` returns -1 for all, since
+    # getting object locations is disabled.
+    assert hits == 0
+    assert misses == 0
+    assert unknowns == 0
+
+    ctx = ray.data.DataContext.get_current()
+    prev_enable_get_object_locations_for_metrics = (
+        ctx.enable_get_object_locations_for_metrics
+    )
+    try:
+        ctx.enable_get_object_locations_for_metrics = True
+        hits, misses, unknowns = _calculate_ref_hits(refs)
+        assert hits == 2
+        assert misses == 0
+        assert unknowns == 0
+    finally:
+        ctx.enable_get_object_locations_for_metrics = (
+            prev_enable_get_object_locations_for_metrics
+        )
+
+
+def _identity(it: Iterator[int]) -> Iterator[int]:
+    return it
+
+
+def _duplicate_each(it: Iterator[int]) -> Iterator[int]:
+    for item in it:
+        yield item
+        yield item
+
+
+class TestIterThreaded:
+    """Unit tests for ``iter_threaded``."""
+
+    @pytest.mark.parametrize("num_workers", [1, 2, 4])
+    @pytest.mark.parametrize("output_buffer_size", [1, 2, 4])
+    @pytest.mark.parametrize(
+        "fn,multiplier",
+        [(_identity, 1), (_duplicate_each, 2)],
+        ids=["identity", "duplicate"],
+    )
+    def test_processes_all_exactly_once(
+        self,
+        num_workers: int,
+        output_buffer_size: int,
+        fn: Callable[[Iterator[int]], Iterator[int]],
+        multiplier: int,
+    ):
+        """Every input item is consumed and produced exactly the expected
+        number of times across the worker pool (no losses, no duplicates).
+        Output ordering is not required."""
+        items = list(range(50))
+        output = list(
+            iter_threaded(
+                iter(items),
+                fn,
+                num_workers=num_workers,
+                output_buffer_size=output_buffer_size,
+            )
+        )
+        assert len(output) == len(items) * multiplier
+        assert Counter(output) == Counter(items * multiplier)
+
+    def test_stateful_base_iterator_thread_safe(self):
+        """Python generators are not thread-safe; concurrent ``next()``
+        calls raise ``ValueError: generator already executing`` without
+        a lock. This test passes only if ``iter_threaded`` serializes
+        the underlying ``next()`` properly."""
+
+        def stateful_gen():
+            for i in range(200):
+                # Encourage interleaving across workers.
+                time.sleep(0.001)
+                yield i
+
+        output = list(iter_threaded(stateful_gen(), _identity, num_workers=4))
+        assert sorted(output) == list(range(200))
+
+    @pytest.mark.parametrize("num_workers", [1, 4])
+    def test_fn_exception_propagates(self, num_workers: int):
+        """An exception raised inside ``fn`` is surfaced to the consumer
+        rather than silently swallowed or hanging the iterator."""
+
+        def fn(it: Iterator[int]) -> Iterator[int]:
+            for i, item in enumerate(it):
+                if i >= 3:
+                    raise ValueError("boom")
+                yield item
+
+        it = iter_threaded(iter(range(100)), fn, num_workers=num_workers)
+        with pytest.raises(ValueError, match="boom"):
+            list(it)
+
+    @pytest.mark.parametrize("num_workers", [1, 4])
+    def test_non_generator_fn_construction_raises(self, num_workers: int):
+        """When ``fn`` is a non-generator function that raises during
+        construction (e.g., setup code before returning the iterator), the
+        exception must surface to the consumer rather than hang. Regression
+        for the case where ``fn(_locked_iter())`` was called outside the
+        worker's try/finally."""
+
+        def fn(it: Iterator[int]) -> Iterator[int]:
+            # Body runs eagerly at call time (not a generator function).
+            raise ValueError("boom in fn construction")
+
+        it = iter_threaded(iter(range(100)), fn, num_workers=num_workers)
+        with pytest.raises(ValueError, match="boom in fn construction"):
+            list(it)
+
+    @pytest.mark.parametrize("num_workers", [1, 4])
+    def test_consumer_break_stops_workers(self, num_workers: int):
+        """When the consumer breaks early and the iterator is no longer
+        referenced, CPython GCs the generator immediately, which runs the
+        ``finally: stopped.set()`` cleanup path. Worker threads should
+        terminate within the ``_put`` poll interval (~100ms) rather than
+        leak."""
+
+        def slow_fn(it: Iterator[int]) -> Iterator[int]:
+            for item in it:
+                time.sleep(0.05)
+                yield item
+
+        # Inline so `break` drops the last reference → GC → finally.
+        for i, _ in enumerate(
+            iter_threaded(iter(range(10_000)), slow_fn, num_workers=num_workers)
+        ):
+            if i >= 5:
+                break
+
+        # Workers poll `stopped` every 100ms inside `_put`; give a generous
+        # margin for CI under load.
+        deadline = time.time() + 5.0
+        while time.time() < deadline:
+            alive = [t for t in threading.enumerate() if t.name == "iter_threaded"]
+            if not alive:
+                break
+            time.sleep(0.05)
+        else:
+            pytest.fail(
+                f"iter_threaded workers did not exit within 5s: "
+                f"{[t.name for t in threading.enumerate() if t.name == 'iter_threaded']}"
+            )
+
+    def test_num_workers_validation(self):
+        with pytest.raises(ValueError, match="num_workers must be at least 1"):
+            list(iter_threaded(iter([1]), _identity, num_workers=0))
+
+    def test_output_buffer_size_validation(self):
+        with pytest.raises(ValueError, match="output_buffer_size must be at least 1"):
+            list(iter_threaded(iter([1]), _identity, output_buffer_size=0))
+
+    def test_empty_base_iterator(self):
+        output = list(iter_threaded(iter([]), _identity, num_workers=4))
+        assert output == []
+
+    @pytest.mark.parametrize("num_workers,output_buffer_size", [(1, 1), (2, 2), (4, 2)])
+    def test_in_flight_items_bounded_by_output_buffer_size(
+        self, num_workers: int, output_buffer_size: int
+    ):
+        """Without consumption, workers must not pull more than
+        ``output_buffer_size`` items from the base iterator. Pulled-but-not-
+        consumed items are 'in flight', and the bound caps them."""
+        pulled = 0
+        pulled_lock = threading.Lock()
+
+        def counting_iter() -> Iterator[int]:
+            nonlocal pulled
+            for i in range(1_000_000):
+                with pulled_lock:
+                    pulled += 1
+                yield i
+
+        it = iter_threaded(
+            counting_iter(),
+            _identity,
+            num_workers=num_workers,
+            output_buffer_size=output_buffer_size,
+        )
+
+        # Trigger the generator body (which starts the workers), then stop
+        # consuming. Workers fill in-flight up to the bound and then block
+        # on _acquire_slot.
+        next(it)
+        time.sleep(0.3)
+
+        with pulled_lock:
+            # Consumer took 1 → in-flight ≤ K. Plus the 1 already consumed.
+            assert (
+                pulled <= 1 + output_buffer_size
+            ), f"Pulled {pulled}, expected <= {1 + output_buffer_size}"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))