diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 3762c5b1471..2b1502f1c45 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -4176,6 +4176,7 @@ dependencies = [ "tokio", "tracing", "twox-hash", + "uuid", ] [[package]] diff --git a/protos/ann.proto b/protos/ann.proto index b5a390b507e..efd97fd3e41 100644 --- a/protos/ann.proto +++ b/protos/ann.proto @@ -61,6 +61,10 @@ message VectorQueryProto { * indexes, such as IVF_RQ. Other index types ignore this setting. */ VectorApproxMode approx_mode = 14; + // Empty selects full precision; otherwise sign, high or full. + string rq_precision = 17; + // Absent scans full codes; present reranks k times this positive factor. + optional uint32 rq_cascade_factor = 18; /* Which covering columns the index must materialize for this query. Absent means * "not computed" -- see CoveringProjection. Carried across the wire so a remote * executor declares the same search output schema the planner did; without it the diff --git a/python/Cargo.lock b/python/Cargo.lock index e6a298cf171..c3656d9f9fe 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4361,6 +4361,7 @@ dependencies = [ "tokio", "tracing", "twox-hash", + "uuid", ] [[package]] diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index 4ce813d4c84..4a9d8785103 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -79,6 +79,8 @@ async-channel = "2.3.1" rand_distr.workspace = true rangemap.workspace = true +uuid.workspace = true + [dev-dependencies] approx.workspace = true criterion.workspace = true @@ -92,7 +94,6 @@ test-log.workspace = true rstest.workspace = true serial_test.workspace = true chrono.workspace = true -uuid.workspace = true [features] geo = ["dep:lance-geo", "lance-geo/geo", "dep:geoarrow-array", "dep:geoarrow-schema", "dep:geo-types"] diff --git a/rust/lance-index/benches/rq.rs b/rust/lance-index/benches/rq.rs index 849740bdce1..03197a40c6a 100644 --- a/rust/lance-index/benches/rq.rs +++ b/rust/lance-index/benches/rq.rs @@ -252,6 +252,7 @@ fn ex_code_storage_load(c: &mut Criterion) { code_dim: LOAD_DIM as u32, num_bits: NUM_BITS, packed: true, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let code_len = LOAD_DIM / 8; @@ -465,7 +466,10 @@ fn ex_bulk_paths(c: &mut Criterion) { 0.0, None, &mut f32_scratch, - DistanceCalculatorOptions { approx_mode }, + DistanceCalculatorOptions { + rq_precision: Default::default(), + approx_mode, + }, ); let mut dists = Vec::new(); let mut u16_scratch = Vec::new(); @@ -613,7 +617,10 @@ fn heap_topk(c: &mut Criterion) { 1.0, None, &mut f32_scratch, - DistanceCalculatorOptions { approx_mode }, + DistanceCalculatorOptions { + rq_precision: Default::default(), + approx_mode, + }, ); let mut heap = BinaryHeap::with_capacity(TOPK_K + 1); let mut dists = Vec::new(); @@ -670,6 +677,7 @@ fn build_transform(c: &mut Criterion) { code_dim: dim as u32, num_bits, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let rq = RabitQuantizer::try_from( diff --git a/rust/lance-index/protos-cache/cache.proto b/rust/lance-index/protos-cache/cache.proto index 80298e96633..be504ae8ecd 100644 --- a/rust/lance-index/protos-cache/cache.proto +++ b/rust/lance-index/protos-cache/cache.proto @@ -217,4 +217,5 @@ message RabitPartitionHeader { optional bytes fast_rotation_signs = 5; // Estimator the RabitQ storage uses at query time (residual vs raw query). RabitQueryEstimator query_estimator = 6; + bool layered = 7; } diff --git a/rust/lance-index/src/vector.rs b/rust/lance-index/src/vector.rs index d99faf60282..5dd3205f77d 100644 --- a/rust/lance-index/src/vector.rs +++ b/rust/lance-index/src/vector.rs @@ -219,6 +219,10 @@ pub struct Query { /// This currently only affects RQ-quantized vector indexes, such as IVF_RQ. /// Other index types ignore this setting. pub approx_mode: ApproxMode, + /// Precision of a layered RaBitQ query; defaults to full. + pub rq_precision: bq::layered::RQPrecision, + /// Candidate oversampling for a high-to-full cascade. None scans full codes. + pub rq_cascade_factor: Option, } impl Query { diff --git a/rust/lance-index/src/vector/bq.rs b/rust/lance-index/src/vector/bq.rs index da7fbf0e549..e61ad93702c 100644 --- a/rust/lance-index/src/vector/bq.rs +++ b/rust/lance-index/src/vector/bq.rs @@ -22,6 +22,7 @@ pub(crate) mod dist_table_quant; pub mod ex_dot; pub mod layered; pub(crate) mod pairwise; +mod plane_cache; pub mod prune; pub mod rotation; pub mod storage; @@ -124,6 +125,8 @@ pub struct RQBuildParams { /// rotates vectors identically. This is transient build-time state and is never /// persisted to the `RabitQuantization` params proto. pub rotation: Option, + /// Store prefix code planes; supported for 3 through 9 bits. Default off. + pub layered: bool, } pub fn validate_rq_num_bits(num_bits: u8) -> Result<()> { @@ -160,11 +163,18 @@ pub fn rabit_ex_code_bytes(rotated_dim: usize, ex_bits: u8) -> Result { } impl RQBuildParams { + /// Opt into the layered format. Requires a reader with layered RQ support. + pub fn with_layered(mut self, layered: bool) -> Self { + self.layered = layered; + self + } + pub fn new(num_bits: u8) -> Self { Self { num_bits, rotation_type: RQRotationType::default(), rotation: None, + layered: false, } } @@ -173,6 +183,7 @@ impl RQBuildParams { num_bits, rotation_type, rotation: None, + layered: false, } } } @@ -182,7 +193,7 @@ impl From<&RQBuildParams> for RabitQuantization { use crate::pb::vector_index_details::rabit_quantization::{RotationType, RowLayout}; Self { num_bits: value.num_bits as u32, - layered: false, + layered: value.layered, row_layout: RowLayout::Columns as i32, rotation_type: match value.rotation_type { RQRotationType::Fast => RotationType::Fast as i32, @@ -204,6 +215,7 @@ impl Default for RQBuildParams { num_bits: RABIT_DEFAULT_NUM_BITS, rotation_type: RQRotationType::default(), rotation: None, + layered: false, } } } diff --git a/rust/lance-index/src/vector/bq/builder.rs b/rust/lance-index/src/vector/bq/builder.rs index 7acca670de3..ade108a2cf6 100644 --- a/rust/lance-index/src/vector/bq/builder.rs +++ b/rust/lance-index/src/vector/bq/builder.rs @@ -63,11 +63,8 @@ pub struct RabitQuantizer { pub(crate) struct RabitQuantizedBatch { pub binary_codes: ArrayRef, pub ex_codes: Option, - #[cfg(test)] pub ex_res_dot_dists: Option>, - #[cfg(test)] pub rotated_residuals: Option>, - #[cfg(test)] pub ex_code_values: Option>, } @@ -542,11 +539,8 @@ impl RabitQuantizer { } else { None }, - #[cfg(test)] rotated_residuals: None, - #[cfg(test)] ex_code_values: None, - #[cfg(test)] ex_res_dot_dists: None, }) } @@ -583,6 +577,7 @@ impl RabitQuantizer { code_dim: code_dim as u32, num_bits, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, } } @@ -594,6 +589,7 @@ impl RabitQuantizer { code_dim: code_dim as u32, num_bits, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }, }; @@ -1035,11 +1031,8 @@ impl RabitQuantizer { .map(|array| Arc::new(array) as ArrayRef) }) .transpose()?, - #[cfg(test)] ex_res_dot_dists, - #[cfg(test)] rotated_residuals: Some(rotated_residuals), - #[cfg(test)] ex_code_values, }) } @@ -1056,6 +1049,9 @@ impl Quantization for RabitQuantizer { params: &Self::BuildParams, ) -> Result { validate_rq_num_bits(params.num_bits)?; + if params.layered { + super::layered::RQLayout::try_new(params.num_bits)?; + } let dim = data.as_fixed_size_list().value_length() as usize; if !dim.is_multiple_of(u8::BITS as usize) { @@ -1063,11 +1059,12 @@ impl Quantization for RabitQuantizer { "vector dimension must be divisible by 8 for IVF_RQ", )); } - if let Some(q) = Self::from_supplied_rotation(params, dim)? { + if let Some(mut q) = Self::from_supplied_rotation(params, dim)? { + q.metadata.layered = params.layered; return Ok(q); } - let q = match data.as_fixed_size_list().value_type() { + let mut q = match data.as_fixed_size_list().value_type() { DataType::Float16 => Self::new_with_rotation::( params.num_bits, data.as_fixed_size_list().value_length(), @@ -1090,6 +1087,7 @@ impl Quantization for RabitQuantizer { ))); } }; + q.metadata.layered = params.layered; Ok(q) } @@ -1152,6 +1150,14 @@ impl Quantization for RabitQuantizer { _: lance_linalg::distance::DistanceType, ) -> Result { validate_rq_num_bits(metadata.num_bits)?; + if metadata.layered { + super::layered::RQLayout::try_new(metadata.num_bits)?; + if metadata.query_estimator != RabitQueryEstimator::RawQuery { + return Err(Error::invalid_input( + "layered IVF_RQ requires the raw-query estimator", + )); + } + } Ok(Quantizer::Rabit(Self { metadata: metadata.clone(), })) @@ -1166,6 +1172,10 @@ impl Quantization for RabitQuantizer { if self.metadata.query_estimator == RabitQueryEstimator::RawQuery { fields.push(ERROR_FACTORS_FIELD.clone()); } + if self.metadata.layered { + return super::layered::storage_fields(self.code_dim(), self.metadata.num_bits, fields) + .expect("layered layout validated at build"); + } if let Some(ex_code_field) = rabit_ex_code_field(self.code_dim(), self.metadata.num_bits) .expect("RabitQ num_bits should be validated") { diff --git a/rust/lance-index/src/vector/bq/layered.rs b/rust/lance-index/src/vector/bq/layered.rs index 4e4b4fb4d5f..81813e1d215 100644 --- a/rust/lance-index/src/vector/bq/layered.rs +++ b/rust/lance-index/src/vector/bq/layered.rs @@ -36,3 +36,663 @@ impl RQLayout { }) } } + +use super::ex_dot::{blocked_ex_code_bytes, pack_blocked_row}; +use super::storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN}; +use super::transform::{EX_ADD_FACTORS_FIELD, EX_SCALE_FACTORS_FIELD}; +use arrow_array::{ArrayRef, FixedSizeListArray, Float32Array, UInt8Array}; +use arrow_schema::{DataType, Field}; +use lance_arrow::FixedSizeListArrayExt; +use std::sync::Arc; + +pub const HIGH_ADD_FACTORS_COLUMN: &str = "__add_factors_ex_hi"; +pub const HIGH_SCALE_FACTORS_COLUMN: &str = "__scale_factors_ex_hi"; +pub const HIGH_BOUNDS_COLUMN: &str = "__rq_bounds_hi"; +pub const FULL_BOUNDS_COLUMN: &str = "__rq_bounds_full"; + +// Each level stores ||w_level-w_sign||, |add_level-add_sign| and a +// coefficient for floating-point/LUT error. These bound the stored estimator, +// rather than its error relative to the unquantized vector. +fn bounds_field(name: &str) -> Field { + Field::new( + name, + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 3), + true, + ) +} + +pub(crate) fn estimator_bounds( + residuals: &[f32], + codes: &[u8], + dim: usize, + bits: u8, + binary: &super::transform::RabitRawQueryFactors, + level: &super::transform::RabitRawQueryFactors, +) -> Result { + let adds = level + .ex_add_factors + .as_ref() + .ok_or_else(|| Error::internal("missing level add factors"))?; + let scales = level + .ex_scale_factors + .as_ref() + .ok_or_else(|| Error::internal("missing level scale factors"))?; + let code_scale = (1u32 << bits) as f64; + let bias = -(code_scale - 0.5); + let mut bounds = Vec::with_capacity(residuals.len() / dim * 3); + for (row, (residual, codes)) in residuals + .chunks_exact(dim) + .zip(codes.chunks_exact(dim)) + .enumerate() + { + let binary_scale = binary.scale_factors.value(row) as f64; + let scale = scales.value(row) as f64; + let norm = residual + .iter() + .zip(codes) + .map(|(&r, &c)| { + let sign = f64::from(u8::from(r.is_sign_positive())); + let delta = + (code_scale * sign + f64::from(c) + bias) * scale - (sign - 0.5) * binary_scale; + delta * delta + }) + .sum::() + .sqrt(); + for value in [ + norm, + (f64::from(adds.value(row)) - f64::from(binary.add_factors.value(row))).abs(), + binary_scale.abs() + code_scale * scale.abs(), + ] { + let rounded = value as f32; + bounds.push(if f64::from(rounded) < value { + rounded.next_up() + } else { + rounded + }); + } + } + Ok(Arc::new(FixedSizeListArray::try_new_from_values( + Float32Array::from(bounds), + 3, + )?)) +} + +pub(crate) fn storage_fields( + dim: usize, + num_bits: u8, + mut fields: Vec, +) -> Result> { + let layout = RQLayout::try_new(num_bits)?; + for (name, bits) in [ + (RABIT_BLOCKED_EX_CODE_COLUMN, layout.high_bits), + (RABIT_BLOCKED_EX_CODE_LO_COLUMN, layout.low_bits), + ] { + fields.push(Field::new( + name, + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::UInt8, true)), + blocked_ex_code_bytes(dim, bits) as i32, + ), + true, + )); + } + fields.extend([ + EX_ADD_FACTORS_FIELD.clone(), + EX_SCALE_FACTORS_FIELD.clone(), + Field::new(HIGH_ADD_FACTORS_COLUMN, DataType::Float32, true), + Field::new(HIGH_SCALE_FACTORS_COLUMN, DataType::Float32, true), + bounds_field(HIGH_BOUNDS_COLUMN), + bounds_field(FULL_BOUNDS_COLUMN), + ]); + Ok(fields) +} + +/// Split already-quantized full codes. No second quantization or scale search. +pub(crate) fn split_codes( + values: &[u8], + dim: usize, + layout: RQLayout, +) -> Result<(ArrayRef, ArrayRef, Vec)> { + let rows = values.len() / dim; + let hi_bytes = blocked_ex_code_bytes(dim, layout.high_bits); + let lo_bytes = blocked_ex_code_bytes(dim, layout.low_bits); + let mut hi = vec![0; rows * hi_bytes]; + let mut lo = vec![0; rows * lo_bytes]; + let hi_values: Vec = values.iter().map(|v| v >> layout.low_bits).collect(); + let mask = ((1u16 << layout.low_bits) - 1) as u8; + let mut low_values = vec![0; dim]; + for row in 0..rows { + for (dst, src) in low_values + .iter_mut() + .zip(&values[row * dim..(row + 1) * dim]) + { + *dst = src & mask; + } + pack_blocked_row( + &hi_values[row * dim..(row + 1) * dim], + layout.high_bits, + &mut hi[row * hi_bytes..(row + 1) * hi_bytes], + ); + pack_blocked_row( + &low_values, + layout.low_bits, + &mut lo[row * lo_bytes..(row + 1) * lo_bytes], + ); + } + Ok(( + Arc::new(FixedSizeListArray::try_new_from_values( + UInt8Array::from(hi), + hi_bytes as i32, + )?), + Arc::new(FixedSizeListArray::try_new_from_values( + UInt8Array::from(lo), + lo_bytes as i32, + )?), + hi_values, + )) +} + +/// Query precision on a layered IVF_RQ index. Other layouts accept only `Full`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RQPrecision { + /// Scan sign codes with the binary estimator. + Sign, + /// Scan sign and high codes with the high-level factor pair. + High, + /// Evaluate all stored bits. + #[default] + Full, +} + +impl std::str::FromStr for RQPrecision { + type Err = Error; + fn from_str(value: &str) -> Result { + match value { + "sign" => Ok(Self::Sign), + "high" => Ok(Self::High), + "full" => Ok(Self::Full), + _ => Err(Error::invalid_input(format!( + "rq_precision must be sign, high or full, got {value}" + ))), + } + } +} + +use arrow_array::RecordBatch; +use lance_core::cache::{CacheCodec, CacheKey}; +use lance_core::deepsize::{Context, DeepSizeOf}; +use std::borrow::Cow; + +/// A plane and its factors share one persistent cache entry. +#[derive(Debug)] +pub struct PlaneBatch(pub RecordBatch); +impl DeepSizeOf for PlaneBatch { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.0.deep_size_of_children(context) + } +} +/// Plane keys live under the immutable index UUID namespace. +pub struct PlaneKey { + pub partition: usize, + pub plane: u8, +} +impl CacheKey for PlaneKey { + type ValueType = PlaneBatch; + fn key(&self) -> Cow<'_, str> { + format!("{}:{}", self.partition, self.plane).into() + } + fn type_name() -> &'static str { + "RQPlane" + } + fn codec_for_key(&self) -> Option { + Self::codec().map(|codec| { + codec + .with_plane_tag(self.plane) + .with_memory_priority(3 - self.plane) + }) + } + fn codec() -> Option { + Some(CacheCodec::from_impl::()) + } +} + +/// Column projection for sign (0), high (1), or low (2), including factors. +pub fn plane_columns(plane: u8) -> &'static [&'static str] { + use super::storage::RABIT_CODE_COLUMN; + use super::transform::{ + ADD_FACTORS_COLUMN, ERROR_FACTORS_COLUMN, EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + }; + match plane { + 0 => &[ + lance_core::ROW_ID, + RABIT_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + HIGH_BOUNDS_COLUMN, + FULL_BOUNDS_COLUMN, + ], + 1 => &[ + RABIT_BLOCKED_EX_CODE_COLUMN, + HIGH_ADD_FACTORS_COLUMN, + HIGH_SCALE_FACTORS_COLUMN, + ], + 2 => &[ + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + ], + _ => &[], + } +} + +/// A quantized candidate tied to an immutable index and its physical partition. +/// +/// Distributed executors may perform a global distance cut on these values, then +/// return the survivors to the same index for full-code reranking with the same query. +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct QuantizedCandidate { + pub index_uuid: uuid::Uuid, + pub partition_id: usize, + pub row_offset: u32, + pub row_id: u64, + pub distance: f32, + pub centroid_distance: f32, + /// Original-partition binary LUT value, including its SIMD rounding policy. + pub binary_inner_product: f32, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::vector::bq::{ + RQBuildParams, builder::RabitQuantizer, storage::RabitQuantizationStorage, + transform::RQTransformer, + }; + use crate::vector::quantizer::{Quantization, QuantizerStorage}; + use crate::vector::storage::{ + DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore, + }; + use crate::vector::transform::Transformer; + use crate::vector::{ApproxMode, CENTROID_DIST_COLUMN, PART_ID_COLUMN}; + use arrow_array::types::Float32Type; + use arrow_array::{Float32Array, UInt32Array, UInt64Array, cast::AsArray}; + use lance_arrow::RecordBatchExt; + use lance_linalg::distance::DistanceType; + use rstest::rstest; + + #[rstest] + #[case::rq3_l2(3, DistanceType::L2)] + #[case::rq4_l2(4, DistanceType::L2)] + #[case::rq5_l2(5, DistanceType::L2)] + #[case::rq6_l2(6, DistanceType::L2)] + #[case::rq8_l2(8, DistanceType::L2)] + #[case::rq7_l2(7, DistanceType::L2)] + #[case::rq9_l2(9, DistanceType::L2)] + #[case::rq3_dot(3, DistanceType::Dot)] + #[case::rq4_dot(4, DistanceType::Dot)] + #[case::rq5_dot(5, DistanceType::Dot)] + #[case::rq6_dot(6, DistanceType::Dot)] + #[case::rq8_dot(8, DistanceType::Dot)] + #[case::rq7_dot(7, DistanceType::Dot)] + #[case::rq9_dot(9, DistanceType::Dot)] + fn layered_roundtrip_and_levels( + #[case] bits: u8, + #[case] distance_type: DistanceType, + #[values(64, 72)] dim: usize, + #[values(0.0, 0.25)] centroid_value: f32, + ) { + const ROWS: usize = 37; + let values: Vec = (0..dim * ROWS) + .map(|i| ((i * 17 % 101) as f32 - 50.) / 50.) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values.clone()), dim as i32) + .unwrap(); + let rq = RabitQuantizer::build( + &vectors, + distance_type, + &RQBuildParams::new(bits).with_layered(true), + ) + .unwrap(); + let norms = Float32Array::from( + values + .chunks(dim) + .map(|r| r.iter().map(|v| v * v).sum::()) + .collect::>(), + ); + let input = RecordBatch::try_from_iter(vec![ + ("vector", Arc::new(vectors.clone()) as ArrayRef), + ( + lance_core::ROW_ID, + Arc::new(UInt64Array::from((0..ROWS as u64).collect::>())) as ArrayRef, + ), + ( + PART_ID_COLUMN, + Arc::new(UInt32Array::from(vec![0; ROWS])) as ArrayRef, + ), + (CENTROID_DIST_COLUMN, Arc::new(norms) as ArrayRef), + ]) + .unwrap(); + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![centroid_value; dim]), + dim as i32, + ) + .unwrap(); + let rotated_centroid = rq + .quantize_split(¢roids) + .unwrap() + .rotated_residuals + .unwrap(); + let batch = RQTransformer::new(rq.clone(), distance_type, centroids.clone(), "vector") + .unwrap() + .transform(&input) + .unwrap(); + let full = RabitQuantizationStorage::try_from_batch( + batch.clone(), + rq.metadata_ref(), + distance_type, + None, + ) + .unwrap(); + let mut metadata = rq.metadata_ref().clone(); + metadata.layered = false; + let native_rq = RabitQuantizer::try_from( + RabitQuantizer::from_metadata(&metadata, distance_type).unwrap(), + ) + .unwrap(); + let layered_codes = rq.quantize_split(&vectors).unwrap(); + let native_codes = native_rq.quantize_split(&vectors).unwrap(); + assert_eq!( + layered_codes.binary_codes.as_ref(), + native_codes.binary_codes.as_ref() + ); + assert_eq!(layered_codes.ex_code_values, native_codes.ex_code_values); + assert_eq!( + layered_codes.ex_res_dot_dists, + native_codes.ex_res_dot_dists + ); + let single_batch = RQTransformer::new(native_rq, distance_type, centroids, "vector") + .unwrap() + .transform(&input) + .unwrap(); + let single = + RabitQuantizationStorage::try_from_batch(single_batch, &metadata, distance_type, None) + .unwrap(); + let dist_q_c = match distance_type { + DistanceType::L2 => values[..dim] + .iter() + .map(|v| (v - centroid_value).powi(2)) + .sum(), + DistanceType::Dot => 1.0 - values[..dim].iter().sum::() * centroid_value, + _ => unreachable!(), + }; + let query: ArrayRef = Arc::new(Float32Array::from(values[..dim].to_vec())); + assert_eq!( + full.dist_calculator(query.clone(), dist_q_c) + .distance_all(ROWS), + single + .dist_calculator(query.clone(), dist_q_c) + .distance_all(ROWS) + ); + for invalid in [f32::NAN, -1.0] { + let hints = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![invalid; ROWS * 3]), + 3, + ) + .unwrap(); + let invalid_batch = batch + .replace_column_by_name(HIGH_BOUNDS_COLUMN, Arc::new(hints)) + .unwrap(); + let store = RabitQuantizationStorage::try_from_batch_at_precision( + invalid_batch, + rq.metadata_ref(), + distance_type, + None, + RQPrecision::High, + ) + .unwrap(); + let calc = store.dist_calculator(query.clone(), dist_q_c); + assert_eq!( + calc.distance_all(ROWS), + full.dist_calculator_with_scratch( + query.clone(), + dist_q_c, + None, + &mut Vec::new(), + DistanceCalculatorOptions { + approx_mode: ApproxMode::Normal, + rq_precision: RQPrecision::High, + }, + ) + .distance_all(ROWS) + ); + for (row, ip) in calc.binary_inner_products().into_iter().enumerate() { + assert_eq!(calc.raw_query_lower_bound(row, ip), Some(f32::NEG_INFINITY)); + } + } + for precision in [RQPrecision::Sign, RQPrecision::High, RQPrecision::Full] { + let projected = RabitQuantizationStorage::try_from_batch_at_precision( + batch.clone(), + rq.metadata_ref(), + distance_type, + None, + precision, + ) + .unwrap(); + for centered in [false, true] { + let context = || { + centered.then_some(QueryResidual::RabitRawQuery { + rotated_centroid: Some(&rotated_centroid), + query: None, + }) + }; + for approx_mode in [ApproxMode::Normal, ApproxMode::Accurate] { + let mut scratch = Vec::new(); + let mut projected_scratch = Vec::new(); + let calc = full.dist_calculator_with_scratch( + query.clone(), + dist_q_c, + context(), + &mut scratch, + DistanceCalculatorOptions { + approx_mode, + rq_precision: precision, + }, + ); + let projected_calc = projected.dist_calculator_with_scratch( + query.clone(), + dist_q_c, + context(), + &mut projected_scratch, + DistanceCalculatorOptions { + approx_mode, + rq_precision: RQPrecision::Full, + }, + ); + assert_eq!(calc.distance_all(ROWS), projected_calc.distance_all(ROWS)); + if precision == RQPrecision::Full { + let mut native_scratch = Vec::new(); + let native = single.dist_calculator_with_scratch( + query.clone(), + dist_q_c, + context(), + &mut native_scratch, + DistanceCalculatorOptions { + approx_mode, + rq_precision: RQPrecision::Full, + }, + ); + let binary = calc.binary_inner_products(); + assert_eq!(binary, native.binary_inner_products()); + assert_eq!(calc.distance_all(ROWS), native.distance_all(ROWS)); + for (row, &ip) in binary.iter().enumerate() { + assert_eq!( + calc.raw_query_lower_bound(row, ip), + native.raw_query_lower_bound(row, ip), + "full pruning differs from native at row {row}, mode={approx_mode:?}", + ); + } + } + if precision == RQPrecision::High && approx_mode == ApproxMode::Accurate { + let binary = calc.binary_inner_products(); + let distances = calc.distance_all(ROWS); + for (row, &ip) in binary.iter().enumerate() { + let bound = calc.raw_query_lower_bound(row, ip).unwrap(); + assert!( + bound <= distances[row], + "invalid bound at row {row}: {bound} > {}", + distances[row] + ); + } + let best = distances.iter().copied().min_by(f32::total_cmp).unwrap(); + assert!( + binary + .iter() + .enumerate() + .any(|(row, &ip)| calc.raw_query_lower_bound(row, ip).unwrap() + > best), + "fixture must exercise pruning" + ); + } + // Exercise SIMD plus tail rows, both context forms, and range cuts. + for (lower, upper) in [(None, None), (Some(-0.5), Some(20.0))] { + let distances = calc.distance_all(ROWS); + let mut expected: Vec<_> = distances + .iter() + .copied() + .enumerate() + .filter(|(_, d)| { + lower.is_none_or(|v| *d >= v) && upper.is_none_or(|v| *d < v) + }) + .map(|(id, d)| (id as u64, d)) + .collect(); + expected.sort_by(|a, b| a.1.total_cmp(&b.1)); + expected.truncate(3); + let mut heap = std::collections::BinaryHeap::new(); + calc.accumulate_topk_with_scratch( + 3, + lower, + upper, + u64::from, + &mut heap, + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new(), + ); + let actual: Vec<_> = heap + .into_sorted_vec() + .into_iter() + .map(|n| (n.id, n.dist.0)) + .collect(); + let case = format!( + "bits={bits} precision={precision:?} dim={dim} centered={centered} mode={approx_mode:?} range={lower:?}..{upper:?}" + ); + if precision == RQPrecision::Full { + let mut native_scratch = Vec::new(); + let native = single.dist_calculator_with_scratch( + query.clone(), + dist_q_c, + context(), + &mut native_scratch, + DistanceCalculatorOptions { + approx_mode, + rq_precision: RQPrecision::Full, + }, + ); + let mut native_heap = std::collections::BinaryHeap::new(); + native.accumulate_topk_with_scratch( + 3, + lower, + upper, + u64::from, + &mut native_heap, + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new(), + &mut Vec::new(), + ); + let native_actual: Vec<_> = native_heap + .into_sorted_vec() + .into_iter() + .map(|n| (n.id, n.dist.0)) + .collect(); + assert_eq!(actual, native_actual, "{case}"); + } + // Normal mode (and Full in both modes) prunes with native's + // statistical confidence bound, which a random rotation can violate + // for a true top-k row. Pruned top-k is only guaranteed exact when + // every expected row's bound holds. + let binary = calc.binary_inner_products(); + let bounds_hold = expected.iter().all(|&(id, d)| { + calc.raw_query_lower_bound(id as usize, binary[id as usize]) + .is_none_or(|bound| bound <= d) + }); + if bounds_hold { + // Equal distances have no defined order in the heap output. + let by_distance_then_id = |a: &(u64, f32), b: &(u64, f32)| { + a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)) + }; + let mut actual = actual; + actual.sort_by(by_distance_then_id); + expected.sort_by(by_distance_then_id); + assert_eq!(actual, expected, "{case}"); + } + } + } + } + } + let missing = batch.drop_column(RABIT_BLOCKED_EX_CODE_LO_COLUMN).unwrap(); + assert!( + RabitQuantizationStorage::try_from_batch( + missing, + rq.metadata_ref(), + distance_type, + None + ) + .unwrap_err() + .to_string() + .contains("missing column") + ); + assert_eq!( + batch[HIGH_ADD_FACTORS_COLUMN] + .as_primitive::() + .len(), + ROWS + ); + } + + #[test] + fn plane_key_codec_tags_plane_and_priority() { + for plane in 0..=2u8 { + let codec = PlaneKey { + partition: 7, + plane, + } + .codec_for_key() + .unwrap(); + assert_eq!(codec.plane_tag(), Some(plane)); + assert_eq!(codec.memory_priority(), 3 - plane); + assert!(codec.supports_row_selection()); + } + assert_eq!(PlaneKey::codec().unwrap().plane_tag(), None); + } + + #[rstest] + #[case::zero(0)] + #[case::sign_only(1)] + #[case::no_low_plane(2)] + #[case::too_wide(10)] + #[case::maximum(255)] + fn rejects_unsupported_layout(#[case] bits: u8) { + let error = RQLayout::try_new(bits).unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!( + error.to_string().contains("requires num_bits in 3..=9"), + "{error}" + ); + } +} diff --git a/rust/lance-index/src/vector/bq/plane_cache.rs b/rust/lance-index/src/vector/bq/plane_cache.rs new file mode 100644 index 00000000000..29cf795fe00 --- /dev/null +++ b/rust/lance-index/src/vector/bq/plane_cache.rs @@ -0,0 +1,274 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Row-addressable persistent ex planes. Sign codes retain their transposed IPC body. +use std::sync::Arc; + +use arrow_array::{Array, Float32Array, RecordBatch, UInt8Array, cast::AsArray, types::UInt8Type}; +use arrow_schema::{DataType, Field, Schema}; +use lance_arrow::FixedSizeListArrayExt; +use lance_core::cache::{CacheCodecImpl, CacheEntryReader, CacheEntryWriter, CacheRangeReader}; +use lance_core::{Error, Result}; + +use super::layered::{PlaneBatch, plane_columns}; +use super::storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN}; + +const HEADER_BYTES: usize = 13; +const COALESCE_GAP_BYTES: usize = 4096; + +struct Header { + plane: u8, + rows: usize, + width: usize, +} +impl Header { + fn parse(bytes: &[u8]) -> Result { + if bytes.len() < HEADER_BYTES || !matches!(bytes[0], 1 | 2) { + return Err(Error::invalid_input("invalid ex-plane cache header")); + } + let rows = usize::try_from(u64::from_le_bytes(bytes[1..9].try_into().unwrap())) + .map_err(|_| Error::invalid_input("plane row count overflow"))?; + let width = u32::from_le_bytes(bytes[9..13].try_into().unwrap()) as usize; + if width == 0 || width > i32::MAX as usize { + return Err(Error::invalid_input("invalid cached code width")); + } + Ok(Self { + plane: bytes[0], + rows, + width, + }) + } + fn stride(&self) -> usize { + self.width + 8 + } + fn batch(&self, codes: Vec, adds: Vec, scales: Vec) -> Result { + let names = plane_columns(self.plane); + let codes = arrow_array::FixedSizeListArray::try_new_from_values( + UInt8Array::from(codes), + self.width as i32, + )?; + let fields = vec![ + Field::new(names[0], codes.data_type().clone(), true), + Field::new(names[1], DataType::Float32, true), + Field::new(names[2], DataType::Float32, true), + ]; + Ok(PlaneBatch(RecordBatch::try_new( + Arc::new(Schema::new(fields)), + vec![ + Arc::new(codes), + Arc::new(Float32Array::from(adds)), + Arc::new(Float32Array::from(scales)), + ], + )?)) + } + fn decode(&self, bytes: &[u8]) -> Result { + let size = self + .rows + .checked_mul(self.stride()) + .ok_or_else(|| Error::invalid_input("plane size overflow"))?; + if bytes.len() != size { + return Err(Error::invalid_input("truncated ex-plane cache body")); + } + let mut codes = Vec::with_capacity(self.rows * self.width); + let mut adds = Vec::with_capacity(self.rows); + let mut scales = Vec::with_capacity(self.rows); + for row in bytes.chunks_exact(self.stride()) { + self.append(row, &mut codes, &mut adds, &mut scales); + } + self.batch(codes, adds, scales) + } + fn append(&self, row: &[u8], codes: &mut Vec, adds: &mut Vec, scales: &mut Vec) { + codes.extend_from_slice(&row[..self.width]); + adds.push(f32::from_le_bytes( + row[self.width..self.width + 4].try_into().unwrap(), + )); + scales.push(f32::from_le_bytes( + row[self.width + 4..self.width + 8].try_into().unwrap(), + )); + } +} + +impl CacheCodecImpl for PlaneBatch { + const TYPE_ID: &'static str = "lance.vector.rq.plane"; + const CURRENT_VERSION: u32 = 2; + const SUPPORTS_ROW_SELECTION: bool = true; + + fn serialize(&self, writer: &mut CacheEntryWriter<'_>) -> Result<()> { + let plane = if self + .0 + .column_by_name(RABIT_BLOCKED_EX_CODE_LO_COLUMN) + .is_some() + { + 2 + } else if self + .0 + .column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN) + .is_some() + { + 1 + } else { + 0 + }; + writer.write_u8(plane)?; + if plane == 0 { + return writer.write_ipc(&self.0); + } + let names = plane_columns(plane); + let codes = self.0[names[0]].as_fixed_size_list(); + let adds = self.0[names[1]].as_primitive::(); + let scales = self.0[names[2]].as_primitive::(); + if codes.null_count() != 0 || adds.null_count() != 0 || scales.null_count() != 0 { + return Err(Error::invalid_input("null cached ex-plane row")); + } + let width = codes.value_length() as usize; + let values = codes.values().as_primitive::().values(); + let writer = writer.raw_writer(); + writer.write_all(&(self.0.num_rows() as u64).to_le_bytes())?; + writer.write_all(&(width as u32).to_le_bytes())?; + for row in 0..self.0.num_rows() { + writer.write_all(&values[row * width..(row + 1) * width])?; + writer.write_all(&adds.value(row).to_le_bytes())?; + writer.write_all(&scales.value(row).to_le_bytes())?; + } + Ok(()) + } + fn deserialize(reader: &mut CacheEntryReader<'_>) -> Result { + if reader.version() == 1 { + return Ok(Self(reader.read_ipc()?)); + } + if reader.version() != Self::CURRENT_VERSION { + return Err(Error::invalid_input("unsupported plane cache version")); + } + let body = reader.body(); + if body.first() == Some(&0) { + reader.read_u8()?; + return Ok(Self(reader.read_ipc()?)); + } + let header = Header::parse(&body)?; + header.decode(&body[HEADER_BYTES..]) + } + fn deserialize_rows( + reader: &dyn CacheRangeReader, + offset: usize, + 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 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 { + return Err(Error::invalid_input("short plane range read")); + } + for &row in &rows[start..end] { + let local = (row as usize - first) * stride; + header.append( + &bytes[local..local + width + 8], + &mut codes, + &mut adds, + &mut scales, + ); + } + start = end; + } + header.batch(codes, adds, scales) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use bytes::Bytes; + use lance_arrow::RecordBatchExt; + use lance_core::cache::{CacheCodec, CacheDecode}; + use std::cell::Cell; + + #[test] + fn sparse_plane_cache_reads_match_full_decode() { + for plane in [1, 2] { + let header = Header { + plane, + rows: 4096, + width: 128, + }; + let original = header + .batch( + (0..4096 * 128).map(|v| (v % 251) as u8).collect(), + (0..4096).map(|v| v as f32).collect(), + vec![2.; 4096], + ) + .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 CacheDecode::Hit(full) = codec.deserialize(&encoded) else { + panic!("full decode failed") + }; + let full = full.downcast::().unwrap(); + let bytes_read = Cell::new(0); + let read = |range: std::ops::Range| -> Result { + bytes_read.set(bytes_read.get() + range.len()); + Ok(encoded.slice(range)) + }; + let rows = [0, 1, 2048, 4095]; + let CacheDecode::Hit(selected) = codec.deserialize_rows(&read, &rows) else { + panic!("range decode failed") + }; + assert_eq!( + selected.downcast::().unwrap().0, + full.0 + .take(&arrow_array::UInt32Array::from(rows.to_vec())) + .unwrap() + ); + assert!(bytes_read.get() < encoded.len() / 10); + assert!(matches!( + codec.deserialize_rows(&read, &[4096]), + CacheDecode::Miss(_) + )); + assert!(matches!( + codec.deserialize_rows(&read, &[1, 1]), + CacheDecode::Miss(_) + )); + } + } +} diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index d1140862087..1c50cedc029 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -48,8 +48,8 @@ use crate::vector::bq::dist_table_quant::{ DistTableDequant, quantize_dist_table_into, quantize_dist_table_u16_into, }; use crate::vector::bq::ex_dot::{ - EX_DOT_BLOCK_DIMS, ExDotFn, blocked_ex_code_bytes, ex_dot_kernel, pad_query_into, - padded_query_len, repack_sequential_row, sequential_matches_blocked, + EX_DOT_BLOCK_DIMS, ExDotFn, ExDotLayeredFn, blocked_ex_code_bytes, ex_dot_kernel, + pad_query_into, padded_query_len, repack_sequential_row, sequential_matches_blocked, }; use crate::vector::bq::prune::{PRUNE_LANES, ScaledLowerBoundTerms, scaled_prune_mask_kernel}; use crate::vector::bq::rotation::{apply_fast_rotation, apply_fast_rotation_in_place}; @@ -78,6 +78,11 @@ pub const RABIT_EX_CODE_COLUMN: &str = "__ex_codes"; /// older versions, which fail with a missing-column error instead of /// misinterpreting the bytes. pub const RABIT_BLOCKED_EX_CODE_COLUMN: &str = "__blocked_ex_codes"; +/// Low-bit plane of a layered index. [`RABIT_BLOCKED_EX_CODE_COLUMN`] holds +/// the high bits; their widths come from [`super::layered::RQLayout`]. High +/// precision omits this plane, while full precision combines both planes +/// without storing the full codes twice. +pub const RABIT_BLOCKED_EX_CODE_LO_COLUMN: &str = "__blocked_ex_codes_lo"; pub const SEGMENT_LENGTH: usize = 4; pub const SEGMENT_NUM_CODES: usize = 1 << SEGMENT_LENGTH; const RABIT_PRUNE_STATS_ENV: &str = "LANCE_RQ_PRUNE_STATS"; @@ -104,6 +109,12 @@ static RABIT_PRUNE_BYPASS_STATS: OnceLock = OnceLock::new static RABIT_PRUNE_STATS_ENABLED: OnceLock = OnceLock::new(); static RABIT_PRUNE_STATS_INTERVAL: OnceLock = OnceLock::new(); +// Enclose the LUT reduction, ex-dot, bias and final score reconstruction. +// The gamma_n form below also disables pruning if an enormous dimension makes +// the floating-point accumulation bound unrepresentable. +const LAYERED_ROUNDING_OPERATIONS_PER_DIM: f64 = 32.0; +const LAYERED_NORMAL_ROUNDING_OPERATIONS_PER_DIM: f64 = 8.0; + fn rabit_prune_stats_enabled() -> bool { *RABIT_PRUNE_STATS_ENABLED.get_or_init(|| match std::env::var(RABIT_PRUNE_STATS_ENV) { Ok(value) => str_is_truthy(value.trim()), @@ -263,6 +274,9 @@ pub struct RabitQuantizationMetadata { pub code_dim: u32, pub num_bits: u8, pub packed: bool, + /// Whether the ex codes use fixed high and low prefix planes. + #[serde(default)] + pub layered: bool, #[serde(default = "default_query_estimator_compat")] pub query_estimator: RabitQueryEstimator, } @@ -494,9 +508,13 @@ pub struct RabitQuantizationStorage { // sequential column with the repacked one, so rewrites emit the blocked // format). ex_codes: Option, + // Low-bit plane of a layered index (see RABIT_BLOCKED_EX_CODE_LO_COLUMN). + ex_codes_lo: Option, packed_ex_codes: Option, ex_add_factors: Option, ex_scale_factors: Option, + high_add_factors: Option, + high_scale_factors: Option, } impl DeepSizeOf for RabitQuantizationStorage { @@ -512,6 +530,217 @@ impl DeepSizeOf for RabitQuantizationStorage { } impl RabitQuantizationStorage { + fn from_batch_with_remapper( + batch: RecordBatch, + metadata: &RabitQuantizationMetadata, + distance_type: DistanceType, + fri: Option>, + prepack_ex: bool, + ) -> Result { + let distance_type = match (metadata.query_estimator, distance_type) { + (RabitQueryEstimator::RawQuery, DistanceType::Cosine) => DistanceType::L2, + _ => distance_type, + }; + validate_rq_num_bits(metadata.num_bits)?; + // The FastScan LUT is `4 * rotated_dim` bytes while the kernels index it + // as `BATCH_SIZE * rotated_dim.div_ceil(8)`, so the two agree only when + // the dimension is a multiple of 8. `RabitQuantizer::build` has rejected + // a non-multiple since #6024, but an index written before that still + // loads here, and the AVX-512, AVX2 and NEON kernels read the LUT + // through unchecked raw pointers. Only the scalar fallback panics. + // + // This has to go through `rotated_dim()`, not `metadata.code_dim`: + // `code_dim` was added by #6024 itself, so it deserializes to 0 for the + // very indices this rejects, and `rotated_dim()` recovers the real + // dimension from the rotation matrix that `parse_buffer` backfills. + let rotated_dim = metadata.rotated_dim(); + if !rotated_dim.is_multiple_of(8) { + return Err(Error::invalid_input(format!( + "RabitQ vector dimension must be divisible by 8, got {rotated_dim}. \ + Rebuild the index." + ))); + } + let row_ids = batch[ROW_ID].as_primitive::().clone(); + let codes = batch[RABIT_CODE_COLUMN].as_fixed_size_list().clone(); + // `rotated_dim() == 0` means the metadata never recorded a code + // dimension, so the width check below could only ever report that the + // column needs 0 bytes. Reject up front with the real cause. + if metadata.rotated_dim() == 0 { + return Err(Error::corrupt_file_named( + "rabitq metadata", + format!( + "no code dimension: code_dim is 0 and the rotation matrix is not loaded \ + (rotation_type={:?}, rotate_mat_position={:?})", + metadata.rotation_type, metadata.rotate_mat_position + ), + )); + } + let expected_code_bytes = metadata.binary_code_bytes(); + if codes.value_length() as usize != expected_code_bytes { + return Err(Error::invalid_input(format!( + "RabitQ code byte width mismatch: column {} has {} bytes, metadata rotated_dim={} requires {} bytes", + RABIT_CODE_COLUMN, + codes.value_length(), + metadata.rotated_dim(), + expected_code_bytes + ))); + } + let add_factors = batch[ADD_FACTORS_COLUMN] + .as_primitive::() + .clone(); + let scale_factors = batch[SCALE_FACTORS_COLUMN] + .as_primitive::() + .clone(); + // Full precision preserves native binary pruning regardless of how + // the ex codes are stored. Projected prefixes remove these factors and + // use their own estimator-difference bounds instead. + let error_factors = batch + .column_by_name(ERROR_FACTORS_COLUMN) + .map(|factors| factors.as_primitive::().clone()); + let ex_bits = rabit_ex_bits(metadata.num_bits)?; + let mut batch = batch; + let mut ex_codes = None; + let mut ex_add_factors = None; + let mut ex_scale_factors = None; + let mut ex_codes_lo = None; + for name in [ + super::layered::HIGH_BOUNDS_COLUMN, + super::layered::FULL_BOUNDS_COLUMN, + ] { + if let Some(column) = batch.column_by_name(name) { + let bounds = column.as_fixed_size_list_opt().ok_or_else(|| { + Error::invalid_input(format!("invalid RaBitQ bounds column {name}")) + })?; + if bounds.value_length() != 3 + || bounds.value_type() != DataType::Float32 + || bounds.null_count() != 0 + || bounds.values().null_count() != 0 + { + return Err(Error::invalid_input(format!( + "RaBitQ bounds column {name} requires three non-null float32 values per row" + ))); + } + } + } + if metadata.layered { + super::layered::RQLayout::try_new(metadata.num_bits)?; + for name in [ + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + super::layered::HIGH_ADD_FACTORS_COLUMN, + super::layered::HIGH_SCALE_FACTORS_COLUMN, + super::layered::HIGH_BOUNDS_COLUMN, + super::layered::FULL_BOUNDS_COLUMN, + ] { + if batch.column_by_name(name).is_none() { + return Err(Error::invalid_input(format!( + "RabitQ layered index missing column {name}" + ))); + } + } + } else if batch + .column_by_name(RABIT_BLOCKED_EX_CODE_LO_COLUMN) + .is_some() + { + return Err(Error::invalid_input( + "RabitQ low plane requires layered metadata", + )); + } + if ex_bits != 0 { + let (normalized_batch, codes, lo) = + load_ex_code_planes(batch, metadata.rotated_dim(), metadata.num_bits)?; + batch = normalized_batch; + ex_codes = Some(codes); + ex_codes_lo = lo; + ex_add_factors = Some( + batch + .column_by_name(EX_ADD_FACTORS_COLUMN) + .ok_or_else(|| { + Error::invalid_input(format!( + "RabitQ num_bits={} requires {} column", + metadata.num_bits, EX_ADD_FACTORS_COLUMN + )) + })? + .as_primitive::() + .clone(), + ); + ex_scale_factors = Some( + batch + .column_by_name(EX_SCALE_FACTORS_COLUMN) + .ok_or_else(|| { + Error::invalid_input(format!( + "RabitQ num_bits={} requires {} column", + metadata.num_bits, EX_SCALE_FACTORS_COLUMN + )) + })? + .as_primitive::() + .clone(), + ); + } else if metadata.query_estimator == RabitQueryEstimator::RawQuery { + if batch.column_by_name(EX_ADD_FACTORS_COLUMN).is_some() + || batch.column_by_name(EX_SCALE_FACTORS_COLUMN).is_some() + || batch.column_by_name(RABIT_EX_CODE_COLUMN).is_some() + || batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_some() + { + return Err(Error::invalid_input( + "RabitQ num_bits=1 raw-query indexes must not contain ex-code columns" + .to_string(), + )); + } + } else if batch.column_by_name(RABIT_EX_CODE_COLUMN).is_some() + || batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_some() + { + return Err(Error::invalid_input(format!( + "RabitQ num_bits={} does not support ex-code columns", + metadata.num_bits + ))); + } + + let (batch, codes) = if !metadata.packed { + let codes = pack_codes(&codes); + let batch = batch.replace_column_by_name(RABIT_CODE_COLUMN, Arc::new(codes))?; + let codes = batch[RABIT_CODE_COLUMN].as_fixed_size_list().clone(); + (batch, codes) + } else { + (batch, codes) + }; + + let mut metadata = metadata.clone(); + metadata.packed = true; + // The FastScan transpose only understands single-plane widths. + let packed_ex_codes = if ex_codes_lo.is_some() || !prepack_ex { + None + } else { + maybe_pack_ex_codes(ex_codes.as_ref(), ex_bits, error_factors.as_ref()) + }; + + let storage = Self { + high_add_factors: batch + .column_by_name(super::layered::HIGH_ADD_FACTORS_COLUMN) + .map(|a| a.as_primitive::().clone()), + high_scale_factors: batch + .column_by_name(super::layered::HIGH_SCALE_FACTORS_COLUMN) + .map(|a| a.as_primitive::().clone()), + metadata, + batch, + distance_type, + row_ids, + codes, + add_factors, + scale_factors, + error_factors, + ex_codes, + ex_codes_lo, + packed_ex_codes, + ex_add_factors, + ex_scale_factors, + }; + + match build_frag_reuse_mapping(fri.as_deref(), &storage.row_ids) { + Some(mapping) => storage.remap(&RowAddrRemap::direct(mapping)), + None => Ok(storage), + } + } + fn code_dim(&self) -> usize { self.metadata.code_dim() } @@ -534,8 +763,9 @@ impl RabitQuantizationStorage { rotated_centroid: Option<&[f32]>, ) -> RabitQueryFactors { let gated = self.metadata.num_bits > 1 && self.error_factors.is_some(); - match self.distance_type { + let mut factors = match self.distance_type { DistanceType::L2 | DistanceType::Cosine => RabitQueryFactors { + centered_query: None, add_scale: 1.0, add_offset: 0.0, add: if self.distance_type == DistanceType::Cosine { @@ -560,6 +790,7 @@ impl RabitQuantizationStorage { (ip, if alpha.is_finite() { alpha } else { 0.0 }) }); RabitQueryFactors { + centered_query: None, add_scale: alpha, // Center the stored factor before scaling: splitting the // constant as alpha*A + (1-alpha-ip) loses it for large alpha. @@ -609,7 +840,29 @@ impl RabitQuantizationStorage { "RabitQ does not support distance type: {}", self.distance_type ), + }; + if self.distance_type != DistanceType::Dot + && self.metadata.num_bits > 1 + && (self.metadata.layered + || self + .batch + .column_by_name(super::layered::FULL_BOUNDS_COLUMN) + .is_some()) + && let Some(centroid) = rotated_centroid + && centroid.len() == rotated_query.len() + { + let alpha = f64::from(factors.add_scale); + let (norm_square, centroid_l1) = + rotated_query + .iter() + .zip(centroid) + .fold((0.0, 0.0), |(norm, l1), (&q, &c)| { + let residual = f64::from(q) - alpha * f64::from(c); + (norm + residual * residual, l1 + f64::from(c).abs()) + }); + factors.centered_query = Some((norm_square.sqrt(), alpha.abs() * centroid_l1)); } + factors } fn distance_calculator_from_parts<'a>( @@ -623,6 +876,7 @@ impl RabitQuantizationStorage { sum_q, query_factors, approx_mode, + rq_precision, } = parts; let ex_code_len = self .ex_codes @@ -633,31 +887,62 @@ impl RabitQuantizationStorage { .ex_codes .as_ref() .map(|codes| codes.values().as_primitive::().values().as_ref()); + let ex_codes_lo = self + .ex_codes_lo + .as_ref() + .map(|codes| codes.values().as_primitive::().values().as_ref()); let packed_ex_codes = self .packed_ex_codes .as_ref() .map(|codes| codes.values().as_primitive::().values().as_ref()); + let rq_precision = if self.metadata.layered { + rq_precision + } 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 => ( + self.metadata.num_bits, + ex_codes, + ex_codes_lo, + packed_ex_codes, + self.ex_add_factors.as_ref(), + self.ex_scale_factors.as_ref(), + ), + super::layered::RQPrecision::Sign => (1, None, None, None, None, None), + super::layered::RQPrecision::High => ( + match self.metadata.num_bits { + 3 => 2, + 4 | 5 => 3, + 6..=9 => 5, + _ => unreachable!("validated layered width"), + }, + ex_codes, + None, + None, + self.high_add_factors.as_ref(), + self.high_scale_factors.as_ref(), + ), + }; let mut calculator = RabitDistCalculator::new( dim, - self.metadata.num_bits, + num_bits, self.metadata.query_estimator, dist_table, ex_query, sum_q, self.codes.values().as_primitive::().values(), ex_codes, + ex_codes_lo, ex_code_len, self.add_factors.values(), self.scale_factors.values(), self.error_factors .as_ref() .map(|factors| factors.values().as_ref()), - self.ex_add_factors - .as_ref() - .map(|factors| factors.values().as_ref()), - self.ex_scale_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()), packed_ex_codes, query_factors.add, query_factors.error, @@ -665,6 +950,22 @@ impl RabitQuantizationStorage { ); calculator.add_factor_scale = query_factors.add_scale; calculator.add_factor_offset = query_factors.add_offset; + if num_bits > 1 + && approx_mode != ApproxMode::Fast + && (rq_precision == super::layered::RQPrecision::High || self.error_factors.is_none()) + { + let column = match rq_precision { + super::layered::RQPrecision::High => super::layered::HIGH_BOUNDS_COLUMN, + _ => super::layered::FULL_BOUNDS_COLUMN, + }; + if let Some(bounds) = self.batch.column_by_name(column) { + calculator.prepare_layered_bounds( + bounds.as_fixed_size_list(), + query_factors.centered_query, + self.distance_type, + ); + } + } calculator } @@ -836,6 +1137,8 @@ struct RabitQueryFactors { add_offset: f32, add: f32, error: f32, + // Norm of q - a*c, and ||a*c||_1 for the centroid-correction rounding margin. + centered_query: Option<(f64, f64)>, } struct RabitDistCalculatorParts<'a> { @@ -845,6 +1148,7 @@ struct RabitDistCalculatorParts<'a> { sum_q: f32, query_factors: RabitQueryFactors, approx_mode: ApproxMode, + rq_precision: super::layered::RQPrecision, } /// Loop-invariant inputs of the raw-query multi-bit top-k scans: the row @@ -877,8 +1181,11 @@ pub struct RabitDistCalculator<'a> { query_estimator: RabitQueryEstimator, // n * d / 8 binary-code bytes codes: &'a [u8], - // per-row ex codes in the blocked kernel layout + // per-row ex codes in the blocked kernel layout (the high-bit plane + // for a layered index) ex_codes: Option<&'a [u8]>, + // Low-bit plane; its row width is the full width minus `ex_code_len`. + ex_codes_lo: Option<&'a [u8]>, // bytes per ex-code row; legacy rows for layout-compatible widths may be // shorter than the blocked size, which the kernels treat as zero padding ex_code_len: usize, @@ -890,9 +1197,11 @@ pub struct RabitDistCalculator<'a> { // the source for the FastScan ex LUT on the legacy bypass path ex_query: Cow<'a, [f32]>, ex_dot: Option, + ex_dot_layered: Option, add_factors: &'a [f32], scale_factors: &'a [f32], error_factors: Option<&'a [f32]>, + layered_errors: Option>, ex_add_factors: Option<&'a [f32]>, ex_scale_factors: Option<&'a [f32]>, packed_ex_codes: Option<&'a [u8]>, @@ -907,6 +1216,123 @@ pub struct RabitDistCalculator<'a> { } impl<'a> RabitDistCalculator<'a> { + fn prepare_layered_bounds( + &mut self, + bounds: &FixedSizeListArray, + centered_query: Option<(f64, f64)>, + metric: DistanceType, + ) { + let values = bounds.values().as_primitive::().values(); + // Retain the wider enclosure for signed Dot scores: tightening that + // statistical bound changed a top-k boundary under a different scan order. + let use_projection = self.approx_mode == ApproxMode::Normal && metric != DistanceType::Dot; + let centered_query = centered_query.filter(|_| use_projection); + let norm = centered_query.map_or_else( + || { + self.ex_query + .iter() + .map(|&q| f64::from(q).powi(2)) + .sum::() + .sqrt() + }, + |(norm, _)| norm, + ); + let l1 = self + .ex_query + .iter() + .map(|&q| f64::from(q).abs()) + .sum::(); + let (min, max) = self + .dist_table + .iter() + .fold((f32::INFINITY, f32::NEG_INFINITY), |(min, max), &v| { + (min.min(v), max.max(v)) + }); + // The sign LUT may be quantized even when the ex-dot is exact. Include + // its worst-case error and dot-product/reconstruction rounding; both + // estimators use the same LUT value with different coefficients. + let rounding_operations = if use_projection { + LAYERED_NORMAL_ROUNDING_OPERATIONS_PER_DIM + } else { + LAYERED_ROUNDING_OPERATIONS_PER_DIM + }; + let rounding_epsilon = rounding_operations * f64::from(f32::EPSILON); + let accumulated_error = rounding_epsilon * (self.dim + 8) as f64; + let rounding_gamma = if accumulated_error < 1.0 { + accumulated_error / (1.0 - accumulated_error) + } else { + f64::INFINITY + }; + let lut_error = (self.dist_table.len() / SEGMENT_NUM_CODES) as f64 + * (f64::from(max) - f64::from(min)) + / f64::from(u8::MAX); + let rounding_error = + rounding_gamma * (l1 + centered_query.map_or(0.0, |(_, centroid_l1)| centroid_l1)); + let code_scale = f64::from(1u32 << (self.num_bits - 1)); + // Normal uses the same angular confidence policy as native RaBitQ, + // applied to this level's actual estimator difference. Accurate retains + // the Cauchy-Schwarz bound; it trades more ex-dot work for no statistical + // candidate loss. Quantization and arithmetic margins are never shrunk. + let angular_scale = if self.approx_mode == ApproxMode::Accurate || self.dim <= 1 { + 1.0 + } else { + (f64::from(super::transform::RABIT_ERROR_EPSILON) / ((self.dim - 1) as f64).sqrt()) + .min(1.0) + }; + let errors = values + .as_chunks::<3>() + .0 + .iter() + .enumerate() + .map(|(row, b)| { + let add_scale = f64::from(self.add_factor_scale).abs(); + let rounding = rounding_epsilon + * (add_scale + * (f64::from(self.add_factors[row]).abs() + + f64::from(self.add_factor_offset).abs() + + f64::from(b[1])) + + f64::from(self.query_factor).abs()); + // Both estimators reconstruct from the same sign LUT value. + // Its error is multiplied by the coefficient difference, not + // the sum of absolute coefficients used for reduction rounding. + let lut_coefficient = self.ex_scale_factors.map_or(f64::from(b[2]), |scales| { + (code_scale * f64::from(scales[row]) - f64::from(self.scale_factors[row])).abs() + }); + // add_level-add_sign cancels the centroid dot of the weight + // difference. Normal uses q-a*c and encloses the encoder's + // floating correction in the centroid rounding margin above. + let add_difference = if centered_query.is_some() { + 0.0 + } else { + f64::from(b[1]) * add_scale + }; + let numeric_margin = if metric == DistanceType::Dot { + f64::from(b[2]) * (lut_error + rounding_error) + } else { + lut_coefficient * lut_error + f64::from(b[2]) * rounding_error + }; + let margin = f64::from(b[0]) * norm * angular_scale + + add_difference + + numeric_margin + + rounding; + // Invalid optional hints fail closed: score the row in full. + // Fold this check into preparation instead of rescanning both + // levels' columns on every reconstructed partition. + if margin.is_finite() && b.iter().all(|v| *v >= 0.0) { + (margin as f32).next_up() + } else { + f32::INFINITY + } + }) + .collect(); + self.layered_errors = Some(errors); + self.query_error = 1.0; + } + + fn pruning_errors(&self) -> Option<&[f32]> { + self.layered_errors.as_deref().or(self.error_factors) + } + #[allow(clippy::too_many_arguments)] pub fn new( dim: usize, @@ -917,6 +1343,7 @@ impl<'a> RabitDistCalculator<'a> { sum_q: f32, codes: &'a [u8], ex_codes: Option<&'a [u8]>, + ex_codes_lo: Option<&'a [u8]>, ex_code_len: usize, add_factors: &'a [f32], scale_factors: &'a [f32], @@ -928,21 +1355,27 @@ impl<'a> RabitDistCalculator<'a> { query_error: f32, approx_mode: ApproxMode, ) -> Self { - let ex_dot = (num_bits > 1).then(|| ex_dot_kernel(num_bits - 1)); + let ex_dot = (num_bits > 1 && ex_codes_lo.is_none()).then(|| ex_dot_kernel(num_bits - 1)); + let ex_dot_layered = ex_codes_lo + .is_some() + .then(|| super::ex_dot::ex_dot_prefix_kernel(num_bits)); Self { dim, num_bits, query_estimator, codes, ex_codes, + ex_codes_lo, ex_code_len, dist_table, ex_query, ex_dot, + ex_dot_layered, add_factors, scale_factors, error_factors, ex_add_factors, + layered_errors: None, ex_scale_factors, packed_ex_codes, query_factor, @@ -958,13 +1391,79 @@ impl<'a> RabitDistCalculator<'a> { /// `sum_d query[d] * ex_code[d]` for the candidate's packed ex codes. #[inline] fn ex_code_dot(&self, ex_codes: &[u8], id: usize) -> f32 { + let row = id * self.ex_code_len..(id + 1) * self.ex_code_len; + if let (Some(ex_codes_lo), Some(ex_dot_layered)) = (self.ex_codes_lo, self.ex_dot_layered) { + // Both planes pad to the same 64-dimensional boundary, so their + // row widths add up to the native full-code row width. + let lo_len = blocked_ex_code_bytes(self.dim, self.num_bits - 1) - self.ex_code_len; + return ex_dot_layered( + self.ex_query.as_ref(), + &ex_codes[row], + &ex_codes_lo[id * lo_len..(id + 1) * lo_len], + ); + } let ex_dot = self .ex_dot .expect("raw-query multi-bit RQ requires an ex-dot kernel"); - ex_dot( - self.ex_query.as_ref(), - &ex_codes[id * self.ex_code_len..(id + 1) * self.ex_code_len], - ) + ex_dot(self.ex_query.as_ref(), &ex_codes[row]) + } + + /// Binary inner products in the original partition's SIMD/tail layout. + /// Carry these through a candidate cut so reranking does not change LUT rounding + /// when a SIMD row becomes a scalar remainder in a smaller candidate batch. + pub fn binary_inner_products(&self) -> Vec { + let code_len = rabit_binary_code_bytes(self.dim); + let n = self.codes.len() / code_len; + let mut dists = Vec::new(); + let mut u16_scratch = Vec::new(); + let mut u8_scratch = Vec::new(); + let mut u32_scratch = Vec::new(); + self.binary_distances_with_scratch( + n, + code_len, + &mut dists, + &mut u16_scratch, + &mut u8_scratch, + &mut u32_scratch, + ); + dists + } + + /// Score a valid row with its original partition's binary inner product. + /// `id` is a local row offset and `dist` must come from the same query rotation. + pub fn distance_with_binary_inner_product(&self, id: u32, dist: f32) -> f32 { + let id = id as usize; + match self.query_estimator { + RabitQueryEstimator::ResidualQuery => { + // distance between quantized residual vector and residual query vector + let dist_vq_qr = (2.0 * dist - self.sum_q) / self.sqrt_d; + dist_vq_qr * self.scale_factors[id] + self.add_factors[id] + self.query_factor + } + RabitQueryEstimator::RawQuery => { + let ex_bits = self.num_bits - 1; + if ex_bits == 0 || self.approx_mode == ApproxMode::Fast { + return self.raw_query_binary_distance(id, dist); + } + + let ex_codes = self + .ex_codes + .expect("raw-query multi-bit RQ requires ex codes"); + let ex_add_factors = self + .ex_add_factors + .expect("raw-query multi-bit RQ requires ex add factors"); + let ex_scale_factors = self + .ex_scale_factors + .expect("raw-query multi-bit RQ requires ex scale factors"); + self.raw_query_multi_bit_exact_distance( + id, + dist, + ex_bits, + ex_codes, + ex_add_factors, + ex_scale_factors, + ) + } + } } /// Fill `dists[0..n]` with exact per-row binary distances computed @@ -1290,8 +1789,22 @@ impl<'a> RabitDistCalculator<'a> { } #[inline] - fn raw_query_lower_bound(&self, id: usize, binary_ip: f32) -> Option { - let error_factors = self.error_factors?; + /// Lower bound for this precision's score using the original partition's + /// binary LUT value. Normal mode uses RaBitQ's statistical angular bound; + /// Accurate layered mode uses the conservative estimator-difference norm. + /// Returns `None` when this calculator cannot prune. + pub fn lower_bound_with_binary_inner_product(&self, id: u32, binary_ip: f32) -> Option { + if self + .raw_query_lower_bound_gating_disabled_reason() + .is_some() + { + return None; + } + self.raw_query_lower_bound(id as usize, binary_ip) + } + + pub(crate) fn raw_query_lower_bound(&self, id: usize, binary_ip: f32) -> Option { + let error_factors = self.pruning_errors()?; Some(self.raw_query_binary_distance(id, binary_ip) - error_factors[id] * self.query_error) } @@ -1496,7 +2009,7 @@ impl<'a> RabitDistCalculator<'a> { let scale_factors = &self.scale_factors[..ctx.n]; let add_factors = &self.add_factors[..ctx.n]; let error_factors = &self - .error_factors + .pruning_errors() .expect("raw-query lower-bound gating requires error factors")[..ctx.n]; // Same expression as `raw_query_lower_bound` with `error_factors` // already resolved; the masks below match it bit for bit. @@ -1584,7 +2097,7 @@ impl<'a> RabitDistCalculator<'a> { Some("residual_query_estimator") } else if self.num_bits <= 1 { Some("num_bits_le_one") - } else if self.error_factors.is_none() { + } else if self.pruning_errors().is_none() { Some("missing_error_factors") } else { None @@ -1793,6 +2306,57 @@ fn blocked_ex_codes_from_sequential( )?) } +/// Load the ex-code planes of an index batch. A layered batch carries +/// [`RABIT_BLOCKED_EX_CODE_LO_COLUMN`]; both blocked plane widths are derived +/// from its total bit width. Otherwise this is [`load_blocked_ex_codes`] with +/// no low plane. +pub(crate) fn load_ex_code_planes( + batch: RecordBatch, + rotated_dim: usize, + num_bits: u8, +) -> Result<(RecordBatch, FixedSizeListArray, Option)> { + let Some(lo_column) = batch.column_by_name(RABIT_BLOCKED_EX_CODE_LO_COLUMN) else { + let (batch, codes) = load_blocked_ex_codes(batch, rotated_dim, num_bits)?; + return Ok((batch, codes, None)); + }; + let layout = super::layered::RQLayout::try_new(num_bits)?; + let hi_column = batch + .column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN) + .ok_or_else(|| { + Error::invalid_input(format!( + "RabitQ layered index requires {} alongside {}", + RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN + )) + })?; + let hi = hi_column + .as_fixed_size_list_opt() + .ok_or_else(|| Error::invalid_input("RabitQ high plane must be a fixed-size list"))? + .clone(); + let lo = lo_column + .as_fixed_size_list_opt() + .ok_or_else(|| Error::invalid_input("RabitQ low plane must be a fixed-size list"))? + .clone(); + if hi.value_type() != DataType::UInt8 || lo.value_type() != DataType::UInt8 { + return Err(Error::invalid_input("RabitQ plane elements must be uint8")); + } + for (name, plane, bits) in [ + (RABIT_BLOCKED_EX_CODE_COLUMN, &hi, layout.high_bits), + (RABIT_BLOCKED_EX_CODE_LO_COLUMN, &lo, layout.low_bits), + ] { + let expected_bytes = blocked_ex_code_bytes(rotated_dim, bits); + if plane.value_length() as usize != expected_bytes { + return Err(Error::invalid_input(format!( + "RabitQ layered plane byte width mismatch: column {} has {} bytes, rotated_dim={} requires {} bytes for its plane", + name, + plane.value_length(), + rotated_dim, + expected_bytes + ))); + } + } + Ok((batch, hi, Some(lo))) +} + /// Load the ex-code column of an index batch into the blocked kernel layout, /// accepting both the blocked format and the legacy sequential format. Legacy /// batches are normalized in place (the sequential column is replaced by the @@ -1855,37 +2419,7 @@ impl DistCalculator for RabitDistCalculator<'_> { let dist = compute_single_rq_distance(self.codes, id, num_vectors, code_len, &self.dist_table); - match self.query_estimator { - RabitQueryEstimator::ResidualQuery => { - // distance between quantized residual vector and residual query vector - let dist_vq_qr = (2.0 * dist - self.sum_q) / self.sqrt_d; - dist_vq_qr * self.scale_factors[id] + self.add_factors[id] + self.query_factor - } - RabitQueryEstimator::RawQuery => { - let ex_bits = self.num_bits - 1; - if ex_bits == 0 || self.approx_mode == ApproxMode::Fast { - return self.raw_query_binary_distance(id, dist); - } - - let ex_codes = self - .ex_codes - .expect("raw-query multi-bit RQ requires ex codes"); - let ex_add_factors = self - .ex_add_factors - .expect("raw-query multi-bit RQ requires ex add factors"); - let ex_scale_factors = self - .ex_scale_factors - .expect("raw-query multi-bit RQ requires ex scale factors"); - self.raw_query_multi_bit_exact_distance( - id, - dist, - ex_bits, - ex_codes, - ex_add_factors, - ex_scale_factors, - ) - } - } + self.distance_with_binary_inner_product(id as u32, dist) } #[inline(always)] @@ -2157,6 +2691,7 @@ impl VectorStore for RabitQuantizationStorage { let dist_table = build_dist_table_direct::(&rotated_qr); let query_factors = match self.metadata.query_estimator { RabitQueryEstimator::ResidualQuery => RabitQueryFactors { + centered_query: None, add_scale: 1.0, add_offset: 0.0, add: self.residual_query_factor(dist_q_c), @@ -2182,6 +2717,7 @@ impl VectorStore for RabitQuantizationStorage { sum_q, query_factors, approx_mode: ApproxMode::Normal, + rq_precision: super::layered::RQPrecision::Full, }) } @@ -2205,7 +2741,7 @@ impl VectorStore for RabitQuantizationStorage { ) = (self.metadata.query_estimator, residual) { debug_assert_eq!(raw_query.code_dim, code_dim); - debug_assert_eq!(raw_query.ex_bits, self.metadata.num_bits - 1); + // 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 { @@ -2218,6 +2754,7 @@ impl VectorStore for RabitQuantizationStorage { sum_q: raw_query.sum_q, query_factors, approx_mode: options.approx_mode, + rq_precision: options.rq_precision, }); } @@ -2251,6 +2788,7 @@ impl VectorStore for RabitQuantizationStorage { } 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), @@ -2284,6 +2822,7 @@ impl VectorStore for RabitQuantizationStorage { sum_q, query_factors, approx_mode: options.approx_mode, + rq_precision: options.rq_precision, }) } @@ -2386,6 +2925,47 @@ pub fn pack_codes(codes: &FixedSizeListArray) -> FixedSizeListArray { FixedSizeListArray::try_new_from_values(UInt8Array::from(blocks), code_len as i32).unwrap() } +/// Gather only selected physical rows from partition-transposed sign codes. +/// The result keeps the packed layout expected by the selected-row calculator. +pub fn take_packed_codes(codes: &FixedSizeListArray, rows: &[u32]) -> Result { + let width = codes.value_length() as usize; + let total = codes.len(); + let packed_rows = total / BATCH_SIZE * BATCH_SIZE; + let remainder = total - packed_rows; + let source = codes.values().as_primitive::().values(); + let mut selected = Vec::with_capacity(rows.len() * width); + for &row in rows { + let row = row as usize; + if row >= total { + return Err(Error::invalid_input("candidate sign offset out of bounds")); + } + if row < packed_rows { + let lane = row % BATCH_SIZE; + let index = PERM0_INVERSE[lane % 16]; + let base = row / BATCH_SIZE * BATCH_SIZE * width; + for column in 0..width { + let offset = base + column * BATCH_SIZE + index; + let low = source[offset]; + let high = source[offset + 16]; + selected.push(if lane < 16 { + (low & 15) | (high << 4) + } else { + (low >> 4) | (high & 240) + }); + } + } else { + let offset = packed_rows * width; + for column in 0..width { + selected.push(source[offset + column * remainder + row - packed_rows]); + } + } + } + Ok(pack_codes(&FixedSizeListArray::try_new_from_values( + UInt8Array::from(selected), + width as i32, + )?)) +} + // Inverse of pack_codes pub fn unpack_codes(codes: &FixedSizeListArray) -> FixedSizeListArray { let code_len = codes.value_length() as usize; @@ -2474,6 +3054,73 @@ fn build_frag_reuse_mapping( #[async_trait] impl QuantizerStorage for RabitQuantizationStorage { + fn try_from_batch_at_precision( + batch: RecordBatch, + metadata: &Self::Metadata, + distance_type: DistanceType, + remapper: Option>, + precision: super::layered::RQPrecision, + ) -> Result { + use super::layered::{ + HIGH_ADD_FACTORS_COLUMN, HIGH_SCALE_FACTORS_COLUMN, RQLayout, RQPrecision, + }; + if precision == RQPrecision::Full { + return Self::try_from_batch_with_remapper(batch, metadata, distance_type, remapper); + } + if !metadata.layered { + return Err(Error::invalid_input( + "rq_precision requires a layered IVF_RQ index", + )); + } + let mut metadata = metadata.clone(); + let layout = RQLayout::try_new(metadata.num_bits)?; + metadata.layered = false; + metadata.num_bits = if precision == RQPrecision::Sign { + 1 + } else { + 1 + layout.high_bits + }; + let mut fields = Vec::new(); + let mut columns = Vec::new(); + for (field, column) in batch.schema().fields().iter().zip(batch.columns()) { + let name = field.name().as_str(); + if [ + ERROR_FACTORS_COLUMN, + super::layered::FULL_BOUNDS_COLUMN, + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + ] + .contains(&name) + { + continue; + } + if precision == RQPrecision::Sign + && [ + RABIT_BLOCKED_EX_CODE_COLUMN, + HIGH_ADD_FACTORS_COLUMN, + HIGH_SCALE_FACTORS_COLUMN, + ] + .contains(&name) + { + continue; + } + let renamed = match name { + super::layered::HIGH_BOUNDS_COLUMN => super::layered::FULL_BOUNDS_COLUMN, + HIGH_ADD_FACTORS_COLUMN => EX_ADD_FACTORS_COLUMN, + HIGH_SCALE_FACTORS_COLUMN => EX_SCALE_FACTORS_COLUMN, + _ => name, + }; + fields.push(field.as_ref().clone().with_name(renamed)); + columns.push(column.clone()); + } + let batch = RecordBatch::try_new(Arc::new(arrow_schema::Schema::new(fields)), columns)?; + // 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) + } + type Metadata = RabitQuantizationMetadata; fn try_from_batch( @@ -2492,150 +3139,7 @@ impl QuantizerStorage for RabitQuantizationStorage { distance_type: DistanceType, fri: Option>, ) -> Result { - let distance_type = match (metadata.query_estimator, distance_type) { - (RabitQueryEstimator::RawQuery, DistanceType::Cosine) => DistanceType::L2, - _ => distance_type, - }; - validate_rq_num_bits(metadata.num_bits)?; - // The FastScan LUT is `4 * rotated_dim` bytes while the kernels index it - // as `BATCH_SIZE * rotated_dim.div_ceil(8)`, so the two agree only when - // the dimension is a multiple of 8. `RabitQuantizer::build` has rejected - // a non-multiple since #6024, but an index written before that still - // loads here, and the AVX-512, AVX2 and NEON kernels read the LUT - // through unchecked raw pointers. Only the scalar fallback panics. - // - // This has to go through `rotated_dim()`, not `metadata.code_dim`: - // `code_dim` was added by #6024 itself, so it deserializes to 0 for the - // very indices this rejects, and `rotated_dim()` recovers the real - // dimension from the rotation matrix that `parse_buffer` backfills. - let rotated_dim = metadata.rotated_dim(); - if rotated_dim % 8 != 0 { - return Err(Error::invalid_input(format!( - "RabitQ vector dimension must be divisible by 8, got {rotated_dim}. \ - Rebuild the index." - ))); - } - let row_ids = batch[ROW_ID].as_primitive::().clone(); - let codes = batch[RABIT_CODE_COLUMN].as_fixed_size_list().clone(); - // `rotated_dim() == 0` means the metadata never recorded a code - // dimension, so the width check below could only ever report that the - // column needs 0 bytes. Reject up front with the real cause. - if metadata.rotated_dim() == 0 { - return Err(Error::corrupt_file_named( - "rabitq metadata", - format!( - "no code dimension: code_dim is 0 and the rotation matrix is not loaded \ - (rotation_type={:?}, rotate_mat_position={:?})", - metadata.rotation_type, metadata.rotate_mat_position - ), - )); - } - let expected_code_bytes = metadata.binary_code_bytes(); - if codes.value_length() as usize != expected_code_bytes { - return Err(Error::invalid_input(format!( - "RabitQ code byte width mismatch: column {} has {} bytes, metadata rotated_dim={} requires {} bytes", - RABIT_CODE_COLUMN, - codes.value_length(), - metadata.rotated_dim(), - expected_code_bytes - ))); - } - let add_factors = batch[ADD_FACTORS_COLUMN] - .as_primitive::() - .clone(); - let scale_factors = batch[SCALE_FACTORS_COLUMN] - .as_primitive::() - .clone(); - let error_factors = batch - .column_by_name(ERROR_FACTORS_COLUMN) - .map(|factors| factors.as_primitive::().clone()); - let ex_bits = rabit_ex_bits(metadata.num_bits)?; - let mut batch = batch; - let mut ex_codes = None; - let mut ex_add_factors = None; - let mut ex_scale_factors = None; - if ex_bits != 0 { - let (normalized_batch, codes) = - load_blocked_ex_codes(batch, metadata.rotated_dim(), metadata.num_bits)?; - batch = normalized_batch; - ex_codes = Some(codes); - ex_add_factors = Some( - batch - .column_by_name(EX_ADD_FACTORS_COLUMN) - .ok_or_else(|| { - Error::invalid_input(format!( - "RabitQ num_bits={} requires {} column", - metadata.num_bits, EX_ADD_FACTORS_COLUMN - )) - })? - .as_primitive::() - .clone(), - ); - ex_scale_factors = Some( - batch - .column_by_name(EX_SCALE_FACTORS_COLUMN) - .ok_or_else(|| { - Error::invalid_input(format!( - "RabitQ num_bits={} requires {} column", - metadata.num_bits, EX_SCALE_FACTORS_COLUMN - )) - })? - .as_primitive::() - .clone(), - ); - } else if metadata.query_estimator == RabitQueryEstimator::RawQuery { - if batch.column_by_name(EX_ADD_FACTORS_COLUMN).is_some() - || batch.column_by_name(EX_SCALE_FACTORS_COLUMN).is_some() - || batch.column_by_name(RABIT_EX_CODE_COLUMN).is_some() - || batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_some() - { - return Err(Error::invalid_input( - "RabitQ num_bits=1 raw-query indexes must not contain ex-code columns" - .to_string(), - )); - } - } else if batch.column_by_name(RABIT_EX_CODE_COLUMN).is_some() - || batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_some() - { - return Err(Error::invalid_input(format!( - "RabitQ num_bits={} does not support ex-code columns", - metadata.num_bits - ))); - } - - let (batch, codes) = if !metadata.packed { - let codes = pack_codes(&codes); - let batch = batch.replace_column_by_name(RABIT_CODE_COLUMN, Arc::new(codes))?; - let codes = batch[RABIT_CODE_COLUMN].as_fixed_size_list().clone(); - (batch, codes) - } else { - (batch, codes) - }; - - let mut metadata = metadata.clone(); - metadata.packed = true; - let packed_ex_codes = - maybe_pack_ex_codes(ex_codes.as_ref(), ex_bits, error_factors.as_ref()); - - let storage = Self { - metadata, - batch, - distance_type, - row_ids, - codes, - add_factors, - scale_factors, - error_factors, - ex_codes, - packed_ex_codes, - ex_add_factors, - ex_scale_factors, - }; - - match build_frag_reuse_mapping(fri.as_deref(), &storage.row_ids) { - Some(mapping) => storage.remap(&RowAddrRemap::direct(mapping)), - None => Ok(storage), - } + Self::from_batch_with_remapper(batch, metadata, distance_type, fri, true) } fn metadata(&self) -> &Self::Metadata { @@ -2700,21 +3204,26 @@ impl QuantizerStorage for RabitQuantizationStorage { let scale_factors = batch[SCALE_FACTORS_COLUMN] .as_primitive::() .clone(); - let error_factors = batch - .column_by_name(ERROR_FACTORS_COLUMN) - .map(|factors| factors.as_primitive::().clone()); + let error_factors = self.error_factors.as_ref().and_then(|_| { + batch + .column_by_name(ERROR_FACTORS_COLUMN) + .map(|factors| factors.as_primitive::().clone()) + }); let ex_bits = rabit_ex_bits(self.metadata.num_bits)?; - let (batch, ex_codes) = if ex_bits == 0 { - (batch, None) + let (batch, ex_codes, ex_codes_lo) = if ex_bits == 0 { + (batch, None, None) } else { // `self.batch` is already normalized at load, so this is a // zero-copy column lookup. - let (batch, codes) = - load_blocked_ex_codes(batch, self.metadata.rotated_dim(), self.metadata.num_bits)?; - (batch, Some(codes)) + let (batch, codes, lo) = + load_ex_code_planes(batch, self.metadata.rotated_dim(), self.metadata.num_bits)?; + (batch, Some(codes), lo) + }; + let packed_ex_codes = if ex_codes_lo.is_some() { + None + } else { + maybe_pack_ex_codes(ex_codes.as_ref(), ex_bits, error_factors.as_ref()) }; - let packed_ex_codes = - maybe_pack_ex_codes(ex_codes.as_ref(), ex_bits, error_factors.as_ref()); let ex_add_factors = batch .column_by_name(EX_ADD_FACTORS_COLUMN) .map(|factors| factors.as_primitive::().clone()); @@ -2725,12 +3234,19 @@ impl QuantizerStorage for RabitQuantizationStorage { Ok(Self { metadata: self.metadata.clone(), distance_type: self.distance_type, + high_add_factors: batch + .column_by_name(super::layered::HIGH_ADD_FACTORS_COLUMN) + .map(|a| a.as_primitive::().clone()), + high_scale_factors: batch + .column_by_name(super::layered::HIGH_SCALE_FACTORS_COLUMN) + .map(|a| a.as_primitive::().clone()), batch, codes, add_factors, scale_factors, error_factors, ex_codes, + ex_codes_lo, packed_ex_codes, ex_add_factors, ex_scale_factors, @@ -2850,6 +3366,31 @@ fn get_rq_code( #[cfg(test)] mod tests { use super::*; + + #[test] + fn test_sparse_packed_sign_gather_matches_full_unpack() { + for count in [1usize, 31, 32, 33, 63, 64, 65, 511] { + for width in [8usize, 128, 256] { + let values: Vec = (0..count * width) + .map(|i| ((i * 71 + 13) % 251) as u8) + .collect(); + let original = + FixedSizeListArray::try_new_from_values(UInt8Array::from(values), width as i32) + .unwrap(); + let packed = pack_codes(&original); + let rows: Vec = (0..count) + .filter(|row| row % 3 == 0 || *row == count - 1) + .map(|row| row as u32) + .collect(); + let selected = take_packed_codes(&packed, &rows).unwrap(); + let expected = + arrow_select::take::take(&original, &UInt32Array::from(rows), None).unwrap(); + assert_eq!(unpack_codes(&selected), *expected.as_fixed_size_list()); + assert!(take_packed_codes(&packed, &[count as u32]).is_err()); + } + } + } + use rstest::rstest; use std::collections::{BinaryHeap, HashMap}; @@ -3298,6 +3839,7 @@ mod tests { code_dim: code_dim as u32, num_bits: 2, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = @@ -3424,6 +3966,7 @@ mod tests { code_dim: code_dim as u32, num_bits, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = FixedSizeListArray::try_new_from_values( @@ -3588,6 +4131,7 @@ mod tests { code_dim: code_dim as u32, num_bits: 2, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = @@ -3613,6 +4157,7 @@ mod tests { None, &mut f32_scratch, DistanceCalculatorOptions { + rq_precision: Default::default(), approx_mode: ApproxMode::Fast, }, ); @@ -3691,6 +4236,7 @@ mod tests { None, &mut f32_scratch, DistanceCalculatorOptions { + rq_precision: Default::default(), approx_mode: ApproxMode::Accurate, }, ); @@ -3759,6 +4305,7 @@ mod tests { code_dim: code_dim as u32, num_bits, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let code_len = rabit_binary_code_bytes(code_dim); @@ -3963,6 +4510,7 @@ mod tests { code_dim: code_dim as u32, num_bits, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = FixedSizeListArray::try_new_from_values( @@ -4047,6 +4595,7 @@ mod tests { code_dim: code_dim as u32, num_bits, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = FixedSizeListArray::try_new_from_values( @@ -4215,6 +4764,7 @@ mod tests { 0.7, &self.codes, Some(&self.ex_codes), + None, blocked_ex_code_bytes(CRAFTED_TOPK_DIM, CRAFTED_TOPK_NUM_BITS - 1), &self.add_factors, &self.scale_factors, @@ -4365,6 +4915,7 @@ mod tests { code_dim: code_dim as u32, num_bits: 1, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = @@ -4399,6 +4950,7 @@ mod tests { code_dim: code_dim as u32, num_bits: 2, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = @@ -4496,6 +5048,7 @@ mod tests { assert!((factors.error - expected_error).abs() < 1e-4); } let calc = storage.distance_calculator_from_parts(RabitDistCalculatorParts { + rq_precision: Default::default(), dim: DIM, dist_table: Cow::Owned(build_dist_table_direct::(&scaled)), ex_query: Cow::Borrowed(&scaled), @@ -4588,6 +5141,7 @@ mod tests { let centroid = vec![centroid_value; DIM]; let expected = 1.0 - dot(&query, ¢roid); let calc = storage.distance_calculator_from_parts(RabitDistCalculatorParts { + rq_precision: Default::default(), dim: DIM, dist_table: Cow::Owned(build_dist_table_direct::(&query)), ex_query: Cow::Borrowed(&query), @@ -4652,6 +5206,7 @@ mod tests { assert!(factors.error as f64 >= residual_norm); assert!((factors.error as f64 - residual_norm).abs() < 1e-10); let calc = storage.distance_calculator_from_parts(RabitDistCalculatorParts { + rq_precision: Default::default(), dim: DIM, dist_table: Cow::Owned(build_dist_table_direct::(&query)), ex_query: Cow::Borrowed(&query), diff --git a/rust/lance-index/src/vector/bq/transform.rs b/rust/lance-index/src/vector/bq/transform.rs index 75b20dd67f8..0b82f1c852c 100644 --- a/rust/lance-index/src/vector/bq/transform.rs +++ b/rust/lance-index/src/vector/bq/transform.rs @@ -17,6 +17,9 @@ use rayon::prelude::*; use tracing::instrument; use crate::vector::bq::builder::RabitQuantizer; +use crate::vector::bq::layered::{ + HIGH_ADD_FACTORS_COLUMN, HIGH_SCALE_FACTORS_COLUMN, RQLayout, split_codes, storage_fields, +}; use crate::vector::bq::rabit_ex_bits; use crate::vector::bq::storage::{ RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_CODE_COLUMN, RabitQueryEstimator, @@ -33,7 +36,7 @@ pub const EX_ADD_FACTORS_COLUMN: &str = "__add_factors_ex"; pub const EX_SCALE_FACTORS_COLUMN: &str = "__scale_factors_ex"; pub const ERROR_FACTORS_COLUMN: &str = "__error_factors"; -const RABIT_ERROR_EPSILON: f32 = 1.9; +pub(super) const RABIT_ERROR_EPSILON: f32 = 1.9; pub static ADD_FACTORS_FIELD: LazyLock = LazyLock::new(|| { arrow_schema::Field::new(ADD_FACTORS_COLUMN, arrow_schema::DataType::Float32, true) @@ -446,7 +449,13 @@ impl Transformer for RQTransformer { let has_split_codes = self.rq.num_bits() == 1 || (batch.column_by_name(RABIT_BLOCKED_EX_CODE_COLUMN).is_some() && batch.column_by_name(EX_ADD_FACTORS_COLUMN).is_some() - && batch.column_by_name(EX_SCALE_FACTORS_COLUMN).is_some()); + && batch.column_by_name(EX_SCALE_FACTORS_COLUMN).is_some() + && (!self.rq.metadata_ref().layered + || (batch + .column_by_name(super::storage::RABIT_BLOCKED_EX_CODE_LO_COLUMN) + .is_some() + && batch.column_by_name(HIGH_ADD_FACTORS_COLUMN).is_some() + && batch.column_by_name(HIGH_SCALE_FACTORS_COLUMN).is_some()))); if batch.column_by_name(RABIT_CODE_COLUMN).is_some() && has_split_codes { return Ok(batch.clone()); } @@ -485,7 +494,10 @@ impl Transformer for RQTransformer { } }; - if self.rq.metadata_ref().query_estimator == RabitQueryEstimator::RawQuery { + // Layered encoding also needs per-dimension codes for its high/low planes. + if self.rq.metadata_ref().query_estimator == RabitQueryEstimator::RawQuery + && !self.rq.metadata_ref().layered + { return self.transform_raw_query(batch, residual_vectors, &res_norm_square); } @@ -575,6 +587,147 @@ impl Transformer for RQTransformer { batch = batch .try_with_column(ADD_FACTORS_FIELD.clone(), Arc::new(add_factors))? .try_with_column(SCALE_FACTORS_FIELD.clone(), Arc::new(scale_factors))?; + } else { + // New RQ indexes use the RaBitQ-Library raw-query estimator. + let ex_bits = rabit_ex_bits(self.rq.num_bits())?; + let ex_codes = rq_codes.ex_codes; + let ex_res_dot_dists = rq_codes.ex_res_dot_dists; + let rotated_residuals = rq_codes.rotated_residuals.ok_or_else(|| { + Error::internal("RabitQ quantization did not return rotated residuals".to_string()) + })?; + let ex_code_values = rq_codes.ex_code_values; + if ex_bits != 0 + && (ex_codes.is_none() || ex_res_dot_dists.is_none() || ex_code_values.is_none()) + { + return Err(Error::internal( + "RabitQ multi-bit quantization did not return split-code values".to_string(), + )); + } + + let part_ids = batch[PART_ID_COLUMN].as_primitive::(); + let rotated_centroids = self.rotated_centroids.as_ref().ok_or_else(|| { + Error::internal("RabitQ raw-query transformer is missing rotated centroids") + })?; + let raw_query_factors = compute_raw_query_factors( + self.distance_type, + &res_norm_square, + &rotated_residuals, + rotated_centroids, + part_ids, + ex_code_values.as_deref(), + ex_res_dot_dists.as_deref(), + ex_bits, + self.rq.dim(), + )?; + + let layered = if self.rq.metadata_ref().layered { + let layout = RQLayout::try_new(self.rq.num_bits())?; + let values = ex_code_values + .as_deref() + .ok_or_else(|| Error::internal("missing full ex codes"))?; + let (hi, lo, hi_values) = split_codes(values, self.rq.dim(), layout)?; + let bias = -((1u32 << layout.high_bits) as f32 - 0.5); + let hi_res_dot: Vec = rotated_residuals + .chunks_exact(self.rq.dim()) + .zip(hi_values.chunks_exact(self.rq.dim())) + .map(|(residual, codes)| { + residual + .iter() + .zip(codes) + .map(|(&r, &c)| { + let sign = u32::from(r.is_sign_positive()); + r * (((sign << layout.high_bits) + u32::from(c)) as f32 + bias) + }) + .sum() + }) + .collect(); + let factors = compute_raw_query_factors( + self.distance_type, + &res_norm_square, + &rotated_residuals, + rotated_centroids, + part_ids, + Some(&hi_values), + Some(&hi_res_dot), + layout.high_bits, + self.rq.dim(), + )?; + let high_bounds = super::layered::estimator_bounds( + &rotated_residuals, + &hi_values, + self.rq.dim(), + layout.high_bits, + &raw_query_factors, + &factors, + )?; + let full_bounds = super::layered::estimator_bounds( + &rotated_residuals, + values, + self.rq.dim(), + ex_bits, + &raw_query_factors, + &raw_query_factors, + )?; + Some((hi, lo, factors, high_bounds, full_bounds)) + } else { + None + }; + + batch = batch + .try_with_column( + ADD_FACTORS_FIELD.clone(), + Arc::new(raw_query_factors.add_factors), + )? + .try_with_column( + SCALE_FACTORS_FIELD.clone(), + Arc::new(raw_query_factors.scale_factors), + )? + .try_with_column( + ERROR_FACTORS_FIELD.clone(), + Arc::new(raw_query_factors.error_factors), + )?; + + if let Some((hi, lo, factors, high_bounds, full_bounds)) = layered { + let fields = storage_fields(self.rq.dim(), self.rq.num_bits(), Vec::new())?; + batch = batch + .try_with_column(fields[0].clone(), hi)? + .try_with_column(fields[1].clone(), lo)? + .try_with_column(fields[6].clone(), high_bounds)? + .try_with_column(fields[7].clone(), full_bounds)? + .try_with_column( + fields[4].clone(), + Arc::new( + factors + .ex_add_factors + .ok_or_else(|| Error::internal("missing high add factors"))?, + ), + )? + .try_with_column( + fields[5].clone(), + Arc::new( + factors + .ex_scale_factors + .ok_or_else(|| Error::internal("missing high scale factors"))?, + ), + )?; + } else if let Some(ex_codes) = ex_codes { + batch = batch.try_with_column( + crate::vector::bq::storage::rabit_ex_code_field( + self.rq.dim(), + self.rq.num_bits(), + )? + .expect("ex-code field should exist for num_bits > 1"), + ex_codes, + )?; + } + if let Some(ex_add_factors) = raw_query_factors.ex_add_factors { + batch = batch + .try_with_column(EX_ADD_FACTORS_FIELD.clone(), Arc::new(ex_add_factors))?; + } + if let Some(ex_scale_factors) = raw_query_factors.ex_scale_factors { + batch = batch + .try_with_column(EX_SCALE_FACTORS_FIELD.clone(), Arc::new(ex_scale_factors))?; + } } let batch = batch @@ -599,6 +752,7 @@ mod tests { use crate::vector::bq::RQRotationType; use crate::vector::bq::builder::RabitQuantizer; + use crate::vector::bq::ex_dot::blocked_ex_code_bytes; use crate::vector::bq::storage::RABIT_BLOCKED_EX_CODE_COLUMN; use crate::vector::transform::Transformer; diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 370011eb8bd..8227a1e9d36 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/rust/lance-index/src/vector/distributed/index_merger.rs @@ -378,7 +378,15 @@ pub async fn init_writer_for_rq( if rq_meta.query_estimator == RabitQueryEstimator::RawQuery { fields.push(ERROR_FACTORS_FIELD.clone()); } - if let Some(ex_code_field) = rabit_ex_code_field(rq_meta.rotated_dim(), rq_meta.num_bits)? { + if rq_meta.layered { + fields = crate::vector::bq::layered::storage_fields( + rq_meta.rotated_dim(), + rq_meta.num_bits, + fields, + )?; + } else if let Some(ex_code_field) = + rabit_ex_code_field(rq_meta.rotated_dim(), rq_meta.num_bits)? + { fields.push(ex_code_field); fields.push(EX_ADD_FACTORS_FIELD.clone()); fields.push(EX_SCALE_FACTORS_FIELD.clone()); @@ -1169,17 +1177,20 @@ async fn merge_partial_vector_auxiliary_files_inner( if let Some(existing_rq) = rq_meta.as_ref() && (existing_rq.code_dim != rq_meta_parsed.code_dim || existing_rq.num_bits != rq_meta_parsed.num_bits + || existing_rq.layered != rq_meta_parsed.layered || existing_rq.rotation_type != rq_meta_parsed.rotation_type || existing_rq.query_estimator != rq_meta_parsed.query_estimator || existing_rq.fast_rotation_signs != rq_meta_parsed.fast_rotation_signs) { return Err(Error::index(format!( - "Distributed RQ merge: structural mismatch across shards; first(code_dim={}, num_bits={}, rotation_type={:?}), current(code_dim={}, num_bits={}, rotation_type={:?})", + "Distributed RQ merge: structural mismatch across shards; first(code_dim={}, num_bits={}, layered={}, rotation_type={:?}), current(code_dim={}, num_bits={}, layered={}, rotation_type={:?})", existing_rq.code_dim, existing_rq.num_bits, + existing_rq.layered, existing_rq.rotation_type, rq_meta_parsed.code_dim, rq_meta_parsed.num_bits, + rq_meta_parsed.layered, rq_meta_parsed.rotation_type ))); } @@ -1628,12 +1639,12 @@ async fn merge_partial_vector_auxiliary_files_inner( Some(meta) if meta.num_bits > 1 => batches .into_iter() .map(|batch| { - crate::vector::bq::storage::load_blocked_ex_codes( + crate::vector::bq::storage::load_ex_code_planes( batch, meta.rotated_dim(), meta.num_bits, ) - .map(|(batch, _)| batch) + .map(|(batch, _, _)| batch) }) .collect::>>()?, _ => batches, @@ -2428,7 +2439,14 @@ mod tests { distance_type: DistanceType, ) -> Result { let num_bytes = (metadata.code_dim as usize).div_ceil(u8::BITS as usize); - let ex_code_field = rabit_ex_code_field(metadata.code_dim as usize, metadata.num_bits)?; + let layout = metadata + .layered + .then(|| crate::vector::bq::layered::RQLayout::try_new(metadata.num_bits)) + .transpose()?; + let ex_code_field = rabit_ex_code_field( + metadata.code_dim as usize, + layout.map_or(metadata.num_bits, |layout| layout.high_bits + 1), + )?; let ex_code_bytes = ex_code_field.as_ref().map(|field| { let DataType::FixedSizeList(_, num_bytes) = field.data_type() else { panic!("RQ ex-code field should be FixedSizeList"); @@ -2451,7 +2469,13 @@ mod tests { if metadata.query_estimator == RabitQueryEstimator::RawQuery { fields.push(ERROR_FACTORS_FIELD.clone()); } - if let Some(field) = ex_code_field { + if metadata.layered { + fields = crate::vector::bq::layered::storage_fields( + metadata.code_dim as usize, + metadata.num_bits, + fields, + )?; + } else if let Some(field) = ex_code_field { fields.push(field); fields.push(EX_ADD_FACTORS_FIELD.clone()); fields.push(EX_SCALE_FACTORS_FIELD.clone()); @@ -2529,8 +2553,22 @@ mod tests { UInt8Array::from(ex_codes), ex_code_bytes as i32, )?)); + if let Some(layout) = layout { + let bytes = crate::vector::bq::ex_dot::blocked_ex_code_bytes( + metadata.code_dim as usize, + layout.low_bits, + ); + columns.push(Arc::new(FixedSizeListArray::try_new_from_values( + UInt8Array::from(vec![0x55; total_rows * bytes]), + bytes as i32, + )?)); + } columns.push(Arc::new(Float32Array::from(ex_add_factors))); columns.push(Arc::new(Float32Array::from(ex_scale_factors))); + if metadata.layered { + columns.push(Arc::new(Float32Array::from(vec![100.; total_rows]))); + columns.push(Arc::new(Float32Array::from(vec![10.; total_rows]))); + } } let batch = RecordBatch::try_new(Arc::new(arrow_schema), columns)?; @@ -2752,6 +2790,7 @@ mod tests { code_dim: 16, num_bits: 1, packed: false, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; @@ -2884,6 +2923,7 @@ mod tests { code_dim: 16, num_bits: 1, packed: true, + layered: false, query_estimator: RabitQueryEstimator::RawQuery, }; @@ -2925,8 +2965,16 @@ mod tests { } } + #[rstest::rstest] + #[case::single(4, false)] + #[case::rq5(5, true)] + #[case::rq7(7, true)] + #[case::rq9(9, true)] #[tokio::test] - async fn test_merge_ivf_rq_multi_bit_preserves_split_columns() { + async fn test_merge_ivf_rq_multi_bit_preserves_split_columns( + #[case] num_bits: u8, + #[case] layered: bool, + ) { let object_store = ObjectStore::memory(); let index_dir = Path::from("index/uuid_rq_multi_bit"); @@ -2944,8 +2992,9 @@ mod tests { fast_rotation_signs: Some(vec![0xAA; 2]), rotation_type: RQRotationType::Fast, code_dim: 16, - num_bits: 4, + num_bits, packed: false, + layered, query_estimator: RabitQueryEstimator::RawQuery, }; @@ -3000,7 +3049,8 @@ mod tests { let meta = reader.metadata(); let rq_meta_json = meta.file_schema.metadata.get(RABIT_METADATA_KEY).unwrap(); let merged_rq_meta: RabitQuantizationMetadata = serde_json::from_str(rq_meta_json).unwrap(); - assert_eq!(merged_rq_meta.num_bits, 4); + assert_eq!(merged_rq_meta.num_bits, num_bits); + assert_eq!(merged_rq_meta.layered, layered); assert!(merged_rq_meta.packed); let mut total_rows = 0usize; @@ -3025,7 +3075,32 @@ mod tests { panic!("RQ ex-code field should be FixedSizeList"); }; // code_dim=16 padded to one 64-dim block at ex_bits=3. - assert_eq!(*ex_code_bytes, 24); + let high_bits = if layered { + crate::vector::bq::layered::RQLayout::try_new(num_bits) + .unwrap() + .high_bits + } else { + num_bits - 1 + }; + assert_eq!( + *ex_code_bytes, + crate::vector::bq::ex_dot::blocked_ex_code_bytes(16, high_bits) as i32 + ); + if layered { + assert!( + schema + .field_with_name( + crate::vector::bq::storage::RABIT_BLOCKED_EX_CODE_LO_COLUMN + ) + .is_ok() + ); + assert_eq!( + batch[crate::vector::bq::layered::HIGH_ADD_FACTORS_COLUMN] + .as_primitive::() + .value(0), + 100. + ); + } assert!(schema.field_with_name(ERROR_FACTORS_FIELD.name()).is_ok()); assert!(schema.field_with_name(EX_ADD_FACTORS_COLUMN).is_ok()); assert!(schema.field_with_name(EX_SCALE_FACTORS_COLUMN).is_ok()); @@ -3042,6 +3117,49 @@ mod tests { assert_eq!(total_rows, expected_total); } + #[tokio::test] + async fn test_merge_rq_rejects_mixed_layouts() { + let store = ObjectStore::memory(); + let root = Path::from("mixed-rq"); + let a = root + .clone() + .join("partial_0") + .join(INDEX_AUXILIARY_FILE_NAME); + let b = root + .clone() + .join("partial_1") + .join(INDEX_AUXILIARY_FILE_NAME); + let mut metadata = RabitQuantizationMetadata { + rotate_mat: None, + rotate_mat_position: None, + fast_rotation_signs: Some(vec![0xAA; 2]), + rotation_type: RQRotationType::Fast, + code_dim: 16, + num_bits: 5, + packed: false, + layered: false, + query_estimator: RabitQueryEstimator::RawQuery, + }; + write_rq_partial_aux(&store, &a, &metadata, &[2], 0, DistanceType::L2) + .await + .unwrap(); + metadata.layered = true; + write_rq_partial_aux(&store, &b, &metadata, &[2], 10, DistanceType::L2) + .await + .unwrap(); + let error = merge_partial_vector_auxiliary_files( + &store, + &[a, b], + &root, + crate::progress::noop_progress(), + ) + .await + .unwrap_err(); + assert!(matches!(error, Error::Index { .. })); + assert!(error.to_string().contains("structural mismatch")); + assert!(error.to_string().contains("layered=true")); + } + #[tokio::test] async fn test_merge_ivf_pq_codebook_mismatch() { let object_store = ObjectStore::memory(); diff --git a/rust/lance-index/src/vector/flat/index.rs b/rust/lance-index/src/vector/flat/index.rs index 6c80c612ccb..a6b8b7ebf7d 100644 --- a/rust/lance-index/src/vector/flat/index.rs +++ b/rust/lance-index/src/vector/flat/index.rs @@ -87,6 +87,7 @@ pub struct FlatQueryParams { upper_bound: Option, dist_q_c: f32, approx_mode: ApproxMode, + rq_precision: super::super::bq::layered::RQPrecision, } impl From<&Query> for FlatQueryParams { @@ -96,6 +97,7 @@ impl From<&Query> for FlatQueryParams { upper_bound: q.upper_bound, dist_q_c: q.dist_q_c, approx_mode: q.approx_mode, + rq_precision: q.rq_precision, } } } @@ -158,6 +160,7 @@ impl IvfSubIndex for FlatIndex { residual, &mut scratch.query_f32, DistanceCalculatorOptions { + rq_precision: params.rq_precision, approx_mode: params.approx_mode, }, ); @@ -303,6 +306,7 @@ impl IvfSubIndex for FlatIndex { residual, &mut scratch.query_f32, DistanceCalculatorOptions { + rq_precision: params.rq_precision, approx_mode: params.approx_mode, }, ); diff --git a/rust/lance-index/src/vector/quantizer.rs b/rust/lance-index/src/vector/quantizer.rs index 13477ae660f..221b066d1b2 100644 --- a/rust/lance-index/src/vector/quantizer.rs +++ b/rust/lance-index/src/vector/quantizer.rs @@ -286,6 +286,22 @@ pub trait QuantizerStorage: Clone + Sized + DeepSizeOf + VectorStore { Self::try_from_batch(batch, metadata, distance_type, None) } + /// Materialize a projected precision level. Non-RQ stores accept only full. + fn try_from_batch_at_precision( + batch: RecordBatch, + metadata: &Self::Metadata, + distance_type: DistanceType, + remapper: Option>, + precision: super::bq::layered::RQPrecision, + ) -> Result { + if precision != super::bq::layered::RQPrecision::Full { + return Err(Error::invalid_input( + "rq_precision requires a layered IVF_RQ index", + )); + } + Self::try_from_batch_with_remapper(batch, metadata, distance_type, remapper) + } + 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 35e15ab64a7..8da3e6e448b 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -3,6 +3,7 @@ //! Vector Storage, holding (quantized) vectors and providing distance calculation. +use crate::vector::bq::layered::{PlaneBatch, PlaneKey}; use crate::vector::quantizer::QuantizerStorage; use arrow::compute::concat_batches; use arrow_array::{ @@ -11,6 +12,7 @@ use arrow_array::{ use arrow_schema::SchemaRef; use futures::stream::{self, Stream, StreamExt, TryStreamExt}; use lance_arrow::{FixedSizeListArrayExt, RecordBatchExt}; +use lance_core::cache::WeakLanceCache; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_core::{Error, ROW_ID, Result}; @@ -25,10 +27,10 @@ use prost::Message; use std::{ any::Any, borrow::Cow, - collections::BinaryHeap, + collections::{BinaryHeap, HashMap}, mem::size_of, ops::{Deref, DerefMut}, - sync::Arc, + sync::{Arc, Mutex}, }; use crossbeam_queue::ArrayQueue; @@ -337,6 +339,7 @@ impl QueryScratchCapacity { #[derive(Clone, Copy, Debug, Default)] pub struct DistanceCalculatorOptions { pub approx_mode: ApproxMode, + pub rq_precision: super::bq::layered::RQPrecision, } #[derive(Debug)] @@ -601,6 +604,34 @@ impl StorageBuilder { } } +/// Avoid loading whole cold planes for isolated candidate reads. Repeated reads +/// covering half a partition amortize promotion, with a cooldown after admission pressure. +const PLANE_PROMOTION_READS: usize = 3; +const PLANE_PROMOTION_COOLDOWN: std::time::Duration = std::time::Duration::from_secs(60); + +#[derive(Debug, Default)] +struct PlaneAccess { + reads: usize, + rows: usize, + last_attempt: Option, +} +impl DeepSizeOf for PlaneAccess { + fn deep_size_of_children(&self, _: &mut lance_core::deepsize::Context) -> usize { + 0 + } +} + +/// 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>>); + +impl DeepSizeOf for PlaneAccessTracker { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.0.deep_size_of_children(context) + } +} + /// Loader to load partitioned PQ storage from disk. #[derive(Debug)] pub struct IvfQuantizationStorage { @@ -615,11 +646,14 @@ pub struct IvfQuantizationStorage { frag_reuse_index: Option>, /// Asynchronous batch remapper (tagged histories). batch_remapper: Option>, + plane_access: PlaneAccessTracker, } impl DeepSizeOf for IvfQuantizationStorage { fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - self.metadata.deep_size_of_children(context) + self.ivf.deep_size_of_children(context) + self.metadata.deep_size_of_children(context) + + self.ivf.deep_size_of_children(context) + + self.plane_access.deep_size_of_children(context) } } @@ -700,6 +734,7 @@ impl IvfQuantizationStorage { ivf, frag_reuse_index, batch_remapper: None, + plane_access: Default::default(), }) } @@ -732,6 +767,7 @@ impl IvfQuantizationStorage { ivf, frag_reuse_index, batch_remapper: None, + plane_access: Default::default(), } } @@ -746,6 +782,17 @@ impl IvfQuantizationStorage { self } + /// Promotion history to retain in the index's cached runtime state. + pub fn plane_access_tracker(&self) -> &PlaneAccessTracker { + &self.plane_access + } + + /// Reuse promotion history when binding new readers to a cached index. + pub fn with_plane_access_tracker(mut self, tracker: PlaneAccessTracker) -> Self { + self.plane_access = tracker; + self + } + pub fn reader(&self) -> &FileReader { &self.reader } @@ -842,6 +889,388 @@ impl IvfQuantizationStorage { ) } + /// Warm every plane entry. Backend admission enforces its byte budget. + /// + /// A backend that gates lower planes on their sign plane is warmed plane + /// by plane: all sign planes first, then the high and low planes of + /// partitions whose sign plane stayed resident. Other backends admit plane + /// entries like any entry, so they are warmed partition by partition, + /// which loads (and persists) every plane whatever fits in RAM, as the + /// native partition prewarm does. + pub async fn prewarm_planes(&self, cache: &WeakLanceCache) -> Result<()> { + if !cache.plane_admission_gated() { + for part_id in 0..self.num_partitions() { + for plane in 0..=2 { + cache + .get_or_insert_with_key( + PlaneKey { + partition: part_id, + plane, + }, + || async { + Ok(PlaneBatch( + self.read_plane(part_id, plane, None, None).await?, + )) + }, + ) + .await?; + } + } + return Ok(()); + } + for plane in 0..=2 { + for part_id in 0..self.num_partitions() { + if plane > 0 + && cache + .get_resident_with_key(&PlaneKey { + partition: part_id, + plane: 0, + }) + .await + .is_none() + { + continue; + } + cache + .get_or_insert_with_key( + PlaneKey { + partition: part_id, + plane, + }, + || async { + Ok(PlaneBatch( + self.read_plane(part_id, plane, None, None).await?, + )) + }, + ) + .await?; + } + } + Ok(()) + } + + /// Whether this loader owns an opt-in layered RaBitQ index. + pub fn is_layered_rq(&self) -> bool { + matches!(self.quantizer(), Ok(Quantizer::Rabit(rq)) if rq.metadata_ref().layered) + } + + /// Read independent plane entries through the existing persistent cache. + pub async fn load_partition_at_precision( + &self, + part_id: usize, + precision: super::bq::layered::RQPrecision, + cache: &WeakLanceCache, + io_stats: Option, + ) -> Result { + use super::bq::layered::RQPrecision; + if !self.is_layered_rq() { + if precision != RQPrecision::Full { + return Err(Error::invalid_input( + "rq_precision requires a layered IVF_RQ index", + )); + } + return self.load_partition(part_id, io_stats).await; + } + let last = match precision { + RQPrecision::Sign => 0, + RQPrecision::High => 1, + 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 mut fields = Vec::new(); + let mut columns = Vec::new(); + for batch in std::iter::once(sign).chain(ex) { + fields.extend(batch.0.schema().fields().iter().cloned()); + columns.extend(batch.0.columns().iter().cloned()); + } + let batch = RecordBatch::try_new(Arc::new(arrow_schema::Schema::new(fields)), columns)?; + let (batch, remapper) = if let Some(remapping) = &self.batch_remapper { + lance_index_core::remapping::check_batch_remapping_entry()?; + let row_id_idx = batch.schema().index_of(ROW_ID)?; + let (batch, remapper) = + remap_row_ids_preserving_layout_async(remapping.as_ref(), batch, row_id_idx) + .await?; + (batch, Some(remapper)) + } else { + (batch, self.frag_reuse_index.clone()) + }; + Q::Storage::try_from_batch_at_precision( + batch, + self.metadata(), + self.distance_type, + remapper, + precision, + ) + } + + /// Remappers can remove physical rows, so candidate offsets require an identity mapping. + pub fn supports_candidate_reads(&self) -> bool { + self.frag_reuse_index.is_none() && self.batch_remapper.is_none() + } + + /// Load one whole plane entry of a layered partition. + /// + /// A backend that gates lower planes on their sign plane admits a lower + /// plane only while the partition's sign plane is resident; otherwise the + /// 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( + &self, + part_id: usize, + plane: u8, + cache: &WeakLanceCache, + io_stats: Option, + ) -> Result> { + let key = PlaneKey { + partition: part_id, + plane, + }; + if !cache.plane_admission_gated() { + return cache + .get_or_insert_with_key(key, || async { + Ok(PlaneBatch( + self.read_plane(part_id, plane, None, io_stats.clone()) + .await?, + )) + }) + .await; + } + let sign_resident = plane == 0 + || cache + .get_resident_with_key(&PlaneKey { + partition: part_id, + plane: 0, + }) + .await + .is_some(); + let batch = if sign_resident { + cache + .get_or_insert_with_key_hit(key, || async { + Ok(PlaneBatch( + self.read_plane(part_id, plane, None, io_stats.clone()) + .await?, + )) + }) + .await? + .0 + } else if let Some(batch) = cache.get_without_promotion_with_key(&key).await { + batch + } else { + // An oversized sign plane falls back to partition streaming; do + // not admit smaller ex entries without their sign dependency. + Arc::new(PlaneBatch( + self.read_plane(part_id, plane, None, io_stats.clone()) + .await?, + )) + }; + Ok(batch) + } + + /// Assemble a full-precision store from candidate rows only. + pub async fn load_candidates( + &self, + part_id: usize, + rows: Vec, + cache: &lance_core::cache::WeakLanceCache, + io_stats: Option, + ) -> Result { + use super::bq::layered::RQPrecision; + use super::bq::storage::{RABIT_CODE_COLUMN, take_packed_codes}; + use arrow_array::cast::AsArray; + if !self.supports_candidate_reads() { + return Err(Error::invalid_input( + "candidate reads require an index without a row-id remapper", + )); + } + if rows.windows(2).any(|w| w[0] >= w[1]) { + return Err(Error::invalid_input( + "candidate offsets must be sorted and unique", + )); + } + let indices = arrow_array::UInt32Array::from(rows.clone()); + // The three planes are independent. Overlap their cache/origin reads, + // while preserving plane order when assembling the full row schema. + let batches = futures::future::try_join_all((0..=2).map(|plane| { + let rows = &rows; + let indices = &indices; + let io_stats = io_stats.clone(); + async move { + let key = PlaneKey { + partition: part_id, + plane, + }; + let resident = cache.get_resident_with_key(&key).await; + let was_resident = resident.is_some(); + let cached = if plane == 0 && resident.is_none() { + cache.get_without_promotion_with_key(&key).await + } else { + resident + }; + let selected_cached = if plane > 0 && cached.is_none() { + cache.get_rows_with_key(&key, rows).await + } else { + None + }; + let batch = if plane == 0 { + // Sign codes are transposed across rows. Gather their physical + // offsets directly, then pack only the selected rows. + let raw = if let Some(value) = cached { + value.0.clone() + } else { + self.read_plane(part_id, plane, None, io_stats.clone()) + .await? + }; + let codes = raw + .column_by_name(RABIT_CODE_COLUMN) + .ok_or_else(|| Error::invalid_input("missing sign codes"))?; + let selected_codes = take_packed_codes(codes.as_fixed_size_list(), rows)?; + // The transposed sign column cannot be gathered with Arrow + // take. Avoid copying it only to replace that copy immediately. + let field = raw.schema().field_with_name(RABIT_CODE_COLUMN)?.clone(); + raw.drop_column(RABIT_CODE_COLUMN)? + .take(indices)? + .try_with_column(field, Arc::new(selected_codes))? + } else if let Some(value) = cached { + value.0.take(indices)? + } else if let Some(value) = selected_cached { + value.0.clone() + } else { + self.read_plane(part_id, plane, Some(rows.clone()), io_stats.clone()) + .await? + }; + if plane > 0 + && !was_resident + && cache + .get_resident_with_key(&PlaneKey { + partition: part_id, + plane: 0, + }) + .await + .is_some() + { + let promote = { + let mut accesses = self + .plane_access + .0 + .lock() + .unwrap_or_else(|e| e.into_inner()); + let access = accesses.entry((part_id, plane)).or_default(); + access.reads = access.reads.saturating_add(1); + access.rows = access.rows.saturating_add(rows.len()); + let cooled = access + .last_attempt + .is_none_or(|last| last.elapsed() >= PLANE_PROMOTION_COOLDOWN); + if cooled + && access.reads >= PLANE_PROMOTION_READS + && access.rows >= self.partition_size(part_id).div_ceil(2) + { + access.reads = 0; + access.rows = 0; + access.last_attempt = Some(std::time::Instant::now()); + true + } else { + false + } + }; + if promote { + cache + .get_or_insert_with_key( + PlaneKey { + partition: part_id, + plane, + }, + || async { + Ok(super::bq::layered::PlaneBatch( + self.read_plane(part_id, plane, None, io_stats.clone()) + .await?, + )) + }, + ) + .await?; + } + } + Ok::<_, Error>(batch) + } + })) + .await?; + let mut fields = Vec::new(); + let mut columns = Vec::new(); + for batch in batches { + fields.extend(batch.schema().fields().iter().cloned()); + columns.extend(batch.columns().iter().cloned()); + } + let batch = RecordBatch::try_new(Arc::new(arrow_schema::Schema::new(fields)), columns)?; + Q::Storage::try_from_batch_at_precision( + batch, + self.metadata(), + self.distance_type, + self.frag_reuse_index.clone(), + RQPrecision::Full, + ) + } + + /// Read only selected rows of a plane; sorted indices are coalesced by the reader. + pub async fn read_plane( + &self, + part_id: usize, + plane: u8, + rows: Option>, + io_stats: Option, + ) -> Result { + let projection = lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + super::bq::layered::plane_columns(plane), + )?; + let schema = Arc::new(arrow_schema::Schema::from(projection.schema.as_ref())); + let range = self.ivf.row_range(part_id); + if range.is_empty() || rows.as_ref().is_some_and(Vec::is_empty) { + return Ok(RecordBatch::new_empty(schema)); + } + let params = if let Some(rows) = rows { + let mut ranges: Vec> = Vec::new(); + for row in rows { + if row as usize >= range.len() { + return Err(Error::invalid_input( + "candidate offset exceeds partition length", + )); + } + let start = range.start as u64 + u64::from(row); + if let Some(last) = ranges.last_mut() + && last.end == start + { + last.end += 1; + } else { + ranges.push(start..start + 1); + } + } + ReadBatchParams::Ranges(ranges.into()) + } else { + ReadBatchParams::Range(range) + }; + let reader = match &io_stats { + Some(stats) => Cow::Owned(self.reader.with_io_stats(stats.recorder())), + None => Cow::Borrowed(&self.reader), + }; + let batches = reader + .read_stream_projected( + params, + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + concat_batches(&schema, batches.iter()).map_err(Into::into) + } + /// Read only `columns` of the given rows; the batch keeps file schema order. async fn read_vector_columns( &self, diff --git a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs index ad942457463..f8ceab1d9fb 100644 --- a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs +++ b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs @@ -581,6 +581,8 @@ async fn run_checkpoint( let q_vec = &queries[q * DIM..(q + 1) * DIM]; let bf_set = >_sets[q]; let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: VECTOR_COL.to_string(), key: Arc::new(Float32Array::from(q_vec.to_vec())), k, diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index b72574da974..f69a4b515d9 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -2416,6 +2416,7 @@ impl Scanner { self.probe_params = ProbeParams::default(); self.nearest = Some(Query { + rq_cascade_factor: None, column: column.to_string(), key, k, @@ -2431,6 +2432,7 @@ impl Scanner { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + rq_precision: Default::default(), }); self.nearest_query_count = query_count; self.is_batch_nearest = is_batch_nearest; @@ -2634,11 +2636,30 @@ impl Scanner { self } + /// Rerank at most k * factor high-level candidates using full code planes. + pub fn rq_cascade_factor(&mut self, factor: u32) -> Result<&mut Self> { + if factor == 0 { + return Err(Error::invalid_input("rq_cascade_factor must be positive")); + } + if let Some(query) = self.nearest.as_mut() { + query.rq_cascade_factor = Some(factor); + } + Ok(self) + } + + /// Select the precision of a layered IVF_RQ index (full by default). + pub fn rq_precision( + &mut self, + precision: lance_index::vector::bq::layered::RQPrecision, + ) -> &mut Self { + if let Some(query) = self.nearest.as_mut() { + query.rq_precision = precision; + } + self + } + /// Configure the speed / accuracy tradeoff for approximate vector search. - /// - /// This setting is currently used by RQ-quantized indexes (such as - /// IVF_RQ) and by prefiltered search on HNSW indexes, where `Fast` - /// enables the ACORN traversal. Other index types ignore this setting. + /// RQ uses this for estimation; HNSW Fast mode enables ACORN traversal. pub fn approx_mode(&mut self, approx_mode: ApproxMode) -> &mut Self { if let Some(q) = self.nearest.as_mut() { q.approx_mode = approx_mode; diff --git a/rust/lance/src/dataset/tests/dataset_scanner.rs b/rust/lance/src/dataset/tests/dataset_scanner.rs index 30493dec248..2cd058160a4 100644 --- a/rust/lance/src/dataset/tests/dataset_scanner.rs +++ b/rust/lance/src/dataset/tests/dataset_scanner.rs @@ -268,6 +268,8 @@ async fn test_vector_filter_fts_search() { let query_vector = Arc::new(Float32Array::from(vec![300f32, 300f32, 300f32, 300f32])); let vector_query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key: query_vector, k: 5, diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 2f17ee88042..09ad23cb3b5 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -2221,6 +2221,7 @@ fn derive_sq_params(sq_quantizer: &ScalarQuantizer) -> SQBuildParams { fn derive_rabit_params(rabit_quantizer: &RabitQuantizer) -> RQBuildParams { RQBuildParams { num_bits: rabit_quantizer.num_bits(), + layered: rabit_quantizer.metadata_ref().layered, rotation_type: rabit_quantizer.rotation_type(), rotation: None, } diff --git a/rust/lance/src/index/vector/details.rs b/rust/lance/src/index/vector/details.rs index 6356e32f04b..1bef3f39947 100644 --- a/rust/lance/src/index/vector/details.rs +++ b/rust/lance/src/index/vector/details.rs @@ -294,7 +294,8 @@ pub fn vector_params_from_details(details: &prost_types::Any) -> Option VectorIndexParams::with_ivf_hnsw_pq_params( @@ -688,8 +689,8 @@ async fn convert_v3_metadata_to_details( } }; Some(Compression::Rq(RabitQuantization { - layered: false, - row_layout: rabit_quantization::RowLayout::Columns.into(), + layered: rq.layered, + row_layout: rabit_quantization::RowLayout::Columns as i32, num_bits: rq.num_bits as u32, rotation_type: rotation_type.into(), })) diff --git a/rust/lance/src/index/vector/fixture_test.rs b/rust/lance/src/index/vector/fixture_test.rs index c887a37210f..01a01eb5305 100644 --- a/rust/lance/src/index/vector/fixture_test.rs +++ b/rust/lance/src/index/vector/fixture_test.rs @@ -252,6 +252,8 @@ mod test { key = normalize_arrow(&key).unwrap().0; }; let q = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "test".to_string(), key, k: 1, diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 45e898babb8..68e9e30ddc6 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -5965,6 +5965,8 @@ mod tests { for row_id in row_ids_to_test { let row = self.get_vector(row_id as u32); let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: Self::COLUMN.to_string(), key: Arc::new(row), k: 5, diff --git a/rust/lance/src/index/vector/ivf/partition_serde.rs b/rust/lance/src/index/vector/ivf/partition_serde.rs index a1869357980..b1230078b17 100644 --- a/rust/lance/src/index/vector/ivf/partition_serde.rs +++ b/rust/lance/src/index/vector/ivf/partition_serde.rs @@ -414,6 +414,7 @@ impl CacheCodecImpl for PartitionEntry { fn serialize(&self, w: &mut CacheEntryWriter<'_>) -> Result<()> { let metadata = self.storage.metadata(); let header = RabitPartitionHeader { + layered: metadata.layered, distance_type: distance_type_to_proto(self.storage.distance_type()) as i32, num_bits: metadata.num_bits as u32, code_dim: metadata.code_dim, @@ -469,6 +470,7 @@ impl CacheCodecImpl for PartitionEntry { num_bits: header.num_bits as u8, // The storage batch already has packed codes; skip re-packing. packed: true, + layered: header.layered, query_estimator, }; let storage = ::Storage::try_from_batch( @@ -1143,6 +1145,7 @@ mod tests { index_file_size: 1024, aux_file_size: 512, rq_search_cache: empty_rabit_search_cache_cell(), + plane_access: Default::default(), }; let entry = IvfStateEntryBox(Arc::new(state)); diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index e0af653793a..4027cb1962e 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -62,8 +62,8 @@ use lance_index::vector::quantizer::{ }; use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::storage::{ - QueryResidual, QueryScratch, QueryScratchCapacity, QueryScratchPool, RabitRawQueryContext, - VectorStore, + PlaneAccessTracker, QueryResidual, QueryScratch, QueryScratchCapacity, QueryScratchPool, + RabitRawQueryContext, VectorStore, }; use lance_index::vector::v3::subindex::SubIndexType; use lance_index::{ @@ -127,6 +127,8 @@ pub(crate) struct IvfIndexState { pub(crate) aux_file_size: u64, /// Runtime-only cache, intentionally excluded from the CacheCodec wire format. pub(crate) rq_search_cache: RabitSearchCacheCell, + /// Runtime-only counters survive reader reconstruction between queries. + pub(crate) plane_access: PlaneAccessTracker, } /// Number of prepared partitions handed to a single `spawn_cpu` dispatch on the @@ -621,6 +623,7 @@ impl DeepSizeOf for IvfIndexState { + self.aux_ivf.deep_size_of_children(context) + self.sub_index_metadata.deep_size_of_children(context) + self.metadata.deep_size_of_children(context) + + self.plane_access.deep_size_of_children(context) + self .rq_search_cache .lock() @@ -735,6 +738,7 @@ impl CacheCodecImpl for IvfStateEntryBox { index_file_size: header.index_file_size, aux_file_size: header.aux_file_size, rq_search_cache: empty_rabit_search_cache_cell(), + plane_access: PlaneAccessTracker::default(), }))) } @@ -918,6 +922,7 @@ pub struct PartitionEntry { pub storage: Q::Storage, partition_rows: OnceLock>, partition_rows_accounted: AtomicBool, + cache_whole_partition: bool, /// Memoized size of the immutable parts (this struct, the sub-index and the /// quantized storage): every query that prepares this partition needs its /// size to budget scoring chunks, and the walk over the storage's arrays @@ -938,6 +943,7 @@ impl PartitionEntry { storage, partition_rows: OnceLock::new(), partition_rows_accounted: AtomicBool::new(false), + cache_whole_partition: true, immutable_bytes: OnceLock::new(), coverage_bytes: OnceLock::new(), } @@ -1112,7 +1118,9 @@ impl IVFIndex { partition: &Arc>, ) -> Result> { let rows = partition.partition_rows(); - if !partition.partition_rows_accounted.load(Ordering::Acquire) { + if partition.cache_whole_partition + && !partition.partition_rows_accounted.load(Ordering::Acquire) + { let cache_key = IVFPartitionKey::::new(partition_id); if index_cache .insert_with_key(&cache_key, partition.clone()) @@ -1225,11 +1233,13 @@ impl IVFIndex { raw_query_context: Option>, ) -> Result> { let _stage_timer = IndexTimer::new(metrics, IndexTiming::PartitionPrepare); - let (part_entry, ()) = - tokio::try_join!(self.load_partition(partition_id, true, metrics), async { + let (part_entry, ()) = tokio::try_join!( + self.load_initial_partition(partition_id, query, metrics), + async { let _wait_timer = IndexTimer::new(metrics, IndexTiming::PrefilterWait); pre_filter.wait_for_ready().await - },)?; + }, + )?; let pre_filter = Self::prefilter_for_partition(&self.index_cache, partition_id, &part_entry, pre_filter) .await?; @@ -1255,7 +1265,9 @@ impl IVFIndex { raw_query_context: Option>, ) -> Result> { let _stage_timer = IndexTimer::new(metrics, IndexTiming::PartitionPrepare); - let part_entry = self.load_partition(partition_id, true, metrics).await?; + let part_entry = self + .load_initial_partition(partition_id, query, metrics) + .await?; let pre_filter = Self::prefilter_for_partition(&self.index_cache, partition_id, &part_entry, pre_filter) .await?; @@ -1762,6 +1774,400 @@ impl IVFIndex { } } + /// Cut candidates across all probed partitions before fetching lower planes. + pub async fn quantized_candidates( + &self, + query: &Query, + partitions: &UInt32Array, + centroid_dists: &Float32Array, + range: Range, + pre_filter: Arc, + metrics: &dyn MetricsCollector, + ) -> Result> { + use lance_index::vector::storage::DistanceCalculatorOptions; + let start = range.start; + let end = range.end; + self.initial_precision(query)?; + if !self.storage.is_layered_rq() || !self.storage.supports_candidate_reads() { + return Err(Error::invalid_input( + "quantized candidates require a layered index without a row-id remapper", + )); + } + if partitions.len() != centroid_dists.len() || start > end || end > partitions.len() { + return Err(Error::invalid_input( + "invalid quantized candidate partition range", + )); + } + let factor = query + .rq_cascade_factor + .ok_or_else(|| Error::internal("missing cascade factor"))?; + let k = query + .k + .checked_mul(query.refine_factor.unwrap_or(1) as usize) + .ok_or_else(|| Error::invalid_input("cascade k overflow"))?; + let limit = k + .checked_mul(factor as usize) + .ok_or_else(|| Error::invalid_input("cascade candidate count overflow"))?; + if limit == 0 { + return Ok(Vec::new()); + } + pre_filter.wait_for_ready().await?; + let raw_query = self.prepare_rq_raw_query_context(&query.key)?; + let mut coarse_heap: BinaryHeap< + OrderedNode<(usize, u32, u64, lance_index::vector::graph::OrderedFloat)>, + > = BinaryHeap::with_capacity(limit); + let mut distances = HashMap::new(); + let load_parallelism = self + .io_parallelism + .max(1) + .min(get_num_compute_intensive_cpus().max(1)); + // Keep the probe order for heap ties while overlapping independent + // partition reads. The window bounds how many loaded stores stay pinned. + let mut prepared = stream::iter(start..end) + .map(|index| async move { + let part_id = partitions.value(index) as usize; + let mut local_query = query.clone(); + local_query.dist_q_c = centroid_dists.value(index); + let entry = self + .load_initial_partition(part_id, &local_query, metrics) + .await?; + Result::Ok((part_id, local_query, entry)) + }) + .buffered(load_parallelism); + let mut pending = prepared.try_next().await?; + while let Some((part_id, local_query, entry)) = pending { + distances.insert(part_id, local_query.dist_q_c); + let filter = pre_filter.clone(); + let raw_query = raw_query.clone(); + let cache = self.rq_search_cache.clone(); + let scoring = spawn_cpu(move || -> Result<_> { + let mut scratch = Vec::new(); + let context = QueryResidual::RabitRawQuery { + rotated_centroid: rotated_partition_centroid_slice(cache.as_deref(), part_id), + query: raw_query.as_deref(), + }; + let storage = entry + .storage + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::internal("layered candidate storage is not RaBitQ"))?; + let calc = storage.dist_calculator_with_scratch( + local_query.key, + local_query.dist_q_c, + Some(context), + &mut scratch, + DistanceCalculatorOptions { + approx_mode: local_query.approx_mode, + ..Default::default() + }, + ); + let binary_inner_products = calc.binary_inner_products(); + for row in filter.filter_row_ids(Box::new(entry.storage.row_ids())) { + if coarse_heap.len() == limit + && calc + .lower_bound_with_binary_inner_product( + row as u32, + binary_inner_products[row as usize], + ) + .is_some_and(|bound| { + coarse_heap.peek().is_some_and(|top| bound >= top.dist.0) + }) + { + continue; + } + let node = OrderedNode::new( + ( + part_id, + row as u32, + entry.storage.row_id(row as u32), + binary_inner_products[row as usize].into(), + ), + calc.distance_with_binary_inner_product( + row as u32, + binary_inner_products[row as usize], + ) + .into(), + ); + if coarse_heap.len() < limit { + coarse_heap.push(node); + } else if coarse_heap.peek().is_some_and(|top| node.dist < top.dist) + && let Some(mut top) = coarse_heap.peek_mut() + { + *top = node; + } + } + Ok(coarse_heap) + }); + let (scored, next) = tokio::join!(scoring, prepared.try_next()); + coarse_heap = scored?; + pending = next?; + } + Ok(coarse_heap + .into_sorted_vec() + .into_iter() + .map( + |node| lance_index::vector::bq::layered::QuantizedCandidate { + index_uuid: self.uuid, + partition_id: node.id.0, + row_offset: node.id.1, + row_id: node.id.2, + distance: node.dist.0, + centroid_distance: distances[&node.id.0], + binary_inner_product: node.id.3.0, + }, + ) + .collect()) + } + + /// Rerank a global cut of quantized candidates with the original query, without data-file reads. + pub async fn rerank_quantized_candidates( + &self, + query: &Query, + input: Vec, + metrics: &dyn MetricsCollector, + ) -> Result { + use lance_index::vector::storage::DistanceCalculatorOptions; + let k = query + .k + .checked_mul(query.refine_factor.unwrap_or(1) as usize) + .ok_or_else(|| Error::invalid_input("candidate rerank k overflow"))?; + if k == 0 { + return Ok(RecordBatch::new_empty(VECTOR_RESULT_SCHEMA.clone())); + } + let raw_query = self.prepare_rq_raw_query_context(&query.key)?; + let mut candidates: std::collections::BTreeMap< + usize, + Vec, + > = std::collections::BTreeMap::new(); + let mut distances = HashMap::new(); + for candidate in input { + if candidate.index_uuid != self.uuid + || candidate.partition_id >= self.ivf.num_partitions() + { + return Err(Error::invalid_input( + "quantized candidate belongs to a different index or invalid partition", + )); + } + distances.insert(candidate.partition_id, candidate.centroid_distance); + candidates + .entry(candidate.partition_id) + .or_default() + .push(candidate); + } + let mut full_heap: BinaryHeap> = BinaryHeap::with_capacity(k); + let load_parallelism = self + .io_parallelism + .max(1) + .min(get_num_compute_intensive_cpus().max(1)); + // Match the dense scan's probe order, including centroid-distance ties. + // Statistical pruning must see the same evolving top-k threshold when + // the candidate cut retains every row. Physical partition order is not + // probe order. Recompute the small centroid ordering for public rerank + // callers, whose candidate list need not retain its original ordering. + let mut candidates: Vec<_> = candidates.into_iter().collect(); + if candidates.len() > 1 { + let (probes, _) = self.find_partitions(query)?; + let mut ranks = vec![usize::MAX; self.ivf.num_partitions()]; + for (rank, &partition) in probes.values().iter().enumerate() { + ranks[partition as usize] = rank; + } + candidates.sort_by_key(|(part_id, _)| ranks[*part_id]); + } + let mut prepared = stream::iter(candidates) + .map(|(part_id, mut rows)| async move { + rows.sort_unstable_by_key(|row| row.row_offset); + if rows + .windows(2) + .any(|pair| pair[0].row_offset == pair[1].row_offset) + { + return Err(Error::invalid_input("duplicate quantized candidate offset")); + } + let offsets = rows.iter().map(|row| row.row_offset).collect(); + let storage = self + .storage + .load_candidates(part_id, offsets, &self.index_cache, metrics.io_stats()) + .await?; + Result::Ok((part_id, rows, storage)) + }) + .buffered(load_parallelism); + let mut pending = prepared.try_next().await?; + while let Some((part_id, rows, storage)) = pending { + let dist_q_c = distances[&part_id]; + let key = query.key.clone(); + let cache = self.rq_search_cache.clone(); + let raw_query = raw_query.clone(); + let approx_mode = query.approx_mode; + let lower_bound = query.lower_bound; + let upper_bound = query.upper_bound; + let scoring = spawn_cpu(move || -> Result<_> { + let mut scratch = Vec::new(); + let context = QueryResidual::RabitRawQuery { + rotated_centroid: rotated_partition_centroid_slice(cache.as_deref(), part_id), + query: raw_query.as_deref(), + }; + let rq_storage = storage + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::internal("layered rerank storage is not RaBitQ"))?; + let calc = rq_storage.dist_calculator_with_scratch( + key, + dist_q_c, + Some(context), + &mut scratch, + DistanceCalculatorOptions { + approx_mode, + ..Default::default() + }, + ); + for (row, candidate) in rows.iter().enumerate() { + if storage.row_id(row as u32) != candidate.row_id { + return Err(Error::invalid_input( + "candidate row id does not match its physical offset", + )); + } + if let Some(bound) = calc.lower_bound_with_binary_inner_product( + row as u32, + candidate.binary_inner_product, + ) && (upper_bound.is_some_and(|upper| bound >= upper) + || (full_heap.len() == k + && full_heap.peek().is_some_and(|top| bound >= top.dist.0))) + { + continue; + } + let distance = calc.distance_with_binary_inner_product( + row as u32, + candidate.binary_inner_product, + ); + if lower_bound.is_some_and(|bound| distance < bound) + || upper_bound.is_some_and(|bound| distance >= bound) + { + continue; + } + let node = OrderedNode::new(storage.row_id(row as u32), distance.into()); + if full_heap.len() < k { + full_heap.push(node); + } else if full_heap.peek().is_some_and(|top| node.dist < top.dist) { + // Match dense top-k's replacement order for equal scores. + full_heap.pop(); + full_heap.push(node); + } + } + Ok(full_heap) + }); + let (scored, next) = tokio::join!(scoring, prepared.try_next()); + full_heap = scored?; + pending = next?; + } + let queued = Instant::now(); + let (batch, result_metrics) = spawn_cpu(move || -> Result<_> { + let result_metrics = LocalMetricsCollector::default(); + result_metrics.record_timing(IndexTiming::CpuQueueWait, queued.elapsed()); + let batch = Self::global_heap_to_batch(full_heap, &result_metrics)?; + Ok((batch, result_metrics)) + }) + .await?; + result_metrics.dump_into(metrics); + Ok(batch) + } + + async fn cascade_partitions( + &self, + query: &Query, + partitions: &UInt32Array, + centroid_dists: &Float32Array, + range: Range, + pre_filter: Arc, + metrics: &dyn MetricsCollector, + ) -> Result { + let candidates = self + .quantized_candidates( + query, + partitions, + centroid_dists, + range, + pre_filter, + metrics, + ) + .await?; + self.rerank_quantized_candidates(query, candidates, metrics) + .await + } + + async fn load_initial_partition( + &self, + partition_id: usize, + query: &Query, + metrics: &dyn MetricsCollector, + ) -> Result>> { + use lance_index::vector::bq::layered::{PlaneKey, RQPrecision}; + let mut precision = self.initial_precision(query)?; + if query.rq_cascade_factor.is_some() + && precision == RQPrecision::High + && self.storage.supports_candidate_reads() + && self + .index_cache + .get_resident_with_key(&PlaneKey { + partition: partition_id, + plane: 1, + }) + .await + .is_none() + { + precision = RQPrecision::Sign; + } + self.load_query_partition(partition_id, precision, metrics) + .await + } + + fn initial_precision( + &self, + query: &Query, + ) -> Result { + use lance_index::vector::bq::layered::RQPrecision; + if let Some(factor) = query.rq_cascade_factor { + if factor == 0 + || query.rq_precision != RQPrecision::Full + || !self.storage.is_layered_rq() + || query.approx_mode == lance_index::vector::ApproxMode::Fast + { + return Err(Error::invalid_input( + "rq_cascade_factor requires a positive factor, full precision and normal/accurate mode on a layered IVF_RQ index", + )); + } + if self.storage.supports_candidate_reads() + && query.lower_bound.is_none() + && query.upper_bound.is_none() + { + return Ok(RQPrecision::High); + } + } + Ok(query.rq_precision) + } + + async fn load_query_partition( + &self, + partition_id: usize, + precision: lance_index::vector::bq::layered::RQPrecision, + metrics: &dyn MetricsCollector, + ) -> Result>> { + if !self.storage.is_layered_rq() { + if precision != lance_index::vector::bq::layered::RQPrecision::Full { + return Err(Error::invalid_input( + "rq_precision requires a layered IVF_RQ index", + )); + } + return self.load_partition(partition_id, true, metrics).await; + } + if partition_id >= self.ivf.num_partitions() { + return Err(Error::invalid_input("partition id out of range")); + } + let mut entry = self + .load_partition_entry_at_precision(partition_id, metrics.io_stats(), Some(precision)) + .await?; + entry.cache_whole_partition = false; + Ok(Arc::new(entry)) + } + async fn load_partition_entry( &self, partition_id: usize, @@ -1779,6 +2185,34 @@ impl IVFIndex { Ok(PartitionEntry::new(idx, storage)) } + async fn load_partition_entry_at_precision( + &self, + partition_id: usize, + io_stats: Option, + precision: Option, + ) -> Result> { + let batch = self + .read_sub_index_batch( + partition_id, + self.read_projection.as_ref(), + io_stats.clone(), + ) + .await?; + let idx = S::load(batch)?; + let storage = if let Some(precision) = precision { + Box::pin(self.storage.load_partition_at_precision( + partition_id, + precision, + &self.index_cache, + io_stats, + )) + .await? + } else { + self.load_partition_storage(partition_id, io_stats).await? + }; + Ok(PartitionEntry::new(idx, storage)) + } + /// Read the serialized sub-index of a partition, with its metadata attached. /// /// `projection` narrows the read to the columns search needs; `None` reads @@ -2029,6 +2463,7 @@ impl IVFIndex { index_file_size: self.reader.metadata().file_size(), aux_file_size: self.storage.reader().metadata().file_size(), rq_search_cache: rabit_search_cache_cell(self.rq_search_cache.clone()), + plane_access: self.storage.plane_access_tracker().clone(), })) } } @@ -2044,6 +2479,9 @@ impl Index for IVFIndex Result<()> { + if self.storage.is_layered_rq() { + return self.storage.prewarm_planes(&self.index_cache).await; + } let cpu_parallelism = get_num_compute_intensive_cpus(); let target_bytes = prewarm_window_size_bytes()?; let parallelism = prewarm_parallelism(self.io_parallelism, cpu_parallelism); @@ -2219,10 +2657,29 @@ impl VectorIndex for IVFInd pre_filter: Arc, metrics: &dyn MetricsCollector, ) -> Result { + self.initial_precision(query)?; + if query.rq_cascade_factor.is_some() + && self.storage.supports_candidate_reads() + && query.lower_bound.is_none() + && query.upper_bound.is_none() + { + return self + .cascade_partitions( + query, + &UInt32Array::from(vec![partition_id as u32]), + &Float32Array::from(vec![query.dist_q_c]), + 0..1, + pre_filter, + metrics, + ) + .await; + } let (part_entry, pre_filter) = { // Match split preparation without counting CPU queueing or search. let _prepare_timer = IndexTimer::new(metrics, IndexTiming::PartitionPrepare); - let part_entry = self.load_partition(partition_id, true, metrics).await?; + let part_entry = self + .load_initial_partition(partition_id, query, metrics) + .await?; { let _wait_timer = IndexTimer::new(metrics, IndexTiming::PrefilterWait); pre_filter.wait_for_ready().await?; @@ -2361,6 +2818,31 @@ impl VectorIndex for IVFInd ))); } + self.initial_precision(&query)?; + if query.rq_cascade_factor.is_some() + && self.storage.supports_candidate_reads() + && query.lower_bound.is_none() + && query.upper_bound.is_none() + { + let batch = self + .cascade_partitions( + &query, + &partitions, + &q_c_dists, + start_idx..end_idx, + pre_filter, + metrics.as_ref(), + ) + .await?; + if let Some(control) = control { + control.record_batch(&batch); + } + return Ok(Box::pin(RecordBatchStreamAdapter::new( + VECTOR_RESULT_SCHEMA.clone(), + stream::once(async { Ok(batch) }), + ))); + } + let prepare_parallelism = get_num_compute_intensive_cpus().max(1); let raw_query_context = self.prepare_rq_raw_query_context(&query.key)?; @@ -2692,6 +3174,30 @@ impl VectorIndex for IVFInd ))); } let dim = query.key.len() / query_count; + self.initial_precision(&query)?; + if query.rq_cascade_factor.is_some() + && self.storage.supports_candidate_reads() + && query.lower_bound.is_none() + && query.upper_bound.is_none() + { + let mut results = Vec::with_capacity(query_count); + for query_idx in 0..query_count { + let mut single = query.clone(); + single.key = query.key.slice(query_idx * dim, dim); + results.push( + self.cascade_partitions( + &single, + &partitions_per_query[query_idx], + &q_c_dists_per_query[query_idx], + 0..partitions_per_query[query_idx].len(), + pre_filter.clone(), + metrics.as_ref(), + ) + .await?, + ); + } + return Ok(results); + } // Per-query immutable search state: the query vector slice and the // optional Rabit raw-query context both depend only on the query vector, @@ -2759,13 +3265,14 @@ impl VectorIndex for IVFInd let load_parallelism = get_num_compute_intensive_cpus().max(1); let load_index = self.clone(); let load_metrics = metrics.clone(); + let precision = query.rq_precision; let mut loaded_chunks = stream::iter(assignment_list) .map(move |(part_id, probing_queries)| { let index = load_index.clone(); let metrics = load_metrics.clone(); async move { let part_entry = index - .load_partition(part_id as usize, true, metrics.as_ref()) + .load_query_partition(part_id as usize, precision, metrics.as_ref()) .await?; Result::Ok((part_id as usize, part_entry, probing_queries)) } @@ -3021,6 +3528,7 @@ async fn reconstruct_typed( state.distance_type, legacy, ) + .with_plane_access_tracker(state.plane_access.clone()) }; let storage = match frag_reuse_index { Some(ResolvedRemapping::V0(remapper)) => make_storage(Some(remapper)), @@ -3072,6 +3580,7 @@ mod tests { use arrow_schema::{DataType, Field, Schema, SchemaRef}; use itertools::Itertools; use lance_arrow::FixedSizeListArrayExt; + use lance_index::vector::bq::builder::RabitQuantizer; use lance_index::vector::bq::{ RQBuildParams, RQRotationType, ex_dot::{blocked_ex_code_bytes, padded_query_len}, @@ -3079,7 +3588,7 @@ mod tests { transform::{EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN}, }; use lance_index::vector::ivf::storage::IvfModel; - use lance_index::vector::storage::VectorStore; + use lance_index::vector::storage::{DistCalculator, IvfQuantizationStorage, VectorStore}; use lance_index::vector::v3::subindex::IvfSubIndex; use crate::dataset::{InsertBuilder, UpdateBuilder, WriteMode, WriteParams}; @@ -6629,6 +7138,601 @@ mod tests { test_recall::(params, 4, 0.5, "vector", &dataset, vectors).await; } + #[tokio::test] + async fn test_layered_plane_promotion_survives_query_reconstruction() { + use lance_index::vector::bq::layered::PlaneKey; + + let dir = TempStrDir::default(); + let (mut dataset, vectors) = + generate_test_dataset::(dir.as_str(), 0.0..1.0).await; + 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(); + let cold = crate::DatasetBuilder::from_uri(dir.as_str()) + .with_session(Arc::new(crate::session::Session::new( + 64 * 1024 * 1024, + 1024 * 1024, + Arc::new(lance_io::object_store::ObjectStoreRegistry::default()), + ))) + .load() + .await + .unwrap(); + let index_id = cold.load_indices().await.unwrap()[0].uuid; + let cache = cold.index_cache.for_index(&index_id, None); + let query = vectors.value(0); + let num_rows = cold.count_rows(None).await.unwrap(); + for pass in 0..3 { + let result = cold + .scan() + .nearest("vector", query.as_ref(), num_rows) + .unwrap() + .nprobes(4) + .rq_cascade_factor(1) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), num_rows); + for partition in 0..4 { + for plane in [1, 2] { + let resident = cache + .get_resident_with_key(&PlaneKey { partition, plane }) + .await; + assert_eq!( + resident.is_some(), + pass == 2, + "pass={pass}, partition={partition}, plane={plane}" + ); + } + } + } + } + + /// Prewarm and full-precision loads of a layered index over a two-tier + /// cache, with and without sign-gated plane admission. + mod layered_plane_gating { + use super::*; + + use bytes::Bytes; + use lance_core::cache::{CacheCodec, CacheEntry, InternalCacheKey, QuickCacheBackend}; + use lance_index::vector::ApproxMode; + use lance_index::vector::bq::layered::{PlaneBatch, RQPrecision}; + use lance_io::assert_io_eq; + use lance_io::object_store::ObjectStoreRegistry; + + use crate::session::Session; + + const PARTITIONS: usize = 8; + const PLANES: usize = 3; + const BITS: u8 = 7; + const K: usize = 10; + /// Holds every entry, so the planes' total size can be measured. + const LARGE_RAM_BYTES: usize = 256 * 1024 * 1024; + const METADATA_CACHE_BYTES: usize = 64 * 1024 * 1024; + /// The small RAM tier holds 1 / this of the plane bytes. + const SMALL_RAM_DIVISOR: usize = 10; + + type DiskEntries = HashMap; + + /// RAM tier plus an unbounded serialized "disk" tier, like a tiered + /// cache: loads are written through and disk hits are admitted to + /// RAM. `gated` is what the backend reports for plane admission. + struct PlaneTierTestBackend { + ram: QuickCacheBackend, + disk: std::sync::Mutex, + gated: bool, + } + + impl std::fmt::Debug for PlaneTierTestBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PlaneTierTestBackend") + .field("gated", &self.gated) + .finish_non_exhaustive() + } + } + + impl PlaneTierTestBackend { + fn new(ram_bytes: usize, gated: bool) -> Self { + Self { + ram: QuickCacheBackend::with_capacity(ram_bytes), + disk: Default::default(), + gated, + } + } + + 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)) + } + + /// Persisted plane entries and their accounted bytes. + fn persisted_planes(&self) -> (usize, usize) { + self.disk + .lock() + .unwrap() + .values() + .filter(|(_, codec, _)| { + codec.type_id() == ::TYPE_ID + }) + .fold((0, 0), |(entries, bytes), (_, _, size)| { + (entries + 1, bytes + size) + }) + } + } + + #[async_trait::async_trait] + impl CacheBackend for PlaneTierTestBackend { + async fn get_resident(&self, key: &InternalCacheKey) -> Option { + self.ram.get_resident(key).await + } + + 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( + &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, + ) { + if let Some(codec) = codec { + let mut bytes = Vec::new(); + codec.serialize(&entry, &mut bytes).unwrap(); + self.disk + .lock() + .unwrap() + .insert(*key, (Bytes::from(bytes), codec, size_bytes)); + } + 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)> { + if let Some(entry) = self.get(key, codec).await { + return Ok((entry, true)); + } + let (entry, size) = loader.await?; + self.insert(key, entry.clone(), size, codec).await; + Ok((entry, false)) + } + + 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 + } + } + + async fn open_prewarmed( + uri: &str, + backend: Arc, + ) -> (Dataset, Arc) { + let session = Session::with_index_cache_backend( + backend, + METADATA_CACHE_BYTES, + Arc::new(ObjectStoreRegistry::default()), + ); + 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(); + index.prewarm().await.unwrap(); + (dataset, index) + } + + fn full_query(key: ArrayRef) -> Query { + Query { + search_effort: 0.5, + column: "vector".to_string(), + key, + k: K, + lower_bound: None, + upper_bound: None, + minimum_nprobes: PARTITIONS, + maximum_nprobes: Some(PARTITIONS), + 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, + rq_precision: RQPrecision::Full, + rq_cascade_factor: None, + } + } + + /// Row ids and distance bits of every probed partition, in result order. + async fn search_all_partitions( + index: &Arc, + query: &Query, + ) -> (Vec, Vec) { + let (partitions, dists) = index.find_partitions(query).unwrap(); + let probes = partitions.len(); + let batches = index + .clone() + .search_partitions( + query.clone(), + Arc::new(partitions), + Arc::new(dists), + 0, + probes, + Arc::new(NoFilter), + None, + Arc::new(NoOpMetricsCollector), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut row_ids = Vec::new(); + let mut distance_bits = Vec::new(); + for batch in &batches { + row_ids.extend(batch[ROW_ID].as_primitive::().values()); + distance_bits.extend( + batch[DIST_COL] + .as_primitive::() + .values() + .iter() + .map(|dist| dist.to_bits()), + ); + } + (row_ids, distance_bits) + } + + /// A backend that admits planes like any entry is warmed partition by + /// partition, so every plane is persisted although RAM holds only a + /// fraction of them, and warm full-precision queries read nothing from + /// the index file. A sign-gated backend of the same size persists only + /// the lower planes of partitions whose sign plane stayed resident. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_ungated_prewarm_persists_every_plane() { + let dir = TempStrDir::default(); + let (mut dataset, vectors) = + generate_test_dataset::(dir.as_str(), 0.0..1.0).await; + let params = VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::new(PARTITIONS), + RQBuildParams::new(BITS).with_layered(true), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + let query = full_query(vectors.value(0)); + + let resident = Arc::new(PlaneTierTestBackend::new(LARGE_RAM_BYTES, false)); + let (_, index) = open_prewarmed(dir.as_str(), resident.clone()).await; + let (entries, plane_bytes) = resident.persisted_planes(); + assert_eq!(entries, PLANES * PARTITIONS); + let expected = search_all_partitions(&index, &query).await; + assert_eq!(expected.0.len(), K); + + let small_ram = plane_bytes / SMALL_RAM_DIVISOR; + let ungated = Arc::new(PlaneTierTestBackend::new(small_ram, false)); + let (dataset, index) = open_prewarmed(dir.as_str(), ungated.clone()).await; + assert_eq!(ungated.persisted_planes(), (entries, plane_bytes)); + assert!(ungated.ram.size_bytes().await <= small_ram); + assert_eq!(search_all_partitions(&index, &query).await, expected); + dataset.object_store.as_ref().io_stats_incremental(); + for _ in 0..2 { + assert_eq!(search_all_partitions(&index, &query).await, expected); + } + let io = dataset.object_store.as_ref().io_stats_incremental(); + assert_io_eq!(io, read_iops, 0, "warm layered queries read no plane"); + + let gated = Arc::new(PlaneTierTestBackend::new(small_ram, true)); + let (_, index) = open_prewarmed(dir.as_str(), gated.clone()).await; + let (gated_entries, _) = gated.persisted_planes(); + assert!( + (PARTITIONS..PLANES * PARTITIONS).contains(&gated_entries), + "gated prewarm persisted {gated_entries} plane entries" + ); + assert_eq!(search_all_partitions(&index, &query).await, expected); + } + } + + #[tokio::test] + async fn test_layered_rq_batch_remapping_preserves_precision() { + use lance_index::scalar::BatchRowIdRemapper; + use lance_index::vector::bq::layered::RQPrecision; + + const ROW_ID_OFFSET: u64 = 1 << 40; + #[derive(Debug)] + struct TestBatchRemapper; + + #[async_trait::async_trait] + impl BatchRowIdRemapper for TestBatchRemapper { + async fn remap_row_ids(&self, row_ids: &[u64]) -> lance_core::Result>> { + Ok(row_ids + .iter() + .map(|&id| (!id.is_multiple_of(2)).then_some(id + ROW_ID_OFFSET)) + .collect()) + } + } + + let dir = TempStrDir::default(); + let (mut dataset, vectors) = + generate_test_dataset::(dir.as_str(), 0.0..1.0).await; + 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(); + let index = &dataset.load_indices().await.unwrap()[0]; + let scheduler = ScanScheduler::new( + Arc::new(ObjectStore::local()), + SchedulerConfig::default_for_testing(), + ); + let reader = open_rq_aux_reader(&dataset, scheduler.clone(), &index.uuid.to_string()).await; + let loader = IvfQuantizationStorage::::try_new(reader, None) + .await + .unwrap(); + let reader = open_rq_aux_reader(&dataset, scheduler, &index.uuid.to_string()).await; + let remapped_loader = IvfQuantizationStorage::::try_new(reader, None) + .await + .unwrap() + .with_row_id_remapping(Arc::new(TestBatchRemapper)); + assert!(!remapped_loader.supports_candidate_reads()); + let cache = LanceCache::with_capacity(16 * 1024 * 1024); + let cache = WeakLanceCache::from(&cache); + assert!( + remapped_loader + .load_candidates(0, vec![0], &cache, None) + .await + .is_err() + ); + let query = vectors.value(0); + for precision in [RQPrecision::Sign, RQPrecision::High, RQPrecision::Full] { + let original = loader + .load_partition_at_precision(0, precision, &cache, None) + .await + .unwrap(); + let original_calc = original.dist_calculator(query.clone(), 0.0); + let expected = (0..original.len() as u32) + .filter(|&offset| !original.row_id(offset).is_multiple_of(2)) + .map(|offset| { + ( + original.row_id(offset) + ROW_ID_OFFSET, + original_calc.distance(offset).to_bits(), + ) + }) + .collect::>(); + assert!(!expected.is_empty()); + assert!(expected.len() < original.len()); + let remapped = remapped_loader + .load_partition_at_precision(0, precision, &cache, None) + .await + .unwrap(); + let remapped_calc = remapped.dist_calculator(query.clone(), 0.0); + let actual = (0..remapped.len() as u32) + .map(|offset| { + ( + remapped.row_id(offset), + remapped_calc.distance(offset).to_bits(), + ) + }) + .collect::>(); + assert_eq!(actual, expected, "precision={precision:?}"); + // Cached planes stay in their source address space. + let reloaded = loader + .load_partition_at_precision(0, precision, &cache, None) + .await + .unwrap(); + assert_eq!( + reloaded.row_ids().copied().collect::>(), + original.row_ids().copied().collect::>(), + ); + } + } + + #[rstest] + #[case::rq3(3)] + #[case::rq4(4)] + #[case::rq5(5)] + #[case::rq6(6)] + #[case::rq7(7)] + #[case::rq8(8)] + #[case::rq9(9)] + #[tokio::test] + async fn test_layered_rq_create_search_cascade_and_remap( + #[case] bits: u8, + #[values(DistanceType::L2, DistanceType::Cosine, DistanceType::Dot)] + distance_type: DistanceType, + ) { + use lance_index::vector::bq::layered::RQPrecision; + let dir = TempStrDir::default(); + let (mut dataset, vectors) = + generate_test_dataset::(dir.as_str(), 0.0..1.0).await; + let params = VectorIndexParams::with_ivf_rq_params( + distance_type, + IvfBuildParams::new(4), + RQBuildParams::new(bits).with_layered(true), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + let query = vectors.value(0); + let gt = ground_truth(&dataset, "vector", &query, 100, distance_type).await; + let mut full = None; + for precision in [RQPrecision::Sign, RQPrecision::High, RQPrecision::Full] { + let result = dataset + .scan() + .nearest("vector", query.as_ref(), 100) + .unwrap() + .nprobes(4) + .rq_precision(precision) + .with_row_id() + .try_into_batch() + .await + .unwrap(); + let ids = result[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + assert!( + ids.intersection(>).count() >= 50, + "precision={precision:?}" + ); + if precision == RQPrecision::Full { + full = Some(result); + } + } + let cold = crate::DatasetBuilder::from_uri(dir.as_str()) + .with_session(Arc::new(crate::session::Session::new( + 0, + 1024 * 1024, + Arc::new(lance_io::object_store::ObjectStoreRegistry::default()), + ))) + .load() + .await + .unwrap(); + let cascade = cold + .scan() + .nearest("vector", query.as_ref(), 100) + .unwrap() + .nprobes(4) + .rq_cascade_factor(100) + .unwrap() + .with_row_id() + .try_into_batch() + .await + .unwrap(); + let full = full.unwrap(); + assert_eq!(cascade[ROW_ID].as_ref(), full[ROW_ID].as_ref()); + assert_eq!(cascade[DIST_COL].as_ref(), full[DIST_COL].as_ref()); + let upper = full[DIST_COL].as_primitive::().value(50); + let mut range_results = Vec::new(); + for cascade_factor in [None, Some(4)] { + let mut scan = cold.scan(); + scan.nearest("vector", query.as_ref(), 100) + .unwrap() + .nprobes(4) + .distance_range(None, Some(upper)) + .with_row_id(); + if let Some(factor) = cascade_factor { + scan.rq_cascade_factor(factor).unwrap(); + } + range_results.push(scan.try_into_batch().await.unwrap()); + } + assert_eq!( + range_results[0][ROW_ID].as_ref(), + range_results[1][ROW_ID].as_ref() + ); + assert_eq!( + range_results[0][DIST_COL].as_ref(), + range_results[1][DIST_COL].as_ref() + ); + + let index = &dataset.load_indices().await.unwrap()[0]; + let scheduler = ScanScheduler::new( + Arc::new(ObjectStore::local()), + SchedulerConfig::default_for_testing(), + ); + let reader = open_rq_aux_reader(&dataset, scheduler, &index.uuid.to_string()).await; + let loader = IvfQuantizationStorage::::try_new(reader, None) + .await + .unwrap(); + let source = loader.load_partition(0, None).await.unwrap(); + let offsets: Vec = (0..source.len() as u32).step_by(3).take(5).collect(); + assert!(offsets.len() > 1); + let no_cache = LanceCache::no_cache(); + let selected = loader + .load_candidates(0, offsets.clone(), &WeakLanceCache::from(&no_cache), None) + .await + .unwrap(); + let source_calc = source.dist_calculator(query.clone(), 0.0); + let selected_calc = selected.dist_calculator(query.clone(), 0.0); + for (row, offset) in offsets.into_iter().enumerate() { + assert_eq!(selected.row_id(row as u32), source.row_id(offset)); + assert_eq!( + selected_calc.distance(row as u32).to_bits(), + source_calc.distance(offset).to_bits() + ); + } + append_dataset::(&mut dataset, 64, 0.0..1.0).await; + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + dataset + .optimize_indices(&OptimizeOptions::merge(10)) + .await + .unwrap(); + let indices = dataset.load_indices().await.unwrap(); + let scheduler = ScanScheduler::new( + Arc::new(ObjectStore::local()), + SchedulerConfig::default_for_testing(), + ); + for index in indices.iter() { + assert!( + get_rq_metadata(&dataset, scheduler.clone(), &index.uuid.to_string()) + .await + .layered + ); + } + test_remap(params, 4, 0.5).await; + } + #[rstest] #[case::fast(RQRotationType::Fast)] #[case::matrix(RQRotationType::Matrix)] @@ -7582,6 +8686,8 @@ mod tests { .prepared_partitions(); let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key: query_vector, k: K, diff --git a/rust/lance/src/index/vector/pq.rs b/rust/lance/src/index/vector/pq.rs index d65d1e32853..29acd75a583 100644 --- a/rust/lance/src/index/vector/pq.rs +++ b/rust/lance/src/index/vector/pq.rs @@ -1035,6 +1035,8 @@ mod tests { frag_reuse_index: None, }; let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key: Arc::new(Float32Array::from(vec![0.0])) as ArrayRef, k: 2, diff --git a/rust/lance/src/io/exec/ann_proto.rs b/rust/lance/src/io/exec/ann_proto.rs index 652f17168e1..19220b33c1a 100644 --- a/rust/lance/src/io/exec/ann_proto.rs +++ b/rust/lance/src/io/exec/ann_proto.rs @@ -120,6 +120,13 @@ pub fn query_to_proto(query: &Query) -> Result { dist_q_c: Some(query.dist_q_c), query_parallelism: Some(query.query_parallelism), approx_mode: approx_mode_to_proto(query.approx_mode) as i32, + rq_cascade_factor: query.rq_cascade_factor, + rq_precision: match query.rq_precision { + lance_index::vector::bq::layered::RQPrecision::Full => "full", + lance_index::vector::bq::layered::RQPrecision::High => "high", + lance_index::vector::bq::layered::RQPrecision::Sign => "sign", + } + .to_string(), // No planner narrows the covering projection yet, so this is always absent: // "materialize every covering column declared". See `CoveringProjection`. covering_projection: None, @@ -139,6 +146,7 @@ pub fn query_from_proto(proto: pb::VectorQueryProto) -> Result { .transpose()?; let query = Query { + rq_cascade_factor: proto.rq_cascade_factor, column: proto.column, key, k: proto.k as usize, @@ -154,6 +162,11 @@ pub fn query_from_proto(proto: pb::VectorQueryProto) -> Result { query_parallelism: proto.query_parallelism.unwrap_or(DEFAULT_QUERY_PARALLELISM), dist_q_c: proto.dist_q_c.unwrap_or(0.0), approx_mode: approx_mode_from_proto(proto.approx_mode), + rq_precision: if proto.rq_precision.is_empty() { + Default::default() + } else { + proto.rq_precision.parse()? + }, }; query.validate_search_effort()?; Ok(query) @@ -345,9 +358,14 @@ mod tests { } #[rstest::rstest] - fn test_query_roundtrip(#[values(0.0, 0.25, 0.5, 0.75, 1.0)] effort: f64) { + fn test_query_roundtrip( + #[values(0.0, 0.25, 0.5, 0.75, 1.0)] effort: f64, + #[values("sign", "high", "full")] precision: &str, + ) { let key: ArrayRef = Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3])); let query = Query { + rq_cascade_factor: (precision == "full").then_some(8), + rq_precision: precision.parse().unwrap(), column: "vector".to_string(), key, k: 10, @@ -367,7 +385,10 @@ mod tests { let proto = query_to_proto(&query).unwrap(); assert_eq!(proto.search_effort, (effort != 0.5).then_some(effort)); - let back = query_from_proto(proto).unwrap(); + // Exercise field numbers as well as the in-memory conversion. + let bytes = prost::Message::encode_to_vec(&proto); + let decoded = ::decode(bytes.as_slice()).unwrap(); + let back = query_from_proto(decoded).unwrap(); assert_eq!(query.column, back.column); assert_eq!(query.k, back.k); @@ -376,6 +397,8 @@ mod tests { assert_eq!(query.minimum_nprobes, back.minimum_nprobes); assert_eq!(query.maximum_nprobes, back.maximum_nprobes); assert_eq!(query.search_effort, back.search_effort); + assert_eq!(query.rq_precision, back.rq_precision); + assert_eq!(query.rq_cascade_factor, back.rq_cascade_factor); assert_eq!(query.ef, back.ef); assert_eq!(query.refine_factor, back.refine_factor); assert_eq!(query.metric_type, back.metric_type); @@ -391,6 +414,8 @@ mod tests { fn test_query_roundtrip_none_metric() { let key: ArrayRef = Arc::new(Float32Array::from(vec![1.0])); let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "v".to_string(), key, k: 5, @@ -485,6 +510,8 @@ mod tests { let key: ArrayRef = Arc::new(Float32Array::from(vec![0.1f32; 128])); let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key, k: 10, @@ -534,6 +561,8 @@ mod tests { let key: ArrayRef = Arc::new(Float32Array::from(vec![0.1f32; 128])); let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key, k: 10, @@ -599,6 +628,8 @@ mod tests { let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); let key: ArrayRef = Arc::new(Float32Array::from(vec![0.1f32; 128])); let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key, k: 10, diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index eb4da119099..a719717a7c6 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -1647,25 +1647,38 @@ impl DisplayAs for ANNIvfSubIndexExec { .metric_type .map(|m| format!("{:?}", m)) .unwrap_or_else(|| "default".to_string()); + let precision = if self.query.rq_precision + != lance_index::vector::bq::layered::RQPrecision::Full + || self.query.rq_cascade_factor.is_some() + { + format!( + ", rq_precision={:?}, rq_cascade_factor={:?}", + self.query.rq_precision, self.query.rq_cascade_factor + ) + } else { + String::new() + }; match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!( f, - "ANNSubIndex: name={}, k={}, deltas={}, metric={}", + "ANNSubIndex: name={}, k={}, deltas={}, metric={}{}", self.indices[0].name, self.query.k * self.query.refine_factor.unwrap_or(1) as usize, self.indices.len(), - metric_str + metric_str, + precision ) } DisplayFormatType::TreeRender => { write!( f, - "ANNSubIndex\nname={}\nk={}\ndeltas={}\nmetric={}", + "ANNSubIndex\nname={}\nk={}\ndeltas={}\nmetric={}{}", self.indices[0].name, self.query.k * self.query.refine_factor.unwrap_or(1) as usize, self.indices.len(), - metric_str + metric_str, + precision ) } } @@ -3105,6 +3118,8 @@ mod tests { fn base_query() -> Query { Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vec".to_string(), key: Arc::new(Float32Array::from(vec![0.0f32])) as ArrayRef, k: 10, @@ -4741,6 +4756,8 @@ mod tests { #[tokio::test] async fn test_multivector_score() { let query = Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key: Arc::new(generate_random_array(1)), k: 10, diff --git a/rust/lance/src/io/exec/knn/adaptive_probe.rs b/rust/lance/src/io/exec/knn/adaptive_probe.rs index 92e6fde26a1..e4e47700b54 100644 --- a/rust/lance/src/io/exec/knn/adaptive_probe.rs +++ b/rust/lance/src/io/exec/knn/adaptive_probe.rs @@ -339,6 +339,8 @@ mod tests { fn query() -> Query { Query { + rq_cascade_factor: None, + rq_precision: Default::default(), column: "vector".to_string(), key: Arc::new(Float32Array::from(vec![1.0, 0.0])), k: 10,