From b731b77fc3f2b44739c824adeb1edad40b21543b Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Sat, 10 Oct 2026 23:08:30 +0800 Subject: [PATCH] feat(index): add origin-aware reads and shared resident columns --- rust/lance-core/src/cache/backend.rs | 8 + rust/lance-core/src/cache/mod.rs | 12 + rust/lance-core/src/cache/quick.rs | 11 + rust/lance-index/src/vector/bq.rs | 1 + rust/lance-index/src/vector/bq/ex_dot.rs | 4 +- rust/lance-index/src/vector/bq/layered.rs | 292 +- .../src/vector/bq/layered_stats.rs | 105 +- rust/lance-index/src/vector/bq/lazy_tests.rs | 4 +- rust/lance-index/src/vector/bq/plane_cache.rs | 60 +- rust/lance-index/src/vector/bq/resident.rs | 1014 +++++ rust/lance-index/src/vector/bq/storage.rs | 29 +- rust/lance-index/src/vector/storage.rs | 1842 +++++++- rust/lance/src/index.rs | 30 +- rust/lance/src/index/vector/ivf/v2.rs | 3789 ++++++++++++++++- .../src/index/vector/ivf/v2/lazy_full.rs | 380 +- rust/lance/src/session.rs | 82 + 16 files changed, 7270 insertions(+), 393 deletions(-) create mode 100644 rust/lance-index/src/vector/bq/resident.rs diff --git a/rust/lance-core/src/cache/backend.rs b/rust/lance-core/src/cache/backend.rs index cae4cc119ec..ab0527c311d 100644 --- a/rust/lance-core/src/cache/backend.rs +++ b/rust/lance-core/src/cache/backend.rs @@ -162,6 +162,14 @@ pub trait CacheBackend: Send + Sync + std::fmt::Debug { codec: Option, ) -> Result<(CacheEntry, bool)>; + /// Largest entry size the backend admits to RAM, in bytes. Defaults to + /// the total capacity; sharded backends may report a smaller limit. + /// `None` means the backend does not expose an admission limit. + /// Admission never guarantees that an entry will remain cached. + fn max_entry_bytes(&self) -> Option { + self.capacity_bytes().map(|bytes| bytes as u64) + } + /// Remove all entries. async fn clear(&self); diff --git a/rust/lance-core/src/cache/mod.rs b/rust/lance-core/src/cache/mod.rs index b6b22aea548..559a6d90c18 100644 --- a/rust/lance-core/src/cache/mod.rs +++ b/rust/lance-core/src/cache/mod.rs @@ -348,6 +348,18 @@ impl LanceCache { self.state.backend.capacity_bytes() } + /// Largest admissible RAM entry, if reported by the backend. An admitted + /// entry remains subject to ordinary cache eviction. + /// + /// ``` + /// use lance_core::cache::LanceCache; + /// let cache = LanceCache::with_capacity(1024); + /// assert_eq!(cache.max_entry_bytes(), Some(1024)); + /// ``` + pub fn max_entry_bytes(&self) -> Option { + self.state.backend.max_entry_bytes() + } + // -- Stats / clear -------------------------------------------------------- pub async fn stats(&self) -> CacheStats { diff --git a/rust/lance-core/src/cache/quick.rs b/rust/lance-core/src/cache/quick.rs index 20ce7b7766a..fcc4006b887 100644 --- a/rust/lance-core/src/cache/quick.rs +++ b/rust/lance-core/src/cache/quick.rs @@ -330,6 +330,10 @@ impl CacheBackend for QuickCacheBackend { Some(self.capacity) } + fn max_entry_bytes(&self) -> Option { + Some((self.capacity / self.cache.num_shards()) as u64) + } + fn approx_num_entries(&self) -> usize { self.cache.len() + self @@ -598,6 +602,13 @@ mod tests { QuickCacheBackend::with_capacity(1 << 20).capacity_bytes(), Some(1 << 20) ); + for capacity in [0, 1 << 20, 16 << 30] { + let backend = QuickCacheBackend::with_capacity(capacity); + assert_eq!( + backend.max_entry_bytes(), + Some((capacity / backend.cache.num_shards()) as u64) + ); + } } fn keys_in_shard( diff --git a/rust/lance-index/src/vector/bq.rs b/rust/lance-index/src/vector/bq.rs index d5bfec5a431..e1fb3927d83 100644 --- a/rust/lance-index/src/vector/bq.rs +++ b/rust/lance-index/src/vector/bq.rs @@ -27,6 +27,7 @@ mod lazy_tests; pub(crate) mod pairwise; mod plane_cache; pub mod prune; +pub mod resident; pub mod rotation; pub mod storage; pub mod transform; diff --git a/rust/lance-index/src/vector/bq/ex_dot.rs b/rust/lance-index/src/vector/bq/ex_dot.rs index e922f1c2088..f5b8334ad8a 100644 --- a/rust/lance-index/src/vector/bq/ex_dot.rs +++ b/rust/lance-index/src/vector/bq/ex_dot.rs @@ -490,7 +490,7 @@ fn select_prefix_kernel() -> ExDotLayeredFn { fn prefix_scalar(q: &[f32], hi: &[u8], lo: &[u8]) -> f32 { let mut sum = 0.0; - for (block, query) in q.chunks_exact(64).enumerate() { + for (block, query) in q.as_chunks::<64>().0.iter().enumerate() { let mut high = [0; 64]; let mut low = [0; 64]; unpack_group(2, &lo[block * 16..(block + 1) * 16], &mut low); @@ -566,7 +566,7 @@ fn ex_dot_layered_scalar(ex_query: &[f32], hi: &[u8], lo: &[u8]) -> f32 { let mut sum = 0.0f32; let mut hi_codes = [0u8; 64]; let mut lo_codes = [0u8; 64]; - for (group, query) in ex_query.chunks_exact(16).enumerate() { + for (group, query) in ex_query.as_chunks::<16>().0.iter().enumerate() { let start = group * LAYERED_GROUP_BYTES; if start + LAYERED_GROUP_BYTES > hi.len() { break; diff --git a/rust/lance-index/src/vector/bq/layered.rs b/rust/lance-index/src/vector/bq/layered.rs index 753e3fbc509..4e16d79dcd3 100644 --- a/rust/lance-index/src/vector/bq/layered.rs +++ b/rust/lance-index/src/vector/bq/layered.rs @@ -33,8 +33,13 @@ impl RQLayout { } use super::ex_dot::{blocked_ex_code_bytes, pack_blocked_row}; -use super::storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN}; -use super::transform::{EX_ADD_FACTORS_FIELD, EX_SCALE_FACTORS_FIELD}; +use super::storage::{ + RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN, RABIT_CODE_COLUMN, +}; +use super::transform::{ + ADD_FACTORS_COLUMN, ERROR_FACTORS_COLUMN, EX_ADD_FACTORS_COLUMN, EX_ADD_FACTORS_FIELD, + EX_SCALE_FACTORS_COLUMN, EX_SCALE_FACTORS_FIELD, SCALE_FACTORS_COLUMN, +}; use arrow_array::{ArrayRef, FixedSizeListArray, Float32Array, UInt8Array}; use arrow_schema::{DataType, Field}; use lance_arrow::FixedSizeListArrayExt; @@ -227,15 +232,67 @@ impl DeepSizeOf for PlaneBatch { self.0.deep_size_of_children(context) } } + +/// Plane entry holding a layered partition's [`HIGH_BOUNDS_COLUMN`] and +/// [`FULL_BOUNDS_COLUMN`] under [`SignBounds::Lazy`]. Only High precision, +/// and full precision on a file without error factors, prune with them. +pub const SIGN_BOUNDS_PLANE: u8 = 3; + +/// Memory priority of the sign plane. Priority-aware tiers keep the sign +/// plane over the high plane over the low plane. +const SIGN_PLANE_MEMORY_PRIORITY: u8 = 3; + +/// Where a layered index's cache keeps the estimator bounds columns +/// ([`HIGH_BOUNDS_COLUMN`], [`FULL_BOUNDS_COLUMN`]). The placement decides +/// what the sign plane entry holds, so it is part of that entry's key. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum SignBounds { + /// In their own plane entry, [`SIGN_BOUNDS_PLANE`], read only by the + /// scans that prune with them. Full-precision scans prune with the + /// native error factors instead, so they read them only on a file + /// without error factors. + #[default] + Lazy, + /// In the sign plane entry, so every scan reads them with the sign + /// codes. There is no bounds plane. + Eager, +} + +impl SignBounds { + /// The spelling of `LANCE_RQ_SIGN_BOUNDS`: `lazy` or `eager`. + pub fn as_str(self) -> &'static str { + match self { + Self::Lazy => "lazy", + Self::Eager => "eager", + } + } +} + +impl std::fmt::Display for SignBounds { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + /// Plane keys live under the immutable index UUID namespace. pub struct PlaneKey { pub partition: usize, pub plane: u8, + /// The index's bounds placement. Only the sign plane's key depends on + /// it: the other planes hold the same columns under either placement. + pub sign_bounds: SignBounds, } impl CacheKey for PlaneKey { type ValueType = PlaneBatch; fn key(&self) -> Cow<'_, str> { - format!("{}:{}", self.partition, self.plane).into() + // Persisted entries outlive the process, so each sign-plane + // composition needs its own key. The sign plane with bounds keeps + // the key its entries had before the bounds plane existed. + if self.plane == 0 && self.sign_bounds == SignBounds::Lazy { + format!("{}:0-bounds", self.partition).into() + } else { + format!("{}:{}", self.partition, self.plane).into() + } } fn type_name() -> &'static str { "RQPlane" @@ -244,7 +301,7 @@ impl CacheKey for PlaneKey { Self::codec().map(|codec| { codec .with_plane_tag(self.plane) - .with_memory_priority(3 - self.plane) + .with_memory_priority(plane_memory_priority(self.plane)) }) } fn codec() -> Option { @@ -252,33 +309,50 @@ impl CacheKey for PlaneKey { } } -/// Column projection for sign (0), high (1), or low (2), including factors. -pub fn plane_columns(plane: u8) -> &'static [&'static str] { - use super::storage::RABIT_CODE_COLUMN; - use super::transform::{ - ADD_FACTORS_COLUMN, ERROR_FACTORS_COLUMN, EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN, - SCALE_FACTORS_COLUMN, - }; +/// The bounds plane is read with the high plane, by High precision, so it +/// ranks with it. +fn plane_memory_priority(plane: u8) -> u8 { match plane { - 0 => &[ - lance_core::ROW_ID, - RABIT_CODE_COLUMN, - ADD_FACTORS_COLUMN, - SCALE_FACTORS_COLUMN, - ERROR_FACTORS_COLUMN, - HIGH_BOUNDS_COLUMN, - FULL_BOUNDS_COLUMN, - ], - 1 => &[ + SIGN_BOUNDS_PLANE => SIGN_PLANE_MEMORY_PRIORITY - 1, + _ => SIGN_PLANE_MEMORY_PRIORITY.saturating_sub(plane), + } +} + +/// The sign plane's columns under [`SignBounds::Eager`]: those it keeps +/// under [`SignBounds::Lazy`], then the bounds plane's. +static SIGN_PLANE_WITH_BOUNDS_COLUMNS: [&str; 7] = [ + lance_core::ROW_ID, + RABIT_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + HIGH_BOUNDS_COLUMN, + FULL_BOUNDS_COLUMN, +]; +/// Where the bounds columns start in [`SIGN_PLANE_WITH_BOUNDS_COLUMNS`]. +const BOUNDS_COLUMNS_START: usize = 5; + +/// Column projection of a layered partition's plane entry, including its +/// factors: sign (0), high (1), low (2) or [`SIGN_BOUNDS_PLANE`], with the +/// bounds columns placed by `sign_bounds`. Empty for a plane the placement +/// does not have. +pub fn plane_columns(plane: u8, sign_bounds: SignBounds) -> &'static [&'static str] { + match (plane, sign_bounds) { + (0, SignBounds::Lazy) => &SIGN_PLANE_WITH_BOUNDS_COLUMNS[..BOUNDS_COLUMNS_START], + (0, SignBounds::Eager) => &SIGN_PLANE_WITH_BOUNDS_COLUMNS, + (1, _) => &[ RABIT_BLOCKED_EX_CODE_COLUMN, HIGH_ADD_FACTORS_COLUMN, HIGH_SCALE_FACTORS_COLUMN, ], - 2 => &[ + (2, _) => &[ RABIT_BLOCKED_EX_CODE_LO_COLUMN, EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN, ], + (SIGN_BOUNDS_PLANE, SignBounds::Lazy) => { + &SIGN_PLANE_WITH_BOUNDS_COLUMNS[BOUNDS_COLUMNS_START..] + } _ => &[], } } @@ -654,20 +728,190 @@ mod tests { #[test] fn plane_key_codec_tags_plane_and_priority() { - for plane in 0..=2u8 { + for (plane, priority) in [(0, 3), (1, 2), (2, 1), (SIGN_BOUNDS_PLANE, 2)] { let codec = PlaneKey { partition: 7, plane, + sign_bounds: SignBounds::Lazy, } .codec_for_key() .unwrap(); assert_eq!(codec.plane_tag(), Some(plane)); - assert_eq!(codec.memory_priority(), 3 - plane); + assert_eq!(codec.memory_priority(), priority); assert!(codec.supports_row_selection()); } assert_eq!(PlaneKey::codec().unwrap().plane_tag(), None); } + /// Only the sign plane's key depends on the bounds placement, and the + /// sign plane with bounds keeps its original key. + #[test] + fn plane_key_separates_sign_plane_compositions() { + let key = |plane, sign_bounds| { + let plane_key = PlaneKey { + partition: 7, + plane, + sign_bounds, + }; + plane_key.key().into_owned() + }; + assert_eq!(key(0, SignBounds::Eager), "7:0"); + assert_ne!(key(0, SignBounds::Lazy), key(0, SignBounds::Eager)); + for plane in [1, 2, SIGN_BOUNDS_PLANE] { + assert_eq!(key(plane, SignBounds::Lazy), format!("7:{plane}")); + assert_eq!(key(plane, SignBounds::Eager), format!("7:{plane}")); + } + } + + /// The bounds plane holds exactly the columns the lazy placement leaves + /// out of the sign plane, in the eager sign plane's order. + #[test] + fn plane_columns_place_bounds_by_sign_bounds() { + let lazy_sign = plane_columns(0, SignBounds::Lazy); + let bounds = plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Lazy); + assert_eq!(bounds, [HIGH_BOUNDS_COLUMN, FULL_BOUNDS_COLUMN]); + assert_eq!( + [lazy_sign, bounds].concat(), + plane_columns(0, SignBounds::Eager) + ); + assert!(!lazy_sign.iter().any(|name| bounds.contains(name))); + assert!(plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Eager).is_empty()); + for plane in [1, 2] { + assert_eq!( + plane_columns(plane, SignBounds::Lazy), + plane_columns(plane, SignBounds::Eager) + ); + assert_eq!(plane_columns(plane, SignBounds::Lazy).len(), 3); + } + assert!(plane_columns(SIGN_BOUNDS_PLANE + 1, SignBounds::Lazy).is_empty()); + } + + /// A layered store needs the full bounds only to prune full precision on + /// a file without error factors. Without the high bounds, a High scan of + /// a layered store scores every row instead of pruning with the error + /// factors, which bound only the full-precision estimator. + #[rstest] + fn layered_bounds_required_only_when_read( + #[values(DistanceType::L2, DistanceType::Dot)] distance_type: DistanceType, + ) { + const ROWS: usize = 37; + const DIM: usize = 64; + let values: Vec = (0..DIM * ROWS) + .map(|i| ((i * 13 % 97) as f32 - 48.) / 48.) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values.clone()), DIM as i32) + .unwrap(); + let rq = RabitQuantizer::build( + &vectors, + distance_type, + &RQBuildParams::new(7).with_layered(true), + ) + .unwrap(); + let norms = Float32Array::from( + values + .chunks(DIM) + .map(|r| r.iter().map(|v| v * v).sum::()) + .collect::>(), + ); + let input = RecordBatch::try_from_iter(vec![ + ("vector", Arc::new(vectors) as ArrayRef), + ( + lance_core::ROW_ID, + Arc::new(UInt64Array::from((0..ROWS as u64).collect::>())) as ArrayRef, + ), + ( + PART_ID_COLUMN, + Arc::new(UInt32Array::from(vec![0; ROWS])) as ArrayRef, + ), + (CENTROID_DIST_COLUMN, Arc::new(norms) as ArrayRef), + ]) + .unwrap(); + let centroids = + FixedSizeListArray::try_new_from_values(Float32Array::from(vec![0.0; DIM]), DIM as i32) + .unwrap(); + let batch = RQTransformer::new(rq.clone(), distance_type, centroids, "vector") + .unwrap() + .transform(&input) + .unwrap(); + let metadata = rq.metadata_ref(); + let open = |batch: RecordBatch| { + RabitQuantizationStorage::try_from_batch(batch, metadata, distance_type, None) + }; + let lean_batch = batch + .drop_column(HIGH_BOUNDS_COLUMN) + .unwrap() + .drop_column(FULL_BOUNDS_COLUMN) + .unwrap(); + let bounded = open(batch.clone()).unwrap(); + let lean = open(lean_batch.clone()).unwrap(); + let query: ArrayRef = Arc::new(Float32Array::from(values[..DIM].to_vec())); + let dist_q_c = match distance_type { + DistanceType::L2 => values[..DIM].iter().map(|v| v * v).sum(), + _ => 1.0, + }; + fn lower_bounds(calc: &super::super::storage::RabitDistCalculator<'_>) -> Vec> { + calc.binary_inner_products() + .into_iter() + .enumerate() + .map(|(row, ip)| calc.raw_query_lower_bound(row, ip).map(f32::to_bits)) + .collect() + } + for approx_mode in [ApproxMode::Normal, ApproxMode::Accurate] { + for rq_precision in [RQPrecision::High, RQPrecision::Full] { + let options = DistanceCalculatorOptions { + approx_mode, + rq_precision, + }; + let case = format!("{distance_type:?} {approx_mode:?} {rq_precision:?}"); + let (mut bounded_scratch, mut lean_scratch) = (Vec::new(), Vec::new()); + let bounded_calc = bounded.dist_calculator_with_scratch( + query.clone(), + dist_q_c, + None, + &mut bounded_scratch, + options, + ); + let lean_calc = lean.dist_calculator_with_scratch( + query.clone(), + dist_q_c, + None, + &mut lean_scratch, + options, + ); + assert_eq!( + lean_calc.distance_all(ROWS), + bounded_calc.distance_all(ROWS), + "{case}" + ); + let bounded_bounds = lower_bounds(&bounded_calc); + assert!(bounded_bounds.iter().all(Option::is_some), "{case}"); + if rq_precision == RQPrecision::Full { + assert_eq!(lower_bounds(&lean_calc), bounded_bounds, "{case}"); + } else { + assert!( + lower_bounds(&lean_calc).iter().all(Option::is_none), + "{case}" + ); + } + } + } + + let error = open(lean_batch.drop_column(ERROR_FACTORS_COLUMN).unwrap()).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains("missing column"), "{error}"); + assert!(message.contains(FULL_BOUNDS_COLUMN), "{error}"); + open( + batch + .drop_column(ERROR_FACTORS_COLUMN) + .unwrap() + .drop_column(HIGH_BOUNDS_COLUMN) + .unwrap(), + ) + .unwrap(); + } + #[rstest] #[case(1)] #[case(2)] diff --git a/rust/lance-index/src/vector/bq/layered_stats.rs b/rust/lance-index/src/vector/bq/layered_stats.rs index cb3df0d3074..4806a9d5d6e 100644 --- a/rust/lance-index/src/vector/bq/layered_stats.rs +++ b/rust/lance-index/src/vector/bq/layered_stats.rs @@ -5,7 +5,10 @@ //! //! The counters only advance while a query runs the lazy scan (or is checked //! for it), so production queries with the scan disabled touch at most the -//! `ineligible_disabled` counter. Readers take deltas with +//! `ineligible_disabled` counter. The `resident_columns_*`, +//! `resident_store_evictions` counters are the +//! exception: they advance when an IVF_RQ index, layered or not, opens with +//! or loads its resident columns. Readers take deltas with //! [`snapshot_and_reset`]. use std::sync::atomic::{AtomicU64, Ordering}; @@ -164,6 +167,18 @@ layered_lazy_counters! { /// Sparse gathers whose rows came from the origin file because the /// cache had no persistent entry for the plane. origin_row_reads, + /// Origin requests of those reads after coalescing within the lazy + /// origin gap (`LANCE_RQ_LAZY_ORIGIN_GAP_BYTES`) and splitting, which + /// are GETs on an object store. A load of the resident columns that + /// such a read starts is counted in `resident_columns_load_requests` + /// instead. + origin_sparse_requests, + /// Bytes those requests read, the gaps they span included. + origin_sparse_bytes, + /// Rows those reads gathered, which the gap does not change. Over one + /// plane, `origin_sparse_bytes / (origin_sparse_rows * row width)` is + /// the byte amplification of the gap. + origin_sparse_rows, /// Sum over gathered probes of probes still unscored at issue. staleness_sum, /// Issues deferred because the heap was not yet full at the gate. @@ -213,6 +228,77 @@ layered_lazy_counters! { /// Bytes of planes loaded by promotions, resident afterwards or not, /// apart from critical-path reads. promotion_bytes, + /// Bytes of values each load of an index's resident columns keeps in + /// memory (see `LANCE_RQ_RESIDENT_COLUMNS`). An index loads them + /// once, so over every snapshot this sums to the store's size. + resident_columns_bytes, + /// Memory the arrays of those loads hold: the capacity of their + /// buffers, counted per column. It exceeds `resident_columns_bytes` + /// when an array keeps a buffer larger than its values. + resident_columns_alloc_bytes, + /// Origin requests of those loads after coalescing and splitting, + /// which are GETs on an object store. + resident_columns_load_requests, + /// Bytes those loads read from the origin. + resident_columns_load_bytes, + /// Time those loads took. + resident_columns_load_ns, + /// Classifying staged probes' ex planes by the cache tier that holds + /// them (only RAM residency on a low-latency origin). + tier_peek_ns, + /// Gathers issued beyond the ordinary window with a permit of their + /// index's pool, because their probe reads an ex plane from a + /// high-latency origin; see `LANCE_RQ_LAZY_FAR_WINDOW`. + far_early_issues, + /// Gathers beyond the ordinary window that found no free permit and + /// waited for one or for the ordinary window, whichever came first. + far_permit_waits, + /// Lazy scans that published the scoring of their first probe; see + /// `lazy_rank0_scored_ns`. + lazy_rank0_scored_queries, + /// Time from the start of each of those scans until it published the + /// scoring of its first probe (rank 0) to the gathers: once the probe + /// was scored or, earlier, once its rows filled the heap. No gather + /// that waits for the threshold or its turn is released before. + lazy_rank0_scored_ns, + /// Lazy scans that issued a gather of a probe that is not + /// `certain_dense`, whose issue waits on scoring progress; see + /// `lazy_first_gather_issue_ns`. Certain-dense gathers, the first + /// probe's always among them, are issued as soon as staged. + lazy_first_gather_issue_queries, + /// Time from the start of each of those scans until the first such + /// gather was issued, its gate waits and permit included. + lazy_first_gather_issue_ns, + /// Gathers whose probe was further ahead of scoring than their + /// staleness window (the far window for a probe that reads an ex + /// plane from a high-latency origin) when they reached the gate. + lazy_window_waits, + /// Time those gathers waited for scoring to bring their probe within + /// the window. Part of `gate_wait_ns`, as are + /// `lazy_release_wait_ns` and `far_permit_wait_ns`. + lazy_window_wait_ns, + /// Time the `deferred_issues` waited for the heap to fill, their turn + /// or, with eager-before-full, the scoring of the probes holding `k` + /// rows. + lazy_release_wait_ns, + /// Time the `far_permit_waits` waited for a permit or the ordinary + /// window. + far_permit_wait_ns, + /// Loads of an index file's resident store: one per store while the + /// index cache keeps it, and one more after it evicted the store idle. + resident_columns_loads, + /// Index opens that bound a resident store a live index of the file + /// had loaded, found through the process's weak registry of stores. + resident_columns_registry_reuses, + /// Index opens that bound a resident store charged in the index cache. + resident_columns_binds, + /// Index opens whose resident store would take more of the index + /// cache's largest admissible entry than `auto` allows, so `auto` + /// kept the small columns in the file (`on` keeps them resident). + resident_columns_oversize, + /// Resident store entries the index cache dropped: evicted while + /// idle, cleared, or refused at admission. + resident_store_evictions, } ranked { /// Probes gathered lazily. @@ -227,11 +313,18 @@ layered_lazy_counters! { empty, /// Probes of non-empty partitions whose ex planes were not both /// resident, predicted from `k` and the partition sizes to be - /// gathered whole, and so loaded and scored by the eager scan; see - /// `LANCE_RQ_LAZY_DENSE_TO_EAGER`. + /// gathered whole, and so loaded and scored by the eager scan, as + /// `LANCE_RQ_LAZY_DENSE_TO_EAGER` routes them for their planes' tiers. dense_to_eager, + /// Probes of non-empty partitions with a high or low plane that no + /// cache tier held when they were staged, on an index whose origin + /// latency class is high: reading the plane is a request to the + /// origin, an object store such as S3. Counted whether the probe was + /// then gathered lazily or routed to the eager scan; always 0 on a + /// low-latency origin. + s3_bound_probes, /// Lazy probes issued at once because the rows of all earlier probes - /// cannot fill the heap. The probes of queries that + /// cannot fill the heap. The probes that /// `LANCE_RQ_LAZY_DENSE_TO_EAGER` routes are counted in /// `dense_to_eager` instead. certain_dense, @@ -266,6 +359,10 @@ layered_lazy_counters! { maxima { /// Largest number of probes still unscored at a gather's issue. staleness_max, + /// Most permits of one index's pool taken at once, observed as each + /// gather beyond the ordinary window takes one; at most + /// `LANCE_RQ_LAZY_FAR_INFLIGHT`. + far_in_flight_max, } } diff --git a/rust/lance-index/src/vector/bq/lazy_tests.rs b/rust/lance-index/src/vector/bq/lazy_tests.rs index 9d38db9a5f9..2e78ea3c780 100644 --- a/rust/lance-index/src/vector/bq/lazy_tests.rs +++ b/rust/lance-index/src/vector/bq/lazy_tests.rs @@ -24,7 +24,7 @@ use super::builder::RabitQuantizer; use super::dist_table_quant::{ DistTableDequant, quantize_dist_table_into, quantize_dist_table_u16_into, }; -use super::layered::{PlaneBatch, RQPrecision, plane_columns}; +use super::layered::{PlaneBatch, RQPrecision, SignBounds, plane_columns}; use super::storage::{ ExRows, RabitPruneCounters, RabitQuantizationStorage, SignStage, StagePruneCounts, SurvivorRow, select_full_survivors, take_captured_prune_counters, @@ -245,7 +245,7 @@ impl Fixture { let (batch, rotated_centroid, dist_q_c) = self.transform(residuals, centroid, row_id_base); let plane = |plane: u8| { let schema = batch.schema(); - let indices: Vec = plane_columns(plane) + let indices: Vec = plane_columns(plane, SignBounds::default()) .iter() .map(|name| schema.index_of(name).unwrap()) .collect(); diff --git a/rust/lance-index/src/vector/bq/plane_cache.rs b/rust/lance-index/src/vector/bq/plane_cache.rs index a0aa8fffed2..a2e2fb02e94 100644 --- a/rust/lance-index/src/vector/bq/plane_cache.rs +++ b/rust/lance-index/src/vector/bq/plane_cache.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -//! Row-addressable persistent ex planes. Sign codes retain their transposed IPC body. +//! Row-addressable persistent ex planes. Sign codes retain their transposed +//! IPC body, and so do the bounds columns, which scans read whole. use std::ops::Range; use std::sync::Arc; @@ -11,7 +12,7 @@ use lance_arrow::FixedSizeListArrayExt; use lance_core::cache::{CacheCodecImpl, CacheEntryReader, CacheEntryWriter, CacheRangeReader}; use lance_core::{Error, Result}; -use super::layered::{PlaneBatch, plane_columns}; +use super::layered::{PlaneBatch, SignBounds, plane_columns}; use super::storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN}; const HEADER_BYTES: usize = 13; @@ -43,7 +44,8 @@ impl Header { self.width + 8 } fn batch(&self, codes: Vec, adds: Vec, scales: Vec) -> Result { - let names = plane_columns(self.plane); + // Ex planes hold the same columns under either bounds placement. + let names = plane_columns(self.plane, SignBounds::default()); let codes = arrow_array::FixedSizeListArray::try_new_from_values( UInt8Array::from(codes), self.width as i32, @@ -114,7 +116,7 @@ impl CacheCodecImpl for PlaneBatch { if plane == 0 { return writer.write_ipc(&self.0); } - let names = plane_columns(plane); + let names = plane_columns(plane, SignBounds::default()); let codes = self.0[names[0]].as_fixed_size_list(); let adds = self.0[names[1]].as_primitive::(); let scales = self.0[names[2]].as_primitive::(); @@ -244,11 +246,59 @@ fn row_runs( #[cfg(test)] mod tests { use super::*; + use crate::vector::bq::layered::{PlaneKey, SIGN_BOUNDS_PLANE}; use bytes::Bytes; use lance_arrow::RecordBatchExt; - use lance_core::cache::{CacheCodec, CacheDecode}; + use lance_core::cache::{CacheCodec, CacheDecode, CacheKey}; use std::cell::Cell; + /// Bounds plane entries keep an IPC body, as sign planes do: they round + /// trip whole and have no row directory, so a row read misses. + #[test] + fn bounds_plane_round_trips_through_ipc() { + const ROWS: usize = 70; + let names = plane_columns(SIGN_BOUNDS_PLANE, SignBounds::Lazy); + let bounds = |offset: f32| { + let values: Vec = (0..ROWS * 3).map(|v| v as f32 + offset).collect(); + arrow_array::FixedSizeListArray::try_new_from_values(Float32Array::from(values), 3) + .unwrap() + }; + let (high, full) = (bounds(0.5), bounds(0.25)); + let original = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(names[0], high.data_type().clone(), true), + Field::new(names[1], full.data_type().clone(), true), + ])), + vec![Arc::new(high), Arc::new(full)], + ) + .unwrap(); + let codec = PlaneKey { + partition: 0, + plane: SIGN_BOUNDS_PLANE, + sign_bounds: SignBounds::Lazy, + } + .codec_for_key() + .unwrap(); + let mut encoded = Vec::new(); + codec + .serialize( + &(Arc::new(PlaneBatch(original.clone())) as Arc), + &mut encoded, + ) + .unwrap(); + let encoded = Bytes::from(encoded); + let CacheDecode::Hit(decoded) = codec.deserialize(&encoded) else { + panic!("bounds plane decode failed") + }; + assert_eq!(decoded.downcast::().unwrap().0, original); + let read = |range: Range| -> Result { Ok(encoded.slice(range)) }; + assert!(matches!( + codec.deserialize_rows(&read, &[0, 5]), + CacheDecode::Miss(_) + )); + assert!(codec.plan_rows(&read, &[0, 5]).is_none()); + } + #[test] fn sparse_plane_cache_reads_match_full_decode() { for plane in [1, 2] { diff --git a/rust/lance-index/src/vector/bq/resident.rs b/rust/lance-index/src/vector/bq/resident.rs new file mode 100644 index 00000000000..c60df58d8c1 --- /dev/null +++ b/rust/lance-index/src/vector/bq/resident.rs @@ -0,0 +1,1014 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! The small columns of an IVF_RQ storage file, kept in memory. +//! +//! A read of a partition or plane from the file costs at least one request +//! per column, yet every column but the codes and the estimator bounds holds +//! a few bytes per row: row ids and factors. With those columns read once +//! per index, later reads fetch only the codes (and bounds) from the file: +//! on an object store a native partition miss takes 2 requests instead of +//! 8, and a layered one 3 instead of 13. Reads attach copies of the resident +//! rows, so cache entries hold, and are charged, what a file read returns. +//! +//! The store is an entry of the index cache ([`ResidentColumnsKey`]), +//! charged what its arrays allocate and evictable under ordinary cache +//! pressure. Each live index holds an `Arc` to its store independently of +//! the cache. A weak registry shares that allocation across index handles, +//! even after eviction; concurrent handles load it once. An idle store is +//! freed when the cache evicts it and no index holds it anymore. + +use std::borrow::Cow; +use std::collections::HashMap; +use std::sync::{Arc, LazyLock, Mutex, Weak}; +use std::time::Instant; + +use arrow::compute::concat_batches; +use arrow_array::{Array, ArrayRef, UInt64Array, new_empty_array}; +use arrow_schema::{Field, Schema}; +use arrow_select::concat::concat; +use arrow_select::take::take; +use futures::TryStreamExt; +use lance_core::cache::{ + CacheKey, CacheKeySchema, InternalCacheKey, KeyBuilder, LanceCache, WeakLanceCache, +}; +use lance_core::deepsize::{Context, DeepSizeOf}; +use lance_core::{Error, Result}; +use lance_encoding::decoder::{FilterExpression, PageInfo}; +use lance_file::reader::FileReader; +use lance_io::ReadBatchParams; +use lance_io::scheduler::IoStats; +use tokio::sync::OnceCell; + +use super::layered::{FULL_BOUNDS_COLUMN, HIGH_BOUNDS_COLUMN}; +use super::layered_stats; +use super::storage::{ + RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN, RABIT_CODE_COLUMN, + RABIT_EX_CODE_COLUMN, +}; +use crate::vector::storage::{IndexFileKey, ResidentLifetime, WeakRegistry, shared_by_key}; + +/// Columns that reads always fetch from the file: the codes, which hold most +/// of its bytes, and the estimator bounds, which only some scans read. +const FILE_COLUMNS: [&str; 6] = [ + RABIT_CODE_COLUMN, + RABIT_EX_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + HIGH_BOUNDS_COLUMN, + FULL_BOUNDS_COLUMN, +]; + +/// Byte width of `field`'s values when the store keeps the column: every +/// fixed-width primitive column but [`FILE_COLUMNS`]. +fn resident_width(field: &Field) -> Option { + if FILE_COLUMNS.contains(&field.name().as_str()) { + return None; + } + field.data_type().primitive_width() +} + +/// Whether the store keeps column `field` of an IVF_RQ storage file. +pub(crate) fn is_resident(field: &Field) -> bool { + resident_width(field).is_some() +} + +/// Bytes of values the resident store holds for an IVF_RQ storage file with +/// `schema` and `num_rows` rows: every row of every column it keeps, at the +/// column's fixed width. It reads nothing, so an open can size the store +/// against the index cache before loading it. The store's allocations, which +/// the cache charges, can exceed it by a few bytes per column. +pub fn resident_columns_bytes(schema: &Schema, num_rows: u64) -> u64 { + schema + .fields() + .iter() + .filter_map(|field| resident_width(field)) + .map(|width| width as u64 * num_rows) + .sum() +} + +/// The stores of the index files with a live handle or cache entry, held +/// weakly: only dedupes loads, as the handles and the index cache own them. +static RESIDENT_SLOTS: LazyLock> = + LazyLock::new(Default::default); + +/// Handles kept for the life of the process under +/// [`ResidentLifetime::Process`], one strong reference per index file. +static PROCESS_STORES: LazyLock>>> = + LazyLock::new(Default::default); + +/// One index file's store, shared by every handle of the file and by its +/// index cache entry, and loaded on first use. +#[derive(Default)] +struct ResidentSlot { + store: OnceCell, +} + +/// Prints the store's columns, not its values. +impl std::fmt::Debug for ResidentSlot { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResidentSlot") + .field("store", &self.store.get()) + .finish() + } +} + +impl ResidentSlot { + /// The store of `reader`'s file, loaded on first use. Concurrent first + /// callers share one load, whose I/O is added to the loading caller's + /// `io_stats`; a failed or dropped load leaves the store for the next + /// caller to load. The load runs once per store, so it is boxed rather + /// than inlined into every read's future. + async fn get_or_load( + &self, + reader: &FileReader, + io_stats: Option<&IoStats>, + ) -> Result<&ResidentColumnStore> { + self.store + .get_or_try_init(|| Box::pin(ResidentColumnStore::load(reader, io_stats))) + .await + } + + fn is_loaded(&self) -> bool { + self.store.initialized() + } +} + +/// The live store of `file`, or a new one registered for the next caller, +/// and whether it was live. +fn shared_slot(file: &IndexFileKey) -> (Arc, bool) { + let mut created = false; + let slot = shared_by_key(&RESIDENT_SLOTS, file, || { + created = true; + Arc::default() + }); + (slot, !created) +} + +/// The live stores that hold their columns. +fn live_slots() -> Vec> { + RESIDENT_SLOTS + .lock() + .unwrap_or_else(|e| e.into_inner()) + .values() + .filter_map(Weak::upgrade) + .filter(|slot| slot.is_loaded()) + .collect() +} + +/// Loaded resident stores alive in the process, cached or held by a live +/// index. 0 once every index of every file is dropped and the index caches +/// evicted or cleared the stores, unless the `process` lifetime keeps them. +pub fn resident_store_count() -> usize { + live_slots().len() +} + +/// Whether a loaded resident store of `file` is alive in the process. +pub fn resident_store_is_live(file: &IndexFileKey) -> bool { + RESIDENT_SLOTS + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(file) + .and_then(Weak::upgrade) + .is_some_and(|slot| slot.is_loaded()) +} + +/// Cache key of the resident store of an IVF_RQ index file. An index keeps +/// it in its namespace of the index cache without a fragment reuse segment, +/// so that a new fragment reuse index, which changes the namespace of the +/// index's other entries, keeps the store. +pub struct ResidentColumnsKey<'a> { + file: &'a IndexFileKey, +} + +impl<'a> ResidentColumnsKey<'a> { + /// The key of the store of `file`. + pub fn new(file: &'a IndexFileKey) -> Self { + Self { file } + } +} + +impl CacheKey for ResidentColumnsKey<'_> { + type ValueType = ResidentColumnsEntry; + + fn key(&self) -> Cow<'_, str> { + format!( + "rq-resident-columns/{}/{}/{}", + self.file.index_uuid(), + self.file.store_prefix(), + self.file.path() + ) + .into() + } + + fn type_name() -> &'static str { + "IvfRqResidentColumns" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.index.rq-resident-columns-key", 1) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_str(self.file.index_uuid()); + builder.write_str(self.file.store_prefix()); + builder.write_str(self.file.path()); + } +} + +/// The index cache's entry of a loaded resident store: RAM only, charged +/// what the store's arrays allocate. Eviction releases only this reference; +/// live index handles continue to own the store. +pub struct ResidentColumnsEntry { + slot: Arc, +} + +impl ResidentColumnsEntry { + fn new(slot: Arc) -> Self { + Self { slot } + } + + /// Bytes of values the store holds; see [`resident_columns_bytes`]. + pub fn loaded_bytes(&self) -> Option { + self.slot.store.get().map(|store| store.bytes) + } +} + +impl std::fmt::Debug for ResidentColumnsEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("ResidentColumnsEntry") + .field(&self.slot) + .finish() + } +} + +impl DeepSizeOf for ResidentColumnsEntry { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + if !context.mark_seen(Arc::as_ptr(&self.slot) as usize) { + return 0; + } + std::mem::size_of::() + + self + .slot + .store + .get() + .map_or(0, ResidentColumnStore::heap_bytes) + } +} + +/// An entry drops once the index cache no longer holds it: evicted, +/// cleared or refused at admission. +impl Drop for ResidentColumnsEntry { + fn drop(&mut self) { + if self.slot.is_loaded() { + layered_stats::counters().resident_store_evictions.incr(); + } + } +} + +/// Where a handle charges its store. +#[derive(Debug, Clone)] +struct ResidentCharge { + file: IndexFileKey, + /// The index's namespace of the index cache without a fragment reuse + /// segment; see [`ResidentColumnsKey`]. + cache: WeakLanceCache, + lifetime: ResidentLifetime, +} + +/// A storage's handle on the resident columns of its file, loaded on first +/// use. A handle bound to an index cache +/// ([`in_index_cache`](Self::in_index_cache)) shares the store with every +/// live handle of the file and caches it once loaded. The handle's strong +/// reference keeps the data alive through cache eviction. A default handle +/// holds a store of its own, charged nowhere. +#[derive(Debug, Clone, Default)] +pub struct ResidentColumns { + slot: Arc, + /// `None` for a store charged nowhere. + charge: Option, + /// Cache the loaded store once per handle, without consulting the cache + /// on every partition read. Reopening an index can cache it again. + cached: OnceCell<()>, +} + +impl ResidentColumns { + /// The store of index file `file` that every live handle of it shares, + /// charged nowhere, for an index opened without an index cache. + pub fn shared(file: &IndexFileKey) -> Self { + Self { + slot: shared_slot(file).0, + charge: None, + cached: OnceCell::new(), + } + } + + /// The handle of an index opening index file `file` with its small + /// columns resident: the store every live handle of the file shares, + /// charged in `cache`, the index's namespace of the index cache without + /// a fragment reuse segment. `preopen` holds the shared store through + /// cache accesses performed while opening the index. Loaded stores are + /// cached now; otherwise the first read loads and caches the store. + /// Under [`ResidentLifetime::Process`] a separate strong reference keeps + /// the loaded store alive until the process exits. + pub async fn in_index_cache( + cache: &LanceCache, + file: &IndexFileKey, + preopen: Option, + lifetime: ResidentLifetime, + ) -> Self { + // Keep the same allocation even if opening the index evicted it. + let (slot, live) = match &preopen { + Some(preopen) => (preopen.slot.clone(), true), + None => shared_slot(file), + }; + let stats = layered_stats::counters(); + stats.resident_columns_binds.incr(); + if live && slot.is_loaded() { + stats.resident_columns_registry_reuses.incr(); + } + let handle = Self { + slot, + charge: Some(ResidentCharge { + file: file.clone(), + cache: WeakLanceCache::from(cache), + lifetime, + }), + cached: OnceCell::new(), + }; + handle.cache_loaded().await; + handle + } + + /// Offer the loaded store to the ordinary cache once per handle. Cache + /// admission and eviction do not control the store's lifetime. + async fn cache_loaded(&self) { + let Some(charge) = &self.charge else { + return; + }; + if !self.slot.is_loaded() { + return; + } + self.cached + .get_or_init(|| async { + let key = ResidentColumnsKey::new(&charge.file); + if charge.cache.get_with_key(&key).await.is_none() { + charge + .cache + .insert_with_key( + &key, + Arc::new(ResidentColumnsEntry::new(self.slot.clone())), + ) + .await; + } + if charge.lifetime == ResidentLifetime::Process { + PROCESS_STORES + .lock() + .unwrap_or_else(|e| e.into_inner()) + .entry(charge.file.clone()) + .or_insert_with(|| self.slot.clone()); + } + }) + .await; + } + + /// Bytes of values the store holds, `None` until it has loaded. Equals + /// [`resident_columns_bytes`] of its file. + pub fn loaded_bytes(&self) -> Option { + self.slot.store.get().map(|store| store.bytes) + } + + /// The store of `reader`'s file, loaded on first use: see + /// [`ResidentSlot::get_or_load`]. Cache the loaded allocation without + /// making subsequent reads depend on whether it stays in the cache. + pub(crate) async fn get_or_load( + &self, + reader: &FileReader, + io_stats: Option<&IoStats>, + ) -> Result<&ResidentColumnStore> { + let store = self.slot.get_or_load(reader, io_stats).await?; + self.cache_loaded().await; + Ok(store) + } +} + +/// Every row of the resident columns of one file. +pub(crate) struct ResidentColumnStore { + columns: HashMap, + num_rows: u64, + /// See [`ResidentColumns::loaded_bytes`]. + bytes: u64, +} + +/// Names the columns instead of printing hundreds of megabytes of values. +impl std::fmt::Debug for ResidentColumnStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut columns: Vec<&String> = self.columns.keys().collect(); + columns.sort(); + f.debug_struct("ResidentColumnStore") + .field("columns", &columns) + .field("num_rows", &self.num_rows) + .field("bytes", &self.bytes) + .finish() + } +} + +impl ResidentColumnStore { + async fn load(reader: &FileReader, io_stats: Option<&IoStats>) -> Result { + let started = Instant::now(); + let schema = Schema::from(reader.schema().as_ref()); + let fields: Vec<(&Field, usize)> = schema + .fields() + .iter() + .filter_map(|field| Some((field.as_ref(), resident_width(field)?))) + .collect(); + let num_rows = reader.num_rows(); + let load_stats = IoStats::new(); + let batch = if fields.is_empty() || num_rows == 0 { + None + } else { + let names: Vec<&str> = fields + .iter() + .map(|(field, _)| field.name().as_str()) + .collect(); + let projection = lance_file::versions::reader_projection_from_column_names( + reader.metadata().version(), + reader.schema(), + &names, + )?; + let projected = Arc::new(Schema::from(projection.schema.as_ref())); + let stats_reader = reader.with_io_stats(load_stats.recorder()); + let batches = stats_reader + .read_stream_projected( + ReadBatchParams::Range(0..num_rows as usize), + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + Some(concat_batches(&projected, batches.iter())?) + }; + let mut columns = HashMap::with_capacity(fields.len()); + let mut bytes = 0u64; + // What the arrays hold, which a zero-copy slice of a larger decoded + // or read buffer would push past `bytes`. + let mut alloc_bytes = 0u64; + for (field, width) in fields { + let name = field.name(); + let values = match &batch { + Some(batch) => batch.column_by_name(name).cloned().ok_or_else(|| { + Error::internal(format!("resident column {name} is missing from its read")) + })?, + None => new_empty_array(field.data_type()), + }; + if values.len() as u64 != num_rows { + return Err(Error::internal(format!( + "resident column {name} read {} of the file's {num_rows} rows", + values.len() + ))); + } + bytes += width as u64 * num_rows; + alloc_bytes += values.get_buffer_memory_size() as u64; + let page_ends = column_page_ends(reader, name)?; + columns.insert(name.clone(), ResidentColumn { values, page_ends }); + } + let loaded = load_stats.snapshot(); + let stats = layered_stats::counters(); + stats.resident_columns_loads.incr(); + stats.resident_columns_bytes.add(bytes); + stats.resident_columns_alloc_bytes.add(alloc_bytes); + stats.resident_columns_load_requests.add(loaded.iops); + stats.resident_columns_load_bytes.add(loaded.bytes_read); + stats.resident_columns_load_ns.add_elapsed(started); + if let Some(io_stats) = io_stats { + io_stats.add_scan_stats(&loaded); + } + Ok(Self { + columns, + num_rows, + bytes, + }) + } + + /// Resident column `name`, `None` for a column reads fetch from the file. + pub(crate) fn column(&self, name: &str) -> Option<&ResidentColumn> { + self.columns.get(name) + } + + /// Rows of the file. + pub(crate) fn num_rows(&self) -> u64 { + self.num_rows + } + + /// Heap memory the store holds: its arrays' buffers, which the + /// `resident_columns_alloc_bytes` of its load count, plus its pages and + /// its map of columns. + fn heap_bytes(&self) -> usize { + self.columns.capacity() * std::mem::size_of::<(String, ResidentColumn)>() + + self + .columns + .iter() + .map(|(name, column)| { + name.capacity() + + column.values.get_buffer_memory_size() + + column.page_ends.capacity() * std::mem::size_of::() + }) + .sum::() + } +} + +/// The pages of column `name` of `reader`'s file. +fn column_pages<'a>(reader: &'a FileReader, name: &str) -> Result<&'a [PageInfo]> { + let metadata = reader.metadata(); + let projection = lance_file::versions::reader_projection_from_column_names( + metadata.version(), + reader.schema(), + &[name], + )?; + match projection.column_indices.as_slice() { + [column] => metadata.column_infos.get(*column as usize), + _ => None, + } + .map(|column| column.page_infos.as_ref()) + .ok_or_else(|| { + Error::internal(format!( + "resident column {name} is not one column of the file: {:?}", + projection.column_indices + )) + }) +} + +/// The row at which each page of column `name` of `reader`'s file ends, one +/// per page in a vector of exactly that capacity, as +/// [`resident_store_charge`] counts them. +fn column_page_ends(reader: &FileReader, name: &str) -> Result> { + let pages = column_pages(reader, name)?; + let mut page_ends = Vec::with_capacity(pages.len()); + let mut end = 0u64; + for page in pages { + end += page.num_rows; + page_ends.push(end); + } + Ok(page_ends) +} + +/// Bytes an index cache charges for the resident store of `reader`'s file +/// once loaded, known without reading it: what its cache entry +/// ([`ResidentColumnsEntry`]) takes, which is the store's columns in buffers +/// of exactly their values ([`resident_columns_bytes`]), their page ends, +/// names and map, the shared slot and the entry's `Arc`, and the cost a +/// backend adds for the entry's key. An open compares it with the cache's +/// largest admissible entry before automatically keeping a store resident. +pub fn resident_store_charge(reader: &FileReader) -> Result { + let schema = Schema::from(reader.schema().as_ref()); + let columns = schema + .fields() + .iter() + .filter_map(|field| Some((field.name(), resident_width(field)?))) + .map(|(name, width)| Ok((name.as_str(), width, column_pages(reader, name)?.len()))) + .collect::>>()?; + Ok(store_charge(&columns, reader.num_rows())) +} + +/// [`resident_store_charge`] of a store of `num_rows` rows of `columns`, +/// each a name, a value width and a number of pages. +fn store_charge(columns: &[(&str, usize, usize)], num_rows: u64) -> u64 { + let map_bytes = HashMap::::with_capacity(columns.len()).capacity() + * std::mem::size_of::<(String, ResidentColumn)>(); + let column_bytes: u64 = columns + .iter() + .map(|&(name, width, pages)| { + (name.len() + pages * std::mem::size_of::()) as u64 + width as u64 * num_rows + }) + .sum(); + let fixed_bytes = std::mem::size_of::() + + std::mem::size_of::() + + 2 * std::mem::size_of::() + + std::mem::size_of::(); + (map_bytes + fixed_bytes) as u64 + column_bytes +} + +/// Every row of one resident column, with the column's pages in the file. +pub(crate) struct ResidentColumn { + values: ArrayRef, + /// The row at which each page of the column ends, ascending. + page_ends: Vec, +} + +impl ResidentColumn { + /// A copy of the ascending file rows `rows`, laid out as a file read of + /// them so that a cache entry holding it is charged the same bytes: the + /// reader decodes each page a read touches into a buffer of exactly its + /// rows and concatenates the buffers when the read touches several + /// pages. A slice would instead keep, and be charged, the whole column. + pub(crate) fn copy_rows(&self, rows: &UInt64Array) -> Result { + let offsets = rows.values(); + // The page split below relies on the order. + if !offsets.is_sorted() { + return Err(Error::invalid_input( + "rows read from resident columns must ascend", + )); + } + if let Some(&last) = offsets.last() + && last >= self.values.len() as u64 + { + return Err(Error::invalid_input(format!( + "resident row {last} is past the {} rows of the file", + self.values.len() + ))); + } + let mut pieces = Vec::new(); + let mut start = 0; + let mut page = 0; + while start < offsets.len() { + let first = offsets[start]; + page += self.page_ends[page..].partition_point(|&end| end <= first); + let page_end = self.page_ends.get(page).copied().unwrap_or(u64::MAX); + let len = offsets[start..].partition_point(|&row| row < page_end); + let piece = take(self.values.as_ref(), &rows.slice(start, len), None)?; + pieces.push(piece); + start += len; + } + match pieces.len() { + 0 => Ok(new_empty_array(self.values.data_type())), + 1 => Ok(pieces.swap_remove(0)), + _ => { + let pieces: Vec<&dyn Array> = pieces.iter().map(|piece| piece.as_ref()).collect(); + Ok(concat(&pieces)?) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Float32Array, cast::AsArray, types::Float32Type}; + use arrow_schema::DataType; + use lance_core::cache::QuickCacheBackend; + + use crate::vector::bq::layered::{HIGH_ADD_FACTORS_COLUMN, HIGH_SCALE_FACTORS_COLUMN}; + use crate::vector::bq::transform::{ + ADD_FACTORS_COLUMN, ERROR_FACTORS_COLUMN, EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + }; + + fn list_field(name: &str, item: DataType, width: i32) -> Field { + Field::new( + name, + DataType::FixedSizeList(Arc::new(Field::new("item", item, true)), width), + true, + ) + } + + fn factor_field(name: &str) -> Field { + Field::new(name, DataType::Float32, true) + } + + #[test] + fn resident_columns_bytes_counts_the_small_fixed_width_columns() { + let native = vec![ + Field::new(lance_core::ROW_ID, DataType::UInt64, true), + list_field(RABIT_CODE_COLUMN, DataType::UInt8, 16), + factor_field(ADD_FACTORS_COLUMN), + factor_field(SCALE_FACTORS_COLUMN), + factor_field(ERROR_FACTORS_COLUMN), + list_field(RABIT_BLOCKED_EX_CODE_COLUMN, DataType::UInt8, 96), + factor_field(EX_ADD_FACTORS_COLUMN), + factor_field(EX_SCALE_FACTORS_COLUMN), + ]; + let mut layered = native.clone(); + layered.extend([ + list_field(RABIT_BLOCKED_EX_CODE_LO_COLUMN, DataType::UInt8, 64), + factor_field(HIGH_ADD_FACTORS_COLUMN), + factor_field(HIGH_SCALE_FACTORS_COLUMN), + list_field(HIGH_BOUNDS_COLUMN, DataType::Float32, 3), + list_field(FULL_BOUNDS_COLUMN, DataType::Float32, 3), + ]); + // Row ids (8 bytes) and four-byte factors: 3 for the sign codes and + // 2 for each ex level. + assert_eq!( + resident_columns_bytes(&Schema::new(native.clone()), 1000), + 28_000 + ); + assert_eq!(resident_columns_bytes(&Schema::new(layered), 1000), 36_000); + assert_eq!(resident_columns_bytes(&Schema::new(native.clone()), 0), 0); + // Legacy ex codes and columns without a fixed width stay in the file. + let mut other = native; + other.push(list_field(RABIT_EX_CODE_COLUMN, DataType::UInt8, 96)); + other.push(Field::new("label", DataType::Utf8, true)); + assert_eq!(resident_columns_bytes(&Schema::new(other), 10), 280); + } + + /// A file key unique to one test, so tests that share the process-wide + /// registry do not share stores. + fn test_file(test: &str) -> IndexFileKey { + let path = format!("t.lance/_indices/{test}/auxiliary.idx"); + IndexFileKey::new(test, "s3$bucket", &path) + } + + /// A loaded store of `rows` rows of one four-byte column in one page. + fn test_store(rows: usize) -> ResidentColumnStore { + let values: ArrayRef = Arc::new(Float32Array::from(vec![1.0; rows])); + let column = ResidentColumn { + values, + page_ends: vec![rows as u64], + }; + ResidentColumnStore { + columns: HashMap::from([("factor".to_string(), column)]), + num_rows: rows as u64, + bytes: 4 * rows as u64, + } + } + + /// The slot of `file`, loaded with a test store, and held by the caller. + fn loaded_slot(file: &IndexFileKey, rows: usize) -> Arc { + let (slot, _) = shared_slot(file); + slot.store.set(test_store(rows)).unwrap(); + slot + } + + fn test_cache() -> LanceCache { + LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(1 << 20))) + } + + /// Handles of one file share its store while any holds it, and a store + /// no handle or cache holds is freed, so the next handle gets a new one; + /// the same path in another bucket is another file. + #[test] + fn slot_registry_shares_live_store_and_frees_dropped() { + let file = test_file("slot-registry"); + let store = ResidentColumns::shared(&file); + let shared = ResidentColumns::shared(&file); + assert!(Arc::ptr_eq(&store.slot, &shared.slot)); + let path = file.path().to_string(); + let other = IndexFileKey::new(file.index_uuid(), "s3$other", &path); + assert!(!Arc::ptr_eq( + &store.slot, + &ResidentColumns::shared(&other).slot + )); + store.slot.store.set(test_store(10)).unwrap(); + assert!(resident_store_is_live(&file)); + let held = Arc::downgrade(&store.slot); + drop((store, shared)); + assert!(held.upgrade().is_none()); + assert!(!resident_store_is_live(&file)); + let reopened = ResidentColumns::shared(&file); + assert_eq!(reopened.loaded_bytes(), None); + } + + /// The cache entry of a store is charged its heap memory: what its + /// arrays allocate, and less than 64 KiB more. + #[test] + fn store_deep_size_matches_alloc() { + const OVERHEAD_BYTES: usize = 64 * 1024; + let slot = Arc::new(ResidentSlot::default()); + let rows = 10_000; + slot.store.set(test_store(rows)).unwrap(); + let entry = ResidentColumnsEntry::new(slot.clone()); + let alloc = rows * 4; + let size = entry.deep_size_of(); + assert!(size >= alloc, "{size} < {alloc}"); + assert!(size <= alloc + OVERHEAD_BYTES, "{size}"); + // Two entries of one store are counted once. + let twin = ResidentColumnsEntry::new(slot); + let mut context = Context::new(); + let both = + entry.deep_size_of_children(&mut context) + twin.deep_size_of_children(&mut context); + assert_eq!(both, entry.deep_size_of_children(&mut Context::new())); + } + + #[tokio::test] + async fn in_index_cache_retains_idle_store() { + let file = test_file("in-index-cache"); + let slot = loaded_slot(&file, 1000); + let held = Arc::downgrade(&slot); + let cache = test_cache(); + let key = ResidentColumnsKey::new(&file); + let handle = + ResidentColumns::in_index_cache(&cache, &file, None, ResidentLifetime::Index).await; + assert!(cache.get_with_key(&key).await.is_some()); + let clone = handle.clone(); + drop((handle, clone, slot)); + assert!( + held.upgrade().is_some(), + "the idle cache entry owns the store" + ); + cache.clear().await; + assert!( + held.upgrade().is_none(), + "the weak registry must not retain data" + ); + } + + #[tokio::test] + async fn live_store_is_evictable_and_reused_after_eviction() { + let file = test_file("active-eviction"); + let slot = loaded_slot(&file, 1000); + let held = Arc::downgrade(&slot); + let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(8192))); + let key = ResidentColumnsKey::new(&file); + let handle = + ResidentColumns::in_index_cache(&cache, &file, None, ResidentLifetime::Index).await; + drop(slot); + assert!(cache.peek_resident_with_key(&key).await); + for id in 0..16 { + let other = test_file(&format!("eviction-pressure-{id}")); + let slot = loaded_slot(&other, 1000); + cache + .insert_with_key( + &ResidentColumnsKey::new(&other), + Arc::new(ResidentColumnsEntry::new(slot)), + ) + .await; + } + assert!(!cache.peek_resident_with_key(&key).await); + assert!(cache.size_bytes().await <= 8192); + assert_eq!(handle.loaded_bytes(), Some(4000)); + let reopened = + ResidentColumns::in_index_cache(&cache, &file, None, ResidentLifetime::Index).await; + assert!(Arc::ptr_eq(&handle.slot, &reopened.slot)); + assert_eq!(reopened.loaded_bytes(), Some(4000)); + let values = reopened.slot.store.get().unwrap().columns["factor"] + .values + .as_any() + .downcast_ref::() + .unwrap(); + assert!(values.values().iter().all(|value| *value == 1.0)); + drop((handle, reopened)); + cache.clear().await; + assert!(held.upgrade().is_none()); + } + + /// A store of `rows` rows of `columns`, each a name and a value width + /// of 4 or 8 bytes, in `pages` pages, laid out as a load lays it out. + fn store_of(columns: &[(&str, usize)], rows: usize, pages: usize) -> ResidentColumnStore { + let mut map = HashMap::with_capacity(columns.len()); + let mut bytes = 0; + for &(name, width) in columns { + let values: ArrayRef = match width { + 4 => Arc::new(Float32Array::from(vec![1.0; rows])), + _ => Arc::new(UInt64Array::from(vec![1; rows])), + }; + let mut page_ends = Vec::with_capacity(pages); + page_ends.extend((1..=pages).map(|page| (rows * page / pages) as u64)); + map.insert(name.to_string(), ResidentColumn { values, page_ends }); + bytes += (width * rows) as u64; + } + ResidentColumnStore { + columns: map, + num_rows: rows as u64, + bytes, + } + } + + /// The slot of `file`, loaded with `store`, and held by the caller. + fn slot_with(file: &IndexFileKey, store: ResidentColumnStore) -> Arc { + let (slot, _) = shared_slot(file); + slot.store.set(store).unwrap(); + slot + } + + const CHARGE_TEST_COLUMNS: [(&str, usize); 3] = + [("_rowid", 8), ("__add_factors", 4), ("__scale_factors", 4)]; + + /// What [`resident_store_charge`] estimates without reading the store + /// is what an ordinary cache charges for it once loaded. + #[tokio::test] + async fn store_charge_is_what_the_cache_charges() { + for (rows, pages) in [(1, 1), (1000, 1), (10_000, 7), (100_000, 64)] { + let columns: Vec<(&str, usize, usize)> = CHARGE_TEST_COLUMNS + .iter() + .map(|&(name, width)| (name, width, pages)) + .collect(); + let charge = store_charge(&columns, rows as u64); + let file = test_file(&format!("store-charge-{rows}-{pages}")); + let _slot = slot_with(&file, store_of(&CHARGE_TEST_COLUMNS, rows, pages)); + let cache = + LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(64 << 20))); + let _handle = + ResidentColumns::in_index_cache(&cache, &file, None, ResidentLifetime::Index).await; + assert_eq!(cache.size_bytes().await as u64, charge, "{rows} {pages}"); + } + } + + #[tokio::test] + async fn preopen_handle_keeps_a_store_evicted_during_the_open() { + let file = test_file("preopen-evicted"); + let cache = test_cache(); + let key = ResidentColumnsKey::new(&file); + let slot = loaded_slot(&file, 1000); + let handle = + ResidentColumns::in_index_cache(&cache, &file, None, ResidentLifetime::Index).await; + let held = Arc::downgrade(&slot); + drop((handle, slot)); + let preopen = ResidentColumns::shared(&file); + cache.clear().await; + assert!(!cache.peek_resident_with_key(&key).await); + let handle = + ResidentColumns::in_index_cache(&cache, &file, Some(preopen), ResidentLifetime::Index) + .await; + assert_eq!(handle.loaded_bytes(), Some(4000)); + assert!(Arc::ptr_eq(&handle.slot, &held.upgrade().unwrap())); + assert!(cache.peek_resident_with_key(&key).await); + drop(handle); + cache.clear().await; + assert!(held.upgrade().is_none()); + } + + #[tokio::test] + async fn resident_lifetime_process_keeps_store_alive() { + let file = test_file("process-lifetime"); + let slot = loaded_slot(&file, 1000); + let held = Arc::downgrade(&slot); + let cache = test_cache(); + let handle = + ResidentColumns::in_index_cache(&cache, &file, None, ResidentLifetime::Process).await; + drop((handle, slot)); + cache.clear().await; + let slot = held + .upgrade() + .expect("the process owns the store after eviction"); + let reopened = + ResidentColumns::in_index_cache(&cache, &file, None, ResidentLifetime::Process).await; + assert!(Arc::ptr_eq(&reopened.slot, &slot)); + PROCESS_STORES.lock().unwrap().remove(&file); + drop((slot, reopened)); + cache.clear().await; + assert!(held.upgrade().is_none()); + } + + /// A 100-row column in pages ending at rows 40, 40 (an empty page) and 100. + fn column() -> ResidentColumn { + let values: Vec = (0..100).map(|row| row as f32).collect(); + ResidentColumn { + values: Arc::new(Float32Array::from(values)), + page_ends: vec![40, 40, 100], + } + } + + fn copy(column: &ResidentColumn, rows: &[u64]) -> ArrayRef { + column.copy_rows(&UInt64Array::from(rows.to_vec())).unwrap() + } + + fn values(array: &ArrayRef) -> Vec { + array.as_primitive::().values().to_vec() + } + + fn as_values(rows: &[u64]) -> Vec { + rows.iter().map(|&row| row as f32).collect() + } + + #[test] + fn copy_rows_takes_the_rows_into_their_own_buffers() { + let column = column(); + let store = column.values.as_primitive::(); + let store_values = store.values().as_ptr(); + let rows: Vec = (10..30).collect(); + let copied = copy(&column, &rows); + assert_eq!(values(&copied), as_values(&rows)); + assert!(copied.nulls().is_none()); + // One page: a buffer of exactly the rows, apart from the store's. + assert_eq!(copied.get_buffer_memory_size(), rows.len() * 4); + let copied_values = copied.as_primitive::().values().as_ptr(); + assert_ne!(copied_values, store_values); + + let sparse = copy(&column, &[3, 7, 8, 39]); + assert_eq!(values(&sparse), as_values(&[3, 7, 8, 39])); + assert_eq!(sparse.get_buffer_memory_size(), 4 * 4); + assert_eq!(copy(&column, &[]).len(), 0); + for rows in [vec![5, 100], vec![7, 5]] { + let error = column.copy_rows(&UInt64Array::from(rows)).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + } + } + + #[test] + fn copy_rows_across_pages_matches_a_file_read() { + let column = column(); + for rows in [ + (30..60).collect::>(), + vec![5, 39, 40, 99], + (0..100).collect(), + ] { + let copied = copy(&column, &rows); + assert_eq!(values(&copied), as_values(&rows)); + assert!(copied.nulls().is_none()); + // The reader concatenates the pages it decodes, each into a + // buffer of exactly its rows. + let (first, second): (Vec, Vec) = rows.iter().partition(|&&row| row < 40); + let pages = [first, second] + .map(|page| take(column.values.as_ref(), &UInt64Array::from(page), None).unwrap()); + let read = concat(&[pages[0].as_ref(), pages[1].as_ref()]).unwrap(); + assert_eq!(copied.as_ref(), read.as_ref(), "{rows:?}"); + assert_eq!( + copied.get_buffer_memory_size(), + read.get_buffer_memory_size(), + "{rows:?}" + ); + } + } +} diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index fda65451361..efbbc3d1ccc 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -699,13 +699,20 @@ impl RabitQuantizationStorage { super::layered::RQLayout::try_new(metadata.num_bits)?; } else if metadata.layered { super::layered::RQLayout::try_new(metadata.num_bits)?; - for name in [ + // The bounds columns are loaded only for the scans that prune + // with them. Full precision prunes with the error factors when + // the file has them; a High scan of a store without the high + // bounds scores every row (see + // `distance_calculator_from_parts_with_ex`). + let mut required = vec![ RABIT_BLOCKED_EX_CODE_LO_COLUMN, super::layered::HIGH_ADD_FACTORS_COLUMN, super::layered::HIGH_SCALE_FACTORS_COLUMN, - super::layered::HIGH_BOUNDS_COLUMN, - super::layered::FULL_BOUNDS_COLUMN, - ] { + ]; + if error_factors.is_none() { + required.push(super::layered::FULL_BOUNDS_COLUMN); + } + for name in required { if batch.column_by_name(name).is_none() { return Err(Error::invalid_input(format!( "RabitQ layered index missing column {name}" @@ -1068,6 +1075,11 @@ impl RabitQuantizationStorage { query_factors.centered_query, self.distance_type, ); + } else if rq_precision == super::layered::RQPrecision::High { + // A layered store loaded for full precision may leave the + // high bounds out. The error factors bound only the + // full-precision estimator, so fail closed: score every row. + calculator.error_factors = None; } } calculator @@ -1708,7 +1720,9 @@ impl<'a> RabitDistCalculator<'a> { .min(1.0) }; let errors = values - .chunks_exact(3) + .as_chunks::<3>() + .0 + .iter() .enumerate() .map(|(row, b)| { let add_scale = f64::from(self.add_factor_scale).abs(); @@ -3609,8 +3623,9 @@ impl QuantizerStorage for RabitQuantizationStorage { } /// Load only the sign plane (row ids, packed sign codes, binary and error - /// factors, bound columns) of a layered multi-bit partition. The result - /// supports [`RabitDistCalculator::full_sign_stage_with_scratch`] and, with + /// factors, and the bounds columns when it holds them) of a layered + /// multi-bit partition. The result supports + /// [`RabitDistCalculator::full_sign_stage_with_scratch`] and, with /// gathered rows, [`RabitQuantizationStorage::dist_calculator_with_ex_rows`]; /// it never scores rows through the eager top-k scans. fn try_from_sign_plane_for_full( diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index ab3bc6125ab..7025245a3b0 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -4,18 +4,20 @@ //! Vector Storage, holding (quantized) vectors and providing distance calculation. use crate::vector::bq::ex_dot::blocked_ex_code_bytes; -use crate::vector::bq::layered::{PlaneBatch, PlaneKey, RQLayout}; +use crate::vector::bq::layered::{PlaneBatch, PlaneKey, RQLayout, SIGN_BOUNDS_PLANE, SignBounds}; use crate::vector::bq::layered_stats; +use crate::vector::bq::resident::{ResidentColumnStore, is_resident}; +use crate::vector::bq::transform::ERROR_FACTORS_COLUMN; use crate::vector::quantizer::QuantizerStorage; use arrow::compute::concat_batches; use arrow_array::{ - Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array, cast::AsArray, - types::UInt8Type, + Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array, UInt64Array, + cast::AsArray, types::UInt8Type, }; use arrow_schema::SchemaRef; use futures::stream::{self, Stream, StreamExt, TryStreamExt}; use lance_arrow::{FixedSizeListArrayExt, RecordBatchExt}; -use lance_core::cache::WeakLanceCache; +use lance_core::cache::{CacheTier, WeakLanceCache}; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::{Error, ROW_ID, Result}; @@ -23,6 +25,7 @@ use lance_encoding::decoder::FilterExpression; use lance_file::reader::FileReader; use lance_index_core::remapping::{BatchRowIdRemapper, remap_row_ids_preserving_layout_async}; use lance_io::ReadBatchParams; +use lance_io::object_store::{DEFAULT_CLOUD_IO_PARALLELISM, ObjectStore}; use lance_io::scheduler::IoStats; use lance_io::spill::SpillStore; use lance_linalg::distance::DistanceType; @@ -31,9 +34,10 @@ use std::{ any::Any, borrow::Cow, collections::{BinaryHeap, HashMap, HashSet}, + hash::Hash, mem::size_of, ops::{Deref, DerefMut}, - sync::{Arc, LazyLock, Mutex, OnceLock}, + sync::{Arc, LazyLock, Mutex, OnceLock, Weak}, time::Instant, }; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; @@ -58,6 +62,11 @@ use super::pairwise::{ use super::quantizer::{Quantizer, QuantizerMetadata}; use super::{ApproxMode, DISTANCE_TYPE_KEY}; +pub use crate::vector::bq::resident::{ + ResidentColumns, ResidentColumnsEntry, ResidentColumnsKey, resident_columns_bytes, + resident_store_charge, resident_store_count, resident_store_is_live, +}; + /// Coalesce source-index reads independently of the scoring vector batch size. const PAIRWISE_READ_BATCH_SIZE: usize = 8192; // Bound each scoring/spill batch in code space; very wide vectors still need @@ -626,13 +635,82 @@ impl DeepSizeOf for PlaneAccess { } } +/// Identifies an index file across the process: the index's UUID, and the +/// prefix of the object store that holds the file (its scheme and bucket, or +/// whatever else tells the store apart; see [`ObjectStore::store_prefix`]) +/// with the file's path in that store. Every open of the file binds to the +/// runtime handles registered under it, so that indexes opened at once, a +/// re-open while an older index still runs, and a state read back from a +/// persistent cache tier share them: the resident store +/// ([`ResidentColumns::in_index_cache`]) and the lazy scan's far gather +/// permits ([`IvfQuantizationStorage::with_index_file`]). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct IndexFileKey(Arc); + +#[derive(Debug, PartialEq, Eq, Hash)] +struct IndexFileLocation { + index_uuid: String, + store_prefix: String, + path: String, +} + +impl IndexFileKey { + /// The file at `path` in the object store with prefix `store_prefix`, + /// of the index with UUID `index_uuid`. + pub fn new(index_uuid: &str, store_prefix: &str, path: &str) -> Self { + Self(Arc::new(IndexFileLocation { + index_uuid: index_uuid.to_owned(), + store_prefix: store_prefix.to_owned(), + path: path.to_owned(), + })) + } + + /// The UUID of the index the file belongs to. + pub(crate) fn index_uuid(&self) -> &str { + &self.0.index_uuid + } + + /// The prefix of the object store that holds the file. + pub(crate) fn store_prefix(&self) -> &str { + &self.0.store_prefix + } + + /// The file's path in its object store. + pub(crate) fn path(&self) -> &str { + &self.0.path + } +} + +/// Values shared process-wide by key and held weakly: a value lives while +/// anything outside the registry holds it. +pub(crate) type WeakRegistry = Mutex>>; + +/// The value `registry` holds for `key` while anything else still holds it, +/// else the one `create` makes, registered for the next caller. Entries of +/// dropped values are pruned whenever a value is registered. +pub(crate) fn shared_by_key( + registry: &LazyLock>, + key: &K, + create: impl FnOnce() -> Arc, +) -> Arc { + let mut values = registry.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(value) = values.get(key).and_then(Weak::upgrade) { + return value; + } + let value = create(); + values.retain(|_, held| held.strong_count() > 0); + values.insert(key.clone(), Arc::downgrade(&value)); + value +} + /// Runtime promotion history shared by reconstructions of the same cached index. /// This contains no readers or object-store handles and is not persisted to disk. #[derive(Debug, Clone, Default)] pub struct PlaneAccessTracker { /// Cascade candidate reads, see [`PLANE_PROMOTION_READS`]. accesses: Arc>>, - /// Background promotions of the lazy full-precision scan. + /// Background promotions and far gather permits of the lazy + /// full-precision scan. lazy: Arc, } @@ -654,6 +732,9 @@ struct LazyPromotionState { history: Mutex, /// Bounds concurrent promotions; sized by the first promoting query's config. permits: OnceLock>, + /// Bounds the lazy scan's gathers in flight beyond the ordinary window + /// across queries; see [`IvfQuantizationStorage::lazy_far_permits`]. + far_permits: Mutex>, } impl DeepSizeOf for LazyPromotionState { @@ -686,6 +767,78 @@ impl Drop for LazyPromotionTicket { } } +/// The lazy far gather pools of the open index files, by file and pool size; +/// see [`LazyFarPermits::for_file`]. +static FAR_PERMIT_POOLS: LazyLock> = + LazyLock::new(Default::default); + +/// The permits of an index's lazy gathers issued beyond the ordinary window +/// ([`LAZY_FAR_WINDOW_ENV`]), shared by the index's queries and by the +/// reconstructions of its cached state, and by every open of its file once +/// bound to it ([`IvfQuantizationStorage::with_index_file`]). A gather holds +/// its permit until its reads return; dropping the permit, also when the +/// gather is cancelled, returns it to the pool. +#[derive(Debug, Clone)] +pub struct LazyFarPermits { + semaphore: Arc, + /// Permits the pool holds while none is taken. + size: usize, +} + +impl LazyFarPermits { + /// A pool of `size` permits, at most [`Semaphore::MAX_PERMITS`]. + pub fn try_new(size: usize) -> Result { + if size > Semaphore::MAX_PERMITS { + return Err(Error::invalid_input(format!( + "a lazy far gather pool of {size} permits exceeds the {} a pool can hold", + Semaphore::MAX_PERMITS + ))); + } + Ok(Self { + semaphore: Arc::new(Semaphore::new(size)), + size, + }) + } + + /// The pool of `size` permits that every open of index file `file` in + /// the process shares: the live one while an index, a cached state or a + /// taken permit holds it, else a new one. Keyed by size too, so that an + /// index whose config sets another size (tests replace it) never draws + /// from a pool of the old size. + fn for_file(file: &IndexFileKey, size: usize) -> Result { + let new = Self::try_new(size)?; + let semaphore = shared_by_key(&FAR_PERMIT_POOLS, &(file.clone(), size), || new.semaphore); + Ok(Self { semaphore, size }) + } + + /// A permit, if one is free now. + pub fn try_acquire(&self) -> Option { + self.semaphore.clone().try_acquire_owned().ok() + } + + /// Wait until a permit is free and take it. Dropping the future before + /// then takes none. + pub async fn acquire(&self) -> Result { + self.semaphore + .clone() + .acquire_owned() + .await + .map_err(|_| Error::internal("the lazy far gather permit pool was closed")) + } + + /// Permits the pool holds while none is taken. + pub fn size(&self) -> usize { + self.size + } + + /// Permits taken now, i.e. gathers issued beyond the ordinary window + /// whose reads have not returned. + pub fn in_flight(&self) -> usize { + // Permits only return to the pool when dropped, so at most `size` are free. + self.size - self.semaphore.available_permits() + } +} + /// Enables the lazy layered full-precision scan (`0` or `1`). pub const LAZY_FULL_ENV: &str = "LANCE_RQ_LAZY_FULL"; /// Staleness window W: a gather may be issued while up to W earlier probes are unscored. @@ -696,8 +849,13 @@ pub const LAZY_DENSE_ENV: &str = "LANCE_RQ_LAZY_DENSE"; pub const LAZY_MAX_RUNS_ENV: &str = "LANCE_RQ_LAZY_MAX_RUNS"; /// Aligned sparse bytes, as a fraction of the plane, from which a whole read is cheaper. pub const LAZY_DENSE_BYTES_FRACTION_ENV: &str = "LANCE_RQ_LAZY_DENSE_BYTES_FRACTION"; -/// Background promotion policy: `off` or `bg:N` (promote after N sparse -/// gathers). Backends that gate plane admission never promote. +/// Background promotion policy: `off` (default), `auto` or `bg:N` (promote +/// after N sparse gathers). `auto` is `off` for an +/// [`OriginLatencyClass::High`] origin and `bg:1` otherwise, resolved when an +/// index opens (see [`LazyPromotion::resolve`]). The default is `off` on every +/// origin because the V10 workload measured no gain from promotions on local +/// NVMe either, at the cost of extra plane reads. Backends that gate plane +/// admission never promote. pub const LAZY_PROMOTE_ENV: &str = "LANCE_RQ_LAZY_PROMOTE"; /// Most background promotions running at once. pub const LAZY_PROMOTE_INFLIGHT_ENV: &str = "LANCE_RQ_LAZY_PROMOTE_INFLIGHT"; @@ -710,30 +868,62 @@ pub const LAZY_INLINE_ROWS_ENV: &str = "LANCE_RQ_LAZY_INLINE_ROWS"; /// dense anyway (large `k`), or when the heap is still not full after the /// probes that hold `k` rows were scored (filters), where waiting would issue /// the gathers one probe at a time. `0` is an ablation knob: such gathers -/// wait for the heap to fill or for every earlier probe to be scored. With -/// [`LAZY_DENSE_TO_EAGER_ENV`] on, the large-`k` probes of the queries it -/// routes are scored eagerly instead and never gathered. +/// wait for the heap to fill or for every earlier probe to be scored. The +/// large-`k` probes that [`LAZY_DENSE_TO_EAGER_ENV`] routes are scored +/// eagerly instead and never gathered. pub const LAZY_EAGER_BEFORE_FULL_ENV: &str = "LANCE_RQ_LAZY_EAGER_BEFORE_FULL"; -/// Whether a probe that `k` and the partition sizes predict to be gathered -/// whole is loaded and scored by the eager scan instead (`0` or `1`, default -/// `1`): a probe scored before the probes ahead of it can hold `k` rows, -/// whose gather selects every row, and every probe when the gathers issued -/// once the heap fills are expected to read whole planes (large `k`). The -/// eager load reads the partition's planes together (all at once on backends -/// that do not gate plane admission), with as many probes in flight as the -/// eager scan prepares, instead of the sign plane at staging and the ex -/// planes a round trip later within the gather window. Only queries without -/// a prefilter (deletions included) or upper distance bound are routed: a -/// gather with an infinite threshold selects only the accepted rows below -/// the bound, which the probe's sign plane tells. The `sparse` gather policy -/// ([`LAZY_DENSE_ENV`]) never reads a whole plane and routes no probe. `0` is -/// an ablation knob: such probes take the lazy pipeline. +/// Which probes that `k` and the partition sizes predict to be gathered +/// whole are loaded and scored by the eager scan instead: `off`, `origin` +/// (default) or `all` (see [`DenseToEager`]; the flag's former `0` and `1` +/// mean `off` and `all`). A probe is predicted dense when it is scored +/// before the probes ahead of it can hold `k` rows, so its gather selects +/// every row, and every probe is when the gathers issued once the heap fills +/// are expected to read whole planes (large `k`). The eager load reads the +/// partition's planes together (all at once on backends that do not gate +/// plane admission), with as many probes in flight as the eager scan +/// prepares, instead of the sign plane at staging and the ex planes a round +/// trip later within the gather window. Only queries without a prefilter +/// (deletions included) or upper distance bound are routed: a gather with an +/// infinite threshold selects only the accepted rows below the bound, which +/// the probe's sign plane tells. The `sparse` gather policy +/// ([`LAZY_DENSE_ENV`]) never reads a whole plane and routes no probe. pub const LAZY_DENSE_TO_EAGER_ENV: &str = "LANCE_RQ_LAZY_DENSE_TO_EAGER"; /// Most coalesced row runs a sparse gather reads from the origin file when the /// persistent tier does not hold the plane; with more, it loads (and admits) /// the whole plane instead. Unlimited by default, and ignored by backends that /// gate plane admission, which would not admit the plane. pub const LAZY_ORIGIN_MAX_RUNS_ENV: &str = "LANCE_RQ_LAZY_ORIGIN_MAX_RUNS"; +/// Largest gap in bytes between the row runs of a sparse gather that reads the +/// origin file, because the persistent tier does not hold the plane, that one +/// request still spans: `auto` (default) or a byte count. Each run is cut back +/// out of the request before decoding, so only the requests and the bytes +/// read change, not the rows returned. `auto` is +/// [`HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES`] for an [`OriginLatencyClass::High`] +/// origin and the object store's block size, the gap of every other read, +/// otherwise (see [`LazyOriginGap`]). +pub const LAZY_ORIGIN_GAP_BYTES_ENV: &str = "LANCE_RQ_LAZY_ORIGIN_GAP_BYTES"; +/// The gap `auto` gives the sparse origin reads of an +/// [`OriginLatencyClass::High`] origin. S3 latency per request is about flat +/// from 60 KiB to 700 KiB, so merging runs up to 256 KiB apart makes about a +/// third fewer S3 requests than the 64 KiB block size. Lazy scans on S3 were +/// faster with it than with 1 MiB, which makes fewer requests but reads more +/// bytes between the runs. +pub const HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES: u64 = 256 * 1024; +/// Staleness window of the gathers of probes that read an ex plane from an +/// [`OriginLatencyClass::High`] origin, because no cache tier held their high +/// or low plane when they were staged (default 64): such a gather may be +/// issued while up to this many earlier probes are unscored, instead of +/// [`LAZY_WINDOW_ENV`]'s. A gather issued beyond the ordinary window holds a +/// permit of the index's pool ([`LAZY_FAR_INFLIGHT_ENV`]) until its reads +/// return; when none is free it waits for a permit or for the ordinary +/// window, whichever comes first, so it is never issued later than without +/// this window. A value no larger than [`LAZY_WINDOW_ENV`]'s turns it off, and +/// it never applies to a low-latency origin. Results are the same either way. +pub const LAZY_FAR_WINDOW_ENV: &str = "LANCE_RQ_LAZY_FAR_WINDOW"; +/// Most gathers of one index, over all its queries, issued beyond the +/// ordinary window under [`LAZY_FAR_WINDOW_ENV`] whose reads have not +/// returned (default 64, at least 1). The pool is sized once per index. +pub const LAZY_FAR_INFLIGHT_ENV: &str = "LANCE_RQ_LAZY_FAR_INFLIGHT"; /// Benchmark knob (`0` or `1`, default `0`): `1` loads a layered partition's /// sign plane before its ex planes on backends that do not gate plane /// admission too, as gated backends always do, instead of loading all three @@ -747,6 +937,12 @@ const DEFAULT_LAZY_PROMOTE_READS: u32 = 1; const DEFAULT_LAZY_PROMOTE_INFLIGHT: usize = 4; /// About the rows one ~100µs CPU dispatch would score. const DEFAULT_LAZY_INLINE_ROWS: usize = 512; +/// Every probe of a 64-probe query. On S3, a window of 64 for every gather +/// cut mean latency to 0.69–0.81x of the ordinary window's at k=100 and +/// k=1000, while 32 gave MS MARCO no gain. +const DEFAULT_LAZY_FAR_WINDOW: usize = 64; +/// As many as Lance's default I/O parallelism for cloud object stores. +const DEFAULT_LAZY_FAR_INFLIGHT: usize = DEFAULT_CLOUD_IO_PARALLELISM; /// [`SEQUENTIAL_PLANE_LOADS_ENV`], read once per process. static SEQUENTIAL_PLANE_LOADS: LazyLock> = LazyLock::new(|| { @@ -768,6 +964,316 @@ fn sequential_plane_loads_from(value: Option<&str>) -> Result { }) } +/// Latency class of an IVF_RQ index's origin reads, those that no cache tier +/// serves: `auto` (default), `low` or `high`. `auto` takes the class the +/// session opening the index declares, and without one the class of its +/// object store: high for a cloud object store ([`ObjectStore::is_cloud`]) and +/// low for local files and memory. A process whose object store is wrapped in +/// a local cache that also caches index files declares `low` on its session +/// instead, since the wrapper hides that from the store's scheme; `low` or +/// `high` here overrides every session. Read once per process; the class is +/// resolved when an index opens, see [`OriginLatencyClass::resolve_with_hint`]. +pub const ORIGIN_LATENCY_ENV: &str = "LANCE_RQ_ORIGIN_LATENCY"; + +/// How slow an index's origin reads are. Reader policies that trade extra +/// bytes or background work for fewer origin requests default to on only for +/// [`Self::High`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum OriginLatencyClass { + /// Local files and memory. + #[default] + Low, + /// Cloud object stores, where every request costs a network round trip. + High, +} + +impl OriginLatencyClass { + /// The class `auto` resolves to for an index read from `object_store`. + pub fn of_store(object_store: &ObjectStore) -> Self { + if object_store.is_cloud() { + Self::High + } else { + Self::Low + } + } + + /// The class of an index read from `object_store` under `setting`, the + /// value of [`ORIGIN_LATENCY_ENV`] with `None` for `auto`, when no session + /// declares one. + pub fn resolve(setting: Option, object_store: &ObjectStore) -> Self { + Self::resolve_with_hint(setting, None, object_store) + } + + /// The class of an index read from `object_store` under `setting`, the + /// value of [`ORIGIN_LATENCY_ENV`] with `None` for `auto`, opened by a + /// session that declares `hint` for the indexes it opens (`None` when it + /// declares none). An explicit setting wins over the hint, and the hint + /// over the store's own class ([`Self::of_store`]): only the serving + /// process knows whether a wrapper around the store serves index reads + /// from a local cache. + pub fn resolve_with_hint( + setting: Option, + hint: Option, + object_store: &ObjectStore, + ) -> Self { + setting + .or(hint) + .unwrap_or_else(|| Self::of_store(object_store)) + } + + /// The knob's spelling of the class: `low` or `high`. + pub fn as_str(self) -> &'static str { + match self { + Self::Low => "low", + Self::High => "high", + } + } + + /// Whether reading planes that the index cache holds in `tiers` makes a + /// request to a slow origin: this class is [`Self::High`] and some plane + /// is [`CacheTier::Absent`]. A lazy probe whose ex planes are such is + /// counted in `s3_bound_probes`. + pub fn reads_slow_origin(self, tiers: &[CacheTier]) -> bool { + self == Self::High && tiers.contains(&CacheTier::Absent) + } +} + +impl std::fmt::Display for OriginLatencyClass { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// [`ORIGIN_LATENCY_ENV`], read once per process. +static ORIGIN_LATENCY: LazyLock, String>> = + LazyLock::new(|| { + origin_latency_from(std::env::var(ORIGIN_LATENCY_ENV).ok().as_deref()) + .map_err(|err| err.to_string()) + }); + +/// The class [`ORIGIN_LATENCY_ENV`] sets, `None` for `auto`. The variable is +/// read once per process; an invalid value fails here and every IVF_RQ index +/// open. +pub fn origin_latency_setting() -> Result> { + ORIGIN_LATENCY.clone().map_err(Error::invalid_input) +} + +fn origin_latency_from(value: Option<&str>) -> Result> { + let Some(value) = value else { + return Ok(None); + }; + match value.trim() { + "auto" => Ok(None), + "low" => Ok(Some(OriginLatencyClass::Low)), + "high" => Ok(Some(OriginLatencyClass::High)), + _ => Err(Error::invalid_input(format!( + "{ORIGIN_LATENCY_ENV}={value:?} is invalid, expected auto, low or high" + ))), + } +} + +/// Where a layered index's cache keeps the estimator bounds columns: +/// `lazy` (default) in their own plane entry, read only by High precision and +/// by full precision on a file without error factors, or `eager` in the sign +/// plane, which every scan reads (see [`SignBounds`]). Results are the same +/// either way. Read once per process; the placement is resolved when a +/// layered index opens. +pub const SIGN_BOUNDS_ENV: &str = "LANCE_RQ_SIGN_BOUNDS"; + +/// [`SIGN_BOUNDS_ENV`], read once per process. +static SIGN_BOUNDS: LazyLock> = LazyLock::new(|| { + sign_bounds_from(std::env::var(SIGN_BOUNDS_ENV).ok().as_deref()).map_err(|err| err.to_string()) +}); + +/// The bounds placement [`SIGN_BOUNDS_ENV`] sets. The variable is read once +/// per process; an invalid value fails here and every layered IVF_RQ index +/// open. +pub fn sign_bounds_setting() -> Result { + SIGN_BOUNDS.clone().map_err(Error::invalid_input) +} + +fn sign_bounds_from(value: Option<&str>) -> Result { + let Some(value) = value else { + return Ok(SignBounds::default()); + }; + match value.trim() { + "lazy" => Ok(SignBounds::Lazy), + "eager" => Ok(SignBounds::Eager), + _ => Err(Error::invalid_input(format!( + "{SIGN_BOUNDS_ENV}={value:?} is invalid, expected lazy or eager" + ))), + } +} + +/// Whether an IVF_RQ index, native or layered, keeps the small columns of its +/// storage file (row ids and factors) in memory, so that partition and plane +/// reads fetch only the code and bounds columns from the file: `auto` +/// (default), `on` or `off`. `auto` is on for [`OriginLatencyClass::High`], +/// where each column read is a request, when the cache can admit the store +/// (see [`ResidentColumnsSetting::admits`]). The store is an entry of the +/// index cache, charged in its budget and kept in RAM while an index of the +/// file is live (see [`ResidentColumns`]); [`resident_columns_bytes`] gives +/// its size without I/O. Results are the same either way. Read once per +/// process; resolved when an index opens. +pub const RESIDENT_COLUMNS_ENV: &str = "LANCE_RQ_RESIDENT_COLUMNS"; + +/// The size of an IVF_RQ storage file's resident store, known without +/// reading it. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ResidentStoreSize { + /// Bytes of values the store holds ([`resident_columns_bytes`]). + pub bytes: u64, + /// Bytes an index cache charges for the loaded store + /// ([`resident_store_charge`]). + pub charge: u64, +} + +/// Whether a resident store an index cache charges `charge` bytes for fits +/// the largest entry the cache admits, `max_entry_bytes` +/// ([`lance_core::cache::CacheBackend::max_entry_bytes`]; `None`: the backend reports no limit). Live index handles keep their allocation even when +/// the cache evicts the entry, so this is not a bound on process memory. +pub fn resident_store_fits(charge: u64, max_entry_bytes: Option) -> bool { + max_entry_bytes.is_none_or(|max| charge <= max) +} + +/// The value of [`RESIDENT_COLUMNS_ENV`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum ResidentColumnsSetting { + /// On for an index whose origin reads are [`OriginLatencyClass::High`]. + #[default] + Auto, + /// On whatever the origin. + On, + /// Off: reads fetch every column they return from the file. + Off, +} + +impl ResidentColumnsSetting { + /// Whether an index whose origin reads are `class` keeps its small + /// columns resident, before [`Self::admits`] sizes the store. + pub fn resolve(self, class: OriginLatencyClass) -> bool { + match self { + Self::Auto => class == OriginLatencyClass::High, + Self::On => true, + Self::Off => false, + } + } + + /// Whether a resident store of `store`'s size may be kept in an index + /// cache whose largest admissible entry is `max_entry_bytes`: never when + /// `off`, and only a store with some columns otherwise; `auto` also + /// requires the bytes the cache charges for it to fit its entry limit + /// ([`resident_store_fits`]). An index keeps its small columns resident + /// when both this and [`Self::resolve`] hold. + pub fn admits(self, store: ResidentStoreSize, max_entry_bytes: Option) -> bool { + match self { + Self::Off => false, + Self::On => store.bytes > 0, + Self::Auto => store.bytes > 0 && resident_store_fits(store.charge, max_entry_bytes), + } + } + + /// The knob's spelling of the setting: `auto`, `on` or `off`. + pub fn as_str(self) -> &'static str { + match self { + Self::Auto => "auto", + Self::On => "on", + Self::Off => "off", + } + } +} + +impl std::fmt::Display for ResidentColumnsSetting { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// [`RESIDENT_COLUMNS_ENV`], read once per process. +static RESIDENT_COLUMNS: LazyLock> = + LazyLock::new(|| { + resident_columns_from(std::env::var(RESIDENT_COLUMNS_ENV).ok().as_deref()) + .map_err(|err| err.to_string()) + }); + +/// The setting of [`RESIDENT_COLUMNS_ENV`]. The variable is read once per +/// process; an invalid value fails here and every IVF_RQ index open. +pub fn resident_columns_setting() -> Result { + RESIDENT_COLUMNS.clone().map_err(Error::invalid_input) +} + +fn resident_columns_from(value: Option<&str>) -> Result { + let Some(value) = value else { + return Ok(ResidentColumnsSetting::default()); + }; + match value.trim() { + "auto" => Ok(ResidentColumnsSetting::Auto), + "on" => Ok(ResidentColumnsSetting::On), + "off" => Ok(ResidentColumnsSetting::Off), + _ => Err(Error::invalid_input(format!( + "{RESIDENT_COLUMNS_ENV}={value:?} is invalid, expected auto, on or off" + ))), + } +} + +/// How long an IVF_RQ index file's resident store is kept in RAM: `index` +/// (default) or `process`. Read once per process; an invalid value fails +/// every IVF_RQ index open. +pub const RESIDENT_LIFETIME_ENV: &str = "LANCE_RQ_RESIDENT_LIFETIME"; + +/// The value of [`RESIDENT_LIFETIME_ENV`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum ResidentLifetime { + /// Held by the live indexes of the file and by an ordinary cache entry. + /// Eviction frees the store after its last index handle drops. + #[default] + Index, + /// Also held by a process-wide strong reference after loading. The + /// cache entry remains evictable; process ownership is outside its budget. + Process, +} + +impl ResidentLifetime { + /// The knob's spelling of the lifetime: `index` or `process`. + pub fn as_str(self) -> &'static str { + match self { + Self::Index => "index", + Self::Process => "process", + } + } +} + +impl std::fmt::Display for ResidentLifetime { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// [`RESIDENT_LIFETIME_ENV`], read once per process. +static RESIDENT_LIFETIME: LazyLock> = + LazyLock::new(|| { + resident_lifetime_from(std::env::var(RESIDENT_LIFETIME_ENV).ok().as_deref()) + .map_err(|err| err.to_string()) + }); + +/// The setting of [`RESIDENT_LIFETIME_ENV`]. +pub fn resident_lifetime_setting() -> Result { + RESIDENT_LIFETIME.clone().map_err(Error::invalid_input) +} + +fn resident_lifetime_from(value: Option<&str>) -> Result { + let Some(value) = value else { + return Ok(ResidentLifetime::default()); + }; + match value.trim() { + "index" => Ok(ResidentLifetime::Index), + "process" => Ok(ResidentLifetime::Process), + _ => Err(Error::invalid_input(format!( + "{RESIDENT_LIFETIME_ENV}={value:?} is invalid, expected index or process" + ))), + } +} + /// Parse a `0`/`1` (or `false`/`true`) environment flag. fn parse_flag(name: &str, value: &str) -> Result { match value.trim() { @@ -804,25 +1310,143 @@ pub enum DenseGatherMode { } /// Whether sparse gathers promote whole planes into RAM in the background. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum LazyPromotion { + /// [`Self::Off`] for an [`OriginLatencyClass::High`] origin and `bg:1` + /// otherwise; see [`Self::resolve`]. An index resolves it when it opens, + /// so a scan never sees it (and would promote nothing). + Auto, /// Only whole-plane gathers admit planes. + #[default] Off, /// Promote a plane after this many sparse gathers of it while not /// resident, on backends that do not gate plane admission. Background { reads: u32 }, } +impl LazyPromotion { + /// The policy of an index whose origin is `class`: [`Self::Auto`] + /// resolves, an explicit policy applies whatever the origin. + /// + /// A promotion reads the whole plane again, from the origin when no + /// cache tier holds it, and its admission evicts sign planes that later + /// queries then miss. On S3 that costs more origin requests than the + /// promoted planes save: without promotions the lazy scan of MS MARCO at + /// k=100 made 175 instead of 235 requests per query and read 12.0 + /// instead of 21.6 MiB. + pub fn resolve(self, class: OriginLatencyClass) -> Self { + match (self, class) { + (Self::Auto, OriginLatencyClass::High) => Self::Off, + (Self::Auto, OriginLatencyClass::Low) => Self::Background { + reads: DEFAULT_LAZY_PROMOTE_READS, + }, + (policy, _) => policy, + } + } +} + +/// The knob's spelling of the policy: `auto`, `off` or `bg:N`. +impl std::fmt::Display for LazyPromotion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Auto => f.write_str("auto"), + Self::Off => f.write_str("off"), + Self::Background { reads } => write!(f, "bg:{reads}"), + } + } +} + +/// The value of [`LAZY_DENSE_TO_EAGER_ENV`]: which probes predicted to be +/// gathered whole, of non-empty partitions whose high and low planes are not +/// both resident, the eager scan loads and scores instead. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)] +pub enum DenseToEager { + /// None: every such probe takes the lazy pipeline. + Off, + /// Those whose high or low plane no cache tier holds on an + /// [`OriginLatencyClass::High`] origin, where the eager load saves a + /// round trip to the origin (see + /// [`OriginLatencyClass::reads_slow_origin`]). Where the planes are + /// local, loading them eagerly saves little and delays scoring: routing + /// every such probe on NVMe raised the p99 of Coyo k=10000 queries from + /// 68.4 to 90.5 ms. + #[default] + Origin, + /// Every one, wherever its planes are. + All, +} + +impl DenseToEager { + /// Whether a predicted-dense probe of a non-empty partition whose high + /// and low planes are in `tiers`, not both resident, is loaded and + /// scored by the eager scan on an origin of `class`. + pub fn routes(self, class: OriginLatencyClass, tiers: &[CacheTier]) -> bool { + match self { + Self::Off => false, + Self::Origin => class.reads_slow_origin(tiers), + Self::All => true, + } + } + + /// The knob's spelling of the mode: `off`, `origin` or `all`. + pub fn as_str(self) -> &'static str { + match self { + Self::Off => "off", + Self::Origin => "origin", + Self::All => "all", + } + } +} + +impl std::fmt::Display for DenseToEager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// The value of [`LAZY_ORIGIN_GAP_BYTES_ENV`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum LazyOriginGap { + /// [`HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES`] for an + /// [`OriginLatencyClass::High`] origin, where every request is a round + /// trip; the object store's block size otherwise. + #[default] + Auto, + /// This many bytes, whatever the origin. + Bytes(u64), +} + +impl LazyOriginGap { + /// The gap within which the sparse origin reads of an index whose origin + /// is `class` merge row runs into one request, or `None` for the object + /// store's block size, which every other read of the file uses. + pub fn coalesce_gap(self, class: OriginLatencyClass) -> Option { + match (self, class) { + (Self::Bytes(bytes), _) => Some(bytes), + (Self::Auto, OriginLatencyClass::High) => Some(HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES), + (Self::Auto, OriginLatencyClass::Low) => None, + } + } + + /// The gap in bytes for an origin of `class` whose object store merges + /// ranges within `block_size` bytes by default. + pub fn resolve(self, class: OriginLatencyClass, block_size: u64) -> u64 { + self.coalesce_gap(class).unwrap_or(block_size) + } +} + /// Settings of the lazy layered full-precision scan, which bounds every row /// from the sign plane and reads the ex planes of the survivors only. #[derive(Debug, Clone, Copy, PartialEq)] pub struct LayeredLazyConfig { pub enabled: bool, - /// Probes a gather may run ahead of scoring; see [`LAZY_WINDOW_ENV`]. + /// Probes a gather may run ahead of scoring; see [`LAZY_WINDOW_ENV`], + /// and [`Self::far_window`] for gathers that read a slow origin. pub window: usize, pub dense: DenseGatherMode, pub max_runs: usize, pub dense_bytes_fraction: f64, + /// See [`LAZY_PROMOTE_ENV`]; an index holds it resolved for its origin. pub promote: LazyPromotion, pub promote_inflight: usize, pub inline_rows: usize, @@ -830,8 +1454,14 @@ pub struct LayeredLazyConfig { pub eager_before_full: bool, /// See [`LAZY_ORIGIN_MAX_RUNS_ENV`]; `usize::MAX` never falls back. pub origin_max_runs: usize, + /// See [`LAZY_ORIGIN_GAP_BYTES_ENV`]. + pub origin_gap: LazyOriginGap, /// See [`LAZY_DENSE_TO_EAGER_ENV`]. - pub dense_to_eager: bool, + pub dense_to_eager: DenseToEager, + /// See [`LAZY_FAR_WINDOW_ENV`] and [`Self::active_far_window`]. + pub far_window: usize, + /// See [`LAZY_FAR_INFLIGHT_ENV`]. + pub far_inflight: usize, } impl Default for LayeredLazyConfig { @@ -842,14 +1472,15 @@ impl Default for LayeredLazyConfig { dense: DenseGatherMode::Cost, max_runs: DEFAULT_LAZY_MAX_RUNS, dense_bytes_fraction: DEFAULT_LAZY_DENSE_BYTES_FRACTION, - promote: LazyPromotion::Background { - reads: DEFAULT_LAZY_PROMOTE_READS, - }, + promote: LazyPromotion::Off, promote_inflight: DEFAULT_LAZY_PROMOTE_INFLIGHT, inline_rows: DEFAULT_LAZY_INLINE_ROWS, eager_before_full: true, origin_max_runs: usize::MAX, - dense_to_eager: true, + origin_gap: LazyOriginGap::Auto, + dense_to_eager: DenseToEager::Origin, + far_window: DEFAULT_LAZY_FAR_WINDOW, + far_inflight: DEFAULT_LAZY_FAR_INFLIGHT, } } } @@ -910,13 +1541,16 @@ impl LayeredLazyConfig { } if let Some(value) = lookup(LAZY_PROMOTE_ENV) { config.promote = match value.trim() { + "auto" => LazyPromotion::Auto, "off" => LazyPromotion::Off, other => other .strip_prefix("bg:") .and_then(|reads| reads.parse::().ok()) .filter(|reads| *reads > 0) .map(|reads| LazyPromotion::Background { reads }) - .ok_or_else(|| invalid(LAZY_PROMOTE_ENV, &value, "off or bg:N with N >= 1"))?, + .ok_or_else(|| { + invalid(LAZY_PROMOTE_ENV, &value, "auto, off or bg:N with N >= 1") + })?, }; } config.promote_inflight = count(LAZY_PROMOTE_INFLIGHT_ENV, config.promote_inflight, 1)?; @@ -925,11 +1559,60 @@ impl LayeredLazyConfig { config.eager_before_full = parse_flag(LAZY_EAGER_BEFORE_FULL_ENV, &value)?; } config.origin_max_runs = count(LAZY_ORIGIN_MAX_RUNS_ENV, config.origin_max_runs, 0)?; + if let Some(value) = lookup(LAZY_ORIGIN_GAP_BYTES_ENV) { + config.origin_gap = match value.trim() { + "auto" => LazyOriginGap::Auto, + bytes => bytes + .parse::() + .map(LazyOriginGap::Bytes) + .map_err(|_| { + invalid(LAZY_ORIGIN_GAP_BYTES_ENV, &value, "auto or a byte count") + })?, + }; + } if let Some(value) = lookup(LAZY_DENSE_TO_EAGER_ENV) { - config.dense_to_eager = parse_flag(LAZY_DENSE_TO_EAGER_ENV, &value)?; + // `0`, `1` and their `false`/`true` spellings are the values of + // the former on/off flag. + config.dense_to_eager = match value.trim() { + "off" | "0" | "false" => DenseToEager::Off, + "origin" => DenseToEager::Origin, + "all" | "1" | "true" => DenseToEager::All, + _ => { + return Err(invalid( + LAZY_DENSE_TO_EAGER_ENV, + &value, + "off, origin or all", + )); + } + }; + } + config.far_window = count(LAZY_FAR_WINDOW_ENV, config.far_window, 0)?; + if let Some(value) = lookup(LAZY_FAR_INFLIGHT_ENV) { + // A permit pool holds at most `Semaphore::MAX_PERMITS`. + config.far_inflight = value + .trim() + .parse::() + .ok() + .filter(|permits| (1..=Semaphore::MAX_PERMITS).contains(permits)) + .ok_or_else(|| { + invalid( + LAZY_FAR_INFLIGHT_ENV, + &value, + &format!("an integer from 1 to {}", Semaphore::MAX_PERMITS), + ) + })?; } Ok(config) } + + /// The staleness window of the gathers of probes that read an ex plane + /// from an origin of `class`, [`Self::far_window`] when it applies, or + /// `None` when they keep [`Self::window`]: the origin is not + /// [`OriginLatencyClass::High`], or the far window is no wider. + pub fn active_far_window(&self, class: OriginLatencyClass) -> Option { + (class == OriginLatencyClass::High && self.far_window > self.window) + .then_some(self.far_window) + } } /// Where the lazy scan reads one ex plane of a partition from. @@ -1084,10 +1767,24 @@ pub struct IvfQuantizationStorage { /// Asynchronous batch remapper (tagged histories). batch_remapper: Option>, plane_access: PlaneAccessTracker, + /// See [`Self::origin_latency`]. + origin_latency: OriginLatencyClass, + /// See [`Self::sign_bounds`]. + sign_bounds: SignBounds, + /// See [`Self::resident_columns`]. + resident_columns: ResidentColumns, + /// See [`Self::resident_columns_enabled`]. + resident_columns_enabled: bool, + /// The file the storage reads, when bound to the runtime handles every + /// open of it shares; see [`Self::with_index_file`]. `None` keeps the + /// far gather pools to the plane access tracker. + index_file: Option, } impl DeepSizeOf for IvfQuantizationStorage { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // The resident store is charged as its own entry of the index cache + // (`ResidentColumnsEntry`), not with every storage reading it. self.metadata.deep_size_of_children(context) + self.ivf.deep_size_of_children(context) + self.plane_access.deep_size_of_children(context) @@ -1172,6 +1869,11 @@ impl IvfQuantizationStorage { frag_reuse_index, batch_remapper: None, plane_access: Default::default(), + origin_latency: OriginLatencyClass::default(), + sign_bounds: SignBounds::default(), + resident_columns: ResidentColumns::default(), + resident_columns_enabled: false, + index_file: None, }) } @@ -1205,6 +1907,11 @@ impl IvfQuantizationStorage { frag_reuse_index, batch_remapper: None, plane_access: Default::default(), + origin_latency: OriginLatencyClass::default(), + sign_bounds: SignBounds::default(), + resident_columns: ResidentColumns::default(), + resident_columns_enabled: false, + index_file: None, } } @@ -1230,6 +1937,147 @@ impl IvfQuantizationStorage { self } + /// Set the latency class of reads from this storage's file, resolved + /// when the index opened. + pub fn with_origin_latency(mut self, class: OriginLatencyClass) -> Self { + self.origin_latency = class; + self + } + + /// The latency class of reads from this storage's file, see + /// [`ORIGIN_LATENCY_ENV`]. [`OriginLatencyClass::Low`] unless the index + /// set it with [`Self::with_origin_latency`]. + pub fn origin_latency(&self) -> OriginLatencyClass { + self.origin_latency + } + + /// Set where the index cache keeps a layered index's bounds columns, + /// resolved when the index opened. + pub fn with_sign_bounds(mut self, sign_bounds: SignBounds) -> Self { + self.sign_bounds = sign_bounds; + self + } + + /// Where the index cache keeps a layered index's bounds columns, see + /// [`SIGN_BOUNDS_ENV`]. [`SignBounds::Lazy`] unless the index set it with + /// [`Self::with_sign_bounds`]. + pub fn sign_bounds(&self) -> SignBounds { + self.sign_bounds + } + + /// This storage's handle on the resident store of its file, which owns + /// the store while the storage lives. It loads only while + /// [`Self::resident_columns_enabled`]. + pub fn resident_columns(&self) -> &ResidentColumns { + &self.resident_columns + } + + /// Read the small columns through `store`, the handle an open binds to + /// the store every live index of the file shares + /// ([`ResidentColumns::in_index_cache`]), so that the store loads once. + pub fn with_resident_columns(mut self, store: ResidentColumns) -> Self { + self.resident_columns = store; + self + } + + /// Bind the lazy scan's far gathers to the permit pools that every open + /// of index file `file` in the process shares, one per pool size, so that + /// indexes of the file opened at once, a re-open and a state read back + /// from a persistent cache tier bound their far gathers together. An + /// index binds when it opens and when it is reconstructed. + pub fn with_index_file(mut self, file: IndexFileKey) -> Self { + self.index_file = Some(file); + self + } + + /// Set whether reads take the file's small columns from the resident + /// store, resolved when the index opened. + pub fn with_resident_columns_enabled(mut self, enabled: bool) -> Self { + self.resident_columns_enabled = enabled; + self + } + + /// Whether reads take the file's small columns from the resident store, + /// see [`RESIDENT_COLUMNS_ENV`], and fetch only the others from the file. + /// `false` unless the index set it with + /// [`Self::with_resident_columns_enabled`]. + pub fn resident_columns_enabled(&self) -> bool { + self.resident_columns_enabled + } + + /// Bytes the resident store holds for this storage's file, loaded or + /// not; see [`resident_columns_bytes`]. + pub fn resident_columns_bytes(&self) -> u64 { + resident_columns_bytes(&self.schema(), self.num_rows()) + } + + /// The size of the resident store of this storage's file, loaded or + /// not, and what an index cache charges for it. It reads nothing. + pub fn resident_store_size(&self) -> Result { + let bytes = self.resident_columns_bytes(); + if bytes == 0 { + return Ok(ResidentStoreSize::default()); + } + Ok(ResidentStoreSize { + bytes, + charge: resident_store_charge(&self.reader)?, + }) + } + + /// The cache key of plane `plane` of `partition` under this storage's + /// bounds placement. + pub fn plane_key(&self, partition: usize, plane: u8) -> PlaneKey { + PlaneKey { + partition, + plane, + sign_bounds: self.sign_bounds, + } + } + + /// The tier the index cache would serve plane `plane` of `partition` + /// from, checked without reading the plane or counting as an access. + /// [`CacheTier::Absent`] means a read goes to this storage's file, whose + /// latency is [`Self::origin_latency`]. + pub async fn plane_tier( + &self, + partition: usize, + plane: u8, + cache: &WeakLanceCache, + ) -> CacheTier { + cache + .peek_tier_with_key(&self.plane_key(partition, plane)) + .await + } + + /// The permits of the lazy scan's gathers issued beyond the ordinary + /// window ([`LAZY_FAR_WINDOW_ENV`]), `config.far_inflight` of them, + /// shared by every query of this index and by the reconstructions of its + /// cached state, and, once bound to its file ([`Self::with_index_file`]), + /// by every open of the file in the process. The plane access tracker + /// keeps the pool, so a cached state holds it. An index's config is + /// fixed when it opens, so it keeps one pool; a config of another size + /// (tests replace it) takes another pool, and gathers holding the old + /// one's permits return them there. + pub fn lazy_far_permits(&self, config: &LayeredLazyConfig) -> Result { + let size = config.far_inflight; + let mut pool = self + .plane_access + .lazy + .far_permits + .lock() + .unwrap_or_else(|e| e.into_inner()); + match pool.as_ref() { + Some(permits) if permits.size() == size => Ok(permits.clone()), + _ => { + let permits = match &self.index_file { + Some(file) => LazyFarPermits::for_file(file, size)?, + None => LazyFarPermits::try_new(size)?, + }; + Ok(pool.insert(permits).clone()) + } + } + } + pub fn reader(&self) -> &FileReader { &self.reader } @@ -1275,21 +2123,25 @@ impl IvfQuantizationStorage { /// scheduler also records into the sink (a cheap clone that shares all /// cached metadata, so no file is re-opened). When `None`, the normal /// uninstrumented reader is used. + /// + /// With [`Self::resident_columns_enabled`], only the code and bounds + /// columns are read from the file. pub async fn load_partition( &self, part_id: usize, io_stats: Option, ) -> Result { let range = self.ivf.row_range(part_id); + let schema: SchemaRef = Arc::new(self.reader.schema().as_ref().into()); let batch = if range.is_empty() { - let schema = self.reader.schema(); - let arrow_schema = arrow_schema::Schema::from(schema.as_ref()); - RecordBatch::new_empty(Arc::new(arrow_schema)) + RecordBatch::new_empty(schema) + } else if let Some(store) = self.resident_store(&schema, io_stats.as_ref()).await? { + let params = ReadBatchParams::Range(range.clone()); + let reader = self.file_reader(io_stats.as_ref(), None); + self.read_with_resident_columns(store, schema, range, None, params, &reader) + .await? } else { - let reader = match &io_stats { - Some(io_stats) => Cow::Owned(self.reader.with_io_stats(io_stats.recorder())), - None => Cow::Borrowed(&self.reader), - }; + let reader = self.file_reader(io_stats.as_ref(), None); let batches = reader .read_stream( ReadBatchParams::Range(range), @@ -1300,7 +2152,6 @@ impl IvfQuantizationStorage { .await? .try_collect::>() .await?; - let schema = Arc::new(self.reader.schema().as_ref().into()); concat_batches(&schema, batches.iter())? }; if let Some(remapping) = &self.batch_remapper { @@ -1326,60 +2177,50 @@ impl IvfQuantizationStorage { ) } - /// Warm every plane entry. Backend admission enforces its byte budget. + /// Warm every plane entry a full-precision scan reads. Backend admission + /// enforces its byte budget. /// /// A backend that gates lower planes on their sign plane is warmed plane - /// by plane: all sign planes first, then the high and low planes of - /// partitions whose sign plane stayed resident. Other backends admit plane - /// entries like any entry, so they are warmed partition by partition, - /// which loads (and persists) every plane whatever fits in RAM, as the - /// native partition prewarm does. + /// by plane: all sign planes first, then the other planes of partitions + /// whose sign plane stayed resident. Other backends admit plane entries + /// like any entry, so they are warmed partition by partition, which loads + /// (and persists) every plane whatever fits in RAM, as the native + /// partition prewarm does. pub async fn prewarm_planes(&self, cache: &WeakLanceCache) -> Result<()> { + use super::bq::layered::RQPrecision; + let planes: Vec = std::iter::once(0) + .chain(self.planes_after_sign(RQPrecision::Full)) + .collect(); if !cache.plane_admission_gated() { for part_id in 0..self.num_partitions() { - for plane in 0..=2 { + for &plane in &planes { cache - .get_or_insert_with_key( - PlaneKey { - partition: part_id, - plane, - }, - || async { - Ok(PlaneBatch( - self.read_plane(part_id, plane, None, None).await?, - )) - }, - ) + .get_or_insert_with_key(self.plane_key(part_id, plane), || async { + Ok(PlaneBatch( + self.read_plane(part_id, plane, None, None).await?, + )) + }) .await?; } } return Ok(()); } - for plane in 0..=2 { + for &plane in &planes { for part_id in 0..self.num_partitions() { if plane > 0 && cache - .get_resident_with_key(&PlaneKey { - partition: part_id, - plane: 0, - }) + .get_resident_with_key(&self.plane_key(part_id, 0)) .await .is_none() { continue; } cache - .get_or_insert_with_key( - PlaneKey { - partition: part_id, - plane, - }, - || async { - Ok(PlaneBatch( - self.read_plane(part_id, plane, None, None).await?, - )) - }, - ) + .get_or_insert_with_key(self.plane_key(part_id, plane), || async { + Ok(PlaneBatch( + self.read_plane(part_id, plane, None, None).await?, + )) + }) .await?; } } @@ -1391,6 +2232,28 @@ impl IvfQuantizationStorage { matches!(self.quantizer(), Ok(Quantizer::Rabit(rq)) if rq.metadata_ref().layered) } + /// The planes a layered partition load at `precision` reads besides the + /// sign plane, in the order the eager bounds placement assembles its + /// columns: the bounds plane when the scan prunes with the bounds and the + /// sign plane leaves them out, then the ex planes the precision scores. + /// High precision prunes with the high bounds; full precision prunes + /// with the error factors, or with the full bounds on a file without them. + fn planes_after_sign(&self, precision: super::bq::layered::RQPrecision) -> Vec { + use super::bq::layered::RQPrecision; + let has_error_factors = self.reader.schema().field(ERROR_FACTORS_COLUMN).is_some(); + let (reads_bounds, last_ex_plane) = match precision { + RQPrecision::Sign => (false, 0), + RQPrecision::High => (true, 1), + RQPrecision::Full => (!has_error_factors, 2), + }; + let mut planes = Vec::with_capacity(usize::from(last_ex_plane) + 1); + if reads_bounds && self.sign_bounds == SignBounds::Lazy { + planes.push(SIGN_BOUNDS_PLANE); + } + planes.extend(1..=last_ex_plane); + planes + } + /// Read independent plane entries through the existing persistent cache. pub async fn load_partition_at_precision( &self, @@ -1408,23 +2271,21 @@ impl IvfQuantizationStorage { } return self.load_partition(part_id, io_stats).await; } - let last = match precision { - RQPrecision::Sign => 0, - RQPrecision::High => 1, - RQPrecision::Full => 2, - }; + let after_sign = self.planes_after_sign(precision); let load_plane = |plane| self.load_plane_entry(part_id, plane, cache, io_stats.clone()); let planes = if sequential_plane_loads()? || cache.plane_admission_gated() { - // Admit the sign dependency first. The high and low reads can then + // Admit the sign dependency first. The other reads can then // overlap without changing admission policy or assembled column order. let sign = load_plane(0).await?; - let ex = futures::future::try_join_all((1..=last).map(load_plane)).await?; - std::iter::once(sign).chain(ex).collect() + let rest = + futures::future::try_join_all(after_sign.iter().copied().map(load_plane)).await?; + std::iter::once(sign).chain(rest).collect() } else { // Admission does not depend on the sign plane's residency, so a // missed partition reads its planes in one round trip. - // `try_join_all` returns them in plane order. - futures::future::try_join_all((0..=last).map(load_plane)).await? + // `try_join_all` returns them in the order requested. + futures::future::try_join_all(std::iter::once(0).chain(after_sign).map(load_plane)) + .await? }; let mut fields = Vec::new(); let mut columns = Vec::new(); @@ -1471,10 +2332,7 @@ impl IvfQuantizationStorage { cache: &WeakLanceCache, io_stats: Option, ) -> Result> { - let key = PlaneKey { - partition: part_id, - plane, - }; + let key = self.plane_key(part_id, plane); if !cache.plane_admission_gated() { return cache .get_or_insert_with_key(key, || async { @@ -1487,10 +2345,7 @@ impl IvfQuantizationStorage { } let sign_resident = plane == 0 || cache - .get_resident_with_key(&PlaneKey { - partition: part_id, - plane: 0, - }) + .get_resident_with_key(&self.plane_key(part_id, 0)) .await .is_some(); let batch = if sign_resident { @@ -1665,10 +2520,7 @@ impl IvfQuantizationStorage { ) -> Result { let stats = layered_stats::counters(); let started = Instant::now(); - let key = PlaneKey { - partition: part_id, - plane, - }; + let key = self.plane_key(part_id, plane); if source == PlaneSource::Resident && let Some(batch) = cache.get_resident_with_key(&key).await { @@ -1686,7 +2538,7 @@ impl IvfQuantizationStorage { let (batch, from_origin) = match persisted { Some(batch) => (batch.0.clone(), false), None => ( - self.read_plane(part_id, plane, Some(rows.to_vec()), io_stats) + self.read_origin_rows(part_id, plane, rows, config, io_stats.as_ref()) .await?, true, ), @@ -1798,25 +2650,26 @@ impl IvfQuantizationStorage { )); } let indices = arrow_array::UInt32Array::from(rows.clone()); - // The three planes are independent. Overlap their cache/origin reads, + // The planes are independent. Overlap their cache/origin reads, // while preserving plane order when assembling the full row schema. - let batches = futures::future::try_join_all((0..=2).map(|plane| { + let planes = std::iter::once(0).chain(self.planes_after_sign(RQPrecision::Full)); + let batches = futures::future::try_join_all(planes.map(|plane| { let rows = &rows; let indices = &indices; let io_stats = io_stats.clone(); async move { - let key = PlaneKey { - partition: part_id, - plane, - }; + let key = self.plane_key(part_id, plane); let resident = cache.get_resident_with_key(&key).await; let was_resident = resident.is_some(); - let cached = if plane == 0 && resident.is_none() { + // Sign codes and bounds keep an IPC body, which has no row + // directory, so their persisted entries are read whole. + let row_addressable = plane != 0 && plane != SIGN_BOUNDS_PLANE; + let cached = if !row_addressable && resident.is_none() { cache.get_without_promotion_with_key(&key).await } else { resident }; - let selected_cached = if plane > 0 && cached.is_none() { + let selected_cached = if row_addressable && cached.is_none() { cache.get_rows_with_key(&key, rows).await } else { None @@ -1851,10 +2704,7 @@ impl IvfQuantizationStorage { if plane > 0 && !was_resident && cache - .get_resident_with_key(&PlaneKey { - partition: part_id, - plane: 0, - }) + .get_resident_with_key(&self.plane_key(part_id, 0)) .await .is_some() { @@ -1884,18 +2734,12 @@ impl IvfQuantizationStorage { }; if promote { cache - .get_or_insert_with_key( - PlaneKey { - partition: part_id, - plane, - }, - || async { - Ok(super::bq::layered::PlaneBatch( - self.read_plane(part_id, plane, None, io_stats.clone()) - .await?, - )) - }, - ) + .get_or_insert_with_key(self.plane_key(part_id, plane), || async { + Ok(super::bq::layered::PlaneBatch( + self.read_plane(part_id, plane, None, io_stats.clone()) + .await?, + )) + }) .await?; } } @@ -1920,6 +2764,8 @@ impl IvfQuantizationStorage { } /// Read only selected rows of a plane; sorted indices are coalesced by the reader. + /// With [`Self::resident_columns_enabled`], only the plane's code (or + /// bounds) column is read from the file. pub async fn read_plane( &self, part_id: usize, @@ -1927,19 +2773,103 @@ impl IvfQuantizationStorage { rows: Option>, io_stats: Option, ) -> Result { + self.read_plane_with_coalesce_gap(part_id, plane, rows, io_stats, None) + .await + } + + /// [`Self::read_plane`], reading the file through + /// [`FileReader::with_coalesce_gap`] when `coalesce_gap` is set: the + /// ranges of the selected rows that are at most that many bytes apart + /// share one request, instead of those within the object store's block + /// size. The batch is the same for every gap; only the requests and the + /// bytes read differ. + pub async fn read_plane_with_coalesce_gap( + &self, + part_id: usize, + plane: u8, + rows: Option>, + io_stats: Option, + coalesce_gap: Option, + ) -> Result { + self.read_plane_with_stats( + part_id, + plane, + rows, + coalesce_gap, + io_stats.as_ref(), + io_stats.as_ref(), + ) + .await + } + + /// Read the sorted partition offsets `rows` of ex plane `plane` from the + /// origin file for a sparse gather, merging row runs within `config`'s + /// origin gap for the index's origin, and count the read's requests and + /// bytes after coalescing and the rows it gathers. A first-use load of the + /// resident store that the read starts is not counted; its I/O, like the + /// read's, is added to `io_stats`. + async fn read_origin_rows( + &self, + part_id: usize, + plane: u8, + rows: &[u32], + config: &LayeredLazyConfig, + io_stats: Option<&IoStats>, + ) -> Result { + let origin = IoStats::new(); + let batch = self + .read_plane_with_stats( + part_id, + plane, + Some(rows.to_vec()), + config.origin_gap.coalesce_gap(self.origin_latency), + Some(&origin), + io_stats, + ) + .await; + let read = origin.snapshot(); + let stats = layered_stats::counters(); + stats.origin_sparse_requests.add(read.iops); + stats.origin_sparse_bytes.add(read.bytes_read); + stats.origin_sparse_rows.add(rows.len() as u64); + if let Some(io_stats) = io_stats { + io_stats.add_scan_stats(&read); + } + batch + } + + /// [`Self::read_plane_with_coalesce_gap`], recording the plane's reads + /// from the file into `read_stats` and a first-use load of the resident + /// store into `load_stats`, so that a caller can count them apart. + async fn read_plane_with_stats( + &self, + part_id: usize, + plane: u8, + rows: Option>, + coalesce_gap: Option, + read_stats: Option<&IoStats>, + load_stats: Option<&IoStats>, + ) -> Result { + let columns = super::bq::layered::plane_columns(plane, self.sign_bounds); + if columns.is_empty() { + return Err(Error::invalid_input(format!( + "a layered partition has no plane {plane} with {} bounds", + self.sign_bounds + ))); + } let projection = lance_file::versions::reader_projection_from_column_names( self.reader.metadata().version(), self.reader.schema(), - super::bq::layered::plane_columns(plane), + columns, )?; let schema = Arc::new(arrow_schema::Schema::from(projection.schema.as_ref())); let range = self.ivf.row_range(part_id); if range.is_empty() || rows.as_ref().is_some_and(Vec::is_empty) { return Ok(RecordBatch::new_empty(schema)); } - let params = if let Some(rows) = rows { + let params = if let Some(rows) = &rows { let mut ranges: Vec> = Vec::new(); - for row in rows { + for &row in rows { if row as usize >= range.len() { return Err(Error::invalid_input( "candidate offset exceeds partition length", @@ -1956,12 +2886,14 @@ impl IvfQuantizationStorage { } ReadBatchParams::Ranges(ranges.into()) } else { - ReadBatchParams::Range(range) - }; - let reader = match &io_stats { - Some(stats) => Cow::Owned(self.reader.with_io_stats(stats.recorder())), - None => Cow::Borrowed(&self.reader), + ReadBatchParams::Range(range.clone()) }; + let reader = self.file_reader(read_stats, coalesce_gap); + if let Some(store) = self.resident_store(&schema, load_stats).await? { + return self + .read_with_resident_columns(store, schema, range, rows.as_deref(), params, &reader) + .await; + } let batches = reader .read_stream_projected( params, @@ -1976,6 +2908,112 @@ impl IvfQuantizationStorage { concat_batches(&schema, batches.iter()).map_err(Into::into) } + /// This storage's file reader, also recording its I/O into `io_stats`, + /// and merging requested ranges at most `coalesce_gap` bytes apart + /// instead of those within the object store's block size. + fn file_reader( + &self, + io_stats: Option<&IoStats>, + coalesce_gap: Option, + ) -> Cow<'_, FileReader> { + let reader = match io_stats { + Some(stats) => Cow::Owned(self.reader.with_io_stats(stats.recorder())), + None => Cow::Borrowed(&self.reader), + }; + match coalesce_gap { + Some(gap) => Cow::Owned(reader.with_coalesce_gap(gap)), + None => reader, + } + } + + /// The resident store, when [`Self::resident_columns_enabled`] and it + /// keeps some of `schema`'s columns. The first read loads it, and the + /// load's I/O is added to that read's `io_stats`. + async fn resident_store( + &self, + schema: &arrow_schema::Schema, + io_stats: Option<&IoStats>, + ) -> Result> { + let reads_resident = schema.fields().iter().any(|field| is_resident(field)); + if !self.resident_columns_enabled || !reads_resident { + return Ok(None); + } + self.resident_columns + .get_or_load(&self.reader, io_stats) + .await + .map(Some) + } + + /// Read `schema`'s columns at the file rows `range` of a partition, + /// either every row or the sorted partition offsets `rows`, as `params` + /// selects them. Columns `store` keeps are copies of its rows, and only + /// the others are read from the file, through `reader`, so the batch is + /// the one a read of every column from the file returns. + async fn read_with_resident_columns( + &self, + store: &ResidentColumnStore, + schema: SchemaRef, + range: std::ops::Range, + rows: Option<&[u32]>, + params: ReadBatchParams, + reader: &FileReader, + ) -> Result { + if range.end as u64 > store.num_rows() { + return Err(Error::invalid_input(format!( + "partition rows {range:?} exceed the {} rows of the storage file", + store.num_rows() + ))); + } + let file_columns: Vec<&str> = schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .filter(|name| store.column(name).is_none()) + .collect(); + let file_batch = if file_columns.is_empty() { + None + } else { + let projection = lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + &file_columns, + )?; + let file_schema = Arc::new(arrow_schema::Schema::from(projection.schema.as_ref())); + let batches = reader + .read_stream_projected( + params, + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + Some(concat_batches(&file_schema, batches.iter())?) + }; + let start = range.start as u64; + let file_rows = match rows { + Some(rows) => { + UInt64Array::from_iter_values(rows.iter().map(|&row| start + u64::from(row))) + } + None => UInt64Array::from_iter_values(start..range.end as u64), + }; + let mut columns = Vec::with_capacity(schema.fields().len()); + for field in schema.fields() { + let column = match store.column(field.name()) { + Some(column) => column.copy_rows(&file_rows)?, + None => file_batch + .as_ref() + .and_then(|batch| batch.column_by_name(field.name())) + .cloned() + .ok_or_else(|| Error::internal(format!("unread column {}", field.name())))?, + }; + columns.push(column); + } + Ok(RecordBatch::try_new(schema, columns)?) + } + /// Read only `columns` of the given rows; the batch keeps file schema order. async fn read_vector_columns( &self, @@ -2296,19 +3334,30 @@ impl IvfQuantizationStorage { #[cfg(test)] mod tests { use super::{ - DenseGatherMode, LAZY_DENSE_BYTES_FRACTION_ENV, LAZY_DENSE_ENV, LAZY_DENSE_TO_EAGER_ENV, - LAZY_EAGER_BEFORE_FULL_ENV, LAZY_FULL_ENV, LAZY_INLINE_ROWS_ENV, LAZY_MAX_RUNS_ENV, + DenseGatherMode, DenseToEager, HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, IndexFileKey, + LAZY_DENSE_BYTES_FRACTION_ENV, LAZY_DENSE_ENV, LAZY_DENSE_TO_EAGER_ENV, + LAZY_EAGER_BEFORE_FULL_ENV, LAZY_FAR_INFLIGHT_ENV, LAZY_FAR_WINDOW_ENV, LAZY_FULL_ENV, + LAZY_INLINE_ROWS_ENV, LAZY_MAX_RUNS_ENV, LAZY_ORIGIN_GAP_BYTES_ENV, LAZY_ORIGIN_MAX_RUNS_ENV, LAZY_PROMOTE_ENV, LAZY_PROMOTE_INFLIGHT_ENV, LAZY_WINDOW_ENV, - LayeredLazyConfig, LazyPromotion, PlaneSource, QueryScratchCapacity, QueryScratchPool, - SEQUENTIAL_PLANE_LOADS_ENV, compact_prewarm_batches, origin_reads_whole_plane, - plan_plane_gather, sequential_plane_loads, sequential_plane_loads_from, - spawn_prewarm_materialization, + LayeredLazyConfig, LazyFarPermits, LazyOriginGap, LazyPromotion, ORIGIN_LATENCY_ENV, + OriginLatencyClass, PlaneSource, QueryScratchCapacity, QueryScratchPool, + RESIDENT_COLUMNS_ENV, RESIDENT_LIFETIME_ENV, ResidentColumnsSetting, ResidentLifetime, + ResidentStoreSize, SEQUENTIAL_PLANE_LOADS_ENV, SIGN_BOUNDS_ENV, SignBounds, + compact_prewarm_batches, origin_latency_from, origin_latency_setting, + origin_reads_whole_plane, plan_plane_gather, resident_columns_from, + resident_columns_setting, resident_lifetime_from, resident_lifetime_setting, + resident_store_fits, sequential_plane_loads, sequential_plane_loads_from, sign_bounds_from, + sign_bounds_setting, spawn_prewarm_materialization, }; use arrow_array::{Array, ArrayRef, RecordBatch, UInt64Array}; + use futures::FutureExt; use lance_core::Error; + use lance_core::cache::CacheTier; use lance_core::deepsize::DeepSizeOf; + use lance_io::object_store::ObjectStore; use std::collections::HashMap; use std::sync::Arc; + use tokio::sync::Semaphore; #[tokio::test] async fn test_prewarm_materialization_uses_cpu_pool() { @@ -2451,7 +3500,10 @@ mod tests { (LAZY_INLINE_ROWS_ENV, "0"), (LAZY_EAGER_BEFORE_FULL_ENV, "0"), (LAZY_ORIGIN_MAX_RUNS_ENV, "5"), + (LAZY_ORIGIN_GAP_BYTES_ENV, " 65536 "), (LAZY_DENSE_TO_EAGER_ENV, "0"), + (LAZY_FAR_WINDOW_ENV, "8"), + (LAZY_FAR_INFLIGHT_ENV, " 2 "), ]); let config = LayeredLazyConfig::from_lookup(|name| env.get(name).map(|value| value.to_string())) @@ -2469,7 +3521,10 @@ mod tests { inline_rows: 0, eager_before_full: false, origin_max_runs: 5, - dense_to_eager: false, + origin_gap: LazyOriginGap::Bytes(65536), + dense_to_eager: DenseToEager::Off, + far_window: 8, + far_inflight: 2, } ); assert_eq!( @@ -2479,9 +3534,14 @@ mod tests { assert!(!LayeredLazyConfig::default().enabled); assert!(LayeredLazyConfig::default().eager_before_full); assert_eq!(LayeredLazyConfig::default().origin_max_runs, usize::MAX); - assert!(LayeredLazyConfig::default().dense_to_eager); + assert_eq!(LayeredLazyConfig::default().promote, LazyPromotion::Off); + assert_eq!( + LayeredLazyConfig::default().dense_to_eager, + DenseToEager::Origin + ); // The origin run cap does not depend on the eager-before-full switch, - // and dense probes go to the eager scan unless switched off. + // and dense probes whose planes are on a slow origin go to the eager + // scan unless set otherwise. let env = HashMap::from([(LAZY_FULL_ENV, "1"), (LAZY_ORIGIN_MAX_RUNS_ENV, "2")]); let config = LayeredLazyConfig::from_lookup(|name| env.get(name).map(|value| value.to_string())) @@ -2492,7 +3552,7 @@ mod tests { config.origin_max_runs, config.dense_to_eager ), - (true, 2, true) + (true, 2, DenseToEager::Origin) ); let env = HashMap::from([(LAZY_FULL_ENV, "1"), (LAZY_DENSE_TO_EAGER_ENV, "false")]); let config = @@ -2502,7 +3562,7 @@ mod tests { config, LayeredLazyConfig { enabled: true, - dense_to_eager: false, + dense_to_eager: DenseToEager::Off, ..Default::default() } ); @@ -2514,12 +3574,18 @@ mod tests { for (name, value) in [ (LAZY_DENSE_ENV, "dense"), (LAZY_PROMOTE_ENV, "bg:0"), + (LAZY_PROMOTE_ENV, "on"), (LAZY_PROMOTE_INFLIGHT_ENV, "0"), (LAZY_DENSE_BYTES_FRACTION_ENV, "-1"), (LAZY_WINDOW_ENV, "many"), (LAZY_EAGER_BEFORE_FULL_ENV, "on"), (LAZY_ORIGIN_MAX_RUNS_ENV, "-1"), + (LAZY_ORIGIN_GAP_BYTES_ENV, "1MiB"), + (LAZY_ORIGIN_GAP_BYTES_ENV, "-1"), (LAZY_DENSE_TO_EAGER_ENV, "2"), + (LAZY_DENSE_TO_EAGER_ENV, "on"), + (LAZY_FAR_WINDOW_ENV, "wide"), + (LAZY_FAR_INFLIGHT_ENV, "0"), ] { let enabled = |key: &str| -> Option { if key == LAZY_FULL_ENV { @@ -2621,6 +3687,270 @@ mod tests { assert!(!origin_reads_whole_plane(&[], row_bytes, &none)); } + #[test] + fn test_lazy_origin_gap_knob() { + let parse = |gap: &str| { + LayeredLazyConfig::from_lookup(|key| match key { + LAZY_FULL_ENV => Some("1".to_string()), + LAZY_ORIGIN_GAP_BYTES_ENV => Some(gap.to_string()), + _ => None, + }) + .map(|config| config.origin_gap) + }; + let defaults = LayeredLazyConfig::default(); + assert_eq!(defaults.origin_gap, LazyOriginGap::Auto); + assert_eq!(parse(" auto ").unwrap(), LazyOriginGap::Auto); + for bytes in [0, 64 * 1024, 256 * 1024, 1024 * 1024, u64::MAX] { + assert_eq!( + parse(&bytes.to_string()).unwrap(), + LazyOriginGap::Bytes(bytes) + ); + } + for value in ["max", "AUTO", "64k", "256k", "-1", ""] { + let error = parse(value).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(LAZY_ORIGIN_GAP_BYTES_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + // `auto` widens the gap only where every request is a round trip and + // keeps the object store's block size elsewhere; a byte count is the + // gap whatever the origin. + let block_size = 4096; + let classes = [OriginLatencyClass::Low, OriginLatencyClass::High]; + assert_eq!( + classes.map(|class| LazyOriginGap::Auto.coalesce_gap(class)), + [None, Some(HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES)] + ); + assert_eq!( + classes.map(|class| LazyOriginGap::Auto.resolve(class, block_size)), + [block_size, HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES] + ); + for bytes in [0, 64 * 1024, 256 * 1024, 1024 * 1024, u64::MAX] { + let gap = LazyOriginGap::Bytes(bytes); + assert_eq!( + classes.map(|class| gap.coalesce_gap(class)), + [Some(bytes); 2] + ); + assert_eq!( + classes.map(|class| gap.resolve(class, block_size)), + [bytes; 2] + ); + } + } + + #[test] + fn test_lazy_promote_knob() { + let parse = |promote: &str| { + LayeredLazyConfig::from_lookup(|key| match key { + LAZY_FULL_ENV => Some("1".to_string()), + LAZY_PROMOTE_ENV => Some(promote.to_string()), + _ => None, + }) + .map(|config| config.promote) + }; + assert_eq!(LayeredLazyConfig::default().promote, LazyPromotion::Off); + for (value, expected) in [ + (" auto ", LazyPromotion::Auto), + ("off", LazyPromotion::Off), + ("bg:1", LazyPromotion::Background { reads: 1 }), + ("bg:3", LazyPromotion::Background { reads: 3 }), + ] { + assert_eq!(parse(value).unwrap(), expected, "{value:?}"); + } + for value in ["AUTO", "on", "bg:", "bg:0", "bg:-1", ""] { + let error = parse(value).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(LAZY_PROMOTE_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + // `auto` promotes after one sparse gather unless every origin read + // is a round trip; an explicit policy applies whatever the origin. + let classes = [OriginLatencyClass::Low, OriginLatencyClass::High]; + assert_eq!( + classes.map(|class| LazyPromotion::Auto.resolve(class)), + [LazyPromotion::Background { reads: 1 }, LazyPromotion::Off] + ); + for policy in [LazyPromotion::Off, LazyPromotion::Background { reads: 2 }] { + assert_eq!(classes.map(|class| policy.resolve(class)), [policy; 2]); + } + for (policy, spelling) in [ + (LazyPromotion::Auto, "auto"), + (LazyPromotion::Off, "off"), + (LazyPromotion::Background { reads: 2 }, "bg:2"), + ] { + assert_eq!(policy.to_string(), spelling); + assert_eq!(parse(spelling).unwrap(), policy); + } + } + + #[test] + fn test_lazy_dense_to_eager_knob() { + let parse = |mode: &str| { + LayeredLazyConfig::from_lookup(|key| match key { + LAZY_FULL_ENV => Some("1".to_string()), + LAZY_DENSE_TO_EAGER_ENV => Some(mode.to_string()), + _ => None, + }) + .map(|config| config.dense_to_eager) + }; + for (value, expected) in [ + ("off", DenseToEager::Off), + (" origin ", DenseToEager::Origin), + ("all", DenseToEager::All), + // The former flag's values keep their meaning. + ("0", DenseToEager::Off), + ("false", DenseToEager::Off), + ("1", DenseToEager::All), + ("true", DenseToEager::All), + ] { + assert_eq!(parse(value).unwrap(), expected, "{value:?}"); + } + for mode in [DenseToEager::Off, DenseToEager::Origin, DenseToEager::All] { + assert_eq!(mode.to_string(), mode.as_str()); + assert_eq!(parse(mode.as_str()).unwrap(), mode); + } + for value in ["2", "ORIGIN", "on", "auto", ""] { + let error = parse(value).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(LAZY_DENSE_TO_EAGER_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + + // Tiers of the high and low planes of probes whose planes are not + // both resident: some plane on no cache tier, or every plane local. + let (resident, local, absent) = (CacheTier::Resident, CacheTier::Local, CacheTier::Absent); + let (low, high) = (OriginLatencyClass::Low, OriginLatencyClass::High); + let slow = [[resident, absent], [absent, local], [absent, absent]]; + let local_only = [[resident, local], [local, resident], [local, local]]; + for tiers in slow { + assert!(high.reads_slow_origin(&tiers), "{tiers:?}"); + assert!(!low.reads_slow_origin(&tiers), "{tiers:?}"); + assert!(DenseToEager::Origin.routes(high, &tiers), "{tiers:?}"); + assert!(!DenseToEager::Origin.routes(low, &tiers), "{tiers:?}"); + } + for tiers in local_only { + for class in [low, high] { + let context = format!("{class} {tiers:?}"); + assert!(!class.reads_slow_origin(&tiers), "{context}"); + assert!(!DenseToEager::Origin.routes(class, &tiers), "{context}"); + } + } + for tiers in slow.iter().chain(&local_only) { + for class in [low, high] { + let context = format!("{class} {tiers:?}"); + assert!(!DenseToEager::Off.routes(class, tiers), "{context}"); + assert!(DenseToEager::All.routes(class, tiers), "{context}"); + } + } + } + + #[test] + fn test_lazy_far_window_knob() { + let parse = |window: Option<&str>, inflight: Option<&str>| { + LayeredLazyConfig::from_lookup(|key| match key { + LAZY_FULL_ENV => Some("1".to_string()), + LAZY_FAR_WINDOW_ENV => window.map(String::from), + LAZY_FAR_INFLIGHT_ENV => inflight.map(String::from), + _ => None, + }) + }; + let defaults = LayeredLazyConfig::default(); + assert_eq!((defaults.far_window, defaults.far_inflight), (64, 64)); + let config = parse(Some(" 0 "), Some("1")).unwrap(); + assert_eq!((config.far_window, config.far_inflight), (0, 1)); + let most = Semaphore::MAX_PERMITS.to_string(); + assert_eq!( + parse(None, Some(most.as_str())).unwrap().far_inflight, + Semaphore::MAX_PERMITS + ); + // A pool cannot hold more permits, and a pool without any would + // never issue a gather beyond the ordinary window. + let too_many = (Semaphore::MAX_PERMITS + 1).to_string(); + for (name, window, inflight) in [ + (LAZY_FAR_WINDOW_ENV, Some("wide"), None), + (LAZY_FAR_WINDOW_ENV, Some("-1"), None), + (LAZY_FAR_INFLIGHT_ENV, None, Some("0")), + (LAZY_FAR_INFLIGHT_ENV, None, Some(too_many.as_str())), + (LAZY_FAR_INFLIGHT_ENV, None, Some("")), + ] { + let error = parse(window, inflight).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + let value = window.or(inflight).unwrap(); + assert!(message.contains(name), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + + // The far window applies only where every origin read is a round + // trip, and only when wider than the ordinary window. + let (low, high) = (OriginLatencyClass::Low, OriginLatencyClass::High); + assert_eq!(defaults.active_far_window(low), None); + assert_eq!(defaults.active_far_window(high), Some(64)); + for (window, far_window, expected) in [ + (16, 17, Some(17)), + (16, 16, None), + (16, 8, None), + (0, usize::MAX, Some(usize::MAX)), + ] { + let config = LayeredLazyConfig { + window, + far_window, + ..defaults + }; + let context = format!("window={window} far_window={far_window}"); + assert_eq!(config.active_far_window(high), expected, "{context}"); + assert_eq!(config.active_far_window(low), None, "{context}"); + } + } + + #[tokio::test] + async fn test_lazy_far_permits_count_permits_in_flight() { + let error = LazyFarPermits::try_new(Semaphore::MAX_PERMITS + 1).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let permits = LazyFarPermits::try_new(2).unwrap(); + let shared = permits.clone(); + assert_eq!((permits.size(), permits.in_flight()), (2, 0)); + let first = permits.try_acquire().unwrap(); + let second = shared.try_acquire().unwrap(); + // Clones share one pool. + assert!(permits.try_acquire().is_none()); + assert_eq!(shared.in_flight(), 2); + // A wait abandoned before a permit frees up takes none. + assert!(permits.acquire().now_or_never().is_none()); + drop(first); + assert_eq!(permits.in_flight(), 1); + let third = permits.acquire().await.unwrap(); + assert_eq!(shared.in_flight(), 2); + drop((second, third)); + assert_eq!(permits.in_flight(), 0); + } + + /// Every open of an index file draws its far gathers from one pool per + /// size, which lives while an open or a taken permit holds it. The same + /// path in another bucket is another file. + #[test] + fn test_lazy_far_permits_shared_per_index_file() { + let path = "t.lance/_indices/far-permits-test/auxiliary.idx"; + let file = IndexFileKey::new("far-permits-test", "s3$bucket", path); + let pool = LazyFarPermits::for_file(&file, 2).unwrap(); + let held = pool.try_acquire().unwrap(); + let shared = LazyFarPermits::for_file(&file, 2).unwrap(); + assert_eq!((shared.size(), shared.in_flight()), (2, 1)); + assert_eq!(LazyFarPermits::for_file(&file, 3).unwrap().in_flight(), 0); + let other = IndexFileKey::new("far-permits-test", "s3$other", path); + assert_eq!(LazyFarPermits::for_file(&other, 2).unwrap().in_flight(), 0); + // The taken permit keeps the pool once no open holds it. + drop((pool, shared)); + assert_eq!(LazyFarPermits::for_file(&file, 2).unwrap().in_flight(), 1); + drop(held); + assert_eq!(LazyFarPermits::for_file(&file, 2).unwrap().in_flight(), 0); + let error = LazyFarPermits::for_file(&file, Semaphore::MAX_PERMITS + 1).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + } + #[test] fn test_sequential_plane_loads_flag() { // The accessor reports the process environment's value. @@ -2641,6 +3971,234 @@ mod tests { ); } + #[test] + fn test_origin_latency_knob() { + // The accessor reports the process environment's value. + let process = std::env::var(ORIGIN_LATENCY_ENV).ok(); + assert_eq!( + origin_latency_setting().ok(), + origin_latency_from(process.as_deref()).ok() + ); + assert_eq!(origin_latency_from(None).unwrap(), None); + assert_eq!(origin_latency_from(Some("auto")).unwrap(), None); + assert_eq!( + origin_latency_from(Some("low")).unwrap(), + Some(OriginLatencyClass::Low) + ); + assert_eq!( + origin_latency_from(Some(" high ")).unwrap(), + Some(OriginLatencyClass::High) + ); + for class in [OriginLatencyClass::Low, OriginLatencyClass::High] { + assert_eq!( + origin_latency_from(Some(class.as_str())).unwrap(), + Some(class) + ); + assert_eq!(class.to_string(), class.as_str()); + } + assert_eq!(OriginLatencyClass::default(), OriginLatencyClass::Low); + for value in ["cloud", "HIGH", "1", ""] { + let error = origin_latency_from(Some(value)).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(ORIGIN_LATENCY_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + } + + /// An explicit setting wins over a session's hint, and the hint over the + /// class of the store, low for local files and memory; without a hint the + /// class resolves as [`OriginLatencyClass::resolve`] resolves it. Lance's + /// IVF tests cover a cloud store, which needs a URL to build. + #[test] + fn test_origin_latency_resolve_with_hint() { + let classes = [ + None, + Some(OriginLatencyClass::Low), + Some(OriginLatencyClass::High), + ]; + for store in [ObjectStore::local(), ObjectStore::memory()] { + let scheme = store.scheme().to_string(); + let auto = OriginLatencyClass::of_store(&store); + assert_eq!(auto, OriginLatencyClass::Low, "{scheme}"); + for setting in classes { + assert_eq!( + OriginLatencyClass::resolve_with_hint(setting, None, &store), + OriginLatencyClass::resolve(setting, &store), + "{scheme} setting={setting:?}" + ); + for hint in classes { + let expected = match (setting, hint) { + (Some(class), _) | (None, Some(class)) => class, + (None, None) => auto, + }; + assert_eq!( + OriginLatencyClass::resolve_with_hint(setting, hint, &store), + expected, + "{scheme} setting={setting:?} hint={hint:?}" + ); + } + } + } + } + + #[test] + fn test_sign_bounds_knob() { + // The accessor reports the process environment's value. + let process = std::env::var(SIGN_BOUNDS_ENV).ok(); + assert_eq!( + sign_bounds_setting().ok(), + sign_bounds_from(process.as_deref()).ok() + ); + assert_eq!(sign_bounds_from(None).unwrap(), SignBounds::Lazy); + assert_eq!( + sign_bounds_from(Some(" eager ")).unwrap(), + SignBounds::Eager + ); + for sign_bounds in [SignBounds::Lazy, SignBounds::Eager] { + assert_eq!( + sign_bounds_from(Some(sign_bounds.as_str())).unwrap(), + sign_bounds + ); + assert_eq!(sign_bounds.to_string(), sign_bounds.as_str()); + } + assert_eq!(SignBounds::default(), SignBounds::Lazy); + for value in ["auto", "EAGER", "1", ""] { + let error = sign_bounds_from(Some(value)).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(SIGN_BOUNDS_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + } + + #[test] + fn test_resident_columns_knob() { + // The accessor reports the process environment's value. + let process = std::env::var(RESIDENT_COLUMNS_ENV).ok(); + assert_eq!( + resident_columns_setting().ok(), + resident_columns_from(process.as_deref()).ok() + ); + assert_eq!( + resident_columns_from(None).unwrap(), + ResidentColumnsSetting::Auto + ); + assert_eq!( + resident_columns_from(Some(" on ")).unwrap(), + ResidentColumnsSetting::On + ); + let settings = [ + ResidentColumnsSetting::Auto, + ResidentColumnsSetting::On, + ResidentColumnsSetting::Off, + ]; + for setting in settings { + assert_eq!( + resident_columns_from(Some(setting.as_str())).unwrap(), + setting + ); + assert_eq!(setting.to_string(), setting.as_str()); + } + assert_eq!( + ResidentColumnsSetting::default(), + ResidentColumnsSetting::Auto + ); + // `auto` keeps the columns resident only when every column read is + // a round trip to a slow origin. + let classes = [OriginLatencyClass::Low, OriginLatencyClass::High]; + for (setting, expected) in [ + (ResidentColumnsSetting::Auto, [false, true]), + (ResidentColumnsSetting::On, [true, true]), + (ResidentColumnsSetting::Off, [false, false]), + ] { + assert_eq!( + classes.map(|class| setting.resolve(class)), + expected, + "{setting}" + ); + } + for value in ["1", "ON", "true", ""] { + let error = resident_columns_from(Some(value)).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(RESIDENT_COLUMNS_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + } + + /// `off` never keeps a store; `on` keeps any store with columns, however + /// large; `auto` keeps one whose charge fits the entry limit of the + /// cache's largest admissible entry, whatever its values' + /// bytes, and any size on a cache without that limit. + #[test] + fn admits_follows_setting_charge_and_entry_limit() { + const MAX_ENTRY: u64 = (1 << 20) + 1; + let cap = MAX_ENTRY; + let size = |bytes, charge| ResidentStoreSize { bytes, charge }; + let cases = [ + (size(0, 0), Some(MAX_ENTRY)), + (size(1, 100), Some(MAX_ENTRY)), + (size(cap - 100, cap), Some(MAX_ENTRY)), + // Values at the cap, charged past it. + (size(cap, cap + 100), Some(MAX_ENTRY)), + (size(cap - 100, cap + 1), Some(MAX_ENTRY)), + (size(MAX_ENTRY + 1, MAX_ENTRY + 100), Some(MAX_ENTRY)), + (size(1, 100), Some(0)), + (size(u64::MAX, u64::MAX), None), + (size(0, 0), None), + ]; + for (setting, expected) in [ + ( + ResidentColumnsSetting::Auto, + [false, true, true, false, false, false, false, true, false], + ), + ( + ResidentColumnsSetting::On, + [false, true, true, true, true, true, true, true, false], + ), + (ResidentColumnsSetting::Off, [false; 9]), + ] { + let admitted = cases.map(|(store, max)| setting.admits(store, max)); + assert_eq!(admitted, expected, "{setting}"); + } + assert!(resident_store_fits(cap, Some(MAX_ENTRY))); + assert!(!resident_store_fits(cap + 1, Some(MAX_ENTRY))); + assert!(resident_store_fits(u64::MAX, None)); + } + + #[test] + fn resident_lifetime_parses_and_rejects_unknown() { + let process = std::env::var(RESIDENT_LIFETIME_ENV).ok(); + assert_eq!( + resident_lifetime_setting().ok(), + resident_lifetime_from(process.as_deref()).ok() + ); + assert_eq!( + resident_lifetime_from(None).unwrap(), + ResidentLifetime::Index + ); + assert_eq!(ResidentLifetime::default(), ResidentLifetime::Index); + for lifetime in [ResidentLifetime::Index, ResidentLifetime::Process] { + assert_eq!( + resident_lifetime_from(Some(lifetime.as_str())).unwrap(), + lifetime + ); + assert_eq!(lifetime.to_string(), lifetime.as_str()); + } + assert_eq!( + resident_lifetime_from(Some(" process ")).unwrap(), + ResidentLifetime::Process + ); + for value in ["query", "PROCESS", "1", ""] { + let error = resident_lifetime_from(Some(value)).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + let message = error.to_string(); + assert!(message.contains(RESIDENT_LIFETIME_ENV), "{error}"); + assert!(message.contains(&format!("{value:?}")), "{error}"); + } + } + #[test] fn test_query_scratch_pool_initializes_buffer_capacity() { let pool = diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 8f35da55c8d..e871f0ef71d 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -44,6 +44,7 @@ use lance_index::vector::hnsw::HNSW; use lance_index::vector::pq::ProductQuantizer; use lance_index::vector::quantizer::Quantization; use lance_index::vector::sq::ScalarQuantizer; +use lance_index::vector::storage::ResidentColumns; use lance_index::vector::v3::subindex::IvfSubIndex; use lance_index::{ FtsPrewarmDiagnostics, FtsPrewarmOptions, FtsPrewarmResult, FtsPrewarmSegmentStatus, @@ -75,7 +76,7 @@ use vector::details::{ vector_details_as_json, }; pub(crate) use vector::details::{vector_index_details, vector_index_details_default}; -use vector::ivf::v2::{IVFIndex, IvfStateEntryBox}; +use vector::ivf::v2::{IVFIndex, IvfOpenContext, IvfStateEntryBox, aux_file_key}; use vector::utils::get_vector_type; mod api; @@ -4169,6 +4170,22 @@ impl Dataset { ) -> Result> { let uuid = &index_meta.uuid; let object_store = self.object_store_for_index(index_meta).await?; + let index_dir = self.indice_files_dir(index_meta)?; + // Hold the shared allocation before cache accesses can evict it. + // The handle is dropped if this index does not use resident columns. + let file_cache = self.index_cache.for_index(uuid, None); + let resident_columns_handle = Some(ResidentColumns::shared(&aux_file_key( + &object_store, + &index_dir, + uuid, + ))); + // What the session declares about the index, for whichever IVF open + // or reconstruction below runs. + let open_context = IvfOpenContext { + origin_latency_hint: self.session.index_origin_latency(), + file_cache: Some(file_cache), + resident_columns_handle, + }; let resolved = frag_reuse::open_row_id_remapping_with_plan(self, index_meta, staged, purpose, metrics) .await?; @@ -4193,6 +4210,7 @@ impl Dataset { self.metadata_cache.as_ref(), partition_cache, remapping, + open_context, ) .await; } @@ -4214,7 +4232,6 @@ impl Dataset { } else { self.open_frag_reuse_index(metrics).await? }; - let index_dir = self.indice_files_dir(index_meta)?; let index_file = index_dir .clone() .join(uuid.to_string()) @@ -4335,6 +4352,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4348,6 +4366,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4367,6 +4386,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4381,6 +4401,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4395,6 +4416,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4410,6 +4432,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4423,6 +4446,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4438,6 +4462,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) @@ -4452,6 +4477,7 @@ impl Dataset { self.metadata_cache.as_ref(), index_cache, file_sizes, + open_context, ) .await?; Ok(wrap_ivf(ivf)) diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index fc977bf000d..6e71fc6711e 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -50,6 +50,7 @@ use lance_index::prefilter::NoFilter; use lance_index::vector::VectorIndexCacheEntry; use lance_index::vector::bq::builder::RabitQuantizer; use lance_index::vector::bq::ex_dot::{blocked_ex_code_bytes, padded_query_len}; +use lance_index::vector::bq::layered::SignBounds; use lance_index::vector::bq::layered_stats; use lance_index::vector::bq::rabit_ex_bits; use lance_index::vector::bq::storage::{RabitQueryEstimator, SEGMENT_NUM_CODES}; @@ -63,8 +64,10 @@ use lance_index::vector::quantizer::{ }; use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::storage::{ - LayeredLazyConfig, PlaneAccessTracker, QueryResidual, QueryScratch, QueryScratchCapacity, - QueryScratchPool, RabitRawQueryContext, VectorStore, + IndexFileKey, LayeredLazyConfig, OriginLatencyClass, PlaneAccessTracker, QueryResidual, + QueryScratch, QueryScratchCapacity, QueryScratchPool, RabitRawQueryContext, ResidentColumns, + ResidentColumnsSetting, VectorStore, origin_latency_setting, resident_columns_setting, + resident_lifetime_setting, resident_store_fits, sign_bounds_setting, }; use lance_index::vector::v3::subindex::SubIndexType; use lance_index::{ @@ -131,6 +134,12 @@ pub(crate) struct IvfIndexState { /// Runtime-only cache, intentionally excluded from the CacheCodec wire format. pub(crate) rq_search_cache: RabitSearchCacheCell, /// Runtime-only counters survive reader reconstruction between queries. + /// Its far gather pools are the ones every open of the index file + /// shares, bound by reconstructions (see `IndexFileKey`), so a state + /// read back from a persistent tier with a fresh tracker shares them too. + /// The state holds no resident store: every reconstruction binds its + /// storage to the store the index cache charges (see + /// [`IvfOpenContext::file_cache`]), so a cached state pins nothing. pub(crate) plane_access: PlaneAccessTracker, } @@ -659,6 +668,7 @@ pub(crate) trait IvfStateEntry: DeepSizeOf + Send + Sync + 'static { file_metadata_cache: &'a LanceCache, index_cache: LanceCache, frag_reuse_index: Option, + context: IvfOpenContext, ) -> BoxFuture<'a, Result>>; } @@ -836,6 +846,7 @@ impl IvfStateEntry for IvfIndexState { file_metadata_cache: &'a LanceCache, index_cache: LanceCache, frag_reuse_index: Option, + context: IvfOpenContext, ) -> BoxFuture<'a, Result>> { Box::pin(async move { match self.sub_index_type { @@ -846,6 +857,7 @@ impl IvfStateEntry for IvfIndexState { file_metadata_cache, index_cache, frag_reuse_index, + context, ) .await } @@ -856,6 +868,7 @@ impl IvfStateEntry for IvfIndexState { file_metadata_cache, index_cache, frag_reuse_index, + context, ) .await } @@ -1050,6 +1063,57 @@ impl CacheKey for IVFPartit } } +/// What the session opening an IVF index declares about it, passed to +/// [`IVFIndex::try_new`] and to the reconstruction from a cached +/// [`IvfIndexState`]. The default declares nothing and charges no cache. +#[derive(Debug, Clone, Default)] +pub(crate) struct IvfOpenContext { + /// The origin latency class the session declares for the IVF_RQ indexes + /// it opens (`Session::index_origin_latency`), `None` when it declares + /// none; see [`IVFIndex::origin_latency_at_open`]. + pub(crate) origin_latency_hint: Option, + /// The index's namespace of the index cache without a fragment reuse + /// segment, where an IVF_RQ index with resident small columns charges + /// their store (see [`ResidentColumns::in_index_cache`]), so + /// that a new fragment reuse index keeps the store. `None` keeps the + /// store charged nowhere, shared with the file's other live indexes. + pub(crate) file_cache: Option, + /// A strong reference acquired before opening the index can evict the + /// cached store. Dropped if the index does not use resident columns. + pub(crate) resident_columns_handle: Option, +} + +/// The key of the auxiliary (storage) file of index `uuid` in `index_dir` +/// on `object_store`, which IVF_RQ runtime handles, the resident store and +/// the lazy scan's far gather pools, are shared by. +pub(crate) fn aux_file_key( + object_store: &ObjectStore, + index_dir: &Path, + uuid: &Uuid, +) -> IndexFileKey { + let uuid = uuid.to_string(); + let aux_path = index_dir + .clone() + .join(uuid.as_str()) + .join(INDEX_AUXILIARY_FILE_NAME); + IndexFileKey::new(&uuid, &object_store.store_prefix, aux_path.as_ref()) +} + +/// Warn, once per process, that `LANCE_RQ_RESIDENT_COLUMNS=on` keeps a +/// resident store that the index cache charges `charge` bytes for, more than +/// its largest admissible entry, `max_entry_bytes`. Index handles keep this +/// allocation alive independently of cache capacity. +fn warn_resident_store_oversize(charge: u64, max_entry_bytes: Option) { + static WARNED: std::sync::Once = std::sync::Once::new(); + WARNED.call_once(|| { + log::warn!( + "LANCE_RQ_RESIDENT_COLUMNS=on keeps a resident store the index cache charges \ + {charge} bytes for, exceeding its {max_entry_bytes:?} byte entry limit; \ + live index handles retain this memory outside the cache" + ); + }); +} + /// IVF Index. #[derive(Debug)] pub struct IVFIndex { @@ -1092,6 +1156,13 @@ pub struct IVFIndex { /// Lazy full-precision scan settings of a layered index; see /// [`LayeredLazyConfig`]. Locked only to copy it once per query. layered_lazy: Mutex, + /// Latency class of reads from the index's files, resolved when the + /// index opens; the storage carries the same class. + origin_latency: OriginLatencyClass, + /// Block size of the object store the index's files are read from, the + /// gap within which reads merge requested ranges unless a lazy sparse + /// origin read sets its own; see [`Self::lazy_origin_gap_bytes`]. + origin_block_size: u64, /// Staging parallelism of the lazy scan when non-zero; see /// `set_lazy_prepare_parallelism_for_test`. #[cfg(test)] @@ -1550,6 +1621,7 @@ impl IVFIndex { } /// Create a new IVF index. + #[allow(clippy::too_many_arguments)] pub(crate) async fn try_new( object_store: Arc, index_dir: Path, @@ -1558,12 +1630,21 @@ impl IVFIndex { file_metadata_cache: &LanceCache, index_cache: LanceCache, file_sizes: HashMap, + context: IvfOpenContext, ) -> Result { let io_parallelism = object_store.io_parallelism(); + let origin_latency = + Self::origin_latency_at_open(&object_store, context.origin_latency_hint)?; + let origin_block_size = object_store.block_size() as u64; + let uuid_str = uuid.to_string(); + let aux_path = index_dir + .clone() + .join(uuid_str.as_str()) + .join(INDEX_AUXILIARY_FILE_NAME); + let index_file = aux_file_key(&object_store, &index_dir, &uuid); let scheduler_config = SchedulerConfig::max_bandwidth(&object_store); let scheduler = ScanScheduler::new(object_store, scheduler_config); - let uuid_str = uuid.to_string(); let uri = index_dir .clone() .join(uuid_str.as_str()) @@ -1612,15 +1693,7 @@ impl IVFIndex { .map(|&size| CachedFileSize::new(size)) .unwrap_or_else(CachedFileSize::unknown); let storage_reader = FileReader::try_open( - scheduler - .open_file( - &index_dir - .clone() - .join(uuid_str.as_str()) - .join(INDEX_AUXILIARY_FILE_NAME), - &aux_cached_size, - ) - .await?, + scheduler.open_file(&aux_path, &aux_cached_size).await?, None, Arc::::default(), file_metadata_cache, @@ -1643,6 +1716,15 @@ impl IVFIndex { IvfQuantizationStorage::try_new_with_remapper(storage_reader, None).await? } }; + let resident_columns = + Self::resident_columns_at_open(origin_latency, &storage, context.file_cache.as_ref())?; + let resident_store = + Self::resident_store_at_open(resident_columns, &index_file, context).await?; + let storage = storage + .with_origin_latency(origin_latency) + .with_resident_columns(resident_store) + .with_resident_columns_enabled(resident_columns) + .with_index_file(index_file); // Cache file metadata so reconstructions from IvfIndexState can skip // footer reads. @@ -1650,10 +1732,6 @@ impl IVFIndex { .with_key_prefix(uri.as_ref()) .insert_with_key(&FileMetadataCacheKey, index_reader.metadata().clone()) .await; - let aux_path = index_dir - .clone() - .join(uuid_str.as_str()) - .join(INDEX_AUXILIARY_FILE_NAME); file_metadata_cache .with_key_prefix(aux_path.as_ref()) .insert_with_key(&FileMetadataCacheKey, storage.reader().metadata().clone()) @@ -1673,7 +1751,8 @@ impl IVFIndex { let read_projection = Self::read_projection(&index_reader)?; let layered_rq = Q::quantization_type() == QuantizationType::Rabit && storage.is_layered_rq(); - let layered_lazy = Mutex::new(Self::layered_lazy_config_at_open(layered_rq)?); + let layered_lazy = Self::layered_lazy_config_at_open(layered_rq, origin_latency)?; + let storage = storage.with_sign_bounds(Self::sign_bounds_at_open(layered_rq)?); Ok(Self { uri: to_local_path(&uri), index_path: uri.as_ref().to_string(), @@ -1693,14 +1772,17 @@ impl IVFIndex { open_io_stats, prepared_partitions: Arc::default(), layered_rq, - layered_lazy, + layered_lazy: Mutex::new(layered_lazy), + origin_latency, + origin_block_size, #[cfg(test)] lazy_prepare_parallelism_for_test: AtomicUsize::new(0), _marker: PhantomData, }) } - /// Reconstruct an IVFIndex from pre-parsed state without any I/O. + /// Reconstruct an IVFIndex from pre-parsed state without any I/O, over + /// a storage already bound to its resident store. #[allow(clippy::too_many_arguments)] pub(crate) fn from_cached_state( uri: String, @@ -1714,14 +1796,18 @@ impl IVFIndex { index_cache: LanceCache, io_parallelism: usize, rq_search_cache: Option>, + origin_latency: OriginLatencyClass, + origin_block_size: u64, ) -> Result { + let storage = storage.with_origin_latency(origin_latency); let scratch_pool = Arc::new(Self::query_scratch_pool(&ivf, &storage)); let use_query_residual = Self::use_query_residual(&storage, distance_type); let use_residual_scratch = Self::use_residual_scratch(&ivf, use_query_residual); let read_projection = Self::read_projection(&reader)?; let layered_rq = Q::quantization_type() == QuantizationType::Rabit && storage.is_layered_rq(); - let layered_lazy = Mutex::new(Self::layered_lazy_config_at_open(layered_rq)?); + let layered_lazy = Self::layered_lazy_config_at_open(layered_rq, origin_latency)?; + let storage = storage.with_sign_bounds(Self::sign_bounds_at_open(layered_rq)?); Ok(Self { uri, index_path, @@ -1744,13 +1830,195 @@ impl IVFIndex { open_io_stats: ScanStats::default(), prepared_partitions: Arc::default(), layered_rq, - layered_lazy, + layered_lazy: Mutex::new(layered_lazy), + origin_latency, + origin_block_size, #[cfg(test)] lazy_prepare_parallelism_for_test: AtomicUsize::new(0), _marker: PhantomData, }) } + /// The origin latency class of an index opened on `object_store` by a + /// session that declares `hint`. Only IVF_RQ indexes read (and validate) + /// `LANCE_RQ_ORIGIN_LATENCY` and take the hint, the setting first (see + /// [`OriginLatencyClass::resolve_with_hint`]); any other index takes the + /// class of its store, since no reader policy of theirs depends on it. + fn origin_latency_at_open( + object_store: &ObjectStore, + hint: Option, + ) -> Result { + if Q::quantization_type() != QuantizationType::Rabit { + return Ok(OriginLatencyClass::of_store(object_store)); + } + Ok(OriginLatencyClass::resolve_with_hint( + origin_latency_setting()?, + hint, + object_store, + )) + } + + /// Latency class of reads from the index's files, resolved when the index + /// opened; see [`OriginLatencyClass`]. + pub fn origin_latency(&self) -> OriginLatencyClass { + self.origin_latency + } + + /// The gap in bytes within which the lazy scan's sparse reads from the + /// index's files merge row runs into one request: the index's + /// `LANCE_RQ_LAZY_ORIGIN_GAP_BYTES` resolved for its origin latency class + /// and object store; see [`lance_index::vector::storage::LazyOriginGap`]. + pub fn lazy_origin_gap_bytes(&self) -> u64 { + let origin_gap = self + .layered_lazy + .lock() + .unwrap_or_else(|err| err.into_inner()) + .origin_gap; + origin_gap.resolve(self.origin_latency, self.origin_block_size) + } + + /// The index's lazy scan settings: `LANCE_RQ_LAZY_*` read when it + /// opened, with the background promotion policy resolved for its origin + /// latency class; see [`LayeredLazyConfig`]. + pub fn layered_lazy_config(&self) -> LayeredLazyConfig { + *self + .layered_lazy + .lock() + .unwrap_or_else(|err| err.into_inner()) + } + + /// Replace the origin latency class resolved at open, and what the + /// index resolves from it when it opens (residency and the lazy scan + /// settings), so tests can open both classes over one store in one + /// process. + #[cfg(test)] + pub(crate) fn with_origin_latency_for_test( + mut self, + origin_latency: OriginLatencyClass, + ) -> Result { + let resident_columns = Self::resident_columns_at_open(origin_latency, &self.storage, None)?; + let layered_lazy = Self::layered_lazy_config_at_open(self.layered_rq, origin_latency)?; + self.origin_latency = origin_latency; + self.storage = self + .storage + .with_origin_latency(origin_latency) + .with_resident_columns_enabled(resident_columns); + self.layered_lazy = Mutex::new(layered_lazy); + Ok(self) + } + + /// Whether an index opened with origin latency `class` keeps the small + /// columns of `storage`'s file resident in `file_cache`: when the setting + /// resolves for the class and admits the store for the cache's largest + /// entry (see [`ResidentColumnsSetting::admits`]), which `auto` sizes by + /// what the cache would charge for the store against its entry limit. + /// Only IVF_RQ indexes read (and validate) `LANCE_RQ_RESIDENT_COLUMNS` + /// and `LANCE_RQ_RESIDENT_LIFETIME`: the resident columns are told apart + /// from the RaBitQ code columns, so other storages keep none. + fn resident_columns_at_open( + class: OriginLatencyClass, + storage: &IvfQuantizationStorage, + file_cache: Option<&LanceCache>, + ) -> Result { + if Q::quantization_type() != QuantizationType::Rabit { + return Ok(false); + } + resident_lifetime_setting()?; + let setting = resident_columns_setting()?; + if !setting.resolve(class) { + return Ok(false); + } + let store = storage.resident_store_size()?; + let max_entry_bytes = file_cache.and_then(LanceCache::max_entry_bytes); + if store.bytes > 0 && !resident_store_fits(store.charge, max_entry_bytes) { + layered_stats::counters().resident_columns_oversize.incr(); + if setting == ResidentColumnsSetting::On { + warn_resident_store_oversize(store.charge, max_entry_bytes); + } + } + Ok(setting.admits(store, max_entry_bytes)) + } + + /// The resident store handle of an index opening on index file `file`: + /// when the index keeps its small columns `resident`, the store every + /// live index of the file shares and offered to the ordinary file cache. + /// The opener's strong reference bridges eviction during index opening. + /// Otherwise use an independent store that loads only if residency is + /// enabled later. + async fn resident_store_at_open( + resident: bool, + file: &IndexFileKey, + context: IvfOpenContext, + ) -> Result { + if !resident { + return Ok(ResidentColumns::default()); + } + Ok(match &context.file_cache { + Some(cache) => { + ResidentColumns::in_index_cache( + cache, + file, + context.resident_columns_handle, + resident_lifetime_setting()?, + ) + .await + } + None => ResidentColumns::shared(file), + }) + } + + /// Whether reads keep the index's small columns resident and fetch only + /// its code and bounds columns from the file, resolved when the index + /// opened; see `LANCE_RQ_RESIDENT_COLUMNS`. + pub fn resident_columns_enabled(&self) -> bool { + self.storage.resident_columns_enabled() + } + + /// Bytes the index's resident columns take, whether or not it keeps + /// them; see [`lance_index::vector::storage::resident_columns_bytes`]. + pub fn resident_columns_bytes(&self) -> u64 { + self.storage.resident_columns_bytes() + } + + /// Replace whether reads keep the small columns resident, so tests can + /// open both in one process. + #[cfg(test)] + pub(crate) fn with_resident_columns_for_test(mut self, enabled: bool) -> Self { + self.storage = self.storage.with_resident_columns_enabled(enabled); + self + } + + /// Where a newly opened index's cache keeps the bounds columns. Only + /// layered indexes read (and validate) `LANCE_RQ_SIGN_BOUNDS`. + fn sign_bounds_at_open(layered_rq: bool) -> Result { + if layered_rq { + sign_bounds_setting() + } else { + Ok(SignBounds::default()) + } + } + + /// Where the index cache keeps a layered index's bounds columns, + /// resolved when the index opened; see [`SignBounds`]. + pub fn sign_bounds(&self) -> SignBounds { + self.storage.sign_bounds() + } + + /// Whether the index is a layered IVF_RQ index, the only kind that reads + /// `LANCE_RQ_SIGN_BOUNDS` and `LANCE_RQ_LAZY_*`; any other index keeps + /// their defaults. + pub fn is_layered_rq(&self) -> bool { + self.layered_rq + } + + /// Replace where the index cache keeps the bounds columns, so tests can + /// open both placements in one process. + #[cfg(test)] + pub(crate) fn with_sign_bounds_for_test(mut self, sign_bounds: SignBounds) -> Self { + self.storage = self.storage.with_sign_bounds(sign_bounds); + self + } + /// Prepared-partition accounting for this index instance; see /// [`PreparedPartitionTracker`]. #[cfg(test)] @@ -2132,17 +2400,14 @@ impl IVFIndex { query: &Query, metrics: &dyn MetricsCollector, ) -> Result>> { - use lance_index::vector::bq::layered::{PlaneKey, RQPrecision}; + use lance_index::vector::bq::layered::RQPrecision; let mut precision = self.initial_precision(query)?; if query.rq_cascade_factor.is_some() && precision == RQPrecision::High && self.storage.supports_candidate_reads() && self .index_cache - .get_resident_with_key(&PlaneKey { - partition: partition_id, - plane: 1, - }) + .get_resident_with_key(&self.storage.plane_key(partition_id, 1)) .await .is_none() { @@ -3551,8 +3816,12 @@ async fn reconstruct_typed( file_metadata_cache: &LanceCache, index_cache: LanceCache, frag_reuse_index: Option, + context: IvfOpenContext, ) -> Result> { let io_parallelism = object_store.io_parallelism(); + let origin_latency = + IVFIndex::::origin_latency_at_open(&object_store, context.origin_latency_hint)?; + let origin_block_size = object_store.block_size() as u64; let index_path = Path::parse(&state.index_file_path) .map_err(|e| Error::io(format!("invalid index path: {e}")))?; @@ -3562,6 +3831,7 @@ async fn reconstruct_typed( parts.pop(); let dir: Path = parts.into_iter().collect(); let aux_path = dir.clone().join(INDEX_AUXILIARY_FILE_NAME); + let index_file = IndexFileKey::new(&state.uuid, &object_store.store_prefix, aux_path.as_ref()); // Readers carry a scheduler bound to an object store, so they cannot be // shared across dataset opens. Reuse only portable file metadata and bind @@ -3600,6 +3870,20 @@ async fn reconstruct_typed( } Some(ResolvedRemapping::V1Identity) | None => make_storage(None), }; + // Bind the store as an open does: the state holds none, so a state read + // back from a persistent tier binds to the one the index cache holds or + // the file's live indexes share, and any state pins nothing. + let resident = IVFIndex::::resident_columns_at_open( + origin_latency, + &storage, + context.file_cache.as_ref(), + )?; + let resident_columns = + IVFIndex::::resident_store_at_open(resident, &index_file, context).await?; + let storage = storage + .with_resident_columns(resident_columns) + .with_resident_columns_enabled(resident) + .with_index_file(index_file); let rq_search_cache = IVFIndex::::rq_search_cache_from_state(state, &storage)?; let parsed_uuid = Uuid::parse_str(&state.uuid) @@ -3616,6 +3900,8 @@ async fn reconstruct_typed( index_cache, io_parallelism, rq_search_cache, + origin_latency, + origin_block_size, )?; Ok(Arc::new(index)) } @@ -3658,7 +3944,8 @@ mod tests { use crate::index::DatasetIndexExt; use crate::index::DatasetIndexInternalExt; use crate::index::vector::ivf::v2::{ - IVFPartitionKey, IvfFlatIndex, IvfHnswSqIndex, IvfPq, IvfStateEntryBox, PartitionEntry, + IVFPartitionKey, IvfFlatIndex, IvfHnswSqIndex, IvfOpenContext, IvfPq, IvfStateEntryBox, + PartitionEntry, }; use crate::utils::test::copy_test_data_to_tmp; use crate::{ @@ -7203,7 +7490,7 @@ mod tests { #[tokio::test] async fn test_layered_plane_promotion_survives_query_reconstruction() { - use lance_index::vector::bq::layered::PlaneKey; + use lance_index::vector::bq::layered::{PlaneKey, SignBounds}; let dir = TempStrDir::default(); let (mut dataset, vectors) = @@ -7244,9 +7531,13 @@ mod tests { assert_eq!(result.num_rows(), num_rows); for partition in 0..4 { for plane in [1, 2] { - let resident = cache - .get_resident_with_key(&PlaneKey { partition, plane }) - .await; + // Ex plane keys do not depend on the bounds placement. + let key = PlaneKey { + partition, + plane, + sign_bounds: SignBounds::default(), + }; + let resident = cache.get_resident_with_key(&key).await; assert_eq!( resident.is_some(), pass == 2, @@ -7263,7 +7554,9 @@ mod tests { use super::*; use bytes::Bytes; - use lance_core::cache::{CacheCodec, CacheEntry, InternalCacheKey, QuickCacheBackend}; + use lance_core::cache::{ + CacheCodec, CacheEntry, CacheTier, InternalCacheKey, QuickCacheBackend, + }; use lance_index::vector::ApproxMode; use lance_index::vector::bq::layered::{PlaneBatch, RQPrecision}; use lance_io::assert_io_eq; @@ -7335,6 +7628,17 @@ mod tests { self.ram.get_resident(key).await } + /// RAM as `peek_resident` reports it, then the serialized tier. + async fn peek_tier(&self, key: &InternalCacheKey) -> CacheTier { + if self.peek_resident(key).await { + CacheTier::Resident + } else if self.disk.lock().unwrap().contains_key(key) { + CacheTier::Local + } else { + CacheTier::Absent + } + } + fn plane_admission_gated(&self) -> bool { self.gated } @@ -7797,21 +8101,38 @@ mod tests { mod layered_lazy { use super::*; use std::sync::LazyLock; + use std::sync::atomic::AtomicU64; use std::time::Duration; use arrow::compute::concat_batches; use bytes::Bytes; - use lance_core::cache::{CacheCodec, CacheEntry, InternalCacheKey, QuickCacheBackend}; + use lance_core::cache::{ + CacheCodec, CacheEntry, CacheTier, InternalCacheKey, QuickCacheBackend, + }; + use lance_encoding::decoder::FilterExpression; + use lance_file::version::LanceFileVersion; + use lance_file::writer::FileWriterOptions; use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; - use lance_index::vector::bq::layered::{PlaneBatch, PlaneKey, RQPrecision}; + use lance_index::vector::bq::layered::{ + FULL_BOUNDS_COLUMN, PlaneBatch, RQPrecision, SIGN_BOUNDS_PLANE, SignBounds, + plane_columns, + }; use lance_index::vector::bq::layered_stats::{self, LayeredLazyStats, RANK_BUCKETS}; use lance_index::vector::bq::storage::{ RabitPruneStatsSnapshot, rabit_prune_stats_snapshot, }; - use lance_index::vector::storage::{DenseGatherMode, LayeredLazyConfig, LazyPromotion}; + use lance_index::vector::storage::{ + DenseGatherMode, DenseToEager, HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, IndexFileKey, + LayeredLazyConfig, LazyOriginGap, LazyPromotion, OriginLatencyClass, + ResidentColumnsKey, ResidentColumnsSetting, ResidentLifetime, ResidentStoreSize, + origin_latency_setting, resident_columns_setting, resident_lifetime_setting, + resident_store_is_live, sign_bounds_setting, + }; use lance_index::vector::{ApproxMode, PartitionSearchControl, VECTOR_RESULT_SCHEMA}; + use lance_io::ReadBatchParams; use lance_io::assert_io_eq; use lance_io::object_store::ObjectStoreRegistry; + use lance_io::scheduler::IoStats; use rand::seq::SliceRandom; use crate::session::Session; @@ -7862,6 +8183,23 @@ mod tests { /// here, and each yields once after it starts so that loads its /// caller issues together overlap. plane_loads: std::sync::Mutex>>, + /// How long every row read waits before it is served or misses, + /// like a slow tier, so that sparse gathers stay in flight. + row_read_delay: std::sync::Mutex, + /// Keys of the cached index states written through, so that a + /// test can evict the states alone from RAM. + state_keys: std::sync::Mutex>, + /// Keys whose RAM copy `get` treats as evicted until it admits + /// the persistent copy again. + evicted_from_ram: std::sync::Mutex>, + /// Bytes of churn the next admission of an entry evicted with + /// `evict_states_from_ram` runs first, as another query's planes + /// admitted meanwhile would. + churn_on_admission: AtomicUsize, + /// Admissions of evicted entries that ran churn first. + churned_admissions: AtomicUsize, + /// Source of the keys churned entries take. + churn_keys: AtomicU64, } /// The loader a backend's `get_or_insert` receives. @@ -7894,9 +8232,48 @@ mod tests { hide_empty_planes: AtomicBool::new(false), spill: Default::default(), plane_loads: Default::default(), + row_read_delay: Default::default(), + state_keys: Default::default(), + evicted_from_ram: Default::default(), + churn_on_admission: AtomicUsize::new(0), + churned_admissions: AtomicUsize::new(0), + churn_keys: AtomicU64::new(0), + } + } + + /// Evict the cached index states from RAM alone, keeping their + /// persistent copies: the next open reads its state back. + fn evict_states_from_ram(&self) { + let states = self.state_keys.lock().unwrap().clone(); + self.evicted_from_ram.lock().unwrap().extend(states); + } + + /// Run `bytes` of churn before the next admission of an entry + /// evicted with `evict_states_from_ram`. + fn churn_on_next_admission(&self, bytes: usize) { + self.churn_on_admission.store(bytes, Ordering::SeqCst); + } + + /// Admit `bytes` of entries to RAM, each read back so that it is + /// promoted to the hot ring and pushes older entries out, as the + /// planes of another index's queries do. + async fn churn_ram(&self, bytes: usize) { + const CHURN_ENTRY_BYTES: usize = 16 * 1024; + for _ in 0..bytes / CHURN_ENTRY_BYTES { + let id = self.churn_keys.fetch_add(1, Ordering::Relaxed); + let mut key = [0xC5; 16]; + key[..8].copy_from_slice(&id.to_le_bytes()); + let key = InternalCacheKey::from_bytes(key); + let entry = Arc::new(vec![0u8; CHURN_ENTRY_BYTES]); + self.ram.insert(&key, entry, CHURN_ENTRY_BYTES, None).await; + self.ram.get(&key, None).await; } } + fn set_row_read_delay(&self, delay: Duration) { + *self.row_read_delay.lock().unwrap() = delay; + } + fn record_plane_loads(&self) { *self.plane_loads.lock().unwrap() = Some(Vec::new()); } @@ -7996,6 +8373,18 @@ mod tests { Some(entry) } + /// RAM as `peek_resident` reports it, hidden empty planes + /// included, then the serialized tier. + async fn peek_tier(&self, key: &InternalCacheKey) -> CacheTier { + if self.peek_resident(key).await { + CacheTier::Resident + } else if self.disk.lock().unwrap().contains_key(key) { + CacheTier::Local + } else { + CacheTier::Absent + } + } + fn plane_admission_gated(&self) -> bool { self.gated } @@ -8017,6 +8406,10 @@ mod tests { rows: &[u32], _codec: Option, ) -> Option { + let delay = *self.row_read_delay.lock().unwrap(); + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } if !self.serve_rows.load(Ordering::Relaxed) { return None; } @@ -8030,10 +8423,20 @@ mod tests { key: &InternalCacheKey, _codec: Option, ) -> Option { - if let Some(entry) = self.ram.get(key, None).await { + let evicted = self.evicted_from_ram.lock().unwrap().remove(key); + if !evicted && let Some(entry) = self.ram.get(key, None).await { return Some(entry); } let (entry, size) = self.read_disk(key)?; + let churn = if evicted { + self.churn_on_admission.swap(0, Ordering::SeqCst) + } else { + 0 + }; + if churn > 0 { + self.churn_ram(churn).await; + self.churned_admissions.fetch_add(1, Ordering::SeqCst); + } self.ram.insert(key, entry.clone(), size, None).await; Some(entry) } @@ -8045,10 +8448,17 @@ mod tests { size_bytes: usize, codec: Option, ) { + if entry.downcast_ref::().is_some() { + self.state_keys.lock().unwrap().insert(*key); + } self.persist(key, &entry, size_bytes, codec); self.ram.insert(key, entry, size_bytes, None).await; } + fn max_entry_bytes(&self) -> Option { + self.ram.max_entry_bytes() + } + async fn get_or_insert<'a>( &self, key: &InternalCacheKey, @@ -8194,18 +8604,34 @@ mod tests { ] } - fn lazy_test_configs() -> Vec { + /// Lazy settings the parity tests rotate through, each with the origin + /// latency class of the index it runs on. On a low-latency origin the + /// default dense-to-eager routing (`origin`) routes no probe, so some + /// settings route every predicted-dense probe (`all`), and the far + /// window never applies. Promotions are off unless set, so some + /// settings that gather sparse rows promote a plane after one read. + /// The count of settings is prime, so over the parity test's cases + /// the rotation pairs every setting with every `k`, filter, + /// approximation mode and bounds choice. + fn lazy_test_configs() -> Vec<(OriginLatencyClass, LayeredLazyConfig)> { let enabled = LayeredLazyConfig { enabled: true, ..Default::default() }; - vec![ - enabled, + let promote_once = LazyPromotion::Background { reads: 1 }; + let low = [ + // The cost policy's defaults, whose sparse gathers promote. + LayeredLazyConfig { + promote: promote_once, + ..enabled + }, + // Sparse origin reads merge every run of a plane's column page. LayeredLazyConfig { window: 0, dense: DenseGatherMode::Sparse, promote: LazyPromotion::Off, inline_rows: 0, + origin_gap: LazyOriginGap::Bytes(u64::MAX), ..enabled }, // Whole gathers predict every probe dense, which would @@ -8215,42 +8641,87 @@ mod tests { dense: DenseGatherMode::Whole, promote: LazyPromotion::Background { reads: 2 }, inline_rows: usize::MAX, - dense_to_eager: false, + dense_to_eager: DenseToEager::Off, ..enabled }, LayeredLazyConfig { window: 8, max_runs: 0, dense_bytes_fraction: 0.1, + dense_to_eager: DenseToEager::All, ..enabled }, + // Sparse origin reads merge only touching runs. LayeredLazyConfig { window: 64, max_runs: 2, promote: LazyPromotion::Off, inline_rows: 0, + origin_gap: LazyOriginGap::Bytes(0), ..enabled }, + // Every gather is sparse, and promotes. LayeredLazyConfig { max_runs: usize::MAX, dense_bytes_fraction: f64::INFINITY, + dense_to_eager: DenseToEager::All, + promote: promote_once, ..enabled }, LayeredLazyConfig { window: 4, eager_before_full: false, + dense_to_eager: DenseToEager::All, ..enabled }, + // Sparse gathers, which promote the planes they read. LayeredLazyConfig { dense: DenseGatherMode::Sparse, origin_max_runs: 0, + promote: promote_once, ..enabled }, LayeredLazyConfig { - dense_to_eager: false, + dense_to_eager: DenseToEager::Off, ..enabled }, - ] + ]; + // On a high-latency origin a probe whose high or low plane no + // cache tier holds is gathered within the far window. + let high = [ + // The defaults: a far window of 64 probes with 64 permits, + // and predicted-dense probes on the slow origin routed to the + // eager scan. + enabled, + // Far gathers one at a time, beyond a window of one probe. + LayeredLazyConfig { + window: 1, + far_window: LAZY_PARTITIONS, + far_inflight: 1, + ..enabled + }, + // Far gathers as soon as the heap fills, beyond no window. + LayeredLazyConfig { + window: 0, + eager_before_full: false, + dense_to_eager: DenseToEager::Off, + far_window: usize::MAX, + ..enabled + }, + // The defaults with promotions after one sparse read: the far + // window and the wide origin gap while promotions run. + LayeredLazyConfig { + promote: promote_once, + ..enabled + }, + ]; + let low = low + .into_iter() + .map(|config| (OriginLatencyClass::Low, config)); + let high = high + .into_iter() + .map(|config| (OriginLatencyClass::High, config)); + low.chain(high).collect() } /// Clustered vectors around seeded centroids, in shuffled row order, @@ -8315,6 +8786,16 @@ mod tests { uri: &str, bits: u8, distance_type: DistanceType, + ) -> (Dataset, RecordBatch) { + write_rq_test_dataset(uri, bits, distance_type, true).await + } + + /// The lazy test data with a layered or a native IVF_RQ index. + async fn write_rq_test_dataset( + uri: &str, + bits: u8, + distance_type: DistanceType, + layered: bool, ) -> (Dataset, RecordBatch) { let (batch, centroids) = lazy_test_data(); let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); @@ -8322,7 +8803,7 @@ mod tests { let params = VectorIndexParams::with_ivf_rq_params( distance_type, IvfBuildParams::try_with_centroids(LAZY_PARTITIONS, centroids).unwrap(), - RQBuildParams::new(bits).with_layered(true), + RQBuildParams::new(bits).with_layered(layered), ); dataset .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) @@ -8531,9 +9012,86 @@ mod tests { ); } assert_eq!(stats.needed_not_fetched, 0, "{context}"); + assert_lazy_chain_timing(&stats, context); stats } + /// Check the release-chain timings of `stats`, taken over lazy scans + /// that all completed. An eager scan times nothing; every lazy scan + /// times its first probe's scoring once; a wait is timed only when it + /// is counted, and the waits of a gather are part of its gate wait. + /// Within one scan, the first probe's scoring is published no later + /// than the heap first fills, and when a gather waited for the + /// threshold or its turn, no gather that waits on scoring was issued + /// before that publish. + fn assert_lazy_chain_timing(stats: &LayeredLazyStats, context: &str) { + let timings = [ + ( + "rank 0 scored", + stats.lazy_rank0_scored_queries, + stats.lazy_rank0_scored_ns, + ), + ( + "first gather issue", + stats.lazy_first_gather_issue_queries, + stats.lazy_first_gather_issue_ns, + ), + ( + "window wait", + stats.lazy_window_waits, + stats.lazy_window_wait_ns, + ), + ( + "release wait", + stats.deferred_issues, + stats.lazy_release_wait_ns, + ), + ( + "far permit wait", + stats.far_permit_waits, + stats.far_permit_wait_ns, + ), + ]; + for (timing, count, ns) in timings { + if stats.lazy_queries == 0 { + assert_eq!((count, ns), (0, 0), "{timing}: {context}"); + } + if count == 0 { + assert_eq!(ns, 0, "{timing}: {context}"); + } + } + assert_eq!( + stats.lazy_rank0_scored_queries, stats.lazy_queries, + "{context}" + ); + assert!( + stats.lazy_first_gather_issue_queries <= stats.lazy_queries, + "{context}" + ); + let waits = + stats.lazy_window_wait_ns + stats.lazy_release_wait_ns + stats.far_permit_wait_ns; + assert!( + waits <= stats.gate_wait_ns, + "waits {waits} ns beyond the gate wait {} ns: {context}", + stats.gate_wait_ns + ); + if stats.lazy_queries == 1 { + if stats.time_to_first_full_ns > 0 { + assert!( + stats.lazy_rank0_scored_ns <= stats.time_to_first_full_ns, + "{stats:?} {context}" + ); + } + if stats.deferred_issues > 0 { + assert_eq!(stats.lazy_first_gather_issue_queries, 1, "{context}"); + assert!( + stats.lazy_rank0_scored_ns <= stats.lazy_first_gather_issue_ns, + "{stats:?} {context}" + ); + } + } + } + /// Deltas of the promotion gauge are not reset by snapshots; wait /// until background promotions have finished. async fn wait_for_promotions() { @@ -8575,6 +9133,7 @@ mod tests { #[case] bits: u8, #[case] distance_type: DistanceType, ) { + const STAGING_STEPS: usize = 8; let _serial = LAZY_TEST_LOCK.lock().await; let dir = TempStrDir::default(); let (_, batch) = write_lazy_test_dataset(dir.as_str(), bits, distance_type).await; @@ -8587,36 +9146,66 @@ mod tests { let filters = lazy_test_filters(); let configs = lazy_test_configs(); let mut case = 0usize; + // Single scans in which a gather waited for the threshold or its + // turn, whose release chain `assert_lazy_chain_timing` orders. + let mut deferred_scans = 0usize; + let mut count_deferred = |stats: &LayeredLazyStats| { + deferred_scans += usize::from(stats.lazy_queries == 1 && stats.deferred_issues > 0); + }; for cache in [ LazyTestCache::Origin, LazyTestCache::Small, LazyTestCache::Resident, LazyTestCache::Ungated, ] { - let (_dataset, index, _) = open_lazy_test_index(dir.as_str(), cache).await; - let lengths = &lazy_index(&index).storage.ivf().lengths; + let (dataset, opened, _) = open_lazy_test_index(dir.as_str(), cache).await; + let lengths = &lazy_index(&opened).storage.ivf().lengths; assert!(lengths.contains(&0), "the fixture needs an empty partition"); assert!( lengths.iter().any(|&rows| rows > 0 && rows < 32), "the fixture needs a partition below one FastScan block" ); + // Each setting runs on an index pinned to its origin latency + // class, over the caches the dataset's index warmed. + let (store, index_dir) = index_files(&dataset).await; + let open = |class| { + open_with_origin_latency(&dataset, store.clone(), index_dir.clone(), class) + }; + let low = open(OriginLatencyClass::Low).await; + let high = open(OriginLatencyClass::High).await; + // Without a cache every plane of the high-latency twin is on + // the slow origin, so its gathers may run ahead of scoring + // within the far window. Staging several probes at once, + // whatever the host's cores, makes some do so, as asserted + // below. + let origin_far = cache == LazyTestCache::Origin; + if origin_far { + lazy_index(&high).set_lazy_prepare_parallelism_for_test(STAGING_STEPS); + } + let mut far_early_issues = 0; + let pinned = |class: OriginLatencyClass| match class { + OriginLatencyClass::Low => &low, + OriginLatencyClass::High => &high, + }; for key in &keys { for nprobes in [1, 4, LAZY_PARTITIONS] { for k in [1, 10, 100, 5000, LAZY_ROWS + 1] { case += 1; let (filter_name, filter) = &filters[case % filters.len()]; - let config = configs[case % configs.len()]; + let (class, config) = configs[case % configs.len()]; + let index = pinned(class); let mut query = lazy_test_query(key.clone(), k, nprobes); if case.is_multiple_of(3) { query.approx_mode = ApproxMode::Accurate; } let context = format!( - "bits={bits} {distance_type:?} cache={cache:?} nprobes={nprobes} k={k} approx={:?} filter={filter_name} config={config:?}", + "bits={bits} {distance_type:?} cache={cache:?} nprobes={nprobes} k={k} approx={:?} filter={filter_name} class={class} config={config:?}", query.approx_mode ); let stats = - assert_lazy_matches_eager(&index, &query, filter, config, &context) + assert_lazy_matches_eager(index, &query, filter, config, &context) .await; + count_deferred(&stats); match cache { LazyTestCache::Resident => { assert_eq!(stats.lazy_all_resident_skips, 1, "{context}"); @@ -8635,18 +9224,25 @@ mod tests { | LazyTestCache::ColdUngated | LazyTestCache::ColdGated => {} } + if class == OriginLatencyClass::High { + far_early_issues += stats.far_early_issues; + } if case % 2 == 1 { let eager = - search_global(&index, &query, filter.clone()).await.unwrap(); + search_global(index, &query, filter.clone()).await.unwrap(); if let Some(bounded) = with_bounds(&query, &eager) { - assert_lazy_matches_eager( - &index, + let stats = assert_lazy_matches_eager( + index, &bounded, filter, config, &format!("{context} bounds"), ) .await; + count_deferred(&stats); + if class == OriginLatencyClass::High { + far_early_issues += stats.far_early_issues; + } } } } @@ -8654,16 +9250,33 @@ mod tests { } // Every setting against every filter on the widest probe set. let query = lazy_test_query(keys[0].clone(), 100, LAZY_PARTITIONS); - for config in &configs { + for &(class, config) in &configs { + let index = pinned(class); for (filter_name, filter) in &filters { let context = format!( - "bits={bits} {distance_type:?} cache={cache:?} filter={filter_name} config={config:?}" + "bits={bits} {distance_type:?} cache={cache:?} filter={filter_name} class={class} config={config:?}" ); - assert_lazy_matches_eager(&index, &query, filter, *config, &context).await; + let stats = + assert_lazy_matches_eager(index, &query, filter, config, &context) + .await; + count_deferred(&stats); + if class == OriginLatencyClass::High { + far_early_issues += stats.far_early_issues; + } } } + if origin_far { + assert!( + far_early_issues > 0, + "bits={bits} {distance_type:?}: no gather ran ahead of scoring on the high-latency origin" + ); + } wait_for_promotions().await; } + assert!( + deferred_scans > 0, + "bits={bits} {distance_type:?}: no scan deferred a gather, so no release chain was ordered" + ); } /// An ungated backend is warmed partition by partition, so every plane @@ -8671,7 +9284,8 @@ mod tests { /// warm queries read nothing from the origin file, whether the first /// probe (the only one this query gathers rows of: its threshold is /// +inf) is gathered sparsely, as the sparse gather policy always - /// does, or loaded by the eager scan, as the cost policy routes it. + /// does, or loaded by the eager scan, as the cost policy routes it + /// when every predicted-dense probe is routed. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_layered_lazy_ungated_prewarm_persists_every_plane() { let _serial = LAZY_TEST_LOCK.lock().await; @@ -8701,6 +9315,7 @@ mod tests { enabled: true, dense, promote: LazyPromotion::Off, + dense_to_eager: DenseToEager::All, ..Default::default() }, &context, @@ -8846,10 +9461,7 @@ mod tests { .map(|(partition, _)| *partition) .unwrap(); assert_eq!(ivf.storage.partition_size(empty), 0); - let high = PlaneKey { - partition: empty, - plane: 1, - }; + let high = ivf.storage.plane_key(empty, 1); assert!(ivf.index_cache.peek_resident_with_key(&high).await); tiered.hide_empty_planes.store(true, Ordering::Relaxed); assert!(!ivf.index_cache.peek_resident_with_key(&high).await); @@ -8874,60 +9486,512 @@ mod tests { ); } - /// A row-id remapper can drop physical rows, so offsets are not row - /// positions and the lazy scan must not run. + /// `auto` makes cloud object stores class high and local files and + /// memory class low, unless the session declares a class; an explicit + /// class applies to every store, whatever the session declares. + #[test] + fn test_origin_latency_class_resolves_by_store() { + let s3 = ObjectStore::new( + Arc::new(object_store::memory::InMemory::new()), + url::Url::parse("s3://bucket/").unwrap(), + None, + None, + false, + true, + 1, + lance_io::object_store::DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + ); + for (store, auto) in [ + (ObjectStore::local(), OriginLatencyClass::Low), + (ObjectStore::memory(), OriginLatencyClass::Low), + (s3, OriginLatencyClass::High), + ] { + let scheme = store.scheme().to_string(); + assert_eq!(OriginLatencyClass::of_store(&store), auto, "{scheme}"); + assert_eq!(OriginLatencyClass::resolve(None, &store), auto, "{scheme}"); + assert_eq!( + OriginLatencyClass::resolve_with_hint(None, None, &store), + auto, + "{scheme}" + ); + for class in [OriginLatencyClass::Low, OriginLatencyClass::High] { + assert_eq!( + OriginLatencyClass::resolve(Some(class), &store), + class, + "{scheme}" + ); + assert_eq!( + OriginLatencyClass::resolve_with_hint(None, Some(class), &store), + class, + "{scheme} hint={class}" + ); + for hint in [OriginLatencyClass::Low, OriginLatencyClass::High] { + assert_eq!( + OriginLatencyClass::resolve_with_hint(Some(class), Some(hint), &store), + class, + "{scheme} setting={class} hint={hint}" + ); + } + } + } + } + + /// The class is resolved when the index opens and carried into its + /// storage, and through a reconstruction from the cached state. The + /// test dataset is a local file, so the class is what the environment + /// sets, or low. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn test_layered_lazy_remapped_index_is_ineligible() { + async fn test_layered_origin_latency_survives_reconstruction() { let _serial = LAZY_TEST_LOCK.lock().await; let dir = TempStrDir::default(); - let (batch, schema) = generate_batch::(1000, None, 0.0..1.0, false); - let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], schema); - let mut dataset = Dataset::write( - reader, - dir.as_str(), - Some(WriteParams { - max_rows_per_file: 200, - ..Default::default() - }), - ) - .await - .unwrap(); - let params = VectorIndexParams::with_ivf_rq_params( - DistanceType::L2, - IvfBuildParams::new(4), - RQBuildParams::new(7).with_layered(true), + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Resident).await; + let expected = OriginLatencyClass::resolve( + origin_latency_setting().unwrap(), + &ObjectStore::local(), ); - dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + let ivf = lazy_index(&index); + assert_eq!(ivf.origin_latency(), expected); + assert_eq!(ivf.storage.origin_latency(), expected); + + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + assert!( + dataset.index_cache.get_with_key(&state_key).await.is_some(), + "the reopen must reconstruct from the cached state" + ); + let reopened = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) .await .unwrap(); - dataset.delete("id % 7 = 0").await.unwrap(); - compact_files( - &mut dataset, - CompactionOptions { - target_rows_per_fragment: 500, - defer_index_remap: true, - ..Default::default() - }, - None, - ) - .await - .unwrap(); - let uuid = dataset - .load_indices() - .await - .unwrap() - .iter() - .find(|index| index.name != FRAG_REUSE_INDEX_NAME) - .unwrap() - .uuid; - let index = dataset - .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + let reconstructed = lazy_index(&reopened); + assert_eq!(reconstructed.origin_latency(), expected); + assert_eq!(reconstructed.storage.origin_latency(), expected); + } + + /// A session's origin latency hint is the class of the IVF_RQ indexes + /// it opens, when they open and when they are reconstructed from the + /// cached state, and residency, the lazy origin gap and the promotion + /// policy follow it. Without a hint an index takes its store's class, + /// low for the test's local files, and an IVF_PQ index the session + /// opens keeps its store's class whatever the hint. Results do not + /// depend on the class. A class the environment sets overrides every + /// hint, so the hint takes effect only where it sets none. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_session_origin_latency_hint_reaches_rq_index() { + let _serial = LAZY_TEST_LOCK.lock().await; + let rq_dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(rq_dir.as_str(), 7, DistanceType::L2).await; + let pq_dir = TempStrDir::default(); + let (mut pq_dataset, _) = + generate_test_dataset::(pq_dir.as_str(), 0.0..1.0).await; + let pq_params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + IvfBuildParams::new(LIGHTWEIGHT_PQ_PARTITIONS), + lightweight_pq_params(), + ); + pq_dataset + .create_index(&["vector"], IndexType::Vector, None, &pq_params, true) .await .unwrap(); - assert!(!lazy_index(&index).storage.supports_candidate_reads()); - let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(1), 20, 4); - let stats = assert_lazy_matches_eager( + + let setting = origin_latency_setting().unwrap(); + let lazy_setting = LayeredLazyConfig::from_env().unwrap(); + let resident_setting = resident_columns_setting().unwrap(); + let vectors = batch["vector"].as_fixed_size_list(); + let queries = [ + lazy_test_query(vectors.value(0), 10, 8), + lazy_test_query(vectors.value(777), 100, LAZY_PARTITIONS), + ]; + let filter: Arc = Arc::new(NoFilter); + let lazy = LayeredLazyConfig { + enabled: true, + ..Default::default() + }; + let mut first_results = vec![None; queries.len()]; + for hint in [ + Some(OriginLatencyClass::High), + Some(OriginLatencyClass::Low), + None, + ] { + let session = Arc::new( + Session::new( + LAZY_LARGE_CACHE_BYTES, + LAZY_METADATA_CACHE_BYTES, + Arc::new(ObjectStoreRegistry::default()), + ) + .with_index_origin_latency(hint), + ); + let dataset = crate::DatasetBuilder::from_uri(rq_dir.as_str()) + .with_session(session.clone()) + .load() + .await + .unwrap(); + let store_class = OriginLatencyClass::of_store(&dataset.object_store); + assert_eq!(store_class, OriginLatencyClass::Low); + let expected = setting.or(hint).unwrap_or(store_class); + let block_size = dataset.object_store.block_size() as u64; + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let opened = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + assert!( + dataset.index_cache.get_with_key(&state_key).await.is_some(), + "the reopen must reconstruct from the cached state" + ); + let reconstructed = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + for (path, index) in [("open", &opened), ("reconstruct", &reconstructed)] { + let context = format!("hint={hint:?} {path}"); + let ivf = lazy_index(index); + assert_eq!(ivf.origin_latency(), expected, "{context}"); + assert_eq!(ivf.storage.origin_latency(), expected, "{context}"); + assert_eq!( + ivf.resident_columns_enabled(), + resident_setting.resolve(expected), + "{context}" + ); + assert_eq!( + ivf.lazy_origin_gap_bytes(), + lazy_setting.origin_gap.resolve(expected, block_size), + "{context}" + ); + assert_eq!( + ivf.layered_lazy_config().promote, + lazy_setting.promote.resolve(expected), + "{context}" + ); + for (position, query) in queries.iter().enumerate() { + let context = format!("{context} query={position}"); + assert_lazy_matches_eager(index, query, &filter, lazy, &context).await; + let eager = search_global(index, query, filter.clone()).await.unwrap(); + let bits = result_bits(&eager); + let first = first_results[position].get_or_insert_with(|| bits.clone()); + assert_eq!(&bits, first, "{context}"); + } + } + + let pq = crate::DatasetBuilder::from_uri(pq_dir.as_str()) + .with_session(session) + .load() + .await + .unwrap(); + let pq_class = OriginLatencyClass::of_store(&pq.object_store); + let pq_uuid = pq.load_indices().await.unwrap()[0].uuid; + let pq_frag_reuse_uuid = pq.frag_reuse_index_uuid().await; + let pq_state_key = + crate::index::IvfIndexStateCacheKey::new(&pq_uuid, pq_frag_reuse_uuid.as_ref()); + for path in ["open", "reconstruct"] { + let cached = pq.index_cache.get_with_key(&pq_state_key).await.is_some(); + assert_eq!(cached, path == "reconstruct", "hint={hint:?} {path}"); + let index = pq + .open_vector_index("vector", &pq_uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let ivf_pq = index + .as_any() + .downcast_ref::() + .expect("IVF_PQ index"); + assert_eq!(ivf_pq.origin_latency(), pq_class, "hint={hint:?} {path}"); + } + } + wait_for_promotions().await; + } + + /// A plane's tier follows where the backend holds it: RAM, then only + /// the serialized tier, then nowhere. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_plane_tier_follows_cache_tiers() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (_dataset, index, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::ColdUngated).await; + let tiered = tiered.unwrap(); + let ivf = lazy_index(&index); + let cache = &ivf.index_cache; + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 0) + .unwrap(); + let high = ivf.storage.plane_tier(partition, 1, cache).await; + assert_eq!(high, CacheTier::Absent); + + ivf.storage + .load_plane_entry(partition, 1, cache, None) + .await + .unwrap(); + let high = ivf.storage.plane_tier(partition, 1, cache).await; + let low = ivf.storage.plane_tier(partition, 2, cache).await; + assert_eq!((high, low), (CacheTier::Resident, CacheTier::Absent)); + + // A tier peek of a persisted plane does not admit it to RAM. + tiered.ram.clear().await; + let high = ivf.storage.plane_tier(partition, 1, cache).await; + assert_eq!(high, CacheTier::Local); + assert_eq!(tiered.ram.num_entries().await, 0); + + tiered.clear().await; + let high = ivf.storage.plane_tier(partition, 1, cache).await; + assert_eq!(high, CacheTier::Absent); + } + + /// Open the index of `dataset` over the dataset's caches, keeping its + /// bounds columns where `sign_bounds` places them. + async fn open_with_sign_bounds( + dataset: &Dataset, + sign_bounds: SignBounds, + ) -> Arc { + let index = dataset.load_indices().await.unwrap()[0].clone(); + let ivf = IvfRq::try_new( + dataset.object_store.clone(), + dataset.indice_files_dir(&index).unwrap(), + index.uuid, + None, + &dataset.metadata_cache, + dataset.index_cache.for_index(&index.uuid, None), + index.file_size_map(), + IvfOpenContext::default(), + ) + .await + .unwrap() + .with_sign_bounds_for_test(sign_bounds); + Arc::new(ivf) + } + + /// Every precision in both approximation modes, the cascade, and the + /// lazy scan over every partition, labelled for assertion messages. + fn sign_bounds_test_queries(key: &ArrayRef) -> Vec<(String, Query, LayeredLazyConfig)> { + let eager_scan = LayeredLazyConfig::default(); + let mut queries = Vec::new(); + for rq_precision in [RQPrecision::Sign, RQPrecision::High, RQPrecision::Full] { + for approx_mode in [ApproxMode::Normal, ApproxMode::Accurate] { + let mut query = lazy_test_query(key.clone(), 100, 8); + query.rq_precision = rq_precision; + query.approx_mode = approx_mode; + let label = format!("precision={rq_precision:?} approx={approx_mode:?}"); + queries.push((label, query, eager_scan)); + } + } + let mut cascade = lazy_test_query(key.clone(), 100, 8); + cascade.rq_cascade_factor = Some(4); + queries.push(("cascade".to_string(), cascade, eager_scan)); + let lazy_scan = LayeredLazyConfig { + enabled: true, + ..Default::default() + }; + let every_probe = lazy_test_query(key.clone(), 100, LAZY_PARTITIONS); + queries.push(("lazy scan".to_string(), every_probe, lazy_scan)); + queries + } + + /// Prewarm both placements into their shared cache, where each finds + /// its own sign plane: only the eager placement's holds the bounds. + async fn assert_sign_planes_coexist(placements: &[Arc]) { + for index in placements { + index.prewarm().await.unwrap(); + } + for index in placements { + let ivf = lazy_index(index); + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 0) + .unwrap(); + let sign = ivf + .index_cache + .get_resident_with_key(&ivf.storage.plane_key(partition, 0)) + .await + .expect("prewarmed sign plane"); + let bounds_in_sign = sign.0.column_by_name(FULL_BOUNDS_COLUMN).is_some(); + assert_eq!(bounds_in_sign, ivf.sign_bounds() == SignBounds::Eager); + } + } + + /// Keeping the bounds columns in their own plane changes only which + /// entries hold them: every precision, the cascade and the lazy scan + /// return the batch of the sign plane with bounds, bit for bit, from + /// the origin file and from one cache that holds both placements. + #[rstest] + #[case::rq5_l2(5, DistanceType::L2)] + #[case::rq7_cosine(7, DistanceType::Cosine)] + #[case::rq9_dot(9, DistanceType::Dot)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_sign_bounds_placements_match( + #[case] bits: u8, + #[case] distance_type: DistanceType, + ) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), bits, distance_type).await; + let vectors = batch["vector"].as_fixed_size_list(); + let keys = [vectors.value(0), vectors.value(777)]; + for cache in [LazyTestCache::Origin, LazyTestCache::TieredResident] { + let (dataset, _, _) = open_lazy_test_index(dir.as_str(), cache).await; + let placements = [ + open_with_sign_bounds(&dataset, SignBounds::Eager).await, + open_with_sign_bounds(&dataset, SignBounds::Lazy).await, + ]; + if cache == LazyTestCache::TieredResident { + assert_sign_planes_coexist(&placements).await; + } + for key in &keys { + for (label, query, scan) in sign_bounds_test_queries(key) { + let context = + format!("bits={bits} {distance_type:?} cache={cache:?} {label}"); + let mut results = Vec::new(); + for index in &placements { + let ivf = lazy_index(index); + ivf.set_layered_lazy_config_for_test(scan); + let result = search_global(index, &query, Arc::new(NoFilter)).await; + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + results.push(result_bits(&result.unwrap())); + } + assert!(!results[0].0.is_empty(), "{context}"); + assert_eq!(results[1], results[0], "{context}"); + } + } + } + } + + /// The bounds plane is read only by the precisions that prune with + /// the bounds, and only when the sign plane leaves them out. Full + /// precision prunes with the error factors every layered file has. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_bounds_plane_loads_only_when_read() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, _, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::ColdUngated).await; + let tiered = tiered.unwrap(); + for (sign_bounds, high) in [ + (SignBounds::Lazy, vec![0, 1, SIGN_BOUNDS_PLANE]), + (SignBounds::Eager, vec![0, 1]), + ] { + let index = open_with_sign_bounds(&dataset, sign_bounds).await; + let ivf = lazy_index(&index); + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 0) + .unwrap(); + for (precision, expected) in [ + (RQPrecision::Sign, vec![0]), + (RQPrecision::High, high.clone()), + (RQPrecision::Full, vec![0, 1, 2]), + ] { + tiered.clear().await; + tiered.record_plane_loads(); + ivf.storage + .load_partition_at_precision(partition, precision, &ivf.index_cache, None) + .await + .unwrap(); + let mut loaded: Vec = tiered + .take_plane_loads() + .into_iter() + .filter_map(|load| match load { + PlaneLoad::Started(plane) => Some(plane), + PlaneLoad::Finished(_) => None, + }) + .collect(); + loaded.sort_unstable(); + assert_eq!(loaded, expected, "{sign_bounds} {precision:?}"); + } + } + } + + /// The bounds placement is what the environment sets, resolved when + /// the index opens and again when it is reconstructed from the + /// cached state. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_sign_bounds_survive_reconstruction() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Resident).await; + let expected = sign_bounds_setting().unwrap(); + let ivf = lazy_index(&index); + assert!(ivf.is_layered_rq()); + assert_eq!(ivf.sign_bounds(), expected); + assert_eq!(ivf.storage.sign_bounds(), expected); + + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + assert!( + dataset.index_cache.get_with_key(&state_key).await.is_some(), + "the reopen must reconstruct from the cached state" + ); + let reopened = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(lazy_index(&reopened).is_layered_rq()); + assert_eq!(lazy_index(&reopened).sign_bounds(), expected); + } + + /// A row-id remapper can drop physical rows, so offsets are not row + /// positions and the lazy scan must not run. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_remapped_index_is_ineligible() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (batch, schema) = generate_batch::(1000, None, 0.0..1.0, false); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], schema); + let mut dataset = Dataset::write( + reader, + dir.as_str(), + Some(WriteParams { + max_rows_per_file: 200, + ..Default::default() + }), + ) + .await + .unwrap(); + let params = VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::new(4), + RQBuildParams::new(7).with_layered(true), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset.delete("id % 7 = 0").await.unwrap(); + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 500, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + let uuid = dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|index| index.name != FRAG_REUSE_INDEX_NAME) + .unwrap() + .uuid; + let index = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + assert!(!lazy_index(&index).storage.supports_candidate_reads()); + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(1), 20, 4); + let stats = assert_lazy_matches_eager( &index, &query, &(Arc::new(NoFilter) as Arc), @@ -9012,6 +10076,7 @@ mod tests { let stats = layered_stats::snapshot_and_reset(); assert_eq!(stats.lazy_queries, 0); assert_eq!(stats.lazy_all_resident_skips, 0); + assert_lazy_chain_timing(&stats, "other search paths"); let mut output = vec![result_bits(&in_partition), result_bits(&cascaded)]; output.extend(streamed.iter().map(result_bits)); output.extend(batched.iter().map(result_bits)); @@ -9185,8 +10250,11 @@ mod tests { }); tiered.set_spill(Some(spill_tx)); ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); - for (eager_before_full, dense_to_eager) in [(false, false), (true, false), (true, true)] - { + for (eager_before_full, dense_to_eager) in [ + (false, DenseToEager::Off), + (true, DenseToEager::Off), + (true, DenseToEager::All), + ] { ivf.set_layered_lazy_config_for_test(LayeredLazyConfig { enabled: true, window: 1, @@ -9223,7 +10291,7 @@ mod tests { assert_eq!(stats.needed_not_fetched, 0, "{setting} {stats:?}"); let whole_planes: u64 = stats.high_whole.iter().chain(&stats.low_whole).sum(); let eager_loads: u64 = stats.dense_to_eager.iter().sum(); - if dense_to_eager { + if dense_to_eager == DenseToEager::All { assert!( eager_loads > 0 && whole_planes == 0, "staging must load every probe eagerly: {setting} {stats:?}" @@ -9256,6 +10324,9 @@ mod tests { origin_row_reads: u64, sparse_planes: u64, whole_planes: u64, + rank0_scored_queries: u64, + first_gather_issue_queries: u64, + release_wait_ns: u64, } impl IssueTotals { @@ -9281,6 +10352,9 @@ mod tests { .chain(&stats.low_sparse) .sum::(); self.whole_planes += stats.high_whole.iter().chain(&stats.low_whole).sum::(); + self.rank0_scored_queries += stats.lazy_rank0_scored_queries; + self.first_gather_issue_queries += stats.lazy_first_gather_issue_queries; + self.release_wait_ns += stats.lazy_release_wait_ns; } } @@ -9327,7 +10401,7 @@ mod tests { dense, eager_before_full, origin_max_runs, - dense_to_eager: false, + dense_to_eager: DenseToEager::Off, ..Default::default() }; let query = lazy_test_query(vectors.value(row), k, nprobes); @@ -9373,13 +10447,19 @@ mod tests { /// Loading and scoring predicted-dense probes with the eager scan /// changes only how they are read: results match the eager scan with - /// the routing on or off, for `k` within the first probe and beyond - /// the rows of the first probes, with prefilters that keep most or few - /// rows and with distance bounds, whether the planes come from the - /// origin, an ungated persistent tier or a tier that gates admission. - /// Only queries without a prefilter or upper bound are routed: those - /// drop rows of every probe before its gather, even at `k` beyond - /// every probed row, so their probes stay in the lazy pipeline. + /// the routing off, on for every such probe (`all`) or on for those + /// that read a plane from a high-latency origin (`origin`) of either + /// class, for `k` within the first probe and beyond the rows of the + /// first probes, with prefilters that keep most or few rows and with + /// distance bounds, whether the planes come from the origin, an + /// ungated persistent tier or a tier that gates admission. Only + /// queries without a prefilter or upper bound are routed: those drop + /// rows of every probe before its gather, even at `k` beyond every + /// probed row, so their probes stay in the lazy pipeline. `origin` + /// routes no probe on a low-latency origin, and on a high-latency one + /// only probes with a plane on no cache tier: every probe `all` + /// routes when nothing is cached, none once the persistent tier holds + /// every plane. #[rstest] #[case::origin(LazyTestCache::Origin)] #[case::ungated(LazyTestCache::Ungated)] @@ -9389,13 +10469,13 @@ mod tests { let _serial = LAZY_TEST_LOCK.lock().await; let dir = TempStrDir::default(); let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; - let (_dataset, index, _) = open_lazy_test_index(dir.as_str(), cache).await; + let (dataset, index, _) = open_lazy_test_index(dir.as_str(), cache).await; let vectors = batch["vector"].as_fixed_size_list(); let filters: Vec<_> = lazy_test_filters() .into_iter() .filter(|(name, _)| matches!(*name, "none" | "sparse" | "dense")) .collect(); - let mut routed = [0u64; 2]; + let mut cases = Vec::new(); for row in [0, 777] { for nprobes in [1, 8, LAZY_PARTITIONS] { for k in [10, 1000, LAZY_ROWS + 1] { @@ -9404,35 +10484,94 @@ mod tests { let eager = search_global(&index, &query, filter.clone()).await.unwrap(); let bounded = with_bounds(&query, &eager); - for query in std::iter::once(&query).chain(&bounded) { - for dense_to_eager in [false, true] { - let config = LayeredLazyConfig { - enabled: true, - dense_to_eager, - ..Default::default() - }; - let context = format!( - "cache={cache:?} row={row} nprobes={nprobes} k={k} filter={filter_name} bounds={} dense_to_eager={dense_to_eager}", - query.upper_bound.is_some() - ); - let stats = assert_lazy_matches_eager( - &index, query, filter, config, &context, - ) - .await; - let probes_routed: u64 = stats.dense_to_eager.iter().sum(); - if *filter_name != "none" || query.upper_bound.is_some() { - assert_eq!(probes_routed, 0, "{context} {stats:?}"); - } - routed[usize::from(dense_to_eager)] += probes_routed; - } + for query in std::iter::once(query).chain(bounded) { + let label = format!( + "cache={cache:?} row={row} nprobes={nprobes} k={k} filter={filter_name} bounds={}", + query.upper_bound.is_some() + ); + cases.push((label, *filter_name == "none", query, filter.clone())); } } } } } + + let (store, index_dir) = index_files(&dataset).await; + let (low, high) = (OriginLatencyClass::Low, OriginLatencyClass::High); + let low_index = + open_with_origin_latency(&dataset, store.clone(), index_dir.clone(), low).await; + let high_index = open_with_origin_latency(&dataset, store, index_dir, high).await; + // `off` and `all` do not depend on the class. + let settings = [ + (DenseToEager::Off, low), + (DenseToEager::All, low), + (DenseToEager::Origin, low), + (DenseToEager::Origin, high), + ]; + let mut routed = [0u64; 4]; + for (label, unfiltered, query, filter) in &cases { + let mut routed_by_all = None; + for (position, &(mode, class)) in settings.iter().enumerate() { + let index = match class { + OriginLatencyClass::Low => &low_index, + OriginLatencyClass::High => &high_index, + }; + let config = LayeredLazyConfig { + enabled: true, + dense_to_eager: mode, + ..Default::default() + }; + let context = format!("{label} dense_to_eager={mode} class={class}"); + let stats = + assert_lazy_matches_eager(index, query, filter, config, &context).await; + let context = format!("{context} {stats:?}"); + let probes_routed: u64 = stats.dense_to_eager.iter().sum(); + let slow_probes: u64 = stats.s3_bound_probes.iter().sum(); + if !unfiltered || query.upper_bound.is_some() { + assert_eq!(probes_routed, 0, "{context}"); + } + if class == low { + assert_eq!(slow_probes, 0, "{context}"); + } + // Every probe `origin` routes reads a plane from the origin. + if mode == DenseToEager::Origin { + let mut buckets = stats.dense_to_eager.iter().zip(&stats.s3_bound_probes); + assert!(buckets.all(|(routed, slow)| routed <= slow), "{context}"); + } + if mode == DenseToEager::All { + routed_by_all = Some(stats.dense_to_eager); + } else if class == high && stats.lazy_queries == 1 { + match cache { + // Nothing is cached: every staged probe reads the origin. + LazyTestCache::Origin => { + let empty_probes: u64 = stats.empty.iter().sum(); + assert_eq!( + slow_probes + empty_probes, + query.minimum_nprobes as u64, + "{context}" + ); + assert_eq!(Some(stats.dense_to_eager), routed_by_all, "{context}"); + } + // Prewarm persisted every plane. + LazyTestCache::Ungated => { + assert_eq!((probes_routed, slow_probes), (0, 0), "{context}"); + } + _ => {} + } + } + routed[position] += probes_routed; + } + } wait_for_promotions().await; - assert_eq!(routed[0], 0, "switched off, no probe is routed"); - assert!(routed[1] > 0, "{routed:?}"); + let [off, all, origin_low, origin_high] = routed; + assert_eq!(off, 0, "switched off, no probe is routed"); + assert!(all > 0, "{routed:?}"); + assert_eq!(origin_low, 0, "a low-latency origin routes no probe"); + match cache { + LazyTestCache::Origin => assert_eq!(origin_high, all, "{routed:?}"), + LazyTestCache::Ungated => assert_eq!(origin_high, 0, "{routed:?}"), + _ => {} + } } /// `count` queries for `k` rows over `nprobes` probes, keyed by rows @@ -9464,7 +10603,10 @@ mod tests { /// Lazy settings without promotions, with the given issue and routing /// policies. - fn issue_policy(eager_before_full: bool, dense_to_eager: bool) -> LayeredLazyConfig { + fn issue_policy( + eager_before_full: bool, + dense_to_eager: DenseToEager, + ) -> LayeredLazyConfig { LayeredLazyConfig { enabled: true, promote: LazyPromotion::Off, @@ -9503,10 +10645,11 @@ mod tests { /// races the first probes. /// /// When `k` spans several probes, the gathers issued once the heap - /// fills would read whole planes anyway, so every probe is loaded and - /// scored by the eager scan and no ex plane is gathered whole. With - /// that switched off, the gathers read the planes whole and are issued - /// at once, and none waits for the heap or its turn. + /// fills would read whole planes anyway, so with every predicted-dense + /// probe routed (`all`) every probe is loaded and scored by the eager + /// scan and no ex plane is gathered whole. With the routing off, the + /// gathers read the planes whole and are issued at once, and none + /// waits for the heap or its turn. /// /// When the first probe alone fills the heap (small `k`), only that /// probe, scored with an infinite threshold, goes to the eager scan, @@ -9536,9 +10679,9 @@ mod tests { }; let small = queries_filled_by_first_probe(&index, vectors, SMALL_K, PROBES, QUERIES); let configs = [ - issue_policy(true, true), - issue_policy(true, false), - issue_policy(false, false), + issue_policy(true, DenseToEager::All), + issue_policy(true, DenseToEager::Off), + issue_policy(false, DenseToEager::Off), ]; ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); let [large_routed, large_on, large_off] = @@ -9602,6 +10745,36 @@ mod tests { "{small:?}" ); } + + // Every scan times its first probe's scoring, and a wait for the + // threshold or a turn is timed exactly when one happens. Each + // small-`k` query gathers later probes, whose issue waits on the + // first probe's threshold (see `assert_lazy_chain_timing`), + // whether that probe is routed or gathered as certain dense; the + // routed large-`k` queries gather nothing. + let queries = QUERIES as u64; + for totals in [ + large_routed, + large_on, + large_off, + small_routed, + small_on, + small_off, + ] { + assert_eq!(totals.rank0_scored_queries, queries, "{totals:?}"); + assert_eq!( + totals.release_wait_ns > 0, + totals.deferred_issues > 0, + "{totals:?}" + ); + } + assert_eq!( + large_routed.first_gather_issue_queries, 0, + "{large_routed:?}" + ); + for small in [small_routed, small_on, small_off] { + assert_eq!(small.first_gather_issue_queries, queries, "{small:?}"); + } } /// A selective prefilter keeps the heap from filling after the probes @@ -9629,7 +10802,10 @@ mod tests { // filter accepts. let queries = queries_filled_by_first_probe(&index, vectors, K, PROBES, QUERIES); ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); - let configs = [issue_policy(true, false), issue_policy(false, false)]; + let configs = [ + issue_policy(true, DenseToEager::Off), + issue_policy(false, DenseToEager::Off), + ]; let [on, off] = lazy_issue_totals(&index, &queries, &filter, configs).await; ivf.set_lazy_prepare_parallelism_for_test(0); assert!(on.eager_before_full > 0, "{on:?}"); @@ -9754,6 +10930,2319 @@ mod tests { assert_eq!(result_bits(&result), expected[1], "{cache:?}"); } } + + /// Rows per page of a storage file that [`rewrite_rq_storage`] cuts + /// into small pages: prime, so page ends fall inside partitions. + const RESIDENT_TEST_PAGE_ROWS: usize = 97; + /// First reads issued at once against a storage whose resident + /// columns have not loaded. + const RESIDENT_TEST_CONCURRENT_READS: usize = 32; + + /// A reader of `source`'s storage file rewritten as a `version` file: + /// in pages of `page_rows` rows when set, so that partitions straddle + /// pages, or else in the writer's default pages, as the index writer + /// lays it out. + async fn rewrite_rq_storage( + source: &IvfQuantizationStorage, + version: LanceFileVersion, + page_rows: Option, + ) -> FileReader { + let schema = source.schema(); + let batches = source + .reader() + .read_stream( + ReadBatchParams::RangeFull, + u32::MAX, + 1, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = concat_batches(&schema, batches.iter()).unwrap(); + let store = Arc::new(ObjectStore::memory()); + let path = object_store::path::Path::from("resident/storage.lance"); + // Without a cache, every write flushes a page of each column. + let options = FileWriterOptions { + data_cache_bytes: page_rows.map(|_| 1), + ..Default::default() + }; + let mut writer = lance_file::versions::create_writer( + version.resolve(), + store.create(&path).await.unwrap(), + lance_core::datatypes::Schema::try_from(schema.as_ref()).unwrap(), + options, + ) + .unwrap(); + let rows = batch.num_rows(); + let step = page_rows.unwrap_or(rows).max(1); + for offset in (0..rows).step_by(step) { + let slice = batch.slice(offset, step.min(rows - offset)); + writer.write_batch(&slice).await.unwrap(); + } + writer.finish().await.unwrap(); + let scheduler = ScanScheduler::new(store, SchedulerConfig::default_for_testing()); + let file = scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + FileReader::try_open( + file, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap() + } + + /// A storage over `reader`, a file of `source`'s rows, that takes the + /// small columns from its resident store when `resident`. + fn rq_storage_over( + source: &IvfQuantizationStorage, + reader: &FileReader, + resident: bool, + ) -> IvfQuantizationStorage { + IvfQuantizationStorage::::from_cached_with_remapper( + reader.clone(), + source.ivf().clone(), + source.metadata().clone(), + source.distance_type(), + None, + ) + .with_sign_bounds(source.sign_bounds()) + .with_resident_columns_enabled(resident) + } + + /// Whether the file rows `range` span more than one page of the first + /// column of `reader`'s file. + fn straddles_page(reader: &FileReader, range: Range) -> bool { + let pages = &reader.metadata().column_infos[0].page_infos; + let mut end = 0; + pages.iter().any(|page| { + end += page.num_rows as usize; + range.start < end && end < range.end + }) + } + + /// Ascending offsets into a partition of `rows` rows, some in runs and + /// some alone, as a sparse gather reads them. + fn resident_test_rows(rows: usize) -> Vec { + (0..rows as u32) + .filter(|row| row % 5 != 3 && row % 7 != 1) + .collect() + } + + /// Open the index of `dataset` over the dataset's caches, with its + /// small columns resident or read from the file. + async fn open_with_resident_columns( + dataset: &Dataset, + resident: bool, + ) -> Arc { + let index = dataset.load_indices().await.unwrap()[0].clone(); + let ivf = IvfRq::try_new( + dataset.object_store.clone(), + dataset.indice_files_dir(&index).unwrap(), + index.uuid, + None, + &dataset.metadata_cache, + dataset.index_cache.for_index(&index.uuid, None), + index.file_size_map(), + IvfOpenContext::default(), + ) + .await + .unwrap() + .with_resident_columns_for_test(resident); + Arc::new(ivf) + } + + /// Reads that take the small columns from the resident store return + /// the batches that reads of the file return, bit for bit: whole + /// partitions and planes, and planes at selected rows, of every + /// partition, empty ones included, in the writer's pages and in pages + /// that partitions straddle, in either file version. They are charged + /// the same bytes, except that selected rows of a v2.0 file may be + /// charged less. + #[rstest] + #[case::native_v2_0(false, LanceFileVersion::V2_0)] + #[case::native_v2_2(false, LanceFileVersion::V2_2)] + #[case::layered_v2_0(true, LanceFileVersion::V2_0)] + #[case::layered_v2_2(true, LanceFileVersion::V2_2)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_columns_reads_match_file_reads( + #[case] layered: bool, + #[case] version: LanceFileVersion, + ) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, layered).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let source = &lazy_index(&index).storage; + // A native file has the sign plane's columns too. + let mut planes = vec![0]; + if layered { + planes.extend([1, 2]); + if source.sign_bounds() == SignBounds::Lazy { + planes.push(SIGN_BOUNDS_PLANE); + } + } + for page_rows in [None, Some(RESIDENT_TEST_PAGE_ROWS)] { + let reader = rewrite_rq_storage(source, version, page_rows).await; + let file = rq_storage_over(source, &reader, false); + let resident = rq_storage_over(source, &reader, true); + let straddled = (0..source.num_partitions()) + .any(|partition| straddles_page(&reader, source.ivf().row_range(partition))); + assert_eq!(straddled, page_rows.is_some(), "pages={page_rows:?}"); + for partition in 0..source.num_partitions() { + let context = format!( + "layered={layered} {version:?} pages={page_rows:?} partition={partition}" + ); + let expected = file.load_partition(partition, None).await.unwrap(); + let actual = resident.load_partition(partition, None).await.unwrap(); + assert_eq!( + actual.to_batches().unwrap().collect::>(), + expected.to_batches().unwrap().collect::>(), + "{context}" + ); + assert_eq!(actual.deep_size_of(), expected.deep_size_of(), "{context}"); + let rows = resident_test_rows(source.partition_size(partition)); + for &plane in &planes { + for rows in [None, Some(rows.clone())] { + let sparse = rows.is_some(); + let context = format!("{context} plane={plane} sparse={sparse}"); + let expected = file + .read_plane(partition, plane, rows.clone(), None) + .await + .unwrap(); + let actual = resident + .read_plane(partition, plane, rows, None) + .await + .unwrap(); + assert_eq!(actual, expected, "{context}"); + let (actual_bytes, expected_bytes) = + (actual.deep_size_of(), expected.deep_size_of()); + if sparse && version == LanceFileVersion::V2_0 { + // A v2.0 read of several row runs in a page copies a + // fixed-width column's runs into a buffer that reserves + // the bytes after the first run only, so it reallocates + // to about twice their bytes, while the resident copies + // hold exactly theirs. The two reads lay out every + // other column, and pages, the same way. Only whole + // planes are cached as plane entries, where the charge + // must match; selected rows are scored and dropped. + assert!( + actual_bytes <= expected_bytes, + "{context}: {actual_bytes} > {expected_bytes}" + ); + } else { + assert_eq!(actual_bytes, expected_bytes, "{context}"); + } + } + } + } + assert_eq!( + resident.resident_columns().loaded_bytes(), + Some(source.resident_columns_bytes()) + ); + } + } + + /// With its small columns resident, a storage reads only the code and + /// bounds columns from the file. In the writer's pages, as the + /// benchmark indexes are laid out, a native partition miss makes 2 + /// requests instead of 8, and a layered whole plane 1. The read that + /// loads the store counts the load's requests too. + #[rstest] + #[case::native(false)] + #[case::layered(true)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_columns_cut_origin_requests(#[case] layered: bool) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, layered).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + assert_eq!(lazy_index(&index).is_layered_rq(), layered); + let source = &lazy_index(&index).storage; + let reader = rewrite_rq_storage(source, LanceFileVersion::V2_0, None).await; + // Lazy bounds whatever `LANCE_RQ_SIGN_BOUNDS` sets: a whole sign + // plane then reads its code column alone from the file, where + // eager bounds would add the two bounds columns, never resident. + let file = rq_storage_over(source, &reader, false).with_sign_bounds(SignBounds::Lazy); + let resident = + rq_storage_over(source, &reader, true).with_sign_bounds(SignBounds::Lazy); + let schema = source.schema(); + // The code and bounds columns are those without a fixed width. + let file_columns = schema + .fields() + .iter() + .filter(|field| field.data_type().primitive_width().is_none()) + .count() as u64; + let columns = schema.fields().len() as u64; + if !layered { + assert_eq!((columns, file_columns), (8, 2)); + } + let partition = (0..source.num_partitions()) + .find(|&partition| source.partition_size(partition) > 1) + .unwrap(); + let requests = |stats: &IoStats| stats.snapshot().iops; + + layered_stats::snapshot_and_reset(); + let first = IoStats::new(); + resident + .load_partition(partition, Some(first.clone())) + .await + .unwrap(); + let load = layered_stats::snapshot_and_reset(); + assert!(load.resident_columns_load_requests > 0, "{load:?}"); + assert_eq!( + requests(&first), + load.resident_columns_load_requests + file_columns + ); + + for (storage, expected) in [(&file, columns), (&resident, file_columns)] { + let stats = IoStats::new(); + storage + .load_partition(partition, Some(stats.clone())) + .await + .unwrap(); + let enabled = storage.resident_columns_enabled(); + assert_eq!(requests(&stats), expected, "resident={enabled}"); + } + if layered { + let rows = resident_test_rows(source.partition_size(partition)); + for plane in [0, 1, 2] { + let columns = plane_columns(plane, SignBounds::Lazy).len() as u64; + let mut sparse = Vec::new(); + for (storage, expected) in [(&file, columns), (&resident, 1)] { + let enabled = storage.resident_columns_enabled(); + let whole = IoStats::new(); + storage + .read_plane(partition, plane, None, Some(whole.clone())) + .await + .unwrap(); + assert_eq!( + requests(&whole), + expected, + "plane={plane} resident={enabled}" + ); + let stats = IoStats::new(); + storage + .read_plane(partition, plane, Some(rows.clone()), Some(stats.clone())) + .await + .unwrap(); + sparse.push(requests(&stats)); + } + // A sparse read fetches the row runs of the code column alone. + assert!(sparse[1] < sparse[0], "plane={plane} {sparse:?}"); + } + } + // Only the first read loaded the store. + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_bytes, 0, "{stats:?}"); + } + + /// First reads issued at once, as by queries that miss together, load + /// the resident store once, and a storage that shares the store, as a + /// reconstruction of the index does, does not load it again. The load + /// holds the bytes that `resident_columns_bytes` gives without I/O. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_columns_load_once() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let source = &lazy_index(&index).storage; + // Row ids and the add, scale and error factors of the sign codes + // and the add and scale factors of both ex planes. + let per_row: u64 = source + .schema() + .fields() + .iter() + .map(|field| match field.data_type() { + DataType::UInt64 => 8, + DataType::Float32 => 4, + _ => 0, + }) + .sum(); + assert_eq!(per_row, 36); + let bytes = source.resident_columns_bytes(); + assert_eq!(bytes, per_row * LAZY_ROWS as u64); + + let storage = Arc::new(rq_storage_over(source, source.reader(), true)); + layered_stats::snapshot_and_reset(); + let reads: Vec<_> = (0..RESIDENT_TEST_CONCURRENT_READS) + .map(|read| { + let storage = storage.clone(); + tokio::spawn(async move { + let partition = read % LAZY_PARTITIONS; + if read % 2 == 0 { + storage + .read_plane(partition, 0, None, None) + .await + .map(|plane| plane.num_rows()) + } else { + storage + .load_partition(partition, None) + .await + .map(|loaded| loaded.len()) + } + }) + }) + .collect(); + for read in reads { + read.await.unwrap().unwrap(); + } + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_bytes, bytes, "{stats:?}"); + // The arrays hold at least their values. + assert!(stats.resident_columns_alloc_bytes >= bytes, "{stats:?}"); + assert!(stats.resident_columns_load_requests > 0, "{stats:?}"); + assert!(stats.resident_columns_load_bytes > 0, "{stats:?}"); + assert_eq!(storage.resident_columns().loaded_bytes(), Some(bytes)); + + let shared = rq_storage_over(source, source.reader(), true) + .with_resident_columns(storage.resident_columns().clone()); + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| source.partition_size(partition) > 0) + .unwrap(); + shared.load_partition(partition, None).await.unwrap(); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_bytes, 0, "{stats:?}"); + } + + /// Searches return the same batches whether the small columns are + /// resident or read from the origin file: every precision, the + /// cascade and the lazy scan of a layered index, and a native index in + /// both approximation modes. + #[rstest] + #[case::native(false)] + #[case::layered(true)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_columns_search_parity(#[case] layered: bool) { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = + write_rq_test_dataset(dir.as_str(), 7, DistanceType::L2, layered).await; + let vectors = batch["vector"].as_fixed_size_list(); + let (dataset, _, _) = open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let file = open_with_resident_columns(&dataset, false).await; + let resident = open_with_resident_columns(&dataset, true).await; + layered_stats::snapshot_and_reset(); + for key in [vectors.value(0), vectors.value(777)] { + let queries = if layered { + sign_bounds_test_queries(&key) + } else { + [ApproxMode::Normal, ApproxMode::Accurate] + .into_iter() + .map(|approx_mode| { + let mut query = lazy_test_query(key.clone(), 100, 8); + query.approx_mode = approx_mode; + let label = format!("approx={approx_mode:?}"); + (label, query, LayeredLazyConfig::default()) + }) + .collect() + }; + for (label, query, scan) in queries { + let mut results = Vec::new(); + for index in [&file, &resident] { + let ivf = lazy_index(index); + ivf.set_layered_lazy_config_for_test(scan); + let result = search_global(index, &query, Arc::new(NoFilter)).await; + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + results.push(result_bits(&result.unwrap())); + } + let context = format!("layered={layered} {label}"); + assert!(!results[0].0.is_empty(), "{context}"); + assert_eq!(results[1], results[0], "{context}"); + } + } + let stats = layered_stats::snapshot_and_reset(); + let bytes = lazy_index(&resident).resident_columns_bytes(); + assert_eq!(stats.resident_columns_bytes, bytes, "{stats:?}"); + let file_store = lazy_index(&file).storage.resident_columns(); + assert_eq!(file_store.loaded_bytes(), None); + } + + /// Whether the small columns are resident is what the environment + /// resolves for the index's origin and index cache, when the index + /// opens and when it is reconstructed from the cached state, and a + /// reconstruction while the index lives shares its store. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_columns_survive_reconstruction() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Resident).await; + let ivf = lazy_index(&index); + let bytes = ivf.resident_columns_bytes(); + let size = ivf.storage.resident_store_size().unwrap(); + assert_eq!(size.bytes, bytes); + let setting = resident_columns_setting().unwrap(); + let expected = setting.resolve(ivf.origin_latency()) + && setting.admits(size, dataset.index_cache.max_entry_bytes()); + assert_eq!(ivf.resident_columns_enabled(), expected); + assert_eq!(ivf.storage.resident_columns_enabled(), expected); + // Load the store through the index's handle. + let partition = (0..LAZY_PARTITIONS) + .find(|&partition| ivf.storage.partition_size(partition) > 0) + .unwrap(); + rq_storage_over(&ivf.storage, ivf.storage.reader(), true) + .with_resident_columns(ivf.storage.resident_columns().clone()) + .load_partition(partition, None) + .await + .unwrap(); + assert_eq!(ivf.storage.resident_columns().loaded_bytes(), Some(bytes)); + + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + assert!( + dataset.index_cache.get_with_key(&state_key).await.is_some(), + "the reopen must reconstruct from the cached state" + ); + let reopened = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let reopened = lazy_index(&reopened); + assert_eq!(reopened.resident_columns_enabled(), expected); + // A resident store is the one the live index loaded; an index + // that keeps none resident reads the file. + let store = reopened.storage.resident_columns(); + assert_eq!( + store.loaded_bytes(), + expected.then_some(bytes), + "shared store" + ); + } + + /// Scheme of the stores [`CloudTestStoreProvider`] makes. + const CLOUD_TEST_SCHEME: &str = "rqcloud"; + + /// In-memory stores over one backend, under a scheme that `auto` + /// takes for a cloud object store: a dataset written through a + /// registry holding the provider opens as class high, from any + /// session over that registry. + #[derive(Debug)] + struct CloudTestStoreProvider(Arc); + + #[async_trait::async_trait] + impl lance_io::object_store::ObjectStoreProvider for CloudTestStoreProvider { + async fn new_store( + &self, + base_path: url::Url, + _params: &ObjectStoreParams, + ) -> Result { + Ok(ObjectStore::new( + self.0.clone(), + base_path, + None, + None, + false, + true, + lance_io::object_store::DEFAULT_CLOUD_IO_PARALLELISM, + lance_io::object_store::DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )) + } + } + + /// A registry that resolves [`CLOUD_TEST_SCHEME`] URIs to a new + /// in-memory backend. + fn cloud_test_registry() -> Arc { + let registry = Arc::new(ObjectStoreRegistry::default()); + let backend = Arc::new(object_store::memory::InMemory::new()); + registry.insert(CLOUD_TEST_SCHEME, Arc::new(CloudTestStoreProvider(backend))); + registry + } + + /// The lazy test data with a layered index, written through + /// `registry` to a new dataset under [`CLOUD_TEST_SCHEME`]: the + /// dataset's URI and the data. + async fn write_cloud_lazy_test_dataset( + registry: Arc, + ) -> (String, RecordBatch) { + let uri = format!("{CLOUD_TEST_SCHEME}://bucket/{}", Uuid::new_v4()); + let (batch, centroids) = lazy_test_data(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let params = WriteParams { + session: Some(Arc::new(Session::new( + 0, + LAZY_METADATA_CACHE_BYTES, + registry, + ))), + ..Default::default() + }; + let mut dataset = Dataset::write(reader, uri.as_str(), Some(params)) + .await + .unwrap(); + let params = VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::try_with_centroids(LAZY_PARTITIONS, centroids).unwrap(), + RQBuildParams::new(7).with_layered(true), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + (uri, batch) + } + + /// The dataset at `uri` over `session`. + async fn open_with_session(uri: &str, session: Session) -> Dataset { + crate::DatasetBuilder::from_uri(uri) + .with_session(Arc::new(session)) + .load() + .await + .unwrap() + } + + /// Open the index of `dataset` with UUID `uuid`, as a query does, and + /// search it for `query` with the lazy scan settings `config`. + async fn open_and_search( + dataset: &Dataset, + uuid: &Uuid, + query: &Query, + config: LayeredLazyConfig, + ) -> Result> { + let index = dataset + .open_vector_index("vector", uuid, &NoOpMetricsCollector) + .await?; + lazy_index(&index).set_layered_lazy_config_for_test(config); + search_global(&index, query, Arc::new(NoFilter)).await?; + Ok(index) + } + + /// Indexes of one file opened at once on a fresh session, as a + /// benchmark's workers open a cold index, bind to one resident store + /// and one far gather pool: their first searches load the store once, + /// the index cache charges and pins it while they live, and a permit + /// taken through one index is taken from every index's pool. The file + /// is on a store that `auto` takes for a cloud object store, so the + /// index opens as class high and keeps its small columns resident, + /// unless the environment sets otherwise; an index cache with no room + /// admits no store, so there `auto` keeps them in the file. + #[rstest] + #[case::charged(LAZY_LARGE_CACHE_BYTES)] + #[case::zero_capacity(0)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_concurrent_cold_opens_load_once(#[case] cache_bytes: usize) { + const OPENS: usize = 8; + const FAR_INFLIGHT: usize = 2; + let _serial = LAZY_TEST_LOCK.lock().await; + let registry = cloud_test_registry(); + let (uri, batch) = write_cloud_lazy_test_dataset(registry.clone()).await; + let session = Session::new(cache_bytes, LAZY_METADATA_CACHE_BYTES, registry); + let dataset = open_with_session(&uri, session).await; + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let key = batch["vector"].as_fixed_size_list().value(0); + let query = lazy_test_query(key, 10, 8); + let config = LayeredLazyConfig { + enabled: true, + far_inflight: FAR_INFLIGHT, + ..Default::default() + }; + + layered_stats::snapshot_and_reset(); + let indexes = futures::future::try_join_all( + (0..OPENS).map(|_| open_and_search(&dataset, &uuid, &query, config)), + ) + .await + .unwrap(); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.lazy_queries, OPENS as u64, "{stats:?}"); + let first = lazy_index(&indexes[0]); + if origin_latency_setting().unwrap().is_none() { + assert_eq!(first.origin_latency(), OriginLatencyClass::High); + } + let setting = resident_columns_setting().unwrap(); + let bytes = first.resident_columns_bytes(); + let max_entry_bytes = dataset.index_cache.max_entry_bytes(); + assert_eq!(max_entry_bytes, Some(cache_bytes as u64)); + let resident = first.resident_columns_enabled(); + assert_eq!( + resident, + setting.resolve(first.origin_latency()) + && setting.admits( + first.storage.resident_store_size().unwrap(), + max_entry_bytes + ) + ); + if cache_bytes == 0 && setting == ResidentColumnsSetting::Auto { + assert!(!resident); + assert!(stats.resident_columns_oversize > 0, "{stats:?}"); + } + let loaded = resident.then_some(bytes); + // One load, however many indexes read the file at once. + assert_eq!( + stats.resident_columns_loads, + u64::from(resident), + "{stats:?}" + ); + assert_eq!( + stats.resident_columns_bytes, + loaded.unwrap_or(0), + "{stats:?}" + ); + assert_eq!( + stats.resident_columns_load_requests > 0, + resident, + "{stats:?}" + ); + for index in &indexes { + let store = lazy_index(index).storage.resident_columns(); + assert_eq!(store.loaded_bytes(), loaded); + } + let pools: Vec<_> = indexes + .iter() + .map(|index| lazy_index(index).storage.lazy_far_permits(&config).unwrap()) + .collect(); + let held = pools[0].try_acquire().unwrap(); + for pool in &pools { + assert_eq!((pool.size(), pool.in_flight()), (FAR_INFLIGHT, 1)); + } + drop(held); + assert!(pools.iter().all(|pool| pool.in_flight() == 0)); + } + + /// A cached state lost while an older index of its file still runs + /// rebinds to that index's resident store and far gather pool, both + /// when the files are opened anew and when the state is read back + /// from the persistent tier, so neither loads the store again, and + /// the open admits the store the cache lost again. The cache then + /// holds it idle, though no state holds it: it stays loaded once + /// every index of the file is dropped. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_lost_state_rebinds_index_file_handles() { + let _serial = LAZY_TEST_LOCK.lock().await; + let registry = cloud_test_registry(); + let (uri, batch) = write_cloud_lazy_test_dataset(registry.clone()).await; + let tiered = Arc::new(TieredPlaneTestBackend::new(LAZY_LARGE_CACHE_BYTES, false)); + let session = Session::with_index_cache_backend( + tiered.clone(), + LAZY_METADATA_CACHE_BYTES, + registry, + ); + let dataset = open_with_session(&uri, session).await; + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + let key = batch["vector"].as_fixed_size_list().value(0); + let query = lazy_test_query(key, 10, 8); + let config = LayeredLazyConfig::default(); + let far = LayeredLazyConfig { + far_inflight: 2, + ..config + }; + + layered_stats::snapshot_and_reset(); + let old = open_and_search(&dataset, &uuid, &query, config) + .await + .unwrap(); + let ivf = lazy_index(&old); + let loaded = ivf + .resident_columns_enabled() + .then(|| ivf.resident_columns_bytes()); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!( + stats.resident_columns_bytes, + loaded.unwrap_or(0), + "{stats:?}" + ); + assert_eq!(ivf.storage.resident_columns().loaded_bytes(), loaded); + let pool = ivf.storage.lazy_far_permits(&far).unwrap(); + let held = pool.try_acquire().unwrap(); + + // Evicted from every tier: the next open reads the files anew. + tiered.clear().await; + assert!(dataset.index_cache.get_with_key(&state_key).await.is_none()); + let reopened = open_and_search(&dataset, &uuid, &query, config) + .await + .unwrap(); + // Evicted from RAM only: the next open reads the state back. + assert_eq!(tiered.persisted_entries(IvfStateEntryBox::TYPE_ID), 1); + tiered.ram.clear().await; + let read_back = open_and_search(&dataset, &uuid, &query, config) + .await + .unwrap(); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_bytes, 0, "{stats:?}"); + for (label, index) in [("reopened", &reopened), ("read back", &read_back)] { + let storage = &lazy_index(index).storage; + assert_eq!(storage.resident_columns().loaded_bytes(), loaded, "{label}"); + let permits = storage.lazy_far_permits(&far).unwrap(); + assert_eq!((permits.size(), permits.in_flight()), (2, 1), "{label}"); + } + drop(held); + + drop((old, reopened, read_back, pool)); + let again = open_and_search(&dataset, &uuid, &query, config) + .await + .unwrap(); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_bytes, 0, "{stats:?}"); + let store = lazy_index(&again).storage.resident_columns(); + assert_eq!(store.loaded_bytes(), loaded); + } + + /// RAM of the tiered backend the resident store lifecycle tests use: + /// room for the lazy test index's store under `auto`, which takes at + /// most its entry limit, and for a query's planes, and small enough that + /// churn evicts idle entries quickly. + const RESIDENT_TEST_RAM_BYTES: usize = 4 * 1024 * 1024; + /// Bytes of the lazy test index's resident store: row ids and seven + /// four-byte factors per row. + const LAZY_RESIDENT_BYTES: u64 = 36 * LAZY_ROWS as u64; + /// Room a store's cache entry may take beyond its arrays' buffers. + const RESIDENT_ENTRY_OVERHEAD_BYTES: u64 = 64 * 1024; + + /// Whether the lazy test index opened as class high, as on S3, over + /// an index cache admitting entries up to `max_entry_bytes`, keeps + /// its small columns resident as the environment resolves it, for + /// the default `index` lifetime. The resident store lifecycle tests + /// have nothing to check when the environment keeps the columns in + /// the file, and check the `process` lifetime elsewhere. + fn resident_as_high(max_entry_bytes: Option) -> bool { + let class = origin_latency_setting() + .unwrap() + .unwrap_or(OriginLatencyClass::High); + let setting = resident_columns_setting().unwrap(); + let index_lifetime = resident_lifetime_setting().unwrap() == ResidentLifetime::Index; + index_lifetime + && setting.resolve(class) + && setting.admits( + ResidentStoreSize { + bytes: LAZY_RESIDENT_BYTES, + charge: LAZY_RESIDENT_BYTES + RESIDENT_ENTRY_OVERHEAD_BYTES, + }, + max_entry_bytes, + ) + } + + /// The layered lazy test index at `uri`, opened by a session that + /// declares class high over the index cache `backend`: the dataset, + /// the index's UUID and the key of its storage file. + struct ResidentTestIndex { + dataset: Dataset, + uuid: Uuid, + file: IndexFileKey, + } + + impl ResidentTestIndex { + async fn open(uri: &str, backend: Arc) -> Self { + let session = Session::with_index_cache_backend( + backend, + LAZY_METADATA_CACHE_BYTES, + Arc::new(ObjectStoreRegistry::default()), + ) + .with_index_origin_latency(Some(OriginLatencyClass::High)); + let dataset = open_with_session(uri, session).await; + let index = dataset.load_indices().await.unwrap()[0].clone(); + let index_dir = dataset.indice_files_dir(&index).unwrap(); + let file = + super::super::aux_file_key(&dataset.object_store, &index_dir, &index.uuid); + Self { + dataset, + uuid: index.uuid, + file, + } + } + + async fn open_index(&self) -> Arc { + self.dataset + .open_vector_index("vector", &self.uuid, &NoOpMetricsCollector) + .await + .unwrap() + } + + /// Open the index as a query does, search it and drop it. + async fn search(&self, query: &Query) -> (Vec, Vec) { + let index = self.open_index().await; + let result = search_global(&index, query, Arc::new(NoFilter)) + .await + .unwrap(); + result_bits(&result) + } + + /// The index's namespace of the index cache that charges the store. + fn file_cache(&self) -> LanceCache { + self.dataset.index_cache.for_index(&self.uuid, None) + } + + /// Whether RAM holds the store, checked without an access. + async fn store_cached(&self) -> bool { + self.file_cache() + .peek_resident_with_key(&ResidentColumnsKey::new(&self.file)) + .await + } + } + + /// Polls of a condition that tasks a search spawned settle, which may + /// hold its index for a moment after the search returned. + const DRAIN_POLLS: usize = 500; + const DRAIN_POLL_INTERVAL: Duration = Duration::from_millis(10); + + /// Wait until no loaded store of `file` is alive. + async fn wait_until_store_freed(file: &IndexFileKey) -> bool { + for _ in 0..DRAIN_POLLS { + if !resident_store_is_live(file) { + return true; + } + tokio::time::sleep(DRAIN_POLL_INTERVAL).await; + } + false + } + + /// The store is an entry of the index's namespace of the index + /// cache, charged what its arrays allocate; no storage charges it + /// again. Index handles keep it alive independently of cache eviction. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_store_charged_in_index_cache() { + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend.clone()).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 10, 8); + + layered_stats::snapshot_and_reset(); + let index = test.open_index().await; + search_global(&index, &query, Arc::new(NoFilter)) + .await + .unwrap(); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert_eq!( + stats.resident_columns_bytes, LAZY_RESIDENT_BYTES, + "{stats:?}" + ); + // The store holds exactly its values. + let alloc = stats.resident_columns_alloc_bytes; + assert_eq!(alloc, LAZY_RESIDENT_BYTES, "{stats:?}"); + let entry = test + .file_cache() + .get_resident_with_key(&ResidentColumnsKey::new(&test.file)) + .await + .expect("the store is cached"); + // Measure the entry through ordinary cache admission, including + // the key and Arc overhead, independently of the other entries. + let accounting_cache = LanceCache::with_capacity(RESIDENT_TEST_RAM_BYTES); + accounting_cache + .insert_with_key(&ResidentColumnsKey::new(&test.file), entry.clone()) + .await; + let charged = accounting_cache.size_bytes().await as u64; + assert!(charged >= alloc, "{charged} < {alloc}"); + assert!( + charged <= alloc + RESIDENT_ENTRY_OVERHEAD_BYTES, + "{charged}" + ); + let ivf = lazy_index(&index); + assert!((ivf.storage.deep_size_of() as u64) < LAZY_RESIDENT_BYTES); + assert_eq!(charged, ivf.storage.resident_store_size().unwrap().charge); + assert!(test.dataset.index_cache.size_bytes().await as u64 >= charged); + + drop((entry, index)); + assert!(test.store_cached().await); + } + + /// A cached state read back from a persistent tier is admitted to RAM + /// before the index binds its store, and that admission can evict: + /// the open holds a strong reference first. With the state on the persistent + /// tier alone, the store in RAM idle and no index live, ten times the + /// RAM of churn during the state's admission evicts the cache entry + /// without reloading its store. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_state_promotion_from_nvme_keeps_store() { + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend.clone()).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 10, 8); + + layered_stats::snapshot_and_reset(); + let expected = test.search(&query).await; + assert!(test.store_cached().await); + backend.evict_states_from_ram(); + backend.churn_on_next_admission(10 * RESIDENT_TEST_RAM_BYTES); + let index = test.open_index().await; + assert_eq!( + lazy_index(&index).storage.resident_columns().loaded_bytes(), + Some(LAZY_RESIDENT_BYTES), + ); + let actual = result_bits( + &search_global(&index, &query, Arc::new(NoFilter)) + .await + .unwrap(), + ); + assert_eq!(actual, expected); + assert_eq!(backend.churned_admissions.load(Ordering::SeqCst), 1); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert!(stats.resident_store_evictions > 0, "{stats:?}"); + assert!(resident_store_is_live(&test.file)); + drop(index); + backend.ram.clear().await; + assert!(wait_until_store_freed(&test.file).await); + } + + /// The V12 regression: a cached state evicted while no index of its + /// file is live leaves the store cached, since the state holds none, + /// and the next open reads the state back and binds the same store. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_state_eviction_keeps_store() { + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend.clone()).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 100, 16); + + layered_stats::snapshot_and_reset(); + let expected = test.search(&query).await; + for _ in 0..3 { + backend.evict_states_from_ram(); + assert_eq!(test.search(&query).await, expected); + } + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert_eq!(stats.resident_store_evictions, 0, "{stats:?}"); + } + + /// Cached index state does not own resident buffers: after every query + /// finishes, evicting the store frees it even while that state is held. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_cached_state_does_not_own_resident_store() { + const CONCURRENT_QUERIES: usize = 4; + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend.clone()).await; + let vectors = batch["vector"].as_fixed_size_list(); + let lazy = LayeredLazyConfig { + enabled: true, + ..Default::default() + }; + let queries: Vec<_> = (0..CONCURRENT_QUERIES) + .map(|query| lazy_test_query(vectors.value(query * 97), 10, 8)) + .collect(); + futures::future::try_join_all( + queries + .iter() + .map(|query| open_and_search(&test.dataset, &test.uuid, query, lazy)), + ) + .await + .unwrap(); + test.search(&queries[0]).await; + + let entry = test + .file_cache() + .get_resident_with_key(&ResidentColumnsKey::new(&test.file)) + .await + .expect("the idle store stays cached"); + let frag_reuse_uuid = test.dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&test.uuid, frag_reuse_uuid.as_ref()); + let state = test + .dataset + .index_cache + .get_with_key(&state_key) + .await + .unwrap(); + drop(entry); + backend.churn_ram(10 * RESIDENT_TEST_RAM_BYTES).await; + assert!(wait_until_store_freed(&test.file).await); + drop(state); + } + + /// An idle store is an ordinary entry: another index's churn evicts + /// it once no index of its file is live, freeing it, and the next + /// query that reads the file loads it once more, with the same + /// results. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_idle_store_evicted_by_other_index_reloads_once() { + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend.clone()).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 10, 8); + + layered_stats::snapshot_and_reset(); + let expected = test.search(&query).await; + backend.churn_ram(10 * RESIDENT_TEST_RAM_BYTES).await; + assert!(!test.store_cached().await); + assert!(wait_until_store_freed(&test.file).await); + // The churn pushed the index's planes off the persistent tier + // too, so the next query reads them from the file. + backend.disk.lock().unwrap().clear(); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert_eq!(stats.resident_store_evictions, 1, "{stats:?}"); + + assert_eq!(test.search(&query).await, expected); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert!(test.store_cached().await); + } + + /// A state read back from the persistent tier after RAM lost both it + /// and the store, while an older index of the file still runs, binds + /// that index's store through the registry and admits it again, + /// without loading its columns again. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_state_from_persistent_tier_binds_resident_store() { + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend.clone()).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 10, 8); + + let old = test.open_index().await; + let expected = result_bits( + &search_global(&old, &query, Arc::new(NoFilter)) + .await + .unwrap(), + ); + backend.ram.clear().await; + assert!(!test.store_cached().await); + layered_stats::snapshot_and_reset(); + let read_back = test.open_index().await; + let actual = search_global(&read_back, &query, Arc::new(NoFilter)) + .await + .unwrap(); + assert_eq!(result_bits(&actual), expected); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 0, "{stats:?}"); + assert_eq!(stats.resident_columns_registry_reuses, 1, "{stats:?}"); + assert!(test.store_cached().await); + assert_eq!( + lazy_index(&read_back) + .storage + .resident_columns() + .loaded_bytes(), + lazy_index(&old).storage.resident_columns().loaded_bytes(), + ); + } + + /// The store is charged in the index's namespace without a fragment + /// reuse segment, so opens under two fragment reuse namespaces, as + /// before and after a compaction with deferred remapping, share one + /// cached store and load it once. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_store_survives_frag_reuse_change() { + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend.clone()).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 10, 8); + let index = test.dataset.load_indices().await.unwrap()[0].clone(); + + layered_stats::snapshot_and_reset(); + let mut results = Vec::new(); + for frag_reuse_uuid in [Uuid::new_v4(), Uuid::new_v4()] { + let context = IvfOpenContext { + origin_latency_hint: Some(OriginLatencyClass::High), + file_cache: Some(test.file_cache()), + resident_columns_handle: None, + }; + let opened: Arc = Arc::new( + IvfRq::try_new( + test.dataset.object_store.clone(), + test.dataset.indice_files_dir(&index).unwrap(), + index.uuid, + None, + &test.dataset.metadata_cache, + test.dataset + .index_cache + .for_index(&index.uuid, Some(&frag_reuse_uuid)), + index.file_size_map(), + context, + ) + .await + .unwrap(), + ); + let result = search_global(&opened, &query, Arc::new(NoFilter)) + .await + .unwrap(); + results.push(result_bits(&result)); + } + assert_eq!(results[0], results[1]); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.resident_columns_loads, 1, "{stats:?}"); + assert_eq!(stats.resident_columns_binds, 2, "{stats:?}"); + assert!(test.store_cached().await); + } + + /// The store lives no longer than its session: once every index, the + /// dataset and the session are dropped, it is freed. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_resident_store_freed_on_session_drop() { + let _serial = LAZY_TEST_LOCK.lock().await; + let backend = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + if !resident_as_high(backend.max_entry_bytes()) { + return; + } + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let test = ResidentTestIndex::open(dir.as_str(), backend).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 10, 8); + test.search(&query).await; + assert!(resident_store_is_live(&test.file)); + let file = test.file.clone(); + drop(test); + assert!(wait_until_store_freed(&file).await); + } + + /// A store that exceeds the index cache's + /// largest admissible entry stays in the file under `auto`, which + /// counts it, with the same results as a resident store. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_oversize_store_resolves_off_under_auto() { + const SMALL_RAM_BYTES: usize = 512 * 1024; + let _serial = LAZY_TEST_LOCK.lock().await; + let small = Arc::new(TieredPlaneTestBackend::new(SMALL_RAM_BYTES, false)); + let large = Arc::new(TieredPlaneTestBackend::new(RESIDENT_TEST_RAM_BYTES, false)); + let setting = resident_columns_setting().unwrap(); + if setting != ResidentColumnsSetting::Auto || !resident_as_high(large.max_entry_bytes()) + { + return; + } + assert!(LAZY_RESIDENT_BYTES > small.max_entry_bytes().unwrap()); + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(0), 100, 16); + let mut results = Vec::new(); + for (backend, resident) in [(small, false), (large, true)] { + let test = ResidentTestIndex::open(dir.as_str(), backend).await; + layered_stats::snapshot_and_reset(); + let index = test.open_index().await; + assert_eq!(lazy_index(&index).resident_columns_enabled(), resident); + let result = search_global(&index, &query, Arc::new(NoFilter)) + .await + .unwrap(); + results.push(result_bits(&result)); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!( + stats.resident_columns_oversize, + u64::from(!resident), + "{stats:?}" + ); + assert_eq!( + stats.resident_columns_loads, + u64::from(resident), + "{stats:?}" + ); + } + assert_eq!(results[0], results[1]); + } + + /// Gaps that coalesced origin reads are checked at, in ascending + /// order so that neighbours compare a narrower gap with a wider one: + /// touching ranges only, V11's S3 block size, a narrower gap for a + /// high-latency origin, the gap `auto` gives one, and every range of + /// a column page. + const COALESCE_TEST_GAPS: [u64; 5] = [ + 0, + 64 * 1024, + 256 * 1024, + HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, + u64::MAX, + ]; + /// Spacing of the single rows of a sparse selection. + const COALESCE_TEST_ROW_STEP: usize = 41; + + /// Sparse plane reads return the same batch at every coalescing gap: + /// every plane of every partition, empty ones included, with the + /// small columns read from the file or resident, in the writer's + /// pages and in pages that partitions straddle, in either file + /// version. In a v2.0 file, as the benchmark indexes are laid out, a + /// wider gap makes no more requests, gap 0 makes more than a gap that + /// covers the partition, and that makes no more than a whole read, + /// which is one request per column page. + #[rstest] + #[case::v2_0(LanceFileVersion::V2_0)] + #[case::v2_2(LanceFileVersion::V2_2)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_read_plane_matches_across_coalesce_gaps( + #[case] version: LanceFileVersion, + ) { + assert!(COALESCE_TEST_GAPS.is_sorted(), "{COALESCE_TEST_GAPS:?}"); + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let source = &lazy_index(&index).storage; + let mut planes = vec![0, 1, 2]; + if source.sign_bounds() == SignBounds::Lazy { + planes.push(SIGN_BOUNDS_PLANE); + } + // A v2.0 file reads each column page with one request per range + // group and no page initialization, so its requests are exact. + let counts_requests = version == LanceFileVersion::V2_0; + let requests = |stats: &IoStats| stats.snapshot().iops; + for (page_rows, resident) in [ + (None, false), + (Some(RESIDENT_TEST_PAGE_ROWS), false), + (Some(RESIDENT_TEST_PAGE_ROWS), true), + ] { + let reader = rewrite_rq_storage(source, version, page_rows).await; + let storage = rq_storage_over(source, &reader, resident); + let straddled = (0..source.num_partitions()) + .any(|partition| straddles_page(&reader, source.ivf().row_range(partition))); + assert_eq!(straddled, page_rows.is_some(), "pages={page_rows:?}"); + for &plane in &planes { + let layout = format!( + "{version:?} pages={page_rows:?} resident={resident} plane={plane}" + ); + // Requests at gap 0 and at the covering gap, summed over + // the partitions and row selections. + let (mut separate, mut covering) = (0, 0); + for partition in 0..source.num_partitions() { + let partition_rows = source.partition_size(partition); + let selections: [Vec; 2] = [ + resident_test_rows(partition_rows), + (0..partition_rows as u32) + .step_by(COALESCE_TEST_ROW_STEP) + .collect(), + ]; + let whole = IoStats::new(); + if counts_requests { + storage + .read_plane(partition, plane, None, Some(whole.clone())) + .await + .unwrap(); + } + for rows in selections { + let context = + format!("{layout} partition={partition} rows={}", rows.len()); + let expected = storage + .read_plane(partition, plane, Some(rows.clone()), None) + .await + .unwrap(); + let mut gap_requests = Vec::with_capacity(COALESCE_TEST_GAPS.len()); + for gap in COALESCE_TEST_GAPS { + let stats = IoStats::new(); + let actual = storage + .read_plane_with_coalesce_gap( + partition, + plane, + Some(rows.clone()), + Some(stats.clone()), + Some(gap), + ) + .await + .unwrap(); + assert_eq!(actual, expected, "{context} gap={gap}"); + gap_requests.push(requests(&stats)); + } + if counts_requests { + let context = format!( + "{context} requests={gap_requests:?} whole={}", + requests(&whole) + ); + assert!( + gap_requests.windows(2).all(|pair| pair[1] <= pair[0]), + "{context}" + ); + let last = *gap_requests.last().unwrap(); + assert!(last <= requests(&whole), "{context}"); + separate += gap_requests[0]; + covering += last; + } + } + } + if counts_requests { + assert!(separate > covering, "{layout}: {separate} vs {covering}"); + } + } + } + } + + /// The lazy scan returns the eager results whatever gap its sparse + /// reads of the origin file merge row runs within, whether it reads + /// every plane by selected rows or chooses by cost. Every gap counts + /// the same origin rows, and reading them, a wider gap makes no more + /// origin requests, for no fewer bytes; `auto` on a high-latency + /// origin reads as [`HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES`] does. + #[rstest] + #[case::rq7_l2(7, DistanceType::L2)] + #[case::rq9_dot(9, DistanceType::Dot)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_origin_gap_parity( + #[case] bits: u8, + #[case] distance_type: DistanceType, + ) { + assert!(COALESCE_TEST_GAPS.is_sorted(), "{COALESCE_TEST_GAPS:?}"); + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), bits, distance_type).await; + let vectors = batch["vector"].as_fixed_size_list(); + let (dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let queries: Vec = [vectors.value(0), vectors.value(777)] + .into_iter() + .flat_map(|key| { + [(4, 10), (16, 100), (LAZY_PARTITIONS, 1000)] + .map(|(nprobes, k)| lazy_test_query(key.clone(), k, nprobes)) + }) + .collect(); + let filter: Arc = Arc::new(NoFilter); + let gaps = std::iter::once(LazyOriginGap::Auto) + .chain(COALESCE_TEST_GAPS.map(LazyOriginGap::Bytes)); + for origin_gap in gaps { + for dense in [DenseGatherMode::Sparse, DenseGatherMode::Cost] { + let config = LayeredLazyConfig { + enabled: true, + dense, + promote: LazyPromotion::Off, + origin_gap, + ..Default::default() + }; + let label = format!("{distance_type:?} bits={bits} {origin_gap:?} {dense:?}"); + let mut origin_row_reads = 0; + for query in &queries { + let context = + format!("{label} nprobes={} k={}", query.minimum_nprobes, query.k); + let stats = + assert_lazy_matches_eager(&index, query, &filter, config, &context) + .await; + origin_row_reads += stats.origin_row_reads; + } + // Nothing is cached, so every sparse read goes to the origin. + if dense == DenseGatherMode::Sparse { + assert!(origin_row_reads > 0, "{label}"); + } + } + } + + let label = format!("{distance_type:?} bits={bits}"); + let mut reads = Vec::with_capacity(COALESCE_TEST_GAPS.len()); + for gap in COALESCE_TEST_GAPS { + reads.push( + origin_sparse_reads(&index, &queries, LazyOriginGap::Bytes(gap), &label).await, + ); + } + let requests = reads.iter().map(|read| read.requests).collect::>(); + let bytes = reads.iter().map(|read| read.bytes).collect::>(); + let rows = reads.iter().map(|read| read.rows).collect::>(); + // Every gap gathers the same rows, so the rows counted are the + // denominator of each gap's byte amplification. + let context = format!("{label} requests={requests:?} bytes={bytes:?} rows={rows:?}"); + assert!(rows[0] > 0, "{context}"); + assert!( + rows.iter().all(|&gap_rows| gap_rows == rows[0]), + "{context}" + ); + // Gap 0 merges only touching row runs, and `u64::MAX` every run + // of a read. + assert!( + requests.windows(2).all(|pair| pair[1] <= pair[0]), + "{context}" + ); + assert!(bytes.windows(2).all(|pair| pair[1] >= pair[0]), "{context}"); + // A wider gap merges row runs that are apart, reading the bytes + // between them, so it makes fewer requests exactly when it reads + // more bytes. Every sparse read of the rq9 Dot data is one run per + // column, so only the rq7 L2 data must merge. + let last = requests.len() - 1; + assert_eq!( + requests[0] > requests[last], + bytes[0] < bytes[last], + "{context}" + ); + if distance_type == DistanceType::L2 { + assert!(requests[0] > requests[last], "{context}"); + } + + let (store, index_dir) = index_files(&dataset).await; + let high = + open_with_origin_latency(&dataset, store, index_dir, OriginLatencyClass::High) + .await; + let auto = origin_sparse_reads(&high, &queries, LazyOriginGap::Auto, &label).await; + let fixed = origin_sparse_reads( + &high, + &queries, + LazyOriginGap::Bytes(HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES), + &label, + ) + .await; + assert!(auto.requests > 0, "{label}: {auto:?}"); + assert_eq!(auto, fixed, "{label}"); + assert_eq!(auto.rows, rows[0], "{label}: {auto:?}"); + } + + /// What the lazy scan's sparse reads of the origin file added up to. + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] + struct OriginSparseReads { + /// Requests after coalescing and splitting. + requests: u64, + /// Bytes read, the gaps between row runs included. + bytes: u64, + /// Rows gathered. + rows: u64, + } + + /// The origin requests, bytes and rows of the lazy scan's sparse + /// reads over `queries` on `index` with `origin_gap`, each query's + /// results checked against the eager scan's. Every lazy probe is + /// gathered by selected rows once every earlier probe was scored, so + /// the rows each gather reads are the same for every gap. Nothing is + /// cached, so the origin serves each survivor's row of both ex planes. + async fn origin_sparse_reads( + index: &Arc, + queries: &[Query], + origin_gap: LazyOriginGap, + label: &str, + ) -> OriginSparseReads { + let filter: Arc = Arc::new(NoFilter); + let config = LayeredLazyConfig { + enabled: true, + window: 0, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Off, + origin_gap, + dense_to_eager: DenseToEager::Off, + far_window: 0, + ..Default::default() + }; + // A gather reads its survivors' rows of the high and the low plane. + const EX_PLANES: u64 = 2; + let mut reads = OriginSparseReads::default(); + for query in queries { + let context = format!( + "{label} {origin_gap:?} nprobes={} k={}", + query.minimum_nprobes, query.k + ); + let stats = + assert_lazy_matches_eager(index, query, &filter, config, &context).await; + let survivors = stats.rows_fetched.iter().sum::(); + assert_eq!( + stats.origin_sparse_rows, + EX_PLANES * survivors, + "{context}: {stats:?}" + ); + reads.requests += stats.origin_sparse_requests; + reads.bytes += stats.origin_sparse_bytes; + reads.rows += stats.origin_sparse_rows; + } + reads + } + + /// The lazy origin gap is the environment's setting resolved for the + /// index's origin, when the index opens and when it is reconstructed + /// from the cached state; a replaced setting resolves the same way. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_origin_gap_resolves_at_open() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Resident).await; + let ivf = lazy_index(&index); + let block_size = dataset.object_store.block_size() as u64; + let setting = LayeredLazyConfig::from_env().unwrap().origin_gap; + let expected = setting.resolve(ivf.origin_latency(), block_size); + assert_eq!(ivf.lazy_origin_gap_bytes(), expected); + + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + assert!( + dataset.index_cache.get_with_key(&state_key).await.is_some(), + "the reopen must reconstruct from the cached state" + ); + let reopened = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + assert_eq!(lazy_index(&reopened).lazy_origin_gap_bytes(), expected); + + let auto = match ivf.origin_latency() { + OriginLatencyClass::High => HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, + OriginLatencyClass::Low => block_size, + }; + for (origin_gap, expected) in [ + (LazyOriginGap::Auto, auto), + (LazyOriginGap::Bytes(0), 0), + (LazyOriginGap::Bytes(u64::MAX), u64::MAX), + ] { + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig { + origin_gap, + ..Default::default() + }); + assert_eq!(ivf.lazy_origin_gap_bytes(), expected, "{origin_gap:?}"); + } + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + } + + /// The object store holding the files of `dataset`'s index, and their + /// directory. + async fn index_files(dataset: &Dataset) -> (Arc, object_store::path::Path) { + let index = dataset.load_indices().await.unwrap()[0].clone(); + ( + dataset.object_store.clone(), + dataset.indice_files_dir(&index).unwrap(), + ) + } + + /// A copy of the files of `dataset`'s index in an in-memory store + /// under an `s3` URL, which `auto` makes a high-latency origin, and + /// their directory there. + async fn cloud_index_files( + dataset: &Dataset, + ) -> (Arc, object_store::path::Path) { + let (local, local_dir) = index_files(dataset).await; + let uuid = dataset.load_indices().await.unwrap()[0].uuid.to_string(); + let store = Arc::new(ObjectStore::new( + Arc::new(object_store::memory::InMemory::new()), + url::Url::parse("s3://bucket/").unwrap(), + None, + None, + false, + true, + lance_io::object_store::DEFAULT_CLOUD_IO_PARALLELISM, + lance_io::object_store::DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + let dir = object_store::path::Path::from("indices"); + for file in [lance_index::INDEX_FILE_NAME, INDEX_AUXILIARY_FILE_NAME] { + let bytes = local + .read_one_all(&local_dir.clone().join(uuid.as_str()).join(file)) + .await + .unwrap(); + store + .put(&dir.clone().join(uuid.as_str()).join(file), &bytes) + .await + .unwrap(); + } + (store, dir) + } + + /// Open the index of `dataset` from its files in `index_dir` of + /// `store`, over the dataset's caches. + async fn open_rq_index( + dataset: &Dataset, + store: Arc, + index_dir: object_store::path::Path, + ) -> IvfRq { + let index = dataset.load_indices().await.unwrap()[0].clone(); + IvfRq::try_new( + store, + index_dir, + index.uuid, + None, + &dataset.metadata_cache, + dataset.index_cache.for_index(&index.uuid, None), + index.file_size_map(), + IvfOpenContext::default(), + ) + .await + .unwrap() + } + + /// [`open_rq_index`], as if the index's origin were of `class`. + async fn open_with_origin_latency( + dataset: &Dataset, + store: Arc, + index_dir: object_store::path::Path, + class: OriginLatencyClass, + ) -> Arc { + let ivf = open_rq_index(dataset, store, index_dir) + .await + .with_origin_latency_for_test(class) + .unwrap(); + Arc::new(ivf) + } + + /// Hold the high and low planes of every non-empty partition where + /// `tiers` places them: on the persistent tier only for + /// [`CacheTier::Local`], nowhere for [`CacheTier::Absent`]. Nothing + /// else stays cached. + async fn lay_out_ex_planes( + index: &Arc, + tiered: &TieredPlaneTestBackend, + tiers: impl Fn(usize) -> [CacheTier; 2], + ) { + let ivf = lazy_index(index); + tiered.clear().await; + for partition in 0..LAZY_PARTITIONS { + if ivf.storage.partition_size(partition) == 0 { + continue; + } + for (plane, tier) in [1u8, 2].into_iter().zip(tiers(partition)) { + if tier == CacheTier::Local { + ivf.storage + .load_plane_entry(partition, plane, &ivf.index_cache, None) + .await + .unwrap(); + } + } + } + tiered.ram.clear().await; + } + + /// Background promotions are off unless set, and `auto` is off for an + /// index whose origin latency class is high and after one sparse read + /// otherwise. The index resolves the environment's policy when it + /// opens and when it is reconstructed from the cached state, and + /// resolves a replaced class or setting the same way. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_promotion_resolves_at_open() { + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Resident).await; + let setting = LayeredLazyConfig::from_env().unwrap().promote; + let ivf = lazy_index(&index); + let expected = setting.resolve(ivf.origin_latency()); + assert_ne!(expected, LazyPromotion::Auto); + assert_eq!(ivf.layered_lazy_config().promote, expected); + + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + assert!( + dataset.index_cache.get_with_key(&state_key).await.is_some(), + "the reopen must reconstruct from the cached state" + ); + let reopened = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let reconstructed = lazy_index(&reopened); + assert_eq!(reconstructed.layered_lazy_config().promote, expected); + + let (store, index_dir) = index_files(&dataset).await; + for class in [OriginLatencyClass::Low, OriginLatencyClass::High] { + let auto = match class { + OriginLatencyClass::Low => LazyPromotion::Background { reads: 1 }, + OriginLatencyClass::High => LazyPromotion::Off, + }; + let index = + open_with_origin_latency(&dataset, store.clone(), index_dir.clone(), class) + .await; + let ivf = lazy_index(&index); + assert_eq!(ivf.origin_latency(), class); + assert_eq!(ivf.storage.origin_latency(), class); + let promote = ivf.layered_lazy_config().promote; + assert_eq!(promote, setting.resolve(class), "{class}"); + for promote in [ + LazyPromotion::Auto, + LazyPromotion::Off, + LazyPromotion::Background { reads: 2 }, + ] { + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig { + promote, + ..Default::default() + }); + let expected = if promote == LazyPromotion::Auto { + auto + } else { + promote + }; + let resolved = ivf.layered_lazy_config().promote; + assert_eq!(resolved, expected, "{class} {promote}"); + } + } + } + + /// With `auto` promotions, an index read from a cloud object store + /// gathers sparse rows from it without ever promoting a plane: no read + /// reaches the store once a query has returned, and results are the + /// eager scan's. The same gathers on a local file, whose class is + /// low, promote planes after one read. Both indexes share one ungated + /// backend that would admit promoted planes, emptied before every + /// query so that each sparse gather reads the origin. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_promotion_follows_origin_latency() { + const QUERIES: usize = 3; + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, _, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::ColdUngated).await; + let tiered = tiered.unwrap(); + let (cloud_store, cloud_dir) = cloud_index_files(&dataset).await; + let cloud = open_rq_index(&dataset, cloud_store.clone(), cloud_dir).await; + // `auto` makes the cloud store's class high unless the environment + // sets it; pin it high for the promotions below. + let class = + OriginLatencyClass::resolve(origin_latency_setting().unwrap(), &cloud_store); + assert_eq!(cloud.origin_latency(), class); + let cloud = cloud + .with_origin_latency_for_test(OriginLatencyClass::High) + .unwrap(); + let cloud: Arc = Arc::new(cloud); + let (local_store, local_dir) = index_files(&dataset).await; + let local = open_with_origin_latency( + &dataset, + local_store.clone(), + local_dir, + OriginLatencyClass::Low, + ) + .await; + + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + let queries: Vec = (0..QUERIES) + .map(|row| lazy_test_query(vectors.value(row * 101), 100, LAZY_PARTITIONS)) + .collect(); + lazy_index(&local).set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let mut expected = Vec::with_capacity(QUERIES); + for query in &queries { + let eager = search_global(&local, query, filter.clone()).await.unwrap(); + expected.push(result_bits(&eager)); + } + // Promotions are off unless set; `auto` resolves by class. + let config = LayeredLazyConfig { + enabled: true, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Auto, + ..Default::default() + }; + for (index, store, class) in [ + (&cloud, &cloud_store, OriginLatencyClass::High), + (&local, &local_store, OriginLatencyClass::Low), + ] { + let ivf = lazy_index(index); + ivf.set_layered_lazy_config_for_test(config); + let promote = ivf.layered_lazy_config().promote; + assert_eq!(promote, LazyPromotion::Auto.resolve(class), "{class}"); + layered_stats::snapshot_and_reset(); + let mut late_reads = 0; + for (query, expected) in queries.iter().zip(&expected) { + tiered.clear().await; + let result = search_global(index, query, filter.clone()).await.unwrap(); + assert_eq!(&result_bits(&result), expected, "{class}"); + store.io_stats_incremental(); + wait_for_promotions().await; + late_reads += store.io_stats_incremental().read_iops; + } + let stats = layered_stats::snapshot_and_reset(); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let context = format!("{class} {stats:?}"); + assert_eq!(stats.lazy_queries, QUERIES as u64, "{context}"); + assert_eq!(stats.needed_not_fetched, 0, "{context}"); + assert!(stats.origin_row_reads > 0, "{context}"); + match class { + OriginLatencyClass::High => { + assert_eq!( + (stats.promotions_issued, stats.promotion_bytes), + (0, 0), + "{context}" + ); + assert_eq!(late_reads, 0, "{context}"); + } + OriginLatencyClass::Low => { + assert!(stats.promotions_issued > 0, "{context}"); + assert!(stats.promotion_bytes > 0, "{context}"); + } + } + } + } + + /// Under `origin`, staging routes a predicted-dense probe only when + /// the index's origin is high latency and no cache tier holds its + /// high or low plane; `all` routes it wherever its planes are, and + /// `off` never does. Before every query the partitions hold both ex + /// planes on the persistent tier, only the high plane there, or + /// neither anywhere, and nothing in RAM. One query's `k` exceeds every + /// probed row, so every probe is predicted dense; the others' first + /// probe alone fills the heap, so only that one is. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_dense_to_eager_follows_plane_tiers() { + const SMALL_K: usize = 10; + const SMALL_K_QUERIES: usize = 4; + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, _, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::ColdUngated).await; + let tiered = tiered.unwrap(); + let (store, index_dir) = index_files(&dataset).await; + let (low, high) = (OriginLatencyClass::Low, OriginLatencyClass::High); + let low_index = + open_with_origin_latency(&dataset, store.clone(), index_dir.clone(), low).await; + let high_index = open_with_origin_latency(&dataset, store, index_dir, high).await; + let (local, absent) = (CacheTier::Local, CacheTier::Absent); + let layout = |partition: usize| match partition % 3 { + 0 => [local, local], + 1 => [local, absent], + _ => [absent, absent], + }; + + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + let every_row = lazy_test_query(vectors.value(0), LAZY_ROWS + 1, LAZY_PARTITIONS); + let mut queries = vec![every_row]; + queries.extend(queries_filled_by_first_probe( + &low_index, + vectors, + SMALL_K, + LAZY_PARTITIONS, + SMALL_K_QUERIES, + )); + lazy_index(&low_index).set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let mut expected = Vec::with_capacity(queries.len()); + for query in &queries { + let eager = search_global(&low_index, query, filter.clone()) + .await + .unwrap(); + expected.push(result_bits(&eager)); + } + + // The layout is what the tier checks of a high-latency index see. + lay_out_ex_planes(&low_index, &tiered, layout).await; + let ivf = lazy_index(&high_index); + let cache = &ivf.index_cache; + for partition in 0..LAZY_PARTITIONS { + if ivf.storage.partition_size(partition) > 0 { + let tiers = [ + ivf.storage.plane_tier(partition, 1, cache).await, + ivf.storage.plane_tier(partition, 2, cache).await, + ]; + assert_eq!(tiers, layout(partition), "partition {partition}"); + } + } + + for mode in [DenseToEager::Off, DenseToEager::Origin, DenseToEager::All] { + for (index, class) in [(&low_index, low), (&high_index, high)] { + let ivf = lazy_index(index); + for (position, (query, eager)) in queries.iter().zip(&expected).enumerate() { + lay_out_ex_planes(index, &tiered, layout).await; + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig { + enabled: true, + promote: LazyPromotion::Off, + dense_to_eager: mode, + ..Default::default() + }); + layered_stats::snapshot_and_reset(); + let result = search_global(index, query, filter.clone()).await.unwrap(); + let stats = layered_stats::snapshot_and_reset(); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let context = format!("{mode} {class} query={position} {stats:?}"); + assert_eq!(&result_bits(&result), eager, "{context}"); + assert_eq!(stats.lazy_queries, 1, "{context}"); + assert_eq!(stats.needed_not_fetched, 0, "{context}"); + + let (partitions, _) = index.find_partitions(query).unwrap(); + let mut routed = [0; RANK_BUCKETS]; + let mut slow = [0; RANK_BUCKETS]; + for (rank, &partition) in partitions.values().iter().enumerate() { + let partition = partition as usize; + if ivf.storage.partition_size(partition) == 0 { + continue; + } + // A plane on no cache tier is a slow read on a slow origin only. + let reads_slow_origin = class == high && !partition.is_multiple_of(3); + let routes = match mode { + DenseToEager::Off => false, + DenseToEager::Origin => reads_slow_origin, + DenseToEager::All => true, + }; + let predicted_dense = query.k > LAZY_ROWS || rank == 0; + let bucket = layered_stats::rank_bucket(rank); + routed[bucket] += u64::from(predicted_dense && routes); + slow[bucket] += u64::from(reads_slow_origin); + } + assert_eq!(stats.dense_to_eager, routed, "{context}"); + assert_eq!(stats.s3_bound_probes, slow, "{context}"); + } + } + } + } + + /// Only a probe whose high or low plane no cache tier holds, on a + /// high-latency origin, may be gathered beyond the ordinary window, + /// and only while the far window is wider. Every other gather is + /// issued within the ordinary window: on a low-latency origin, with + /// every ex plane on the persistent tier, or with the far window no + /// wider. Row reads are slowed so that scoring lags the staging of + /// later probes, which leaves far gathers room to run ahead. On the + /// slow origin far gathers also run ahead for larger `k`, under a + /// prefilter and distance bounds, whether gathers wait for the heap to + /// fill or not, and while predicted-dense probes are routed to the + /// eager scan. Results are the eager scan's in every setting. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_far_window_by_tier() { + const PROBES: usize = 32; + const SMALL_K: usize = 10; + /// Heaps that about one probe fills, and that several probes fill. + const WIDE_KS: [usize; 2] = [100, 1000]; + const QUERIES: usize = 2; + const STAGING_STEPS: usize = 8; + const WINDOW: usize = 1; + const ROW_READ_DELAY: Duration = Duration::from_millis(5); + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, _, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::ColdUngated).await; + let tiered = tiered.unwrap(); + let (store, index_dir) = index_files(&dataset).await; + let (low, high) = (OriginLatencyClass::Low, OriginLatencyClass::High); + let low_index = + open_with_origin_latency(&dataset, store.clone(), index_dir.clone(), low).await; + let high_index = open_with_origin_latency(&dataset, store, index_dir, high).await; + let (local, absent) = (CacheTier::Local, CacheTier::Absent); + let mixed = |partition: usize| match partition % 3 { + 0 => [local, local], + 1 => [local, absent], + _ => [absent, absent], + }; + + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + let queries = + queries_filled_by_first_probe(&low_index, vectors, SMALL_K, PROBES, QUERIES); + lazy_index(&low_index).set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let mut expected = Vec::with_capacity(QUERIES); + for query in &queries { + let eager = search_global(&low_index, query, filter.clone()) + .await + .unwrap(); + expected.push(result_bits(&eager)); + } + // The same queries with heaps that fill after one probe or + // several, without and with a prefilter, and without and with + // distance bounds, which keep the heap from filling. + let (prefilter_name, prefilter) = lazy_test_filters() + .into_iter() + .find(|(name, _)| *name == "alternating") + .unwrap(); + let mut wide = Vec::new(); + for query in &queries { + for k in WIDE_KS { + let mut query = query.clone(); + query.k = k; + for (filter_name, filter) in [ + ("none", filter.clone()), + (prefilter_name, prefilter.clone()), + ] { + let eager = search_global(&low_index, &query, filter.clone()) + .await + .unwrap(); + let bounded = with_bounds(&query, &eager).unwrap(); + let bounded_eager = search_global(&low_index, &bounded, filter.clone()) + .await + .unwrap(); + wide.push(( + query.clone(), + filter_name, + filter.clone(), + result_bits(&eager), + )); + wide.push((bounded, filter_name, filter, result_bits(&bounded_eager))); + } + } + } + + let far = LayeredLazyConfig { + enabled: true, + window: WINDOW, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Off, + far_window: PROBES, + ..Default::default() + }; + let far_off = LayeredLazyConfig { + far_window: WINDOW, + ..far + }; + tiered.set_row_read_delay(ROW_READ_DELAY); + // (setting, index, whether every ex plane is on the persistent + // tier, config, whether far gathers may run ahead) + for (setting, index, all_local, config, runs_ahead) in [ + ("low", &low_index, false, far, false), + ("high, planes local", &high_index, true, far, false), + ("high, far window off", &high_index, false, far_off, false), + ("high", &high_index, false, far, true), + ] { + let ivf = lazy_index(index); + ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); + let mut far_early_issues = 0; + let mut staleness_max = 0; + for (position, (query, eager)) in queries.iter().zip(&expected).enumerate() { + if all_local { + lay_out_ex_planes(index, &tiered, |_| [local, local]).await; + } else { + lay_out_ex_planes(index, &tiered, mixed).await; + } + ivf.set_layered_lazy_config_for_test(config); + layered_stats::snapshot_and_reset(); + let result = search_global(index, query, filter.clone()).await.unwrap(); + let stats = layered_stats::snapshot_and_reset(); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let context = format!("{setting} query={position} {stats:?}"); + assert_eq!(&result_bits(&result), eager, "{context}"); + assert_eq!(stats.lazy_queries, 1, "{context}"); + assert_eq!(stats.needed_not_fetched, 0, "{context}"); + assert_lazy_chain_timing(&stats, &context); + // A probe on the slow origin issues at most one far gather. + let slow_probes: u64 = stats.s3_bound_probes.iter().sum(); + assert!(stats.far_early_issues <= slow_probes, "{context}"); + if !runs_ahead { + assert_eq!( + (stats.far_early_issues, stats.far_permit_waits), + (0, 0), + "{context}" + ); + assert!(stats.staleness_max <= WINDOW as u64, "{context}"); + } + far_early_issues += stats.far_early_issues; + staleness_max = staleness_max.max(stats.staleness_max); + } + ivf.set_lazy_prepare_parallelism_for_test(0); + if runs_ahead { + assert!(far_early_issues > 0, "{setting}"); + assert!(staleness_max > WINDOW as u64, "{setting}: {staleness_max}"); + } + } + + // The wider queries run ahead on the slow origin too: with gathers + // issued before the heap fills or held until it does, and with the + // predicted-dense probes whose planes are on the slow origin + // loaded by the eager scan. Under the cost policy every probe of an + // unfiltered, unbounded `k = 1000` query is predicted dense, so + // that setting routes some; the sparse policy routes none. + // (setting, config, whether it routes probes to the eager scan) + let wide_settings = [ + ("sparse", far, false), + ( + "sparse, held until the heap fills", + LayeredLazyConfig { + eager_before_full: false, + ..far + }, + false, + ), + ( + "cost, dense probes on the slow origin routed", + LayeredLazyConfig { + dense: DenseGatherMode::Cost, + dense_to_eager: DenseToEager::Origin, + ..far + }, + true, + ), + ]; + let ivf = lazy_index(&high_index); + ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); + for (setting, config, routes) in wide_settings { + let mut far_early_issues = 0; + let mut routed = 0; + for (position, (query, filter_name, filter, eager)) in wide.iter().enumerate() { + lay_out_ex_planes(&high_index, &tiered, mixed).await; + ivf.set_layered_lazy_config_for_test(config); + layered_stats::snapshot_and_reset(); + let result = search_global(&high_index, query, filter.clone()) + .await + .unwrap(); + let stats = layered_stats::snapshot_and_reset(); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let context = format!( + "{setting} query={position} k={} filter={filter_name} bounds={} {stats:?}", + query.k, + query.upper_bound.is_some() + ); + assert_eq!(&result_bits(&result), eager, "{context}"); + assert_eq!(stats.lazy_queries, 1, "{context}"); + assert_eq!(stats.needed_not_fetched, 0, "{context}"); + assert_lazy_chain_timing(&stats, &context); + let slow_probes: u64 = stats.s3_bound_probes.iter().sum(); + assert!(stats.far_early_issues <= slow_probes, "{context}"); + far_early_issues += stats.far_early_issues; + routed += stats.dense_to_eager.iter().sum::(); + } + assert!(far_early_issues > 0, "{setting}"); + assert_eq!(routed > 0, routes, "{setting}: {routed} probes routed"); + } + ivf.set_lazy_prepare_parallelism_for_test(0); + tiered.set_row_read_delay(Duration::ZERO); + } + + /// Gathers beyond the ordinary window share their index's permits + /// across queries: concurrent queries hold at most one at a time with + /// one permit and at most four with four, and every permit is back + /// once the queries finish or are dropped mid-scan. The pool lives in + /// the state the cached index is reconstructed from. Nothing is + /// cached, so every probe reads its ex planes from the high-latency + /// origin, and row reads are slowed so that gathers overlap. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_far_permits_bound_in_flight() { + const PROBES: usize = 32; + const SMALL_K: usize = 10; + const QUERIES: usize = 3; + const STAGING_STEPS: usize = 8; + const ROW_READ_DELAY: Duration = Duration::from_millis(5); + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, opened, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::ColdUngated).await; + let tiered = tiered.unwrap(); + let one_permit = LayeredLazyConfig { + enabled: true, + window: 0, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Off, + far_window: PROBES, + far_inflight: 1, + ..Default::default() + }; + + // A reconstruction of the index shares its pool. + let pool = lazy_index(&opened) + .storage + .lazy_far_permits(&one_permit) + .unwrap(); + let held = pool.try_acquire().unwrap(); + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let frag_reuse_uuid = dataset.frag_reuse_index_uuid().await; + let state_key = + crate::index::IvfIndexStateCacheKey::new(&uuid, frag_reuse_uuid.as_ref()); + assert!( + dataset.index_cache.get_with_key(&state_key).await.is_some(), + "the reopen must reconstruct from the cached state" + ); + let reopened = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let shared = lazy_index(&reopened) + .storage + .lazy_far_permits(&one_permit) + .unwrap(); + assert_eq!((shared.size(), shared.in_flight()), (1, 1)); + assert!(shared.try_acquire().is_none()); + drop(held); + assert_eq!(shared.in_flight(), 0); + + let (store, index_dir) = index_files(&dataset).await; + let high = OriginLatencyClass::High; + let index = open_with_origin_latency(&dataset, store, index_dir, high).await; + let ivf = lazy_index(&index); + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + let queries = queries_filled_by_first_probe(&index, vectors, SMALL_K, PROBES, QUERIES); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let mut expected = Vec::with_capacity(QUERIES); + for query in &queries { + let eager = search_global(&index, query, filter.clone()).await.unwrap(); + expected.push(result_bits(&eager)); + } + + tiered.set_row_read_delay(ROW_READ_DELAY); + ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); + for far_inflight in [1, 4] { + let config = LayeredLazyConfig { + far_inflight, + ..one_permit + }; + let permits = ivf.storage.lazy_far_permits(&config).unwrap(); + ivf.set_layered_lazy_config_for_test(config); + tiered.clear().await; + layered_stats::snapshot_and_reset(); + let results = futures::future::try_join_all( + queries + .iter() + .map(|query| search_global(&index, query, filter.clone())), + ) + .await + .unwrap(); + let stats = layered_stats::snapshot_and_reset(); + let context = format!("far_inflight={far_inflight} {stats:?}"); + for (result, expected) in results.iter().zip(&expected) { + assert_eq!(&result_bits(result), expected, "{context}"); + } + assert_eq!(stats.lazy_queries, QUERIES as u64, "{context}"); + assert_eq!(stats.needed_not_fetched, 0, "{context}"); + assert_lazy_chain_timing(&stats, &context); + assert!(stats.far_early_issues > 0, "{context}"); + // Observed as each far gather takes its permit. + let most = far_inflight as u64; + assert!((1..=most).contains(&stats.far_in_flight_max), "{context}"); + if far_inflight == 1 { + // The queries' far gathers queue for the one permit, and + // the queue is timed. + assert!(stats.far_permit_waits > 0, "{context}"); + assert!(stats.far_permit_wait_ns > 0, "{context}"); + } else { + assert!(stats.far_in_flight_max > 1, "{context}"); + } + assert_eq!(permits.in_flight(), 0, "{context}"); + } + + // A query dropped while a far gather holds a permit returns it. + let permits = ivf.storage.lazy_far_permits(&one_permit).unwrap(); + ivf.set_layered_lazy_config_for_test(one_permit); + tiered.clear().await; + let permit_taken = async { + while permits.in_flight() == 0 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }; + let finished = tokio::select! { + result = search_global(&index, &queries[0], filter.clone()) => Some(result), + () = permit_taken => None, + }; + assert!( + finished.is_none(), + "the query finished before a far gather took a permit" + ); + for _ in 0..500 { + if permits.in_flight() == 0 { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!( + permits.in_flight(), + 0, + "a dropped query must return its permits" + ); + tiered.set_row_read_delay(Duration::ZERO); + ivf.set_lazy_prepare_parallelism_for_test(0); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + } } #[rstest] diff --git a/rust/lance/src/index/vector/ivf/v2/lazy_full.rs b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs index 3a447fbc392..e4a5e29b017 100644 --- a/rust/lance/src/index/vector/ivf/v2/lazy_full.rs +++ b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs @@ -33,30 +33,46 @@ //! whether or not the producer is polling the buffer that holds it; see //! `spawn_scan_step`. //! +//! On a high-latency origin, a probe whose high or low plane no cache tier +//! held when it was staged reads it from the origin, a round trip that the +//! ordinary window cannot hide. Its gather may be issued up to +//! `LANCE_RQ_LAZY_FAR_WINDOW` probes ahead of scoring instead, which only +//! takes an older threshold. A gather beyond the ordinary window holds a +//! permit from a pool shared by the index's queries until its reads return; +//! one that finds no free permit is issued at the ordinary window at the +//! latest. No gather is therefore issued later than without the far window, +//! and the next probe to score, always within the ordinary window, never +//! waits for a permit. +//! //! Probes whose ex planes are both resident, probes that cannot gate on the //! lower bound, and probes that `k` and the partition sizes predict to be //! gathered whole (`LANCE_RQ_LAZY_DENSE_TO_EAGER`, only for queries without a //! prefilter or upper distance bound, see `LazyDenseForecast`) are scored by -//! the eager scan in their probe position. A predicted-dense probe is loaded -//! at staging as the eager scan loads it, with its planes read together and -//! as many probes in flight as the eager scan prepares, instead of its sign -//! plane first and its ex planes within the gather window. +//! the eager scan in their probe position. By default a predicted-dense probe +//! is routed only when the index's origin is high latency and no cache tier +//! holds one of its ex planes; staging tells from one tier check per probe. +//! A routed probe is loaded at staging as the eager scan loads it, with its +//! planes read together and as many probes in flight as the eager scan +//! prepares, instead of its sign plane first and its ex planes within the +//! gather window. use std::collections::BinaryHeap; use std::future::Future; use std::ops::Range; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Instant; use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array}; use futures::StreamExt; use futures::prelude::stream; +use lance_core::cache::CacheTier; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::{Error, Result}; use lance_index::metrics::{IndexTiming, MetricsCollector}; use lance_index::prefilter::PreFilter; -use lance_index::vector::bq::layered::{PlaneKey, RQPrecision}; +use lance_index::vector::bq::layered::RQPrecision; use lance_index::vector::bq::layered_stats; use lance_index::vector::bq::storage::{ ExRows, RabitQuantizationStorage, SignStage, StagePruneCounts, SurvivorRow, @@ -65,13 +81,14 @@ use lance_index::vector::bq::storage::{ use lance_index::vector::graph::OrderedNode; use lance_index::vector::quantizer::Quantization; use lance_index::vector::storage::{ - DenseGatherMode, DistanceCalculatorOptions, ExIndex, GatherPlan, GatheredEx, LayeredExLayout, - LayeredLazyConfig, LazyPromotionTicket, PlaneSource, QueryScratchPool, RabitRawQueryContext, - VectorStore, origin_reads_whole_plane, plan_plane_gather, + DenseGatherMode, DenseToEager, DistanceCalculatorOptions, ExIndex, GatherPlan, GatheredEx, + LayeredExLayout, LayeredLazyConfig, LazyFarPermits, LazyPromotionTicket, OriginLatencyClass, + PlaneSource, QueryScratchPool, RabitRawQueryContext, VectorStore, origin_reads_whole_plane, + plan_plane_gather, }; use lance_index::vector::v3::subindex::IvfSubIndex; use lance_index::vector::{ApproxMode, Query}; -use tokio::sync::{mpsc, watch}; +use tokio::sync::{OwnedSemaphorePermit, mpsc, watch}; use super::{ GLOBAL_TOPK_INLINE_HEAP_LEN, IVFIndex, LAYERED_LAZY_CONFIG, PartitionEntry, @@ -121,6 +138,9 @@ struct LazySignProbe { stage: SignStage, /// Prefilter acceptance by partition offset; `None` accepts every row. accept: Option>, + /// Where the cache held the high and low planes when the probe was + /// staged; see `ex_plane_tiers`. + tiers: [CacheTier; 2], _in_flight: PreparedPartitionGuard, } @@ -159,8 +179,8 @@ struct LazyDenseForecast { /// Whether the gathers issued once the heap fills are expected to read /// whole planes anyway, so issuing them before it fills reads no more. dense_at_fill: bool, - /// Whether probes predicted to read their ex planes whole are loaded and - /// scored by the eager scan instead; see [`Self::new`]. + /// Whether probes predicted to read their ex planes whole may be loaded + /// and scored by the eager scan instead; see [`Self::new`]. route_dense: bool, } @@ -207,7 +227,7 @@ impl LazyDenseForecast { Self { fill_rank, dense_at_fill, - route_dense: config.dense_to_eager + route_dense: config.dense_to_eager != DenseToEager::Off && every_row_candidate && config.dense != DenseGatherMode::Sparse, } @@ -220,16 +240,65 @@ impl LazyDenseForecast { rank < self.fill_rank } - /// Whether the probe at `rank` is loaded and scored by the eager scan - /// because it is expected to read its ex planes whole: it is certain to - /// be dense, or every gather is expected to be dense once the heap fills. - fn routes_to_eager(&self, rank: usize) -> bool { + /// Whether the probe at `rank` may be loaded and scored by the eager + /// scan because it is expected to read its ex planes whole: it is + /// certain to be dense, or every gather is expected to be dense once the + /// heap fills. Staging routes it when `LANCE_RQ_LAZY_DENSE_TO_EAGER` does + /// for its planes' tiers ([`DenseToEager::routes`]). + fn routable(&self, rank: usize) -> bool { self.route_dense && (self.certain_dense(rank) || self.dense_at_fill) } } +/// How far the gathers of probes that read an ex plane from a high-latency +/// origin may run ahead of scoring; see `LANCE_RQ_LAZY_FAR_WINDOW`. +struct LazyFarIssue { + /// Their staleness window, wider than the ordinary one. + window: usize, + /// Permits of the index's gathers beyond the ordinary window. + permits: LazyFarPermits, +} + +impl LazyFarIssue { + /// Wait until the gather of the probe at `rank`, beyond the ordinary + /// `window` and without a free permit, may be issued: with the first + /// permit that frees up, which it then holds, or without one once + /// scoring brings the probe within the ordinary window, whichever comes + /// first. It is thus never issued later than without the far window. + /// Returns the progress then, the latest to prune against: the heap's + /// top never increases, so it prunes at least as many rows as the + /// progress at the far gate. + async fn wait_for_permit_or_window( + &self, + progress: &mut watch::Receiver, + window: usize, + rank: usize, + ) -> Result<(LazyProgress, Option)> { + let within_window = + |progress: &LazyProgress| progress.scored.saturating_add(window) >= rank; + // The window goes first: within it the gather needs no permit. + let permit = tokio::select! { + biased; + gate = progress.wait_for(within_window) => { + gate.map_err(|_| scan_cancelled())?; + None + } + permit = self.permits.acquire() => Some(permit?), + }; + Ok((*progress.borrow(), permit)) + } +} + /// Query-wide inputs of the gathers. struct LazyFetchContext { + /// When the scan started, which its chain timings count from. + started: Instant, + /// Whether a gather that waits on scoring progress was issued; the first + /// times `lazy_first_gather_issue_ns`. + first_gather_issued: AtomicBool, + /// `None` when every gather keeps the ordinary window: the index's origin + /// is low latency, or the far window is no wider. + far: Option, config: LayeredLazyConfig, forecast: LazyDenseForecast, upper_bound: Option, @@ -256,6 +325,25 @@ struct LazyScorer { scratch_pool: Arc, metrics: Arc, progress: Arc>, + /// When the scan started, which its chain timings count from. + started: Instant, +} + +impl LazyScorer { + /// Publish scoring progress to the gathers. The scan's first publish + /// comes from its first probe, once scored or once its rows fill the + /// heap, and is the earliest a gather waiting for the threshold or its + /// turn can be released; it is timed before any gather can see it. + fn publish(&self, update: impl FnOnce(&mut LazyProgress)) { + self.progress.send_modify(|progress| { + if progress.scored == 0 && !progress.full { + let stats = layered_stats::counters(); + stats.lazy_rank0_scored_queries.incr(); + stats.lazy_rank0_scored_ns.add_elapsed(self.started); + } + update(progress); + }); + } } /// Aborts a task of the scan when its owner is dropped: the producer when the @@ -302,19 +390,28 @@ fn spawn_scan_step( } impl IVFIndex { - /// The lazy scan settings for a newly opened index. Only layered indexes - /// read (and validate) the environment. - pub(super) fn layered_lazy_config_at_open(layered_rq: bool) -> Result { - if !layered_rq { - return Ok(LayeredLazyConfig::default()); - } - LAYERED_LAZY_CONFIG.clone().map_err(Error::invalid_input) + /// The lazy scan settings for a newly opened index whose origin is + /// `origin_latency`, with the background promotion policy resolved for + /// it. Only layered indexes read (and validate) the environment. + pub(super) fn layered_lazy_config_at_open( + layered_rq: bool, + origin_latency: OriginLatencyClass, + ) -> Result { + let mut config = if layered_rq { + LAYERED_LAZY_CONFIG.clone().map_err(Error::invalid_input)? + } else { + LayeredLazyConfig::default() + }; + config.promote = config.promote.resolve(origin_latency); + Ok(config) } /// Replace this index's lazy scan settings, avoiding the process-wide - /// environment in concurrent tests. + /// environment in concurrent tests. The promotion policy resolves for the + /// index's origin as it does at open. #[cfg(test)] - pub(crate) fn set_layered_lazy_config_for_test(&self, config: LayeredLazyConfig) { + pub(crate) fn set_layered_lazy_config_for_test(&self, mut config: LayeredLazyConfig) { + config.promote = config.promote.resolve(self.origin_latency); *self .layered_lazy .lock() @@ -450,8 +547,24 @@ impl IVFIndex { scratch_pool: self.scratch_pool.clone(), metrics: metrics.clone(), progress: Arc::new(progress_tx), + started, }; + let far = match config.active_far_window(self.origin_latency) { + Some(window) => Some(LazyFarIssue { + window, + permits: self.storage.lazy_far_permits(&config)?, + }), + None => None, + }; + // Room for every gather the widest window lets run ahead of scoring. + let fetch_slots = far + .as_ref() + .map_or(config.window, |far| far.window) + .saturating_add(LAZY_FETCH_EXTRA_SLOTS); let fetch_context = Arc::new(LazyFetchContext { + started, + first_gather_issued: AtomicBool::new(false), + far, config, forecast, upper_bound: query.upper_bound, @@ -467,6 +580,7 @@ impl IVFIndex { let stage_index = self.clone(); let fetch_index = self.clone(); let probe_start = probes.start; + let dense_to_eager = config.dense_to_eager; let result_metrics = metrics.clone(); let producer = tokio::spawn(async move { // `map` spawns a step only when `buffered` pulls it, so at most @@ -477,11 +591,16 @@ impl IVFIndex { let rank = idx - probe_start; let mut query = query.clone(); query.dist_q_c = q_c_dists.value(idx); + let dense_to_eager = if forecast.routable(rank) { + dense_to_eager + } else { + DenseToEager::Off + }; spawn_scan_step(stage_index.clone().lazy_stage_probe( rank, partitions.value(idx) as usize, query, - forecast.routes_to_eager(rank), + dense_to_eager, pre_filter.clone(), metrics.clone(), raw_query_context.clone(), @@ -493,7 +612,7 @@ impl IVFIndex { let context = fetch_context.clone(); spawn_scan_step(async move { index.lazy_fetch(staged?, &context).await }) }) - .buffered(config.window.saturating_add(LAZY_FETCH_EXTRA_SLOTS)); + .buffered(fetch_slots); futures::pin_mut!(fetched); while let Some(probe) = fetched.next().await { let failed = probe.is_err(); @@ -565,28 +684,38 @@ impl IVFIndex { } /// Load a probe's sign plane and run stage 1, or prepare it for the eager - /// scan when both ex planes are resident or when it is predicted to be - /// gathered whole (`dense_to_eager`). + /// scan when both ex planes are resident or when `dense_to_eager` routes + /// it for its planes' tiers. `dense_to_eager` is [`DenseToEager::Off`] + /// for a probe the forecast does not predict to be gathered whole. #[allow(clippy::too_many_arguments)] async fn lazy_stage_probe( self: Arc, rank: usize, partition_id: usize, query: Query, - dense_to_eager: bool, + dense_to_eager: DenseToEager, pre_filter: Arc, metrics: Arc, raw_query_context: Option>, ) -> Result> { let stats = layered_stats::counters(); - let eager = if self.ex_planes_resident(partition_id).await { - Some(if self.storage.partition_size(partition_id) == 0 { - &stats.empty - } else { - &stats.eager_resident - }) + let (tiers, eager) = if self.storage.partition_size(partition_id) == 0 { + // No ex rows to gather, so the planes need no check; see + // `ex_planes_resident`. + ([CacheTier::Resident; 2], Some(&stats.empty)) } else { - dense_to_eager.then_some(&stats.dense_to_eager) + let tiers = self.ex_plane_tiers(partition_id).await; + if self.origin_latency.reads_slow_origin(&tiers) { + stats.s3_bound_probes.incr(rank); + } + let eager = if tiers == [CacheTier::Resident; 2] { + Some(&stats.eager_resident) + } else { + dense_to_eager + .routes(self.origin_latency, &tiers) + .then_some(&stats.dense_to_eager) + }; + (tiers, eager) }; if let Some(counter) = eager { counter.incr(rank); @@ -658,10 +787,43 @@ impl IVFIndex { entry, stage, accept, + tiers, _in_flight: in_flight, }))) } + /// Where the cache holds a non-empty partition's high and low planes, + /// checked without reading them or counting as accesses. Only a + /// high-latency origin tells [`CacheTier::Local`] from + /// [`CacheTier::Absent`]: every read of a low-latency origin is local, + /// from a cache tier or the origin file, so only RAM residency is checked + /// there and a plane outside RAM is recorded as `Local`. + async fn ex_plane_tiers(&self, partition: usize) -> [CacheTier; 2] { + let classifying = Instant::now(); + let mut tiers = [CacheTier::Resident; 2]; + for (slot, plane) in [1u8, 2].into_iter().enumerate() { + tiers[slot] = match self.origin_latency { + OriginLatencyClass::High => { + self.storage + .plane_tier(partition, plane, &self.index_cache) + .await + } + OriginLatencyClass::Low => { + let key = self.storage.plane_key(partition, plane); + if self.index_cache.peek_resident_with_key(&key).await { + CacheTier::Resident + } else { + CacheTier::Local + } + } + }; + } + layered_stats::counters() + .tier_peek_ns + .add_elapsed(classifying); + tiers + } + /// Whether a partition's high and low planes are in RAM, checked without /// counting as accesses. An empty partition has no ex rows to gather, so /// it counts as resident: the lazy scan never touches its (empty) plane @@ -674,7 +836,7 @@ impl IVFIndex { for plane in [1, 2] { if !self .index_cache - .peek_resident_with_key(&PlaneKey { partition, plane }) + .peek_resident_with_key(&self.storage.plane_key(partition, plane)) .await { return false; @@ -769,19 +931,33 @@ impl IVFIndex { }; let stats = layered_stats::counters(); let rank = probe.rank; + // Only a gather that reads an ex plane from a slow origin runs further ahead. + let far = context + .far + .as_ref() + .filter(|_| self.origin_latency.reads_slow_origin(&probe.tiers)); let gating = Instant::now(); let mut progress = context.progress.clone(); let forecast = &context.forecast; - let issue = if forecast.certain_dense(rank) { + let (issue, far_permit) = if forecast.certain_dense(rank) { // The earlier probes cannot fill the heap, so the threshold is +inf at scoring. stats.certain_dense.incr(rank); - *progress.borrow() + (*progress.borrow(), None) } else { let window = context.config.window; + let gate_window = far.map_or(window, |far| far.window); + let within_gate = + |progress: &LazyProgress| progress.scored.saturating_add(gate_window) >= rank; + let window_closed = !within_gate(&progress.borrow()); + let windowing = Instant::now(); let gate = *progress - .wait_for(|progress| progress.scored.saturating_add(window) >= rank) + .wait_for(within_gate) .await .map_err(|_| scan_cancelled())?; + if window_closed { + stats.lazy_window_waits.incr(); + stats.lazy_window_wait_ns.add_elapsed(windowing); + } // An infinite threshold gathers every accepted row, so wait for a // finite one, or for this probe's turn. With eager-before-full, // stop waiting where it saves no reads: when the gathers after the @@ -799,19 +975,52 @@ impl IVFIndex { gate } else { stats.deferred_issues.incr(); + let releasing = Instant::now(); let issue = *progress .wait_for(released) .await .map_err(|_| scan_cancelled())?; + stats.lazy_release_wait_ns.add_elapsed(releasing); if !issue.full && issue.scored >= rank { stats.serial_waits.incr(); } issue }; + // A gather beyond the ordinary window needs a permit. + let (issue, far_permit) = match far { + Some(far) if issue.scored.saturating_add(window) < rank => { + let (issue, permit) = match far.permits.try_acquire() { + Some(permit) => (issue, Some(permit)), + None => { + stats.far_permit_waits.incr(); + let waiting = Instant::now(); + let waited = far + .wait_for_permit_or_window(&mut progress, window, rank) + .await?; + stats.far_permit_wait_ns.add_elapsed(waiting); + waited + } + }; + if permit.is_some() { + stats.far_early_issues.incr(); + stats + .far_in_flight_max + .observe(far.permits.in_flight() as u64); + } + (issue, permit) + } + _ => (issue, None), + }; if !issue.full && issue.scored < rank { stats.eager_before_full.incr(); } - issue + if !context.first_gather_issued.swap(true, Ordering::Relaxed) { + stats.lazy_first_gather_issue_queries.incr(); + stats + .lazy_first_gather_issue_ns + .add_elapsed(context.started); + } + (issue, far_permit) }; stats.gate_wait_ns.add_elapsed(gating); let staleness = rank.saturating_sub(issue.scored) as u64; @@ -846,10 +1055,7 @@ impl IVFIndex { let row_bytes = context.layout.row_bytes[slot]; plan.planes[slot] = if self .index_cache - .peek_resident_with_key(&PlaneKey { - partition: probe.partition_id, - plane, - }) + .peek_resident_with_key(&self.storage.plane_key(probe.partition_id, plane)) .await { PlaneSource::Resident @@ -870,6 +1076,8 @@ impl IVFIndex { context.metrics.io_stats(), ) .await?; + // The permit bounds the far gathers' reads, not their scoring. + drop(far_permit); record_gather_sources(rank, &gathered); stats.origin_row_reads.add(gathered.origin_row_reads as u64); for ticket in std::mem::take(&mut gathered.promotions) { @@ -900,10 +1108,9 @@ impl IVFIndex { .add(batch.as_ref().deep_size_of() as u64); let resident = index .index_cache - .peek_resident_with_key(&PlaneKey { - partition: ticket.partition, - plane: ticket.plane, - }) + .peek_resident_with_key( + &index.storage.plane_key(ticket.partition, ticket.plane), + ) .await; if resident { stats.promotions_completed.incr(); @@ -955,7 +1162,7 @@ impl IVFIndex { .then(|| heap.peek().map(|node| node.dist.0)) .flatten(); let scored = *scored; - scorer.progress.send_modify(|progress| { + scorer.publish(|progress| { progress.scored = scored; if let Some(top) = top { progress.full = true; @@ -1120,7 +1327,11 @@ fn record_gather_sources(rank: usize, gathered: &GatheredEx) { #[cfg(test)] mod tests { - use super::{DenseGatherMode, LayeredLazyConfig, LazyDenseForecast, count_missing}; + use super::{ + DenseGatherMode, DenseToEager, LayeredLazyConfig, LazyDenseForecast, LazyFarIssue, + LazyFarPermits, LazyProgress, count_missing, + }; + use tokio::sync::watch; #[test] fn dense_forecast_follows_k_and_partition_sizes() { @@ -1132,7 +1343,7 @@ mod tests { let rows = [100, 0, 50, 200]; let predicted = |forecast: LazyDenseForecast| { (0..rows.len()) - .map(|rank| (forecast.certain_dense(rank), forecast.routes_to_eager(rank))) + .map(|rank| (forecast.certain_dense(rank), forecast.routable(rank))) .collect::>() }; @@ -1189,11 +1400,25 @@ mod tests { let config = LayeredLazyConfig { dense, ..config }; let forecast = LazyDenseForecast::new(rows, k, true, &config); assert_eq!(forecast.dense_at_fill, dense_at_fill, "{dense:?} k={k}"); - let routed = (0..rows.len()).all(|rank| forecast.routes_to_eager(rank)); + let routed = (0..rows.len()).all(|rank| forecast.routable(rank)); assert_eq!(routed, dense == DenseGatherMode::Whole, "{dense:?} k={k}"); } } + // Either routing mode lets staging route the predicted-dense probes, + // whose planes' tiers then decide. + for dense_to_eager in [DenseToEager::Origin, DenseToEager::All] { + let config = LayeredLazyConfig { + dense_to_eager, + ..config + }; + assert_eq!( + LazyDenseForecast::new(rows, 120, true, &config), + large, + "{dense_to_eager}" + ); + } + // No probe is routed under the sparse policy, which never reads a // whole plane, when a prefilter or an upper bound may drop rows of // any probe before its gather, or with the routing switched off; @@ -1207,7 +1432,7 @@ mod tests { ..config }; let off = LayeredLazyConfig { - dense_to_eager: false, + dense_to_eager: DenseToEager::Off, ..config }; for k in [10, 120, 1000] { @@ -1226,6 +1451,51 @@ mod tests { } } + /// A gather beyond the ordinary window without a free permit waits for + /// whichever comes first: scoring bringing it within the ordinary + /// window, where it needs no permit, or a freed permit, which it then + /// holds. Either way it prunes against the latest progress. + #[tokio::test] + async fn far_gather_waits_for_a_permit_or_the_ordinary_window() { + const WINDOW: usize = 1; + let far = LazyFarIssue { + window: 8, + permits: LazyFarPermits::try_new(1).unwrap(), + }; + let (progress_tx, mut progress) = watch::channel(LazyProgress { + scored: 0, + full: true, + threshold: 1.0, + }); + let held = far.permits.try_acquire().unwrap(); + + { + let waiting = far.wait_for_permit_or_window(&mut progress, WINDOW, 5); + tokio::pin!(waiting); + assert!(futures::poll!(&mut waiting).is_pending()); + progress_tx.send_modify(|progress| { + progress.scored = 4; + progress.threshold = 0.5; + }); + let (issue, permit) = waiting.await.unwrap(); + assert_eq!((issue.scored, issue.threshold), (4, 0.5)); + assert!(permit.is_none()); + assert_eq!(far.permits.in_flight(), 1); + } + + { + let waiting = far.wait_for_permit_or_window(&mut progress, WINDOW, 9); + tokio::pin!(waiting); + assert!(futures::poll!(&mut waiting).is_pending()); + drop(held); + let (issue, permit) = waiting.await.unwrap(); + assert_eq!(issue.scored, 4); + assert!(permit.is_some()); + assert_eq!(far.permits.in_flight(), 1); + } + assert_eq!(far.permits.in_flight(), 0); + } + #[test] fn count_missing_compares_sorted_offsets() { assert_eq!(count_missing(&[], &[1, 2]), 0); diff --git a/rust/lance/src/session.rs b/rust/lance/src/session.rs index 83cde73be48..7ee69d14141 100644 --- a/rust/lance/src/session.rs +++ b/rust/lance/src/session.rs @@ -11,6 +11,8 @@ use lance_index::IndexType; use lance_io::object_store::ObjectStoreRegistry; use lance_io::spill::{LocalSpillStore, SpillStore}; +pub use lance_index::vector::storage::OriginLatencyClass; + use crate::dataset::{DEFAULT_INDEX_CACHE_SIZE, DEFAULT_METADATA_CACHE_SIZE}; use crate::session::caches::GlobalMetadataCache; use crate::session::index_caches::GlobalIndexCache; @@ -74,6 +76,11 @@ pub struct Session { store_registry: Arc, spill_store: Arc, + + /// The origin latency class the process declares for the IVF_RQ indexes + /// this session opens; see [`Self::with_index_origin_latency`]. `None` + /// leaves each index the class of its object store. + index_origin_latency: Option, } impl DeepSizeOf for Session { @@ -104,6 +111,7 @@ impl std::fmt::Debug for Session { "index_extensions", &self.index_extensions.keys().collect::>(), ) + .field("index_origin_latency", &self.index_origin_latency) .finish() } } @@ -138,6 +146,7 @@ impl Session { index_extensions: HashMap::new(), store_registry, spill_store: Arc::new(LocalSpillStore::default()), + index_origin_latency: None, } } @@ -159,6 +168,7 @@ impl Session { index_extensions: HashMap::new(), store_registry, spill_store: Arc::new(LocalSpillStore::default()), + index_origin_latency: None, } } @@ -179,6 +189,37 @@ impl Session { self } + /// Declare the origin latency class of the IVF_RQ indexes this session + /// opens: how slow reads are that no Lance cache tier serves. `None` (the + /// default) leaves each index the class of its object store, high for a + /// cloud object store and low for local files and memory. + /// + /// Declare it when the object store is wrapped in something the store's + /// scheme does not show, such as a local page cache that also caches index + /// files in front of a cloud object store (declare + /// [`OriginLatencyClass::Low`]). The class sets reader policies that trade + /// bytes or background work for fewer origin requests; results do not + /// depend on it. Other index types ignore it, and `LANCE_RQ_ORIGIN_LATENCY` + /// set to `low` or `high` overrides it. + /// + /// ```rust,no_run + /// # use lance::session::{OriginLatencyClass, Session}; + /// let session = + /// Session::default().with_index_origin_latency(Some(OriginLatencyClass::Low)); + /// assert_eq!(session.index_origin_latency(), Some(OriginLatencyClass::Low)); + /// ``` + pub fn with_index_origin_latency(mut self, class: Option) -> Self { + self.index_origin_latency = class; + self + } + + /// The origin latency class this session declares for the IVF_RQ indexes + /// it opens, `None` when it declares none; see + /// [`Self::with_index_origin_latency`]. + pub fn index_origin_latency(&self) -> Option { + self.index_origin_latency + } + /// Return a reference to the session's spill store. /// /// Callers use this to obtain reclaimable scratch space for intermediate @@ -238,6 +279,7 @@ impl Session { index_extensions: HashMap::new(), store_registry, spill_store: Arc::new(LocalSpillStore::default()), + index_origin_latency: None, } } @@ -565,6 +607,46 @@ mod tests { .unwrap(); } + /// A session declares no origin latency class unless built with one, + /// whichever constructor built it; its clones keep the class, and its + /// debug output shows it. + #[test] + fn test_index_origin_latency_hint() { + let sessions = [ + Session::default(), + Session::new(0, 0, Default::default()), + Session::with_index_cache_backend( + Arc::new(QuickCacheBackend::with_capacity(0)), + 0, + Default::default(), + ), + Session::with_cache_backends( + CacheSpec::Default, + CacheSpec::Default, + Default::default(), + ), + ]; + for session in sessions { + assert_eq!(session.index_origin_latency(), None); + assert!( + format!("{session:?}").contains("index_origin_latency: None"), + "{session:?}" + ); + for class in [OriginLatencyClass::Low, OriginLatencyClass::High] { + let hinted = session.clone().with_index_origin_latency(Some(class)); + assert_eq!(hinted.index_origin_latency(), Some(class)); + assert_eq!(hinted.clone().index_origin_latency(), Some(class)); + assert!( + format!("{hinted:?}") + .contains(&format!("index_origin_latency: Some({class:?})")), + "{hinted:?}" + ); + let cleared = hinted.with_index_origin_latency(None); + assert_eq!(cleared.index_origin_latency(), None); + } + } + } + #[tokio::test] async fn test_custom_spill_store_injected() { let capped = Arc::new(LocalSpillStore::with_cap(50).unwrap());