chore: import upstream snapshot with attribution

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__]))