a

python/python/lancedb/streaming.py

@@ -34,6 +34,7 @@ from copy import deepcopy
 from multiprocessing import RawArray
 from typing import Any, Callable, cast, Iterator, Literal, NamedTuple, Optional, Union
 
+import numpy as np
 import pyarrow as pa
 import pyarrow.compute as pc
 import torch
@@ -55,18 +56,177 @@ logger = logging.getLogger(__name__)
 # distinct seeds for any practically encountered epoch count.
 _EPOCH_PRIME = 100003
 
-# Multiplier used to give each split its own in-RAM row permutation seed in
-# 2-phase mode (distinct from _EPOCH_PRIME so the two offsets can't collide).
-_SPLIT_PRIME = 1000003
-
-# Number of blocks kept in RAM per split at once for 2-phase's row-shuffle
-# window.  Fixed for now; will become a real sliding window later.
-_TWO_PHASE_WINDOW_BLOCKS = 3
+# Offsets added when seeding each split's 2-phase bounded shuffle
+# permutation (distinct from _EPOCH_PRIME and from each other so they
+# can't collide).
+_SIGMA_PRIME = 1000003
+_SIGMA_SPLIT_PRIME = 3000017
 
 DEFAULT_READ_BATCH_SIZE = 64
 DEFAULT_PREFETCH_BATCHES = 4
 
 
