From 1b16057d23ae12ef9ff56dbfa3306c3eafd7914f Mon Sep 17 00:00:00 2001 From: jasonz-lance Date: Thu, 8 Oct 2026 18:35:59 -0700 Subject: [PATCH] perf(python): halve 2-phase's per-row plan-array memory footprint _TwoPhaseSplitReader retains _plan_block_pos/_plan_local_offset for its whole lifetime, one entry per row in the entire split (not bounded by window_blocks -- see the open whole-split-memory review thread), so their dtype is a real per-row memory cost, not a rounding choice. Both were int64 despite holding small values: block_pos is a block index (bounded by this split's block count) and local_offset is an in-block row offset (bounded by block_size), both far smaller than the row-rank range sigma/order need int64 for. Switch to int32, falling back to int64 only if a single split's block count or block size itself exceeds what int32 can represent (practically never). Halves this bookkeeping's memory cost for any realistic config, as a stopgap short of the full lazy/windowed redesign that would bound it by window_blocks instead of split size -- deliberately deferred, since that touches the same hot-path indexing logic as the prefetch-loop regression fixed earlier and deserves its own dedicated pass. Added a regression test pinning the dtype (so a future change can't silently regress it back to int64) and checking values are still correct, not just narrower. Full suite: 211 passed. Co-Authored-By: Claude Sonnet 5 --- python/python/lancedb/streaming.py | 21 +++++++++++++++---- .../python/tests/test_elastic_dataloader.py | 20 ++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 93cd218b9..030b3abbb 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -161,12 +161,25 @@ class _TwoPhaseSplitReader: raise ValueError(f"skip ({skip}) out of range for {total_rows} row(s)") # 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) + # describe the row whose dense natural rank is exactly k. These two + # arrays are retained for the reader's whole lifetime, one entry per + # row in the *entire* split (not just the resident window -- see the + # whole-split-memory review thread, still open) -- so their dtype is + # the full per-row memory cost, not a rounding choice. block_pos + # values are block indices (bounded by num_blocks) and local_offset + # values are in-block row offsets (bounded by this split's largest + # block), both far smaller than the row-rank range sigma/order below + # need int64 for, so int32 halves that cost in the realistic case; + # falls back to int64 only if a single split's block count or block + # size itself exceeds what int32 can represent. + max_count = int(counts.max()) if num_blocks else 0 + pos_dtype = np.int32 if num_blocks < 2**31 else np.int64 + local_dtype = np.int32 if max_count < 2**31 else np.int64 + block_pos = np.repeat(np.arange(num_blocks, dtype=pos_dtype), counts) local_offset = ( - np.concatenate([np.arange(c, dtype=np.int64) for c in counts]) + np.concatenate([np.arange(c, dtype=local_dtype) for c in counts]) if num_blocks - else np.array([], dtype=np.int64) + else np.array([], dtype=local_dtype) ) if shuffle and total_rows: diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 3941974a9..6cb9940b8 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -36,6 +36,7 @@ import threading from unittest.mock import patch import lancedb +import numpy as np import pyarrow as pa import pytest from utils import ( @@ -3476,6 +3477,25 @@ def test_two_phase_memory_table_already_at_a_later_version(): assert sorted(ids) == list(range(8)) +def test_two_phase_reader_plan_arrays_are_int32(lance_table): + """_plan_block_pos/_plan_local_offset are retained for the reader's + whole lifetime, one entry per row in the entire split (the open + whole-split-memory review thread) -- so their dtype is a real, constant + memory-footprint decision, not an implementation detail. Pinned here at + int32 (half of int64) so a future change can't silently regress it.""" + ds = StreamingDataset( + lance_table, num_splits=1, shuffle_seed=SHUFFLE_SEED, block_size=10 + ) + reader = ds._make_two_phase_reader(0, ds._open_pinned_table(), None) + assert reader._plan_block_pos.dtype == np.int32 + assert reader._plan_local_offset.dtype == np.int32 + # Values must still be correct, not just narrower -- e.g. a block index + # silently wrapping around int32 would be a far worse regression than + # the memory this dtype choice saves. + assert reader._plan_block_pos.max() < ds._num_blocks // ds._num_splits + assert reader._plan_local_offset.max() < 10 # block_size + + def test_two_phase_rejects_lsm_write_spec_table(tmp_path): """An LSM write spec means some acknowledged writes live only in the MemWAL, invisible to a version-pinned read; block_size must refuse the