diff --git a/diskann-benchmark/src/flat/search.rs b/diskann-benchmark/src/flat/search.rs index f2f0d422d8..47e14b0151 100644 --- a/diskann-benchmark/src/flat/search.rs +++ b/diskann-benchmark/src/flat/search.rs @@ -223,29 +223,25 @@ impl Strategy { } /// The visitor that iterates over all vectors in the provider. -struct Visitor<'a, T> { +struct Visitor<'a, T: VectorRepr> { data: &'a Matrix, + computer: T::QueryDistance, } impl HasId for Visitor<'_, T> { type Id = u32; } -impl DistancesUnordered for Visitor<'_, T> { - type ElementRef<'a> = &'a [T]; +impl DistancesUnordered for Visitor<'_, T> { type Error = diskann::error::Infallible; - fn distances_unordered( - &mut self, - computer: &T::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { async move { for (i, vector) in self.data.row_iter().enumerate() { - let dist = computer.evaluate_similarity(vector); + let dist = self.computer.evaluate_similarity(vector); f(i as u32, dist); } Ok(()) @@ -253,33 +249,21 @@ impl DistancesUnordered for Visitor<'_, T> { } } -impl SearchStrategy, &[T]> for Strategy { - type ElementRef<'a> = &'a [T]; - type QueryComputer = T::QueryDistance; - type QueryComputerError = diskann::error::Infallible; - type Visitor<'a> - = Visitor<'a, T> - where - Self: 'a, - InMemProvider: 'a; +impl<'a, T: VectorRepr> SearchStrategy<'a, InMemProvider, &'a [T]> for Strategy { + type Visitor = Visitor<'a, T>; type Error = diskann::error::Infallible; - fn create_visitor<'a>( + fn create_visitor( &'a self, provider: &'a InMemProvider, _context: &'a DefaultContext, - ) -> Result, Self::Error> { + query: &'a [T], + ) -> Result { Ok(Visitor { data: &provider.data, + computer: T::query_distance(query, self.metric), }) } - - fn build_query_computer( - &self, - query: &[T], - ) -> Result { - Ok(T::query_distance(query, self.metric)) - } } ////////////////////////////////////////// diff --git a/diskann-disk/src/search/pq/mod.rs b/diskann-disk/src/search/pq/mod.rs index ef9ffa1d5d..06bfd6aabd 100644 --- a/diskann-disk/src/search/pq/mod.rs +++ b/diskann-disk/src/search/pq/mod.rs @@ -11,4 +11,5 @@ pub use pq_scratch::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..19b770b204 100644 --- a/diskann-disk/src/search/pq/pq_scratch.rs +++ b/diskann-disk/src/search/pq/pq_scratch.rs @@ -13,8 +13,7 @@ use crate::error::{diskann_error, ErrorKind}; #[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 +24,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 +42,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 +76,7 @@ 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. pub(crate) fn max_vectors(&self) -> usize { self.aligned_dist_scratch.len() } @@ -120,7 +115,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..467c7f4a99 100644 --- a/diskann-disk/src/search/pq/quantizer_preprocess.rs +++ b/diskann-disk/src/search/pq/quantizer_preprocess.rs @@ -14,16 +14,16 @@ use super::{PQData, 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,44 @@ 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( + pq_scratch: &mut PQScratch, + pq_data: &PQData, + metric: Metric, + query: &[f32], +) -> ANNResult<()> { + pq_scratch.set(query)?; + preprocess_query( + &pq_scratch.query_scratch, + &mut pq_scratch.aligned_pqtable_dist_scratch, + 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 37a3171dfa..4f52f3195f 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -16,15 +16,21 @@ 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::{self, BufferState, IdDistanceAssociatedData}, DiskANNIndex, }, - neighbor::{self, Neighbor, NeighborPriorityQueue}, + neighbor::{self, Neighbor}, provider::{DataProvider, DefaultContext, HasId, NoopGuard}, utils::{IntoUsize, VectorRepr}, ANNError, ANNResult, @@ -38,11 +44,12 @@ 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, PQData, PQScratch}; use diskann_vector::{distance::Metric, DistanceFunction}; use tokio::runtime::Runtime; use tracing::debug; @@ -218,15 +225,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>), @@ -240,9 +243,7 @@ where // Borrowed from `search_internal` so the strategy can be passed by value io_tracker: &'a IOTracker, cache_indexed_vectors: bool, - /// 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. @@ -398,7 +399,8 @@ impl search_output_buffer::SearchOutputBuffer( - accessor: &mut DiskAccessor<'_, Data, VP>, + cache_indexed_vectors: bool, + scratch: &mut DiskSearchScratch, reranked: I, output: &mut B, ) -> ANNResult @@ -411,7 +413,7 @@ where > + Send + ?Sized, { - if !accessor.cache_indexed_vectors { + if !cache_indexed_vectors { return Ok(output.extend(reranked.into_iter().map(|candidate| { let ((id, data), distance) = candidate.as_tuple(); Neighbor::new((id, data, None), distance) @@ -424,14 +426,13 @@ where break; } let ((id, data), distance) = candidate.as_tuple(); - let vector = match accessor - .scratch + let vector = match scratch .distance_cache .remove(&id) .and_then(|(_, _, vector)| vector) { Some(vector) => vector, - None => Box::from(accessor.scratch.vertex_provider.get_vector(&id)?), + None => Box::from(scratch.vertex_provider.get_vector(&id)?), }; count += 1; if output @@ -444,6 +445,63 @@ where Ok(count) } +fn rerank_and_filter( + filter: PostprocessStrategy<'_>, + provider: &DiskProvider, + scratch: &mut DiskSearchScratch, + query: &[Data::VectorDataType], + cache_indexed_vectors: bool, + candidates: I, + output: &mut B, +) -> ANNResult +where + Data: GraphDataType, + VP: VertexProvider, + I: Iterator>, + B: search_output_buffer::SearchOutputBuffer< + SearchPayload, + > + Send + + ?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); + extend_output(cache_indexed_vectors, scratch, reranked, output) +} + impl SearchPostProcess< DiskAccessor<'_, Data, VP>, @@ -469,45 +527,15 @@ 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. - extend_output(accessor, reranked, output) + rerank_and_filter( + self.filter, + accessor.provider, + &mut accessor.scratch, + query, + accessor.cache_indexed_vectors, + candidates, + output, + ) } } @@ -581,7 +609,8 @@ where )?; extend_output( - accessor, + accessor.cache_indexed_vectors, + &mut accessor.scratch, reranked.into_iter().map(|idx| { let id = candidate_ids[idx]; let distance = candidate_distances[idx]; @@ -629,7 +658,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, @@ -655,6 +684,38 @@ where } } +impl<'this, Data, ProviderFactory> + FlatSearchStrategy<'this, DiskProvider, &'this [Data::VectorDataType]> + for DiskSearchStrategy<'this, Data, ProviderFactory> +where + Data: GraphDataType, + ProviderFactory: VertexProviderFactory, +{ + type Visitor = FlatVisitor<'this, Data, ProviderFactory::VertexProviderType>; + type Error = ANNError; + + fn create_visitor( + &'this self, + provider: &'this DiskProvider, + _context: &'this DefaultContext, + query: &'this [Data::VectorDataType], + ) -> Result { + let filter = match self.postprocess_filter { + PostprocessStrategy::AcceptAll => None, + PostprocessStrategy::Apply(filter) => Some(filter), + }; + FlatVisitor::new( + provider, + filter, + self.io_tracker, + query, + self.vertex_provider_factory, + self.scratch_pool, + self.cache_indexed_vectors, + ) + } +} + impl<'this, Data, ProviderFactory> DefaultPostProcessor< 'this, @@ -734,6 +795,70 @@ where } } +impl DiskSearchScratch +where + Data: GraphDataType, + VP: VertexProvider, +{ + fn pooled_for_query( + pool: &Arc>, + provider: &DiskProvider, + io_tracker: &IOTracker, + query: &[Data::VectorDataType], + vertex_provider_factory: &VPF, + ) -> ANNResult> + where + VPF: VertexProviderFactory, + { + let mut scratch = PoolOption::try_pooled( + 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, + }, + )?; + + let timer = Instant::now(); + let query = Data::VectorDataType::as_f32(query).into_ann_result()?; + prepare_query( + &mut scratch.pq_scratch, + &provider.pq_data, + provider.metric, + &query, + )?; + IOTracker::add_time( + &io_tracker.preprocess_time_us, + timer.elapsed().as_micros() as u64, + ); + Ok(scratch) + } + + fn pq_distances(&mut self, pq_data: &PQData, ids: &[u32], mut f: F) -> ANNResult<()> + where + F: FnMut(u32, f32), + { + compute_pq_distance( + ids, + pq_data.get_num_chunks(), + &self.pq_scratch.aligned_pqtable_dist_scratch, + pq_data.pq_compressed_data().as_slice(), + &mut self.pq_scratch.aligned_pq_coord_scratch, + &mut self.pq_scratch.aligned_dist_scratch, + )?; + + for (id, distance) in + std::iter::zip(ids, &self.pq_scratch.aligned_dist_scratch[..ids.len()]) + { + f(*id, *distance); + } + Ok(()) + } +} + pub struct DiskAccessor<'a, Data, VP> where Data: GraphDataType, @@ -757,26 +882,76 @@ where where F: FnMut(f32, u32), { - let pq_scratch = &mut self.scratch.pq_scratch; - compute_pq_distance( - ids, - self.provider.pq_data.get_num_chunks(), - &pq_scratch.aligned_pqtable_dist_scratch, - self.provider.pq_data.pq_compressed_data().as_slice(), - &mut pq_scratch.aligned_pq_coord_scratch, - &mut pq_scratch.aligned_dist_scratch, - )?; + self.scratch + .pq_distances(&self.provider.pq_data, ids, |id, distance| { + f(distance, id); + }) + } +} + +impl HasId for DiskAccessor<'_, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + type Id = u32; +} + +pub struct FlatVisitor<'a, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + provider: &'a DiskProvider, + filter: Option>, + scratch: PoolOption>, + cache_indexed_vectors: bool, +} - for (i, id) in ids.iter().enumerate() { - let distance = self.scratch.pq_scratch.aligned_dist_scratch[i]; - f(distance, *id); +impl<'a, Data, VP> FlatVisitor<'a, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + fn new( + provider: &'a DiskProvider, + filter: Option>, + io_tracker: &'a IOTracker, + query: &'a [Data::VectorDataType], + vertex_provider_factory: &'a VPF, + scratch_pool: &'a Arc>>, + cache_indexed_vectors: bool, + ) -> 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, + )); } - Ok(()) + let scratch = DiskSearchScratch::pooled_for_query( + scratch_pool, + provider, + io_tracker, + query, + vertex_provider_factory, + )?; + + Ok(Self { + provider, + filter, + scratch, + cache_indexed_vectors, + }) } } -impl HasId for DiskAccessor<'_, Data, VP> +impl HasId for FlatVisitor<'_, Data, VP> where Data: GraphDataType, VP: VertexProvider, @@ -784,6 +959,84 @@ where type Id = u32; } +impl DistancesUnordered for FlatVisitor<'_, Data, VP> +where + Data: GraphDataType, + VP: VertexProvider, +{ + type Error = ANNError; + + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> + where + F: Send + FnMut(Self::Id, f32), + { + async move { + let batch_size = self.scratch.pq_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; + } + + self.scratch + .pq_distances(&self.provider.pq_data, &ids, &mut f)?; + } + + Ok(()) + } + } +} + +impl + SearchPostProcess< + FlatVisitor<'_, Data, VP>, + &[Data::VectorDataType], + SearchPayload, + > 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< + SearchPayload, + > + Send + + ?Sized, + { + rerank_and_filter( + self.filter, + visitor.provider, + &mut visitor.scratch, + query, + visitor.cache_indexed_vectors, + candidates, + output, + ) + } +} + impl glue::SearchAccessor for DiskAccessor<'_, Data, VP> where Data: GraphDataType, @@ -860,34 +1113,13 @@ where where VPF: VertexProviderFactory, { - let mut scratch = PoolOption::try_pooled( + let scratch = DiskSearchScratch::pooled_for_query( 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], + provider, + io_tracker, + query, + vertex_provider_factory, )?; - IOTracker::add_time( - &io_tracker.preprocess_time_us, - timer.elapsed().as_micros() as u64, - ); Ok(Self { provider, @@ -1098,7 +1330,6 @@ 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 @@ -1107,56 +1338,27 @@ where SearchPayload, > + 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, }) } @@ -1386,13 +1588,14 @@ where SearchMode::FlatScan { filter } => { let strategy = self.search_strategy( &io_tracker, - PostprocessStrategy::AcceptAll, + filter + .as_deref() + .map_or(PostprocessStrategy::AcceptAll, PostprocessStrategy::Apply), cache_indexed_vectors, ); self.runtime.block_on(self.flat_search( &strategy, query, - filter.as_deref(), l, &mut result_output_buffer, ))? @@ -2742,6 +2945,136 @@ 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); + + let result = search_engine + .search( + &[0.1; 128], + 10, + 10, + None, + SearchMode::flat_filtered(|id| { + calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + id % 64 == 0 + }), + ) + .unwrap(); + + assert_eq!(calls.load(std::sync::atomic::Ordering::Relaxed), 256); + assert_eq!(result.stats.cmps, 4); + assert_eq!(result.stats.result_count, 4); + assert!(result.results.iter().all(|item| item.vertex_id % 64 == 0)); + } + + #[test] + fn graph_and_flat_queries_isolate_pooled_query_state() { + 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 graph_before = search_engine + .search(&[0.1; 128], 10, 10, None, SearchMode::graph()) + .unwrap(); + let flat_before = search_engine + .search(&[0.9; 128], 10, 10, None, SearchMode::flat()) + .unwrap(); + let graph_after = search_engine + .search(&[0.1; 128], 10, 10, None, SearchMode::graph()) + .unwrap(); + let flat_after = search_engine + .search(&[0.9; 128], 10, 10, None, SearchMode::flat()) + .unwrap(); + + let assert_same = |before: &SearchResult<()>, after: &SearchResult<()>| { + assert_eq!(before.stats.cmps, after.stats.cmps); + assert_eq!(before.results.len(), after.results.len()); + for (before, after) in std::iter::zip(&before.results, &after.results) { + assert_eq!(before.vertex_id, after.vertex_id); + assert_eq!(before.distance, after.distance); + } + }; + assert_same(&graph_before, &graph_after); + assert_same(&flat_before, &flat_after); + } + // =========================================================================== // Inline filter + AdaptiveL behavioral tests // =========================================================================== diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index 0c9aa6c81c..436954a0e7 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -35,6 +35,60 @@ pub struct FlatIndex { provider: P, } +/// Brute-force k-nearest-neighbor search over a borrowed provider. +/// +/// Streams every distance produced by the strategy's accessor, keeps the best `k` +/// candidates in a [`NeighborPriorityQueue`], then runs `processor` over the survivors +/// to populate `output`. +/// +/// # Errors +/// +/// Returns an error if accessor construction, 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, P, T> + 'a, + T: Copy + Send + Sync + 'a, + O: Send + 'a, + PP: SearchPostProcess + Send + Sync + 'a, + OB: SearchOutputBuffer + Send + ?Sized + 'a, +{ + async move { + let mut visitor = strategy + .create_visitor(provider, context, query) + .into_ann_result()?; + + let k = k.get(); + let mut queue = NeighborPriorityQueue::new(k); + let mut cmps: u32 = 0; + + visitor + .distances_unordered(|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 { @@ -48,53 +102,44 @@ impl FlatIndex

