diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 1dbec7c39..868a18e22 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -26,6 +26,7 @@ import os import random import threading import time +import uuid import warnings from collections import deque from concurrent.futures import ThreadPoolExecutor @@ -592,11 +593,56 @@ class StreamingDataset(IterableDataset): # row-level permutation. Instead we only shuffle at the # granularity of contiguous row blocks; _perm_table (and its # row-by-row index mapping) is not built at all in this mode. - num_rows = ( - table.count_rows(filter) if filter is not None else table.count_rows() - ) + # + # block_size counts LIVE rows (rows passing `filter`), not raw + # rows, so a block's raw span varies with local filter + # selectivity. Each block therefore remembers its own raw + # [start, end) row range in self._block_ranges rather than being + # derived from block_id via fixed-size arithmetic alone. + if filter is None: + num_rows = table.count_rows() + num_blocks = math.ceil(num_rows / block_size) if num_rows else 0 + block_ranges = [ + (i * block_size, min((i + 1) * block_size, num_rows)) + for i in range(num_blocks) + ] + else: + # One pass, computed once here and never repeated at read + # time: every fragment's raw and live row counts are read + # exactly once (count_rows(filter) touches that fragment's + # 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. + 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 = [] + 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_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)) + num_blocks = len(block_ranges) + self._num_rows = num_rows - self._num_blocks = math.ceil(num_rows / block_size) if num_rows else 0 + self._num_blocks = num_blocks + self._block_ranges: list[tuple[int, int]] = block_ranges + if self._num_blocks < num_splits: + raise ValueError( + f"block_size={block_size} yields only {self._num_blocks} " + f"block(s) from {num_rows} row(s), fewer than num_splits " + f"({num_splits}); use a smaller block_size or fewer splits" + ) block_order = list(range(self._num_blocks)) if shuffle: block_seed = shuffle_seed + epoch * _EPOCH_PRIME @@ -812,40 +858,79 @@ class StreamingDataset(IterableDataset): end = start + blocks_per_split split_blocks.append(self._block_perm[start:end]) - block_size = self._block_size - num_rows = self._num_rows lance_ds = self._table.to_lance() + from . import connect as _connect - def _read_block(block_id: int) -> pa.Table: - row_start = block_id * block_size - row_end = min(row_start + block_size, num_rows) - scanner = lance_ds.scanner( - columns=self._columns, - filter=self._filter, + def _read_block(block_id: int) -> tuple[pa.Table, list[int]]: + # Plain offset/limit, no filter: unambiguous raw-position + # 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. + 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)) + + # Which of these already-in-RAM rows are live? Answered + # without any further object-storage access: tag each row + # with its position, round-trip that through a throwaway + # in-memory Lance table, and apply the real filter engine + # to it locally -- guarantees the exact same filter + # semantics as everywhere else, with no extra I/O. + tagged = raw.append_column( + "__pos", pa.array(range(raw.num_rows), type=pa.int64()) ) - return scanner.to_table() + mem_table = _connect("memory://").create_table( + f"block-{uuid.uuid4().hex}", tagged + ) + live_positions = ( + mem_table.to_lance() + .scanner(filter=self._filter, columns=["__pos"]) + .to_table() + .column("__pos") + .to_pylist() + ) + 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, then shuffle its rows with a permutation unique to - # this split so windows in different splits don't repeat the - # same row order. + # 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] - window = pa.concat_tables([_read_block(b) for b in block_ids]) + 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) - split_idx = my_splits[split_pos] - row_seed = ( - self._shuffle_seed - + self._epoch * _EPOCH_PRIME - + (split_idx + 1) * _SPLIT_PRIME - ) - row_order = list(range(window.num_rows)) - random.Random(row_seed).shuffle(row_order) - return window.take(row_order) + 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( @@ -1421,9 +1506,11 @@ class StreamingDataset(IterableDataset): state["_table"] = _table_to_pickle_state(self._table) # _perm_table: always in-memory; serialise as Arrow data (mirrors # how Permutation.__getstate__ handles its permutation_table). + # None in 2-phase mode, which never builds one. state["_perm_table"] = ( - self._perm_table.name, - self._perm_table.to_arrow(), + (self._perm_table.name, self._perm_table.to_arrow()) + if self._perm_table is not None + else None ) for key in ( "_raw_batches_ref", @@ -1442,17 +1529,23 @@ class StreamingDataset(IterableDataset): table_name = state.pop("_table_name") table_state = state.pop("_table") - perm_name, perm_data = state.pop("_perm_table") + perm_table_state = state.pop("_perm_table") self.__dict__.update(state) self._consumer_iterator_lock = threading.Lock() if self._connection_factory is not None: self._table = self._connection_factory(table_name) else: self._table = _table_from_pickle_state(table_state) - if table_state["kind"] == "memory": + if table_state["kind"] == "memory" and perm_table_state is not None: # Rebuilt from Arrow, so the recorded pin cannot resolve on it. - perm_data = _drop_base_version(perm_data) - self._perm_table = _connect("memory://").create_table(perm_name, perm_data) + perm_name, perm_data = perm_table_state + perm_table_state = (perm_name, _drop_base_version(perm_data)) + if perm_table_state is not None: + perm_name, perm_data = perm_table_state + self._perm_table = _connect("memory://").create_table(perm_name, perm_data) + else: + # 2-phase mode never builds a permutation table. + self._perm_table = None def state_dict(self) -> dict: """Snapshot the dataset's consumption state.