+def _build_bounded_permutation(n: int, max_distance: int, seed: int) -> np.ndarray:
+    """Build a permutation of ``[0, n)`` with ``|sigma(k) - k| < max_distance``.
+
+    Partitions ``[0, n)`` into consecutive chunks of ``max_distance`` entries
+    (the last chunk may be shorter) and fully shuffles each chunk
+    independently.  Since ``k`` and ``sigma(k)`` always fall in the same
+    chunk, this trivially satisfies the displacement bound.  Vectorized
+    (one random key matrix, sorted per row) so it stays cheap even sized to
+    an entire split's live-row count.
+    """
+    if n == 0:
+        return np.array([], dtype=np.int64)
+    rng = np.random.default_rng(seed)
+    sigma = np.arange(n, dtype=np.int64)
+    num_full = n // max_distance
+    full_len = num_full * max_distance
+    if num_full:
+        # Sorting i.i.d. random keys within each row of a
+        # (num_full, max_distance) matrix gives, per row, a uniformly
+        # random permutation of that row's column indices -- i.e. one
+        # independent shuffled chunk per row, all at once.
+        keys = rng.random((num_full, max_distance))
+        chunk_perms = np.argsort(keys, axis=1)
+        bases = (np.arange(num_full, dtype=np.int64) * max_distance)[:, None]
+        sigma[:full_len] = (chunk_perms + bases).reshape(-1)
+    if full_len < n:
+        tail = np.arange(full_len, n, dtype=np.int64)
+        tail_perm = np.argsort(rng.random(n - full_len))
+        sigma[full_len:] = tail[tail_perm]
+    return sigma
+
+
+class _TwoPhaseSplitReader:
+    """Sequential-only stand-in for ``Permutation`` driving 2-phase reads.
+
+    Duck-types ``Permutation.__getitems__`` so 2-phase's shuffled reads flow
+    through the same stage-1 I/O machinery (``io_queue_depth`` prefetch,
+    ``Future``-based fetches) used by 1-phase, instead of a separate ad hoc
+    path.
+
+    This split's assigned blocks, concatenated in order and restricted to
+    live rows, give every row a dense natural rank in ``[0, num_rows)``.  A
+    single permutation ``sigma`` of that rank space -- built once, with
+    ``|sigma(rank) - rank| < max_shuffle_distance`` for every rank -- decides
+    each row's output position.  A row may move anywhere within that bound,
+    including into a different block's range entirely; the only restriction
+    is the distance from its own natural rank, not which block it started
+    in.  A block is evicted, and the next block assigned to this split read
+    in to replace it, as soon as every one of its rows has been read past
+    in the output order -- which happens naturally once ``sigma`` has been
+    applied, without any extra bookkeeping about "generations" or "slots".
+
+    Only one ``__getitems__`` call executes at a time (serialized by
+    ``_lock``): the block cache and read plan are shared, mutable state,
+    and letting ``io_queue_depth`` submit multiple concurrent calls for the
+    same split would race on them.  ``io_queue_depth`` therefore does not
+    add fetch parallelism within a single 2-phase split (a known
+    limitation) -- only across splits, which each get their own reader and
+    lock.  Blocks are also still loaded strictly on demand (the first time
+    a gather references them), not eagerly prefetched ahead of
+    ``window_blocks``, so a single split's own I/O does not yet overlap
+    with itself either -- a separate follow-up from this pass.
+    """
+
+    def __init__(
+        self,
+        block_ids: list[int],
+        block_live_counts: list[int],
+        shuffle: bool,
+        max_shuffle_distance: int,
+        seed: int,
+        read_block_fn: Callable[[int], tuple[pa.Table, list[int]]],
+        columns: Optional[list[str]],
+    ):
+        self._block_ids = block_ids
+        self._read_block_fn = read_block_fn
+        self._columns = columns
+        self._lock = threading.Lock()
+        self._loaded: dict[int, pa.Table] = {}
+
+        counts = np.asarray(block_live_counts, dtype=np.int64)
+        self.num_rows = int(counts.sum())
+        num_blocks = len(block_ids)
+
+        # Natural (block, then row) order: block_pos[k]/local_offset[k]
+        # describe the row whose dense natural rank is exactly k.
+        block_pos = np.repeat(np.arange(num_blocks, dtype=np.int64), counts)
+        local_offset = (
+            np.concatenate([np.arange(c, dtype=np.int64) for c in counts])
+            if num_blocks
+            else np.array([], dtype=np.int64)
+        )
+
+        if shuffle and self.num_rows:
+            sigma = _build_bounded_permutation(
+                self.num_rows, max_shuffle_distance, seed
+            )
+        else:
+            sigma = np.arange(self.num_rows, dtype=np.int64)
+
+        # sigma is itself a permutation of dense ranks, so inverting it
+        # (argsort) directly gives, for each output position, which
+        # natural rank -- and thus which (block, local offset) -- to draw.
+        order = np.argsort(sigma, kind="stable")
+        self._plan_block_pos = block_pos[order]
+        self._plan_local_offset = local_offset[order]
+
+        # Last plan position (inclusive) at which each block is still
+        # needed, so it can be evicted as soon as we've read past it.
+        self._block_last_pos: dict[int, int] = {}
+        for pos, bp in enumerate(self._plan_block_pos.tolist()):
+            self._block_last_pos[bp] = pos
+
+    def __getitems__(self, indices) -> pa.RecordBatch:
+        start = indices[0]
+        fetch = len(indices)
+        with self._lock:
+            block_pos_slice = self._plan_block_pos[start : start + fetch]
+            local_off_slice = self._plan_local_offset[start : start + fetch]
+            needed = sorted(set(block_pos_slice.tolist()))
+            for bp in needed:
+                if bp not in self._loaded:
+                    block_id = self._block_ids[bp]
+                    raw, live_positions = self._read_block_fn(block_id)
+                    self._loaded[bp] = raw.take(
+                        pa.array(live_positions, type=pa.int64())
+                    )
+
+            # Group by source block (stable, so each block's rows keep
+            # their relative order), gather from each, then invert the
+            # grouping to restore this chunk's actual priority order.
+            group_order = np.argsort(block_pos_slice, kind="stable")
+            grouped_block_pos = block_pos_slice[group_order]
+            grouped_local_off = local_off_slice[group_order]
+            tables = []
+            pos = 0
+            n = len(grouped_block_pos)
+            while pos < n:
+                bp = grouped_block_pos[pos]
+                run_end = pos + 1
+                while run_end < n and grouped_block_pos[run_end] == bp:
+                    run_end += 1
+                offs = grouped_local_off[pos:run_end]
+                tables.append(
+                    self._loaded[int(bp)].take(pa.array(offs, type=pa.int64()))
+                )
+                pos = run_end
+            combined = pa.concat_tables(tables)
+            inverse = np.argsort(group_order)
+            result = combined.take(pa.array(inverse, type=pa.int64()))
+
+            watermark = start + fetch - 1
+            for bp in list(self._loaded):
+                if self._block_last_pos.get(bp, -1) <= watermark:
+                    del self._loaded[bp]
+
+        if self._columns is not None:
+            result = result.select(self._columns)
+        return result.combine_chunks().to_batches()[0]
+
+
 class _WorkerSample(NamedTuple):
     data: Any
     dataset: "StreamingDataset"
