From e907ff8c0e7aa9eb90af5bf6d5b8be4ed8d20d72 Mon Sep 17 00:00:00 2001 From: jasonz-lance Date: Thu, 8 Oct 2026 16:35:24 -0700 Subject: [PATCH] perf(python): drop O(max_shuffle_distance) prefetch scan in 2-phase reads The background prefetch added in ecf44fed topped up its queue by walking raw plan positions ahead of the read watermark, bounded by `2 * max_shuffle_distance`, until it had found `prefetch_depth` distinct blocks not already loaded or in flight. That bound is a count of *ranks*, not of blocks, so with the default `max_shuffle_distance` of 16384 the scan window covers 32768 ranks but only ~8 distinct blocks -- fewer than the prefetch depth of 16. The loop could therefore never satisfy its exit condition and ran all 32768 iterations, each one a Python-level scalar extraction out of a numpy array, on *every* `__getitems__` call. Prefetching upcoming blocks regressed 2-phase read throughput by ~5.6x (107 -> 19 MB/s on a local NVMe benchmark). Precompute the distinct blocks in first-needed order once, in __init__, and have prefetch walk that instead: one step per block rather than one per rank, so the scan is bounded by the block count (~24 here) and exits as soon as the queue is full. Blocks whose last plan position is already behind the watermark are skipped -- they can never be needed again. Restores local-NVMe 2-phase read throughput to 78.7 MB/s. Co-Authored-By: Claude Opus 5 --- python/python/lancedb/streaming.py | 43 +++++++++++++++--------------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 930e2063d..0b430c79c 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -153,15 +153,6 @@ class _TwoPhaseSplitReader: self._pending: dict[int, Future] = {} self._prefetch_depth = max(1, prefetch_depth) self._executor = ThreadPoolExecutor(max_workers=self._prefetch_depth) - # How far ahead in the plan to look for new blocks to prefetch. - # _build_bounded_permutation shuffles within fixed, non-overlapping - # chunks of max_shuffle_distance dense ranks, so scanning two such - # chunks ahead is always enough to find prefetch_depth distinct new - # blocks (barring an unusually large prefetch_depth relative to - # block_size); it's just a search bound, not a correctness - # requirement, so an occasional miss just skips prefetching that - # block this round. - self._prefetch_scan_cap = max(1, 2 * max_shuffle_distance) counts = np.asarray(block_live_counts, dtype=np.int64) total_rows = int(counts.sum()) @@ -201,6 +192,16 @@ class _TwoPhaseSplitReader: for pos, bp in enumerate(self._plan_block_pos.tolist()): self._block_last_pos[bp] = pos + # Distinct blocks in the order the plan first needs them. Prefetch + # walks this (one entry per block) instead of rescanning raw plan + # positions, which costs one Python step per *row* scanned and + # dominates read time once max_shuffle_distance is large. + if len(self._plan_block_pos): + _, first_pos = np.unique(self._plan_block_pos, return_index=True) + self._block_visit_order = self._plan_block_pos[np.sort(first_pos)].tolist() + else: + self._block_visit_order = [] + def __getitems__(self, indices) -> pa.RecordBatch: # Fancy-indexed (not sliced) so this also supports non-contiguous # index lists, e.g. the sparse random sampling used by @@ -259,18 +260,18 @@ class _TwoPhaseSplitReader: # the (by then finished) result above instead of blocking. room = self._prefetch_depth - len(self._pending) if room > 0: - scan_pos = watermark + 1 - scan_end = min( - scan_pos + self._prefetch_scan_cap, len(self._plan_block_pos) - ) - while scan_pos < scan_end and room > 0: - bp = int(self._plan_block_pos[scan_pos]) - if bp not in self._loaded and bp not in self._pending: - self._pending[bp] = self._executor.submit( - self._read_block_fn, self._block_ids[bp] - ) - room -= 1 - scan_pos += 1 + for bp in self._block_visit_order: + if room == 0: + break + if bp in self._loaded or bp in self._pending: + continue + if self._block_last_pos.get(bp, -1) <= watermark: + # Already read past: never needed again. + continue + self._pending[bp] = self._executor.submit( + self._read_block_fn, self._block_ids[bp] + ) + room -= 1 if self._columns is not None: result = result.select(self._columns)