diff --git a/rust/lance-index/src/vector/bq.rs b/rust/lance-index/src/vector/bq.rs index e61ad93702c..4d99a4a2d75 100644 --- a/rust/lance-index/src/vector/bq.rs +++ b/rust/lance-index/src/vector/bq.rs @@ -21,6 +21,9 @@ pub mod builder; pub(crate) mod dist_table_quant; pub mod ex_dot; pub mod layered; +pub mod layered_stats; +#[cfg(test)] +mod lazy_tests; pub(crate) mod pairwise; mod plane_cache; pub mod prune; diff --git a/rust/lance-index/src/vector/bq/layered_stats.rs b/rust/lance-index/src/vector/bq/layered_stats.rs new file mode 100644 index 00000000000..cb3df0d3074 --- /dev/null +++ b/rust/lance-index/src/vector/bq/layered_stats.rs @@ -0,0 +1,321 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Process-wide diagnostics of the lazy layered full-precision scan. +//! +//! 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 +//! [`snapshot_and_reset`]. + +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Number of probe-rank buckets: ranks 0, 1-3, 4-15, 16-63 and 64+. +pub const RANK_BUCKETS: usize = 5; + +/// Exclusive upper rank of every bucket but the last. +const RANK_BUCKET_ENDS: [usize; RANK_BUCKETS - 1] = [1, 4, 16, 64]; + +/// The bucket a probe rank (0 = the nearest partition) is reported in. +pub fn rank_bucket(rank: usize) -> usize { + RANK_BUCKET_ENDS + .iter() + .position(|&end| rank < end) + .unwrap_or(RANK_BUCKETS - 1) +} + +/// A cumulative counter. +#[derive(Debug)] +pub struct Counter(AtomicU64); + +impl Counter { + const fn new() -> Self { + Self(AtomicU64::new(0)) + } + + pub fn add(&self, value: u64) { + self.0.fetch_add(value, Ordering::Relaxed); + } + + pub fn incr(&self) { + self.add(1); + } + + /// Add the nanoseconds elapsed since `start`. + pub fn add_elapsed(&self, start: std::time::Instant) { + self.add(u64::try_from(start.elapsed().as_nanos()).unwrap_or(u64::MAX)); + } + + fn take(&self) -> u64 { + self.0.swap(0, Ordering::Relaxed) + } +} + +/// A cumulative counter split by probe-rank bucket (see [`rank_bucket`]). +#[derive(Debug)] +pub struct RankCounter([AtomicU64; RANK_BUCKETS]); + +impl RankCounter { + const fn new() -> Self { + Self([const { AtomicU64::new(0) }; RANK_BUCKETS]) + } + + pub fn add(&self, rank: usize, value: u64) { + self.0[rank_bucket(rank)].fetch_add(value, Ordering::Relaxed); + } + + pub fn incr(&self, rank: usize) { + self.add(rank, 1); + } + + fn take(&self) -> [u64; RANK_BUCKETS] { + std::array::from_fn(|bucket| self.0[bucket].swap(0, Ordering::Relaxed)) + } +} + +/// The largest value observed since the last snapshot. +#[derive(Debug)] +pub struct MaxCounter(AtomicU64); + +impl MaxCounter { + const fn new() -> Self { + Self(AtomicU64::new(0)) + } + + pub fn observe(&self, value: u64) { + self.0.fetch_max(value, Ordering::Relaxed); + } + + fn take(&self) -> u64 { + self.0.swap(0, Ordering::Relaxed) + } +} + +macro_rules! layered_lazy_counters { + ( + counters { $($(#[$counter_doc:meta])* $counter:ident,)* } + ranked { $($(#[$ranked_doc:meta])* $ranked:ident,)* } + maxima { $($(#[$max_doc:meta])* $max:ident,)* } + ) => { + /// Live counters; see [`LayeredLazyStats`] for their meaning. + #[derive(Debug)] + pub struct LayeredLazyCounters { + $(pub $counter: Counter,)* + $(pub $ranked: RankCounter,)* + $(pub $max: MaxCounter,)* + } + + impl LayeredLazyCounters { + const fn new() -> Self { + Self { + $($counter: Counter::new(),)* + $($ranked: RankCounter::new(),)* + $($max: MaxCounter::new(),)* + } + } + + fn take(&self) -> LayeredLazyStats { + LayeredLazyStats { + $($counter: self.$counter.take(),)* + $($ranked: self.$ranked.take(),)* + $($max: self.$max.take(),)* + } + } + } + + /// Counters accumulated since the previous [`snapshot_and_reset`]. + /// Rank-split counters hold one value per [`rank_bucket`]; `*_ns` + /// counters are summed nanoseconds. + #[derive(Debug, Default, Clone, PartialEq, Eq)] + pub struct LayeredLazyStats { + $($(#[$counter_doc])* pub $counter: u64,)* + $($(#[$ranked_doc])* pub $ranked: [u64; RANK_BUCKETS],)* + $($(#[$max_doc])* pub $max: u64,)* + } + }; +} + +layered_lazy_counters! { + counters { + /// Queries that ran the lazy scan. + lazy_queries, + /// Eligible queries whose probed ex planes were all resident, so the + /// eager scan ran unchanged. + lazy_all_resident_skips, + /// Layered queries checked while the lazy scan is disabled. + ineligible_disabled, + /// Layered queries on an index with a row-id remapper. + ineligible_remapper, + /// Layered queries below full precision. + ineligible_precision, + /// Layered queries in Fast approx mode. + ineligible_fast, + /// Layered queries without a raw-query estimator context. + ineligible_residual, + /// Layered queries with a cascade factor. + ineligible_cascade, + /// Layered queries with a refine factor. + ineligible_refine, + /// Layered queries with `k == 0`, which the eager scan serves. + ineligible_k_zero, + /// Rows the live heap still needed that the gather did not fetch. + /// Always 0 unless the lazy scan has a bug; the query then fails. + needed_not_fetched, + /// Sparse gathers whose rows came from the origin file because the + /// cache had no persistent entry for the plane. + origin_row_reads, + /// Sum over gathered probes of probes still unscored at issue. + staleness_sum, + /// Issues deferred because the heap was not yet full at the gate. + deferred_issues, + /// Deferred issues released only when every earlier probe was scored. + serial_waits, + /// Gathers issued before the heap was full and before their turn, + /// selecting every accepted row; see `LANCE_RQ_LAZY_EAGER_BEFORE_FULL`. + eager_before_full, + /// Sparse gathers of planes the persistent tier did not hold that + /// loaded the whole plane instead of more origin row runs than + /// `LANCE_RQ_LAZY_ORIGIN_MAX_RUNS` (unlimited by default). + origin_whole_fallbacks, + /// Loading sign planes. + sign_load_ns, + /// Stage-1 CPU dispatches. + stage1_dispatches, + /// Stage-1 work, measured on the CPU pool. + stage1_cpu_ns, + /// Waiting for the threshold gate before issuing a gather. + gate_wait_ns, + /// Gathering resident ex planes. + fetch_resident_ns, + /// Gathering whole ex planes from the cache or origin. + fetch_whole_ns, + /// Gathering selected ex-plane rows. + fetch_sparse_ns, + /// Scorer idle time waiting for the next probe: I/O on the critical path. + scorer_wait_ns, + /// Stage-2 and eager scoring work. + stage2_cpu_ns, + /// Scoring batches run on the query task. + dispatch_inline, + /// Scoring batches dispatched to the CPU pool. + dispatch_spawn, + /// Time from query start until the heap first held `k` rows. + time_to_first_full_ns, + /// Background whole-plane promotions started. + promotions_issued, + /// Promotions not started because the plane was already in flight. + promotions_deduped, + /// Promotions not started for lack of an in-flight permit, failed, or + /// whose plane was no longer resident once loaded. + promotions_skipped, + /// Promotions whose plane was resident once loaded. + promotions_completed, + /// Bytes of planes loaded by promotions, resident afterwards or not, + /// apart from critical-path reads. + promotion_bytes, + } + ranked { + /// Probes gathered lazily. + lazy_probes, + /// Probes of non-empty partitions scored eagerly because both ex + /// planes were resident. + eager_resident, + /// Probes scored eagerly because lower-bound gating is disabled. + gating_off, + /// Probes of empty partitions, which have no ex rows and are scored + /// by the eager scan. + 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`. + dense_to_eager, + /// Lazy probes issued at once because the rows of all earlier probes + /// cannot fill the heap. The probes of queries that + /// `LANCE_RQ_LAZY_DENSE_TO_EAGER` routes are counted in + /// `dense_to_eager` instead. + certain_dense, + /// High planes served from RAM. + high_resident, + /// High planes read whole. + high_whole, + /// High planes read by selected rows. + high_sparse, + /// Low planes served from RAM. + low_resident, + /// Low planes read whole. + low_whole, + /// Low planes read by selected rows. + low_sparse, + /// Rows accepted by the prefilter. + stage1_candidates, + /// Rows stage 1 pruned against the query's upper bound. + stage1_pruned_ub, + /// Rows stage 1 pruned against the issue-time heap threshold. + stage1_pruned_heap, + /// Stage-1 survivors whose ex rows were gathered. + rows_fetched, + /// Survivors scored exactly. + stage2_exact, + /// Survivors pruned against the live heap: over-fetch from a stale + /// issue-time threshold. + stage2_pruned_live, + /// Exactly scored survivors outside the query's distance range. + stage2_exact_rejected, + } + maxima { + /// Largest number of probes still unscored at a gather's issue. + staleness_max, + } +} + +static COUNTERS: LayeredLazyCounters = LayeredLazyCounters::new(); +static PROMOTIONS_IN_FLIGHT: AtomicU64 = AtomicU64::new(0); + +/// The live counters, for recording. +pub fn counters() -> &'static LayeredLazyCounters { + &COUNTERS +} + +/// Read and zero every counter. +pub fn snapshot_and_reset() -> LayeredLazyStats { + COUNTERS.take() +} + +/// Background promotions currently running. A gauge, not reset by snapshots. +pub fn promotions_in_flight() -> u64 { + PROMOTIONS_IN_FLIGHT.load(Ordering::Relaxed) +} + +pub(crate) fn promotion_started() { + PROMOTIONS_IN_FLIGHT.fetch_add(1, Ordering::Relaxed); +} + +pub(crate) fn promotion_finished() { + PROMOTIONS_IN_FLIGHT.fetch_sub(1, Ordering::Relaxed); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rank_buckets_match_their_documented_ranges() { + let buckets = [0, 1, 3, 4, 15, 16, 63, 64, 10_000].map(rank_bucket); + assert_eq!(buckets, [0, 1, 1, 2, 2, 3, 3, 4, 4]); + } + + #[test] + fn counters_reset_on_snapshot() { + let counters = LayeredLazyCounters::new(); + counters.lazy_queries.add(3); + counters.rows_fetched.add(5, 7); + counters.staleness_max.observe(4); + counters.staleness_max.observe(2); + let stats = counters.take(); + assert_eq!(stats.lazy_queries, 3); + assert_eq!(stats.rows_fetched, [0, 0, 7, 0, 0]); + assert_eq!(stats.staleness_max, 4); + assert_eq!(counters.take(), LayeredLazyStats::default()); + } +} diff --git a/rust/lance-index/src/vector/bq/lazy_tests.rs b/rust/lance-index/src/vector/bq/lazy_tests.rs new file mode 100644 index 00000000000..415ca1efc85 --- /dev/null +++ b/rust/lance-index/src/vector/bq/lazy_tests.rs @@ -0,0 +1,1545 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Parity tests for the lazy full-precision scan of a layered index: stage 1 +//! bounds every row of the sign plane, stage 2 reranks the survivors against +//! gathered ex rows, and together they must reproduce the eager scan's heap +//! (row ids, distance bits, tie outcomes) and prune counters exactly. + +use std::collections::{BinaryHeap, HashSet}; +use std::sync::Arc; + +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, UInt32Array, UInt64Array, +}; +use arrow_schema::Schema; +use bytes::Bytes; +use lance_arrow::{FixedSizeListArrayExt, RecordBatchExt}; +use lance_core::cache::{CacheCodec, CacheDecode}; +use lance_linalg::distance::DistanceType; +use rstest::rstest; + +use super::RQBuildParams; +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::storage::{ + ExRows, RabitPruneCounters, RabitQuantizationStorage, SignStage, StagePruneCounts, SurvivorRow, + select_full_survivors, take_captured_prune_counters, +}; +use super::transform::{ERROR_FACTORS_COLUMN, RQTransformer}; +use crate::vector::graph::OrderedNode; +use crate::vector::quantizer::{Quantization, QuantizerStorage}; +use crate::vector::storage::{ + DistCalculator, DistanceCalculatorOptions, QueryResidual, RabitRawQueryContext, VectorStore, +}; +use crate::vector::transform::Transformer; +use crate::vector::{ApproxMode, CENTROID_DIST_COLUMN, PART_ID_COLUMN}; + +type Heap = BinaryHeap>; + +/// Deterministic generator so every case is reproducible from its seed. +struct SplitMix64(u64); + +impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + + fn below(&mut self, bound: usize) -> usize { + (self.next_u64() % bound as u64) as usize + } + + /// Uniform in `[-scale, scale)`. + fn symmetric(&mut self, scale: f32) -> f32 { + ((self.next_u64() >> 40) as f32 / (1u64 << 24) as f32 * 2.0 - 1.0) * scale + } + + fn vector(&mut self, dim: usize, scale: f32) -> Vec { + (0..dim).map(|_| self.symmetric(scale)).collect() + } +} + +/// One layered partition as the plane cache holds it. +struct Partition { + sign: RecordBatch, + hi: RecordBatch, + lo: RecordBatch, + row_ids: Vec, + rotated_centroid: Vec, + dist_q_c: f32, +} + +impl Partition { + fn len(&self) -> usize { + self.row_ids.len() + } + + /// The first `rows` rows, as fresh (unsliced) arrays. + fn prefix(&self, rows: usize) -> Self { + let indices = UInt32Array::from((0..rows as u32).collect::>()); + Self { + sign: self.sign.take(&indices).unwrap(), + hi: self.hi.take(&indices).unwrap(), + lo: self.lo.take(&indices).unwrap(), + row_ids: self.row_ids[..rows].to_vec(), + rotated_centroid: self.rotated_centroid.clone(), + dist_q_c: self.dist_q_c, + } + } + + /// The eager loader's storage: all three planes, sign columns first. + fn full_batch(&self) -> RecordBatch { + let mut fields = Vec::new(); + let mut columns = Vec::new(); + for batch in [&self.sign, &self.hi, &self.lo] { + fields.extend(batch.schema().fields().iter().cloned()); + columns.extend(batch.columns().iter().cloned()); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + } +} + +struct Fixture { + rq: RabitQuantizer, + /// Metric the storages are opened with; `Cosine` exercises the raw-query + /// mapping to L2. + distance_type: DistanceType, + bits: u8, + query: ArrayRef, + raw_query: RabitRawQueryContext, +} + +impl Fixture { + fn new(dim: usize, bits: u8, distance_type: DistanceType, rng: &mut SplitMix64) -> Self { + let sample = FixedSizeListArray::try_new_from_values( + Float32Array::from(rng.vector(64 * dim, 1.0)), + dim as i32, + ) + .unwrap(); + let rq = RabitQuantizer::build( + &sample, + DistanceType::L2, + &RQBuildParams::new(bits).with_layered(true), + ) + .unwrap(); + let query = rng.vector(dim, 1.0); + Self::with_query(rq, bits, distance_type, query) + } + + fn with_query( + rq: RabitQuantizer, + bits: u8, + distance_type: DistanceType, + query: Vec, + ) -> Self { + let query: ArrayRef = Arc::new(Float32Array::from(query)); + let raw_query = rq.metadata_ref().prepare_raw_query_context(&query).unwrap(); + Self { + rq, + distance_type, + bits, + query, + raw_query, + } + } + + fn dim(&self) -> usize { + self.query.len() + } + + fn rotated_dim(&self) -> usize { + self.rq.metadata_ref().rotated_dim() + } + + fn transform_metric(&self) -> DistanceType { + match self.distance_type { + DistanceType::Dot => DistanceType::Dot, + _ => DistanceType::L2, + } + } + + /// Quantize `residuals` (row-major, `dim` wide) of a partition centered + /// at `centroid`, with row ids `row_id_base..`: the transformer's batch, + /// the rotated centroid and the query's centroid distance. + fn transform( + &self, + residuals: &[f32], + centroid: &[f32], + row_id_base: u64, + ) -> (RecordBatch, Vec, f32) { + let dim = self.dim(); + let rows = residuals.len() / dim; + let query = self.query.as_any().downcast_ref::().unwrap(); + let query = query.values(); + let metric = self.transform_metric(); + let centroid_dists: Vec = residuals + .chunks_exact(dim) + .map(|residual| match metric { + DistanceType::Dot => { + 1.0 - residual + .iter() + .zip(centroid) + .map(|(r, c)| (r + c) * c) + .sum::() + } + _ => residual.iter().map(|r| r * r).sum(), + }) + .collect(); + let row_ids: Vec = (row_id_base..row_id_base + rows as u64).collect(); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from(residuals.to_vec()), + dim as i32, + ) + .unwrap(); + let input = RecordBatch::try_from_iter(vec![ + ("vector", Arc::new(vectors) as ArrayRef), + ( + lance_core::ROW_ID, + Arc::new(UInt64Array::from(row_ids)) as ArrayRef, + ), + ( + PART_ID_COLUMN, + Arc::new(UInt32Array::from(vec![0; rows])) as ArrayRef, + ), + ( + CENTROID_DIST_COLUMN, + Arc::new(Float32Array::from(centroid_dists)) as ArrayRef, + ), + ]) + .unwrap(); + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(centroid.to_vec()), + dim as i32, + ) + .unwrap(); + let rotated_centroid = self + .rq + .quantize_split(¢roids) + .unwrap() + .rotated_residuals + .unwrap(); + let batch = RQTransformer::new(self.rq.clone(), metric, centroids, "vector") + .unwrap() + .transform(&input) + .unwrap(); + let dist_q_c = match metric { + DistanceType::Dot => 1.0 - query.iter().zip(centroid).map(|(q, c)| q * c).sum::(), + _ => query + .iter() + .zip(centroid) + .map(|(q, c)| (q - c).powi(2)) + .sum(), + }; + (batch, rotated_centroid, dist_q_c) + } + + /// Encode a layered partition and split it into its cache planes. + fn encode(&self, residuals: &[f32], centroid: &[f32], row_id_base: u64) -> Partition { + 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) + .iter() + .map(|name| schema.index_of(name).unwrap()) + .collect(); + batch.project(&indices).unwrap() + }; + let sign = plane(0); + let row_ids = sign[lance_core::ROW_ID] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + Partition { + sign, + hi: plane(1), + lo: plane(2), + row_ids, + rotated_centroid, + dist_q_c, + } + } + + fn random_partition(&self, rows: usize, row_id_base: u64, rng: &mut SplitMix64) -> Partition { + let centroid = rng.vector(self.dim(), 0.2); + let residuals = rng.vector(rows * self.dim(), 1.0); + self.encode(&residuals, ¢roid, row_id_base) + } + + fn residual<'a>(&'a self, partition: &'a Partition) -> Option> { + Some(QueryResidual::RabitRawQuery { + rotated_centroid: Some(&partition.rotated_centroid), + query: Some(&self.raw_query), + }) + } + + fn full_storage(&self, partition: &Partition) -> RabitQuantizationStorage { + RabitQuantizationStorage::try_from_batch( + partition.full_batch(), + self.rq.metadata_ref(), + self.distance_type, + None, + ) + .unwrap() + } + + fn sign_storage(&self, partition: &Partition) -> RabitQuantizationStorage { + RabitQuantizationStorage::try_from_sign_plane_for_full( + partition.sign.clone(), + self.rq.metadata_ref(), + self.distance_type, + ) + .unwrap() + } +} + +fn full_options(approx_mode: ApproxMode) -> DistanceCalculatorOptions { + DistanceCalculatorOptions { + approx_mode, + rq_precision: RQPrecision::Full, + } +} + +#[derive(Debug, Clone, Copy)] +struct Scan { + k: usize, + lower: Option, + upper: Option, + approx_mode: ApproxMode, +} + +impl Scan { + fn top_k(k: usize) -> Self { + Self { + k, + lower: None, + upper: None, + approx_mode: ApproxMode::Normal, + } + } +} + +/// Where stage 1 takes its heap threshold from. Every variant but the +/// negative controls' `Fixed` is a legal (possibly stale) heap top. +#[derive(Debug, Clone, Copy)] +enum Threshold { + Unbounded, + Live, + /// The live top as of this many partitions ago. + Stale(usize), + /// A value just above the live top. + NextUp, + /// The top published when the heap first filled. + FirstFull, + /// A caller-chosen threshold, used only to show illegal ones break parity. + Fixed(f32), +} + +/// How the survivors' ex rows are served to stage 2. +#[derive(Debug, Clone, Copy)] +enum Gather { + /// Whole planes; the ex index is the sign-plane offset. + Identity, + /// Arrow `take` of the survivors. + Take, + /// A persistent plane-cache round trip with a row-selective decode. + Codec, +} + +const GATHERS: [Gather; 3] = [Gather::Identity, Gather::Take, Gather::Codec]; + +fn codec_rows(plane: &RecordBatch, rows: &[u32]) -> RecordBatch { + let codec = CacheCodec::from_impl::(); + let mut encoded = Vec::new(); + codec + .serialize( + &(Arc::new(PlaneBatch(plane.clone())) as Arc), + &mut encoded, + ) + .unwrap(); + let encoded = Bytes::from(encoded); + let read = + |range: std::ops::Range| -> lance_core::Result { Ok(encoded.slice(range)) }; + let CacheDecode::Hit(decoded) = codec.deserialize_rows(&read, rows) else { + panic!("row-selective plane decode failed for rows {rows:?}") + }; + decoded.downcast::().unwrap().0.clone() +} + +fn filter_mask(partition: &Partition, filter: Option<&HashSet>) -> Option> { + filter.map(|filter| { + partition + .row_ids + .iter() + .map(|row_id| filter.contains(row_id)) + .collect() + }) +} + +fn stage_one( + fixture: &Fixture, + partition: &Partition, + sign: &RabitQuantizationStorage, + approx_mode: ApproxMode, +) -> std::result::Result { + let mut scratch = Vec::new(); + let calc = sign.dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(partition), + &mut scratch, + full_options(approx_mode), + ); + let mut stage = SignStage::default(); + calc.full_sign_stage_with_scratch( + &mut stage, + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new(), + )?; + Ok(stage) +} + +/// The eager scan exactly as `FlatIndex` drives it; returns the heap and the +/// per-partition prune counters it reported. +fn eager_scan( + fixture: &Fixture, + partitions: &[&Partition], + scan: Scan, + filter: Option<&HashSet>, +) -> (Heap, Vec) { + take_captured_prune_counters(); + let mut heap = BinaryHeap::new(); + for partition in partitions { + let storage = fixture.full_storage(partition); + let mut scratch = Vec::new(); + let calc = storage.dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(partition), + &mut scratch, + full_options(scan.approx_mode), + ); + let (mut dists, mut u16s, mut u8s, mut u32s) = + (Vec::new(), Vec::new(), Vec::new(), Vec::new()); + match filter { + None => calc.accumulate_topk_with_scratch( + scan.k, + scan.lower, + scan.upper, + |id| storage.row_id(id), + &mut heap, + &mut dists, + &mut u16s, + &mut u8s, + &mut u32s, + ), + Some(filter) => calc.accumulate_filtered_topk_with_scratch( + scan.k, + scan.lower, + scan.upper, + storage + .row_ids() + .enumerate() + .map(|(id, &row_id)| (id as u32, row_id)), + |row_id| filter.contains(&row_id), + &mut heap, + &mut dists, + &mut u16s, + &mut u8s, + &mut u32s, + ), + } + } + (heap, take_captured_prune_counters()) +} + +/// Knobs of one lazy replay that the tests vary. +#[derive(Debug, Clone, Copy)] +struct Lazy { + threshold: Threshold, + gather: Gather, +} + +/// The lazy two-stage scan over the same partitions, in the same probe +/// order and on one heap, as the IVF orchestrator drives it. +fn lazy_scan( + fixture: &Fixture, + partitions: &[&Partition], + scan: Scan, + filter: Option<&HashSet>, + lazy: Lazy, +) -> (Heap, Vec, Vec) { + take_captured_prune_counters(); + let mut heap: Heap = BinaryHeap::new(); + let mut returned = Vec::new(); + let mut stage1 = Vec::new(); + let heap_top = |heap: &Heap| { + (scan.k > 0 && heap.len() >= scan.k).then(|| heap.peek().map(|node| node.dist.0)) + }; + // history[j] is the published threshold after j partitions were scored. + let mut history = vec![None]; + let mut first_full = None; + for (probe, partition) in partitions.iter().enumerate() { + let sign = fixture.sign_storage(partition); + let stage = stage_one(fixture, partition, &sign, scan.approx_mode).unwrap(); + let accept = filter_mask(partition, filter); + let live = heap_top(&heap).flatten(); + let threshold = match lazy.threshold { + Threshold::Unbounded => None, + Threshold::Live => live, + Threshold::Stale(staleness) => history[probe.saturating_sub(staleness)], + Threshold::NextUp => live.map(f32::next_up), + Threshold::FirstFull => first_full, + // Published like a legal threshold, only once the heap is full. + Threshold::Fixed(threshold) => live.map(|_| threshold), + }; + let mut survivors = Vec::new(); + let mut counts = StagePruneCounts::default(); + select_full_survivors( + &stage.lower_bounds, + accept.as_deref(), + scan.upper, + threshold, + &mut survivors, + &mut counts, + ); + if !matches!(lazy.threshold, Threshold::Fixed(_)) { + let mut needed = Vec::new(); + select_full_survivors( + &stage.lower_bounds, + accept.as_deref(), + scan.upper, + live, + &mut needed, + &mut StagePruneCounts::default(), + ); + let fetched: HashSet = survivors.iter().copied().collect(); + assert!( + needed.iter().all(|offset| fetched.contains(offset)), + "a legal threshold must fetch every row the live heap can still accept ({lazy:?})" + ); + } + let (hi, lo, ex_indices): (RecordBatch, RecordBatch, Vec) = match lazy.gather { + Gather::Identity => ( + partition.hi.clone(), + partition.lo.clone(), + survivors.clone(), + ), + Gather::Take => { + let indices = UInt32Array::from(survivors.clone()); + ( + partition.hi.take(&indices).unwrap(), + partition.lo.take(&indices).unwrap(), + (0..survivors.len() as u32).collect(), + ) + } + Gather::Codec => ( + codec_rows(&partition.hi, &survivors), + codec_rows(&partition.lo, &survivors), + (0..survivors.len() as u32).collect(), + ), + }; + let ex_rows = + ExRows::from_plane_batches(&hi, &lo, fixture.rotated_dim(), fixture.bits).unwrap(); + let mut scratch = Vec::new(); + let calc = sign + .dist_calculator_with_ex_rows( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(partition), + &mut scratch, + full_options(scan.approx_mode), + ex_rows, + ) + .unwrap(); + let rows: Vec = survivors + .iter() + .zip(ex_indices) + .map(|(&offset, ex_index)| SurvivorRow { + ex_index, + row_id: partition.row_ids[offset as usize], + binary_ip: stage.binary_ips[offset as usize], + lower_bound: stage.lower_bounds[offset as usize], + }) + .collect(); + let was_full = heap_top(&heap).is_some(); + let mut published = None; + let counters = calc.accumulate_survivor_rows( + scan.k, + scan.lower, + scan.upper, + &rows, + counts, + &mut heap, + |top| { + assert!(published.is_none(), "the first-full event fires once"); + published = Some(top); + }, + ); + assert_eq!( + published.is_some(), + !was_full && heap_top(&heap).is_some(), + "first-full fires exactly when this partition fills the heap" + ); + if first_full.is_none() { + first_full = published; + } + if scan.k > 0 && partition.len() > 0 { + returned.push(counters); + stage1.push(counts); + } + history.push(heap_top(&heap).flatten()); + } + let captured = take_captured_prune_counters(); + assert_eq!( + returned, captured, + "returned counters are the reported ones" + ); + (heap, captured, stage1) +} + +fn to_bits(values: &[f32]) -> Vec { + values.iter().map(|value| value.to_bits()).collect() +} + +fn heap_bits(heap: Heap) -> Vec<(u64, u32)> { + heap.into_vec() + .into_iter() + .map(|node| (node.id, node.dist.0.to_bits())) + .collect() +} + +#[track_caller] +fn assert_lazy_matches_eager( + fixture: &Fixture, + partitions: &[&Partition], + scan: Scan, + filter: Option<&HashSet>, + lazy: Lazy, +) -> HeapPrunes { + let (eager, eager_counters) = eager_scan(fixture, partitions, scan, filter); + let (actual, lazy_counters, stage1) = lazy_scan(fixture, partitions, scan, filter, lazy); + assert_eq!( + heap_bits(actual), + heap_bits(eager), + "heap differs: {scan:?} {lazy:?}" + ); + assert_eq!( + lazy_counters, eager_counters, + "prune counters differ: {scan:?} {lazy:?}" + ); + let stage1_heap: usize = stage1.iter().map(|counts| counts.pruned_heap).sum(); + let total_heap: usize = lazy_counters + .iter() + .map(|counters| counters.pruned_heap) + .sum(); + HeapPrunes { + stage1: stage1_heap, + stage2: total_heap - stage1_heap, + } +} + +/// Heap prunes of a lazy replay by stage; stage 2 prunes the rows a stale +/// threshold over-fetched. +#[derive(Debug, Default, Clone, Copy)] +struct HeapPrunes { + stage1: usize, + stage2: usize, +} + +fn legal_thresholds(rng: &mut SplitMix64) -> [Threshold; 5] { + [ + Threshold::Unbounded, + Threshold::Live, + Threshold::Stale(1 + rng.below(8)), + Threshold::NextUp, + Threshold::FirstFull, + ] +} + +/// Stage 1 reproduces the eager full store's binary inner products and native +/// lower bounds bit for bit, across SIMD block boundaries, padded dims, every +/// layered width, both LUT precisions and the exact-dequantization fallback. +#[rstest] +fn sign_stage_matches_full_store( + #[values(64, 96, 1024)] dim: usize, + #[values(3, 4, 5, 6, 7, 8, 9)] bits: u8, +) { + const ROWS: [usize; 14] = [0, 1, 5, 15, 16, 17, 31, 32, 33, 63, 64, 65, 100, 2158]; + let mut rng = SplitMix64(dim as u64 * 131 + bits as u64); + let fixture = Fixture::new(dim, bits, DistanceType::L2, &mut rng); + let all_rows = fixture.random_partition(ROWS[ROWS.len() - 1], 0, &mut rng); + let query = fixture + .query + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + // A query this large makes the affine LUT reconstruction overflow while + // the table itself stays finite, which forces the exact per-row fallback. + // The window depends on the (unseeded) rotation, so sweep scales densely + // from 1e30 to 1e38 in quarter decades. + let exact_fixture = (120..=152) + .map(|quarter_decades| 10f32.powf(quarter_decades as f32 / 4.0)) + .map(|scale| { + Fixture::with_query( + fixture.rq.clone(), + bits, + DistanceType::L2, + query.iter().map(|q| q * scale).collect(), + ) + }) + .find(|fixture| { + let dist_table = &fixture.raw_query.dist_table; + dist_table.iter().all(|value| value.is_finite()) + && quantize_dist_table_into(dist_table, &mut Vec::new()) == DistTableDequant::Exact + }) + .expect("some query scale must force the exact dequantization fallback"); + assert_eq!( + quantize_dist_table_u16_into(&exact_fixture.raw_query.dist_table, &mut Vec::new()), + DistTableDequant::Exact + ); + + for rows in ROWS { + let partition = all_rows.prefix(rows); + for fixture in [&fixture, &exact_fixture] { + let full = fixture.full_storage(&partition); + let sign = fixture.sign_storage(&partition); + assert_eq!(sign.len(), rows); + for approx_mode in [ApproxMode::Normal, ApproxMode::Accurate] { + for residual in [fixture.residual(&partition), None] { + let (mut full_scratch, mut sign_scratch) = (Vec::new(), Vec::new()); + let options = full_options(approx_mode); + let eager = full.dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + residual, + &mut full_scratch, + options, + ); + let calc = sign.dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + residual, + &mut sign_scratch, + options, + ); + let mut stage = SignStage::default(); + calc.full_sign_stage_with_scratch( + &mut stage, + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new(), + ) + .unwrap(); + let expected_ips = eager.binary_inner_products(); + assert_eq!( + to_bits(&stage.binary_ips), + to_bits(&expected_ips), + "binary inner products differ: rows={rows} mode={approx_mode:?}" + ); + let expected_bounds: Vec = expected_ips + .iter() + .enumerate() + .map(|(row, &ip)| eager.raw_query_lower_bound(row, ip).unwrap()) + .collect(); + assert_eq!( + to_bits(&stage.lower_bounds), + to_bits(&expected_bounds), + "lower bounds differ: rows={rows} mode={approx_mode:?}" + ); + } + } + } + } +} + +/// Property test: random probe sequences of layered partitions, all legal +/// threshold schedules and all three ways of serving the ex rows produce the +/// eager heap and prune counters bit for bit. +#[rstest] +fn lazy_replay_equals_dense_heap( + #[values(3, 4, 5, 6, 7, 8, 9)] bits: u8, + #[values(DistanceType::L2, DistanceType::Dot, DistanceType::Cosine)] metric: DistanceType, +) { + const SEEDS: u64 = 200; + const POOL_ROWS: [usize; 9] = [0, 1, 17, 31, 32, 33, 64, 100, 150]; + let metric_seed = match metric { + DistanceType::L2 => 1, + DistanceType::Dot => 2, + _ => 3, + }; + let mut rng = SplitMix64(bits as u64 * 7919 + metric_seed); + let dim = if bits == 7 { 96 } else { 64 }; + let fixture = Fixture::new(dim, bits, metric, &mut rng); + let pool: Vec = POOL_ROWS + .iter() + .enumerate() + .map(|(part, &rows)| fixture.random_partition(rows, part as u64 * 10_000, &mut rng)) + .collect(); + // Range bounds drawn from the spread of full-precision scores. + let mut scores: Vec = pool + .iter() + .filter(|partition| partition.len() > 0) + .flat_map(|partition| { + let storage = fixture.full_storage(partition); + storage + .dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(partition), + &mut Vec::new(), + full_options(ApproxMode::Normal), + ) + .distance_all(partition.len()) + }) + .collect(); + scores.sort_by(f32::total_cmp); + let quantile = |q: f64| scores[(q * (scores.len() - 1) as f64) as usize]; + + let mut prunes = HeapPrunes::default(); + for seed in 0..SEEDS { + let mut rng = SplitMix64(seed ^ 0xA5A5_0000); + let mut order: Vec = (0..pool.len()).collect(); + for i in (1..order.len()).rev() { + order.swap(i, rng.below(i + 1)); + } + let partitions: Vec<&Partition> = order[..1 + rng.below(6)] + .iter() + .map(|&part| &pool[part]) + .collect(); + let total: usize = partitions.iter().map(|partition| partition.len()).sum(); + let approx_mode = if rng.below(2) == 0 { + ApproxMode::Normal + } else { + ApproxMode::Accurate + }; + let (lower, upper) = match rng.below(4) { + 0 => (None, None), + 1 => (Some(quantile(0.1)), None), + 2 => (None, Some(quantile(0.6))), + _ => (Some(quantile(0.05)), Some(quantile(0.7))), + }; + let filter: Option> = (rng.below(3) == 0).then(|| { + partitions + .iter() + .flat_map(|partition| partition.row_ids.iter().copied()) + .filter(|_| rng.below(3) != 0) + .collect() + }); + for k in [0, 1, 7, 100, total + 1] { + let scan = Scan { + k, + lower, + upper, + approx_mode, + }; + for threshold in legal_thresholds(&mut rng) { + let gather = GATHERS[rng.below(GATHERS.len())]; + let replay = assert_lazy_matches_eager( + &fixture, + &partitions, + scan, + filter.as_ref(), + Lazy { threshold, gather }, + ); + prunes.stage1 += replay.stage1; + prunes.stage2 += replay.stage2; + } + } + } + assert!( + prunes.stage1 > 0 && prunes.stage2 > 0, + "the replays must prune in stage 1 and re-prune stale over-fetch in stage 2: {prunes:?}" + ); +} + +/// Duplicate rows (identical codes and factors) tie within a partition, +/// across a 32-row SIMD block boundary and across partitions; the incumbent +/// keeps its slot in both scans, and the outcome depends on probe order. +#[test] +fn ties_keep_incumbents_in_probe_order() { + let mut rng = SplitMix64(0x7135); + let fixture = Fixture::new(64, 7, DistanceType::L2, &mut rng); + let dim = fixture.dim(); + let centroid = rng.vector(dim, 0.2); + let mut residuals = rng.vector(70 * dim, 1.0); + // Rows 31 and 32 straddle the first SIMD block boundary; both copy row 5. + let duplicate = residuals[5 * dim..6 * dim].to_vec(); + residuals[31 * dim..32 * dim].copy_from_slice(&duplicate); + residuals[32 * dim..33 * dim].copy_from_slice(&duplicate); + let first = fixture.encode(&residuals, ¢roid, 0); + let second = fixture.encode(&residuals, ¢roid, 1_000); + + let forward = [&first, &second]; + let backward = [&second, &first]; + let mut order_sensitive = false; + for k in 1..=24 { + for probes in [&forward[..], &backward[..]] { + for gather in GATHERS { + for threshold in [Threshold::Live, Threshold::Stale(1), Threshold::FirstFull] { + assert_lazy_matches_eager( + &fixture, + probes, + Scan::top_k(k), + None, + Lazy { threshold, gather }, + ); + } + } + } + let ids = |probes: &[&Partition]| { + let mut ids: Vec = eager_scan(&fixture, probes, Scan::top_k(k), None) + .0 + .into_iter() + .map(|node| node.id) + .collect(); + ids.sort_unstable(); + ids + }; + order_sensitive |= ids(&forward[..]) != ids(&backward[..]); + } + assert!( + order_sensitive, + "swapping identical partitions must change which tied rows are kept" + ); +} + +/// Negative controls: the parity above depends on carrying the original +/// partition's binary inner products and on taking the threshold only from +/// the full-score heap. +#[test] +fn illegal_inputs_break_parity() { + let mut rng = SplitMix64(0x4e67); + let fixture = Fixture::new(64, 7, DistanceType::L2, &mut rng); + + // Recomputing a row's inner product from a compacted sign plane moves it + // between the quantized SIMD LUT and the float tail LUT. + let partition = fixture.random_partition(33, 0, &mut rng); + let sign = fixture.sign_storage(&partition); + let stage = stage_one(&fixture, &partition, &sign, ApproxMode::Normal).unwrap(); + let ex_rows = ExRows::from_plane_batches( + &partition.hi, + &partition.lo, + fixture.rotated_dim(), + fixture.bits, + ) + .unwrap(); + let mut scratch = Vec::new(); + let calc = sign + .dist_calculator_with_ex_rows( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(&partition), + &mut scratch, + full_options(ApproxMode::Normal), + ex_rows, + ) + .unwrap(); + let changed_scores = (0..partition.len()) + .filter(|&row| { + let single = partition_rows(&partition, &[row as u32]); + let recomputed = stage_one( + &fixture, + &single, + &fixture.sign_storage(&single), + ApproxMode::Normal, + ) + .unwrap() + .binary_ips[0]; + let carried = + calc.distance_with_binary_inner_product(row as u32, stage.binary_ips[row]); + let rescored = calc.distance_with_binary_inner_product(row as u32, recomputed); + carried.to_bits() != rescored.to_bits() + }) + .count(); + assert!( + changed_scores > 0, + "recomputed inner products must change some full-precision score" + ); + + // A threshold below the live heap top prunes a row the eager scan keeps. + let probes: Vec = (0..4) + .map(|part| fixture.random_partition(100, part * 1_000, &mut rng)) + .collect(); + let probes: Vec<&Partition> = probes.iter().collect(); + let scan = Scan::top_k(10); + let (last, entrant_bound, live, eager_bits) = (1..probes.len()) + .rev() + .find_map(|last| { + let prefix = &probes[..=last]; + let (eager, _) = eager_scan(&fixture, prefix, scan, None); + let kept: HashSet = eager.iter().map(|node| node.id).collect(); + let partition = prefix[last]; + let stage = stage_one( + &fixture, + partition, + &fixture.sign_storage(partition), + scan.approx_mode, + ) + .unwrap(); + let entrant_bound = partition + .row_ids + .iter() + .zip(&stage.lower_bounds) + .filter(|(row_id, _)| kept.contains(row_id)) + .map(|(_, &bound)| bound) + .min_by(f32::total_cmp)?; + let (before, _) = eager_scan(&fixture, &prefix[..last], scan, None); + let live = before.peek().unwrap().dist.0; + Some((last, entrant_bound, live, heap_bits(eager))) + }) + .expect("some later probe must place a row in the final top-k"); + assert!( + entrant_bound < live, + "an entrant was reranked below the live top" + ); + let (tightened, _, _) = lazy_scan( + &fixture, + &probes[..=last], + scan, + None, + Lazy { + threshold: Threshold::Fixed(entrant_bound), + gather: Gather::Take, + }, + ); + assert_ne!(heap_bits(tightened), eager_bits); + + // The k-th best 1-bit estimate is not a full-score threshold. + let mut broke = 0; + for k in 1..=30 { + let scan = Scan::top_k(k); + let (eager, _) = eager_scan(&fixture, &probes, scan, None); + let mut estimates: Vec = probes + .iter() + .flat_map(|partition| { + fixture + .full_storage(partition) + .dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(partition), + &mut Vec::new(), + full_options(ApproxMode::Fast), + ) + .distance_all(partition.len()) + }) + .collect(); + estimates.sort_by(f32::total_cmp); + let (estimated, _, _) = lazy_scan( + &fixture, + &probes, + scan, + None, + Lazy { + threshold: Threshold::Fixed(estimates[k - 1]), + gather: Gather::Take, + }, + ); + broke += usize::from(heap_bits(estimated) != heap_bits(eager)); + } + assert!( + broke > 0, + "a 1-bit estimate threshold must lose some true top-k row" + ); +} + +/// The rows of `partition` at `rows`, re-materialized as their own partition. +fn partition_rows(partition: &Partition, rows: &[u32]) -> Partition { + let indices = UInt32Array::from(rows.to_vec()); + Partition { + sign: partition.sign.take(&indices).unwrap(), + hi: partition.hi.take(&indices).unwrap(), + lo: partition.lo.take(&indices).unwrap(), + row_ids: rows + .iter() + .map(|&row| partition.row_ids[row as usize]) + .collect(), + rotated_centroid: partition.rotated_centroid.clone(), + dist_q_c: partition.dist_q_c, + } +} + +/// Prefiltered partitions follow the eager filtered scan; rows the filter +/// rejects are never selected. +#[rstest] +#[case::all(|_: usize, _: usize| true)] +#[case::none(|_: usize, _: usize| false)] +#[case::alternating(|offset: usize, _: usize| offset.is_multiple_of(2))] +#[case::single_row(|offset: usize, _: usize| offset == 3)] +#[case::every_17th(|offset: usize, _: usize| offset.is_multiple_of(17))] +#[case::simd_tail_only(|offset: usize, rows: usize| offset >= rows - rows % 32)] +fn filters_match_eager_filtered_scan(#[case] accept: fn(usize, usize) -> bool) { + let mut rng = SplitMix64(0xF117); + let fixture = Fixture::new(64, 9, DistanceType::L2, &mut rng); + let partitions: Vec = [70, 33, 5, 100] + .into_iter() + .enumerate() + .map(|(part, rows)| fixture.random_partition(rows, part as u64 * 1_000, &mut rng)) + .collect(); + let probes: Vec<&Partition> = partitions.iter().collect(); + let filter: HashSet = partitions + .iter() + .flat_map(|partition| { + let rows = partition.len(); + partition + .row_ids + .iter() + .enumerate() + .filter(move |(offset, _)| accept(*offset, rows)) + .map(|(_, &row_id)| row_id) + }) + .collect(); + for partition in &partitions { + let stage = stage_one( + &fixture, + partition, + &fixture.sign_storage(partition), + ApproxMode::Normal, + ) + .unwrap(); + let mask = filter_mask(partition, Some(&filter)).unwrap(); + let mut survivors = Vec::new(); + let mut counts = StagePruneCounts::default(); + select_full_survivors( + &stage.lower_bounds, + Some(mask.as_slice()), + None, + None, + &mut survivors, + &mut counts, + ); + assert!(survivors.iter().all(|&offset| mask[offset as usize])); + assert_eq!( + counts.candidates, + mask.iter().filter(|&&accepted| accepted).count() + ); + } + for k in [1, 10, 400] { + for gather in GATHERS { + for threshold in [Threshold::Live, Threshold::Stale(2), Threshold::Unbounded] { + assert_lazy_matches_eager( + &fixture, + &probes, + Scan::top_k(k), + Some(&filter), + Lazy { threshold, gather }, + ); + } + } + } +} + +/// Query distance ranges, including an upper bound below every lower bound +/// (no survivors), a lower bound above every score (every exact score is +/// rejected) and a heap that never fills. +#[test] +fn range_bounds_match_eager() { + let mut rng = SplitMix64(0x5A9E); + let fixture = Fixture::new(64, 5, DistanceType::L2, &mut rng); + let partitions: Vec = [40, 65, 12] + .into_iter() + .enumerate() + .map(|(part, rows)| fixture.random_partition(rows, part as u64 * 1_000, &mut rng)) + .collect(); + let probes: Vec<&Partition> = partitions.iter().collect(); + let total: usize = partitions.iter().map(Partition::len).sum(); + let mut bounds = Vec::new(); + let mut scores = Vec::new(); + for partition in &partitions { + bounds.extend( + stage_one( + &fixture, + partition, + &fixture.sign_storage(partition), + ApproxMode::Normal, + ) + .unwrap() + .lower_bounds, + ); + scores.extend( + fixture + .full_storage(partition) + .dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(partition), + &mut Vec::new(), + full_options(ApproxMode::Normal), + ) + .distance_all(partition.len()), + ); + } + let min_bound = bounds.iter().copied().min_by(f32::total_cmp).unwrap(); + let max_score = scores.iter().copied().max_by(f32::total_cmp).unwrap(); + let median = { + let mut sorted = scores.clone(); + sorted.sort_by(f32::total_cmp); + sorted[sorted.len() / 2] + }; + let ranges = [ + (None, None), + (Some(median), None), + (None, Some(median)), + (Some(median - 1.0), Some(median + 1.0)), + (None, Some(min_bound.next_down())), + (Some(max_score.next_up()), None), + ]; + for (lower, upper) in ranges { + for k in [1, 10, total + 5] { + let scan = Scan { + k, + lower, + upper, + approx_mode: ApproxMode::Normal, + }; + for gather in GATHERS { + for threshold in [Threshold::Live, Threshold::Stale(1), Threshold::NextUp] { + assert_lazy_matches_eager( + &fixture, + &probes, + scan, + None, + Lazy { threshold, gather }, + ); + } + } + } + } + let (heap, counters) = eager_scan( + &fixture, + &probes, + Scan { + k: 10, + lower: None, + upper: Some(min_bound.next_down()), + approx_mode: ApproxMode::Normal, + }, + None, + ); + assert!(heap.is_empty()); + assert!( + counters + .iter() + .all(|c| c.pruned_upper_bound == c.candidates) + ); + let (heap, counters) = eager_scan( + &fixture, + &probes, + Scan { + k: 10, + lower: Some(max_score.next_up()), + upper: None, + approx_mode: ApproxMode::Normal, + }, + None, + ); + assert!(heap.is_empty()); + assert!( + counters + .iter() + .all(|c| c.exact_rejected == c.exact && c.exact > 0) + ); +} + +/// NaN bounds are never pruned and infinite ones are pruned as upper-bound +/// rows, in both scans. +#[test] +fn non_finite_bounds_match_eager() { + let mut rng = SplitMix64(0x0A11); + let fixture = Fixture::new(64, 7, DistanceType::L2, &mut rng); + let partitions: Vec = [70, 40] + .into_iter() + .enumerate() + .map(|(part, rows)| { + let mut partition = fixture.random_partition(rows, part as u64 * 1_000, &mut rng); + let errors = partition.sign[ERROR_FACTORS_COLUMN] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .enumerate() + .map(|(row, &error)| match row % 11 { + 3 => f32::NAN, + 7 => f32::NEG_INFINITY, + _ => error, + }) + .collect::>(); + partition.sign = partition + .sign + .replace_column_by_name(ERROR_FACTORS_COLUMN, Arc::new(Float32Array::from(errors))) + .unwrap(); + partition + }) + .collect(); + let probes: Vec<&Partition> = partitions.iter().collect(); + let stage = stage_one( + &fixture, + probes[0], + &fixture.sign_storage(probes[0]), + ApproxMode::Normal, + ) + .unwrap(); + assert!(stage.lower_bounds[3].is_nan()); + assert_eq!(stage.lower_bounds[7], f32::INFINITY); + for k in [1, 5, 50] { + for gather in GATHERS { + for threshold in [Threshold::Live, Threshold::Stale(1), Threshold::FirstFull] { + assert_lazy_matches_eager( + &fixture, + &probes, + Scan::top_k(k), + None, + Lazy { threshold, gather }, + ); + } + } + } + let mut survivors = Vec::new(); + let mut counts = StagePruneCounts::default(); + select_full_survivors( + &stage.lower_bounds, + None, + None, + Some(f32::NAN), + &mut survivors, + &mut counts, + ); + assert_eq!(counts.pruned_heap, 0, "a NaN threshold prunes nothing"); + assert_eq!( + survivors.len() + counts.pruned_upper_bound, + stage.lower_bounds.len() + ); +} + +/// Survivor selection boundaries match the eager per-row step's categories. +#[test] +fn select_full_survivors_boundaries() { + let bounds = [ + 1.0, + 2.0, + 2.0, + 3.0, + f32::NAN, + f32::NEG_INFINITY, + f32::INFINITY, + f32::MAX, + 2.5, + ]; + let mut survivors = vec![99]; + let mut counts = StagePruneCounts { + candidates: 7, + ..Default::default() + }; + // lb == T is a heap prune, lb == U an upper-bound prune; NaN survives. + select_full_survivors( + &bounds, + None, + Some(3.0), + Some(2.0), + &mut survivors, + &mut counts, + ); + assert_eq!(survivors, vec![0, 4, 5]); + assert_eq!( + counts, + StagePruneCounts { + candidates: 9, + pruned_upper_bound: 3, + pruned_heap: 3, + } + ); + // Without a query upper bound the eager scan still compares to f32::MAX. + select_full_survivors(&bounds, None, None, Some(2.0), &mut survivors, &mut counts); + assert_eq!(survivors, vec![0, 4, 5]); + assert_eq!(counts.pruned_upper_bound, 2); + assert_eq!(counts.pruned_heap, 4); + // Upper-bound pruning takes precedence over the heap threshold. + select_full_survivors( + &bounds, + None, + Some(2.0), + Some(1.0), + &mut survivors, + &mut counts, + ); + assert_eq!(survivors, vec![4, 5]); + assert_eq!(counts.pruned_upper_bound, 6); + assert_eq!(counts.pruned_heap, 1); + let accept = [true, false, true, true, false, true, true, true, false]; + select_full_survivors( + &bounds, + Some(&accept[..]), + Some(3.0), + None, + &mut survivors, + &mut counts, + ); + assert_eq!(survivors, vec![0, 2, 5]); + assert_eq!( + counts, + StagePruneCounts { + candidates: 6, + pruned_upper_bound: 3, + pruned_heap: 0, + } + ); + select_full_survivors(&[], None, None, None, &mut survivors, &mut counts); + assert!(survivors.is_empty()); + assert_eq!(counts, StagePruneCounts::default()); +} + +/// Stage 1 refuses calculators that cannot gate on the native lower bound, +/// and the sign-only constructors refuse inputs they cannot serve. +#[test] +fn gating_off_reports_reason() { + let mut rng = SplitMix64(0x6A7E); + let fixture = Fixture::new(64, 7, DistanceType::L2, &mut rng); + let partition = fixture.random_partition(40, 0, &mut rng); + let sign = fixture.sign_storage(&partition); + assert_eq!( + stage_one(&fixture, &partition, &sign, ApproxMode::Fast).unwrap_err(), + "approx_mode_fast" + ); + + let mut residual_metadata = fixture.rq.metadata_ref().clone(); + residual_metadata.query_estimator = super::storage::RabitQueryEstimator::ResidualQuery; + let residual_sign = RabitQuantizationStorage::try_from_sign_plane_for_full( + partition.sign.clone(), + &residual_metadata, + fixture.distance_type, + ) + .unwrap(); + let mut scratch = Vec::new(); + let calc = residual_sign.dist_calculator_with_scratch( + fixture.query.clone(), + partition.dist_q_c, + None, + &mut scratch, + full_options(ApproxMode::Normal), + ); + assert_eq!( + calc.full_sign_stage_with_scratch( + &mut SignStage::default(), + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new() + ) + .unwrap_err(), + "residual_query_estimator" + ); + + // A native index written without error factors scores every row. + let mut native_metadata = fixture.rq.metadata_ref().clone(); + native_metadata.layered = false; + let native_rq = RabitQuantizer::from_metadata(&native_metadata, DistanceType::L2).unwrap(); + let native_fixture = Fixture::with_query( + RabitQuantizer::try_from(native_rq).unwrap(), + 7, + DistanceType::L2, + fixture + .query + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(), + ); + let native_centroid = rng.vector(native_fixture.dim(), 0.2); + let (native_batch, native_rotated_centroid, native_dist_q_c) = native_fixture.transform( + &rng.vector(40 * native_fixture.dim(), 1.0), + &native_centroid, + 0, + ); + let native_storage = RabitQuantizationStorage::try_from_batch( + native_batch.drop_column(ERROR_FACTORS_COLUMN).unwrap(), + &native_metadata, + DistanceType::L2, + None, + ) + .unwrap(); + let mut native_scratch = Vec::new(); + let calc = native_storage.dist_calculator_with_scratch( + native_fixture.query.clone(), + native_dist_q_c, + Some(QueryResidual::RabitRawQuery { + rotated_centroid: Some(&native_rotated_centroid), + query: Some(&native_fixture.raw_query), + }), + &mut native_scratch, + full_options(ApproxMode::Normal), + ); + assert_eq!( + calc.full_sign_stage_with_scratch( + &mut SignStage::default(), + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new() + ) + .unwrap_err(), + "missing_error_factors" + ); + + let no_errors = partition.sign.drop_column(ERROR_FACTORS_COLUMN).unwrap(); + let error = RabitQuantizationStorage::try_from_sign_plane_for_full( + no_errors, + fixture.rq.metadata_ref(), + fixture.distance_type, + ) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains(ERROR_FACTORS_COLUMN), "{error}"); + let error = RabitQuantizationStorage::try_from_sign_plane_for_full( + partition.sign.clone(), + &native_metadata, + fixture.distance_type, + ) + .unwrap_err(); + assert!(error.to_string().contains("layered multi-bit"), "{error}"); + + // Gathered rows must match the index layout and the scan precision. + let error = ExRows::from_plane_batches(&partition.hi, &partition.lo, fixture.rotated_dim(), 9) + .unwrap_err(); + assert!(error.to_string().contains("byte width mismatch"), "{error}"); + let short_lo = partition.lo.slice(0, 10); + let error = + ExRows::from_plane_batches(&partition.hi, &short_lo, fixture.rotated_dim(), 7).unwrap_err(); + assert!(error.to_string().contains("disagree in length"), "{error}"); + let ex_rows = + ExRows::from_plane_batches(&partition.hi, &partition.lo, fixture.rotated_dim(), 7).unwrap(); + for options in [ + full_options(ApproxMode::Fast), + DistanceCalculatorOptions { + approx_mode: ApproxMode::Normal, + rq_precision: RQPrecision::High, + }, + ] { + let error = sign + .dist_calculator_with_ex_rows( + fixture.query.clone(), + partition.dist_q_c, + fixture.residual(&partition), + &mut Vec::new(), + options, + ex_rows, + ) + .err() + .unwrap(); + assert!( + error.to_string().contains("gathered ex rows require"), + "{error}" + ); + } +} diff --git a/rust/lance-index/src/vector/bq/plane_cache.rs b/rust/lance-index/src/vector/bq/plane_cache.rs index 29cf795fe00..a0aa8fffed2 100644 --- a/rust/lance-index/src/vector/bq/plane_cache.rs +++ b/rust/lance-index/src/vector/bq/plane_cache.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors //! Row-addressable persistent ex planes. Sign codes retain their transposed IPC body. +use std::ops::Range; use std::sync::Arc; use arrow_array::{Array, Float32Array, RecordBatch, UInt8Array, cast::AsArray, types::UInt8Type}; @@ -153,48 +154,20 @@ impl CacheCodecImpl for PlaneBatch { version: u32, rows: &[u32], ) -> Result { - if version != Self::CURRENT_VERSION { - return Err(Error::invalid_input( - "plane cache version has no row directory", - )); - } - let header = Header::parse(&reader.read_range(offset..offset + HEADER_BYTES)?)?; - if rows.windows(2).any(|pair| pair[0] >= pair[1]) - || rows.last().is_some_and(|&row| row as usize >= header.rows) - { - return Err(Error::invalid_input("invalid cached candidate offsets")); - } + let header = read_row_header(reader, offset, version)?; let width = header.width; let stride = header.stride(); let mut codes = Vec::new(); let mut adds = Vec::new(); let mut scales = Vec::new(); - let mut start = 0; - while start < rows.len() { - let mut end = start + 1; - while end < rows.len() - && (rows[end] - rows[end - 1] - 1) as usize <= COALESCE_GAP_BYTES / stride - { - end += 1; - } - let first = rows[start] as usize; - let last = rows[end - 1] as usize + 1; - let base = offset - .checked_add(HEADER_BYTES) - .ok_or_else(|| Error::invalid_input("plane offset overflow"))?; - let byte_start = first - .checked_mul(stride) - .and_then(|v| base.checked_add(v)) - .ok_or_else(|| Error::invalid_input("plane offset overflow"))?; - let byte_end = last - .checked_mul(stride) - .and_then(|v| base.checked_add(v)) - .ok_or_else(|| Error::invalid_input("plane offset overflow"))?; - let bytes = reader.read_range(byte_start..byte_end)?; - if bytes.len() != byte_end - byte_start { + for (bytes_range, run) in row_runs(&header, offset, rows)? { + let first = rows[run.start] as usize; + let len = bytes_range.len(); + let bytes = reader.read_range(bytes_range)?; + if bytes.len() != len { return Err(Error::invalid_input("short plane range read")); } - for &row in &rows[start..end] { + for &row in &rows[run] { let local = (row as usize - first) * stride; header.append( &bytes[local..local + width + 8], @@ -203,10 +176,69 @@ impl CacheCodecImpl for PlaneBatch { &mut scales, ); } - start = end; } header.batch(codes, adds, scales) } + fn plan_row_ranges( + reader: &dyn CacheRangeReader, + offset: usize, + version: u32, + rows: &[u32], + ) -> Result>>> { + let header = read_row_header(reader, offset, version)?; + Ok(Some( + row_runs(&header, offset, rows)? + .into_iter() + .map(|(bytes, _)| bytes) + .collect(), + )) + } +} + +fn read_row_header(reader: &dyn CacheRangeReader, offset: usize, version: u32) -> Result
{ + if version != PlaneBatch::CURRENT_VERSION { + return Err(Error::invalid_input( + "plane cache version has no row directory", + )); + } + Header::parse(&reader.read_range(offset..offset + HEADER_BYTES)?) +} + +/// Payload byte ranges for sorted, unique `rows`, each paired with the span of +/// `rows` it covers. Rows closer than [`COALESCE_GAP_BYTES`] share one range. +fn row_runs( + header: &Header, + offset: usize, + rows: &[u32], +) -> Result, Range)>> { + if rows.windows(2).any(|pair| pair[0] >= pair[1]) + || rows.last().is_some_and(|&row| row as usize >= header.rows) + { + return Err(Error::invalid_input("invalid cached candidate offsets")); + } + let stride = header.stride(); + let base = offset + .checked_add(HEADER_BYTES) + .ok_or_else(|| Error::invalid_input("plane offset overflow"))?; + let byte_offset = |row: usize| { + row.checked_mul(stride) + .and_then(|v| base.checked_add(v)) + .ok_or_else(|| Error::invalid_input("plane offset overflow")) + }; + let mut runs = Vec::new(); + let mut start = 0; + while start < rows.len() { + let mut end = start + 1; + while end < rows.len() + && (rows[end] - rows[end - 1] - 1) as usize <= COALESCE_GAP_BYTES / stride + { + end += 1; + } + let bytes = byte_offset(rows[start] as usize)?..byte_offset(rows[end - 1] as usize + 1)?; + runs.push((bytes, start..end)); + start = end; + } + Ok(runs) } #[cfg(test)] @@ -271,4 +303,68 @@ mod tests { )); } } + + #[test] + fn row_plan_matches_ranges_read_by_decode() { + const ROWS: usize = 4096; + const WIDTH: usize = 56; + let header = Header { + plane: 1, + rows: ROWS, + width: WIDTH, + }; + let original = header + .batch( + (0..ROWS * WIDTH).map(|v| (v % 253) as u8).collect(), + (0..ROWS).map(|v| v as f32).collect(), + vec![0.5; ROWS], + ) + .unwrap(); + let codec = CacheCodec::from_impl::(); + let mut encoded = Vec::new(); + codec + .serialize( + &(Arc::new(original) as Arc), + &mut encoded, + ) + .unwrap(); + let encoded = Bytes::from(encoded); + let requested = std::cell::RefCell::new(Vec::new()); + let read = |range: Range| -> Result { + requested.borrow_mut().push(range.clone()); + Ok(encoded.slice(range)) + }; + let stride = WIDTH + 8; + let gap = COALESCE_GAP_BYTES / stride; + let row_sets: Vec> = vec![ + vec![], + vec![0], + vec![ROWS as u32 - 1], + (0..ROWS as u32).collect(), + (0..ROWS as u32).step_by(gap + 1).collect(), + (0..ROWS as u32).step_by(gap + 2).collect(), + vec![1, 2, 3, 900, 901, 3000, 4095], + ]; + for rows in row_sets { + requested.borrow_mut().clear(); + let CacheDecode::Hit(_) = codec.deserialize_rows(&read, &rows) else { + panic!("range decode failed for {rows:?}") + }; + // The envelope and body header come first; the rest are row reads. + let decoded_rows = requested.borrow()[3..].to_vec(); + let plan = codec.plan_rows(&read, &rows).expect("plane codec plans"); + assert_eq!(plan, decoded_rows); + let expected_runs = if rows.is_empty() { + 0 + } else { + 1 + rows + .windows(2) + .filter(|w| (w[1] - w[0] - 1) as usize > gap) + .count() + }; + assert_eq!(plan.len(), expected_runs); + } + assert!(codec.plan_rows(&read, &[5, 5]).is_none()); + assert!(codec.plan_rows(&read, &[ROWS as u32]).is_none()); + } } diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index 1c50cedc029..cd6655c134f 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -149,17 +149,73 @@ fn emit_rabit_prune_stats(message: &str) { } /// Per-scan tallies of the raw-query lower-bound gating, reported through -/// `record_rabit_prune_stats`. -#[derive(Default)] -struct RabitPruneCounters { - candidates: usize, - pruned_upper_bound: usize, - pruned_heap: usize, - exact: usize, - exact_rejected: usize, +/// `record_rabit_prune_stats`. Every candidate row lands in exactly one of +/// `pruned_upper_bound`, `pruned_heap` or `exact`; `exact_rejected` counts +/// the exactly scored rows that fell outside the query's distance range. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct RabitPruneCounters { + pub candidates: usize, + pub pruned_upper_bound: usize, + pub pruned_heap: usize, + pub exact: usize, + pub exact_rejected: usize, +} + +/// Cumulative raw-query prune counters since process start. They advance only +/// while `LANCE_RQ_PRUNE_STATS` is enabled (reported by `enabled`), so callers +/// compare two snapshots instead of reading absolute values. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct RabitPruneStatsSnapshot { + pub enabled: bool, + /// Gated partition scans that reported counters. + pub calls: u64, + pub candidates: u64, + pub pruned_upper_bound: u64, + pub pruned_heap: u64, + pub exact: u64, + pub exact_rejected: u64, + /// Partition scans that could not gate on the lower bound (Fast mode, + /// residual estimator, or no error factors) and scored every row. + pub bypass_calls: u64, +} + +/// Read the process-wide prune counters without resetting them. +pub fn rabit_prune_stats_snapshot() -> RabitPruneStatsSnapshot { + let mut snapshot = RabitPruneStatsSnapshot { + enabled: rabit_prune_stats_enabled(), + ..Default::default() + }; + if let Some(stats) = RABIT_PRUNE_STATS.get() { + snapshot.calls = stats.calls.load(Ordering::Relaxed); + snapshot.candidates = stats.candidates.load(Ordering::Relaxed); + snapshot.pruned_upper_bound = stats.pruned_upper_bound.load(Ordering::Relaxed); + snapshot.pruned_heap = stats.pruned_heap.load(Ordering::Relaxed); + snapshot.exact = stats.exact.load(Ordering::Relaxed); + snapshot.exact_rejected = stats.exact_rejected.load(Ordering::Relaxed); + } + if let Some(stats) = RABIT_PRUNE_BYPASS_STATS.get() { + snapshot.bypass_calls = stats.calls.load(Ordering::Relaxed); + } + snapshot +} + +// Tests compare per-partition counters of two scan paths without enabling the +// process-wide, env-gated statistics that concurrent tests would share. +#[cfg(test)] +thread_local! { + static CAPTURED_PRUNE_COUNTERS: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +/// Drain the per-scan counters reported on this thread, in report order. +#[cfg(test)] +pub(crate) fn take_captured_prune_counters() -> Vec { + CAPTURED_PRUNE_COUNTERS.with(|captured| std::mem::take(&mut *captured.borrow_mut())) } fn record_rabit_prune_stats(counters: &RabitPruneCounters) { + #[cfg(test)] + CAPTURED_PRUNE_COUNTERS.with(|captured| captured.borrow_mut().push(*counters)); if !rabit_prune_stats_enabled() { return; } @@ -515,6 +571,21 @@ pub struct RabitQuantizationStorage { ex_scale_factors: Option, high_add_factors: Option, high_scale_factors: Option, + // Built from the sign plane alone for the lazy full-precision scan (see + // `try_from_sign_plane_for_full`): the metadata still says multi-bit but + // no ex codes are loaded, so only the stage-1 bound scan and a survivor + // rerank over separately gathered `ExRows` are valid. + stage1_only: bool, +} + +/// Which columns `from_batch_with_remapper` requires and loads. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PlaneScope { + /// Every column the metadata calls for. + All, + /// Only the sign plane of a layered multi-bit index; the ex planes are + /// supplied per query as [`ExRows`]. + SignOnlyForFull, } impl DeepSizeOf for RabitQuantizationStorage { @@ -536,6 +607,7 @@ impl RabitQuantizationStorage { distance_type: DistanceType, fri: Option>, prepack_ex: bool, + scope: PlaneScope, ) -> Result { let distance_type = match (metadata.query_estimator, distance_type) { (RabitQueryEstimator::RawQuery, DistanceType::Cosine) => DistanceType::L2, @@ -622,7 +694,9 @@ impl RabitQuantizationStorage { } } } - if metadata.layered { + if metadata.layered && scope == PlaneScope::SignOnlyForFull { + super::layered::RQLayout::try_new(metadata.num_bits)?; + } else if metadata.layered { super::layered::RQLayout::try_new(metadata.num_bits)?; for name in [ RABIT_BLOCKED_EX_CODE_LO_COLUMN, @@ -645,7 +719,9 @@ impl RabitQuantizationStorage { "RabitQ low plane requires layered metadata", )); } - if ex_bits != 0 { + if ex_bits != 0 && scope == PlaneScope::SignOnlyForFull { + // The ex planes are gathered per query; see `ExRows`. + } else if ex_bits != 0 { let (normalized_batch, codes, lo) = load_ex_code_planes(batch, metadata.rotated_dim(), metadata.num_bits)?; batch = normalized_batch; @@ -733,6 +809,7 @@ impl RabitQuantizationStorage { packed_ex_codes, ex_add_factors, ex_scale_factors, + stage1_only: scope == PlaneScope::SignOnlyForFull, }; match build_frag_reuse_mapping(fri.as_deref(), &storage.row_ids) { @@ -868,6 +945,18 @@ impl RabitQuantizationStorage { fn distance_calculator_from_parts<'a>( &'a self, parts: RabitDistCalculatorParts<'a>, + ) -> RabitDistCalculator<'a> { + self.distance_calculator_from_parts_with_ex(parts, None) + } + + /// Build the calculator for `parts`. `Some(ex)` replaces this storage's + /// full-precision ex codes and factors with gathered rows, which the + /// calculator then indexes by ex index; everything else (kernel choice, + /// query factors, add-factor transform) is shared with the eager path. + fn distance_calculator_from_parts_with_ex<'a>( + &'a self, + parts: RabitDistCalculatorParts<'a>, + ex: Option>, ) -> RabitDistCalculator<'a> { let RabitDistCalculatorParts { dim, @@ -900,18 +989,30 @@ impl RabitQuantizationStorage { } else { super::layered::RQPrecision::Full }; - let (num_bits, ex_codes, ex_codes_lo, packed_ex_codes, ex_add, ex_scale) = - match rq_precision { - super::layered::RQPrecision::Full => ( + let (num_bits, ex_codes, ex_codes_lo, ex_code_len, packed_ex_codes, ex_add, ex_scale) = + match (rq_precision, ex) { + (super::layered::RQPrecision::Full, Some(ex)) => ( + self.metadata.num_bits, + Some(ex.hi), + Some(ex.lo), + ex.hi_row_bytes, + None, + Some(ex.ex_add), + Some(ex.ex_scale), + ), + (super::layered::RQPrecision::Full, None) => ( self.metadata.num_bits, ex_codes, ex_codes_lo, + ex_code_len, packed_ex_codes, - self.ex_add_factors.as_ref(), - self.ex_scale_factors.as_ref(), + self.ex_add_factors.as_ref().map(|f| f.values().as_ref()), + self.ex_scale_factors.as_ref().map(|f| f.values().as_ref()), ), - super::layered::RQPrecision::Sign => (1, None, None, None, None, None), - super::layered::RQPrecision::High => ( + (super::layered::RQPrecision::Sign, _) => { + (1, None, None, ex_code_len, None, None, None) + } + (super::layered::RQPrecision::High, _) => ( match self.metadata.num_bits { 3 => 2, 4 | 5 => 3, @@ -920,9 +1021,12 @@ impl RabitQuantizationStorage { }, ex_codes, None, + ex_code_len, None, - self.high_add_factors.as_ref(), - self.high_scale_factors.as_ref(), + self.high_add_factors.as_ref().map(|f| f.values().as_ref()), + self.high_scale_factors + .as_ref() + .map(|f| f.values().as_ref()), ), }; let mut calculator = RabitDistCalculator::new( @@ -941,8 +1045,8 @@ impl RabitQuantizationStorage { self.error_factors .as_ref() .map(|factors| factors.values().as_ref()), - ex_add.map(|factors| factors.values().as_ref()), - ex_scale.map(|factors| factors.values().as_ref()), + ex_add, + ex_scale, packed_ex_codes, query_factors.add, query_factors.error, @@ -950,6 +1054,7 @@ impl RabitQuantizationStorage { ); calculator.add_factor_scale = query_factors.add_scale; calculator.add_factor_offset = query_factors.add_offset; + calculator.stage1_only = self.stage1_only; if num_bits > 1 && approx_mode != ApproxMode::Fast && (rq_precision == super::layered::RQPrecision::High || self.error_factors.is_none()) @@ -969,6 +1074,170 @@ impl RabitQuantizationStorage { calculator } + /// Build the stage-2 calculator of a lazy full-precision scan: the same + /// query preparation and constructor as + /// [`VectorStore::dist_calculator_with_scratch`] (so query factors, the + /// ex-dot kernel and the add-factor transform are bit-identical to the + /// eager scan), scoring survivors against `ex_rows` by ex index. Pass the + /// same `query`, `dist_q_c`, `residual` and `options` as stage 1. + pub fn dist_calculator_with_ex_rows<'a>( + &'a self, + query: Arc, + dist_q_c: f32, + residual: Option>, + f32_scratch: &'a mut Vec, + options: DistanceCalculatorOptions, + ex_rows: ExRows<'a>, + ) -> Result> { + if !self.metadata.layered + || options.rq_precision != super::layered::RQPrecision::Full + || options.approx_mode == ApproxMode::Fast + { + return Err(Error::invalid_input(format!( + "gathered ex rows require a layered full-precision gated scan, got layered={} rq_precision={:?} approx_mode={:?}", + self.metadata.layered, options.rq_precision, options.approx_mode + ))); + } + let layout = super::layered::RQLayout::try_new(self.metadata.num_bits)?; + let code_dim = self.code_dim(); + let hi_row_bytes = blocked_ex_code_bytes(code_dim, layout.high_bits); + let lo_row_bytes = blocked_ex_code_bytes(code_dim, layout.low_bits); + let rows = ex_rows.len(); + if ex_rows.hi_row_bytes != hi_row_bytes + || ex_rows.hi.len() != rows * hi_row_bytes + || ex_rows.lo.len() != rows * lo_row_bytes + { + return Err(Error::invalid_input(format!( + "gathered ex rows do not match rotated_dim={code_dim} num_bits={}: {rows} rows, high plane {} bytes at {} per row (expected {hi_row_bytes}), low plane {} bytes (expected {lo_row_bytes} per row)", + self.metadata.num_bits, + ex_rows.hi.len(), + ex_rows.hi_row_bytes, + ex_rows.lo.len(), + ))); + } + Ok(self.dist_calculator_with_scratch_and_ex( + query, + dist_q_c, + residual, + f32_scratch, + options, + Some(ex_rows), + )) + } + + /// Shared body of [`VectorStore::dist_calculator_with_scratch`] and + /// [`Self::dist_calculator_with_ex_rows`]: identical query preparation, + /// with `ex` optionally overriding the full-precision ex rows. + fn dist_calculator_with_scratch_and_ex<'a>( + &'a self, + qr: Arc, + dist_q_c: f32, + residual: Option>, + f32_scratch: &'a mut Vec, + options: DistanceCalculatorOptions, + ex: Option>, + ) -> RabitDistCalculator<'a> { + let code_dim = self.code_dim(); + if let ( + RabitQueryEstimator::RawQuery, + Some(QueryResidual::RabitRawQuery { + rotated_centroid, + query: Some(raw_query), + }), + ) = (self.metadata.query_estimator, residual) + { + debug_assert_eq!(raw_query.code_dim, code_dim); + // Prefix levels share the rotated query; only stored code factors differ. + let query_factors = + self.raw_query_factors(dist_q_c, &raw_query.rotated_query, rotated_centroid); + return self.distance_calculator_from_parts_with_ex( + RabitDistCalculatorParts { + dim: code_dim, + dist_table: Cow::Borrowed(&raw_query.dist_table), + ex_query: Cow::Borrowed(kernel_query( + &raw_query.rotated_query, + &raw_query.ex_query, + )), + sum_q: raw_query.sum_q, + query_factors, + approx_mode: options.approx_mode, + rq_precision: options.rq_precision, + }, + ex, + ); + } + + let dist_table_len = code_dim * 4; + let ex_bits = self.metadata.num_bits - 1; + // The kernels read the rotated query in place; a zero-padded copy is + // only needed when the rotated dim is not block-aligned. + let ex_query_table_len = if ex_bits == 0 || code_dim.is_multiple_of(EX_DOT_BLOCK_DIMS) { + 0 + } else { + padded_query_len(code_dim) + }; + f32_scratch.resize(code_dim + dist_table_len + ex_query_table_len, 0.0); + + let query_factors; + let sum_q = { + let (rotated_qr, remaining) = f32_scratch.split_at_mut(code_dim); + let (dist_table, ex_query) = remaining.split_at_mut(dist_table_len); + match residual { + Some(QueryResidual::Centroid(residual_centroid)) => { + self.rotate_query_vector_into( + code_dim, + &qr, + Some(residual_centroid), + rotated_qr, + ); + } + Some(QueryResidual::RabitRawQuery { .. }) | None => { + self.rotate_query_vector_into(code_dim, &qr, None, rotated_qr); + } + } + query_factors = match (self.metadata.query_estimator, residual) { + (RabitQueryEstimator::ResidualQuery, _) => RabitQueryFactors { + centered_query: None, + add_scale: 1.0, + add_offset: 0.0, + add: self.residual_query_factor(dist_q_c), + error: 0.0, + }, + ( + RabitQueryEstimator::RawQuery, + Some(QueryResidual::RabitRawQuery { + rotated_centroid, .. + }), + ) => self.raw_query_factors(dist_q_c, rotated_qr, rotated_centroid), + (RabitQueryEstimator::RawQuery, _) => { + self.raw_query_factors(dist_q_c, rotated_qr, None) + } + }; + build_dist_table_direct_into::(rotated_qr, dist_table); + if ex_query_table_len > 0 { + pad_query_into(rotated_qr, ex_query); + } + rotated_qr.iter().copied().sum() + }; + + let ex_query_start = code_dim + dist_table_len; + self.distance_calculator_from_parts_with_ex( + RabitDistCalculatorParts { + dim: code_dim, + dist_table: Cow::Borrowed(&f32_scratch[code_dim..ex_query_start]), + ex_query: Cow::Borrowed(kernel_query( + &f32_scratch[..code_dim], + &f32_scratch[ex_query_start..ex_query_start + ex_query_table_len], + )), + sum_q, + query_factors, + approx_mode: options.approx_mode, + rq_precision: options.rq_precision, + }, + ex, + ) + } + fn rotate_query_vector(&self, code_dim: usize, qr: &dyn Array) -> Vec { let mut output = vec![0.0f32; code_dim]; self.rotate_query_vector_into(code_dim, qr, None, &mut output); @@ -1165,6 +1434,162 @@ struct RawQueryTopkContext<'a> { query_upper_bound: f32, } +/// Full-precision ex rows gathered for the stage-2 rerank of a lazy layered +/// scan: blocked high and low plane codes plus the full ex factors, all +/// indexed by the survivor's ex index. Whole planes keep the sign-plane +/// offset as the ex index; a compact gather uses the position in the gather. +#[derive(Debug, Clone, Copy)] +pub struct ExRows<'a> { + hi: &'a [u8], + hi_row_bytes: usize, + lo: &'a [u8], + ex_add: &'a [f32], + ex_scale: &'a [f32], +} + +impl<'a> ExRows<'a> { + /// Borrow the rows of a high-plane batch (plane 1: high codes) and a + /// low-plane batch (plane 2: low codes and the full ex factors) with the + /// same row order, validating code widths as the eager loader does. + pub fn from_plane_batches( + hi: &'a RecordBatch, + lo: &'a RecordBatch, + rotated_dim: usize, + num_bits: u8, + ) -> Result { + let column = |batch: &'a RecordBatch, name: &str| { + batch.column_by_name(name).ok_or_else(|| { + Error::invalid_input(format!("RabitQ gathered ex rows require column {name}")) + }) + }; + let hi_codes = column(hi, RABIT_BLOCKED_EX_CODE_COLUMN)?; + let lo_codes = column(lo, RABIT_BLOCKED_EX_CODE_LO_COLUMN)?; + let factors = |name: &str| -> Result<&'a [f32]> { + let values = column(lo, name)? + .as_primitive_opt::() + .ok_or_else(|| { + Error::invalid_input(format!("RabitQ column {name} must be float32")) + })? + .values(); + Ok(values.as_ref()) + }; + let ex_add = factors(EX_ADD_FACTORS_COLUMN)?; + let ex_scale = factors(EX_SCALE_FACTORS_COLUMN)?; + let rows = hi.num_rows(); + if lo.num_rows() != rows || ex_add.len() != rows || ex_scale.len() != rows { + return Err(Error::invalid_input(format!( + "RabitQ gathered ex rows disagree in length: high plane {rows}, low plane {}, ex add factors {}, ex scale factors {}", + lo.num_rows(), + ex_add.len(), + ex_scale.len() + ))); + } + let planes = RecordBatch::try_new( + Arc::new(arrow_schema::Schema::new(vec![ + Field::new( + RABIT_BLOCKED_EX_CODE_COLUMN, + hi_codes.data_type().clone(), + true, + ), + Field::new( + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + lo_codes.data_type().clone(), + true, + ), + ])), + vec![hi_codes.clone(), lo_codes.clone()], + )?; + load_ex_code_planes(planes, rotated_dim, num_bits)?; + // Validated above as uint8 fixed-size lists of the layout's widths. + let hi_codes = hi_codes.as_fixed_size_list(); + let lo_codes = lo_codes.as_fixed_size_list(); + Ok(Self { + hi: hi_codes.values().as_primitive::().values(), + hi_row_bytes: hi_codes.value_length() as usize, + lo: lo_codes.values().as_primitive::().values(), + ex_add, + ex_scale, + }) + } + + /// Number of gathered rows. + pub fn len(&self) -> usize { + self.ex_add.len() + } + + pub fn is_empty(&self) -> bool { + self.ex_add.is_empty() + } +} + +/// Stage-1 output of a lazy full-precision scan over every row of the +/// original sign plane: binary inner products in the partition's own +/// SIMD/tail layout and the native lower bounds derived from them. +#[derive(Debug, Default, Clone)] +pub struct SignStage { + pub binary_ips: Vec, + pub lower_bounds: Vec, +} + +/// A stage-1 survivor handed to [`RabitDistCalculator::accumulate_survivor_rows`]. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct SurvivorRow { + /// Row position in the [`ExRows`] the stage-2 calculator was built with. + pub ex_index: u32, + pub row_id: u64, + /// Stage-1 binary inner product of the row's original sign-plane offset. + pub binary_ip: f32, + /// Stage-1 lower bound of the row's original sign-plane offset. + pub lower_bound: f32, +} + +/// Rows a stage-1 selection classified, in the categories of the eager +/// per-row step (see [`RabitPruneCounters`]). +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct StagePruneCounts { + /// Accepted rows: every row without a prefilter, otherwise the rows the + /// prefilter selects. + pub candidates: usize, + pub pruned_upper_bound: usize, + pub pruned_heap: usize, +} + +/// Select the stage-1 survivors of a partition in ascending offset order. +/// +/// Row `i` survives iff it is accepted, `lower_bounds[i] < upper`, and +/// `lower_bounds[i] < threshold`, with the upper-bound check first as in the +/// eager per-row step. `upper` is the query's upper distance bound (none +/// means `f32::MAX`, as in the eager scan). `threshold` must be the top of a +/// full-score heap that already held `k` rows — never an estimate — and +/// `None` while the heap is not full. A NaN bound or threshold compares false +/// and keeps the row. Replaces the contents of `survivors` and `counts`. +pub fn select_full_survivors( + lower_bounds: &[f32], + accept: Option<&[bool]>, + upper: Option, + threshold: Option, + survivors: &mut Vec, + counts: &mut StagePruneCounts, +) { + debug_assert!(accept.is_none_or(|accept| accept.len() == lower_bounds.len())); + let upper = upper.unwrap_or(f32::MAX); + survivors.clear(); + *counts = StagePruneCounts::default(); + for (offset, &lower_bound) in lower_bounds.iter().enumerate() { + if accept.is_some_and(|accept| !accept[offset]) { + continue; + } + counts.candidates += 1; + if lower_bound >= upper { + counts.pruned_upper_bound += 1; + } else if threshold.is_some_and(|threshold| lower_bound >= threshold) { + counts.pruned_heap += 1; + } else { + survivors.push(offset as u32); + } + } +} + /// Pick the query slice the ex-dot kernels consume: the rotated query itself /// when the dim is block-aligned, otherwise a zero-padded copy. fn kernel_query<'a>(rotated_query: &'a [f32], padded: &'a [f32]) -> &'a [f32] { @@ -1210,6 +1635,10 @@ pub struct RabitDistCalculator<'a> { add_factor_scale: f32, add_factor_offset: f32, approx_mode: ApproxMode, + // Built on a sign-plane-only storage: ex arrays, when present, are indexed + // by gathered ex index rather than by sign-plane offset, so the dense and + // filtered top-k scans must not run on it. + stage1_only: bool, sum_q: f32, sqrt_d: f32, @@ -1383,6 +1812,7 @@ impl<'a> RabitDistCalculator<'a> { approx_mode, add_factor_scale: 1.0, add_factor_offset: 0.0, + stage1_only: false, sqrt_d: (dim as f32 * num_bits as f32).sqrt(), sum_q, } @@ -1841,6 +2271,10 @@ impl<'a> RabitDistCalculator<'a> { quantized_dists_table: &mut Vec, hacc_quantized_dists: &mut Vec, ) -> Option> { + debug_assert!( + !self.stage1_only, + "a sign-plane-only calculator scores survivors through accumulate_survivor_rows" + ); let code_len = rabit_binary_code_bytes(self.dim); let n = self.codes.len() / code_len; if n == 0 { @@ -1859,7 +2293,19 @@ impl<'a> RabitDistCalculator<'a> { hacc_quantized_dists, ); - Some(RawQueryTopkContext { + Some(self.raw_query_ex_context(n, k, lower_bound, upper_bound)) + } + + /// The loop-invariant exact-rerank inputs shared by the eager top-k scans + /// and the lazy survivor rerank. + fn raw_query_ex_context( + &self, + n: usize, + k: usize, + lower_bound: Option, + upper_bound: Option, + ) -> RawQueryTopkContext<'_> { + RawQueryTopkContext { n, k, ex_bits: self.num_bits - 1, @@ -1874,7 +2320,113 @@ impl<'a> RabitDistCalculator<'a> { .expect("raw-query multi-bit RQ requires ex scale factors"), query_lower_bound: lower_bound.unwrap_or(f32::MIN), query_upper_bound: upper_bound.unwrap_or(f32::MAX), - }) + } + } + + /// Stage 1 of a lazy full-precision scan: the binary inner products and + /// native lower bounds of every row of the original sign plane, computed + /// exactly as the eager dense scan computes them (same SIMD/tail split and + /// LUT quantization, exact fallback, and Accurate high-accuracy LUT). + /// + /// Returns `Err(reason)` without touching `out` when this calculator + /// cannot gate on the lower bound (the eager scan then scores every row, + /// so the caller must take the eager path for the partition). + pub fn full_sign_stage_with_scratch( + &self, + out: &mut SignStage, + quantized_dists: &mut Vec, + quantized_dists_table: &mut Vec, + hacc_quantized_dists: &mut Vec, + ) -> std::result::Result<(), &'static str> { + if let Some(reason) = self.raw_query_lower_bound_gating_disabled_reason() { + return Err(reason); + } + let Some(error_factors) = self.pruning_errors() else { + return Err("missing_error_factors"); + }; + let code_len = rabit_binary_code_bytes(self.dim); + let n = self.codes.len() / code_len; + out.lower_bounds.clear(); + if n == 0 { + out.binary_ips.clear(); + return Ok(()); + } + self.binary_distances_with_scratch( + n, + code_len, + &mut out.binary_ips, + quantized_dists, + quantized_dists_table, + hacc_quantized_dists, + ); + // Same expression as `raw_query_lower_bound`, so the survivors' + // bounds match the eager scan's (and its prune masks) bit for bit. + out.lower_bounds + .extend(out.binary_ips.iter().enumerate().map(|(id, &binary_ip)| { + self.raw_query_binary_distance(id, binary_ip) - error_factors[id] * self.query_error + })); + Ok(()) + } + + /// Stage 2 of a lazy full-precision scan for one partition. `rows` are + /// the partition's stage-1 survivors in ascending sign-plane offset order; + /// each goes through the eager scan's per-row step (bound checks against + /// the live heap, exact rerank, range check, heap update) with its + /// stage-1 inner product and bound. `stage1` carries the rows stage 1 + /// already pruned so the partition reports the same prune counters as the + /// eager scan; they are returned together with the stage-2 tallies. + /// + /// `on_first_full` receives the heap top once, when this partition first + /// fills the heap to `k` rows. + /// + /// The calculator must come from + /// [`RabitQuantizationStorage::dist_calculator_with_ex_rows`] with the + /// rows the survivors' `ex_index` refers to. + #[allow(clippy::too_many_arguments)] + pub fn accumulate_survivor_rows( + &self, + k: usize, + lower_bound: Option, + upper_bound: Option, + rows: &[SurvivorRow], + stage1: StagePruneCounts, + res: &mut BinaryHeap>, + mut on_first_full: impl FnMut(f32), + ) -> RabitPruneCounters { + let mut counters = RabitPruneCounters { + candidates: stage1.candidates, + pruned_upper_bound: stage1.pruned_upper_bound, + pruned_heap: stage1.pruned_heap, + ..Default::default() + }; + let n = self.codes.len() / rabit_binary_code_bytes(self.dim); + // The eager scans return before reporting for these, too. + if k == 0 || n == 0 { + return counters; + } + let ctx = self.raw_query_ex_context(n, k, lower_bound, upper_bound); + let mut max_dist = res.peek().map(|node| node.dist); + let mut is_full = res.len() >= k; + for row in rows { + self.accumulate_raw_query_multi_bit_row( + &ctx, + row.ex_index as usize, + row.row_id, + row.binary_ip, + row.lower_bound, + res, + &mut max_dist, + &mut counters, + ); + if !is_full && res.len() >= k { + is_full = true; + if let Some(top) = res.peek() { + on_first_full(top.dist.0); + } + } + } + record_rabit_prune_stats(&counters); + counters } /// Process one candidate row given its lower bound: the bound checks, @@ -2731,99 +3283,7 @@ impl VectorStore for RabitQuantizationStorage { f32_scratch: &'a mut Vec, options: DistanceCalculatorOptions, ) -> Self::DistanceCalculator<'a> { - let code_dim = self.code_dim(); - if let ( - RabitQueryEstimator::RawQuery, - Some(QueryResidual::RabitRawQuery { - rotated_centroid, - query: Some(raw_query), - }), - ) = (self.metadata.query_estimator, residual) - { - debug_assert_eq!(raw_query.code_dim, code_dim); - // Prefix levels share the rotated query; only stored code factors differ. - let query_factors = - self.raw_query_factors(dist_q_c, &raw_query.rotated_query, rotated_centroid); - return self.distance_calculator_from_parts(RabitDistCalculatorParts { - dim: code_dim, - dist_table: Cow::Borrowed(&raw_query.dist_table), - ex_query: Cow::Borrowed(kernel_query( - &raw_query.rotated_query, - &raw_query.ex_query, - )), - sum_q: raw_query.sum_q, - query_factors, - approx_mode: options.approx_mode, - rq_precision: options.rq_precision, - }); - } - - let dist_table_len = code_dim * 4; - let ex_bits = self.metadata.num_bits - 1; - // The kernels read the rotated query in place; a zero-padded copy is - // only needed when the rotated dim is not block-aligned. - let ex_query_table_len = if ex_bits == 0 || code_dim.is_multiple_of(EX_DOT_BLOCK_DIMS) { - 0 - } else { - padded_query_len(code_dim) - }; - f32_scratch.resize(code_dim + dist_table_len + ex_query_table_len, 0.0); - - let query_factors; - let sum_q = { - let (rotated_qr, remaining) = f32_scratch.split_at_mut(code_dim); - let (dist_table, ex_query) = remaining.split_at_mut(dist_table_len); - match residual { - Some(QueryResidual::Centroid(residual_centroid)) => { - self.rotate_query_vector_into( - code_dim, - &qr, - Some(residual_centroid), - rotated_qr, - ); - } - Some(QueryResidual::RabitRawQuery { .. }) | None => { - self.rotate_query_vector_into(code_dim, &qr, None, rotated_qr); - } - } - query_factors = match (self.metadata.query_estimator, residual) { - (RabitQueryEstimator::ResidualQuery, _) => RabitQueryFactors { - centered_query: None, - add_scale: 1.0, - add_offset: 0.0, - add: self.residual_query_factor(dist_q_c), - error: 0.0, - }, - ( - RabitQueryEstimator::RawQuery, - Some(QueryResidual::RabitRawQuery { - rotated_centroid, .. - }), - ) => self.raw_query_factors(dist_q_c, rotated_qr, rotated_centroid), - (RabitQueryEstimator::RawQuery, _) => { - self.raw_query_factors(dist_q_c, rotated_qr, None) - } - }; - build_dist_table_direct_into::(rotated_qr, dist_table); - if ex_query_table_len > 0 { - pad_query_into(rotated_qr, ex_query); - } - rotated_qr.iter().copied().sum() - }; - - let ex_query_start = code_dim + dist_table_len; - self.distance_calculator_from_parts(RabitDistCalculatorParts { - dim: code_dim, - dist_table: Cow::Borrowed(&f32_scratch[code_dim..ex_query_start]), - ex_query: Cow::Borrowed(kernel_query( - &f32_scratch[..code_dim], - &f32_scratch[ex_query_start..ex_query_start + ex_query_table_len], - )), - sum_q, - query_factors, - approx_mode: options.approx_mode, - rq_precision: options.rq_precision, - }) + self.dist_calculator_with_scratch_and_ex(qr, dist_q_c, residual, f32_scratch, options, None) } // TODO: implement this @@ -3118,7 +3578,14 @@ impl QuantizerStorage for RabitQuantizationStorage { // The projected store is reconstructed per query. Its blocked kernels // consume codes directly, so rebuilding an ex FastScan transpose here // would cost more than scoring the requested prefix. - Self::from_batch_with_remapper(batch, &metadata, distance_type, remapper, false) + Self::from_batch_with_remapper( + batch, + &metadata, + distance_type, + remapper, + false, + PlaneScope::All, + ) } type Metadata = RabitQuantizationMetadata; @@ -3139,7 +3606,38 @@ impl QuantizerStorage for RabitQuantizationStorage { distance_type: DistanceType, fri: Option>, ) -> Result { - Self::from_batch_with_remapper(batch, metadata, distance_type, fri, true) + Self::from_batch_with_remapper(batch, metadata, distance_type, fri, true, PlaneScope::All) + } + + /// 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 + /// 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( + batch: RecordBatch, + metadata: &Self::Metadata, + distance_type: DistanceType, + ) -> Result { + if !metadata.layered || metadata.num_bits <= 1 { + return Err(Error::invalid_input(format!( + "sign-plane full-precision storage requires a layered multi-bit IVF_RQ index, got layered={} num_bits={}", + metadata.layered, metadata.num_bits + ))); + } + if batch.column_by_name(ERROR_FACTORS_COLUMN).is_none() { + return Err(Error::invalid_input(format!( + "sign-plane full-precision storage requires the {ERROR_FACTORS_COLUMN} column for native lower-bound pruning" + ))); + } + Self::from_batch_with_remapper( + batch, + metadata, + distance_type, + None, + false, + PlaneScope::SignOnlyForFull, + ) } fn metadata(&self) -> &Self::Metadata { @@ -3251,6 +3749,7 @@ impl QuantizerStorage for RabitQuantizationStorage { ex_add_factors, ex_scale_factors, row_ids: new_row_ids, + stage1_only: false, }) } } diff --git a/rust/lance-index/src/vector/quantizer.rs b/rust/lance-index/src/vector/quantizer.rs index 221b066d1b2..b1f434453ed 100644 --- a/rust/lance-index/src/vector/quantizer.rs +++ b/rust/lance-index/src/vector/quantizer.rs @@ -302,6 +302,19 @@ pub trait QuantizerStorage: Clone + Sized + DeepSizeOf + VectorStore { Self::try_from_batch_with_remapper(batch, metadata, distance_type, remapper) } + /// Materialize only the sign plane of a layered partition for the lazy + /// full-precision scan, which bounds every row from it before reading the + /// ex planes of the survivors. Only layered IVF_RQ storage supports it. + fn try_from_sign_plane_for_full( + _batch: RecordBatch, + _metadata: &Self::Metadata, + _distance_type: DistanceType, + ) -> Result { + Err(Error::not_supported( + "sign-plane-only storage requires a layered IVF_RQ index".to_string(), + )) + } + fn metadata(&self) -> &Self::Metadata; fn remap(&self, mapping: &RowAddrRemap) -> Result { diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 8da3e6e448b..ab3bc6125ab 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -3,11 +3,14 @@ //! Vector Storage, holding (quantized) vectors and providing distance calculation. -use crate::vector::bq::layered::{PlaneBatch, PlaneKey}; +use crate::vector::bq::ex_dot::blocked_ex_code_bytes; +use crate::vector::bq::layered::{PlaneBatch, PlaneKey, RQLayout}; +use crate::vector::bq::layered_stats; use crate::vector::quantizer::QuantizerStorage; use arrow::compute::concat_batches; use arrow_array::{ - Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, cast::AsArray, types::UInt8Type, + Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array, cast::AsArray, + types::UInt8Type, }; use arrow_schema::SchemaRef; use futures::stream::{self, Stream, StreamExt, TryStreamExt}; @@ -27,11 +30,13 @@ use prost::Message; use std::{ any::Any, borrow::Cow, - collections::{BinaryHeap, HashMap}, + collections::{BinaryHeap, HashMap, HashSet}, mem::size_of, ops::{Deref, DerefMut}, - sync::{Arc, Mutex}, + sync::{Arc, LazyLock, Mutex, OnceLock}, + time::Instant, }; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use crossbeam_queue::ArrayQueue; @@ -624,14 +629,446 @@ impl DeepSizeOf for PlaneAccess { /// 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(Arc>>); +pub struct PlaneAccessTracker { + /// Cascade candidate reads, see [`PLANE_PROMOTION_READS`]. + accesses: Arc>>, + /// Background promotions of the lazy full-precision scan. + lazy: Arc, +} impl DeepSizeOf for PlaneAccessTracker { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.0.deep_size_of_children(context) + self.accesses.deep_size_of_children(context) + self.lazy.deep_size_of_children(context) + } +} + +#[derive(Debug, Default)] +struct LazyPromotionHistory { + /// Sparse gathers of each non-resident `(partition, plane)` since its last promotion. + sparse_reads: HashMap<(usize, u8), u32>, + in_flight: HashSet<(usize, u8)>, +} + +#[derive(Debug, Default)] +struct LazyPromotionState { + history: Mutex, + /// Bounds concurrent promotions; sized by the first promoting query's config. + permits: OnceLock>, +} + +impl DeepSizeOf for LazyPromotionState { + fn deep_size_of_children(&self, _: &mut lance_core::deepsize::Context) -> usize { + let history = self.history.lock().unwrap_or_else(|e| e.into_inner()); + history.sparse_reads.capacity() * size_of::<((usize, u8), u32)>() + + history.in_flight.capacity() * size_of::<(usize, u8)>() + } +} + +/// A started background promotion of one ex plane. Dropping it, when the +/// promotion finishes or is abandoned, frees its in-flight slot and permit. +#[derive(Debug)] +pub struct LazyPromotionTicket { + pub partition: usize, + pub plane: u8, + state: Arc, + _permit: OwnedSemaphorePermit, +} + +impl Drop for LazyPromotionTicket { + fn drop(&mut self) { + self.state + .history + .lock() + .unwrap_or_else(|e| e.into_inner()) + .in_flight + .remove(&(self.partition, self.plane)); + layered_stats::promotion_finished(); + } +} + +/// 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. +pub const LAZY_WINDOW_ENV: &str = "LANCE_RQ_LAZY_WINDOW"; +/// Ex-plane gather policy: `cost`, `sparse` or `whole`. +pub const LAZY_DENSE_ENV: &str = "LANCE_RQ_LAZY_DENSE"; +/// Most coalesced row runs a sparse gather may issue per plane. +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. +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"; +/// Largest survivor count a scoring batch scores on the query task instead of the CPU pool. +pub const LAZY_INLINE_ROWS_ENV: &str = "LANCE_RQ_LAZY_INLINE_ROWS"; +/// Whether a gather whose window gate opens before the heap holds `k` rows +/// may be issued then, selecting every accepted row as the eager load does +/// (`0` or `1`, default `1`). It is issued early when waiting would not save +/// reads, because the gathers issued once the heap fills are expected to be +/// 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. +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. +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"; +/// 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 +/// planes at once. Read once per process. +pub const SEQUENTIAL_PLANE_LOADS_ENV: &str = "LANCE_RQ_SEQUENTIAL_PLANE_LOADS"; + +const DEFAULT_LAZY_WINDOW: usize = 16; +const DEFAULT_LAZY_MAX_RUNS: usize = 16; +const DEFAULT_LAZY_DENSE_BYTES_FRACTION: f64 = 0.5; +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; + +/// [`SEQUENTIAL_PLANE_LOADS_ENV`], read once per process. +static SEQUENTIAL_PLANE_LOADS: LazyLock> = LazyLock::new(|| { + sequential_plane_loads_from(std::env::var(SEQUENTIAL_PLANE_LOADS_ENV).ok().as_deref()) + .map_err(|err| err.to_string()) +}); + +/// Whether layered partition loads read the sign plane before the ex planes on +/// every backend ([`SEQUENTIAL_PLANE_LOADS_ENV`]). The variable is read once +/// per process; an invalid value fails here and every layered partition load. +/// Call it at startup to fail before serving. +pub fn sequential_plane_loads() -> Result { + SEQUENTIAL_PLANE_LOADS.clone().map_err(Error::invalid_input) +} + +fn sequential_plane_loads_from(value: Option<&str>) -> Result { + value.map_or(Ok(false), |value| { + parse_flag(SEQUENTIAL_PLANE_LOADS_ENV, value) + }) +} + +/// Parse a `0`/`1` (or `false`/`true`) environment flag. +fn parse_flag(name: &str, value: &str) -> Result { + match value.trim() { + "1" | "true" => Ok(true), + "0" | "false" => Ok(false), + _ => Err(Error::invalid_input(format!( + "{name}={value:?} is invalid, expected 0 or 1" + ))), } } +/// Page size the gather cost model aligns row runs to, matching the +/// persistent cache's aligned reads. +const LAZY_GATHER_PAGE_BYTES: usize = 4096; +/// Rows closer than this share one read, as in the plane cache codec. +const LAZY_GATHER_COALESCE_GAP_BYTES: usize = 4096; +/// Estimated cache envelope and plane header ahead of the first row. It only +/// shifts page boundaries in the cost estimate. +const LAZY_GATHER_BODY_OFFSET_BYTES: usize = 64; +/// Bytes a cached ex-plane row stores besides its codes: add and scale factors. +const EX_PLANE_FACTOR_BYTES: usize = 2 * size_of::(); + +/// How the lazy scan chooses between reading selected ex-plane rows and the whole plane. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DenseGatherMode { + /// Read the whole plane when the row runs are too many or cover too much + /// of it (see [`plan_plane_gather`]). + Cost, + /// Always read selected rows; no probe is routed to the eager scan + /// ([`LAZY_DENSE_TO_EAGER_ENV`]). + Sparse, + /// Always read the whole plane. + Whole, +} + +/// Whether sparse gathers promote whole planes into RAM in the background. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LazyPromotion { + /// Only whole-plane gathers admit planes. + 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 }, +} + +/// 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`]. + pub window: usize, + pub dense: DenseGatherMode, + pub max_runs: usize, + pub dense_bytes_fraction: f64, + pub promote: LazyPromotion, + pub promote_inflight: usize, + pub inline_rows: usize, + /// See [`LAZY_EAGER_BEFORE_FULL_ENV`]. + pub eager_before_full: bool, + /// See [`LAZY_ORIGIN_MAX_RUNS_ENV`]; `usize::MAX` never falls back. + pub origin_max_runs: usize, + /// See [`LAZY_DENSE_TO_EAGER_ENV`]. + pub dense_to_eager: bool, +} + +impl Default for LayeredLazyConfig { + fn default() -> Self { + Self { + enabled: false, + window: DEFAULT_LAZY_WINDOW, + 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_inflight: DEFAULT_LAZY_PROMOTE_INFLIGHT, + inline_rows: DEFAULT_LAZY_INLINE_ROWS, + eager_before_full: true, + origin_max_runs: usize::MAX, + dense_to_eager: true, + } + } +} + +impl LayeredLazyConfig { + /// Read the `LANCE_RQ_LAZY_*` variables; unset variables keep their defaults. + /// The other knobs are read only when [`LAZY_FULL_ENV`] enables the scan, + /// so with it off they are ignored exactly as before the scan existed. + pub fn from_env() -> Result { + Self::from_lookup(|name| std::env::var(name).ok()) + } + + fn from_lookup(lookup: impl Fn(&str) -> Option) -> Result { + let invalid = |name: &str, value: &str, expected: &str| { + Error::invalid_input(format!("{name}={value:?} is invalid, expected {expected}")) + }; + let count = |name: &str, default: usize, min: usize| -> Result { + let Some(value) = lookup(name) else { + return Ok(default); + }; + value + .trim() + .parse::() + .ok() + .filter(|parsed| *parsed >= min) + .ok_or_else(|| invalid(name, &value, &format!("an integer >= {min}"))) + }; + let mut config = Self::default(); + if let Some(value) = lookup(LAZY_FULL_ENV) { + config.enabled = parse_flag(LAZY_FULL_ENV, &value)?; + } + if !config.enabled { + return Ok(config); + } + config.window = count(LAZY_WINDOW_ENV, config.window, 0)?; + if let Some(value) = lookup(LAZY_DENSE_ENV) { + config.dense = match value.trim() { + "cost" => DenseGatherMode::Cost, + "sparse" => DenseGatherMode::Sparse, + "whole" => DenseGatherMode::Whole, + _ => return Err(invalid(LAZY_DENSE_ENV, &value, "cost, sparse or whole")), + }; + } + config.max_runs = count(LAZY_MAX_RUNS_ENV, config.max_runs, 0)?; + if let Some(value) = lookup(LAZY_DENSE_BYTES_FRACTION_ENV) { + config.dense_bytes_fraction = value + .trim() + .parse::() + .ok() + .filter(|fraction| fraction.is_finite() && *fraction >= 0.0) + .ok_or_else(|| { + invalid( + LAZY_DENSE_BYTES_FRACTION_ENV, + &value, + "a finite number >= 0", + ) + })?; + } + if let Some(value) = lookup(LAZY_PROMOTE_ENV) { + config.promote = match value.trim() { + "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"))?, + }; + } + config.promote_inflight = count(LAZY_PROMOTE_INFLIGHT_ENV, config.promote_inflight, 1)?; + config.inline_rows = count(LAZY_INLINE_ROWS_ENV, config.inline_rows, 0)?; + if let Some(value) = lookup(LAZY_EAGER_BEFORE_FULL_ENV) { + 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_DENSE_TO_EAGER_ENV) { + config.dense_to_eager = parse_flag(LAZY_DENSE_TO_EAGER_ENV, &value)?; + } + Ok(config) + } +} + +/// Where the lazy scan reads one ex plane of a partition from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PlaneSource { + /// The RAM-resident entry, read like a native memory hit. Falls back to + /// [`Self::Whole`] if the entry was evicted since it was planned. + Resident, + /// The whole plane, admitted like the eager full-precision load. + Whole, + /// Selected rows of the persistent entry (or the origin file), without + /// RAM admission. + Sparse, +} + +/// Sources for the high (index 0) and low (index 1) ex planes of one gather. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct GatherPlan { + pub planes: [PlaneSource; 2], + /// Per plane, whether a [`PlaneSource::Sparse`] read that the persistent + /// tier cannot serve loads the whole plane (see + /// [`origin_reads_whole_plane`]) instead of reading the rows from the + /// origin file. Backends that gate plane admission ignore it: they would + /// read the whole plane without admitting it, so every query would pay + /// for it again. + pub origin_whole: [bool; 2], +} + +/// Row geometry of a layered index's ex planes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct LayeredExLayout { + pub rotated_dim: usize, + pub num_bits: u8, + /// Persistent-cache bytes per row of the high and low planes (codes plus factors). + pub row_bytes: [usize; 2], +} + +/// Choose a sparse or whole read of one non-resident ex plane for sorted, +/// unique `rows` of a partition with `partition_rows` rows. +/// +/// Rows are grouped into runs with the plane cache codec's gap rule and each +/// run is rounded out to whole pages. The plane is read whole when that takes +/// more than `max_runs` runs, or when the pages cover at least +/// `dense_bytes_fraction` of the plane. +pub fn plan_plane_gather( + rows: &[u32], + partition_rows: usize, + row_bytes: usize, + config: &LayeredLazyConfig, +) -> PlaneSource { + match config.dense { + DenseGatherMode::Sparse => return PlaneSource::Sparse, + DenseGatherMode::Whole => return PlaneSource::Whole, + DenseGatherMode::Cost => {} + } + let (runs, pages) = sparse_read_extent(rows, row_bytes); + let plane_bytes = LAZY_GATHER_BODY_OFFSET_BYTES + partition_rows * row_bytes; + if runs > config.max_runs + || (pages * LAZY_GATHER_PAGE_BYTES) as f64 + >= config.dense_bytes_fraction * plane_bytes as f64 + { + PlaneSource::Whole + } else { + PlaneSource::Sparse + } +} + +/// Whether a sparse read of sorted, unique `rows` that the persistent tier +/// cannot serve should load the whole plane instead of reading the rows from +/// the origin file: when the rows take more than +/// [`LayeredLazyConfig::origin_max_runs`] runs (grouped as in +/// [`plan_plane_gather`]). +pub fn origin_reads_whole_plane( + rows: &[u32], + row_bytes: usize, + config: &LayeredLazyConfig, +) -> bool { + config.origin_max_runs != usize::MAX + && sparse_read_extent(rows, row_bytes).0 > config.origin_max_runs +} + +/// Row runs and pages of a sparse read of sorted, unique `rows`: rows closer +/// than the plane cache codec's gap share a run, and every run is rounded out +/// to whole pages. +fn sparse_read_extent(rows: &[u32], row_bytes: usize) -> (usize, usize) { + let gap_rows = LAZY_GATHER_COALESCE_GAP_BYTES / row_bytes.max(1); + let page = + |row: usize| (LAZY_GATHER_BODY_OFFSET_BYTES + row * row_bytes) / LAZY_GATHER_PAGE_BYTES; + let page_end = |row: usize| { + (LAZY_GATHER_BODY_OFFSET_BYTES + row * row_bytes).div_ceil(LAZY_GATHER_PAGE_BYTES) + }; + let mut runs = 0usize; + let mut pages = 0usize; + let mut start = 0; + while start < rows.len() { + let mut end = start + 1; + while end < rows.len() && (rows[end] - rows[end - 1] - 1) as usize <= gap_rows { + end += 1; + } + pages += page_end(rows[end - 1] as usize + 1) - page(rows[start] as usize); + runs += 1; + start = end; + } + (runs, pages) +} + +/// How stage-2 survivors index a [`GatheredEx`]'s batches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ExIndex { + /// Both planes are whole: a survivor's ex index is its partition offset. + Identity, + /// The batches hold the gathered rows only, in gather order. + Compact, +} + +/// The ex-plane rows of one partition's survivors. +#[derive(Debug)] +pub struct GatheredEx { + /// High-plane batch (codes and high factors). + pub high: RecordBatch, + /// Low-plane batch (codes and the full ex factors). + pub low: RecordBatch, + pub index: ExIndex, + /// The source that served each plane, `None` when no rows were needed. + pub sources: [Option; 2], + /// Planes whose sparse rows came from the origin file. + pub origin_row_reads: usize, + /// Background promotions the caller must run (and then drop). + pub promotions: Vec, +} + +/// One ex plane fetched for a gather. +struct FetchedPlane { + batch: RecordBatch, + source: PlaneSource, + from_origin: bool, + promotion: Option, +} + /// Loader to load partitioned PQ storage from disk. #[derive(Debug)] pub struct IvfQuantizationStorage { @@ -977,13 +1414,21 @@ impl IvfQuantizationStorage { RQPrecision::Full => 2, }; let load_plane = |plane| self.load_plane_entry(part_id, plane, cache, io_stats.clone()); - // Admit the sign dependency first. The high and low 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?; + let planes = if sequential_plane_loads()? || cache.plane_admission_gated() { + // Admit the sign dependency first. The high and low 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() + } 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? + }; let mut fields = Vec::new(); let mut columns = Vec::new(); - for batch in std::iter::once(sign).chain(ex) { + for batch in planes { fields.extend(batch.0.schema().fields().iter().cloned()); columns.extend(batch.0.columns().iter().cloned()); } @@ -1019,7 +1464,7 @@ impl IvfQuantizationStorage { /// plane is read without admission, from the persistent entry or the /// origin file. Other backends admit every plane through their ordinary /// policy, as native partitions are. - async fn load_plane_entry( + pub async fn load_plane_entry( &self, part_id: usize, plane: u8, @@ -1071,6 +1516,266 @@ impl IvfQuantizationStorage { Ok(batch) } + /// Ex-plane geometry of a layered index. + pub fn layered_ex_layout(&self) -> Result { + let Quantizer::Rabit(rq) = self.quantizer()? else { + return Err(Error::invalid_input( + "ex-plane layout requires a layered IVF_RQ index", + )); + }; + let metadata = rq.metadata_ref(); + if !metadata.layered { + return Err(Error::invalid_input( + "ex-plane layout requires a layered IVF_RQ index", + )); + } + let layout = RQLayout::try_new(metadata.num_bits)?; + let rotated_dim = metadata.rotated_dim(); + Ok(LayeredExLayout { + rotated_dim, + num_bits: metadata.num_bits, + row_bytes: [layout.high_bits, layout.low_bits] + .map(|bits| blocked_ex_code_bytes(rotated_dim, bits) + EX_PLANE_FACTOR_BYTES), + }) + } + + /// Stage-1 storage of the lazy full-precision scan: the partition's sign + /// plane alone, loaded (and admitted) like the eager load's sign plane. + pub async fn load_sign_stage( + &self, + part_id: usize, + cache: &WeakLanceCache, + io_stats: Option, + ) -> Result { + if !self.supports_candidate_reads() { + return Err(Error::invalid_input( + "the lazy full-precision scan requires an index without a row-id remapper", + )); + } + let sign = self.load_plane_entry(part_id, 0, cache, io_stats).await?; + Q::Storage::try_from_sign_plane_for_full( + sign.0.clone(), + self.metadata(), + self.distance_type, + ) + } + + /// Fetch the high and low ex-plane rows at the sorted, unique partition + /// offsets `rows`, each plane from the source `plan` names. + /// + /// When both planes end up whole the batches are the whole planes and + /// survivors keep their offsets as ex index ([`ExIndex::Identity`]); + /// otherwise whole planes are cut down to `rows` so both batches hold the + /// gathered rows in order ([`ExIndex::Compact`]). No rows means no reads. + #[allow(clippy::too_many_arguments)] + pub async fn gather_ex_rows( + &self, + part_id: usize, + rows: &[u32], + plan: GatherPlan, + config: &LayeredLazyConfig, + cache: &WeakLanceCache, + io_stats: Option, + ) -> Result { + let partition_rows = self.partition_size(part_id); + if rows.windows(2).any(|pair| pair[0] >= pair[1]) + || rows + .last() + .is_some_and(|&row| row as usize >= partition_rows) + { + return Err(Error::invalid_input(format!( + "gathered offsets of partition {part_id} must be sorted, unique and below its {partition_rows} rows" + ))); + } + if rows.is_empty() { + let (high, low) = futures::try_join!( + self.read_plane(part_id, 1, Some(Vec::new()), None), + self.read_plane(part_id, 2, Some(Vec::new()), None), + )?; + return Ok(GatheredEx { + high, + low, + index: ExIndex::Compact, + sources: [None; 2], + origin_row_reads: 0, + promotions: Vec::new(), + }); + } + let (high, low) = futures::try_join!( + self.fetch_ex_plane( + part_id, + 1, + rows, + plan.planes[0], + plan.origin_whole[0], + config, + cache, + io_stats.clone() + ), + self.fetch_ex_plane( + part_id, + 2, + rows, + plan.planes[1], + plan.origin_whole[1], + config, + cache, + io_stats + ), + )?; + let index = if high.source != PlaneSource::Sparse && low.source != PlaneSource::Sparse { + ExIndex::Identity + } else { + ExIndex::Compact + }; + let indices = UInt32Array::from(rows.to_vec()); + let compact = |plane: &FetchedPlane| -> Result { + if index == ExIndex::Compact && plane.source != PlaneSource::Sparse { + Ok(plane.batch.take(&indices)?) + } else { + Ok(plane.batch.clone()) + } + }; + Ok(GatheredEx { + high: compact(&high)?, + low: compact(&low)?, + index, + sources: [Some(high.source), Some(low.source)], + origin_row_reads: usize::from(high.from_origin) + usize::from(low.from_origin), + promotions: [high.promotion, low.promotion] + .into_iter() + .flatten() + .collect(), + }) + } + + /// Fetch one ex plane from `source`; `origin_whole` is the plan's + /// [`GatherPlan::origin_whole`] for the plane. + #[allow(clippy::too_many_arguments)] + async fn fetch_ex_plane( + &self, + part_id: usize, + plane: u8, + rows: &[u32], + source: PlaneSource, + origin_whole: bool, + config: &LayeredLazyConfig, + cache: &WeakLanceCache, + io_stats: Option, + ) -> Result { + let stats = layered_stats::counters(); + let started = Instant::now(); + let key = PlaneKey { + partition: part_id, + plane, + }; + if source == PlaneSource::Resident + && let Some(batch) = cache.get_resident_with_key(&key).await + { + stats.fetch_resident_ns.add_elapsed(started); + return Ok(FetchedPlane { + batch: batch.0.clone(), + source: PlaneSource::Resident, + from_origin: false, + promotion: None, + }); + } + if source == PlaneSource::Sparse { + let persisted = cache.get_rows_with_key(&key, rows).await; + if persisted.is_some() || !origin_whole || cache.plane_admission_gated() { + 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) + .await?, + true, + ), + }; + if batch.num_rows() != rows.len() { + return Err(Error::internal(format!( + "sparse gather of partition {part_id} plane {plane} returned {} rows for {} offsets", + batch.num_rows(), + rows.len() + ))); + } + stats.fetch_sparse_ns.add_elapsed(started); + return Ok(FetchedPlane { + batch, + source: PlaneSource::Sparse, + from_origin, + promotion: self.request_lazy_promotion(part_id, plane, config, cache), + }); + } + // Too many origin row runs: load the whole plane instead, which + // this backend admits like the eager load, so later queries find + // it cached. + stats.origin_whole_fallbacks.incr(); + } + let batch = self + .load_plane_entry(part_id, plane, cache, io_stats) + .await?; + stats.fetch_whole_ns.add_elapsed(started); + Ok(FetchedPlane { + batch: batch.0.clone(), + source: PlaneSource::Whole, + from_origin: false, + promotion: None, + }) + } + + /// Record a sparse gather of a non-resident plane and, once the policy's + /// read count is reached, start a promotion unless one for the plane is + /// already running or every in-flight permit is taken. + /// + /// A gated backend admits ex planes only behind their resident sign + /// plane and evicts the lowest-priority plane first, so a promoted plane + /// would mostly be read whole and dropped again; promotions run only on + /// backends that admit planes through their ordinary policy. + fn request_lazy_promotion( + &self, + part_id: usize, + plane: u8, + config: &LayeredLazyConfig, + cache: &WeakLanceCache, + ) -> Option { + let LazyPromotion::Background { reads } = config.promote else { + return None; + }; + if cache.plane_admission_gated() { + return None; + } + let stats = layered_stats::counters(); + let state = &self.plane_access.lazy; + let key = (part_id, plane); + let mut history = state.history.lock().unwrap_or_else(|e| e.into_inner()); + if history.in_flight.contains(&key) { + stats.promotions_deduped.incr(); + return None; + } + let count = history.sparse_reads.entry(key).or_default(); + *count = count.saturating_add(1); + if *count < reads { + return None; + } + let permits = state + .permits + .get_or_init(|| Arc::new(Semaphore::new(config.promote_inflight))); + let Ok(permit) = permits.clone().try_acquire_owned() else { + stats.promotions_skipped.incr(); + return None; + }; + history.sparse_reads.remove(&key); + history.in_flight.insert(key); + stats.promotions_issued.incr(); + layered_stats::promotion_started(); + Some(LazyPromotionTicket { + partition: part_id, + plane, + state: state.clone(), + _permit: permit, + }) + } + /// Assemble a full-precision store from candidate rows only. pub async fn load_candidates( &self, @@ -1156,7 +1861,7 @@ impl IvfQuantizationStorage { let promote = { let mut accesses = self .plane_access - .0 + .accesses .lock() .unwrap_or_else(|e| e.into_inner()); let access = accesses.entry((part_id, plane)).or_default(); @@ -1591,11 +2296,18 @@ impl IvfQuantizationStorage { #[cfg(test)] mod tests { use super::{ - QueryScratchCapacity, QueryScratchPool, compact_prewarm_batches, + 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, + 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, }; use arrow_array::{Array, ArrayRef, RecordBatch, UInt64Array}; + use lance_core::Error; use lance_core::deepsize::DeepSizeOf; + use std::collections::HashMap; use std::sync::Arc; #[tokio::test] @@ -1726,6 +2438,209 @@ mod tests { assert_eq!(pool.deep_size_of(), idle_size); } + #[test] + fn test_layered_lazy_config_parses_every_knob() { + let env = HashMap::from([ + (LAZY_FULL_ENV, "1"), + (LAZY_WINDOW_ENV, "4"), + (LAZY_DENSE_ENV, "sparse"), + (LAZY_MAX_RUNS_ENV, "0"), + (LAZY_DENSE_BYTES_FRACTION_ENV, "0.25"), + (LAZY_PROMOTE_ENV, "bg:2"), + (LAZY_PROMOTE_INFLIGHT_ENV, "8"), + (LAZY_INLINE_ROWS_ENV, "0"), + (LAZY_EAGER_BEFORE_FULL_ENV, "0"), + (LAZY_ORIGIN_MAX_RUNS_ENV, "5"), + (LAZY_DENSE_TO_EAGER_ENV, "0"), + ]); + let config = + LayeredLazyConfig::from_lookup(|name| env.get(name).map(|value| value.to_string())) + .unwrap(); + assert_eq!( + config, + LayeredLazyConfig { + enabled: true, + window: 4, + dense: DenseGatherMode::Sparse, + max_runs: 0, + dense_bytes_fraction: 0.25, + promote: LazyPromotion::Background { reads: 2 }, + promote_inflight: 8, + inline_rows: 0, + eager_before_full: false, + origin_max_runs: 5, + dense_to_eager: false, + } + ); + assert_eq!( + LayeredLazyConfig::from_lookup(|_| None).unwrap(), + LayeredLazyConfig::default() + ); + 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); + // The origin run cap does not depend on the eager-before-full switch, + // and dense probes go to the eager scan unless switched off. + 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())) + .unwrap(); + assert_eq!( + ( + config.eager_before_full, + config.origin_max_runs, + config.dense_to_eager + ), + (true, 2, true) + ); + let env = HashMap::from([(LAZY_FULL_ENV, "1"), (LAZY_DENSE_TO_EAGER_ENV, "false")]); + let config = + LayeredLazyConfig::from_lookup(|name| env.get(name).map(|value| value.to_string())) + .unwrap(); + assert_eq!( + config, + LayeredLazyConfig { + enabled: true, + dense_to_eager: false, + ..Default::default() + } + ); + let error = + LayeredLazyConfig::from_lookup(|key| (key == LAZY_FULL_ENV).then(|| "yes".into())) + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!(error.to_string().contains(LAZY_FULL_ENV), "{error}"); + for (name, value) in [ + (LAZY_DENSE_ENV, "dense"), + (LAZY_PROMOTE_ENV, "bg:0"), + (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_DENSE_TO_EAGER_ENV, "2"), + ] { + let enabled = |key: &str| -> Option { + if key == LAZY_FULL_ENV { + Some("1".into()) + } else { + (key == name).then(|| value.into()) + } + }; + let error = LayeredLazyConfig::from_lookup(enabled).unwrap_err(); + assert!( + matches!(error, Error::InvalidInput { .. }), + "{name}={value}" + ); + assert!(error.to_string().contains(name), "{error}"); + // With the scan off (unset or 0) the other knobs are ignored. + for full in [None, Some("0")] { + let disabled = |key: &str| -> Option { + if key == LAZY_FULL_ENV { + full.map(String::from) + } else { + (key == name).then(|| value.into()) + } + }; + assert_eq!( + LayeredLazyConfig::from_lookup(disabled).unwrap(), + LayeredLazyConfig::default(), + "{name}={value} with {LAZY_FULL_ENV}={full:?}" + ); + } + } + } + + #[test] + fn test_plan_plane_gather_cost_rule() { + // 520-byte rows: rows at most 7 apart share one run. + let row_bytes = 520; + let config = LayeredLazyConfig { + max_runs: 2, + dense_bytes_fraction: 0.5, + ..Default::default() + }; + assert_eq!( + plan_plane_gather(&[10, 17, 24], 2000, row_bytes, &config), + PlaneSource::Sparse + ); + assert_eq!( + plan_plane_gather(&[10, 100, 1000], 2000, row_bytes, &config), + PlaneSource::Whole, + "three runs exceed max_runs" + ); + assert_eq!( + plan_plane_gather(&(0..1200).collect::>(), 2000, row_bytes, &config), + PlaneSource::Whole, + "one run covering most of the plane" + ); + for (dense, expected) in [ + (DenseGatherMode::Sparse, PlaneSource::Sparse), + (DenseGatherMode::Whole, PlaneSource::Whole), + ] { + let config = LayeredLazyConfig { dense, ..config }; + assert_eq!( + plan_plane_gather(&[10, 100, 1000], 2000, row_bytes, &config), + expected + ); + assert_eq!(plan_plane_gather(&[10], 2000, row_bytes, &config), expected); + } + } + + #[test] + fn test_origin_reads_whole_plane_counts_row_runs() { + // 520-byte rows: rows at most 7 apart share one run. + let row_bytes = 520; + let two_runs = [10, 17, 100]; + let three_runs = [10, 100, 1000]; + // Unlimited by default: origin rows are read whatever their runs. + let unlimited = LayeredLazyConfig::default(); + assert!(!origin_reads_whole_plane( + &three_runs, + row_bytes, + &unlimited + )); + let capped = LayeredLazyConfig { + origin_max_runs: 2, + ..unlimited + }; + assert!(!origin_reads_whole_plane(&two_runs, row_bytes, &capped)); + assert!(origin_reads_whole_plane(&three_runs, row_bytes, &capped)); + // The cap applies whether gathers before the heap fills are deferred. + let deferred = LayeredLazyConfig { + eager_before_full: false, + ..capped + }; + assert!(origin_reads_whole_plane(&three_runs, row_bytes, &deferred)); + let none = LayeredLazyConfig { + origin_max_runs: 0, + ..unlimited + }; + assert!(origin_reads_whole_plane(&[10], row_bytes, &none)); + assert!(!origin_reads_whole_plane(&[], row_bytes, &none)); + } + + #[test] + fn test_sequential_plane_loads_flag() { + // The accessor reports the process environment's value. + let process = std::env::var(SEQUENTIAL_PLANE_LOADS_ENV).ok(); + assert_eq!( + sequential_plane_loads().ok(), + sequential_plane_loads_from(process.as_deref()).ok() + ); + assert!(!sequential_plane_loads_from(None).unwrap()); + assert!(!sequential_plane_loads_from(Some("0")).unwrap()); + assert!(sequential_plane_loads_from(Some("1")).unwrap()); + assert!(sequential_plane_loads_from(Some(" true ")).unwrap()); + let error = sequential_plane_loads_from(Some("yes")).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!( + error.to_string().contains(SEQUENTIAL_PLANE_LOADS_ENV), + "{error}" + ); + } + #[test] fn test_query_scratch_pool_initializes_buffer_capacity() { let pool = diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 4027cb1962e..18fb46336a4 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_stats; use lance_index::vector::bq::rabit_ex_bits; use lance_index::vector::bq::storage::{RabitQueryEstimator, SEGMENT_NUM_CODES}; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex, FlatQuantizer}; @@ -62,8 +63,8 @@ use lance_index::vector::quantizer::{ }; use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::storage::{ - PlaneAccessTracker, QueryResidual, QueryScratch, QueryScratchCapacity, QueryScratchPool, - RabitRawQueryContext, VectorStore, + LayeredLazyConfig, PlaneAccessTracker, QueryResidual, QueryScratch, QueryScratchCapacity, + QueryScratchPool, RabitRawQueryContext, VectorStore, }; use lance_index::vector::v3::subindex::SubIndexType; use lance_index::{ @@ -94,6 +95,8 @@ use uuid::Uuid; use super::{IvfIndexPartitionStatistics, IvfIndexStatistics, maybe_centroids_for_stats}; +mod lazy_full; + pub(crate) type RabitSearchCacheCell = Arc>>>>; /// Serializable state of an IVF index, sufficient to reconstruct the index @@ -206,6 +209,13 @@ pub(crate) const GLOBAL_TOPK_CHUNK_MAX_PARTITIONS: usize = 128; /// instead of a `spawn_cpu` dispatch; see the use site for the rationale. const GLOBAL_TOPK_INLINE_HEAP_LEN: usize = 4096; +/// Lazy layered full-precision scan settings from the `LANCE_RQ_LAZY_*` +/// environment, read once per process. An invalid value fails opening a +/// layered index rather than silently changing its scan; with the scan off +/// only `LANCE_RQ_LAZY_FULL` itself is read. +static LAYERED_LAZY_CONFIG: LazyLock> = + LazyLock::new(|| LayeredLazyConfig::from_env().map_err(|err| err.to_string())); + const IVF_PREWARM_WINDOW_SIZE_ENV: &str = "LANCE_IVF_PREWARM_WINDOW_SIZE_BYTES"; /// Default encoded-byte target of one prewarm read window. /// @@ -1077,6 +1087,15 @@ pub struct IVFIndex { use_residual_scratch: bool, rq_search_cache: Option>, prepared_partitions: Arc, + /// Whether the storage is a layered IVF_RQ index, fixed when opened. + layered_rq: bool, + /// Lazy full-precision scan settings of a layered index; see + /// [`LayeredLazyConfig`]. Locked only to copy it once per query. + layered_lazy: Mutex, + /// Staging parallelism of the lazy scan when non-zero; see + /// `set_lazy_prepare_parallelism_for_test`. + #[cfg(test)] + lazy_prepare_parallelism_for_test: AtomicUsize, _marker: PhantomData<(S, Q)>, } @@ -1652,6 +1671,9 @@ impl IVFIndex { let open_io_stats = scheduler.stats(); 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)?); Ok(Self { uri: to_local_path(&uri), index_path: uri.as_ref().to_string(), @@ -1670,6 +1692,10 @@ impl IVFIndex { io_parallelism, open_io_stats, prepared_partitions: Arc::default(), + layered_rq, + layered_lazy, + #[cfg(test)] + lazy_prepare_parallelism_for_test: AtomicUsize::new(0), _marker: PhantomData, }) } @@ -1693,6 +1719,9 @@ impl IVFIndex { 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)?); Ok(Self { uri, index_path, @@ -1714,6 +1743,10 @@ impl IVFIndex { // and the first open via `try_new` already accounts for it). open_io_stats: ScanStats::default(), prepared_partitions: Arc::default(), + layered_rq, + layered_lazy, + #[cfg(test)] + lazy_prepare_parallelism_for_test: AtomicUsize::new(0), _marker: PhantomData, }) } @@ -2213,6 +2246,13 @@ impl IVFIndex { Ok(PartitionEntry::new(idx, storage)) } + async fn load_sub_index(&self, partition_id: usize, io_stats: Option) -> Result { + let batch = self + .read_sub_index_batch(partition_id, self.read_projection.as_ref(), io_stats) + .await?; + S::load(batch) + } + /// Read the serialized sub-index of a partition, with its metadata attached. /// /// `projection` narrows the read to the columns search needs; `None` reads @@ -2847,6 +2887,29 @@ impl VectorIndex for IVFInd let raw_query_context = self.prepare_rq_raw_query_context(&query.key)?; if control.is_none() && S::supports_global_topk_heap() { + if let Some(config) = self.lazy_full_config(&query, raw_query_context.is_some()) { + if !self + .all_ex_planes_resident(&partitions, start_idx..end_idx) + .await + { + let batch = Box::pin(self.clone().search_partitions_lazy_full( + query, + partitions, + q_c_dists, + start_idx..end_idx, + pre_filter, + metrics, + raw_query_context, + config, + )) + .await?; + return Ok(Box::pin(RecordBatchStreamAdapter::new( + VECTOR_RESULT_SCHEMA.clone(), + stream::once(async move { Ok(batch) }), + ))); + } + layered_stats::counters().lazy_all_resident_skips.incr(); + } let heap_capacity = query.k * query.refine_factor.unwrap_or(1) as usize; { let _wait_timer = IndexTimer::new(metrics.as_ref(), IndexTiming::PrefilterWait); @@ -7733,6 +7796,1972 @@ mod tests { test_remap(params, 4, 0.5).await; } + /// The lazy layered full-precision scan must return exactly the eager + /// scan's batch (row ids, distance bits and heap order) for every cache + /// state, filter, bound and scan setting. + mod layered_lazy { + use super::*; + use std::sync::LazyLock; + use std::time::Duration; + + use arrow::compute::concat_batches; + use bytes::Bytes; + use lance_core::cache::{CacheCodec, CacheEntry, InternalCacheKey, QuickCacheBackend}; + use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; + use lance_index::vector::bq::layered::{PlaneBatch, PlaneKey, RQPrecision}; + 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::{ApproxMode, PartitionSearchControl, VECTOR_RESULT_SCHEMA}; + use lance_io::assert_io_eq; + use lance_io::object_store::ObjectStoreRegistry; + use rand::seq::SliceRandom; + + use crate::session::Session; + + type IvfRq = super::super::IVFIndex; + + const LAZY_DIM: usize = 128; + const LAZY_PARTITIONS: usize = 64; + const LAZY_ROWS: usize = 20_000; + /// `(partition, rows)` of partitions that are empty or smaller than + /// one 32-row FastScan block, or straddle its boundary. + const LAZY_SMALL_PARTITIONS: [(usize, usize); 8] = [ + (3, 0), + (11, 0), + (7, 1), + (13, 5), + (19, 17), + (23, 31), + (29, 32), + (31, 33), + ]; + /// About a tenth of the index's plane bytes, so plane entries churn. + const LAZY_SMALL_CACHE_BYTES: usize = 512 * 1024; + /// A few plane entries, so nearly every plane a query reads is admitted anew. + const LAZY_TINY_CACHE_BYTES: usize = 64 * 1024; + const LAZY_LARGE_CACHE_BYTES: usize = 256 * 1024 * 1024; + const LAZY_METADATA_CACHE_BYTES: usize = 64 * 1024 * 1024; + + /// These tests assert deltas of process-wide counters, so they run one at a time. + static LAZY_TEST_LOCK: LazyLock> = LazyLock::new(Default::default); + + /// RAM tier plus a serialized "disk" tier, like a tiered cache: disk + /// hits are admitted to RAM, loads are written through, and selected + /// rows are read from the serialized entry. + struct TieredPlaneTestBackend { + ram: QuickCacheBackend, + disk: std::sync::Mutex>, + gated: bool, + serve_rows: AtomicBool, + /// Report the plane entries of empty partitions as not resident, + /// as if the tier had evicted them. + hide_empty_planes: AtomicBool, + /// When set, every RAM admission through `get_or_insert` waits + /// for a slot in this bounded channel, like a tiered cache whose + /// admissions queue their evictions for the disk. + spill: std::sync::Mutex>>, + /// When set, plane loads through `get_or_insert` are recorded + /// here, and each yields once after it starts so that loads its + /// caller issues together overlap. + plane_loads: std::sync::Mutex>>, + } + + /// The loader a backend's `get_or_insert` receives. + type EntryLoader<'a> = std::pin::Pin< + Box> + Send + 'a>, + >; + + /// A plane load through the backend's `get_or_insert`, by plane. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum PlaneLoad { + Started(u8), + Finished(u8), + } + + impl std::fmt::Debug for TieredPlaneTestBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TieredPlaneTestBackend") + .field("gated", &self.gated) + .finish_non_exhaustive() + } + } + + impl TieredPlaneTestBackend { + fn new(ram_bytes: usize, gated: bool) -> Self { + Self { + ram: QuickCacheBackend::with_capacity(ram_bytes), + disk: Default::default(), + gated, + serve_rows: AtomicBool::new(true), + hide_empty_planes: AtomicBool::new(false), + spill: Default::default(), + plane_loads: Default::default(), + } + } + + fn record_plane_loads(&self) { + *self.plane_loads.lock().unwrap() = Some(Vec::new()); + } + + /// The plane loads recorded since `record_plane_loads`, which + /// stops recording. + fn take_plane_loads(&self) -> Vec { + self.plane_loads.lock().unwrap().take().unwrap_or_default() + } + + /// Record `load` when recording; returns whether it did. + fn record_plane_load(&self, load: PlaneLoad) -> bool { + match self.plane_loads.lock().unwrap().as_mut() { + Some(loads) => { + loads.push(load); + true + } + None => false, + } + } + + async fn get_or_insert_entry( + &self, + key: &InternalCacheKey, + loader: EntryLoader<'_>, + codec: Option, + ) -> Result<(CacheEntry, bool)> { + if let Some(entry) = self.ram.get(key, None).await { + return Ok((entry, true)); + } + let admitted = match self.read_disk(key) { + Some((entry, size)) => { + self.ram.insert(key, entry.clone(), size, None).await; + (entry, true) + } + None => { + let (entry, size) = loader.await?; + self.insert(key, entry.clone(), size, codec).await; + (entry, false) + } + }; + let spill = self.spill.lock().unwrap().clone(); + if let Some(spill) = spill { + // A closed channel only means the test stopped draining it. + let _ = spill.send(()).await; + } + Ok(admitted) + } + + fn set_spill(&self, spill: Option>) { + *self.spill.lock().unwrap() = spill; + } + + fn persist( + &self, + key: &InternalCacheKey, + entry: &CacheEntry, + size: usize, + codec: Option, + ) { + let Some(codec) = codec else { + return; + }; + let mut bytes = Vec::new(); + codec.serialize(entry, &mut bytes).unwrap(); + self.disk + .lock() + .unwrap() + .insert(*key, (Bytes::from(bytes), codec, size)); + } + + fn read_disk(&self, key: &InternalCacheKey) -> Option<(CacheEntry, usize)> { + let (bytes, codec, size) = self.disk.lock().unwrap().get(key).cloned()?; + codec.deserialize(&bytes).hit().map(|entry| (entry, size)) + } + + fn persisted_entries(&self, type_id: &str) -> usize { + self.disk + .lock() + .unwrap() + .values() + .filter(|(_, codec, _)| codec.type_id() == type_id) + .count() + } + } + + #[async_trait::async_trait] + impl CacheBackend for TieredPlaneTestBackend { + async fn get_resident(&self, key: &InternalCacheKey) -> Option { + let entry = self.ram.get_resident(key).await?; + let empty_plane = entry + .downcast_ref::() + .is_some_and(|plane| plane.0.num_rows() == 0); + if empty_plane && self.hide_empty_planes.load(Ordering::Relaxed) { + return None; + } + Some(entry) + } + + fn plane_admission_gated(&self) -> bool { + self.gated + } + + async fn get_without_promotion( + &self, + key: &InternalCacheKey, + _codec: Option, + ) -> Option { + match self.ram.get_resident(key).await { + Some(entry) => Some(entry), + None => self.read_disk(key).map(|(entry, _)| entry), + } + } + + async fn get_rows( + &self, + key: &InternalCacheKey, + rows: &[u32], + _codec: Option, + ) -> Option { + if !self.serve_rows.load(Ordering::Relaxed) { + return None; + } + let (bytes, codec, _) = self.disk.lock().unwrap().get(key).cloned()?; + let reader = |range: Range| -> Result { Ok(bytes.slice(range)) }; + codec.deserialize_rows(&reader, rows).hit() + } + + async fn get( + &self, + key: &InternalCacheKey, + _codec: Option, + ) -> Option { + if let Some(entry) = self.ram.get(key, None).await { + return Some(entry); + } + let (entry, size) = self.read_disk(key)?; + self.ram.insert(key, entry.clone(), size, None).await; + Some(entry) + } + + async fn insert( + &self, + key: &InternalCacheKey, + entry: CacheEntry, + size_bytes: usize, + codec: Option, + ) { + self.persist(key, &entry, size_bytes, codec); + self.ram.insert(key, entry, size_bytes, None).await; + } + + async fn get_or_insert<'a>( + &self, + key: &InternalCacheKey, + loader: std::pin::Pin< + Box> + Send + 'a>, + >, + codec: Option, + ) -> Result<(CacheEntry, bool)> { + let plane = codec + .and_then(|codec| codec.plane_tag()) + .filter(|&plane| self.record_plane_load(PlaneLoad::Started(plane))); + if plane.is_some() { + tokio::task::yield_now().await; + } + let result = self.get_or_insert_entry(key, loader, codec).await; + if let Some(plane) = plane { + self.record_plane_load(PlaneLoad::Finished(plane)); + } + result + } + + async fn clear(&self) { + self.ram.clear().await; + self.disk.lock().unwrap().clear(); + } + + async fn num_entries(&self) -> usize { + self.ram.num_entries().await + } + + async fn size_bytes(&self) -> usize { + self.ram.size_bytes().await + } + } + + /// Prefilter over a row-address mask. `collapses` makes the scan ask + /// for partition coverage and replace the filter with `NoFilter`. + struct LazyTestFilter { + mask: Arc, + collapses: bool, + } + + #[async_trait::async_trait] + impl PreFilter for LazyTestFilter { + async fn wait_for_ready(&self) -> Result<()> { + Ok(()) + } + + fn is_empty(&self) -> bool { + false + } + + fn needs_partition_row_ids(&self) -> bool { + self.collapses + } + + fn is_empty_for(&self, _rows: &RowAddrTreeMap) -> bool { + self.collapses + } + + fn mask(&self) -> Arc { + self.mask.clone() + } + + fn filter_row_ids<'a>( + &self, + row_ids: Box + 'a>, + ) -> Vec { + self.mask.selected_indices(row_ids) + } + } + + /// A prefilter whose scalar-index load failed. + struct FailingFilter; + + #[async_trait::async_trait] + impl PreFilter for FailingFilter { + async fn wait_for_ready(&self) -> Result<()> { + Err(lance_core::Error::internal("prefilter failed to load")) + } + + fn is_empty(&self) -> bool { + false + } + + fn mask(&self) -> Arc { + Arc::new(RowAddrMask::all_rows()) + } + + fn filter_row_ids<'a>( + &self, + row_ids: Box + 'a>, + ) -> Vec { + row_ids.copied().collect() + } + } + + struct NeverStop; + + impl PartitionSearchControl for NeverStop { + fn should_stop(&self) -> bool { + false + } + } + + fn lazy_test_filters() -> Vec<(&'static str, Arc)> { + let rows = |keep: fn(u64) -> bool| -> RowAddrTreeMap { + (0..LAZY_ROWS as u64) + .filter(|&row| keep(row)) + .collect::>() + .iter() + .collect() + }; + let filter = |mask: RowAddrMask, collapses: bool| -> Arc { + Arc::new(LazyTestFilter { + mask: Arc::new(mask), + collapses, + }) + }; + vec![ + ("none", Arc::new(NoFilter) as Arc), + ( + "sparse", + filter(RowAddrMask::from_allowed(rows(|row| row % 20 == 3)), false), + ), + ( + "dense", + filter(RowAddrMask::from_block(rows(|row| row % 20 == 7)), false), + ), + ( + "alternating", + filter(RowAddrMask::from_allowed(rows(|row| row % 2 == 0)), false), + ), + ( + "single", + filter(RowAddrMask::from_allowed(rows(|row| row == 42)), false), + ), + ( + "empty", + filter(RowAddrMask::from_allowed(RowAddrTreeMap::default()), false), + ), + ("collapse", filter(RowAddrMask::all_rows(), true)), + ] + } + + fn lazy_test_configs() -> Vec { + let enabled = LayeredLazyConfig { + enabled: true, + ..Default::default() + }; + vec![ + enabled, + LayeredLazyConfig { + window: 0, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Off, + inline_rows: 0, + ..enabled + }, + // Whole gathers predict every probe dense, which would + // otherwise send every probe to the eager scan. + LayeredLazyConfig { + window: 1, + dense: DenseGatherMode::Whole, + promote: LazyPromotion::Background { reads: 2 }, + inline_rows: usize::MAX, + dense_to_eager: false, + ..enabled + }, + LayeredLazyConfig { + window: 8, + max_runs: 0, + dense_bytes_fraction: 0.1, + ..enabled + }, + LayeredLazyConfig { + window: 64, + max_runs: 2, + promote: LazyPromotion::Off, + inline_rows: 0, + ..enabled + }, + LayeredLazyConfig { + max_runs: usize::MAX, + dense_bytes_fraction: f64::INFINITY, + ..enabled + }, + LayeredLazyConfig { + window: 4, + eager_before_full: false, + ..enabled + }, + LayeredLazyConfig { + dense: DenseGatherMode::Sparse, + origin_max_runs: 0, + ..enabled + }, + LayeredLazyConfig { + dense_to_eager: false, + ..enabled + }, + ] + } + + /// Clustered vectors around seeded centroids, in shuffled row order, + /// with the partition sizes of [`LAZY_SMALL_PARTITIONS`]. + fn lazy_test_data() -> (RecordBatch, Arc) { + let mut rng = StdRng::seed_from_u64(0x1a2e_7ed5); + let centroids: Vec = (0..LAZY_PARTITIONS * LAZY_DIM) + .map(|_| rng.random_range(-1.0f32..1.0)) + .collect(); + let small_rows: usize = LAZY_SMALL_PARTITIONS.iter().map(|(_, rows)| rows).sum(); + let regular_rows = + (LAZY_ROWS - small_rows) / (LAZY_PARTITIONS - LAZY_SMALL_PARTITIONS.len()); + let mut sizes: Vec = (0..LAZY_PARTITIONS) + .map(|partition| { + LAZY_SMALL_PARTITIONS + .iter() + .find(|(small, _)| *small == partition) + .map_or(regular_rows, |(_, rows)| *rows) + }) + .collect(); + sizes[0] += LAZY_ROWS - sizes.iter().sum::(); + let mut partitions: Vec = sizes + .iter() + .enumerate() + .flat_map(|(partition, &rows)| repeat_n(partition, rows)) + .collect(); + partitions.shuffle(&mut rng); + let mut values = Vec::with_capacity(LAZY_ROWS * LAZY_DIM); + for partition in partitions { + for dim in 0..LAZY_DIM { + values.push( + centroids[partition * LAZY_DIM + dim] + rng.random_range(-0.05f32..0.05), + ); + } + } + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from(values), + LAZY_DIM as i32, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("vector", vectors.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt64Array::from_iter_values(0..LAZY_ROWS as u64)), + Arc::new(vectors), + ], + ) + .unwrap(); + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(centroids), + LAZY_DIM as i32, + ) + .unwrap(); + (batch, Arc::new(centroids)) + } + + async fn write_lazy_test_dataset( + uri: &str, + bits: u8, + distance_type: DistanceType, + ) -> (Dataset, RecordBatch) { + let (batch, centroids) = lazy_test_data(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + let params = VectorIndexParams::with_ivf_rq_params( + distance_type, + IvfBuildParams::try_with_centroids(LAZY_PARTITIONS, centroids).unwrap(), + RQBuildParams::new(bits).with_layered(true), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + (dataset, batch) + } + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + enum LazyTestCache { + /// No capacity: every plane comes from the origin file. + Origin, + /// Lance's default backend at about a tenth of the planes, prewarmed. + Small, + /// Everything resident after prewarm. + Resident, + /// A tiered backend without sign gating, prewarmed to disk. + Ungated, + /// A tiered backend with sign gating at about a tenth of the planes. + Gated, + /// A tiered backend without sign gating that holds every plane. + TieredResident, + /// A tiered backend without sign gating whose RAM holds only a + /// few planes, prewarmed to disk. + Tiny, + /// An empty tiered backend without sign gating that can hold every plane. + ColdUngated, + /// An empty tiered backend with sign gating that can hold every plane. + ColdGated, + } + + async fn open_lazy_test_index( + uri: &str, + cache: LazyTestCache, + ) -> ( + Dataset, + Arc, + Option>, + ) { + let registry = Arc::new(ObjectStoreRegistry::default()); + let (session, tiered) = match cache { + LazyTestCache::Origin => { + (Session::new(0, LAZY_METADATA_CACHE_BYTES, registry), None) + } + LazyTestCache::Small => ( + Session::new(LAZY_SMALL_CACHE_BYTES, LAZY_METADATA_CACHE_BYTES, registry), + None, + ), + LazyTestCache::Resident => ( + Session::new(LAZY_LARGE_CACHE_BYTES, LAZY_METADATA_CACHE_BYTES, registry), + None, + ), + LazyTestCache::Ungated + | LazyTestCache::Gated + | LazyTestCache::TieredResident + | LazyTestCache::Tiny + | LazyTestCache::ColdUngated + | LazyTestCache::ColdGated => { + let (ram_bytes, gated) = match cache { + LazyTestCache::Gated => (LAZY_SMALL_CACHE_BYTES, true), + LazyTestCache::TieredResident | LazyTestCache::ColdUngated => { + (LAZY_LARGE_CACHE_BYTES, false) + } + LazyTestCache::ColdGated => (LAZY_LARGE_CACHE_BYTES, true), + LazyTestCache::Tiny => (LAZY_TINY_CACHE_BYTES, false), + _ => (LAZY_SMALL_CACHE_BYTES, false), + }; + let backend = Arc::new(TieredPlaneTestBackend::new(ram_bytes, gated)); + ( + Session::with_index_cache_backend( + backend.clone(), + LAZY_METADATA_CACHE_BYTES, + registry, + ), + Some(backend), + ) + } + }; + let dataset = crate::DatasetBuilder::from_uri(uri) + .with_session(Arc::new(session)) + .load() + .await + .unwrap(); + let uuid = dataset.load_indices().await.unwrap()[0].uuid; + let index = dataset + .open_vector_index("vector", &uuid, &NoOpMetricsCollector) + .await + .unwrap(); + if !matches!( + cache, + LazyTestCache::Origin | LazyTestCache::ColdUngated | LazyTestCache::ColdGated + ) { + index.prewarm().await.unwrap(); + } + (dataset, index, tiered) + } + + fn lazy_index(index: &Arc) -> &IvfRq { + index + .as_any() + .downcast_ref::() + .expect("layered IVF_RQ index") + } + + fn lazy_test_query(key: ArrayRef, k: usize, nprobes: usize) -> Query { + Query { + search_effort: 0.5, + rq_cascade_factor: None, + rq_precision: RQPrecision::Full, + column: "vector".to_string(), + key, + k, + lower_bound: None, + upper_bound: None, + minimum_nprobes: nprobes, + maximum_nprobes: Some(nprobes), + ef: None, + refine_factor: None, + metric_type: None, + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: ApproxMode::Normal, + } + } + + async fn search_global( + index: &Arc, + query: &Query, + filter: Arc, + ) -> Result { + let (partitions, dists) = index.find_partitions(query)?; + let probes = partitions.len(); + let batches = index + .clone() + .search_partitions( + query.clone(), + Arc::new(partitions), + Arc::new(dists), + 0, + probes, + filter, + None, + Arc::new(NoOpMetricsCollector), + ) + .await? + .try_collect::>() + .await?; + Ok(concat_batches(&VECTOR_RESULT_SCHEMA, batches.iter())?) + } + + /// Row ids and distance bits in heap order. + fn result_bits(batch: &RecordBatch) -> (Vec, Vec) { + ( + batch[ROW_ID].as_primitive::().values().to_vec(), + batch[DIST_COL] + .as_primitive::() + .values() + .iter() + .map(|dist| dist.to_bits()) + .collect(), + ) + } + + fn prune_delta( + after: RabitPruneStatsSnapshot, + before: RabitPruneStatsSnapshot, + ) -> [u64; 7] { + [ + after.calls - before.calls, + after.candidates - before.candidates, + after.pruned_upper_bound - before.pruned_upper_bound, + after.pruned_heap - before.pruned_heap, + after.exact - before.exact, + after.exact_rejected - before.exact_rejected, + after.bypass_calls - before.bypass_calls, + ] + } + + /// Run `query` eagerly and with `config`, assert identical batches and + /// (when `LANCE_RQ_PRUNE_STATS` is on) identical prune counters, and + /// return the lazy run's counters. + async fn assert_lazy_matches_eager( + index: &Arc, + query: &Query, + filter: &Arc, + config: LayeredLazyConfig, + context: &str, + ) -> LayeredLazyStats { + let ivf = lazy_index(index); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let before = rabit_prune_stats_snapshot(); + let eager = search_global(index, query, filter.clone()).await.unwrap(); + let after_eager = rabit_prune_stats_snapshot(); + layered_stats::snapshot_and_reset(); + ivf.set_layered_lazy_config_for_test(config); + let lazy = search_global(index, query, filter.clone()).await.unwrap(); + let after_lazy = rabit_prune_stats_snapshot(); + let stats = layered_stats::snapshot_and_reset(); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + assert_eq!(result_bits(&lazy), result_bits(&eager), "{context}"); + if before.enabled { + assert_eq!( + prune_delta(after_lazy, after_eager), + prune_delta(after_eager, before), + "prune counters: {context}" + ); + } + assert_eq!(stats.needed_not_fetched, 0, "{context}"); + stats + } + + /// Deltas of the promotion gauge are not reset by snapshots; wait + /// until background promotions have finished. + async fn wait_for_promotions() { + for _ in 0..500 { + if layered_stats::promotions_in_flight() == 0 { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!( + "{} lazy promotions still running", + layered_stats::promotions_in_flight() + ); + } + + fn with_bounds(query: &Query, eager: &RecordBatch) -> Option { + let mut dists = eager[DIST_COL] + .as_primitive::() + .values() + .to_vec(); + if dists.len() < 4 { + return None; + } + dists.sort_by(f32::total_cmp); + let mut bounded = query.clone(); + bounded.lower_bound = Some(dists[dists.len() / 10]); + bounded.upper_bound = Some(dists[dists.len() * 3 / 5]); + Some(bounded) + } + + #[rstest] + #[case::rq5_l2(5, DistanceType::L2)] + #[case::rq7_l2(7, DistanceType::L2)] + #[case::rq9_l2(9, DistanceType::L2)] + #[case::rq7_cosine(7, DistanceType::Cosine)] + #[case::rq7_dot(7, DistanceType::Dot)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_full_matches_eager( + #[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 mut rng = StdRng::seed_from_u64(u64::from(bits)); + let random: ArrayRef = Arc::new(Float32Array::from_iter_values( + (0..LAZY_DIM).map(|_| rng.random_range(-1.0f32..1.0)), + )); + let keys = [vectors.value(0), vectors.value(777), random]; + let filters = lazy_test_filters(); + let configs = lazy_test_configs(); + let mut case = 0usize; + 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; + 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" + ); + 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 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:?}", + query.approx_mode + ); + let stats = + assert_lazy_matches_eager(&index, &query, filter, config, &context) + .await; + match cache { + LazyTestCache::Resident => { + assert_eq!(stats.lazy_all_resident_skips, 1, "{context}"); + assert_eq!(stats.lazy_queries, 0, "{context}"); + } + LazyTestCache::Origin => { + assert_eq!(stats.lazy_queries, 1, "{context}"); + } + LazyTestCache::Ungated => { + assert_eq!(stats.origin_row_reads, 0, "{context}"); + } + LazyTestCache::Small + | LazyTestCache::Gated + | LazyTestCache::TieredResident + | LazyTestCache::Tiny + | LazyTestCache::ColdUngated + | LazyTestCache::ColdGated => {} + } + if case % 2 == 1 { + let eager = + search_global(&index, &query, filter.clone()).await.unwrap(); + if let Some(bounded) = with_bounds(&query, &eager) { + assert_lazy_matches_eager( + &index, + &bounded, + filter, + config, + &format!("{context} bounds"), + ) + .await; + } + } + } + } + } + // 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 (filter_name, filter) in &filters { + let context = format!( + "bits={bits} {distance_type:?} cache={cache:?} filter={filter_name} config={config:?}" + ); + assert_lazy_matches_eager(&index, &query, filter, *config, &context).await; + } + } + wait_for_promotions().await; + } + } + + /// An ungated backend is warmed partition by partition, so every plane + /// is persisted even when RAM holds only part of the sign planes, and + /// 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. + #[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; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (dataset, index, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Ungated).await; + let tiered = tiered.unwrap(); + assert_eq!( + tiered.persisted_entries(::TYPE_ID), + 3 * LAZY_PARTITIONS + ); + let key = batch["vector"].as_fixed_size_list().value(5); + let query = lazy_test_query(key, 100, LAZY_PARTITIONS); + let filter: Arc = Arc::new(NoFilter); + search_global(&index, &query, filter.clone()).await.unwrap(); + // The sparse run goes first: the routed eager load admits the + // first probe's planes, after which it would gather nothing. + for dense in [DenseGatherMode::Sparse, DenseGatherMode::Cost] { + let context = format!("ungated prewarm dense={dense:?}"); + dataset.object_store.as_ref().io_stats_incremental(); + let stats = assert_lazy_matches_eager( + &index, + &query, + &filter, + LayeredLazyConfig { + enabled: true, + dense, + promote: LazyPromotion::Off, + ..Default::default() + }, + &context, + ) + .await; + let io = dataset.object_store.as_ref().io_stats_incremental(); + assert_io_eq!( + io, + read_iops, + 0, + "warm layered queries read no origin plane: {context}" + ); + assert_eq!(stats.origin_row_reads, 0, "{context}"); + let sparse_planes: u64 = stats.high_sparse.iter().chain(&stats.low_sparse).sum(); + let eager_loads: u64 = stats.dense_to_eager.iter().sum(); + if dense == DenseGatherMode::Sparse { + assert!(eager_loads == 0 && sparse_planes > 0, "{context} {stats:?}"); + } else { + assert!(eager_loads > 0, "{context} {stats:?}"); + } + } + } + + /// Queries the lazy scan cannot serve run the eager scan and are + /// counted by reason. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_ineligible_queries_fall_back() { + 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(), LazyTestCache::Origin).await; + let key = batch["vector"].as_fixed_size_list().value(9); + let filter: Arc = Arc::new(NoFilter); + let config = LayeredLazyConfig { + enabled: true, + ..Default::default() + }; + + let mut fast = lazy_test_query(key.clone(), 10, 8); + fast.approx_mode = ApproxMode::Fast; + let stats = assert_lazy_matches_eager(&index, &fast, &filter, config, "fast").await; + assert_eq!((stats.ineligible_fast, stats.lazy_queries), (1, 0)); + + let base = lazy_test_query(key.clone(), 100, 8); + let eager = search_global(&index, &base, filter.clone()).await.unwrap(); + let mut cascade = with_bounds(&base, &eager).unwrap(); + cascade.rq_cascade_factor = Some(4); + let stats = + assert_lazy_matches_eager(&index, &cascade, &filter, config, "cascade").await; + assert_eq!((stats.ineligible_cascade, stats.lazy_queries), (1, 0)); + + let mut refine = lazy_test_query(key.clone(), 10, 8); + refine.refine_factor = Some(2); + let stats = assert_lazy_matches_eager(&index, &refine, &filter, config, "refine").await; + assert_eq!((stats.ineligible_refine, stats.lazy_queries), (1, 0)); + + let mut high = lazy_test_query(key.clone(), 10, 8); + high.rq_precision = RQPrecision::High; + let stats = assert_lazy_matches_eager(&index, &high, &filter, config, "high").await; + assert_eq!((stats.ineligible_precision, stats.lazy_queries), (1, 0)); + + // `k == 0` runs the eager scan, which still waits for the + // prefilter and so reports its failure. + let zero = lazy_test_query(key, 0, 8); + let stats = assert_lazy_matches_eager(&index, &zero, &filter, config, "k=0").await; + assert_eq!((stats.ineligible_k_zero, stats.lazy_queries), (1, 0)); + let failing: Arc = Arc::new(FailingFilter); + let ivf = lazy_index(&index); + for lazy in [LayeredLazyConfig::default(), config] { + ivf.set_layered_lazy_config_for_test(lazy); + let error = search_global(&index, &zero, failing.clone()) + .await + .unwrap_err(); + assert!(error.to_string().contains("prefilter failed"), "{error}"); + } + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + } + + /// A gated backend admits ex planes only behind their resident sign + /// plane, so a background promotion would read a whole plane only to + /// drop it; sparse gathers never promote there. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_gated_backend_never_promotes() { + 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(), LazyTestCache::Gated).await; + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + let config = LayeredLazyConfig { + enabled: true, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Background { reads: 1 }, + ..Default::default() + }; + let mut sparse_planes = 0; + for row in 0..4 { + let query = lazy_test_query(vectors.value(row * 101), 100, LAZY_PARTITIONS); + // Repeat each query so a plane is gathered sparsely more than once. + for repeat in 0..2 { + let context = format!("gated row={row} repeat={repeat}"); + let stats = + assert_lazy_matches_eager(&index, &query, &filter, config, &context).await; + assert_eq!(stats.lazy_queries, 1, "{context}"); + assert_eq!( + ( + stats.promotions_issued, + stats.promotions_deduped, + stats.promotions_skipped, + stats.promotion_bytes + ), + (0, 0, 0, 0), + "{context}" + ); + sparse_planes += stats + .high_sparse + .iter() + .chain(&stats.low_sparse) + .sum::(); + } + } + assert_eq!(layered_stats::promotions_in_flight(), 0); + assert!(sparse_planes > 0, "the gated setup must gather sparsely"); + } + + /// Empty partitions have no ex rows, so their plane entries, which the + /// lazy scan never touches, do not keep a query whose other ex planes + /// are all resident off the eager scan. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_empty_partitions_count_as_resident() { + 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, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::TieredResident).await; + let tiered = tiered.unwrap(); + let ivf = lazy_index(&index); + let empty = LAZY_SMALL_PARTITIONS + .iter() + .find(|(_, rows)| *rows == 0) + .map(|(partition, _)| *partition) + .unwrap(); + assert_eq!(ivf.storage.partition_size(empty), 0); + let high = PlaneKey { + partition: empty, + plane: 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); + + let filter: Arc = Arc::new(NoFilter); + let config = LayeredLazyConfig { + enabled: true, + ..Default::default() + }; + let query = lazy_test_query( + batch["vector"].as_fixed_size_list().value(11), + 100, + LAZY_PARTITIONS, + ); + let stats = + assert_lazy_matches_eager(&index, &query, &filter, config, "hidden empty planes") + .await; + assert_eq!( + (stats.lazy_all_resident_skips, stats.lazy_queries), + (1, 0), + "{stats:?}" + ); + } + + /// 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), + LayeredLazyConfig { + enabled: true, + ..Default::default() + }, + "remapper", + ) + .await; + assert_eq!((stats.ineligible_remapper, stats.lazy_queries), (1, 0)); + } + + /// Per-partition, streaming, batch and cascade searches never use the + /// lazy scan. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_leaves_other_search_paths_unchanged() { + 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(), LazyTestCache::Origin).await; + let ivf = lazy_index(&index); + let query = lazy_test_query(batch["vector"].as_fixed_size_list().value(3), 50, 8); + let (partitions, dists) = index.find_partitions(&query).unwrap(); + let partitions = Arc::new(partitions); + let dists = Arc::new(dists); + let filter: Arc = Arc::new(NoFilter); + let mut cascade = query.clone(); + cascade.rq_cascade_factor = Some(4); + + let mut outputs = Vec::new(); + for config in [ + LayeredLazyConfig::default(), + LayeredLazyConfig { + enabled: true, + ..Default::default() + }, + ] { + ivf.set_layered_lazy_config_for_test(config); + layered_stats::snapshot_and_reset(); + let in_partition = index + .search_in_partition( + partitions.value(0) as usize, + &query, + filter.clone(), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let streamed = index + .clone() + .search_partitions( + query.clone(), + partitions.clone(), + dists.clone(), + 0, + partitions.len(), + filter.clone(), + Some(Arc::new(NeverStop)), + Arc::new(NoOpMetricsCollector), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batched = index + .clone() + .search_partitions_batch( + query.clone(), + vec![partitions.clone()], + vec![dists.clone()], + filter.clone(), + Arc::new(NoOpMetricsCollector), + ) + .await + .unwrap(); + let cascaded = search_global(&index, &cascade, filter.clone()) + .await + .unwrap(); + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.lazy_queries, 0); + assert_eq!(stats.lazy_all_resident_skips, 0); + 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)); + outputs.push(output); + } + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + assert_eq!(outputs[0], outputs[1]); + } + + /// Concurrent queries share the promotion tracker, a cleared cache + /// and a missing persistent entry fall back to the origin file, and a + /// dropped query releases what it holds. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_lifecycle() { + 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, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Ungated).await; + let tiered = tiered.unwrap(); + let ivf = lazy_index(&index); + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + let config = LayeredLazyConfig { + enabled: true, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Background { reads: 1 }, + // Gathers the persistent tier cannot serve read origin rows. + origin_max_runs: usize::MAX, + ..Default::default() + }; + let queries: Vec = (0..4) + .map(|row| lazy_test_query(vectors.value(row * 101), 10, LAZY_PARTITIONS)) + .collect(); + + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let mut expected = Vec::new(); + for query in &queries { + expected.push(result_bits( + &search_global(&index, query, filter.clone()).await.unwrap(), + )); + } + + // Identical queries in parallel promote each plane at most once at a time. + ivf.set_layered_lazy_config_for_test(config); + layered_stats::snapshot_and_reset(); + let concurrent = futures::future::try_join_all( + queries + .iter() + .chain(queries.iter()) + .map(|query| search_global(&index, query, filter.clone())), + ) + .await + .unwrap(); + for (position, result) in concurrent.iter().enumerate() { + assert_eq!(result_bits(result), expected[position % queries.len()]); + } + wait_for_promotions().await; + let stats = layered_stats::snapshot_and_reset(); + assert!(stats.promotions_issued > 0); + // A promoted plane the small RAM tier evicted again is not completed. + assert!(stats.promotions_completed > 0, "{stats:?}"); + assert!( + stats.promotions_completed <= stats.promotions_issued, + "{stats:?}" + ); + + // Without persistent rows the gathers read the origin file. The + // promotions above may have left every plane these queries need + // resident, which would skip the gathers, so empty the RAM tier. + tiered.ram.clear().await; + tiered.serve_rows.store(false, Ordering::Relaxed); + for (query, expected) in queries.iter().zip(&expected) { + let result = search_global(&index, query, filter.clone()).await.unwrap(); + assert_eq!(&result_bits(&result), expected); + } + wait_for_promotions().await; + let stats = layered_stats::snapshot_and_reset(); + assert!(stats.origin_row_reads > 0, "{stats:?}"); + tiered.serve_rows.store(true, Ordering::Relaxed); + + // Clearing the cache between queries loses every plane entry. + tiered.clear().await; + for (query, expected) in queries.iter().zip(&expected) { + let result = search_global(&index, query, filter.clone()).await.unwrap(); + assert_eq!(&result_bits(&result), expected); + } + wait_for_promotions().await; + let stats = layered_stats::snapshot_and_reset(); + assert!(stats.origin_row_reads > 0, "{stats:?}"); + + // A dropped query stops its producer and releases its partitions. + for timeout in [1, 5, 50] { + let _ = tokio::time::timeout( + Duration::from_micros(timeout), + search_global(&index, &queries[0], filter.clone()), + ) + .await; + } + wait_for_promotions().await; + for _ in 0..500 { + if ivf.prepared_partitions().in_flight() == 0 { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!(ivf.prepared_partitions().in_flight(), 0); + let result = search_global(&index, &queries[0], filter.clone()) + .await + .unwrap(); + assert_eq!(result_bits(&result), expected[0]); + wait_for_promotions().await; + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + drop(dataset); + } + + /// A staged probe the producer is not polling must not hold a slot of + /// a fair cache resource that the head gather waits for. + /// + /// Every RAM admission of the backend waits for a slot in a bounded + /// spill channel drained slowly, like a tiered cache whose admissions + /// queue their evictions for the disk. Staging steps queue there for + /// sign planes ahead of the gathers' whole ex planes. With one step + /// window the gather buffer holds three probes, and while its head is + /// pending the producer stops polling the staging buffer. When the + /// steps were plain futures in those buffers, the channel handed the + /// next freed slot to its oldest waiter, a staging step nobody polled, + /// which never used it: the head gather and every later admission then + /// waited forever. Steps now run as tasks and complete on their own, + /// whether gathers before the heap fills are deferred or issued at + /// once, and when staging loads predicted-dense probes for the eager + /// scan instead (every probe here, as whole gathers predict them dense). + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_unpolled_staging_cannot_hold_spill_slot() { + const PROBES: usize = 32; + const QUERIES: usize = 16; + // Staged probes in flight; the deadlock needs staging steps + // queued on the spill channel behind the ones that fill the + // gather buffer, whatever the host's core count. + const STAGING_STEPS: usize = 8; + const QUERY_TIMEOUT: Duration = Duration::from_secs(10); + 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, tiered) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Tiny).await; + let tiered = tiered.unwrap(); + let ivf = lazy_index(&index); + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + // A quarter of the probed rows: the first probes are gathered + // before the heap fills, the later ones gate on its threshold. + let k = LAZY_ROWS * PROBES / LAZY_PARTITIONS / 4; + let queries: Vec = (0..QUERIES) + .map(|row| lazy_test_query(vectors.value(row * 211), k, PROBES)) + .collect(); + + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + let mut expected = Vec::with_capacity(QUERIES); + for query in &queries { + expected.push(result_bits( + &search_global(&index, query, filter.clone()).await.unwrap(), + )); + } + + let (spill_tx, mut spill_rx) = tokio::sync::mpsc::channel::<()>(1); + let drainer = tokio::spawn(async move { + while spill_rx.recv().await.is_some() { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }); + 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)] + { + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig { + enabled: true, + window: 1, + dense: DenseGatherMode::Whole, + promote: LazyPromotion::Off, + eager_before_full, + dense_to_eager, + ..Default::default() + }); + let setting = format!( + "eager_before_full={eager_before_full} dense_to_eager={dense_to_eager}" + ); + layered_stats::snapshot_and_reset(); + for (position, (query, expected)) in queries.iter().zip(&expected).enumerate() { + let result = tokio::time::timeout( + QUERY_TIMEOUT, + search_global(&index, query, filter.clone()), + ) + .await + .unwrap_or_else(|_| { + panic!( + "lazy query {position} ({setting}) did not finish within {QUERY_TIMEOUT:?}" + ) + }) + .unwrap(); + assert_eq!( + &result_bits(&result), + expected, + "query {position} {setting}" + ); + } + let stats = layered_stats::snapshot_and_reset(); + assert_eq!(stats.lazy_queries, QUERIES as u64, "{setting} {stats:?}"); + 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 { + assert!( + eager_loads > 0 && whole_planes == 0, + "staging must load every probe eagerly: {setting} {stats:?}" + ); + } else { + assert!( + eager_loads == 0 && whole_planes > 0, + "gathers must admit whole ex planes: {setting} {stats:?}" + ); + } + } + tiered.set_spill(None); + drainer.abort(); + ivf.set_lazy_prepare_parallelism_for_test(0); + ivf.set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + } + + /// Sums of the lazy counters that tell routing, issue and origin + /// policies apart. + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] + struct IssueTotals { + dense_to_eager: [u64; RANK_BUCKETS], + certain_dense: [u64; RANK_BUCKETS], + lazy_probes: u64, + empty: u64, + eager_before_full: u64, + deferred_issues: u64, + serial_waits: u64, + origin_whole_fallbacks: u64, + origin_row_reads: u64, + sparse_planes: u64, + whole_planes: u64, + } + + impl IssueTotals { + fn add(&mut self, stats: &LayeredLazyStats) { + for (totals, counts) in [ + (&mut self.dense_to_eager, stats.dense_to_eager), + (&mut self.certain_dense, stats.certain_dense), + ] { + for (total, count) in totals.iter_mut().zip(counts) { + *total += count; + } + } + self.lazy_probes += stats.lazy_probes.iter().sum::(); + self.empty += stats.empty.iter().sum::(); + self.eager_before_full += stats.eager_before_full; + self.deferred_issues += stats.deferred_issues; + self.serial_waits += stats.serial_waits; + self.origin_whole_fallbacks += stats.origin_whole_fallbacks; + self.origin_row_reads += stats.origin_row_reads; + self.sparse_planes += stats + .high_sparse + .iter() + .chain(&stats.low_sparse) + .sum::(); + self.whole_planes += stats.high_whole.iter().chain(&stats.low_whole).sum::(); + } + } + + /// Issuing gathers before the heap fills, and loading whole planes + /// instead of many origin row runs, change only what is read: results + /// match the eager scan with either switched on or off, including `k` + /// beyond the rows of the first probes and backends that serve no + /// persistent rows. Backends that gate plane admission (Lance's + /// default cache included) never fall back to whole planes. Dense + /// probes stay in the lazy pipeline here, so every probe is gathered. + #[rstest] + #[case::origin(LazyTestCache::Origin, true)] + #[case::small(LazyTestCache::Small, true)] + #[case::gated(LazyTestCache::Gated, true)] + #[case::tiny_without_rows(LazyTestCache::Tiny, false)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_eager_before_full_matches_eager( + #[case] cache: LazyTestCache, + #[case] serve_rows: bool, + ) { + 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, tiered) = open_lazy_test_index(dir.as_str(), cache).await; + if let Some(tiered) = &tiered { + tiered.serve_rows.store(serve_rows, Ordering::Relaxed); + } + let origin_only = tiered.is_none() || !serve_rows; + let gated = lazy_index(&index).index_cache.plane_admission_gated(); + let vectors = batch["vector"].as_fixed_size_list(); + let filters = lazy_test_filters(); + let mut totals: HashMap<(bool, usize), IssueTotals> = HashMap::new(); + let mut case = 0usize; + for row in [0, 777] { + for nprobes in [8, LAZY_PARTITIONS] { + for k in [10, 1000, 5000, LAZY_ROWS + 1] { + for eager_before_full in [true, false] { + for origin_max_runs in [0, 2, usize::MAX] { + for dense in [DenseGatherMode::Cost, DenseGatherMode::Sparse] { + case += 1; + let (filter_name, filter) = &filters[case % filters.len()]; + let config = LayeredLazyConfig { + enabled: true, + dense, + eager_before_full, + origin_max_runs, + dense_to_eager: false, + ..Default::default() + }; + let query = lazy_test_query(vectors.value(row), k, nprobes); + let context = format!( + "cache={cache:?} serve_rows={serve_rows} row={row} nprobes={nprobes} k={k} filter={filter_name} config={config:?}" + ); + let stats = assert_lazy_matches_eager( + &index, &query, filter, config, &context, + ) + .await; + if !eager_before_full { + assert_eq!(stats.eager_before_full, 0, "{context}"); + } + if origin_max_runs == usize::MAX || gated { + assert_eq!(stats.origin_whole_fallbacks, 0, "{context}"); + } + totals + .entry((eager_before_full, origin_max_runs)) + .or_default() + .add(&stats); + } + } + } + } + } + } + wait_for_promotions().await; + let on_uncapped = totals[&(true, usize::MAX)]; + assert!(on_uncapped.eager_before_full > 0, "{totals:?}"); + if origin_only && !gated { + // The origin run cap applies with either issue policy. + for eager_before_full in [true, false] { + let capped = totals[&(eager_before_full, 0)]; + assert!(capped.origin_whole_fallbacks > 0, "{totals:?}"); + assert_eq!(capped.origin_row_reads, 0, "{totals:?}"); + assert!( + totals[&(eager_before_full, usize::MAX)].origin_row_reads > 0, + "{totals:?}" + ); + } + } + } + + /// 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. + #[rstest] + #[case::origin(LazyTestCache::Origin)] + #[case::ungated(LazyTestCache::Ungated)] + #[case::gated(LazyTestCache::Gated)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_dense_to_eager_matches_eager(#[case] cache: LazyTestCache) { + 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 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]; + for row in [0, 777] { + for nprobes in [1, 8, LAZY_PARTITIONS] { + for k in [10, 1000, LAZY_ROWS + 1] { + for (filter_name, filter) in &filters { + let query = lazy_test_query(vectors.value(row), k, nprobes); + 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; + } + } + } + } + } + } + wait_for_promotions().await; + assert_eq!(routed[0], 0, "switched off, no probe is routed"); + assert!(routed[1] > 0, "{routed:?}"); + } + + /// `count` queries for `k` rows over `nprobes` probes, keyed by rows + /// of `vectors`, whose first probe alone holds more than `2 * k` rows: + /// it fills the heap, and the threshold then keeps under half of the + /// rows scored, so later gathers are not expected to be dense. + fn queries_filled_by_first_probe( + index: &Arc, + vectors: &FixedSizeListArray, + k: usize, + nprobes: usize, + count: usize, + ) -> Vec { + let queries: Vec = (0..vectors.len()) + .step_by(211) + .map(|row| lazy_test_query(vectors.value(row), k, nprobes)) + .filter(|query| { + let (partitions, _) = index.find_partitions(query).unwrap(); + lazy_index(index) + .storage + .partition_size(partitions.value(0) as usize) + > 2 * k + }) + .take(count) + .collect(); + assert_eq!(queries.len(), count); + queries + } + + /// Lazy settings without promotions, with the given issue and routing + /// policies. + fn issue_policy(eager_before_full: bool, dense_to_eager: bool) -> LayeredLazyConfig { + LayeredLazyConfig { + enabled: true, + promote: LazyPromotion::Off, + eager_before_full, + dense_to_eager, + ..Default::default() + } + } + + /// Issue counters of the lazy scan under each of `configs`, summed + /// over `queries`. + async fn lazy_issue_totals( + index: &Arc, + queries: &[Query], + filter: &Arc, + configs: [LayeredLazyConfig; N], + ) -> [IssueTotals; N] { + let mut totals = [IssueTotals::default(); N]; + for (config, totals) in configs.into_iter().zip(&mut totals) { + for (position, query) in queries.iter().enumerate() { + let context = format!( + "eager_before_full={} dense_to_eager={} query={position} k={}", + config.eager_before_full, config.dense_to_eager, query.k + ); + let stats = + assert_lazy_matches_eager(index, query, filter, config, &context).await; + assert_eq!(stats.lazy_queries, 1, "{context}"); + totals.add(&stats); + } + } + totals + } + + /// Which probes the eager scan loads, and whether lazy gathers may be + /// issued before the heap fills, depend on `k`, not on how staging + /// 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. + /// + /// When the first probe alone fills the heap (small `k`), only that + /// probe, scored with an infinite threshold, goes to the eager scan, + /// exactly the probe the lazy pipeline otherwise gathers as certain + /// dense; later gathers wait for its finite threshold, as the + /// eager-before-full ablation always does. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_issue_policy_depends_on_k() { + const PROBES: usize = 32; + const QUERIES: usize = 8; + const STAGING_STEPS: usize = 8; + const SMALL_K: usize = 10; + 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(), LazyTestCache::Origin).await; + let ivf = lazy_index(&index); + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + // A quarter of the probed rows, several probes' worth. + let large_k = LAZY_ROWS * PROBES / LAZY_PARTITIONS / 4; + let queries = |k: usize| -> Vec { + (0..QUERIES) + .map(|row| lazy_test_query(vectors.value(row * 211), k, PROBES)) + .collect() + }; + 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), + ]; + ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); + let [large_routed, large_on, large_off] = + lazy_issue_totals(&index, &queries(large_k), &filter, configs).await; + let [small_routed, small_on, small_off] = + lazy_issue_totals(&index, &small, &filter, configs).await; + ivf.set_lazy_prepare_parallelism_for_test(0); + + // Nothing is resident, so every probe of a non-empty partition + // goes to the eager scan, in every rank bucket the probes reach. + let probes = (QUERIES * PROBES) as u64; + let routed: u64 = large_routed.dense_to_eager.iter().sum(); + assert_eq!(routed + large_routed.empty, probes, "{large_routed:?}"); + assert!( + large_routed.dense_to_eager[..4] + .iter() + .all(|&count| count > 0), + "{large_routed:?}" + ); + assert_eq!(large_routed.dense_to_eager[4], 0, "{large_routed:?}"); + assert_eq!( + ( + large_routed.lazy_probes, + large_routed.whole_planes, + large_routed.eager_before_full, + large_routed.deferred_issues, + ), + (0, 0, 0, 0), + "{large_routed:?}" + ); + assert_eq!(large_routed.certain_dense, [0; RANK_BUCKETS]); + for large in [large_on, large_off] { + assert_eq!(large.dense_to_eager, [0; RANK_BUCKETS], "{large:?}"); + assert_eq!(large.lazy_probes + large.empty, probes, "{large:?}"); + assert!(large.whole_planes > 0, "{large:?}"); + assert!(large.certain_dense[0] > 0, "{large:?}"); + } + assert_eq!( + (large_on.deferred_issues, large_on.serial_waits), + (0, 0), + "{large_on:?}" + ); + assert!(large_on.eager_before_full > 0, "{large_on:?}"); + assert_eq!(large_off.eager_before_full, 0, "{large_off:?}"); + assert!(large_off.deferred_issues > 0, "{large_off:?}"); + + let mut first_probe = [0; RANK_BUCKETS]; + first_probe[0] = QUERIES as u64; + assert_eq!(small_routed.dense_to_eager, first_probe, "{small_routed:?}"); + assert_eq!(small_routed.certain_dense, [0; RANK_BUCKETS]); + assert!(small_routed.lazy_probes > 0, "{small_routed:?}"); + for small in [small_on, small_off] { + assert_eq!(small.dense_to_eager, [0; RANK_BUCKETS], "{small:?}"); + assert_eq!(small.certain_dense, first_probe, "{small:?}"); + assert!(small.deferred_issues > 0, "{small:?}"); + } + for small in [small_routed, small_on, small_off] { + assert_eq!( + (small.eager_before_full, small.serial_waits), + (0, 0), + "{small:?}" + ); + } + } + + /// A selective prefilter keeps the heap from filling after the probes + /// that hold `k` rows. Deferred gathers are then issued at once + /// instead of one probe at a time, with identical results; the + /// ablation waits for each probe's turn. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_eager_before_full_when_filters_delay_the_heap() { + const PROBES: usize = 16; + const QUERIES: usize = 8; + const STAGING_STEPS: usize = 8; + const K: usize = 100; + 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(), LazyTestCache::Origin).await; + let ivf = lazy_index(&index); + let vectors = batch["vector"].as_fixed_size_list(); + let (_, filter) = lazy_test_filters() + .into_iter() + .find(|(name, _)| *name == "sparse") + .unwrap(); + // The first probe alone holds `k` rows, too few of which the + // 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 [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:?}"); + assert!(on.serial_waits < off.serial_waits, "{on:?} {off:?}"); + assert_eq!(off.eager_before_full, 0, "{off:?}"); + } + + /// A sparse gather the persistent tier cannot serve loads the whole + /// plane once its origin row runs exceed the cap, whether gathers + /// before the heap fills are deferred or not, with identical results. + /// Without a cap, or on a backend that gates plane admission (which + /// may not admit the plane; Lance's default cache is one), it reads + /// the rows. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_origin_runs_fall_back_to_whole_plane() { + 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 vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + for cache in [ + LazyTestCache::Origin, + LazyTestCache::Ungated, + LazyTestCache::Gated, + ] { + let (_dataset, index, tiered) = open_lazy_test_index(dir.as_str(), cache).await; + if let Some(tiered) = &tiered { + // Every sparse gather that misses RAM goes to the origin. + tiered.serve_rows.store(false, Ordering::Relaxed); + } + let gated = lazy_index(&index).index_cache.plane_admission_gated(); + assert_eq!(gated, cache != LazyTestCache::Ungated, "{cache:?}"); + for eager_before_full in [true, false] { + for origin_max_runs in [0, usize::MAX] { + let config = LayeredLazyConfig { + enabled: true, + dense: DenseGatherMode::Sparse, + promote: LazyPromotion::Off, + eager_before_full, + origin_max_runs, + ..Default::default() + }; + let mut totals = IssueTotals::default(); + for row in 0..4 { + let query = + lazy_test_query(vectors.value(row * 101), 100, LAZY_PARTITIONS); + let context = format!("{cache:?} {config:?} row={row}"); + totals.add( + &assert_lazy_matches_eager( + &index, &query, &filter, config, &context, + ) + .await, + ); + } + let context = format!("{cache:?} {config:?} {totals:?}"); + if origin_max_runs == 0 && !gated { + // Every origin gather takes at least one run. + assert!(totals.origin_whole_fallbacks > 0, "{context}"); + assert_eq!( + (totals.origin_row_reads, totals.sparse_planes), + (0, 0), + "{context}" + ); + } else { + assert_eq!(totals.origin_whole_fallbacks, 0, "{context}"); + assert!(totals.origin_row_reads > 0, "{context}"); + } + } + } + } + } + + /// A backend that admits planes like any entry loads a missed + /// partition's three planes at once, so it pays one round trip; a + /// gated backend admits the sign plane before it starts the ex planes, + /// whose admission depends on it. Results are identical. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_plane_loads_overlap_unless_gated() { + 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 key = batch["vector"].as_fixed_size_list().value(17); + let filter: Arc = Arc::new(NoFilter); + let single = lazy_test_query(key.clone(), 10, 1); + let every = lazy_test_query(key, 100, LAZY_PARTITIONS); + let (_resident_dataset, resident, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Resident).await; + let expected = [ + result_bits( + &search_global(&resident, &single, filter.clone()) + .await + .unwrap(), + ), + result_bits( + &search_global(&resident, &every, filter.clone()) + .await + .unwrap(), + ), + ]; + for cache in [LazyTestCache::ColdUngated, LazyTestCache::ColdGated] { + let (_dataset, index, tiered) = open_lazy_test_index(dir.as_str(), cache).await; + let tiered = tiered.unwrap(); + tiered.record_plane_loads(); + let result = search_global(&index, &single, filter.clone()) + .await + .unwrap(); + let loads = tiered.take_plane_loads(); + assert_eq!(result_bits(&result), expected[0], "{cache:?}"); + assert_eq!(loads.len(), 6, "{cache:?} {loads:?}"); + let position = + |load: PlaneLoad| loads.iter().position(|&seen| seen == load).unwrap(); + let sign_loaded = position(PlaneLoad::Finished(0)); + for plane in [1, 2] { + let started = position(PlaneLoad::Started(plane)); + if cache == LazyTestCache::ColdGated { + assert!(sign_loaded < started, "{cache:?} {loads:?}"); + } else { + assert!(started < sign_loaded, "{cache:?} {loads:?}"); + } + } + let result = search_global(&index, &every, filter.clone()).await.unwrap(); + assert_eq!(result_bits(&result), expected[1], "{cache:?}"); + } + } + } + #[rstest] #[case::fast(RQRotationType::Fast)] #[case::matrix(RQRotationType::Matrix)] diff --git a/rust/lance/src/index/vector/ivf/v2/lazy_full.rs b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs new file mode 100644 index 00000000000..f8dd41b78fa --- /dev/null +++ b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs @@ -0,0 +1,1236 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Lazy full-precision scan of a layered IVF_RQ index. +//! +//! The eager global-heap scan loads the sign, high and low planes of every +//! probed partition before scoring. This scan loads only the sign plane, +//! bounds every row with the native 1-bit lower bound (stage 1), and gathers +//! the ex-plane rows of the survivors only, which it then scores with the +//! eager scan's per-row step on the same global heap (stage 2). +//! +//! Stage 1 prunes against a threshold `T` taken from that heap: the top of a +//! heap holding `k` rows at some point no later than the probe's scoring, or +//! +inf while the heap holds fewer. The heap's top never increases once full, +//! so every row stage 1 prunes is one the eager scan's per-row step would +//! prune too, and stage 2 replays the remaining rows in the same order with +//! the same inputs. Heap contents, distances, ties and prune counters are +//! therefore identical to the eager scan; `T` only decides how many rows are +//! read. +//! +//! Pipeline, in probe order: +//! - a producer task stages probes (sign plane, prefilter, stage 1 on the CPU +//! pool), then gathers each probe's survivors once the probe is at most +//! `window` probes ahead of scoring, reading the threshold the scorer +//! publishes over a `watch` channel. A probe whose gate opens before the +//! heap is full waits until the heap fills or its turn comes, unless +//! waiting would not save reads (`LANCE_RQ_LAZY_EAGER_BEFORE_FULL`): then +//! it is gathered at once with `T` = +inf, selecting every accepted row as +//! the eager load does; +//! - the query task scores ready probes on one heap and publishes progress. +//! +//! Every staging and gather step runs as its own task, so it makes progress +//! whether or not the producer is polling the buffer that holds it; see +//! `spawn_scan_step`. +//! +//! 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. + +use std::collections::BinaryHeap; +use std::future::Future; +use std::ops::Range; +use std::sync::Arc; +use std::time::Instant; + +use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array}; +use futures::StreamExt; +use futures::prelude::stream; +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_stats; +use lance_index::vector::bq::storage::{ + ExRows, RabitQuantizationStorage, SignStage, StagePruneCounts, SurvivorRow, + select_full_survivors, +}; +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, +}; +use lance_index::vector::v3::subindex::IvfSubIndex; +use lance_index::vector::{ApproxMode, Query}; +use tokio::sync::{mpsc, watch}; + +use super::{ + GLOBAL_TOPK_INLINE_HEAP_LEN, IVFIndex, LAYERED_LAZY_CONFIG, PartitionEntry, + PreparedPartitionGuard, PreparedPartitionSearch, RabitSearchCache, + rotated_partition_centroid_slice, +}; + +/// Ready probes buffered between the producer and the scorer. Gathers beyond +/// it wait in the producer, where the staleness gate bounds them. +const LAZY_READY_CHANNEL_CAPACITY: usize = 2; +/// Most ready probes one scoring dispatch takes. +const LAZY_SCORE_BATCH_MAX_PROBES: usize = 64; +/// Gathers kept in flight beyond the staleness window, so the reads of the +/// next probes overlap the scoring of the current one. +const LAZY_FETCH_EXTRA_SLOTS: usize = 2; + +/// Scoring progress the scorer publishes to the gathers. +#[derive(Debug, Clone, Copy)] +struct LazyProgress { + /// Probes scored so far, in probe order. + scored: usize, + /// Whether the heap has held `k` rows; it then stays full. + full: bool, + /// The heap's top when `full`, otherwise +inf. + threshold: f32, +} + +/// A probe after its sign stage. +enum LazyStaged { + Eager(PreparedPartitionSearch), + /// Lower-bound gating is disabled for the partition, so it is scored by + /// the eager scan, which scores every row. + GatingOff { + partition_id: usize, + query: Query, + }, + Empty, + Lazy(Box>), +} + +struct LazySignProbe { + rank: usize, + partition_id: usize, + dist_q_c: f32, + /// The sub-index and the sign-plane-only storage. + entry: Arc>, + stage: SignStage, + /// Prefilter acceptance by partition offset; `None` accepts every row. + accept: Option>, + _in_flight: PreparedPartitionGuard, +} + +/// A probe ready for scoring. +enum LazyProbe { + Eager(Box>), + Empty, + Ready(Box>), +} + +struct LazyReadyProbe { + sign: LazySignProbe, + /// Sorted partition offsets whose ex rows were gathered. + survivors: Vec, + stage1: StagePruneCounts, + gathered: GatheredEx, +} + +impl LazyProbe { + fn survivor_rows(&self) -> usize { + match self { + Self::Ready(ready) => ready.survivors.len(), + Self::Eager(_) | Self::Empty => 0, + } + } +} + +/// What `k` and the partition sizes predict about the probes' gathers before +/// any plane is read. Staging and the gathers decide from the same forecast. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct LazyDenseForecast { + /// First rank whose earlier probes hold `k` rows (the probe count if + /// none): once they are scored the heap is full, unless the prefilter or + /// the distance bounds dropped rows. + fill_rank: usize, + /// 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`]. + route_dense: bool, +} + +impl LazyDenseForecast { + /// The forecast for a query with heap capacity `k` over probes holding + /// `probe_rows` rows each, in probe order. + /// + /// `every_row_candidate` is whether stage 1 may keep every row of a + /// probe: the query has no prefilter (deletions included) and no upper + /// distance bound. Otherwise which rows a gather selects, even with an + /// infinite threshold, is known only once the probe's sign plane is read, + /// so no probe is routed to the eager scan. Nor is any under the `sparse` + /// gather policy, which never reads a whole plane, or with + /// `LANCE_RQ_LAZY_DENSE_TO_EAGER` off. + fn new( + probe_rows: impl IntoIterator, + k: usize, + every_row_candidate: bool, + config: &LayeredLazyConfig, + ) -> Self { + let rows_before = probe_rows + .into_iter() + .scan(0usize, |rows, probe| { + let before = *rows; + *rows += probe; + Some(before) + }) + .collect::>(); + let fill_rank = rows_before + .iter() + .position(|&rows| rows >= k) + .unwrap_or(rows_before.len()); + let dense_at_fill = match config.dense { + DenseGatherMode::Whole => true, + DenseGatherMode::Sparse => false, + // The threshold once the heap fills keeps about `k` of the rows + // scored by then, so about that fraction of a later probe's rows + // survives stage 1 (more, as the lower bound is looser). Survivors + // are scattered, so their aligned bytes are a larger fraction still. + DenseGatherMode::Cost => rows_before + .get(fill_rank) + .is_some_and(|&rows| k as f64 >= config.dense_bytes_fraction * rows as f64), + }; + Self { + fill_rank, + dense_at_fill, + route_dense: config.dense_to_eager + && every_row_candidate + && config.dense != DenseGatherMode::Sparse, + } + } + + /// Whether the probe at `rank` is scored before the heap can hold `k` + /// rows, so its threshold is +inf and its gather selects every accepted + /// row below the upper bound. + fn certain_dense(&self, rank: usize) -> bool { + 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 { + self.route_dense && (self.certain_dense(rank) || self.dense_at_fill) + } +} + +/// Query-wide inputs of the gathers. +struct LazyFetchContext { + config: LayeredLazyConfig, + forecast: LazyDenseForecast, + upper_bound: Option, + layout: LayeredExLayout, + progress: watch::Receiver, + pre_filter: Arc, + metrics: Arc, + raw_query_context: Option>, +} + +/// Query-wide inputs of the scorer. +#[derive(Clone)] +struct LazyScorer { + key: ArrayRef, + heap_capacity: usize, + lower_bound: Option, + upper_bound: Option, + approx_mode: ApproxMode, + layout: LayeredExLayout, + use_query_residual: bool, + use_residual_scratch: bool, + rq_search_cache: Option>, + raw_query_context: Option>, + scratch_pool: Arc, + metrics: Arc, + progress: Arc>, +} + +/// Aborts a task of the scan when its owner is dropped: the producer when the +/// query is dropped or fails, and a staging or gather step when the producer +/// drops its buffers. Promotions are not tracked and keep running on their own. +struct AbortOnDrop(tokio::task::AbortHandle); + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } +} + +fn scan_cancelled() -> Error { + Error::internal("lazy layered scan ended before scoring every probe") +} + +/// Run one staging or gather step as its own task, returning a future of its +/// result that aborts the task when dropped. +/// +/// The producer keeps the steps in nested ordered buffers and stops polling +/// the staging buffer while the gather buffer is full. A step polled only +/// through its buffer would then freeze inside whatever fair resource it was +/// queued on: a cache that hands a freed slot of a bounded spill queue to its +/// oldest waiter can hand it to a frozen staging step, which never uses it, +/// while the head gather waits behind it forever. As a task, every step runs +/// to completion on its own. The buffers still start a step only when they +/// have room, so they bound the steps in flight and keep results in probe +/// order exactly as before. +fn spawn_scan_step( + step: impl Future> + Send + 'static, +) -> impl Future> + Send + 'static { + let task = tokio::spawn(step); + let abort = AbortOnDrop(task.abort_handle()); + async move { + let joined = task.await; + drop(abort); + match joined { + Ok(result) => result, + Err(err) if err.is_panic() => std::panic::resume_unwind(err.into_panic()), + Err(_) => Err(scan_cancelled()), + } + } +} + +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) + } + + /// Replace this index's lazy scan settings, avoiding the process-wide + /// environment in concurrent tests. + #[cfg(test)] + pub(crate) fn set_layered_lazy_config_for_test(&self, config: LayeredLazyConfig) { + *self + .layered_lazy + .lock() + .unwrap_or_else(|err| err.into_inner()) = config; + } + + /// Probes the lazy scan stages at once: the compute pool's width, like + /// the eager scan's prepare window. + fn lazy_prepare_parallelism(&self) -> usize { + #[cfg(test)] + { + let parallelism = self + .lazy_prepare_parallelism_for_test + .load(std::sync::atomic::Ordering::Relaxed); + if parallelism > 0 { + return parallelism; + } + } + get_num_compute_intensive_cpus().max(1) + } + + /// Override the lazy scan's staging parallelism, or restore the default + /// with 0, so tests stage several probes at once on hosts with few cores. + #[cfg(test)] + pub(crate) fn set_lazy_prepare_parallelism_for_test(&self, parallelism: usize) { + self.lazy_prepare_parallelism_for_test + .store(parallelism, std::sync::atomic::Ordering::Relaxed); + } + + /// The settings for `query` when it can run the lazy scan. Queries that + /// cannot run it take the eager scan and are counted by reason. + pub(super) fn lazy_full_config( + &self, + query: &Query, + has_raw_query_context: bool, + ) -> Option { + if !self.layered_rq { + return None; + } + let config = *self + .layered_lazy + .lock() + .unwrap_or_else(|err| err.into_inner()); + let stats = layered_stats::counters(); + let ineligible = if !config.enabled { + &stats.ineligible_disabled + } else if query.k == 0 { + // Nothing to gather; the eager scan still waits for the prefilter + // and loads every probe, and a `k == 0` query must keep doing so. + &stats.ineligible_k_zero + } else if !self.storage.supports_candidate_reads() { + &stats.ineligible_remapper + } else if query.rq_cascade_factor.is_some() { + &stats.ineligible_cascade + } else if query.rq_precision != RQPrecision::Full { + &stats.ineligible_precision + } else if query.refine_factor.is_some() { + &stats.ineligible_refine + } else if query.approx_mode == ApproxMode::Fast { + &stats.ineligible_fast + } else if self.use_query_residual || !has_raw_query_context { + &stats.ineligible_residual + } else { + return Some(config); + }; + ineligible.incr(); + None + } + + /// Whether the high and low planes of every probed partition are in RAM, + /// checked without counting as accesses. + pub(super) async fn all_ex_planes_resident( + &self, + partitions: &UInt32Array, + probes: Range, + ) -> bool { + for idx in probes { + if !self + .ex_planes_resident(partitions.value(idx) as usize) + .await + { + return false; + } + } + true + } + + /// Score probes `probes` of `partitions` into one global heap with the + /// lazy scan and return the heap as a result batch. + #[allow(clippy::too_many_arguments)] + pub(super) async fn search_partitions_lazy_full( + self: Arc, + query: Query, + partitions: Arc, + q_c_dists: Arc, + probes: Range, + pre_filter: Arc, + metrics: Arc, + raw_query_context: Option>, + config: LayeredLazyConfig, + ) -> Result { + let stats = layered_stats::counters(); + stats.lazy_queries.incr(); + let started = Instant::now(); + pre_filter.wait_for_ready().await?; + let heap_capacity = query.k; + let probe_count = probes.len(); + let layout = self.storage.layered_ex_layout()?; + let forecast = LazyDenseForecast::new( + probes + .clone() + .map(|idx| self.storage.partition_size(partitions.value(idx) as usize)), + heap_capacity, + pre_filter.is_empty() && query.upper_bound.is_none(), + &config, + ); + let (progress_tx, progress_rx) = watch::channel(LazyProgress { + scored: 0, + full: false, + threshold: f32::INFINITY, + }); + let scorer = LazyScorer { + key: query.key.clone(), + heap_capacity, + lower_bound: query.lower_bound, + upper_bound: query.upper_bound, + approx_mode: query.approx_mode, + layout, + use_query_residual: self.use_query_residual, + use_residual_scratch: self.use_residual_scratch, + rq_search_cache: self.rq_search_cache.clone(), + raw_query_context: raw_query_context.clone(), + scratch_pool: self.scratch_pool.clone(), + metrics: metrics.clone(), + progress: Arc::new(progress_tx), + }; + let fetch_context = Arc::new(LazyFetchContext { + config, + forecast, + upper_bound: query.upper_bound, + layout, + progress: progress_rx, + pre_filter: pre_filter.clone(), + metrics: metrics.clone(), + raw_query_context: raw_query_context.clone(), + }); + + let (ready_tx, mut ready_rx) = mpsc::channel(LAZY_READY_CHANNEL_CAPACITY); + let prepare_parallelism = self.lazy_prepare_parallelism(); + let stage_index = self.clone(); + let fetch_index = self.clone(); + let probe_start = probes.start; + let result_metrics = metrics.clone(); + let producer = tokio::spawn(async move { + // `map` spawns a step only when `buffered` pulls it, so at most + // `prepare_parallelism` probes are staged or held staged, and the + // gather window bounds the gathers, as with plain futures. + let fetched = stream::iter(probes) + .map(move |idx| { + let rank = idx - probe_start; + let mut query = query.clone(); + query.dist_q_c = q_c_dists.value(idx); + spawn_scan_step(stage_index.clone().lazy_stage_probe( + rank, + partitions.value(idx) as usize, + query, + forecast.routes_to_eager(rank), + pre_filter.clone(), + metrics.clone(), + raw_query_context.clone(), + )) + }) + .buffered(prepare_parallelism) + .map(move |staged| { + let index = fetch_index.clone(); + 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)); + futures::pin_mut!(fetched); + while let Some(probe) = fetched.next().await { + let failed = probe.is_err(); + if ready_tx.send(probe).await.is_err() || failed { + break; + } + } + }); + let _abort_producer = AbortOnDrop(producer.abort_handle()); + + let mut heap = BinaryHeap::with_capacity(heap_capacity); + let mut scored = 0usize; + let mut first_full = false; + loop { + let waiting = Instant::now(); + let Some(probe) = ready_rx.recv().await else { + break; + }; + stats.scorer_wait_ns.add_elapsed(waiting); + let mut batch = vec![probe?]; + while batch.len() < LAZY_SCORE_BATCH_MAX_PROBES { + match ready_rx.try_recv() { + Ok(probe) => batch.push(probe?), + Err(_) => break, + } + } + let survivor_rows: usize = batch.iter().map(LazyProbe::survivor_rows).sum(); + let has_eager = batch + .iter() + .any(|probe| matches!(probe, LazyProbe::Eager(_))); + if !has_eager && survivor_rows <= config.inline_rows { + stats.dispatch_inline.incr(); + Self::score_lazy_batch(&scorer, batch, &mut heap, &mut scored)?; + } else { + stats.dispatch_spawn.incr(); + let batch_scorer = scorer.clone(); + (heap, scored) = spawn_cpu(move || { + let (mut heap, mut scored) = (heap, scored); + Self::score_lazy_batch(&batch_scorer, batch, &mut heap, &mut scored)?; + Ok::<_, Error>((heap, scored)) + }) + .await?; + } + if !first_full && heap.len() >= heap_capacity { + first_full = true; + stats.time_to_first_full_ns.add_elapsed(started); + } + } + match producer.await { + Ok(()) => {} + Err(err) if err.is_panic() => std::panic::resume_unwind(err.into_panic()), + Err(_) => return Err(scan_cancelled()), + } + if scored != probe_count { + return Err(Error::internal(format!( + "lazy layered scan scored {scored} of {probe_count} probes" + ))); + } + if heap.len() <= GLOBAL_TOPK_INLINE_HEAP_LEN { + Self::global_heap_to_batch(heap, result_metrics.as_ref()) + } else { + let queued = Instant::now(); + spawn_cpu(move || { + result_metrics.record_timing(IndexTiming::CpuQueueWait, queued.elapsed()); + Self::global_heap_to_batch(heap, result_metrics.as_ref()) + }) + .await + } + } + + /// 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`). + #[allow(clippy::too_many_arguments)] + async fn lazy_stage_probe( + self: Arc, + rank: usize, + partition_id: usize, + query: Query, + dense_to_eager: bool, + 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 + }) + } else { + dense_to_eager.then_some(&stats.dense_to_eager) + }; + if let Some(counter) = eager { + counter.incr(rank); + return Ok(LazyStaged::Eager( + self.prepare_partition_without_prefilter_wait( + partition_id, + &query, + pre_filter, + metrics.as_ref(), + raw_query_context, + ) + .await?, + )); + } + let loading = Instant::now(); + let io_stats = metrics.io_stats(); + let (index, sign) = tokio::try_join!( + self.load_sub_index(partition_id, io_stats.clone()), + self.storage + .load_sign_stage(partition_id, &self.index_cache, io_stats), + )?; + stats.sign_load_ns.add_elapsed(loading); + let mut entry = PartitionEntry::new(index, sign); + // Like the eager layered load, a query-assembled entry is never cached whole. + entry.cache_whole_partition = false; + let entry = Arc::new(entry); + let pre_filter = + Self::prefilter_for_partition(&self.index_cache, partition_id, &entry, pre_filter) + .await?; + if entry.storage.is_empty() { + stats.empty.incr(rank); + return Ok(LazyStaged::Empty); + } + let in_flight = self.prepared_partitions.track(); + stats.stage1_dispatches.incr(); + let stage_entry = entry.clone(); + let stage_query = query.clone(); + let rq_search_cache = self.rq_search_cache.clone(); + let scratch_pool = self.scratch_pool.clone(); + let stage = spawn_cpu(move || { + let computing = Instant::now(); + let stage = Self::lazy_sign_stage( + &stage_entry, + partition_id, + &stage_query, + pre_filter.as_ref(), + rq_search_cache.as_deref(), + raw_query_context.as_deref(), + &scratch_pool, + ); + layered_stats::counters() + .stage1_cpu_ns + .add_elapsed(computing); + stage + }) + .await?; + let Some((stage, accept)) = stage else { + stats.gating_off.incr(rank); + return Ok(LazyStaged::GatingOff { + partition_id, + query, + }); + }; + stats.lazy_probes.incr(rank); + Ok(LazyStaged::Lazy(Box::new(LazySignProbe { + rank, + partition_id, + dist_q_c: query.dist_q_c, + entry, + stage, + accept, + _in_flight: in_flight, + }))) + } + + /// 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 + /// entries, and once they were evicted it would otherwise never take the + /// eager scan again. + async fn ex_planes_resident(&self, partition: usize) -> bool { + if self.storage.partition_size(partition) == 0 { + return true; + } + for plane in [1, 2] { + if !self + .index_cache + .peek_resident_with_key(&PlaneKey { partition, plane }) + .await + { + return false; + } + } + true + } + + /// Stage 1 of one partition: the binary inner products and lower bounds + /// of every row, computed with the calculator the eager scan builds, and + /// the prefilter's row acceptance. `None` when the partition cannot gate + /// on the lower bound. + #[allow(clippy::type_complexity)] + fn lazy_sign_stage( + entry: &PartitionEntry, + partition_id: usize, + query: &Query, + pre_filter: &dyn PreFilter, + rq_search_cache: Option<&RabitSearchCache>, + raw_query_context: Option<&RabitRawQueryContext>, + scratch_pool: &QueryScratchPool, + ) -> Result>)>> { + let storage = rabit_storage(&entry.storage)?; + let residual = Self::query_context_for_scratch( + false, + false, + partition_id, + None, + rotated_partition_centroid_slice(rq_search_cache, partition_id), + raw_query_context, + )?; + let mut stage = SignStage::default(); + let gated = scratch_pool.with_scratch(|scratch| { + let calc = storage.dist_calculator_with_scratch( + query.key.clone(), + query.dist_q_c, + residual, + &mut scratch.query_f32, + DistanceCalculatorOptions { + rq_precision: query.rq_precision, + approx_mode: query.approx_mode, + }, + ); + calc.full_sign_stage_with_scratch( + &mut stage, + &mut scratch.u16, + &mut scratch.u8, + &mut scratch.u32, + ) + .is_ok() + }); + if !gated { + return Ok(None); + } + // The eager scan takes its filtered path exactly when the prefilter is not empty. + let accept = (!pre_filter.is_empty()).then(|| { + let mask = pre_filter.mask(); + storage + .row_ids() + .map(|&row_id| mask.selected(row_id)) + .collect() + }); + Ok(Some((stage, accept))) + } + + /// Wait until a probe may be issued, select its survivors against the + /// published threshold, and gather their ex rows. + async fn lazy_fetch( + self: Arc, + staged: LazyStaged, + context: &LazyFetchContext, + ) -> Result> { + let probe = match staged { + LazyStaged::Eager(prepared) => return Ok(LazyProbe::Eager(Box::new(prepared))), + LazyStaged::Empty => return Ok(LazyProbe::Empty), + LazyStaged::GatingOff { + partition_id, + query, + } => { + return Ok(LazyProbe::Eager(Box::new( + self.prepare_partition_without_prefilter_wait( + partition_id, + &query, + context.pre_filter.clone(), + context.metrics.as_ref(), + context.raw_query_context.clone(), + ) + .await?, + ))); + } + LazyStaged::Lazy(probe) => *probe, + }; + let stats = layered_stats::counters(); + let rank = probe.rank; + let gating = Instant::now(); + let mut progress = context.progress.clone(); + let forecast = &context.forecast; + let issue = 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() + } else { + let window = context.config.window; + let gate = *progress + .wait_for(|progress| progress.scored.saturating_add(window) >= rank) + .await + .map_err(|_| scan_cancelled())?; + // 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 + // fill are expected to be dense anyway, or once the probes holding + // `k` rows were scored without filling the heap, after which + // waiting would issue gathers one probe at a time. +inf only + // over-selects: stage 2 prunes on the live heap. + let early = context.config.eager_before_full; + let released = |progress: &LazyProgress| { + progress.full + || progress.scored >= rank + || (early && progress.scored >= forecast.fill_rank) + }; + let issue = if released(&gate) || (early && forecast.dense_at_fill) { + gate + } else { + stats.deferred_issues.incr(); + let issue = *progress + .wait_for(released) + .await + .map_err(|_| scan_cancelled())?; + if !issue.full && issue.scored >= rank { + stats.serial_waits.incr(); + } + issue + }; + if !issue.full && issue.scored < rank { + stats.eager_before_full.incr(); + } + issue + }; + stats.gate_wait_ns.add_elapsed(gating); + let staleness = rank.saturating_sub(issue.scored) as u64; + stats.staleness_sum.add(staleness); + stats.staleness_max.observe(staleness); + + let mut survivors = Vec::new(); + let mut stage1 = StagePruneCounts::default(); + select_full_survivors( + &probe.stage.lower_bounds, + probe.accept.as_deref(), + context.upper_bound, + issue.full.then_some(issue.threshold), + &mut survivors, + &mut stage1, + ); + stats.stage1_candidates.add(rank, stage1.candidates as u64); + stats + .stage1_pruned_ub + .add(rank, stage1.pruned_upper_bound as u64); + stats + .stage1_pruned_heap + .add(rank, stage1.pruned_heap as u64); + stats.rows_fetched.add(rank, survivors.len() as u64); + + let partition_rows = probe.entry.storage.len(); + let mut plan = GatherPlan { + planes: [PlaneSource::Sparse; 2], + origin_whole: [false; 2], + }; + for (slot, plane) in [1u8, 2].into_iter().enumerate() { + 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, + }) + .await + { + PlaneSource::Resident + } else { + plan_plane_gather(&survivors, partition_rows, row_bytes, &context.config) + }; + plan.origin_whole[slot] = plan.planes[slot] == PlaneSource::Sparse + && origin_reads_whole_plane(&survivors, row_bytes, &context.config); + } + let mut gathered = self + .storage + .gather_ex_rows( + probe.partition_id, + &survivors, + plan, + &context.config, + &self.index_cache, + context.metrics.io_stats(), + ) + .await?; + 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) { + self.spawn_lazy_promotion(ticket); + } + Ok(LazyProbe::Ready(Box::new(LazyReadyProbe { + sign: probe, + survivors, + stage1, + gathered, + }))) + } + + /// Load a whole ex plane in the background, off the query's critical + /// path, admitting it as the eager load would. + fn spawn_lazy_promotion(self: &Arc, ticket: LazyPromotionTicket) { + let index = self.clone(); + tokio::spawn(async move { + let stats = layered_stats::counters(); + match index + .storage + .load_plane_entry(ticket.partition, ticket.plane, &index.index_cache, None) + .await + { + Ok(batch) => { + stats + .promotion_bytes + .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, + }) + .await; + if resident { + stats.promotions_completed.incr(); + } else { + // Loaded but evicted before it could serve a query. + stats.promotions_skipped.incr(); + } + } + Err(err) => { + stats.promotions_skipped.incr(); + tracing::warn!( + partition = ticket.partition, + plane = ticket.plane, + "lazy layered plane promotion failed: {err}" + ); + } + } + drop(ticket); + }); + } + + /// Score `batch` in probe order into `heap`, publishing progress after + /// every whole probe, matching the native scan threshold boundary. + fn score_lazy_batch( + scorer: &LazyScorer, + batch: Vec>, + heap: &mut BinaryHeap>, + scored: &mut usize, + ) -> Result<()> { + let stats = layered_stats::counters(); + let scoring = Instant::now(); + let result = batch.into_iter().try_for_each(|probe| { + match probe { + LazyProbe::Eager(prepared) => scorer.scratch_pool.with_scratch(|scratch| { + Self::accumulate_prepared_partition_search( + scorer.use_query_residual, + scorer.use_residual_scratch, + *prepared, + heap, + scratch, + scorer.metrics.as_ref(), + ) + })?, + LazyProbe::Empty => {} + LazyProbe::Ready(ready) => Self::score_ready_probe(scorer, *ready, heap)?, + } + *scored += 1; + let top = (heap.len() >= scorer.heap_capacity) + .then(|| heap.peek().map(|node| node.dist.0)) + .flatten(); + let scored = *scored; + scorer.progress.send_modify(|progress| { + progress.scored = scored; + if let Some(top) = top { + progress.full = true; + progress.threshold = top; + } + }); + Ok(()) + }); + stats.stage2_cpu_ns.add_elapsed(scoring); + result + } + + /// Stage 2 of one partition: check that the gather covers every row the + /// live heap can still take, then replay the survivors through the eager + /// scan's per-row step. + fn score_ready_probe( + scorer: &LazyScorer, + ready: LazyReadyProbe, + heap: &mut BinaryHeap>, + ) -> Result<()> { + let stats = layered_stats::counters(); + let LazyReadyProbe { + sign, + survivors, + stage1, + gathered, + } = ready; + let rank = sign.rank; + let storage = rabit_storage(&sign.entry.storage)?; + scorer.metrics.record_comparisons(storage.len()); + + let live_threshold = (heap.len() >= scorer.heap_capacity) + .then(|| heap.peek().map(|node| node.dist.0)) + .flatten(); + let mut needed = Vec::new(); + let mut needed_counts = StagePruneCounts::default(); + select_full_survivors( + &sign.stage.lower_bounds, + sign.accept.as_deref(), + scorer.upper_bound, + live_threshold, + &mut needed, + &mut needed_counts, + ); + let missing = count_missing(&needed, &survivors); + if missing > 0 { + stats.needed_not_fetched.add(missing as u64); + return Err(Error::internal(format!( + "lazy layered scan of partition {} did not gather {missing} rows the heap still needs", + sign.partition_id + ))); + } + + let expected_rows = match gathered.index { + ExIndex::Identity => storage.len(), + ExIndex::Compact => survivors.len(), + }; + if gathered.high.num_rows() != expected_rows || gathered.low.num_rows() != expected_rows { + return Err(Error::internal(format!( + "lazy layered gather of partition {} returned {} high and {} low rows, expected {expected_rows}", + sign.partition_id, + gathered.high.num_rows(), + gathered.low.num_rows() + ))); + } + let rows = survivors + .iter() + .enumerate() + .map(|(position, &offset)| SurvivorRow { + ex_index: match gathered.index { + ExIndex::Identity => offset, + ExIndex::Compact => position as u32, + }, + row_id: storage.row_id(offset), + binary_ip: sign.stage.binary_ips[offset as usize], + lower_bound: sign.stage.lower_bounds[offset as usize], + }) + .collect::>(); + let ex_rows = ExRows::from_plane_batches( + &gathered.high, + &gathered.low, + scorer.layout.rotated_dim, + scorer.layout.num_bits, + )?; + let residual = Self::query_context_for_scratch( + scorer.use_query_residual, + scorer.use_residual_scratch, + sign.partition_id, + None, + rotated_partition_centroid_slice(scorer.rq_search_cache.as_deref(), sign.partition_id), + scorer.raw_query_context.as_deref(), + )?; + let counters = scorer.scratch_pool.with_scratch(|scratch| { + let calc = storage.dist_calculator_with_ex_rows( + scorer.key.clone(), + sign.dist_q_c, + residual, + &mut scratch.query_f32, + DistanceCalculatorOptions { + rq_precision: RQPrecision::Full, + approx_mode: scorer.approx_mode, + }, + ex_rows, + )?; + Ok::<_, Error>(calc.accumulate_survivor_rows( + scorer.heap_capacity, + scorer.lower_bound, + scorer.upper_bound, + &rows, + stage1, + heap, + |_threshold| {}, + )) + })?; + stats.stage2_exact.add(rank, counters.exact as u64); + stats.stage2_pruned_live.add( + rank, + counters.pruned_heap.saturating_sub(stage1.pruned_heap) as u64, + ); + stats + .stage2_exact_rejected + .add(rank, counters.exact_rejected as u64); + Ok(()) + } +} + +fn rabit_storage(storage: &T) -> Result<&RabitQuantizationStorage> { + storage + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::internal("lazy layered scan requires RaBitQ storage")) +} + +/// Rows of the sorted `needed` missing from the sorted `gathered`. +fn count_missing(needed: &[u32], gathered: &[u32]) -> usize { + let mut gathered = gathered.iter().peekable(); + needed + .iter() + .filter(|&&row| { + while gathered.next_if(|&&candidate| candidate < row).is_some() {} + gathered.next_if_eq(&&row).is_none() + }) + .count() +} + +fn record_gather_sources(rank: usize, gathered: &GatheredEx) { + let stats = layered_stats::counters(); + let counters = [ + [&stats.high_resident, &stats.high_whole, &stats.high_sparse], + [&stats.low_resident, &stats.low_whole, &stats.low_sparse], + ]; + for (source, counters) in gathered.sources.iter().zip(counters) { + let counter = match source { + Some(PlaneSource::Resident) => counters[0], + Some(PlaneSource::Whole) => counters[1], + Some(PlaneSource::Sparse) => counters[2], + None => continue, + }; + counter.incr(rank); + } +} + +#[cfg(test)] +mod tests { + use super::{DenseGatherMode, LayeredLazyConfig, LazyDenseForecast, count_missing}; + + #[test] + fn dense_forecast_follows_k_and_partition_sizes() { + let config = LayeredLazyConfig { + dense_bytes_fraction: 0.5, + ..Default::default() + }; + // 0, 100, 100 and 150 rows before each probe. + let rows = [100, 0, 50, 200]; + let predicted = |forecast: LazyDenseForecast| { + (0..rows.len()) + .map(|rank| (forecast.certain_dense(rank), forecast.routes_to_eager(rank))) + .collect::>() + }; + + // The first probe fills the heap, whose threshold then keeps about a + // tenth of the rows scored: later gathers are sparse. + let small = LazyDenseForecast::new(rows, 10, true, &config); + assert_eq!( + small, + LazyDenseForecast { + fill_rank: 1, + dense_at_fill: false, + route_dense: true + } + ); + assert_eq!( + predicted(small), + [(true, true), (false, false), (false, false), (false, false)] + ); + + // The heap fills from rank 3, keeping 120 of 150 rows: every probe is dense. + let large = LazyDenseForecast::new(rows, 120, true, &config); + assert_eq!( + large, + LazyDenseForecast { + fill_rank: 3, + dense_at_fill: true, + route_dense: true + } + ); + assert_eq!( + predicted(large), + [(true, true), (true, true), (true, true), (false, true)] + ); + + // The probes never hold `k` rows, so every threshold is +inf. + let unfilled = LazyDenseForecast::new(rows, 1000, true, &config); + assert_eq!( + unfilled, + LazyDenseForecast { + fill_rank: rows.len(), + dense_at_fill: false, + route_dense: true + } + ); + assert_eq!(predicted(unfilled), [(true, true); 4]); + + // A fixed gather mode overrides the cost rule; whole gathers route + // every probe. + for (dense, dense_at_fill) in [ + (DenseGatherMode::Whole, true), + (DenseGatherMode::Sparse, false), + ] { + for k in [10, 120] { + 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)); + assert_eq!(routed, dense == DenseGatherMode::Whole, "{dense:?} k={k}"); + } + } + + // 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; + // their certain-dense gathers are still issued at once. + let sparse = LayeredLazyConfig { + dense: DenseGatherMode::Sparse, + ..config + }; + let whole = LayeredLazyConfig { + dense: DenseGatherMode::Whole, + ..config + }; + let off = LayeredLazyConfig { + dense_to_eager: false, + ..config + }; + for k in [10, 120, 1000] { + let fill_rank = LazyDenseForecast::new(rows, k, true, &config).fill_rank; + for (every_row_candidate, config) in + [(true, sparse), (true, off), (false, config), (false, whole)] + { + let forecast = LazyDenseForecast::new(rows, k, every_row_candidate, &config); + let context = format!("{config:?} every_row_candidate={every_row_candidate} k={k}"); + assert_eq!(forecast.fill_rank, fill_rank, "{context}"); + assert!( + predicted(forecast).iter().all(|&(_, routed)| !routed), + "{context}" + ); + } + } + } + + #[test] + fn count_missing_compares_sorted_offsets() { + assert_eq!(count_missing(&[], &[1, 2]), 0); + assert_eq!(count_missing(&[1, 3], &[1, 2, 3]), 0); + assert_eq!(count_missing(&[0, 3, 7], &[1, 3, 5]), 2); + assert_eq!(count_missing(&[4], &[]), 1); + } +}