From 8ddd4d437d443ba0285bbcfcf2226b9630f8758e Mon Sep 17 00:00:00 2001 From: Ryan Mack Date: Sun, 16 Aug 2026 14:04:17 -0400 Subject: [PATCH 1/6] perf(scalar): build one decode plan per index scan, not one per batch 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. --- rust/lance-index-core/src/scalar.rs | 24 +++ rust/lance-index/src/scalar/btree.rs | 154 ++++++++++++++++---- rust/lance-index/src/scalar/lance_format.rs | 26 ++++ 3 files changed, 175 insertions(+), 29 deletions(-) diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index 80da1633a26..f74763bfc2f 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -223,6 +223,30 @@ pub trait IndexReader: Send + Sync { futures::stream::once(async move { Ok(batch) }), ))) } + /// Stream the entire file as `batch_size`-row batches under a **single** decode plan, + /// or `None` if this reader has no such fast path (the caller then falls back to + /// per-batch reads). + /// + /// This exists because reading a file as N independent `read_record_batch` calls + /// builds N decode plans, and each plan issues its own small I/O: one read for the + /// page metadata it touches and one for the page data, neither of which coalesces with + /// the neighbouring batch's. A full scan therefore costs O(batches) round trips. A + /// single plan over the whole file schedules every page up front, so the scheduler can + /// coalesce those reads into a handful of large requests. + /// + /// A large index runs to thousands of pages, so this dominates the cost of a full + /// scan. Callers that read a file end to end should prefer this method. + /// + /// [`Self::read_range_stream`] is the same mechanism over a sub-range, but it uses a + /// fixed batch size; callers whose batches must line up with the file's own pages + /// (a btree page is one batch) need this method. + async fn whole_file_stream( + &self, + _batch_size: u32, + _batch_readahead: u32, + ) -> Result>>> { + Ok(None) + } /// Return the number of batches in the file async fn num_batches(&self, batch_size: u64) -> u32; /// Return the number of rows in the file diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index cb2c842b906..4d68659f522 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -55,9 +55,9 @@ use datafusion_physical_expr::{ PhysicalExpr, PhysicalSortExpr, create_physical_expr, expressions::Column, }; use futures::{ - FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, + FutureExt, Stream, StreamExt, TryStreamExt, future::BoxFuture, - stream::{self}, + stream::{self, BoxStream}, }; use lance_core::deepsize::DeepSizeOf; use lance_core::{ @@ -1848,10 +1848,9 @@ impl BTreeIndex { let reader = lazy_reader.get().await?; let new_schema = Arc::new(self.train_schema()); let new_schema_clone = new_schema.clone(); - let reader_stream = IndexReaderStream::new(reader, self.batch_size).await; - let batches = reader_stream - .map(|fut| fut.map_err(DataFusionError::from)) - .buffered(self.store.io_parallelism()) + let batches = stream_index_file(reader, self.batch_size, self.store.io_parallelism()) + .await? + .map_err(DataFusionError::from) .map_ok(move |batch| { RecordBatch::try_new( new_schema.clone(), @@ -2110,9 +2109,12 @@ impl Index for BTreeIndex { let lazy_reader = LazyIndexReader::new(self.store.clone(), self.ranges_to_files.clone()); let sub_index_reader = lazy_reader.get().await?; - let mut reader_stream = IndexReaderStream::new(sub_index_reader, self.batch_size) - .await - .buffered(self.store.io_parallelism()); + let mut reader_stream = stream_index_file( + sub_index_reader, + self.batch_size, + self.store.io_parallelism(), + ) + .await?; while let Some(serialized) = reader_stream.try_next().await? { let page = FlatIndex::try_new(serialized)?; frag_ids |= page.calculate_included_frags()?; @@ -2308,20 +2310,23 @@ impl ScalarIndex for BTreeIndex { let train_schema_clone = train_schema.clone(); let train_schema = train_schema.clone(); - let remapped_stream = IndexReaderStream::new(sub_index_reader, self.batch_size) - .await - .buffered(self.store.io_parallelism()) - .map_err(DataFusionError::from) - .and_then(move |batch| { - // Remap the batch and then convert from the serialized schema to the training input schema - let remapped = - FlatIndex::remap_batch(batch, &mapping).map_err(DataFusionError::from); - let with_train_schema = remapped.and_then(|batch| { - RecordBatch::try_new(train_schema.clone(), batch.columns().to_vec()) - .map_err(DataFusionError::from) - }); - std::future::ready(with_train_schema) + let remapped_stream = stream_index_file( + sub_index_reader, + self.batch_size, + self.store.io_parallelism(), + ) + .await? + .map_err(DataFusionError::from) + .and_then(move |batch| { + // Remap the batch and then convert from the serialized schema to the training input schema + let remapped = + FlatIndex::remap_batch(batch, &mapping).map_err(DataFusionError::from); + let with_train_schema = remapped.and_then(|batch| { + RecordBatch::try_new(train_schema.clone(), batch.columns().to_vec()) + .map_err(DataFusionError::from) }); + std::future::ready(with_train_schema) + }); let remapped_stream = Box::pin(RecordBatchStreamAdapter::new( train_schema_clone, @@ -2816,8 +2821,12 @@ async fn merge_range_partitioned_lookups( for (idx, (part_id, part_lookup_file)) in sorted_part_lookup_files.into_iter().enumerate() { let lookup_reader = store.open_index_file(&part_lookup_file).await?; - let reader_stream = IndexReaderStream::new(lookup_reader.clone(), batch_size).await; - let mut stream = reader_stream.buffered(batch_readhead.unwrap_or(1)).boxed(); + let mut stream = stream_index_file( + lookup_reader.clone(), + batch_size, + batch_readhead.unwrap_or(1), + ) + .await?; while let Some(batch) = stream.next().await { let original_batch = batch?; let modified_batch = add_offset_to_page_idx(&original_batch, num_pages_written)?; @@ -2978,11 +2987,9 @@ async fn merge_pages( let reader = store.open_index_file(&page_file_name).await?; - let reader_stream = IndexReaderStream::new(reader, batch_size).await; - - let stream = reader_stream - .map(|fut| fut.map_err(DataFusionError::from)) - .buffered(batch_readhead.unwrap_or(1)) + let stream = stream_index_file(reader, batch_size, batch_readhead.unwrap_or(1)) + .await? + .map_err(DataFusionError::from) .boxed(); let sendable_stream = @@ -3161,6 +3168,29 @@ pub(crate) fn part_lookup_file_path(partition_id: u64) -> String { format!("part_{}_{}", partition_id, BTREE_LOOKUP_NAME) } +/// Stream an index file end to end as `batch_size`-row batches. +/// +/// Prefers the reader's single-decode-plan path ([`IndexReader::whole_file_stream`]). +/// Falls back to independent per-batch reads for readers that don't offer one, which is +/// what every caller used to do unconditionally — correct, but it builds a decode plan per +/// batch and each one issues its own uncoalesced metadata and data reads. +async fn stream_index_file( + reader: Arc, + batch_size: u64, + batch_readahead: usize, +) -> Result>> { + if let Some(stream) = reader + .whole_file_stream(batch_size as u32, batch_readahead.max(1) as u32) + .await? + { + return Ok(stream.boxed()); + } + Ok(IndexReaderStream::new(reader, batch_size) + .await + .buffered(batch_readahead.max(1)) + .boxed()) +} + /// A stream that reads the original training data back out of the index /// /// This is used for updating the index @@ -3587,6 +3617,72 @@ mod tests { assert_eq!(original_data, remapped_data); } + /// Scanning an index end to end must build ONE decode plan, not one per batch. + /// + /// Reading a file as N independent `read_record_batch` calls builds N decode plans, + /// and each one issues its own small reads: the page metadata it touches and then the + /// page data. Neither coalesces with the neighbouring batch's, so the scan costs two + /// round trips per batch. A single plan over the whole file schedules every page up + /// front and the scheduler merges those reads into a handful of large requests. + #[tokio::test] + async fn test_data_stream_builds_one_decode_plan() { + const BATCH_SIZE: u64 = 64; + const ROWS: u64 = 10_000; + + let tmpdir = TempObjDir::default(); + let object_store = Arc::new(ObjectStore::local()); + let test_store = Arc::new(LanceIndexStore::new( + object_store.clone(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let data = gen_batch() + .col("value", array::step::()) + .col("_rowid", array::step::()) + .into_df_exec(RowCount::from(100), BatchCount::from(100)); + let schema = data.schema(); + let sort_expr = PhysicalSortExpr::new_default(col("value", schema.as_ref()).unwrap()); + let plan = Arc::new(SortExec::new([sort_expr].into(), data)); + let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let stream = break_stream(stream, BATCH_SIZE as usize); + let stream = stream.map_err(DataFusionError::from); + let stream = + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream; + + train_btree_index(stream, test_store.as_ref(), BATCH_SIZE, None, None) + .await + .unwrap(); + + let index = BTreeIndex::load(test_store, None, &LanceCache::no_cache()) + .await + .unwrap(); + + let num_pages = ROWS.div_ceil(BATCH_SIZE); + + // Reset the counters so we measure only the scan. + let _ = object_store.io_stats_incremental(); + + let mut stream = index.data_stream().await.unwrap(); + let mut rows_seen = 0u64; + while let Some(batch) = stream.try_next().await.unwrap() { + rows_seen += batch.num_rows() as u64; + } + assert_eq!(rows_seen, ROWS); + + let iops = object_store.io_stats_incremental().read_iops; + // Measured on this fixture: a handful of IOPs with a single plan, two per batch + // plus one with one plan per batch. The bound is deliberately far below the + // per-batch figure so the test fails loudly if the fast path stops being taken, + // but well above the measured value so it does not tighten into a flake if + // readahead or footer handling changes. + assert!( + iops < 32, + "scanning a {num_pages}-batch index took {iops} read IOPs; a decode plan \ + rebuilt per batch costs about two per batch" + ); + } + #[tokio::test] async fn test_nan_ordering() { let tmpdir = TempObjDir::default(); diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index 69e18fc852c..da5d02d899e 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -377,6 +377,32 @@ impl IndexReader for CurrentIndexReader { .await } + /// Single decode plan for the whole file — see [`IndexReader::whole_file_stream`]. + async fn whole_file_stream( + &self, + batch_size: u32, + batch_readahead: u32, + ) -> Result>>> { + if CurrentFileReader::num_rows(&self.0) == 0 { + return Ok(None); + } + let projection = versions::reader_projection_from_whole_schema( + self.0.schema(), + self.0.metadata().version(), + ); + let stream = self + .0 + .read_stream_projected( + ReadBatchParams::RangeFull, + batch_size, + batch_readahead, + projection, + FilterExpression::no_filter(), + ) + .await?; + Ok(Some(stream)) + } + // V2 format has removed the row group concept, // so here we assume each batch is with 4096 rows. async fn num_batches(&self, batch_size: u64) -> u32 { From 9be0a8d4ed8b7f112134c6c0b04181e8caf1d7b9 Mon Sep 17 00:00:00 2001 From: Ryan Mack Date: Mon, 17 Aug 2026 04:59:43 -0400 Subject: [PATCH 2/6] fix(scalar): clamp btree batch_size and keep shared-store merges paced 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. --- rust/lance-index/src/scalar/btree.rs | 93 ++++++++++++++++++++++------ 1 file changed, 74 insertions(+), 19 deletions(-) diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 4d68659f522..741c34092b4 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -55,7 +55,7 @@ use datafusion_physical_expr::{ PhysicalExpr, PhysicalSortExpr, create_physical_expr, expressions::Column, }; use futures::{ - FutureExt, Stream, StreamExt, TryStreamExt, + FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::BoxFuture, stream::{self, BoxStream}, }; @@ -2821,12 +2821,14 @@ async fn merge_range_partitioned_lookups( for (idx, (part_id, part_lookup_file)) in sorted_part_lookup_files.into_iter().enumerate() { let lookup_reader = store.open_index_file(&part_lookup_file).await?; - let mut stream = stream_index_file( - lookup_reader.clone(), - batch_size, - batch_readhead.unwrap_or(1), - ) - .await?; + // Deliberately NOT `stream_index_file`: this loop and `merge_pages` below build one + // stream per partition against a *shared* `store`, and every stream is constructed + // before the consumer polls any of them. 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 that is being + // drained. These two paths stay demand-paced at `buffered(1)`. + let reader_stream = IndexReaderStream::new(lookup_reader.clone(), batch_size).await; + let mut stream = reader_stream.buffered(batch_readhead.unwrap_or(1)).boxed(); while let Some(batch) = stream.next().await { let original_batch = batch?; let modified_batch = add_offset_to_page_idx(&original_batch, num_pages_written)?; @@ -2987,9 +2989,11 @@ async fn merge_pages( let reader = store.open_index_file(&page_file_name).await?; - let stream = stream_index_file(reader, batch_size, batch_readhead.unwrap_or(1)) - .await? - .map_err(DataFusionError::from) + // Shared `store` across k partitions — see the note in `merge_range_partitioned_lookups`. + let reader_stream = IndexReaderStream::new(reader, batch_size).await; + let stream = reader_stream + .map(|fut| fut.map_err(DataFusionError::from)) + .buffered(batch_readhead.unwrap_or(1)) .boxed(); let sendable_stream = @@ -3179,8 +3183,12 @@ async fn stream_index_file( batch_size: u64, batch_readahead: usize, ) -> Result>> { + // Clamp rather than cast: `batch_size` comes from the index file's own metadata + // (`BATCH_SIZE_META_KEY`), so a `u64` that is a multiple of 2^32 would truncate to a + // `rows_per_batch` of 0 and yield an empty stream — silently retraining an empty index. + let batch_size_u32 = u32::try_from(batch_size).unwrap_or(u32::MAX).max(1); if let Some(stream) = reader - .whole_file_stream(batch_size as u32, batch_readahead.max(1) as u32) + .whole_file_stream(batch_size_u32, batch_readahead.max(1) as u32) .await? { return Ok(stream.boxed()); @@ -3627,7 +3635,9 @@ mod tests { #[tokio::test] async fn test_data_stream_builds_one_decode_plan() { const BATCH_SIZE: u64 = 64; - const ROWS: u64 = 10_000; + // Above `DEFAULT_INLINE_SCHEDULING_THRESHOLD` (16 Ki rows) so this exercises the + // *spawned* scheduling path that production uses, not the inline one. + const ROWS: u64 = 40_000; let tmpdir = TempObjDir::default(); let object_store = Arc::new(ObjectStore::local()); @@ -3640,7 +3650,7 @@ mod tests { let data = gen_batch() .col("value", array::step::()) .col("_rowid", array::step::()) - .into_df_exec(RowCount::from(100), BatchCount::from(100)); + .into_df_exec(RowCount::from(100), BatchCount::from(400)); let schema = data.schema(); let sort_expr = PhysicalSortExpr::new_default(col("value", schema.as_ref()).unwrap()); let plan = Arc::new(SortExec::new([sort_expr].into(), data)); @@ -3671,18 +3681,63 @@ mod tests { assert_eq!(rows_seen, ROWS); let iops = object_store.io_stats_incremental().read_iops; - // Measured on this fixture: a handful of IOPs with a single plan, two per batch - // plus one with one plan per batch. The bound is deliberately far below the - // per-batch figure so the test fails loudly if the fast path stops being taken, - // but well above the measured value so it does not tighten into a flake if - // readahead or footer handling changes. + // Measured on this fixture: 9 IOPs with a single plan, 1251 (two per batch plus + // one) with one plan per batch. The bound is deliberately far below the per-batch + // figure so the test fails loudly if the fast path stops being taken, but well + // above 9 so it does not tighten into a flake if readahead or footer handling + // changes. assert!( - iops < 32, + iops < 64, "scanning a {num_pages}-batch index took {iops} read IOPs; a decode plan \ rebuilt per batch costs about two per batch" ); } + /// `batch_size` is read from the index file's own metadata, so a value that is a + /// multiple of 2^32 used to truncate to a `rows_per_batch` of 0 — an empty stream that + /// silently retrained an empty index and reported success. It must clamp instead. + #[tokio::test] + async fn test_stream_index_file_clamps_oversized_batch_size() { + let tmpdir = TempObjDir::default(); + let object_store = Arc::new(ObjectStore::local()); + let test_store = Arc::new(LanceIndexStore::new( + object_store.clone(), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let data = gen_batch() + .col("value", array::step::()) + .col("_rowid", array::step::()) + .into_df_exec(RowCount::from(100), BatchCount::from(10)); + let schema = data.schema(); + let sort_expr = PhysicalSortExpr::new_default(col("value", schema.as_ref()).unwrap()); + let plan = Arc::new(SortExec::new([sort_expr].into(), data)); + let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let stream = break_stream(stream, 64); + let stream = stream.map_err(DataFusionError::from); + let stream = + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream; + train_btree_index(stream, test_store.as_ref(), 64, None, None) + .await + .unwrap(); + + let reader = test_store.open_index_file(BTREE_PAGES_NAME).await.unwrap(); + + // 2^32 truncates to 0 in a bare `as u32`. + let mut stream = super::stream_index_file(reader, 1u64 << 32, 8) + .await + .unwrap(); + let mut rows = 0usize; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows(); + } + assert_eq!( + rows, 1000, + "oversized batch_size must not yield an empty stream" + ); + } + #[tokio::test] async fn test_nan_ordering() { let tmpdir = TempObjDir::default(); From 0b9b080819d8aa54357ae21d2f17297570a4ada5 Mon Sep 17 00:00:00 2001 From: Ryan Mack Date: Mon, 17 Aug 2026 10:16:53 -0400 Subject: [PATCH 3/6] perf(scalar): split one scan's prefetch budget across merge_segments 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`. --- rust/lance-index-core/src/scalar.rs | 28 ++++++++ rust/lance-index/src/scalar/btree.rs | 72 +++++++++++++++++++++ rust/lance-index/src/scalar/lance_format.rs | 15 +++++ 3 files changed, 115 insertions(+) diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index f74763bfc2f..ed3f2c94b59 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -240,6 +240,21 @@ pub trait IndexReader: Send + Sync { /// [`Self::read_range_stream`] is the same mechanism over a sub-range, but it uses a /// fixed batch size; callers whose batches must line up with the file's own pages /// (a btree page is one batch) need this method. + /// + /// # Memory profile + /// + /// Unlike a per-batch read, whose outstanding data is capped by its own readahead, + /// a whole-file plan is scheduled eagerly and its peak resident encoded bytes are + /// bounded only by the store's scheduler byte budget, which for [`IndexStore`] + /// implementations backed by `SchedulerConfig::max_bandwidth` is 32 MiB per I/O + /// thread (2 GiB at the cloud default of 64, 256 MiB at the local default of 8). + /// `batch_readahead` does not bound this; it bounds decode buffering only. None of + /// it is visible to a DataFusion memory pool. + /// + /// That budget is per store, so a caller that fans out over several stores at once + /// (see `BTreeIndex::merge_segments`, which opens one store per source segment) + /// multiplies it. Such callers should rescope each store first with + /// [`IndexStore::with_io_buffer_size`]. async fn whole_file_stream( &self, _batch_size: u32, @@ -296,6 +311,19 @@ pub trait IndexStore: std::fmt::Debug + Send + Sync + DeepSizeOf { /// Suggested I/O parallelism for the store fn io_parallelism(&self) -> usize; + /// Return an equivalent store whose scheduler holds at most `bytes` of outstanding + /// prefetched data. + /// + /// This is a hint, for callers that drive several stores concurrently and want the + /// aggregate prefetch to stay near what one store would have used on its own. Stores + /// that do not own a scheduler return themselves unchanged. Note that the underlying + /// budget is soft: 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. + fn with_io_buffer_size(&self, _bytes: u64) -> Arc { + self.clone_arc() + } + /// Create a new file and return a writer to store data in the file async fn new_index_file(&self, name: &str, schema: Arc) -> Result>; diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 741c34092b4..9c24e45a92e 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -1842,7 +1842,23 @@ impl BTreeIndex { Schema::new(vec![value_field, row_id_field]) } + /// Clone this index with its store rescoped to hold at most `bytes` of outstanding + /// prefetched data, for callers that drain several segments at once. + /// + /// See [`IndexStore::with_io_buffer_size`]; stores that cannot rescope return + /// themselves and this is then a no-op. + fn with_io_buffer_size(&self, bytes: u64) -> Self { + let mut scoped = self.clone(); + scoped.store = self.store.with_io_buffer_size(bytes); + scoped + } + /// Create a stream of all the data in the index, in the same format used to train the index + /// + /// This is a whole-file plan where the reader supports one (see + /// [`IndexReader::whole_file_stream`]), so its peak resident encoded bytes are bounded + /// by the store's scheduler budget rather than by readahead. Callers that run several + /// of these concurrently should rescope with [`Self::with_io_buffer_size`] first. async fn data_stream(&self) -> Result { let lazy_reader = LazyIndexReader::new(self.store.clone(), self.ranges_to_files.clone()); let reader = lazy_reader.get().await?; @@ -1908,11 +1924,27 @@ impl BTreeIndex { ))); } + // Every `data_stream()` below is a whole-file plan that begins prefetching as soon + // as it is awaited, i.e. before the merge operator below even exists, and each + // source segment carries its own store and therefore its own scheduler budget + // (`open_scalar_index` builds one store per segment). Left alone the merge would + // hold N times a single scan's worth of encoded bytes, none of it visible to the + // DataFusion memory pool. Rescope each source so the aggregate stays near one + // scan's budget. `new_data` is a dataset scan against a different store and keeps + // its own budget. + let num_sources = old_data_filters + .iter() + .filter(|filter| !filter_keeps_nothing(filter)) + .count() + .max(1) as u64; + let io_buffer_size = merge_io_buffer_size(first.store.io_parallelism(), num_sources); + let mut inputs: Vec> = Vec::with_capacity(segments.len() + 1); for (segment, old_data_filter) in segments.iter().zip(old_data_filters) { if filter_keeps_nothing(old_data_filter) { continue; } + let segment = segment.with_io_buffer_size(io_buffer_size); let stream = segment.data_stream().await?; let stream = match segment.frag_reuse_index.clone() { Some(frag_reuse_index) => remap_row_ids(stream, frag_reuse_index), @@ -1977,6 +2009,22 @@ fn filter_row_ids( Box::pin(RecordBatchStreamAdapter::new(schema, filtered)) } +/// Prefetch allowance per I/O thread, mirroring `SchedulerConfig::max_bandwidth`, which +/// sizes an index store's budget assuming a 32 MiB maximum page. +const MERGE_BYTES_PER_IO_THREAD: u64 = 32 * 1024 * 1024; + +/// Divide one scan's worth of prefetch budget across `num_sources` segment streams that +/// will be drained concurrently. +/// +/// The floor keeps a whole page in flight per source, so the aggregate bound is +/// `max(one scan's budget, num_sources * 32 MiB)` rather than a constant. Past +/// `io_parallelism` sources the floor wins and the total starts growing again; that is the +/// point at which a caller wanting a real constant would have to merge in passes. +fn merge_io_buffer_size(io_parallelism: usize, num_sources: u64) -> u64 { + let total = MERGE_BYTES_PER_IO_THREAD * io_parallelism.max(1) as u64; + (total / num_sources.max(1)).max(MERGE_BYTES_PER_IO_THREAD) +} + /// True if `filter` would keep no rows at all (its keep-set is empty), letting /// the merge skip reading the segment entirely. fn filter_keeps_nothing(filter: &Option) -> bool { @@ -3693,6 +3741,30 @@ mod tests { ); } + /// A whole-file plan's prefetch is bounded by its store's scheduler budget, and + /// `merge_segments` drains one store per source at the same time. Splitting the budget + /// is what keeps the merge's aggregate near a single scan's, and the floor is what + /// makes that "near" rather than "equal" past `io_parallelism` sources. + #[test] + fn test_merge_io_buffer_size_splits_one_scans_budget() { + const MIB: u64 = 1024 * 1024; + // Cloud default: 64 threads, so one scan gets 2 GiB. + assert_eq!(super::merge_io_buffer_size(64, 1), 2048 * MIB); + assert_eq!(super::merge_io_buffer_size(64, 4), 512 * MIB); + // Aggregate is conserved while the floor is not binding. + for sources in [1u64, 2, 8, 64] { + assert_eq!( + super::merge_io_buffer_size(64, sources) * sources, + 2048 * MIB + ); + } + // Past that the floor wins and the aggregate grows again, as documented. + assert_eq!(super::merge_io_buffer_size(64, 128), 32 * MIB); + // Local default of 8 threads, and degenerate inputs stay at the floor. + assert_eq!(super::merge_io_buffer_size(8, 1), 256 * MIB); + assert_eq!(super::merge_io_buffer_size(0, 0), 32 * MIB); + } + /// `batch_size` is read from the index file's own metadata, so a value that is a /// multiple of 2^32 used to truncate to a `rows_per_batch` of 0 — an empty stream that /// silently retrained an empty index and reported success. It must clamp instead. diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index da5d02d899e..ad090da3e4b 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -82,6 +82,11 @@ impl LanceIndexStore { metadata_cache: Arc, format_version: ConcreteFileVersion, ) -> Self { + // Deliberately not consulting any process-wide buffer-size env var here: this + // constructor is on the query path too (`open_scalar_index` builds one 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, via + // `IndexStore::with_io_buffer_size`. let scheduler = ScanScheduler::new( object_store.clone(), SchedulerConfig::max_bandwidth(&object_store), @@ -445,6 +450,16 @@ impl IndexStore for LanceIndexStore { self.object_store.io_parallelism() } + fn with_io_buffer_size(&self, bytes: u64) -> Arc { + // The metadata cache is shared with the original store, so files already opened + // through it stay cached; only the scheduler, and therefore the prefetch budget, + // is private to the returned store. + let mut scoped = self.clone(); + scoped.scheduler = + ScanScheduler::new(self.object_store.clone(), SchedulerConfig::new(bytes)); + Arc::new(scoped) + } + async fn new_index_file( &self, name: &str, From 132c615b10347f8637144671ed948c4c1a5b60a3 Mon Sep 17 00:00:00 2001 From: Ryan Mack Date: Tue, 18 Aug 2026 15:34:46 -0400 Subject: [PATCH 4/6] feat(io): report peak bytes in flight and priority-bypass admissions 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. --- rust/lance-index/src/scalar/lance_format.rs | 6 ++ rust/lance-io/src/scheduler.rs | 61 +++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index ad090da3e4b..77c3341dd1c 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -116,6 +116,12 @@ impl LanceIndexStore { self.io_priority } + /// The scheduler backing this store, for tests and diagnostics that need to observe + /// how much prefetch the store actually held (see `ScanScheduler::backpressure_stats`). + pub fn scheduler(&self) -> &Arc { + &self.scheduler + } + fn index_file_path(&self, name: &str) -> Result { let relative_path = Path::parse(name).map_err(|err| { Error::invalid_input(format!("invalid index file path {name:?}: {err}")) diff --git a/rust/lance-io/src/scheduler.rs b/rust/lance-io/src/scheduler.rs index 1876fe3feba..f12943f828d 100644 --- a/rust/lance-io/src/scheduler.rs +++ b/rust/lance-io/src/scheduler.rs @@ -119,6 +119,17 @@ struct IoQueueState { last_warn: AtomicU64, // When true, skip all byte-based backpressure checks (set when io_buffer_size == 0) no_backpressure: bool, + // High-water mark of bytes debited but not yet credited back. + // + // This is the quantity the byte budget is supposed to bound, and until now nothing + // reported it: callers could only observe whole-process RSS, which cannot separate + // scheduler prefetch from decode buffers. + max_bytes_in_flight: u64, + // Requests admitted even though the remaining budget did not cover them, on priority: + // either nothing of lower priority is in flight (the path that guarantees forward + // progress) or the request is a later chunk of one already in flight. Each one is a + // place where the budget is a soft bound rather than a hard one. + priority_bypass_admissions: u64, } impl IoQueueState { @@ -134,9 +145,16 @@ impl IoQueueState { start: Instant::now(), last_warn: AtomicU64::from(0), no_backpressure: io_buffer_size == 0, + max_bytes_in_flight: 0, + priority_bypass_admissions: 0, } } + /// Bytes currently debited and not yet credited back. + fn bytes_in_flight(&self) -> u64 { + (self.io_buffer_size as i64 - self.bytes_avail).max(0) as u64 + } + fn scheduler_state_event(&self) -> Option { if !tracing::enabled!(target: SCHEDULER_STATE_EVENT_TARGET, tracing::Level::TRACE) { return None; @@ -252,7 +270,12 @@ impl IoQueueState { self.priorities_in_flight.push(task.priority); self.iops_avail -= 1; if !skip_bytes_accounting { + if task.num_bytes() as i64 > self.bytes_avail { + // `can_deliver` let this through on priority, not on budget. + self.priority_bypass_admissions += 1; + } self.bytes_avail -= task.num_bytes() as i64; + self.max_bytes_in_flight = self.max_bytes_in_flight.max(self.bytes_in_flight()); if self.bytes_avail < 0 { // This can happen when we admit special priority requests log::debug!( @@ -647,6 +670,25 @@ pub struct ScanStats { pub bytes_read: u64, } +/// What the backpressure budget actually did, as opposed to what it was configured to do. +/// +/// `max_bytes_in_flight` is the peak of bytes debited and not yet credited: the quantity +/// the budget bounds. `priority_bypass_admissions` counts the times a request was admitted +/// over budget on priority: because nothing of lower priority was in flight (which is how +/// the budget stays soft enough to guarantee forward progress) or because it was a later +/// chunk of a request already in flight. A non-zero count means the ceiling was crossed +/// deliberately. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct BackpressureStats { + pub io_buffer_size: u64, + pub max_bytes_in_flight: u64, + pub priority_bypass_admissions: u64, + /// Bytes debited and not yet credited *right now*. Unlike `max_bytes_in_flight` this + /// returns to zero once every outstanding request has been consumed or dropped, which + /// is how a caller checks that an abandoned read gave its budget back. + pub bytes_in_flight: u64, +} + impl ScanStats { fn new(stats: &StatsCollector) -> Self { Self { @@ -1128,6 +1170,25 @@ impl ScanScheduler { self.stats.snapshot() } + /// Peak backpressure usage for this scheduler; see [`BackpressureStats`]. + /// + /// Returns the default (all zero) for the lite scheduler, which accounts for + /// reservations with its own RAII type rather than this queue. + pub fn backpressure_stats(&self) -> BackpressureStats { + match &self.io_queue { + IoQueueType::Standard(queue) => { + let state = queue.state.lock().expect("io queue mutex poisoned"); + BackpressureStats { + io_buffer_size: state.io_buffer_size, + max_bytes_in_flight: state.max_bytes_in_flight, + priority_bypass_admissions: state.priority_bypass_admissions, + bytes_in_flight: state.bytes_in_flight(), + } + } + IoQueueType::Lite(_) => BackpressureStats::default(), + } + } + #[cfg(test)] fn uses_lite_scheduler(&self) -> bool { matches!(self.io_queue, IoQueueType::Lite(_)) From 5f5e98c650e5a4d272ead623f45d13f2231c8b99 Mon Sep 17 00:00:00 2001 From: Ryan Mack Date: Tue, 18 Aug 2026 15:34:46 -0400 Subject: [PATCH 5/6] test(scalar): cover the whole-file scan's budget, drop and query-path 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. --- rust/lance-index/src/scalar/btree.rs | 390 +++++++++++++++++++++++++++ 1 file changed, 390 insertions(+) diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 9c24e45a92e..7c7af9fa38f 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -3497,9 +3497,14 @@ impl ScalarIndexPlugin for BTreeIndexPlugin { #[cfg(test)] mod tests { use lance_core::utils::row_addr_remap::RowAddrRemap; + use std::pin::Pin; use std::sync::atomic::Ordering; use std::{collections::HashMap, sync::Arc}; + use async_trait::async_trait; + + use super::super::IndexReader; + use arrow::datatypes::{Float32Type, Float64Type, Int32Type, UInt64Type}; use arrow_array::{FixedSizeListArray, record_batch}; use datafusion::{ @@ -3508,6 +3513,7 @@ mod tests { }; use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_physical_expr::{PhysicalSortExpr, expressions::col}; + use futures::FutureExt; use futures::TryStreamExt; use futures::stream; use lance_core::cache::LanceCache; @@ -3765,6 +3771,390 @@ mod tests { assert_eq!(super::merge_io_buffer_size(0, 0), 32 * MIB); } + /// Build a btree index of `rows` rows in a fresh store and return the store. + /// + /// Both columns are random so the page file is incompressible and its size is + /// predictable (about 12 bytes per row), which lets a test size it against a budget. + async fn build_index_for_scan( + index_dir: Path, + object_store: Arc, + rows: u64, + batch_size: u64, + ) -> Arc { + let store = Arc::new(LanceIndexStore::new( + object_store, + index_dir, + Arc::new(LanceCache::no_cache()), + )); + let data = gen_batch() + .col("value", array::rand::()) + .col("_rowid", array::rand::()) + .into_df_exec(RowCount::from(100), BatchCount::from((rows / 100) as u32)); + let schema = data.schema(); + let sort_expr = PhysicalSortExpr::new_default(col("value", schema.as_ref()).unwrap()); + let plan = Arc::new(SortExec::new([sort_expr].into(), data)); + let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let stream = break_stream(stream, batch_size as usize); + let stream = stream.map_err(DataFusionError::from); + let stream = + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream; + train_btree_index(stream, store.as_ref(), batch_size, None, None) + .await + .unwrap(); + store + } + + /// On-disk size of the index file `name`, so a test can check its fixture is large + /// enough to exercise the budget it runs against. + async fn index_file_size(store: &LanceIndexStore, name: &str) -> u64 { + store + .list_files_with_sizes() + .await + .unwrap() + .into_iter() + .find(|file| file.path == name) + .unwrap_or_else(|| panic!("index file {name} not found")) + .size_bytes + } + + /// A whole-file scan must respect the store's byte budget. + /// + /// A whole-file decode plan schedules every page eagerly, so its outstanding bytes are + /// bounded only by the scheduler budget. The fixture is several times larger than each + /// budget plus the slack below, and the guard asserts that, so an unenforced budget + /// fails the bound instead of passing it vacuously. + /// + /// The bound is deliberately budget + `max_iop_size`: the budget is soft by design, since + /// `can_deliver` admits a request over budget when its priority is at or below everything + /// in flight, which is what guarantees forward progress, or when it is a later chunk of a + /// request already in flight. The slack is the largest single admission the scheduler + /// can make; several small over-budget admissions stay well inside it. + #[tokio::test] + async fn test_whole_file_scan_respects_the_byte_budget() { + const BATCH_SIZE: u64 = 64; + // About 23 MiB on disk; see `build_index_for_scan`. + const ROWS: u64 = 2_000_000; + const MIB: u64 = 1024 * 1024; + + let tmpdir = TempObjDir::default(); + let object_store = Arc::new(ObjectStore::local()); + let max_iop_size = object_store.max_iop_size(); + let store = + build_index_for_scan(tmpdir.clone(), object_store.clone(), ROWS, BATCH_SIZE).await; + let file_size = index_file_size(store.as_ref(), BTREE_PAGES_NAME).await; + + for budget in [MIB, 4 * MIB] { + // Fixture guard: the plan schedules the whole file, so the bound below can only + // hold because the budget throttled it. + assert!( + file_size > budget + max_iop_size, + "fixture is {file_size} bytes, too small to exercise a {budget} byte budget" + ); + + let scoped = store.with_io_buffer_size(budget); + let scoped_lance = scoped + .as_any() + .downcast_ref::() + .expect("with_io_buffer_size returns a LanceIndexStore"); + let index = BTreeIndex::load(scoped.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + + let mut stream = index.data_stream().await.unwrap(); + let mut rows_seen = 0u64; + let first = stream.try_next().await.unwrap().expect("empty scan"); + rows_seen += first.num_rows() as u64; + + // Pause the consumer and let the scheduler admit as much as the budget allows. + // Nothing is being credited back, so bytes in flight climb until the budget + // blocks admission and then stop; wait for that plateau instead of relying on + // I/O outpacing decode on this machine. + let plateau = tokio::time::timeout(std::time::Duration::from_secs(10), async { + let mut last = scoped_lance + .scheduler() + .backpressure_stats() + .bytes_in_flight; + let mut unchanged = 0; + while unchanged < 10 { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + let now = scoped_lance + .scheduler() + .backpressure_stats() + .bytes_in_flight; + unchanged = if now == last { unchanged + 1 } else { 0 }; + last = now; + } + last + }) + .await + .expect("bytes in flight never settled while the consumer was paused"); + assert!( + plateau >= budget / 2, + "with the consumer paused, a {file_size} byte scan settled at only {plateau} \ + bytes in flight against a {budget} byte budget; backpressure never engaged" + ); + + while let Some(batch) = stream.try_next().await.unwrap() { + rows_seen += batch.num_rows() as u64; + } + assert_eq!(rows_seen, ROWS); + + let bp = scoped_lance.scheduler().backpressure_stats(); + assert_eq!(bp.io_buffer_size, budget, "store was not rescoped"); + assert!( + bp.max_bytes_in_flight <= budget + max_iop_size, + "whole-file scan held {} bytes in flight against a {budget} byte budget \ + ({} priority-bypass admissions); the budget is soft but not this soft", + bp.max_bytes_in_flight, + bp.priority_bypass_admissions + ); + } + } + + /// Abandoning a scan part-way must return the whole budget. + /// + /// A whole-file plan has far more outstanding when it is dropped than a per-batch read + /// did, so any leak in the credit path is proportionally worse here. The file is + /// several times the budget, so the abandoned scan is holding a full budget at the + /// moment it is dropped. Drop at three depths, require the bytes in flight to return to + /// zero, and require that a later scan on the same store still completes. + #[tokio::test] + async fn test_dropping_a_scan_part_way_returns_the_budget() { + const BATCH_SIZE: u64 = 64; + // About 4.6 MiB on disk; see `build_index_for_scan`. + const ROWS: u64 = 400_000; + const BUDGET: u64 = 1024 * 1024; + + let tmpdir = TempObjDir::default(); + let object_store = Arc::new(ObjectStore::local()); + let store = build_index_for_scan(tmpdir.clone(), object_store, ROWS, BATCH_SIZE).await; + let file_size = index_file_size(store.as_ref(), BTREE_PAGES_NAME).await; + assert!( + file_size > 2 * BUDGET, + "fixture is {file_size} bytes, too small to hold a full {BUDGET} byte budget \ + when the scan is dropped" + ); + + for stop_after in [1usize, 32, 256] { + let scoped = store.with_io_buffer_size(BUDGET); + let scoped_lance = scoped + .as_any() + .downcast_ref::() + .expect("with_io_buffer_size returns a LanceIndexStore"); + let index = BTreeIndex::load(scoped.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + + { + let mut stream = index.data_stream().await.unwrap(); + let mut batches = 0usize; + while (stream.try_next().await.unwrap()).is_some() { + batches += 1; + if batches >= stop_after { + break; + } + } + assert_eq!(batches, stop_after, "fixture too small to stop this late"); + } + + // Requests already admitted are credited back when their responses are dropped, + // which happens asynchronously once nothing is waiting for them. + let settled = tokio::time::timeout(std::time::Duration::from_secs(10), async { + loop { + if scoped_lance + .scheduler() + .backpressure_stats() + .bytes_in_flight + == 0 + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await; + assert!( + settled.is_ok(), + "dropping the scan at batch {stop_after} left {} bytes in flight", + scoped_lance + .scheduler() + .backpressure_stats() + .bytes_in_flight + ); + + // The abandoned scan must not have stranded any of the budget: a full scan on + // the same store has to complete. + let mut stream = index.data_stream().await.unwrap(); + let mut rows_seen = 0u64; + while let Some(batch) = + tokio::time::timeout(std::time::Duration::from_secs(30), stream.try_next()) + .await + .expect("scan after an abandoned scan stalled; budget was leaked") + .unwrap() + { + rows_seen += batch.num_rows() as u64; + } + assert_eq!( + rows_seen, ROWS, + "scan after dropping at batch {stop_after} did not complete" + ); + } + } + + /// A reader that refuses to serve a whole-file scan. + struct NoWholeFileReader(Arc); + + impl std::fmt::Debug for NoWholeFileReader { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("NoWholeFileReader") + } + } + + #[async_trait] + impl IndexReader for NoWholeFileReader { + async fn read_record_batch( + &self, + n: u64, + batch_size: u64, + ) -> lance_core::Result { + self.0.read_record_batch(n, batch_size).await + } + async fn read_range( + &self, + range: std::ops::Range, + projection: Option<&[&str]>, + ) -> lance_core::Result { + self.0.read_range(range, projection).await + } + async fn whole_file_stream( + &self, + _batch_size: u32, + _batch_readahead: u32, + ) -> lance_core::Result>>> { + panic!("the search path must not take the whole-file scan path"); + } + async fn num_batches(&self, batch_size: u64) -> u32 { + self.0.num_batches(batch_size).await + } + fn num_rows(&self) -> usize { + self.0.num_rows() + } + fn schema(&self) -> &lance_core::datatypes::Schema { + self.0.schema() + } + } + + /// Wraps a store so every reader it hands out panics on a whole-file scan. + #[derive(Debug)] + struct NoWholeFileStore(Arc); + + impl DeepSizeOf for NoWholeFileStore { + fn deep_size_of_children(&self, _: &mut lance_core::deepsize::Context) -> usize { + 0 + } + } + + #[async_trait] + impl IndexStore for NoWholeFileStore { + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn clone_arc(&self) -> Arc { + Arc::new(Self(self.0.clone())) + } + fn io_parallelism(&self) -> usize { + self.0.io_parallelism() + } + fn with_io_priority(&self, io_priority: u64) -> Arc { + Arc::new(Self(self.0.with_io_priority(io_priority))) + } + fn with_io_buffer_size(&self, bytes: u64) -> Arc { + Arc::new(Self(self.0.with_io_buffer_size(bytes))) + } + async fn new_index_file( + &self, + name: &str, + schema: Arc, + ) -> lance_core::Result> { + self.0.new_index_file(name, schema).await + } + async fn open_index_file(&self, name: &str) -> lance_core::Result> { + let inner = self.0.open_index_file(name).await?; + Ok(Arc::new(NoWholeFileReader(inner))) + } + async fn copy_index_file( + &self, + name: &str, + dest_store: &dyn IndexStore, + ) -> lance_core::Result { + self.0.copy_index_file(name, dest_store).await + } + async fn rename_index_file( + &self, + name: &str, + new_name: &str, + ) -> lance_core::Result { + self.0.rename_index_file(name, new_name).await + } + async fn delete_index_file(&self, name: &str) -> lance_core::Result<()> { + self.0.delete_index_file(name).await + } + async fn list_files_with_sizes(&self) -> lance_core::Result> { + self.0.list_files_with_sizes().await + } + } + + /// The query path must not reach the whole-file scan. + /// + /// `whole_file_stream` is wired into three maintenance call sites; searches do point + /// lookups through `read_range`. Assert that structurally by handing the index a reader + /// that panics if anyone asks it for a whole-file stream, so a future change that routes + /// a search through one fails here instead of in production. + #[tokio::test] + async fn test_search_path_never_takes_the_whole_file_scan() { + const BATCH_SIZE: u64 = 64; + const ROWS: u64 = 40_000; + + let tmpdir = TempObjDir::default(); + let object_store = Arc::new(ObjectStore::local()); + let store = build_index_for_scan(tmpdir.clone(), object_store, ROWS, BATCH_SIZE).await; + + let guarded: Arc = Arc::new(NoWholeFileStore(store)); + let index = BTreeIndex::load(guarded, None, &LanceCache::no_cache()) + .await + .unwrap(); + + // Positive control: a maintenance scan through the same store must hit the guard. + // Without this, the searches below could pass simply because the guard is never + // consulted, and the test would prove nothing. + let reached = std::panic::AssertUnwindSafe(async { + let mut scan = index.data_stream().await.unwrap(); + while (scan.try_next().await.unwrap()).is_some() {} + }) + .catch_unwind() + .await + .is_err(); + assert!( + reached, + "the guard was never consulted; this test is vacuous" + ); + + // Equality, range and IS NULL: the three shapes a btree serves. + for query in [ + SargableQuery::Equals(ScalarValue::Int32(Some(17))), + SargableQuery::Range( + std::ops::Bound::Included(ScalarValue::Int32(Some(10))), + std::ops::Bound::Excluded(ScalarValue::Int32(Some(4000))), + ), + SargableQuery::IsNull(), + ] { + index + .search(&query, &NoOpMetricsCollector) + .await + .expect("search failed"); + } + } + /// `batch_size` is read from the index file's own metadata, so a value that is a /// multiple of 2^32 used to truncate to a `rows_per_batch` of 0 — an empty stream that /// silently retrained an empty index and reported success. It must clamp instead. From 6e6b50e842fb5c91a4499f268038e3d9d6e593d3 Mon Sep 17 00:00:00 2001 From: Ryan Mack Date: Tue, 18 Aug 2026 16:43:27 -0400 Subject: [PATCH 6/6] test(scalar): assert merge_segments hands each source a 1/N prefetch 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. --- rust/lance-index/src/scalar/btree.rs | 250 +++++++++++++++++++++++++-- 1 file changed, 237 insertions(+), 13 deletions(-) diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 7c7af9fa38f..10594c252ec 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -3771,10 +3771,26 @@ mod tests { assert_eq!(super::merge_io_buffer_size(0, 0), 32 * MIB); } - /// Build a btree index of `rows` rows in a fresh store and return the store. + /// A sorted `(value, _rowid)` stream of `rows` rows, the shape `train_btree_index` + /// and `merge_segments` both consume. /// /// Both columns are random so the page file is incompressible and its size is /// predictable (about 12 bytes per row), which lets a test size it against a budget. + fn sorted_training_stream(rows: u64, batch_size: u64) -> SendableRecordBatchStream { + let data = gen_batch() + .col("value", array::rand::()) + .col("_rowid", array::rand::()) + .into_df_exec(RowCount::from(100), BatchCount::from((rows / 100) as u32)); + let schema = data.schema(); + let sort_expr = PhysicalSortExpr::new_default(col("value", schema.as_ref()).unwrap()); + let plan = Arc::new(SortExec::new([sort_expr].into(), data)); + let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let stream = break_stream(stream, batch_size as usize); + let stream = stream.map_err(DataFusionError::from); + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream + } + + /// Build a btree index of `rows` rows in a fresh store and return the store. async fn build_index_for_scan( index_dir: Path, object_store: Arc, @@ -3786,18 +3802,7 @@ mod tests { index_dir, Arc::new(LanceCache::no_cache()), )); - let data = gen_batch() - .col("value", array::rand::()) - .col("_rowid", array::rand::()) - .into_df_exec(RowCount::from(100), BatchCount::from((rows / 100) as u32)); - let schema = data.schema(); - let sort_expr = PhysicalSortExpr::new_default(col("value", schema.as_ref()).unwrap()); - let plan = Arc::new(SortExec::new([sort_expr].into(), data)); - let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); - let stream = break_stream(stream, batch_size as usize); - let stream = stream.map_err(DataFusionError::from); - let stream = - Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream; + let stream = sorted_training_stream(rows, batch_size); train_btree_index(stream, store.as_ref(), batch_size, None, None) .await .unwrap(); @@ -4155,6 +4160,225 @@ mod tests { } } + /// Records every store `with_io_buffer_size` hands out, so a test can reach the + /// schedulers `merge_segments` actually drained. All other calls pass through. + #[derive(Debug)] + struct RescopeRecordingStore { + inner: Arc, + rescoped: Arc>>>, + } + + impl DeepSizeOf for RescopeRecordingStore { + fn deep_size_of_children(&self, _: &mut lance_core::deepsize::Context) -> usize { + 0 + } + } + + #[async_trait] + impl IndexStore for RescopeRecordingStore { + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn clone_arc(&self) -> Arc { + Arc::new(Self { + inner: self.inner.clone(), + rescoped: self.rescoped.clone(), + }) + } + fn io_parallelism(&self) -> usize { + self.inner.io_parallelism() + } + fn with_io_priority(&self, io_priority: u64) -> Arc { + Arc::new(Self { + inner: self.inner.with_io_priority(io_priority), + rescoped: self.rescoped.clone(), + }) + } + fn with_io_buffer_size(&self, bytes: u64) -> Arc { + let scoped = self.inner.with_io_buffer_size(bytes); + self.rescoped.lock().unwrap().push(scoped.clone()); + scoped + } + async fn new_index_file( + &self, + name: &str, + schema: Arc, + ) -> lance_core::Result> { + self.inner.new_index_file(name, schema).await + } + async fn open_index_file(&self, name: &str) -> lance_core::Result> { + self.inner.open_index_file(name).await + } + async fn copy_index_file( + &self, + name: &str, + dest_store: &dyn IndexStore, + ) -> lance_core::Result { + self.inner.copy_index_file(name, dest_store).await + } + async fn rename_index_file( + &self, + name: &str, + new_name: &str, + ) -> lance_core::Result { + self.inner.rename_index_file(name, new_name).await + } + async fn delete_index_file(&self, name: &str) -> lance_core::Result<()> { + self.inner.delete_index_file(name).await + } + async fn list_files_with_sizes(&self) -> lance_core::Result> { + self.inner.list_files_with_sizes().await + } + } + + /// `merge_segments` must bound its aggregate prefetch across N sources. + /// + /// Every source's whole-file plan starts prefetching the moment it is built, each + /// against its own store, so an N-way merge holds N budgets of encoded bytes unless + /// something splits them. `test_whole_file_scan_respects_the_byte_budget` shows one + /// scan respects one store's budget; this shows the merge hands every source a 1/N + /// share, drains all N through those rescoped stores and nothing else, and still + /// merges correctly. Together the two bound the aggregate at + /// `max(one scan's budget, N * 32 MiB)` plus one split request per source. + /// + /// The per-source budgets here (32 MiB and up) are far larger than the 40k-row + /// fixture, so this checks the split and the plumbing, not throttling under it. + #[tokio::test] + async fn test_merge_segments_bounds_aggregate_prefetch() { + const BATCH_SIZE: u64 = 64; + const ROWS: u64 = 40_000; + const NEW_ROWS: u64 = 1_000; + const MIB: u64 = 1024 * 1024; + const MAX_IOP_SIZE: u64 = 16 * MIB; + + for num_segments in [2usize, 4, 8] { + let object_store = Arc::new(ObjectStore::local()); + let io_parallelism = object_store.io_parallelism(); + let expected_budget = super::merge_io_buffer_size(io_parallelism, num_segments as u64); + + let mut tmpdirs = Vec::new(); + let mut originals = Vec::new(); + let mut recorders = Vec::new(); + let mut segments = Vec::new(); + for _ in 0..num_segments { + let tmpdir = TempObjDir::default(); + let store = + build_index_for_scan(tmpdir.clone(), object_store.clone(), ROWS, BATCH_SIZE) + .await; + let recorder = Arc::new(RescopeRecordingStore { + inner: store.clone(), + rescoped: Arc::new(std::sync::Mutex::new(Vec::new())), + }); + let index = BTreeIndex::load( + recorder.clone() as Arc, + None, + &LanceCache::no_cache(), + ) + .await + .unwrap(); + segments.push(index); + originals.push(store); + recorders.push(recorder); + tmpdirs.push(tmpdir); + } + + // Everything the merge reads must go through the rescoped stores, so the + // original schedulers have to end the merge exactly where they start it. + let before: Vec<_> = originals + .iter() + .map(|store| store.scheduler().backpressure_stats().max_bytes_in_flight) + .collect(); + + let dest_dir = TempObjDir::default(); + let dest_store = Arc::new(LanceIndexStore::new( + object_store.clone(), + dest_dir.clone(), + Arc::new(LanceCache::no_cache()), + )); + let filters: Vec> = vec![None; num_segments]; + BTreeIndex::merge_segments( + &segments, + sorted_training_stream(NEW_ROWS, BATCH_SIZE), + dest_store.as_ref(), + &filters, + ) + .await + .unwrap(); + + // Each source was rescoped exactly once, to a 1/N share of one scan's budget. + let mut aggregate_peak = 0u64; + for (i, recorder) in recorders.iter().enumerate() { + let handed_out = recorder.rescoped.lock().unwrap(); + assert_eq!( + handed_out.len(), + 1, + "segment {i} was rescoped {} times, expected exactly once", + handed_out.len() + ); + let scoped = handed_out[0] + .as_any() + .downcast_ref::() + .expect("rescoping a LanceIndexStore yields a LanceIndexStore"); + let bp = scoped.scheduler().backpressure_stats(); + assert_eq!( + bp.io_buffer_size, expected_budget, + "segment {i} of {num_segments} got budget {}, expected {}", + bp.io_buffer_size, expected_budget + ); + // Without this the bounds below pass vacuously if a stream never ran. + assert!( + bp.max_bytes_in_flight > 0, + "segment {i} recorded no bytes in flight; the bound proves nothing" + ); + assert!( + bp.max_bytes_in_flight <= expected_budget + MAX_IOP_SIZE, + "segment {i} held {} bytes against a {} byte budget \ + ({} priority-bypass admissions)", + bp.max_bytes_in_flight, + expected_budget, + bp.priority_bypass_admissions + ); + aggregate_peak += bp.max_bytes_in_flight; + } + + // The documented aggregate bound. Summing per-store peaks overstates the true + // simultaneous peak (they need not coincide), so this can only fail on the + // conservative side. + let one_scan = super::MERGE_BYTES_PER_IO_THREAD * io_parallelism.max(1) as u64; + let bound = one_scan.max(num_segments as u64 * super::MERGE_BYTES_PER_IO_THREAD) + + num_segments as u64 * MAX_IOP_SIZE; + assert!( + aggregate_peak <= bound, + "merge of {num_segments} segments held {aggregate_peak} bytes aggregate, \ + documented bound is {bound}" + ); + + for (i, (store, before)) in originals.iter().zip(&before).enumerate() { + let after = store.scheduler().backpressure_stats().max_bytes_in_flight; + assert_eq!( + after, *before, + "segment {i}'s original scheduler saw merge traffic; the budget \ + split does not cover whatever read this" + ); + } + + // And the merge has to have actually merged. + let merged = BTreeIndex::load( + dest_store.clone() as Arc, + None, + &LanceCache::no_cache(), + ) + .await + .unwrap(); + let mut stream = merged.data_stream().await.unwrap(); + let mut rows = 0u64; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows() as u64; + } + assert_eq!(rows, num_segments as u64 * ROWS + NEW_ROWS); + } + } + /// `batch_size` is read from the index file's own metadata, so a value that is a /// multiple of 2^32 used to truncate to a `rows_per_batch` of 0 — an empty stream that /// silently retrained an empty index and reported success. It must clamp instead.