From 299274a39ac20323ae399ecacc35486e6a356401 Mon Sep 17 00:00:00 2001 From: jasonz-lance Date: Wed, 30 Sep 2026 11:55:57 -0700 Subject: [PATCH] fix(python): resolve easy 2-phase shuffled-read PR review comments - Push columns/projection down into block reads (scanner/take) instead of selecting after the full row is already fetched and cached. - Fix _TwoPhaseSplitReader.__getitems__ to support arbitrary (non-contiguous) index lists via fancy indexing, not just the contiguous ranges the main pipeline happens to use; enables reuse in _estimate_blocks_per_epoch's sparse sampling. - Extract _read_block/_make_two_phase_reader helpers out of _iter_owned so blocks_per_epoch="auto" can sample through the real 2-phase reader instead of crashing on a null _perm_table. - Tie max_shuffle_distance's default to window_blocks * block_size so the default window actually holds window_blocks blocks in RAM, and reject shuffle_clump_size combined with block_size. - Fix a stale comment implying _block_perm is a list for RAM reasons. --- python/python/lancedb/streaming.py | 164 ++++++++++++++++++----------- 1 file changed, 101 insertions(+), 63 deletions(-) diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 8ed099262..d6d1629fe 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -178,11 +178,14 @@ class _TwoPhaseSplitReader: self._block_last_pos[bp] = pos def __getitems__(self, indices) -> pa.RecordBatch: - start = indices[0] - fetch = len(indices) + # Fancy-indexed (not sliced) so this also supports non-contiguous + # index lists, e.g. the sparse random sampling used by + # _estimate_blocks_per_epoch, not just the contiguous + # range(start, start + fetch) that the main I/O pipeline requests. + idx_arr = np.asarray(indices, dtype=np.int64) with self._lock: - block_pos_slice = self._plan_block_pos[start : start + fetch] - local_off_slice = self._plan_local_offset[start : start + fetch] + block_pos_slice = self._plan_block_pos[idx_arr] + local_off_slice = self._plan_local_offset[idx_arr] needed = sorted(set(block_pos_slice.tolist())) for bp in needed: if bp not in self._loaded: @@ -215,7 +218,7 @@ class _TwoPhaseSplitReader: inverse = np.argsort(group_order) result = combined.take(pa.array(inverse, type=pa.int64())) - watermark = start + fetch - 1 + watermark = int(idx_arr.max()) for bp in list(self._loaded): if self._block_last_pos.get(bp, -1) <= watermark: del self._loaded[bp] @@ -441,8 +444,9 @@ 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. Ignored when ``block_size`` is - set (2-phase mode has its own locality/randomness knobs, below). + default) shuffles rows individually. Cannot be combined with + ``block_size`` (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 exactly this many live rows, and each split @@ -468,7 +472,9 @@ class StreamingDataset(IterableDataset): ``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). + ``window_blocks * block_size``, so the default window actually + holds ``window_blocks`` blocks in RAM at once; pass this explicitly + to decouple the shuffle bound from ``window_blocks``. filter: Optional SQL filter expression (e.g. ``"label = 'dog'"``). Only rows that satisfy the predicate are included in the permutation. The filter @@ -640,6 +646,12 @@ class StreamingDataset(IterableDataset): 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 shuffle_clump_size is not None: + raise ValueError( + "shuffle_clump_size is ignored when block_size is set " + "(2-phase mode uses window_blocks/max_shuffle_distance " + "instead); remove one of them" + ) 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: @@ -864,22 +876,24 @@ class StreamingDataset(IterableDataset): if shuffle: block_seed = shuffle_seed + epoch * _EPOCH_PRIME random.Random(block_seed).shuffle(block_order) - # A block permutation is tiny (one int per block, not per row), - # so it is kept in RAM rather than round-tripped through an - # Arrow-backed permutation table like the 1-phase row mapping. + # Kept as a plain Python list for simplicity, unlike the + # 1-phase row mapping below, which is Arrow-backed because it + # is built and consumed through the Arrow-based permutation + # table machinery; a block permutation has no such need. 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. + # window_blocks is the primary residency knob: by default, + # max_shuffle_distance (the displacement bound that actually + # drives which blocks stay resident -- see + # _TwoPhaseSplitReader) is derived from it so that a split + # holds roughly window_blocks blocks in RAM at once. Passing + # max_shuffle_distance explicitly overrides that derivation. 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 self._window_blocks * block_size ) else: #1-phase shuffled read # Build the permutation table once, deterministically. @@ -904,6 +918,60 @@ class StreamingDataset(IterableDataset): range(rank_start, rank_start + splits_per_rank) ) + def _read_block( + self, lance_ds, block_id: int, columns: Optional[list[str]] + ) -> pa.Table: + """Read one 2-phase block's live rows, in natural (ascending) order.""" + if self._filter is None: + # Plain offset/limit: unambiguous raw-position scan that Lance + # can seek to directly. + row_start, row_end = self._block_ranges[block_id] + return lance_ds.scanner( + offset=row_start, + limit=row_end - row_start, + columns=columns, + ).to_table() + + # Filtered: this block's exact live-row dataset offsets were + # already computed once, up front (from the "_rowoffset" system + # column -- see __init__), so a sparse take() fetches just the + # live rows themselves, skipping any dead-row gaps entirely. + offsets = self._block_live_offsets[block_id] + return lance_ds.take(offsets.tolist(), columns=columns) + + def _make_two_phase_reader( + self, split_idx: int, lance_ds, columns: Optional[list[str]] + ) -> "_TwoPhaseSplitReader": + """Build the 2-phase reader for one split's assigned blocks. + + Shared by ``_iter_owned`` (the real read path) and + ``_estimate_blocks_per_epoch`` (which samples token lengths through + the same interface, so its estimate reflects 2-phase's actual block + assignment and shuffle bound rather than a separate approximation). + """ + blocks_per_split = self._num_blocks // self._num_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 + ) + return _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=lambda block_id: self._read_block( + lance_ds, block_id, columns + ), + columns=columns, + ) + def _estimate_blocks_per_epoch(self) -> int: """Estimate a fixed packed-block budget from a bounded token sample.""" # TODO: Replace this fallback with Lance's dedicated exact token-count @@ -926,12 +994,22 @@ class StreamingDataset(IterableDataset): "pass an explicit value for exact epoch sizing", ) + lance_ds = self._table.to_lance() if self._block_size is not None else None for split in range(self._num_splits): - permutation = Permutation.from_tables( - self._table, self._perm_table, split=split - ) - permutation = permutation.select_columns([token_column]) - permutation = permutation.with_transform(Transforms.arrow2arrow) + if self._block_size is not None: + # Sample through the same 2-phase reader used at read time, + # so the estimate reflects 2-phase's actual block + # assignment and shuffle bound rather than a separate, + # unrelated approximation. + permutation = self._make_two_phase_reader( + split, lance_ds, [token_column] + ) + else: + permutation = Permutation.from_tables( + self._table, self._perm_table, split=split + ) + permutation = permutation.select_columns([token_column]) + permutation = permutation.with_transform(Transforms.arrow2arrow) split_rows = permutation.num_rows if split_rows == 0: raise ValueError( @@ -1086,54 +1164,14 @@ class StreamingDataset(IterableDataset): # 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 lance_ds = self._table.to_lance() - def _read_block(block_id: int) -> pa.Table: - if self._filter is None: - # Plain offset/limit: unambiguous raw-position scan - # that Lance can seek to directly. Bytes/time - # accounting happens generically in _io_call, below, - # not here. - row_start, row_end = self._block_ranges[block_id] - return lance_ds.scanner( - offset=row_start, - limit=row_end - row_start, - ).to_table() - - # Filtered: this block's exact live-row dataset offsets - # were already computed once, up front (from the - # "_rowoffset" system column -- see __init__), so a - # sparse take() fetches just the live rows themselves, - # skipping any dead-row gaps entirely, in the requested - # (already-ascending, already-natural) order. - offsets = self._block_live_offsets[block_id] - return lance_ds.take(offsets.tolist()) - 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, - ) + self._make_two_phase_reader(split_idx, lance_ds, self._columns) ) # No resume/skip support yet in 2-phase: every split starts # fresh regardless of any loaded checkpoint.