chore: import upstream snapshot with attribution

python/ray/data/_internal/block_batching/block_batching.py

@@ -0,0 +1,101 @@
+from typing import Callable, Iterator, Optional, TypeVar
+
+from ray.data._internal.block_batching.interfaces import ResolvedBlock
+from ray.data._internal.block_batching.util import (
+    _MappingIterator,
+    blocks_to_batches,
+    collate,
+    format_batches,
+)
+from ray.data._internal.stats import DatasetStats
+from ray.data.block import Block, DataBatch
+
+T = TypeVar("T")
+
+
+def batch_blocks(
+    blocks: Iterator[Block],
+    *,
+    stats: Optional[DatasetStats] = None,
+    batch_size: Optional[int] = None,
+    batch_format: str = "default",
+    drop_last: bool = False,
+    collate_fn: Optional[Callable[[DataBatch], DataBatch]] = None,
+    shuffle_buffer_min_size: Optional[int] = None,
+    shuffle_seed: Optional[int] = None,
+    ensure_copy: bool = False,
+) -> Iterator[DataBatch]:
+    """Create formatted batches of data from 1 or more blocks.
+
+    This function takes in an iterator of already fetched blocks. Consequently, this
+    function doesn't support block prefetching.
+    """
+    # TODO: make stage timings optional at _BatchingIterator so this
+    # shim can be removed. map() avoids holding block references.
+    wrapped_blocks = map(lambda b: ResolvedBlock(block=b), blocks)
+
+    # Build the processing pipeline
+    batch_iter = format_batches(
+        blocks_to_batches(
+            block_iter=wrapped_blocks,
+            stats=stats,
+            batch_size=batch_size,
+            drop_last=drop_last,
+            shuffle_buffer_min_size=shuffle_buffer_min_size,
+            shuffle_seed=shuffle_seed,
+            ensure_copy=ensure_copy,
+        ),
+        batch_format=batch_format,
+        stats=stats,
+        ensure_copy=ensure_copy,
+    )
+
+    if collate_fn is not None:
+        batch_iter = collate(batch_iter, collate_fn=collate_fn, stats=stats)
+
+    return _UserTimingIterator(
+        _MappingIterator(batch_iter, lambda batch: batch.data), stats
+    )
+
+
+class _UserTimingIterator(Iterator[DataBatch]):
+    def __init__(self, iter: Iterator[DataBatch], stats: Optional[DatasetStats]):
+        self._iter = iter
+        self._stats = stats
+        self._active_timer = None
+
+    def __iter__(self) -> Iterator[DataBatch]:
+        return self
+
+    def __next__(self) -> DataBatch:
+        # Since we're tracking time spent in user-code, we stop
+        # the timer immediately when `__next__` is called
+        self._stop_timer()
+
+        try:
+            res = next(self._iter)
+            # Reset timer and return
+            #
+            # NOTE: It's crucial that we reset the timer only after we
+            #       retrieved the result to avoid starting the timer before
+            #       we retrieve the next value
+            self._reset_timer()
+            return res
+        except StopIteration:
+            self._stop_timer()
+            raise
+
+    def _stop_timer(self):
+        if not self._stats:
+            return
+
+        if self._active_timer:
+            self._active_timer.__exit__(None, None, None)
+            self._active_timer = None
+
+    def _reset_timer(self):
+        if not self._stats:
+            return
+
+        self._active_timer = self._stats.iter_user_s.timer()
+        self._active_timer.__enter__()