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