{ /// Brute-force k-nearest-neighbor flat search. /// - /// 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`. + /// Streams every distance produced by the strategy's accessor, keeps the best `k` + /// candidates in a [`NeighborPriorityQueue`], then runs `processor` over the + /// survivors to populate `output`. /// /// The post-processor [`SearchPostProcess::post_process`] outputs the number /// of results that survive, which is returned as `SearchStats::result_count`. - pub fn knn_search( - &self, + /// + /// # Errors + /// + /// Returns an error if visitor construction, distance scanning, or result + /// post-processing fails. + pub fn knn_search<'a, S, T, O, PP, OB>( + &'a self, k: NonZeroUsize, - strategy: &S, + strategy: &'a S, processor: PP, - context: &P::Context, + context: &'a P::Context, query: T, - output: &mut OB, - ) -> impl SendFuture> + output: &'a mut OB, + ) -> impl SendFuture> + 'a where - S: SearchStrategy, - T: Copy + Send + Sync, - O: Send, - PP: for<'a> SearchPostProcess, T, O> + Send + Sync, - OB: SearchOutputBuffer + Send + ?Sized, + S: SearchStrategy<'a, P, T> + 'a, + T: Copy + Send + Sync + 'a, + O: Send + 'a, + PP: SearchPostProcess + Send + Sync + 'a, + OB: SearchOutputBuffer + Send + ?Sized + 'a, { 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)] diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs index 1cdb6957fd..92120c6d3a 100644 --- a/diskann/src/flat/strategy.rs +++ b/diskann/src/flat/strategy.rs @@ -8,96 +8,58 @@ use std::fmt::Debug; use diskann_utils::future::SendFuture; -use diskann_vector::PreprocessedDistanceFunction; use crate::{ error::{StandardError, ToRanked}, provider::{DataProvider, HasId}, }; -/// Fused iterate-and-score primitive over the elements of a flat index. +/// Per-query accessor that drives a complete flat scan. /// -/// Implementations drive an entire scan over the underlying data, scoring each element -/// with the supplied computer `C` and invoking `f` with the resulting `(id, distance)` -/// pair. The associated [`Self::ElementRef`] is the reference shape on which `C` must -/// be able to compute distances. -pub trait DistancesUnordered: HasId + Send + Sync -where - C: for<'a> PreprocessedDistanceFunction, f32>, -{ - /// Lifetime is intentionally unconstrained so it can appear under HRTB without - /// inducing a `'static` bound on `Self`. - type ElementRef<'a>; - +/// The accessor owns the query-specific computation state so implementations can fuse +/// query preprocessing, data access, batching, filtering, and distance computation. +pub trait DistancesUnordered: HasId + Send + Sync { /// The error type for [`Self::distances_unordered`]. type Error: ToRanked + Debug + Send + Sync + 'static; - /// Drive the entire scan, scoring each element with `computer` and invoking `f` - /// with the resulting `(id, distance)` pair. - fn distances_unordered( - &mut self, - computer: &C, - f: F, - ) -> impl SendFuture> + /// Drive the entire scan, invoking `f` with each `(id, distance)` pair. + /// + /// # Errors + /// + /// Returns an error when the backend cannot complete the scan. + fn distances_unordered(&mut self, f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32); } -/// Per-call configuration that knows how to construct a per-query -/// [`DistancesUnordered`] visitor for a provider, and the [`Self::QueryComputer`] used -/// to score each element during the scan. -pub trait SearchStrategy: Send + Sync +/// Per-call configuration that constructs a query-aware [`DistancesUnordered`] visitor. +pub trait SearchStrategy<'a, P, T>: Send + Sync where P: DataProvider, { - /// The reference element shape on which [`Self::QueryComputer`] computes - /// distances. - type ElementRef<'a>; - - /// The concrete query-computer type. - type QueryComputer: for<'a> PreprocessedDistanceFunction, f32> - + Send - + Sync - + 'static; - - /// The error type for [`Self::build_query_computer`]. - type QueryComputerError: StandardError; - - /// The visitor type produced by [`Self::create_visitor`]. - type Visitor<'a>: for<'b> DistancesUnordered< - Self::QueryComputer, - ElementRef<'b> = Self::ElementRef<'b>, - Id = P::InternalId, - > - where - Self: 'a, - P: 'a; + /// The query-aware visitor used to execute the scan. + type Visitor: DistancesUnordered; - /// The error type for [`Self::create_visitor`]. + /// An error that can occur while constructing [`Self::Visitor`]. type Error: StandardError; - /// Construct a fresh visitor over `provider` for the given request `context`. - fn create_visitor<'a>( + /// Construct a fresh visitor for `query`. + /// + /// # Errors + /// + /// Returns an error when query preprocessing or visitor initialization fails. + fn create_visitor( &'a self, provider: &'a P, context: &'a P::Context, - ) -> Result, Self::Error>; - - /// Construct the per-query computer. - fn build_query_computer( - &self, query: T, - ) -> Result; + ) -> Result; } #[cfg(test)] mod tests { - //! Direct [`DistancesUnordered`] impls over a few in-memory fixtures: a - //! happy-path scanner over `&[f32]` elements, a scanner whose `ElementRef<'a>` - //! is a lifetime-carrying non-reference type, and a scanner that fails - //! mid-stream. - - use std::marker::PhantomData; + //! Direct [`DistancesUnordered`] impls over in-memory fixtures, including a + //! happy-path scanner and one that fails mid-stream. use diskann_utils::future::SendFuture; use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; @@ -121,27 +83,23 @@ mod tests { /// Scans `items` in order, scoring each with the supplied computer. struct Scanner { items: Vec<(u32, Vec)>, + computer: ::QueryDistance, } impl HasId for Scanner { type Id = u32; } - impl DistancesUnordered<::QueryDistance> for Scanner { - type ElementRef<'a> = &'a [f32]; + impl DistancesUnordered for Scanner { type Error = Infallible; - fn distances_unordered( - &mut self, - computer: &::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { async move { for (id, v) in &self.items { - let dist = computer.evaluate_similarity(v.as_slice()); + let dist = self.computer.evaluate_similarity(v.as_slice()); f(*id, dist); } Ok(()) @@ -153,25 +111,69 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn distances_unordered_scanner() { let query = vec![0.5_f32, 0.9]; - let computer = f32::query_distance(&query, Metric::L2); + let expected_computer = f32::query_distance(&query, Metric::L2); let expected: Vec<(u32, f32)> = sample_items() .into_iter() - .map(|(id, v)| (id, computer.evaluate_similarity(v.as_slice()))) + .map(|(id, v)| (id, expected_computer.evaluate_similarity(v.as_slice()))) .collect(); let mut scanner = Scanner { items: sample_items(), + computer: f32::query_distance(&query, Metric::L2), }; let mut seen: Vec<(u32, f32)> = Vec::new(); scanner - .distances_unordered(&computer, |id, d| seen.push((id, d))) + .distances_unordered(|id, d| seen.push((id, d))) .await .unwrap(); assert_eq!(seen, expected); } + struct BorrowingScanner<'a> { + items: &'a [(u32, f32)], + query: &'a f32, + } + + impl HasId for BorrowingScanner<'_> { + type Id = u32; + } + + impl DistancesUnordered for BorrowingScanner<'_> { + type Error = Infallible; + + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> + where + F: Send + FnMut(Self::Id, f32), + { + async move { + for (id, value) in self.items { + f(*id, (*value - *self.query).abs()); + } + Ok(()) + } + } + } + + #[tokio::test] + async fn accessor_can_borrow_query_state() { + let items = [(10, 1.0), (11, 4.0)]; + let query = 2.0; + let mut scanner = BorrowingScanner { + items: &items, + query: &query, + }; + let mut seen = Vec::new(); + + scanner + .distances_unordered(|id, distance| seen.push((id, distance))) + .await + .unwrap(); + + assert_eq!(seen, [(10, 1.0), (11, 2.0)]); + } + /////////////////////////// // Failing scanner // /////////////////////////// @@ -189,21 +191,17 @@ mod tests { struct Failing { items: Vec<(u32, Vec)>, fail_after: usize, + computer: ::QueryDistance, } impl HasId for Failing { type Id = u32; } - impl DistancesUnordered<::QueryDistance> for Failing { - type ElementRef<'a> = &'a [f32]; + impl DistancesUnordered for Failing { type Error = Boom; - fn distances_unordered( - &mut self, - computer: &::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { @@ -212,7 +210,7 @@ mod tests { if i == self.fail_after { return Err(Boom(*id)); } - let dist = computer.evaluate_similarity(v.as_slice()); + let dist = self.computer.evaluate_similarity(v.as_slice()); f(*id, dist); } Ok(()) @@ -227,14 +225,12 @@ mod tests { let mut scanner = Failing { items: sample_items(), fail_after: 1, // Yield item 0 successfully, fail on item 1. + computer: f32::query_distance(&[0.0, 0.0], Metric::L2), }; - let query = vec![0.0_f32, 0.0]; - let computer = f32::query_distance(&query, Metric::L2); - let mut seen: Vec = Vec::new(); let err = scanner - .distances_unordered(&computer, |id, _d| seen.push(id)) + .distances_unordered(|id, _d| seen.push(id)) .await .expect_err("Failing scanner must surface its error"); @@ -245,106 +241,4 @@ mod tests { "the closure must only see items yielded before the failure", ); } - - ///////////////////////////////////////////// - // Lifetime-carrying concrete `ElementRef` // - ///////////////////////////////////////////// - - struct View<'a> { - ptr: *const f32, - len: usize, - _phantom: PhantomData<&'a [f32]>, - } - - // SAFETY: `View<'a>` semantically carries a `&'a [f32]`, which is `Send + Sync`. - unsafe impl Send for View<'_> {} - unsafe impl Sync for View<'_> {} - - /// Computer that reconstructs a `&[f32]` from a [`View`]'s ptr+len and - /// computes inner product against a stored query. - struct ViewComputer { - query: Vec, - } - - impl<'a> PreprocessedDistanceFunction, f32> for ViewComputer { - fn evaluate_similarity(&self, v: View<'a>) -> f32 { - // SAFETY: `v.ptr` / `v.len` were produced from a `&'a [f32]` held by the - // scanner that owns the backing `Vec`; the phantom lifetime ties this view - // to that borrow, so the slice is valid for the duration of this call. - let s = unsafe { std::slice::from_raw_parts(v.ptr, v.len) }; - s.iter().zip(&self.query).map(|(a, b)| a * b).sum() - } - } - - /// Scans `rows`, yielding a [`View`] tied (via its phantom lifetime) to the - /// borrow of the underlying `Vec`. - struct ViewScanner { - rows: Vec<(u32, Vec)>, - } - - impl ViewScanner { - fn iter<'a>(&self) -> impl Iterator)> { - self.rows.iter().map(|(x, y)| { - ( - *x, - View { - ptr: y.as_ptr(), - len: y.len(), - _phantom: PhantomData, - }, - ) - }) - } - } - - impl HasId for ViewScanner { - type Id = u32; - } - - impl DistancesUnordered for ViewScanner { - type ElementRef<'a> = View<'a>; - type Error = Infallible; - - fn distances_unordered( - &mut self, - computer: &ViewComputer, - mut f: F, - ) -> impl SendFuture> - where - F: Send + FnMut(Self::Id, f32), - { - async move { - for (id, v) in self.iter() { - f(id, computer.evaluate_similarity(v)); - } - Ok(()) - } - } - } - - #[tokio::test] - async fn distances_unordered_lifetime_carrying_element_ref() { - let mut scanner = ViewScanner { - rows: vec![ - (10, vec![1.0, 0.0]), - (11, vec![0.5, 0.5]), - (12, vec![0.0, 2.0]), - ], - }; - let computer = ViewComputer { - query: vec![1.0, 3.0], - }; - let expected: Vec<(u32, f32)> = vec![ - (10, 1.0 * 1.0 + 0.0 * 3.0), - (11, 0.5 * 1.0 + 0.5 * 3.0), - (12, 0.0 * 1.0 + 2.0 * 3.0), - ]; - - let mut seen: Vec<(u32, f32)> = Vec::new(); - scanner - .distances_unordered(&computer, |id, d| seen.push((id, d))) - .await - .unwrap(); - assert_eq!(seen, expected); - } } diff --git a/diskann/src/flat/test/provider.rs b/diskann/src/flat/test/provider.rs index 3bd0f3d31e..4c730bc8e8 100644 --- a/diskann/src/flat/test/provider.rs +++ b/diskann/src/flat/test/provider.rs @@ -283,25 +283,32 @@ pub struct Visitor<'a> { provider: &'a Provider, transient_ids: Option>>, get_element: LocalCounter<'a>, + computer: ::QueryDistance, } impl<'a> Visitor<'a> { /// Construct a visitor with no fault injection. - pub fn new(provider: &'a Provider) -> Self { + pub fn new(provider: &'a Provider, computer: ::QueryDistance) -> Self { Self { provider, transient_ids: None, get_element: provider.get_element.local(), + computer, } } /// Construct a visitor that returns a [`TransientGetError`] for any id in /// `transient_ids`. Other ids behave normally. - pub fn flaky(provider: &'a Provider, transient_ids: Cow<'a, HashSet>) -> Self { + pub fn flaky( + provider: &'a Provider, + transient_ids: Cow<'a, HashSet>, + computer: ::QueryDistance, + ) -> Self { Self { provider, transient_ids: Some(transient_ids), get_element: provider.get_element.local(), + computer, } } } @@ -319,15 +326,10 @@ impl HasId for Visitor<'_> { type Id = u32; } -impl DistancesUnordered<::QueryDistance> for Visitor<'_> { - type ElementRef<'a> = &'a [f32]; +impl DistancesUnordered for Visitor<'_> { type Error = AccessError; - fn distances_unordered( - &mut self, - computer: &::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { @@ -340,7 +342,7 @@ impl DistancesUnordered<::QueryDistance> for Visitor<'_> { return Err(AccessError::Transient(TransientGetError::new(id))); } self.get_element.increment(); - let dist = computer.evaluate_similarity(vector); + let dist = self.computer.evaluate_similarity(vector); f(id, dist); } Ok(()) @@ -352,8 +354,7 @@ impl DistancesUnordered<::QueryDistance> for Visitor<'_> { // Strategy // ////////////// -/// Error from [`Strategy::create_visitor`] or [`Strategy::build_query_computer`] -/// when dimensions don't match. +/// Error from [`Strategy::create_visitor`] when dimensions don't match. #[derive(Debug, Clone, Error)] #[error("dimension mismatch: strategy expects {expected}, got {actual}")] pub struct StrategyError { @@ -390,18 +391,16 @@ impl Strategy { } } -impl SearchStrategy for Strategy { - type ElementRef<'a> = &'a [f32]; - type QueryComputer = ::QueryDistance; - type QueryComputerError = StrategyError; - type Visitor<'a> = Visitor<'a>; +impl<'a> SearchStrategy<'a, Provider, &'a [f32]> for Strategy { + type Visitor = Visitor<'a>; type Error = StrategyError; - fn create_visitor<'a>( + fn create_visitor( &'a self, provider: &'a Provider, _context: &'a Context, - ) -> Result, Self::Error> { + query: &'a [f32], + ) -> Result { let actual = provider.dim(); if actual != self.dim { return Err(StrategyError { @@ -409,23 +408,16 @@ impl SearchStrategy for Strategy { actual, }); } - let visitor = match &self.transient_ids { - Some(ids) => Visitor::flaky(provider, Cow::Borrowed(ids)), - None => Visitor::new(provider), - }; - Ok(visitor) - } - - fn build_query_computer( - &self, - from: &[f32], - ) -> Result { - if from.len() != self.dim { + if query.len() != self.dim { return Err(StrategyError { expected: self.dim, - actual: from.len(), + actual: query.len(), }); } - Ok(f32::query_distance(from, Metric::L2)) + let computer = f32::query_distance(query, Metric::L2); + Ok(match &self.transient_ids { + Some(ids) => Visitor::flaky(provider, Cow::Borrowed(ids), computer), + None => Visitor::new(provider, computer), + }) } } diff --git a/rfcs/00983-flat-search.md b/rfcs/00983-flat-search.md index 3d56e2db87..c2b813aed7 100644 --- a/rfcs/00983-flat-search.md +++ b/rfcs/00983-flat-search.md @@ -21,7 +21,7 @@ The problem-statement here is simple: provide first-class support for sequential ### 1.3 Goals -1. Define a fused iterate-and-score primitive — `flat::DistancesUnordered` — that +1. Define a fused iterate-and-score primitive — `flat::DistancesUnordered` — that mirrors the role `Accessor` plays for graph search but exposes a sequential scan-and-score operation instead of random access. 2. Provide flat-search algorithm implementations built on the new primitives, so consumers can use this against their own providers / backends. @@ -35,28 +35,22 @@ The module exposes three layers: | Layer | Trait | Role | |-------|-------|------| -| Backend | `DistancesUnordered` | Scan-and-score primitive | -| Factory | `SearchStrategy` | Per-query visitor + computer construction | +| Backend | `DistancesUnordered` | Query-aware scan-and-score primitive | +| Factory | `SearchStrategy<'a, P, T>` | Per-query visitor construction | | Algorithm | `FlatIndex::knn_search` | Brute-force top-k | -### 2.1 `DistancesUnordered` — the core scanning trait +### 2.1 `DistancesUnordered` — the core scanning trait -The single required trait for flat search. It is generic over a **computer type** `C` -rather than a query type — the algorithm supplies a pre-built computer and the visitor -drives the scan. +The single required trait for flat search. The visitor owns any query-specific state and +drives the scan, allowing a backend to fuse query preprocessing, filtering, batching, +data access, and distance computation. ```rust -pub trait DistancesUnordered: Send + Sync -where - C: for<'a> PreprocessedDistanceFunction, f32>, -{ - type ElementRef<'a>; - type Id; +pub trait DistancesUnordered: HasId + Send + Sync { type Error: ToRanked + Debug + Send + Sync + 'static; fn distances_unordered( &mut self, - computer: &C, f: F, ) -> impl SendFuture> where @@ -67,52 +61,36 @@ where Key differences from the graph-side `Accessor` path: - No random access — the visitor drives the entire scan internally. -- `ElementRef<'a>` and `Id` live on `DistancesUnordered` itself, decoupling the - scan-and-score primitive from `HasId` and from any provider-specific id type. A - visitor is free to yield ids derived from but not equal to its provider's - `InternalId`. We expect this constraint to go away once we're able to clean up the `VectorId` trait - and its restrictive bounds - i.e. expects id to be scalar-like. +- The query computer is an implementation detail of the visitor rather than a separate + algorithm-level abstraction. +- The visitor's id is constrained to the provider's `InternalId`, allowing the shared + top-k and post-processing machinery to operate directly on provider ids. -### 2.2 `SearchStrategy` — per-query factory +### 2.2 `SearchStrategy<'a, P, T>` — per-query factory -The strategy owns both visitor construction and query-computer construction: +The strategy constructs a query-aware visitor. The trait-level lifetime allows a visitor +to borrow the query or provider when that is the backend's natural representation. ```rust -pub trait SearchStrategy: Send + Sync +pub trait SearchStrategy<'a, P, T>: Send + Sync where P: DataProvider, { - type ElementRef<'a>; - type Id; - type QueryComputer: for<'a> PreprocessedDistanceFunction, f32> - + Send + Sync + 'static; - type QueryComputerError: StandardError; - - type Visitor<'a>: for<'b> DistancesUnordered< - Self::QueryComputer, - ElementRef<'b> = Self::ElementRef<'b>, - Id = Self::Id, - > - where Self: 'a, P: 'a; - + type Visitor: DistancesUnordered; type Error: StandardError; - fn create_visitor<'a>( + fn create_visitor( &'a self, provider: &'a P, context: &'a P::Context, - ) -> Result, Self::Error>; - - fn build_query_computer( - &self, query: T, - ) -> Result; + ) -> Result; } ``` -`build_query_computer` lives on the **strategy**, not the visitor. This keeps the -visitor free of any distance-computation trait bounds — it only needs to implement -`DistancesUnordered` for the strategy's computer type. +Passing the query during visitor construction keeps all backend-specific execution state +behind one coarse interface. This follows the same ownership model as graph +`SearchAccessor`: the backend receives the query once and computes its own distances. ### 2.3 `FlatIndex::knn_search` @@ -121,35 +99,29 @@ method is the brute-force top-k algorithm: ```rust impl FlatIndex

{ - pub fn knn_search( - &self, + pub fn knn_search<'a, S, T, O, PP, OB>( + &'a self, k: NonZeroUsize, - strategy: &S, - context: &P::Context, + strategy: &'a S, + processor: PP, + context: &'a P::Context, query: T, - output: &mut OB, - ) -> impl SendFuture> + output: &'a mut OB, + ) -> impl SendFuture> + 'a where - S: SearchStrategy, - S::Id: NeighborPriorityQueueIdType, - T: Send + Sync, - OB: SearchOutputBuffer + Send + ?Sized; + S: SearchStrategy<'a, P, T> + 'a, + T: Copy + Send + Sync + 'a, + O: Send + 'a, + PP: SearchPostProcess + Send + Sync + 'a, + OB: SearchOutputBuffer + Send + ?Sized + 'a; } ``` Algorithm: -1. `strategy.create_visitor(&provider, context)` — acquire the scanning visitor. -2. `strategy.build_query_computer(query)` — preprocess the query into a computer. -3. `visitor.distances_unordered(&computer, |id, dist| queue.insert(...))` — full scan. -4. Drain the priority queue into `output` in best-first order. - -**No post-processing parameter (yet).** Currently `knn_search` writes -`(S::Id, f32)` directly into the `SearchOutputBuffer`. Once the graph-search -trait refactor in [PR #1076](https://github.com/microsoft/DiskANN/pull/1076) -lands, `knn_search` will accept an optional `SearchPostProcess` parameter -(the same trait graph search uses), enabling id remapping, re-ranking, and -other transformations as a composable layer. +1. `strategy.create_visitor(&provider, context, query)` — acquire the query-aware visitor. +2. `visitor.distances_unordered(|id, dist| queue.insert(...))` — full scan. +3. Run `SearchPostProcess` over the best candidates to populate `output`. #### Call-chain diagram @@ -158,16 +130,15 @@ other transformations as a composable layer. ───── ──── DiskANNIndex::search FlatIndex::knn_search - │ │ - ▼ ▼ + │ │ + ▼ ▼ graph::glue::SearchStrategy flat::SearchStrategy ::search_accessor ::create_visitor - │ ::build_query_computer - ▼ │ - Accessor + BuildQueryComputer ▼ - → QueryComputer DistancesUnordered - │ ::distances_unordered(&computer, f) - ▼ │ + ▼ │ + SearchAccessor ▼ + │ DistancesUnordered + │ ::distances_unordered(f) + ▼ │ ExpandBeam::expand_beam │ (greedy beam, random access) │ │ │ @@ -175,19 +146,18 @@ other transformations as a composable layer. NeighborPriorityQueue NeighborPriorityQueue │ │ ▼ ▼ - SearchPostProcess SearchPostProcess (planned, PR #1076) + SearchPostProcess SearchPostProcess → SearchOutputBuffer → SearchOutputBuffer ``` ## Trade-offs -### No built-in post-processing (temporary) +### Backend-owned query computation -`knn_search` currently writes `(InternalId, f32)` directly. Once the graph-search -trait refactor in [PR #1076](https://github.com/microsoft/DiskANN/pull/1076) lands -and stabilizes a shared `SearchPostProcess` trait, `knn_search` will gain an optional -post-processor parameter matching the graph-search signature. Until then, callers that -need id remapping or re-ranking compose it externally. +Owning query computation in the visitor gives the backend a coarse optimization boundary +and avoids coupling the algorithm to a particular element-reference or distance-computer +shape. Implementations that only need an element-wise computer store it in the visitor; +backends with bulk kernels can bypass that abstraction entirely. ### Reusing `DataProvider` @@ -195,11 +165,10 @@ The design requires implementations to provide `InternalId` / `ExternalId` conve This is arguably too restrictive for some flat-index consumers, but avoids introducing a second provider trait. -### Expand `ElementRef` and `QueryComputer` to support batched distance computation? +### Batched distance computation -The design for `DistancesUnordered` assumes the computer acts on single vectors. An alternative is to allow the computer to work -over batches, enabling (potentially) better cache utilization. Backends that need this can implement `DistancesUnordered` -directly with an optimized bulk loop. Some refactoring for the bounds on `DistancesUnordered` is needed here. +`DistancesUnordered` deliberately does not prescribe element-wise computation. Backends can +score batches directly to improve cache utilization or fuse storage-specific operations. ### Intra-query parallelism @@ -208,7 +177,5 @@ scan. A parallel variant would need a different trait shape (e.g. splitting the shards). This is left for future work. ## Future Work -- **Post-processing support** — once [PR #1076](https://github.com/microsoft/DiskANN/pull/1076) lands, add a `SearchPostProcess` parameter to `knn_search` so flat search can share the same id-remapping / re-ranking infrastructure as graph search. - Support for other flat-search algorithms like filtered, range, and diverse flat algorithms as additional methods on `FlatIndex`. - Index build — this is just one part of the picture; more work needs to be done around how this fits in with any traits / interface we need for index build. -