diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 868a18e22..6be8e1472 100644 --- a/python/python/lancedb/streaming.py +++ b/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)