Conversation
Scanning a scalar index end to end went through `IndexReaderStream`, which issues an independent `read_record_batch` per batch. Each of those builds its own decode plan and issues its own small I/O: one read for the page metadata it touches, then one for the page data. Neither coalesces with the neighbouring batch's, so a full scan costs two round trips per batch plus a scheduler build per batch, and the scan is only as far ahead of the consumer as `buffered(n)` lets it be. Adds `IndexReader::whole_file_stream`, which returns a stream over the whole file under a single decode plan or `None`, and a `stream_index_file` helper that prefers it and falls back to the per-batch path for readers without one. The v2 `FileReader` implements it as one `read_stream_projected` over `RangeFull` at the index's own `batch_size`, so btree pages still line up with batches. Taken at the three btree scan sites: `data_stream`, `calculate_included_frags` and `remap`. `IndexReader::read_range_stream` is the same mechanism over a sub-range but at a fixed batch size, which is why this is a separate method. Measured on a 40k-row btree index at batch size 64 (625 batches, no cache, local FS): 1251 read IOPs before, 9 after.
Three follow-ups to the whole-file index scan. 1. `batch_size` is read from the index file's own metadata, so a `u64` that is a multiple of 2^32 truncated to a `rows_per_batch` of 0 in the `as u32` cast: an empty stream that silently retrained an empty index and reported success. Clamp instead. The regression test fails without the clamp. 2. `merge_range_partitioned_lookups` and `merge_pages` build one stream per partition against a *shared* store, and construct all of them before the consumer polls any. A whole-file plan per partition would let partitions the consumer is not draining occupy the scheduler's byte budget, which is credited back only when a batch is polled, and stall the partition being drained. Both keep the demand-paced `buffered(1)` path; only the three single-stream sites use `stream_index_file`. 3. The IOP test used 10,000 rows, below `DEFAULT_INLINE_SCHEDULING_THRESHOLD` (16 Ki), so it exercised inline scheduling rather than the spawned path this change actually affects. Raised to 40,000 rows.
…sources Now that `data_stream()` is a whole-file plan, what it keeps in flight is capped by the store's scheduler budget rather than by readahead. `merge_segments` opens a store per source segment and starts a stream on each, so the merge was sitting on N of those budgets at once (and none of it is visible to the DataFusion memory pool). Adds `IndexStore::with_io_buffer_size` (default no-op) and has `merge_segments` split one scan's budget across the segments it actually reads. The rescope happens before the stream is awaited, because that is where prefetch starts, not at `execute_plan`. Two limits, documented rather than left to be discovered. The 32 MiB floor keeps a whole page in flight per source, so the aggregate is max(one scan's budget, sources * 32 MiB) and not a constant (past `io_parallelism` sources the floor wins and it grows again). And the budget is soft: a request is admitted regardless of remaining bytes when nothing of lower priority is in flight, which is what keeps things moving at any budget. `LanceIndexStore` deliberately does not size its default budget from a process-wide buffer-size variable: its constructor is on the query path too (`open_scalar_index` builds one store per scalar index a scan opens), so a knob set for one workload would silently retune the other. Callers that need a smaller budget ask for it explicitly through `with_io_buffer_size`.
…per scheduler A whole-file plan's outstanding bytes are bounded only by the scheduler budget, and nothing reported that number: `IoQueueState` tracks `bytes_avail` but exposes nothing, so the only available measurement was whole-process RSS, which cannot separate scheduler prefetch from decode buffers from an index merge's own sort. Adds a high-water mark of bytes debited and not yet credited, the current value of the same quantity (which must return to zero once an abandoned read has given its budget back), plus a count of requests admitted over budget on priority: because nothing of lower priority was in flight, or because the request was a later chunk of one already in flight. That last number matters: it is exactly where the budget is soft rather than hard, and until now the only evidence it happens at all was a debug log. Surfaced through `ScanScheduler::backpressure_stats()` alongside the existing `stats()`, and `LanceIndexStore::scheduler()` so a caller can reach it. The lite scheduler reports zeroes; it accounts for reservations with its own RAII type rather than this queue.
… properties Three tests for the memory properties a whole-file index scan has to hold. The budget is respected. A whole-file scan of a 2M-row index (about 23 MiB of incompressible page data) runs at a 1 MiB and a 4 MiB budget and must hold no more than budget + max_iop_size in flight. The slack is deliberate: the budget is soft, because `can_deliver` admits a request over budget when its priority is at or below everything in flight (which guarantees forward progress) or when it is a later chunk of a request already in flight; the slack is the largest single such admission. A fixture guard asserts the file is larger than that bound, and after the first batch the consumer pauses until bytes in flight stop growing and asserts they reached at least half the budget, so the bound cannot hold vacuously and does not depend on I/O outpacing decode. The budget survives an abandoned scan. A whole-file plan has far more outstanding when dropped than a per-batch read did, so a leak in the credit path is proportionally worse here. A scan of a 400k-row index at a 1 MiB budget is dropped after 1, 32 and 256 batches; the scheduler's bytes in flight must return to zero and a subsequent full scan on the same store must complete. The query path cannot reach the whole-file scan. The index is loaded through a store whose readers panic if asked for a whole-file stream, then searched with equality, range and IS NULL. This is a structural assertion, not an inspection: if a search is later routed through a whole-file plan it fails here. A positive control first asserts that a maintenance scan through the same store does trip the guard, proving the guard is consulted at all.
…budget The budget split in `merge_segments` exists for the multi-source case, so it needs a test that runs it with more than one source. This one merges 2, 4 and 8 segments through stores that record the rescoped store each source is handed, then checks: - each source got exactly one store, budgeted at 1/N of a scan (the 32 MiB floor engages at N=8 with the local default of 8 I/O threads) - every rescoped scheduler carried traffic and stayed within its budget plus one split request, with zero priority bypasses - the original stores' schedulers saw no merge traffic at all, so no read path escapes the split - the merged index still contains every row from all N sources The per-source budgets are far larger than the fixture, so this covers the split and the plumbing rather than throttling under it; the single-scan test covers the latter. Also lifts the training-stream construction into `sorted_training_stream`, shared by the index fixture and the merge's new-data input.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Scanning a btree index end to end (
data_stream,calculate_included_frags,remap) goes throughIndexReaderStream, which issues an independentread_record_batchper batch. Each of those builds its own decode plan and issues its own small I/O: one read for the page metadata it touches, then one for the page data. Neither coalesces with the neighbouring batch's, so a full scan costs two round trips and a scheduler build per batch, and the scan is only as far ahead of the consumer asbuffered(n)lets it be.Measured on a 40k-row btree index at batch size 64 (625 batches, no cache, local FS): 1251 read IOPs for the per-batch scan, 9 for the whole-file plan below.
Public API
Two new trait methods, both with defaults, so existing implementors compile unchanged (
rust/lance-index-core/src/scalar.rs):IndexReader::whole_file_streamasync fn whole_file_stream(&self, batch_size: u32, batch_readahead: u32) -> Result<Option<Pin<Box<dyn RecordBatchStream>>>>Ok(None), caller falls back to per-batch readsIndexStore::with_io_buffer_sizefn with_io_buffer_size(&self, bytes: u64) -> Arc<dyn IndexStore>Three new public items for observability:
scheduler::BackpressureStats(fieldsio_buffer_size,max_bytes_in_flight,bytes_in_flight,priority_bypass_admissions)lance-ioScanScheduler::backpressure_stats() -> BackpressureStatslance-ioLanceIndexStore::scheduler() -> &Arc<ScanScheduler>lance-indexNothing is removed, no existing signature changes.
Change
IndexReader::whole_file_stream(batch_size, batch_readahead) -> Option<stream>: the whole file under a single decode plan, orNonefor readers without a fast path. The v2FileReaderimplements it as oneread_stream_projectedoverRangeFullat the index's ownbatch_size, so btree pages still line up with batches. Astream_index_filehelper prefers it and falls back to the per-batch path; the three single-stream btree scan sites use it.IndexReader::read_range_streamalready exists and is the same mechanism over a sub-range, but at a fixed batch size (4096). Btree pages must align with the index's ownbatch_size, which is why this is a separate method rather than a caller of that one.merge_range_partitioned_lookupsandmerge_pagesdeliberately stay on the per-batchbuffered(1)path: they build one stream per partition against a shared store and construct all of them before the consumer polls any, so a whole-file plan per partition would let undrained partitions hold the scheduler's byte budget and stall the one being drained.batch_sizecomes from the index file's own metadata and is now clamped rather than cast tou32: au64multiple of 2^32 truncated to 0 and produced an empty stream that silently retrained an empty index.Memory
A whole-file plan's outstanding bytes are bounded by the store's scheduler budget (
SchedulerConfig::max_bandwidth, 32 MiB per I/O thread) rather than by readahead, and that budget is per store.merge_segmentsopens one store per source segment, so it would otherwise hold N budgets at once, none of it visible to the DataFusion memory pool.IndexStore::with_io_buffer_size(bytes)(default: return self) lets a caller rescope a store's prefetch budget.merge_segmentssplits one scan's budget across the segments it actually reads, before each stream is awaited (which is where prefetch starts). A 32 MiB floor keeps a whole page in flight per source, so the aggregate bound ismax(one scan's budget, sources * 32 MiB)plus one split request per source; pastio_parallelismsources the floor wins and it grows again.LanceIndexStoredoes not size its default budget from a process-wide variable: its constructor is on the query path too (open_scalar_indexbuilds one store per scalar index a scan opens), so callers that need a smaller budget ask for it explicitly.ScanScheduler::backpressure_stats()reports the peak and current bytes debited-but-not-credited and the count of those over-budget admissions, so the budget claims above are checkable rather than argued.LanceIndexStore::scheduler()exposes the scheduler so a caller can reach it. The lite scheduler reports zeroes.Testing
cargo test -p lance-index-core -p lance-index -p lance-io --libpasses (1562 / 2 / 306);cargo clippy --all --tests --benches -- -D warningsis clean. New tests inlance-index:data_streamscan of the 40k-row fixture must take fewer than 64 read IOPs (measured 9; the per-batch path takes 1251).batch_size(2^32) still yields every row.merge_io_buffer_sizeconserves one scan's budget across sources until the floor binds.max_iop_size(read from the object store, so it followsLANCE_MAX_IOP_SIZE); after the first batch the consumer pauses until bytes in flight stop growing and asserts they reached at least half the budget (backpressure engaged, independent of I/O speed); then the scan drains and the boundpeak <= budget + max_iop_sizeholds.whole_file_stream, a maintenance scan is shown to trip that guard, then equality, range and IS NULL searches run against it.merge_segmentsof 2, 4 and 8 segments hands each source exactly one rescoped store at 1/N of a scan's budget, every rescoped scheduler carries traffic within its bound with zero bypasses, the original schedulers see no merge traffic, and the merged index has every row. Its per-source budgets (32 MiB and up) are far larger than the 40k-row fixture, so this test checks the split and the plumbing rather than throttling under it; the single-scan test covers throttling.Compatibility
No file format change.
whole_file_streamandwith_io_buffer_sizeare trait methods with defaults, so existingIndexReader/IndexStoreimplementations are unaffected and keep the per-batch behaviour. New public API:lance_io::scheduler::BackpressureStats,ScanScheduler::backpressure_stats,LanceIndexStore::scheduler.Tracking: Ported from rerun-io#63.