@@ -283,7 +443,28 @@ class StreamingDataset(IterableDataset):
         When set, rows are shuffled in contiguous groups of this size rather
         than individually.  Larger clumps improve I/O locality (important on
         object storage) at the cost of reduced randomness.  ``None`` (the
-        default) shuffles rows individually.
+        default) shuffles rows individually.  Ignored when ``block_size`` is
+        set (2-phase mode has its own locality/randomness knobs, below).
+    block_size:
+        Enables 2-phase shuffled reads: rows are shuffled at the granularity
+        of contiguous blocks of (approximately, or exactly when ``filter``
+        is unset) this many rows, and each split reads a rolling window of
+        ``window_blocks`` blocks at a time instead of taking individual rows
+        by row-id.  This trades some randomness for I/O that is almost
+        entirely contiguous range scans rather than row-id takes, which is
+        far cheaper on object storage.  ``None`` (the default) uses 1-phase
+        shuffling, which permutes every row individually up front.
+    window_blocks:
+        Number of blocks kept in RAM per split at once in 2-phase mode.
+        Larger windows shuffle across more rows at the cost of more memory;
+        must be at least 1.  Requires ``block_size``.  Defaults to 4.
+    max_shuffle_distance:
+        Bounds how far, in rows, 2-phase shuffling may displace a row from
+        its position in unshuffled (block, then row) order.  Requires
+        ``block_size``.  Smaller values shuffle less but let a block's RAM
+        be reclaimed and refilled with the next block sooner; larger values
+        shuffle more but hold blocks in RAM longer.  Defaults to
+        ``block_size // 4`` (at least 1).
     filter:
         Optional SQL filter expression (e.g. ``"label = 'dog'"``).  Only rows
         that satisfy the predicate are included in the permutation.  The filter
@@ -407,6 +588,8 @@ class StreamingDataset(IterableDataset):
         columns: Optional[list[str]] = None,
         shuffle_clump_size: Optional[int] = None,
         block_size: Optional[int] = None,
+        window_blocks: Optional[int] = None,
+        max_shuffle_distance: Optional[int] = None,
         filter: Optional[str] = None,
         transform: Optional[Callable] = None,
         transform_parallelism: Optional[int] = None,
@@ -441,6 +624,18 @@ class StreamingDataset(IterableDataset):
             raise ValueError("io_queue_depth must be greater than 0")
         if block_size is not None and block_size <= 1:
             raise ValueError("block_size must be greater than 1")
+        if block_size is None:
+            if window_blocks is not None:
+                raise ValueError("window_blocks requires block_size to be set")
+            if max_shuffle_distance is not None:
+                raise ValueError(
+                    "max_shuffle_distance requires block_size to be set"
+                )
+        else:
+            if window_blocks is not None and window_blocks < 1:
+                raise ValueError("window_blocks must be at least 1")
+            if max_shuffle_distance is not None and max_shuffle_distance <= 0:
+                raise ValueError("max_shuffle_distance must be greater than 0")
         if transform_parallelism is not None and transform_parallelism <= 0:
             raise ValueError("transform_parallelism must be greater than 0")
         if pack_sequences is not None:
@@ -606,6 +801,7 @@ class StreamingDataset(IterableDataset):
                     (i * block_size, min((i + 1) * block_size, num_rows))
                     for i in range(num_blocks)
                 ]
+                block_live_counts = [end - start for start, end in block_ranges]
             else:
                 # One pass, computed once here and never repeated at read
                 # time: every fragment's raw and live row counts are read
@@ -613,30 +809,37 @@ class StreamingDataset(IterableDataset):
                 # data once) to lay out block boundaries on fragment edges,
                 # so later block reads are plain offset/limit scans that
                 # Lance can seek to directly, instead of re-scanning from
-                # the start of the table on every read.
+                # the start of the table on every read.  The exact live
+                # count of each block is also recorded here, from the same
+                # pass, so read time never has to guess how many live rows
+                # a block will actually yield.
                 fragments = table.to_lance().get_fragments()
                 frag_raw_lens = [frag.count_rows() for frag in fragments]
                 frag_live_lens = [frag.count_rows(filter) for frag in fragments]
                 num_rows = sum(frag_live_lens)
 
                 block_ranges = []
