diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index 80da1633a26..ed3f2c94b59 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -223,6 +223,45 @@ 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. + /// + /// # 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, + _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 @@ -272,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 cb2c842b906..10594c252ec 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -57,7 +57,7 @@ use datafusion_physical_expr::{ use futures::{ FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::BoxFuture, - stream::{self}, + stream::{self, BoxStream}, }; use lance_core::deepsize::DeepSizeOf; use lance_core::{ @@ -1842,16 +1842,31 @@ 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?; 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(), @@ -1909,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), @@ -1978,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 { @@ -2110,9 +2157,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 +2358,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,6 +2869,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?; + // 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 { @@ -2978,8 +3037,8 @@ async fn merge_pages( let reader = store.open_index_file(&page_file_name).await?; + // 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)) @@ -3161,6 +3220,33 @@ 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>> { + // 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_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 @@ -3411,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::{ @@ -3422,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; @@ -3587,6 +3679,751 @@ 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; + // 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()); + 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(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)); + 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: 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 < 64, + "scanning a {num_pages}-batch index took {iops} read IOPs; a decode plan \ + rebuilt per batch costs about two per batch" + ); + } + + /// 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); + } + + /// 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, + rows: u64, + batch_size: u64, + ) -> Arc { + let store = Arc::new(LanceIndexStore::new( + object_store, + index_dir, + Arc::new(LanceCache::no_cache()), + )); + let stream = sorted_training_stream(rows, batch_size); + 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"); + } + } + + /// 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. + #[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(); diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index 69e18fc852c..77c3341dd1c 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), @@ -111,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}")) @@ -377,6 +388,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 { @@ -419,6 +456,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, 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(_))