Skip to content

perf(scalar): build one decode plan per index scan, not one per batch - #9289

Draft
amunra wants to merge 6 commits into
lance-format:mainfrom
rerun-io:upstream/btree-single-decode-plan
Draft

amunra wants to merge 6 commits into
lance-format:mainfrom
rerun-io:upstream/btree-single-decode-plan

Conversation

@amunra

@amunra amunra commented Sep 16, 2026

Copy link
Copy Markdown

Problem

Scanning a btree index end to end (data_stream, calculate_included_frags, remap) goes 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 and a scheduler build per batch, and the scan is only as far ahead of the consumer as buffered(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):

item shape default
IndexReader::whole_file_stream async 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 reads
IndexStore::with_io_buffer_size fn with_io_buffer_size(&self, bytes: u64) -> Arc<dyn IndexStore> returns self unchanged

Three new public items for observability:

item crate
scheduler::BackpressureStats (fields io_buffer_size, max_bytes_in_flight, bytes_in_flight, priority_bypass_admissions) lance-io
ScanScheduler::backpressure_stats() -> BackpressureStats lance-io
LanceIndexStore::scheduler() -> &Arc<ScanScheduler> lance-index

Nothing 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, or None for readers without a fast path. 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. A stream_index_file helper prefers it and falls back to the per-batch path; the three single-stream btree scan sites use it.
  • IndexReader::read_range_stream already 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 own batch_size, which is why this is a separate method rather than a caller of that one.
  • merge_range_partitioned_lookups and merge_pages deliberately stay on the per-batch buffered(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_size comes from the index file's own metadata and is now clamped rather than cast to u32: a u64 multiple 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_segments opens 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_segments splits 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 is max(one scan's budget, sources * 32 MiB) plus one split request per source; past io_parallelism sources the floor wins and it grows again.
  • The budget is soft by design: a request is admitted regardless of remaining bytes when nothing of lower priority is in flight, which is what guarantees forward progress at any budget, and when it is a later chunk of a request already in flight.
  • LanceIndexStore does not size its default budget from a process-wide variable: its constructor is on the query path too (open_scalar_index builds 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 --lib passes (1562 / 2 / 306); cargo clippy --all --tests --benches -- -D warnings is clean. New tests in lance-index:

  • IOP bound: a whole-file data_stream scan of the 40k-row fixture must take fewer than 64 read IOPs (measured 9; the per-batch path takes 1251).
  • Oversized batch_size (2^32) still yields every row.
  • merge_io_buffer_size conserves one scan's budget across sources until the floor binds.
  • Budget respected: a 2M-row incompressible index (about 23 MiB) is scanned at 1 MiB and 4 MiB budgets; a fixture guard asserts the file is larger than budget + max_iop_size (read from the object store, so it follows LANCE_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 bound peak <= budget + max_iop_size holds.
  • Budget returned on drop: 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 following full scan must complete.
  • Query path never takes the whole-file scan: the index is loaded through a store whose readers panic on whole_file_stream, a maintenance scan is shown to trip that guard, then equality, range and IS NULL searches run against it.
  • merge_segments of 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_stream and with_io_buffer_size are trait methods with defaults, so existing IndexReader / IndexStore implementations 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.

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.
@github-actions github-actions Bot added A-index Vector index, linalg, tokenizer A-encoding Encoding, IO, file reader/writer performance labels Sep 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-encoding Encoding, IO, file reader/writer A-index Vector index, linalg, tokenizer performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants