diff --git a/diskann-disk/src/search/pq/mod.rs b/diskann-disk/src/search/pq/mod.rs index ef9ffa1d5d..6d22373daf 100644 --- a/diskann-disk/src/search/pq/mod.rs +++ b/diskann-disk/src/search/pq/mod.rs @@ -6,9 +6,11 @@ //! Product quantization types and functions used during disk-based search. mod pq_scratch; -pub use pq_scratch::PQScratch; +pub(crate) use pq_scratch::{PQBatchScratch, PQQueryComputerArgs, PQQueryComputerStorage}; +pub use pq_scratch::{PQQueryComputer, PQScratch}; pub(crate) use crate::storage::quant::pq::PQData; mod quantizer_preprocess; +pub(crate) use quantizer_preprocess::prepare_query; pub use quantizer_preprocess::quantizer_preprocess; diff --git a/diskann-disk/src/search/pq/pq_scratch.rs b/diskann-disk/src/search/pq/pq_scratch.rs index 54408c0e88..726584c2d4 100644 --- a/diskann-disk/src/search/pq/pq_scratch.rs +++ b/diskann-disk/src/search/pq/pq_scratch.rs @@ -7,14 +7,143 @@ use diskann::ANNResult; use diskann_quantization::alloc::{AlignedAllocator, Poly}; +use diskann_utils::object_pool::{ObjectPool, PoolOption, TryAsPooled}; +use diskann_vector::PreprocessedDistanceFunction; +use std::sync::Arc; use crate::error::{diskann_error, ErrorKind}; +#[derive(Clone, Copy, Debug)] +pub(crate) struct PQQueryComputerArgs { + dim: usize, + num_pq_chunks: usize, + num_centers: usize, +} + +impl PQQueryComputerArgs { + pub(crate) fn new(dim: usize, num_pq_chunks: usize, num_centers: usize) -> Self { + Self { + dim, + num_pq_chunks, + num_centers, + } + } +} + +#[derive(Debug)] +pub(crate) struct PQQueryComputerStorage { + aligned_pqtable_dist_scratch: Poly<[f32], AlignedAllocator>, + query_scratch: Vec, + num_pq_chunks: usize, + num_centers: usize, +} + +impl TryAsPooled for PQQueryComputerStorage { + type Error = diskann::ANNError; + + fn try_create(args: PQQueryComputerArgs) -> Result { + let aligned_pqtable_dist_scratch = Poly::broadcast( + 0f32, + args.num_centers * args.num_pq_chunks, + AlignedAllocator::A128, + ) + .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; + + Ok(Self { + aligned_pqtable_dist_scratch, + query_scratch: vec![0.0; args.dim], + num_pq_chunks: args.num_pq_chunks, + num_centers: args.num_centers, + }) + } + + fn try_modify(&mut self, args: PQQueryComputerArgs) -> Result<(), Self::Error> { + if self.query_scratch.len() != args.dim + || self.num_pq_chunks != args.num_pq_chunks + || self.num_centers != args.num_centers + { + *self = Self::try_create(args)?; + } + Ok(()) + } +} + +/// Opaque preprocessed query state created internally by disk search strategies. +#[derive(Debug)] +pub struct PQQueryComputer { + storage: PoolOption, +} + +impl PQQueryComputer { + /// Create an empty query computer for the given PQ schema. + #[cfg(test)] + pub(crate) fn new(dim: usize, num_pq_chunks: usize, num_centers: usize) -> ANNResult { + Ok(Self { + storage: PoolOption::try_non_pooled_create(PQQueryComputerArgs::new( + dim, + num_pq_chunks, + num_centers, + ))?, + }) + } + + pub(crate) fn pooled( + pool: &Arc>, + args: PQQueryComputerArgs, + ) -> ANNResult { + Ok(Self { + storage: PoolOption::try_pooled(pool, args)?, + }) + } + + /// Copy a full-precision query into the preprocessing buffer. + pub(crate) fn set(&mut self, query: &[f32]) -> ANNResult<()> { + let dim = self.storage.query_scratch.len(); + if query.len() != dim { + return Err(diskann_error!( + ErrorKind::DimensionMismatchError, + "PQQueryComputer::set: expected query of length {dim}, got {}", + query.len() + )); + } + self.storage.query_scratch.copy_from_slice(query); + Ok(()) + } + + pub(crate) fn lookup_table(&self) -> &[f32] { + &self.storage.aligned_pqtable_dist_scratch + } + + pub(super) fn preprocessing_buffers(&mut self) -> (&[f32], &mut [f32]) { + let storage = &mut *self.storage; + ( + &storage.query_scratch, + &mut storage.aligned_pqtable_dist_scratch, + ) + } +} + +impl PreprocessedDistanceFunction<&[u8], f32> for PQQueryComputer { + fn evaluate_similarity(&self, code: &[u8]) -> f32 { + assert_eq!( + code.len(), + self.storage.num_pq_chunks, + "PQ code has the wrong number of chunks", + ); + code.iter() + .enumerate() + .map(|(chunk, ¢er)| { + self.storage.aligned_pqtable_dist_scratch + [chunk * self.storage.num_centers + center as usize] + }) + .sum() + } +} + #[derive(Debug)] /// PQ scratch pub struct PQScratch { - /// Aligned pq table distance scratch, the length must be at least [256 * NCHUNKS]. 256 is the number of PQ centroids. - /// This is used to store the distance between each chunk in the query vector to each centroid, which is why the length is num of centroids * num of chunks + /// Aligned PQ table distance scratch. pub aligned_pqtable_dist_scratch: Poly<[f32], AlignedAllocator>, /// Aligned dist scratch, must be at least diskann MAX_DEGREE @@ -25,9 +154,7 @@ pub struct PQScratch { /// This is used to store the pq coordinates of the candidate vectors. pub aligned_pq_coord_scratch: Poly<[u8], AlignedAllocator>, - /// Query scratch buffer stored as `f32`, sized by the PQ table's logical dimension. - /// `set` populates it from a caller-provided `&[f32]`; `PQTable::preprocess_query` can - /// then rotate or otherwise preprocess it. + /// Query scratch buffer stored as `f32`. pub query_scratch: Vec, } @@ -45,18 +172,17 @@ impl PQScratch { let aligned_pq_coord_scratch = Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128) .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; + let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128) + .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; let aligned_pqtable_dist_scratch = Poly::broadcast(0f32, num_centers * num_pq_chunks, AlignedAllocator::A128) .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; - let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128) - .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; - let query_scratch = vec![0.0f32; dim]; Ok(Self { aligned_pqtable_dist_scratch, aligned_dist_scratch, aligned_pq_coord_scratch, - query_scratch, + query_scratch: vec![0.0; dim], }) } @@ -80,8 +206,32 @@ impl PQScratch { Ok(()) } - /// Return the largest number of PQ vectors whose distances can be computed using this - /// scratch data structure. + /// Return the largest number of PQ vectors that fit in the batch scratch. + #[cfg(test)] + pub(crate) fn max_vectors(&self) -> usize { + self.aligned_dist_scratch.len() + } +} + +#[derive(Debug)] +pub(crate) struct PQBatchScratch { + pub(crate) aligned_dist_scratch: Poly<[f32], AlignedAllocator>, + pub(crate) aligned_pq_coord_scratch: Poly<[u8], AlignedAllocator>, +} + +impl PQBatchScratch { + pub(crate) fn new(graph_degree: usize, num_pq_chunks: usize) -> ANNResult { + let aligned_pq_coord_scratch = + Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128) + .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; + let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128) + .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?; + Ok(Self { + aligned_dist_scratch, + aligned_pq_coord_scratch, + }) + } + pub(crate) fn max_vectors(&self) -> usize { self.aligned_dist_scratch.len() } @@ -89,13 +239,53 @@ impl PQScratch { #[cfg(test)] mod tests { + use std::sync::Arc; + use diskann_quantization::num::PowerOfTwo; + use diskann_utils::object_pool::ObjectPool; + use diskann_vector::PreprocessedDistanceFunction; use rstest::rstest; - use super::PQScratch; + use super::{PQQueryComputer, PQQueryComputerArgs, PQQueryComputerStorage, PQScratch}; use crate::error::{error_kind, ErrorKind}; + #[test] + fn query_computer_scores_pq_code() { + let mut computer = PQQueryComputer::new(2, 2, 3).unwrap(); + computer + .preprocessing_buffers() + .1 + .copy_from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + assert_eq!(computer.evaluate_similarity(&[1, 2]), 6.0); + } + + #[test] + fn query_computer_reuses_pooled_storage() { + let args = PQQueryComputerArgs::new(2, 2, 3); + let pool = Arc::new(ObjectPool::::try_new(args, 0, None).unwrap()); + assert!(pool.is_empty()); + + let first_ptr; + { + let computer = PQQueryComputer::pooled(&pool, args).unwrap(); + assert!(pool.is_empty()); + first_ptr = computer.lookup_table().as_ptr(); + } + assert_eq!(pool.len(), 1); + + let second_ptr; + { + let computer = PQQueryComputer::pooled(&pool, args).unwrap(); + assert!(pool.is_empty()); + second_ptr = computer.lookup_table().as_ptr(); + } + assert_eq!(pool.len(), 1); + + assert_eq!(first_ptr, second_ptr); + } + #[rstest] #[case(512, 8, 128, 256)] // default test case #[case(59, 16, 37, 41)] // not multiple of 256 @@ -120,7 +310,6 @@ mod tests { (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(), 0 ); - assert_eq!(pq_scratch.max_vectors(), graph_degree); // Test set() method diff --git a/diskann-disk/src/search/pq/quantizer_preprocess.rs b/diskann-disk/src/search/pq/quantizer_preprocess.rs index c5a2026305..5c4947f476 100644 --- a/diskann-disk/src/search/pq/quantizer_preprocess.rs +++ b/diskann-disk/src/search/pq/quantizer_preprocess.rs @@ -9,21 +9,21 @@ use diskann_vector::distance::Metric; use diskann_providers::model::compute_pq_distance; use diskann_providers::utils::BridgeErr; -use super::{PQData, PQScratch}; +use super::{PQData, PQQueryComputer, PQScratch}; /// Preprocesses the query vector for PQ distance calculations. /// This function rotates the query vector and prepares the PQ table distances /// for efficient computation during search operations. -pub fn quantizer_preprocess( - pq_scratch: &mut PQScratch, +fn preprocess_query( + query: &[f32], + lookup_table: &mut [f32], pq_data: &PQData, metric: Metric, - id_to_calculate_pq_distance: &[u32], ) -> ANNResult<()> { let table = pq_data.pq_table(); let expected_len = table.ncenters() * table.nchunks(); let dst = diskann_utils::views::MutMatrixView::try_from( - &mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len], + &mut lookup_table[..expected_len], table.nchunks(), table.ncenters(), ) @@ -36,21 +36,40 @@ pub fn quantizer_preprocess( // We're keeping that behavior here - treating `Cosine` and `CosineNormalized` // as L2 until a more thorough evaluation can be made. Metric::L2 | Metric::Cosine | Metric::CosineNormalized => { - table.process_into::( - &pq_scratch.query_scratch, - dst, - ); + table.process_into::(query, dst); } Metric::InnerProduct => { - table.process_into::( - &pq_scratch.query_scratch, - dst, - ); + table.process_into::(query, dst); } } - // Compute the pq distance between query vector to all the vertex in the pq - // calculation id scratch. + Ok(()) +} + +pub(crate) fn prepare_query( + computer: &mut PQQueryComputer, + pq_data: &PQData, + metric: Metric, + query: &[f32], +) -> ANNResult<()> { + computer.set(query)?; + let (query, lookup_table) = computer.preprocessing_buffers(); + preprocess_query(query, lookup_table, pq_data, metric) +} + +pub fn quantizer_preprocess( + pq_scratch: &mut PQScratch, + pq_data: &PQData, + metric: Metric, + id_to_calculate_pq_distance: &[u32], +) -> ANNResult<()> { + preprocess_query( + &pq_scratch.query_scratch, + &mut pq_scratch.aligned_pqtable_dist_scratch, + pq_data, + metric, + )?; + compute_pq_distance( id_to_calculate_pq_distance, pq_data.get_num_chunks(), diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 100936fbab..b228ebdb07 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -16,14 +16,20 @@ use std::{ use crate::data_model::GraphDataType; use diskann::{ error::IntoANNResult, + flat::{ + knn_search as flat_knn_search, DistancesUnordered, SearchStats as FlatSearchStats, + SearchStrategy as FlatSearchStrategy, + }, graph::{ self, ext::labeled::{self, QueryLabelProvider}, - glue::{self, DefaultPostProcessor, SearchPostProcess, SearchStrategy}, + glue::{ + self, DefaultPostProcessor, SearchPostProcess, SearchStrategy as GraphSearchStrategy, + }, search::{AdaptiveL, InlineFilterSearch, Knn}, search_output_buffer, DiskANNIndex, }, - neighbor::{self, Neighbor, NeighborPriorityQueue}, + neighbor::{self, Neighbor}, provider::{DataProvider, DefaultContext, HasId, NoopGuard}, utils::{IntoUsize, VectorRepr}, ANNError, ANNResult, @@ -37,11 +43,15 @@ use diskann_providers::{ storage::{get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file, LoadWith}, }; use diskann_utils::{ + future::SendFuture, object_pool::{ObjectPool, PoolOption, TryAsPooled}, views::Matrix, }; -use crate::search::pq::{quantizer_preprocess, PQData, PQScratch}; +use crate::search::pq::{ + prepare_query, PQBatchScratch, PQData, PQQueryComputer, PQQueryComputerArgs, + PQQueryComputerStorage, +}; use diskann_vector::{distance::Metric, DistanceFunction}; use tokio::runtime::Runtime; use tracing::debug; @@ -217,15 +227,11 @@ where /// `clippy::type_complexity`'s default threshold. type PostprocessFilter<'a> = &'a (dyn Fn(&u32) -> bool + Send + Sync); -/// Encodes whether to accept all candidates at rerank time or apply a -/// specific predicate. Used by `RerankAndFilter` and -/// `DeterminantDiversityAndFilter` instead of `Option` -/// so call sites are self-documenting without relying on comments to -/// explain what `None` means. +/// Encodes whether to accept all candidates or apply a specific predicate. +/// Used by `RerankAndFilter`, `DeterminantDiversityAndFilter`, and the flat visitor. #[derive(Clone, Copy)] pub enum PostprocessStrategy<'a> { - /// Accept every candidate — no predicate is called. Used by `FlatScan` - /// (filtered at scan time) and `InlineFilter` (filtered at visit time). + /// Accept every candidate — no predicate is called. AcceptAll, /// Apply the given predicate; non-matching candidates are dropped. Apply(PostprocessFilter<'a>), @@ -238,9 +244,7 @@ where { // Borrowed from `search_internal` so the strategy can be passed by value io_tracker: &'a IOTracker, - /// Consumed only by `default_post_processor()` → `RerankAndFilter`. - /// `FlatScan` and `InlineFilter` filter earlier in their pipelines and - /// pass `AcceptAll` here to avoid a redundant second pass. + /// Used by the flat visitor and the default post-processor. postprocess_filter: PostprocessStrategy<'a>, /// The vertex provider factory is used to create the vertex provider for each search instance. @@ -248,6 +252,10 @@ where /// Scratch pool for disk search operations that need allocations. scratch_pool: &'a Arc>>, + + pq_data: &'a PQData, + metric: Metric, + query_computer_pool: &'a Arc>, } // Struct to track IO. This is used by single thread, but needs to be Atomic as the Strategy has "Send" trait bound. @@ -287,6 +295,29 @@ impl IOTracker { } } +impl DiskSearchStrategy<'_, Data, ProviderFactory> +where + Data: GraphDataType, + ProviderFactory: VertexProviderFactory, +{ + fn prepare_query_computer(&self, query: &[Data::VectorDataType]) -> ANNResult { + let timer = Instant::now(); + let query = Data::VectorDataType::as_f32(query).into_ann_result()?; + let args = PQQueryComputerArgs::new( + self.pq_data.get_dim(), + self.pq_data.get_num_chunks(), + self.pq_data.get_num_centers(), + ); + let mut computer = PQQueryComputer::pooled(self.query_computer_pool, args)?; + prepare_query(&mut computer, self.pq_data, self.metric, &query)?; + IOTracker::add_time( + &self.io_tracker.preprocess_time_us, + timer.elapsed().as_micros() as u64, + ); + Ok(computer) + } +} + #[derive(Clone, Copy)] pub struct RerankAndFilter<'a> { filter: PostprocessStrategy<'a>, @@ -319,6 +350,59 @@ impl<'a> DeterminantDiversityAndFilter<'a> { } } +fn rerank_and_filter( + filter: PostprocessStrategy<'_>, + provider: &DiskProvider, + scratch: &mut DiskSearchScratch, + query: &[Data::VectorDataType], + candidates: I, + output: &mut B, +) -> ANNResult +where + Data: GraphDataType, + VP: VertexProvider, + I: Iterator>, + B: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + ?Sized, +{ + let mut uncached_ids = Vec::new(); + let mut reranked: Vec<_> = { + let mut process = |id: u32| { + if let Some(entry) = scratch.distance_cache.get(&id) { + Some(Neighbor::new((id, entry.1), entry.0)) + } else { + uncached_ids.push(id); + None + } + }; + match filter { + PostprocessStrategy::AcceptAll => candidates + .map(|candidate| *candidate.id()) + .filter_map(&mut process) + .collect(), + PostprocessStrategy::Apply(predicate) => candidates + .map(|candidate| *candidate.id()) + .filter(|id| predicate(id)) + .filter_map(&mut process) + .collect(), + } + }; + + if !uncached_ids.is_empty() { + ensure_vertex_loaded(&mut scratch.vertex_provider, &uncached_ids)?; + for id in uncached_ids { + let vector = scratch.vertex_provider.get_vector(&id)?; + let distance = provider + .distance_comparer + .evaluate_similarity(query, vector); + let data = *scratch.vertex_provider.get_associated_data(&id)?; + reranked.push(Neighbor::new((id, data), distance)); + } + } + + reranked.sort_unstable_by(neighbor::ord::fast_distance); + Ok(output.extend(reranked)) +} + impl SearchPostProcess< DiskAccessor<'_, Data, VP>, @@ -346,45 +430,14 @@ where + Send + ?Sized, { - let provider = accessor.provider; - - let mut uncached_ids = Vec::new(); - let mut reranked: Vec<_> = { - let mut process = |n: u32| { - if let Some(entry) = accessor.scratch.distance_cache.get(&n) { - Some(Neighbor::new((n, entry.1), entry.0)) - } else { - uncached_ids.push(n); - None - } - }; - match self.filter { - PostprocessStrategy::AcceptAll => candidates - .map(|n| *n.id()) - .filter_map(&mut process) - .collect(), - PostprocessStrategy::Apply(f) => candidates - .map(|n| *n.id()) - .filter(|id| f(id)) - .filter_map(&mut process) - .collect(), - } - }; - if !uncached_ids.is_empty() { - ensure_vertex_loaded(&mut accessor.scratch.vertex_provider, &uncached_ids)?; - for n in &uncached_ids { - let v = accessor.scratch.vertex_provider.get_vector(n)?; - let d = provider.distance_comparer.evaluate_similarity(query, v); - let a = accessor.scratch.vertex_provider.get_associated_data(n)?; - reranked.push(Neighbor::new((*n, *a), d)); - } - } - - // Sort the full precision distances. - reranked.sort_unstable_by(neighbor::ord::fast_distance); - - // Store the reranked results. - Ok(output.extend(reranked)) + rerank_and_filter( + self.filter, + accessor.provider, + &mut accessor.scratch, + query, + candidates, + output, + ) } } @@ -506,7 +559,7 @@ where } impl<'this, Data, ProviderFactory> - SearchStrategy<'this, DiskProvider, &'this [Data::VectorDataType]> + GraphSearchStrategy<'this, DiskProvider, &'this [Data::VectorDataType]> for DiskSearchStrategy<'this, Data, ProviderFactory> where Data: GraphDataType, @@ -521,16 +574,60 @@ where _context: &DefaultContext, query: &'this [Data::VectorDataType], ) -> Result { + let query_computer = self.prepare_query_computer(query)?; DiskAccessor::new( provider, self.io_tracker, query, + query_computer, self.vertex_provider_factory, self.scratch_pool, ) } } +impl<'strategy, 'query, Data, ProviderFactory> + FlatSearchStrategy, &'query [Data::VectorDataType]> + for DiskSearchStrategy<'strategy, Data, ProviderFactory> +where + Data: GraphDataType, + ProviderFactory: VertexProviderFactory, +{ + type ElementRef<'a> = &'a [u8]; + type QueryComputer = PQQueryComputer; + type QueryComputerError = ANNError; + type Visitor<'a> + = FlatVisitor<'a, Data, ProviderFactory::VertexProviderType> + where + Self: 'a, + DiskProvider: 'a; + type Error = ANNError; + + fn create_visitor<'a>( + &'a self, + provider: &'a DiskProvider, + _context: &'a DefaultContext, + ) -> Result, Self::Error> { + let filter = match self.postprocess_filter { + PostprocessStrategy::AcceptAll => None, + PostprocessStrategy::Apply(filter) => Some(filter), + }; + FlatVisitor::new( + provider, + filter, + self.vertex_provider_factory, + self.scratch_pool, + ) + } + + fn build_query_computer( + &self, + query: &'query [Data::VectorDataType], + ) -> Result { + self.prepare_query_computer(query) + } +} + impl<'this, Data, ProviderFactory> DefaultPostProcessor< 'this, @@ -560,16 +657,14 @@ where VP: VertexProvider, { distance_cache: HashMap, - pq_scratch: PQScratch, + pq_batch_scratch: PQBatchScratch, vertex_provider: VP, } #[derive(Clone)] struct DiskSearchScratchArgs<'a, ProviderFactory> { graph_degree: usize, - pq_dim: usize, num_pq_chunks: usize, - num_pq_centers: usize, vertex_factory: &'a ProviderFactory, graph_header: &'a GraphHeader, } @@ -583,12 +678,7 @@ where type Error = ANNError; fn try_create(args: &DiskSearchScratchArgs) -> Result { - let pq_scratch = PQScratch::new( - args.graph_degree, - args.pq_dim, - args.num_pq_chunks, - args.num_pq_centers, - )?; + let pq_batch_scratch = PQBatchScratch::new(args.graph_degree, args.num_pq_chunks)?; const DEFAULT_BEAM_WIDTH: usize = 0; // Setting as 0 to avoid preallocation of memory. let vertex_provider = args @@ -597,7 +687,7 @@ where Ok(Self { distance_cache: HashMap::new(), - pq_scratch, + pq_batch_scratch, vertex_provider, }) } @@ -620,6 +710,7 @@ where provider: &'a DiskProvider, io_tracker: &'a IOTracker, scratch: PoolOption>, + query_computer: PQQueryComputer, query: &'a [Data::VectorDataType], } @@ -634,18 +725,18 @@ where where F: FnMut(f32, u32), { - let pq_scratch = &mut self.scratch.pq_scratch; + let pq_scratch = &mut self.scratch.pq_batch_scratch; compute_pq_distance( ids, self.provider.pq_data.get_num_chunks(), - &pq_scratch.aligned_pqtable_dist_scratch, + self.query_computer.lookup_table(), self.provider.pq_data.pq_compressed_data().as_slice(), &mut pq_scratch.aligned_pq_coord_scratch, &mut pq_scratch.aligned_dist_scratch, )?; for (i, id) in ids.iter().enumerate() { - let distance = self.scratch.pq_scratch.aligned_dist_scratch[i]; + let distance = self.scratch.pq_batch_scratch.aligned_dist_scratch[i]; f(distance, *id); } @@ -661,6 +752,159 @@ where type Id = u32; } +pub struct FlatVisitor<'a, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + provider: &'a DiskProvider, + filter: Option>, + scratch: PoolOption>, +} + +impl<'a, Data, VP> FlatVisitor<'a, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + fn new( + provider: &'a DiskProvider, + filter: Option>, + vertex_provider_factory: &'a VPF, + scratch_pool: &'a Arc>>, + ) -> ANNResult + where + VPF: VertexProviderFactory, + { + let pq_points = provider.pq_data.pq_compressed_data().nrows(); + if pq_points != provider.num_points { + return Err(diskann_error!( + ErrorKind::IndexError, + "PQ data contains {pq_points} points, expected {}", + provider.num_points, + )); + } + + let scratch = PoolOption::try_pooled( + scratch_pool, + &DiskSearchScratchArgs { + graph_degree: provider.graph_header.max_degree::()?, + num_pq_chunks: provider.pq_data.get_num_chunks(), + vertex_factory: vertex_provider_factory, + graph_header: &provider.graph_header, + }, + )?; + + Ok(Self { + provider, + filter, + scratch, + }) + } +} + +impl HasId for FlatVisitor<'_, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + type Id = u32; +} + +impl DistancesUnordered for FlatVisitor<'_, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + type ElementRef<'a> = &'a [u8]; + type Error = ANNError; + + fn distances_unordered( + &mut self, + computer: &PQQueryComputer, + mut f: F, + ) -> impl SendFuture> + where + F: Send + FnMut(Self::Id, f32), + { + async move { + let batch_size = self.scratch.pq_batch_scratch.max_vectors(); + if batch_size == 0 { + return Err(diskann_error!( + ErrorKind::IndexError, + "pq scratch must support at least one vector", + )); + } + + let mut ids = Vec::with_capacity(batch_size); + let mut remaining = (0..self.provider.num_points as u32) + .filter(|id| self.filter.is_none_or(|filter| filter(id))); + + loop { + ids.clear(); + ids.extend(remaining.by_ref().take(batch_size)); + if ids.is_empty() { + break; + } + + let scratch = &mut self.scratch.pq_batch_scratch; + compute_pq_distance( + &ids, + self.provider.pq_data.get_num_chunks(), + computer.lookup_table(), + self.provider.pq_data.pq_compressed_data().as_slice(), + &mut scratch.aligned_pq_coord_scratch, + &mut scratch.aligned_dist_scratch, + )?; + + for (id, distance) in + std::iter::zip(&ids, &scratch.aligned_dist_scratch[..ids.len()]) + { + f(*id, *distance); + } + } + + Ok(()) + } + } +} + +impl + SearchPostProcess< + FlatVisitor<'_, Data, VP>, + &[Data::VectorDataType], + (u32, Data::AssociatedDataType), + > for RerankAndFilter<'_> +where + Data: GraphDataType, + VP: VertexProvider, +{ + type Error = ANNError; + + async fn post_process( + &self, + visitor: &mut FlatVisitor<'_, Data, VP>, + query: &[Data::VectorDataType], + candidates: I, + output: &mut B, + ) -> Result + where + I: Iterator> + Send, + B: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + + Send + + ?Sized, + { + rerank_and_filter( + self.filter, + visitor.provider, + &mut visitor.scratch, + query, + candidates, + output, + ) + } +} + impl glue::SearchAccessor for DiskAccessor<'_, Data, VP> where Data: GraphDataType, @@ -730,45 +974,28 @@ where provider: &'a DiskProvider, io_tracker: &'a IOTracker, query: &'a [Data::VectorDataType], + query_computer: PQQueryComputer, vertex_provider_factory: &'a VPF, scratch_pool: &'a Arc>>, ) -> ANNResult where VPF: VertexProviderFactory, { - let mut scratch = PoolOption::try_pooled( + let scratch = PoolOption::try_pooled( scratch_pool, &DiskSearchScratchArgs { graph_degree: provider.graph_header.max_degree::()?, - pq_dim: provider.pq_data.get_dim(), num_pq_chunks: provider.pq_data.get_num_chunks(), - num_pq_centers: provider.pq_data.get_num_centers(), vertex_factory: vertex_provider_factory, graph_header: &provider.graph_header, }, )?; - // Decode caller's native vector representation into `f32`; downstream PQ kernels operate purely on `&[f32]`. - let f32_query = Data::VectorDataType::as_f32(query).into_ann_result()?; - scratch.pq_scratch.set(&f32_query)?; - let start_vertex_id = provider.graph_header.metadata().medoid as u32; - - let timer = Instant::now(); - quantizer_preprocess( - &mut scratch.pq_scratch, - &provider.pq_data, - provider.metric, - &[start_vertex_id], - )?; - IOTracker::add_time( - &io_tracker.preprocess_time_us, - timer.elapsed().as_micros() as u64, - ); - Ok(Self { provider, io_tracker, scratch, + query_computer, query, }) } @@ -817,6 +1044,8 @@ pub struct DiskIndexSearcher< /// Scratch pool for disk search operations that need allocations. scratch_pool: Arc>>, + + query_computer_pool: Arc>, } #[derive(Debug)] @@ -895,13 +1124,17 @@ where let pq_data = disk_index_reader.get_pq_data(); let scratch_pool_args = DiskSearchScratchArgs { graph_degree: graph_header.max_degree::()?, - pq_dim: pq_data.get_dim(), num_pq_chunks: pq_data.get_num_chunks(), - num_pq_centers: pq_data.get_num_centers(), vertex_factory: &vertex_provider_factory, graph_header: &graph_header, }; let scratch_pool = Arc::new(ObjectPool::try_new(&scratch_pool_args, 0, None)?); + let query_computer_args = PQQueryComputerArgs::new( + pq_data.get_dim(), + pq_data.get_num_chunks(), + pq_data.get_num_centers(), + ); + let query_computer_pool = Arc::new(ObjectPool::try_new(query_computer_args, 0, None)?); let disk_provider = DiskProvider::new( disk_index_reader, @@ -917,6 +1150,7 @@ where runtime, vertex_provider_factory, scratch_pool, + query_computer_pool, }) } @@ -926,11 +1160,15 @@ where io_tracker: &'a IOTracker, postprocess_filter: PostprocessStrategy<'a>, ) -> DiskSearchStrategy<'a, Data, ProviderFactory> { + let provider = self.index.provider(); DiskSearchStrategy { io_tracker, postprocess_filter, vertex_provider_factory: &self.vertex_provider_factory, scratch_pool: &self.scratch_pool, + pq_data: &provider.pq_data, + metric: provider.metric, + query_computer_pool: &self.query_computer_pool, } } @@ -946,63 +1184,33 @@ where &self, strategy: &DiskSearchStrategy<'_, Data, ProviderFactory>, query: &[Data::VectorDataType], - vector_filter: Option<&(dyn Fn(&u32) -> bool + Send + Sync)>, neighbors_before_reranking: usize, output: &mut OB, ) -> ANNResult where OB: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + Send, { - let provider = self.index.provider(); - let mut accessor = strategy - .search_accessor(provider, &DefaultContext, query) - .into_ann_result()?; - - // Derive the batch size from the scratch data structure. Providing too many vectors - // will panic. - let batch_size = accessor.scratch.pq_scratch.max_vectors(); - - // This check should always hold since `graph_degree` comes from - // `diskann::graph::Config` and is forced to be non-zero. But this is defensive - // against misconfiguration. - if batch_size == 0 { - return Err(diskann_error!( + let k = NonZeroUsize::new(neighbors_before_reranking).ok_or_else(|| { + diskann_error!( ErrorKind::IndexError, - "pq scratch must support at least one vector", - )); - } - - let mut id_buffer = Vec::with_capacity(batch_size); - - let mut best = NeighborPriorityQueue::new(neighbors_before_reranking); - let mut cmps = 0u32; - - // `None` short-circuits to `true` — no dyn-fn call per node on the - // unfiltered (recall-baseline) path. - let mut iter = - (0..provider.num_points as u32).filter(|id| vector_filter.is_none_or(|f| f(id))); - loop { - id_buffer.clear(); - id_buffer.extend(iter.by_ref().take(batch_size)); - - if id_buffer.is_empty() { - break; - } - - accessor.pq_distances(&id_buffer, |dist, id| best.insert(Neighbor::new(id, dist)))?; - cmps += id_buffer.len() as u32; - } - - let result_count = strategy - .default_post_processor() - .post_process(&mut accessor, query, best.iter(), output) - .await - .into_ann_result()?; + "flat search list size must be greater than zero", + ) + })?; + let FlatSearchStats { cmps, result_count } = flat_knn_search( + self.index.provider(), + k, + strategy, + RerankAndFilter::new(PostprocessStrategy::AcceptAll), + &DefaultContext, + query, + output, + ) + .await?; Ok(graph::index::SearchStats { cmps, hops: 0, - result_count: result_count as u32, + result_count, range_search_second_round: false, }) } @@ -1140,11 +1348,15 @@ where // as the post-processor over the L candidate pool. let stats = match mode { SearchMode::FlatScan { filter } => { - let strategy = self.search_strategy(&io_tracker, PostprocessStrategy::AcceptAll); + let strategy = self.search_strategy( + &io_tracker, + filter + .as_deref() + .map_or(PostprocessStrategy::AcceptAll, PostprocessStrategy::Apply), + ); self.runtime.block_on(self.flat_search( &strategy, query, - filter.as_deref(), l, &mut result_output_buffer, ))? @@ -2376,6 +2588,93 @@ mod disk_provider_tests { ); } + #[test] + fn unfiltered_flat_scan_matches_baseline() { + let storage_provider = Arc::new(VirtualStorageProvider::new_overlay(test_data_root())); + let search_engine = create_disk_index_searcher::( + CreateDiskIndexSearcherParams { + max_thread_num: 1, + pq_pivot_file_path: TEST_PQ_PIVOT_128DIM, + pq_compressed_file_path: TEST_PQ_COMPRESSED_128DIM, + index_path: TEST_INDEX_128DIM, + index_path_prefix: TEST_INDEX_PREFIX_128DIM, + ..Default::default() + }, + &storage_provider, + ); + + let result = search_engine + .search(&[0.1; 128], 10, 10, None, SearchMode::flat()) + .unwrap(); + + let expected = [ + (152, 256101.7), + (115, 256400.48), + (98, 256451.28), + (73, 256572.89), + (20, 256623.28), + (173, 256636.9), + (95, 256661.28), + (137, 256673.7), + (118, 256675.3), + (72, 256709.69), + ]; + assert_eq!(result.results.len(), expected.len()); + for (index, (actual, (expected_id, expected_distance))) in + std::iter::zip(&result.results, expected).enumerate() + { + assert_eq!( + actual.vertex_id, expected_id, + "flat baseline ID mismatch at result {index}", + ); + assert!( + (actual.distance - expected_distance).abs() <= 0.02, + "flat baseline distance mismatch at result {index}: expected \ + {expected_distance}, got {}", + actual.distance, + ); + } + + assert!(result + .results + .windows(2) + .all(|pair| pair[0].distance <= pair[1].distance)); + assert_eq!(result.stats.cmps, 256); + assert_eq!(result.stats.result_count, 10); + } + + #[test] + fn flat_filter_runs_once_per_point() { + let storage_provider = Arc::new(VirtualStorageProvider::new_overlay(test_data_root())); + let search_engine = create_disk_index_searcher::( + CreateDiskIndexSearcherParams { + max_thread_num: 1, + pq_pivot_file_path: TEST_PQ_PIVOT_128DIM, + pq_compressed_file_path: TEST_PQ_COMPRESSED_128DIM, + index_path: TEST_INDEX_128DIM, + index_path_prefix: TEST_INDEX_PREFIX_128DIM, + ..Default::default() + }, + &storage_provider, + ); + let calls = AtomicUsize::new(0); + + search_engine + .search( + &[0.1; 128], + 10, + 10, + None, + SearchMode::flat_filtered(|_| { + calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + true + }), + ) + .unwrap(); + + assert_eq!(calls.load(std::sync::atomic::Ordering::Relaxed), 256); + } + // =========================================================================== // Inline filter + AdaptiveL behavioral tests // =========================================================================== diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index 0c9aa6c81c..99e7d58709 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -35,6 +35,62 @@ pub struct FlatIndex { provider: P, } +/// Brute-force k-nearest-neighbor search over a borrowed provider. +/// +/// Streams every element produced by the strategy's visitor through the query +/// computer, keeps the best `k` candidates in a [`NeighborPriorityQueue`], then runs +/// `processor` over the survivors to populate `output`. +/// +/// # Errors +/// +/// Returns an error if visitor creation, query preprocessing, distance scanning, +/// or result post-processing fails. Distance-scan errors are escalated because a +/// partial flat scan cannot produce correct k-nearest-neighbor results. +pub fn knn_search<'a, P, S, T, O, PP, OB>( + provider: &'a P, + k: NonZeroUsize, + strategy: &'a S, + processor: PP, + context: &'a P::Context, + query: T, + output: &'a mut OB, +) -> impl SendFuture> + 'a +where + P: DataProvider, + S: SearchStrategy + 'a, + T: Copy + Send + Sync + 'a, + O: Send + 'a, + PP: SearchPostProcess, T, O> + Send + Sync + 'a, + OB: SearchOutputBuffer + Send + ?Sized + 'a, +{ + async move { + let mut visitor = strategy + .create_visitor(provider, context) + .into_ann_result()?; + + let computer = strategy.build_query_computer(query).into_ann_result()?; + + let k = k.get(); + let mut queue = NeighborPriorityQueue::new(k); + let mut cmps: u32 = 0; + + visitor + .distances_unordered(&computer, |id, dist| { + cmps += 1; + queue.insert(Neighbor::new(id, dist)); + }) + .await + .escalate("flat scan must complete to produce correct k-NN results")?; + + let result_count = processor + .post_process(&mut visitor, query, queue.iter().take(k), output) + .await + .into_ann_result()? as u32; + + Ok(SearchStats { cmps, result_count }) + } +} + impl FlatIndex