+                block_live_counts = []
                 raw_pos = 0
                 block_start = 0
                 block_live = 0
                 for raw_len, live_len in zip(frag_raw_lens, frag_live_lens):
                     if block_live > 0 and block_live + live_len > block_size:
                         block_ranges.append((block_start, raw_pos))
+                        block_live_counts.append(block_live)
                         block_start = raw_pos
                         block_live = 0
                     block_live += live_len
                     raw_pos += raw_len
                 if block_live > 0:
                     block_ranges.append((block_start, raw_pos))
+                    block_live_counts.append(block_live)
                 num_blocks = len(block_ranges)
 
             self._num_rows = num_rows
             self._num_blocks = num_blocks
             self._block_ranges: list[tuple[int, int]] = block_ranges
+            self._block_live_counts: list[int] = block_live_counts
             if self._num_blocks < num_splits:
                 raise ValueError(
                     f"block_size={block_size} yields only {self._num_blocks} "
@@ -652,6 +855,18 @@ class StreamingDataset(IterableDataset):
             # Arrow-backed permutation table like the 1-phase row mapping.
             self._block_perm: list[int] = block_order
             self._perm_table = None
+
+            # window_blocks currently only sets an expectation for how many
+            # blocks are resident at once; peak residency is actually
+            # driven by max_shuffle_distance relative to block_size (see
+            # _TwoPhaseSplitReader), a known gap between this knob and
+            # actual behavior.
+            self._window_blocks = window_blocks if window_blocks is not None else 4
+            self._max_shuffle_distance = (
+                max_shuffle_distance
+                if max_shuffle_distance is not None
+                else max(1, block_size // 4)
+            )
         else: #1-phase shuffled read
             # Build the permutation table once, deterministically.
             builder = permutation_builder(table)
@@ -851,13 +1066,13 @@ class StreamingDataset(IterableDataset):
             # self._block_perm[k * (b // num_splits) : (k + 1) * (b // num_splits)]
             # -- the surplus b % num_splits blocks are dropped, mirroring how
             # 1-phase drops surplus rows to keep all splits the same length.
+            #
+            # Each split's blocks are wrapped in a _TwoPhaseSplitReader,
+            # which duck-types Permutation.__getitems__ -- so from here on,
+            # 2-phase flows through the exact same stage-1 I/O machinery
+            # (_fill_io/_submit_io/_drain_io, io_queue_depth prefetch) as
+            # 1-phase.  No 2-phase-specific code exists past this point.
             blocks_per_split = self._num_blocks // self._num_splits
-            split_blocks: list[list[int]] = []
-            for split_idx in my_splits:
-                start = split_idx * blocks_per_split
-                end = start + blocks_per_split
-                split_blocks.append(self._block_perm[start:end])
-
             lance_ds = self._table.to_lance()
             from . import connect as _connect
 
@@ -866,14 +1081,13 @@ class StreamingDataset(IterableDataset):
                 # scan that Lance can seek to directly, and it pulls in
                 # dead rows right along with live ones -- that's fine, they
                 # never get selected below and cost nothing extra to fetch.
+                # Bytes/time accounting happens generically in _io_call,
+                # below, not here.
                 row_start, row_end = self._block_ranges[block_id]
-                t0 = time.perf_counter()
                 raw = lance_ds.scanner(
                     offset=row_start,
                     limit=row_end - row_start,
                 ).to_table()
-                self._bytes_loaded += raw.nbytes
-                self._fetch_time += time.perf_counter() - t0
 
                 if self._filter is None:
                     return raw, list(range(raw.num_rows))
@@ -899,66 +1113,38 @@ class StreamingDataset(IterableDataset):
                 )
                 return raw, live_positions
 
-            def _fetch_and_permute_window(split_pos: int) -> pa.Table:
-                # Read this split's first _TWO_PHASE_WINDOW_BLOCKS blocks as
-                # contiguous range reads (each is one object-storage scan,
-                # not a row-id take), concatenate them into one in-RAM
-                # buffer.  Dead rows ride along in that buffer as dead
-                # weight -- they are never removed from it -- and the
-                # permutation below is built only from live positions, so
-                # dead rows simply never get selected by the final take().
-                block_ids = split_blocks[split_pos][:_TWO_PHASE_WINDOW_BLOCKS]
-                raw_tables = []
-                live_positions: list[int] = []
-                raw_offset = 0
-                for block_id in block_ids:
-                    raw, block_live = _read_block(block_id)
-                    live_positions.extend(p + raw_offset for p in block_live)
-                    raw_offset += raw.num_rows
-                    raw_tables.append(raw)
-                window = pa.concat_tables(raw_tables)
-
-                if self._shuffle:
-                    split_idx = my_splits[split_pos]
-                    row_seed = (
-                        self._shuffle_seed
-                        + self._epoch * _EPOCH_PRIME
-                        + (split_idx + 1) * _SPLIT_PRIME
-                    )
-                    random.Random(row_seed).shuffle(live_positions)
-
-                live = window.take(live_positions)
-                if self._columns is not None:
-                    live = live.select(self._columns)
-                return live
-
-            with ThreadPoolExecutor(max_workers=len(my_splits)) as window_pool:
-                windows = list(
-                    window_pool.map(_fetch_and_permute_window, range(len(my_splits)))
+            permutations: list[_TwoPhaseSplitReader] = []
+            initial_samples = []
+            initial_positions = []
+            for split_idx in my_splits:
+                start = split_idx * blocks_per_split
+                end = start + blocks_per_split
+                block_ids = self._block_perm[start:end]
+                live_counts = [self._block_live_counts[bid] for bid in block_ids]
+                sigma_seed = (
+                    self._shuffle_seed
+                    + self._epoch * _EPOCH_PRIME
+                    + _SIGMA_PRIME
+                    + split_idx * _SIGMA_SPLIT_PRIME
                 )
+                permutations.append(
+                    _TwoPhaseSplitReader(
+                        block_ids=block_ids,
+                        block_live_counts=live_counts,
+                        shuffle=self._shuffle,
+                        max_shuffle_distance=self._max_shuffle_distance,
+                        seed=sigma_seed,
+                        read_block_fn=_read_block,
+                        columns=self._columns,
+                    )
+                )
+                # No resume/skip support yet in 2-phase: every split starts
+                # fresh regardless of any loaded checkpoint.
+                initial_samples.append(0)
+                initial_positions.append(0)
 
-            # No resume/skip support yet in 2-phase: every split starts fresh.
-            n = len(my_splits)
-            initial_samples = [0] * n
-            initial_positions = [0] * n
-            # NOTE: for now a split's "size" is just its first window (until
-            # the sliding window is implemented), so a split is exhausted
-            # after its first _TWO_PHASE_WINDOW_BLOCKS blocks -- epochs are
-            # temporarily truncated to that many blocks per split.
-            split_sizes = [window.num_rows for window in windows]
-
-            # Chop each split's permuted window into read_batch_size chunks,
-            # in the same (abs_start, RecordBatch) shape stage 1 normally
-            # produces, so stage 2 and the main loop's round-robin need no
-            # changes at all.
-            two_phase_initial_batches: list[list[tuple[int, pa.RecordBatch]]] = []
-            for window in windows:
-                pos = 0
-                chunks = []
-                for batch in window.to_batches(max_chunksize=self._read_batch_size):
-                    chunks.append((pos, batch))
-                    pos += batch.num_rows
-                two_phase_initial_batches.append(chunks)
+            n = len(permutations)
+            split_sizes = [reader.num_rows for reader in permutations]
         local_consumed = [0] * n
         # Permutation position each split has consumed through (absolute,
         # i.e. counted from the start of the unskipped split).  Runs ahead of
@@ -1007,16 +1193,6 @@ class StreamingDataset(IterableDataset):
         tx_pending = [deque() for _ in range(n)]  # Future[list[(abs_pos, row)]]
         cooked = [deque() for _ in range(n)]  # (abs_pos, row) ready to yield
 
-        if self._block_size is not None:
-            # 2-phase already fetched+permuted each split's window above;
-            # seed raw_batches with it directly and mark stage 1 "done" for
-            # now (fetch_head == split_sizes) so _fill_io/_submit_io, which
-            # only know how to read via the 1-phase `permutations` list,
-            # are never invoked in this mode.
-            for i in range(n):
-                raw_batches[i].extend(two_phase_initial_batches[i])
-                fetch_head[i] = split_sizes[i]
-
         # Limit simultaneous transforms to transform_workers across all splits.
         tx_semaphore = threading.Semaphore(transform_workers)