{ /// Construct a new [`FlatIndex`] around `provider`. pub fn new(provider: P) -> Self { @@ -71,30 +127,16 @@ impl FlatIndex

{ OB: SearchOutputBuffer + Send + ?Sized, { async move { - let mut visitor = strategy - .create_visitor(&self.provider, context) - .into_ann_result()?; - - let computer = strategy.build_query_computer(query).into_ann_result()?; - - let k = k.get(); - let mut queue = NeighborPriorityQueue::new(k); - let mut cmps: u32 = 0; - - visitor - .distances_unordered(&computer, |id, dist| { - cmps += 1; - queue.insert(Neighbor::new(id, dist)); - }) - .await - .escalate("flat scan must complete to produce correct k-NN results")?; - - let result_count = processor - .post_process(&mut visitor, query, queue.iter().take(k), output) - .await - .into_ann_result()? as u32; - - Ok(SearchStats { cmps, result_count }) + knn_search( + &self.provider, + k, + strategy, + processor, + context, + query, + output, + ) + .await } } } diff --git a/diskann/src/flat/mod.rs b/diskann/src/flat/mod.rs index f33d97851d..3106490cb5 100644 --- a/diskann/src/flat/mod.rs +++ b/diskann/src/flat/mod.rs @@ -26,7 +26,7 @@ pub mod index; pub mod strategy; -pub use index::{FlatIndex, SearchStats}; +pub use index::{FlatIndex, SearchStats, knn_search}; pub use strategy::{DistancesUnordered, SearchStrategy}; #[cfg(test)]