From 316e21f921f0ad762118b1c41b4e76586cc94424 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Thu, 23 Apr 2026 19:52:14 -0400 Subject: [PATCH 1/5] Move ObjectPool from diskann to diskann-utils Relocates the object pool module so that it is available to crates that depend on diskann-utils but not diskann (notably diskann-quantization, which will gain pool-aware distance-table allocation in a follow-up). diskann::utils::object_pool stays as a re-export for backwards compatibility. Direct importers in diskann-providers, diskann-disk, and diskann-garnet are switched to use diskann_utils::object_pool directly. Internal diskann users continue to use the re-export. --- diskann-disk/src/search/provider/disk_provider.rs | 6 ++---- diskann-garnet/src/provider.rs | 6 ++---- .../graph/provider/async_/bf_tree/quant_vector_provider.rs | 7 ++----- .../provider/async_/fast_memory_quant_vector_provider.rs | 7 ++----- .../graph/provider/async_/memory_quant_vector_provider.rs | 3 ++- diskann-providers/src/model/pq/distance/common.rs | 2 +- diskann-providers/src/model/pq/distance/dynamic.rs | 3 ++- diskann-providers/src/model/pq/distance/innerproduct.rs | 6 ++---- diskann-providers/src/model/pq/distance/l2.rs | 6 ++---- diskann-utils/src/lib.rs | 1 + {diskann/src/utils => diskann-utils/src}/object_pool.rs | 0 diskann/src/graph/index.rs | 2 +- diskann/src/graph/search/scratch.rs | 3 ++- diskann/src/utils/mod.rs | 2 -- 14 files changed, 21 insertions(+), 33 deletions(-) rename {diskann/src/utils => diskann-utils/src}/object_pool.rs (100%) diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 0403a5019..b7b30e94a 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -31,10 +31,7 @@ use diskann::{ Accessor, BuildQueryComputer, DataProvider, DefaultContext, DelegateNeighbor, HasId, NeighborAccessor, NoopGuard, }, - utils::{ - object_pool::{ObjectPool, PoolOption, TryAsPooled}, - IntoUsize, VectorRepr, - }, + utils::{IntoUsize, VectorRepr}, ANNError, ANNResult, }; use diskann_providers::storage::StorageReadProvider; @@ -42,6 +39,7 @@ use diskann_providers::{ model::{compute_pq_distance, compute_pq_distance_for_pq_coordinates}, storage::{get_compressed_pq_file, get_disk_index_file, get_pq_pivot_file, LoadWith}, }; +use diskann_utils::object_pool::{ObjectPool, PoolOption, TryAsPooled}; use crate::search::pq::{quantizer_preprocess, PQData, PQScratch}; use diskann_vector::{distance::Metric, DistanceFunction, PreprocessedDistanceFunction}; diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 1bb0679ad..e2522b147 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -20,13 +20,11 @@ use diskann::{ Accessor, BuildDistanceComputer, BuildQueryComputer, DataProvider, DelegateNeighbor, Delete, ElementStatus, HasId, NeighborAccessor, NeighborAccessorMut, NoopGuard, SetElement, }, - utils::{ - VectorRepr, - object_pool::{AsPooled, ObjectPool, PooledRef, Undef}, - }, + utils::VectorRepr, }; use diskann_providers::model::graph::provider::async_::common::FullPrecision; use diskann_utils::Reborrow; +use diskann_utils::object_pool::{AsPooled, ObjectPool, PooledRef, Undef}; use diskann_vector::{PreprocessedDistanceFunction, contains::ContainsSimd, distance::Metric}; use std::{ future, mem, diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs index 57e5120d9..0d26fa680 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs @@ -9,12 +9,9 @@ use std::sync::Arc; use bf_tree::{BfTree, Config}; use bytemuck::bytes_of; -use diskann::{ - ANNError, ANNErrorKind, ANNResult, - error::IntoANNResult, - utils::{VectorRepr, object_pool::ObjectPool}, -}; +use diskann::{ANNError, ANNErrorKind, ANNResult, error::IntoANNResult, utils::VectorRepr}; use diskann_quantization::CompressInto; +use diskann_utils::object_pool::ObjectPool; use diskann_vector::distance::Metric; use thiserror::Error; diff --git a/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs index e70776ded..62d9c05e1 100644 --- a/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs @@ -14,12 +14,9 @@ use std::sync::{Arc, Mutex}; use crate::storage::{StorageReadProvider, StorageWriteProvider}; -use diskann::{ - ANNError, ANNResult, - error::IntoANNResult, - utils::{VectorRepr, object_pool::ObjectPool}, -}; +use diskann::{ANNError, ANNResult, error::IntoANNResult, utils::VectorRepr}; use diskann_quantization::CompressInto; +use diskann_utils::object_pool::ObjectPool; use diskann_vector::distance::Metric; use super::common::{AlignedMemoryVectorStore, TestCallCount}; diff --git a/diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs index bdd191f32..dca7ea762 100644 --- a/diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs @@ -14,9 +14,10 @@ use crate::storage::{StorageReadProvider, StorageWriteProvider}; use arc_swap::{ArcSwap, Guard}; #[cfg(test)] use diskann::utils::VectorRepr; -use diskann::{ANNError, ANNResult, utils::object_pool::ObjectPool}; +use diskann::{ANNError, ANNResult}; #[cfg(test)] use diskann_quantization::CompressInto; +use diskann_utils::object_pool::ObjectPool; use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; use super::{VectorGuard, common::TestCallCount}; diff --git a/diskann-providers/src/model/pq/distance/common.rs b/diskann-providers/src/model/pq/distance/common.rs index 2888e6836..f36b5b386 100644 --- a/diskann-providers/src/model/pq/distance/common.rs +++ b/diskann-providers/src/model/pq/distance/common.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann::utils::object_pool::{self, ObjectPool}; +use diskann_utils::object_pool::{self, ObjectPool}; use crate::model::pq::fixed_chunk_pq_table::FixedChunkPQTable; diff --git a/diskann-providers/src/model/pq/distance/dynamic.rs b/diskann-providers/src/model/pq/distance/dynamic.rs index bc9d83b94..cefacacde 100644 --- a/diskann-providers/src/model/pq/distance/dynamic.rs +++ b/diskann-providers/src/model/pq/distance/dynamic.rs @@ -5,7 +5,8 @@ use std::{ops::Deref, sync::Arc}; -use diskann::{ANNResult, utils::object_pool::ObjectPool}; +use diskann::ANNResult; +use diskann_utils::object_pool::ObjectPool; use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; // Concrete implementations diff --git a/diskann-providers/src/model/pq/distance/innerproduct.rs b/diskann-providers/src/model/pq/distance/innerproduct.rs index 8849f30ff..7d9abfc7e 100644 --- a/diskann-providers/src/model/pq/distance/innerproduct.rs +++ b/diskann-providers/src/model/pq/distance/innerproduct.rs @@ -5,10 +5,8 @@ use std::{ops::Deref, sync::Arc}; -use diskann::{ - ANNResult, - utils::object_pool::{self, ObjectPool, PoolOption}, -}; +use diskann::ANNResult; +use diskann_utils::object_pool::{self, ObjectPool, PoolOption}; use diskann_vector::PreprocessedDistanceFunction; use super::common::get_lookup_table_size; diff --git a/diskann-providers/src/model/pq/distance/l2.rs b/diskann-providers/src/model/pq/distance/l2.rs index e551ee35a..543534796 100644 --- a/diskann-providers/src/model/pq/distance/l2.rs +++ b/diskann-providers/src/model/pq/distance/l2.rs @@ -5,10 +5,8 @@ use std::{ops::Deref, sync::Arc}; -use diskann::{ - ANNResult, - utils::object_pool::{self, ObjectPool, PoolOption}, -}; +use diskann::ANNResult; +use diskann_utils::object_pool::{self, ObjectPool, PoolOption}; use diskann_vector::PreprocessedDistanceFunction; use super::common::get_lookup_table_size; diff --git a/diskann-utils/src/lib.rs b/diskann-utils/src/lib.rs index cd8c1b84d..089e0dece 100644 --- a/diskann-utils/src/lib.rs +++ b/diskann-utils/src/lib.rs @@ -15,6 +15,7 @@ pub use lifetime::WithLifetime; pub mod future; pub mod io; +pub mod object_pool; pub mod sampling; // Views diff --git a/diskann/src/utils/object_pool.rs b/diskann-utils/src/object_pool.rs similarity index 100% rename from diskann/src/utils/object_pool.rs rename to diskann-utils/src/object_pool.rs diff --git a/diskann/src/graph/index.rs b/diskann/src/graph/index.rs index cd81ba5f1..eca65b8e6 100644 --- a/diskann/src/graph/index.rs +++ b/diskann/src/graph/index.rs @@ -52,9 +52,9 @@ use crate::{ utils::{ IntoUsize, TryIntoVectorId, VectorId, async_tools::{self, DynamicBalancer}, - object_pool::{ObjectPool, PooledRef}, }, }; +use diskann_utils::object_pool::{ObjectPool, PooledRef}; #[derive(Debug)] pub struct DiskANNIndex { diff --git a/diskann/src/graph/search/scratch.rs b/diskann/src/graph/search/scratch.rs index 2a4706821..98a5b3127 100644 --- a/diskann/src/graph/search/scratch.rs +++ b/diskann/src/graph/search/scratch.rs @@ -11,8 +11,9 @@ use std::collections::VecDeque; use crate::{ neighbor::{Neighbor, NeighborPriorityQueue}, - utils::{VectorId, object_pool::AsPooled}, + utils::VectorId, }; +use diskann_utils::object_pool::AsPooled; use hashbrown::HashSet; /// In-mem index related limits diff --git a/diskann/src/utils/mod.rs b/diskann/src/utils/mod.rs index bf6ac5c40..86b78f0ab 100644 --- a/diskann/src/utils/mod.rs +++ b/diskann/src/utils/mod.rs @@ -3,8 +3,6 @@ * Licensed under the MIT license. */ -pub mod object_pool; - pub mod async_tools; #[allow(clippy::module_inception)] From ca3a5293bbf6cc201491f3824c1abcc0b2c62871 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Fri, 24 Apr 2026 12:28:50 -0400 Subject: [PATCH 2/5] move chunks_offsets to diskann_quantization --- .../src/backend/exhaustive/product.rs | 2 +- .../src/search/pq/quantizer_preprocess.rs | 19 +++++------ diskann-providers/src/model/mod.rs | 2 +- .../src/model/pq/distance/test_utils.rs | 3 +- diskann-providers/src/model/pq/mod.rs | 7 ++-- .../src/model/pq/pq_construction.rs | 30 +---------------- diskann-quantization/src/views.rs | 33 +++++++++++++++++++ 7 files changed, 50 insertions(+), 46 deletions(-) diff --git a/diskann-benchmark/src/backend/exhaustive/product.rs b/diskann-benchmark/src/backend/exhaustive/product.rs index 4723753a0..f339a3c44 100644 --- a/diskann-benchmark/src/backend/exhaustive/product.rs +++ b/diskann-benchmark/src/backend/exhaustive/product.rs @@ -87,7 +87,7 @@ mod imp { 5, ); - let offsets = diskann_providers::model::pq::calculate_chunk_offsets_auto( + let offsets = diskann_quantization::views::calculate_chunk_offsets_auto( data.ncols(), input.num_pq_chunks.get(), ); diff --git a/diskann-disk/src/search/pq/quantizer_preprocess.rs b/diskann-disk/src/search/pq/quantizer_preprocess.rs index cc454ea7b..82f7cc7c5 100644 --- a/diskann-disk/src/search/pq/quantizer_preprocess.rs +++ b/diskann-disk/src/search/pq/quantizer_preprocess.rs @@ -3,6 +3,13 @@ * Licensed under the MIT license. */ +//! PQ quantizer query preprocessing. +//! +//! Prior to the introduction of the [`quantizer_preprocess`] method, the disk index was +//! hard-coded to use L2 distance for comparisons. We're keeping that behavior here - +//! treating `Cosine` and `CosineNormalized` as L2 until a more thorough evaluation can +//! be made. + use diskann::ANNResult; use diskann_vector::distance::Metric; @@ -33,11 +40,7 @@ pub fn quantizer_preprocess( .bridge_err()?; match metric { - // Prior to the introduction of the `quantizer_preprocess` method, the - // disk index was hard-coded to use L2 distance for comparisons. - // - // We're keeping that behavior here - treating `Cosine` and `CosineNormalized` - // as L2 until a more thorough evaluation can be made. + // Cosine and CosineNormalized fall back to L2; see module docs. Metric::L2 | Metric::Cosine | Metric::CosineNormalized => { table.process_into::( &pq_scratch.rotated_query[..dim], @@ -54,11 +57,7 @@ pub fn quantizer_preprocess( } PQTable::Fixed(table) => { match metric { - // Prior to the introduction of the `quantizer_preprocess` method, the - // disk index was hard-coded to use L2 distance for comparisons. - // - // We're keeping that behavior here - treating `Cosine` and `CosineNormalized` - // as L2 until a more thorough evaluation can be made. + // Cosine and CosineNormalized fall back to L2; see module docs. Metric::L2 | Metric::Cosine | Metric::CosineNormalized => { // The scratch only stores the aligned dimension. However, preprocessing // wants the actual dimension used, so we have to shrink the rotated query diff --git a/diskann-providers/src/model/mod.rs b/diskann-providers/src/model/mod.rs index f6ae2be75..e9fceb869 100644 --- a/diskann-providers/src/model/mod.rs +++ b/diskann-providers/src/model/mod.rs @@ -11,7 +11,7 @@ pub use configuration::IndexConfiguration; pub mod pq; pub use pq::{ FixedChunkPQTable, GeneratePivotArguments, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, - NUM_PQ_CENTROIDS, accum_row_inplace, calculate_chunk_offsets_auto, compute_pq_distance, + NUM_PQ_CENTROIDS, accum_row_inplace, compute_pq_distance, compute_pq_distance_for_pq_coordinates, direct_distance_impl, distance, generate_pq_data_from_pivots_from_membuf, generate_pq_data_from_pivots_from_membuf_batch, generate_pq_pivots, generate_pq_pivots_from_membuf, diff --git a/diskann-providers/src/model/pq/distance/test_utils.rs b/diskann-providers/src/model/pq/distance/test_utils.rs index 76c240e1f..677fc32a2 100644 --- a/diskann-providers/src/model/pq/distance/test_utils.rs +++ b/diskann-providers/src/model/pq/distance/test_utils.rs @@ -13,7 +13,8 @@ use diskann_vector::{ use rand::{Rng, distr::Distribution}; use rand_distr::{Normal, Uniform}; -use crate::model::{FixedChunkPQTable, pq::calculate_chunk_offsets_auto}; +use crate::model::FixedChunkPQTable; +use diskann_quantization::views::calculate_chunk_offsets_auto; /// We need a way to generate random queries. /// diff --git a/diskann-providers/src/model/pq/mod.rs b/diskann-providers/src/model/pq/mod.rs index 6338e39ec..6c1962a4a 100644 --- a/diskann-providers/src/model/pq/mod.rs +++ b/diskann-providers/src/model/pq/mod.rs @@ -11,10 +11,9 @@ pub use fixed_chunk_pq_table::{ mod pq_construction; pub use pq_construction::{ MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS, accum_row_inplace, - calculate_chunk_offsets, calculate_chunk_offsets_auto, generate_pq_data_from_pivots, - generate_pq_data_from_pivots_from_membuf, generate_pq_data_from_pivots_from_membuf_batch, - generate_pq_pivots, generate_pq_pivots_from_membuf, get_chunk_from_training_data, - move_train_data_by_centroid, + generate_pq_data_from_pivots, generate_pq_data_from_pivots_from_membuf, + generate_pq_data_from_pivots_from_membuf_batch, generate_pq_pivots, + generate_pq_pivots_from_membuf, get_chunk_from_training_data, move_train_data_by_centroid, }; /// all metadata of individual sub-component files is written in first 4KB for unified files diff --git a/diskann-providers/src/model/pq/pq_construction.rs b/diskann-providers/src/model/pq/pq_construction.rs index b2ab3da85..69b95c92a 100644 --- a/diskann-providers/src/model/pq/pq_construction.rs +++ b/diskann-providers/src/model/pq/pq_construction.rs @@ -19,6 +19,7 @@ use diskann::{ use diskann_quantization::{ CompressInto, product::{BasicTableView, TransposedTable, train::TrainQuantizer}, + views::calculate_chunk_offsets, }; use diskann_utils::{ io::Metadata, @@ -328,35 +329,6 @@ pub fn move_train_data_by_centroid( } } -/// Calculate the number of chunks for the product quantization algorithm. Returns a vector of offsets where -/// each offset corresponds to a chunk based on the index of the chunk in the vector. -/// -/// # Arguments -/// * `dimensions` Number of dimensions of the input data -/// * `num_pq_chunks` - Number of chunks that will be used in the PQ calculation. Each vector will be split into these -/// number of chunks and each chunk will be compressed down to one byte. -/// * `offsets` - An output vector of offsets, where the size is equal to the number of pq chunks + 1. -#[inline] -pub fn calculate_chunk_offsets(dimensions: usize, num_pq_chunks: usize, offsets: &mut [usize]) { - // Calculate each chunk's offset - // If we have 8 dimension and 3 chunks then offsets would be [0,3,6,8] - let mut chunk_offset: usize = 0; - offsets[0] = chunk_offset; - for chunk_index in 0..num_pq_chunks { - chunk_offset += dimensions / num_pq_chunks; - if chunk_index < (dimensions % num_pq_chunks) { - chunk_offset += 1; - } - offsets[chunk_index + 1] = chunk_offset; - } -} - -pub fn calculate_chunk_offsets_auto(dimensions: usize, num_pq_chunks: usize) -> Vec { - let mut offsets = vec![0; num_pq_chunks + 1]; - calculate_chunk_offsets(dimensions, num_pq_chunks, offsets.as_mut_slice()); - offsets -} - /// Add the row `y` to every row in `x`. /// /// # Panics diff --git a/diskann-quantization/src/views.rs b/diskann-quantization/src/views.rs index 6ef928345..f8882e8b4 100644 --- a/diskann-quantization/src/views.rs +++ b/diskann-quantization/src/views.rs @@ -205,6 +205,39 @@ impl<'a> From> for &'a [usize] { } } +/// Calculate the chunk offsets for the product quantization algorithm. Fills `offsets` +/// with the prefix-sum partitioning of `dimensions` into `num_pq_chunks` chunks, where +/// the first `dimensions % num_pq_chunks` chunks are one element larger than the rest. +/// +/// # Arguments +/// * `dimensions` - Number of dimensions of the input data. +/// * `num_pq_chunks` - Number of chunks that will be used in the PQ calculation. Each +/// vector will be split into these number of chunks and each chunk will be compressed +/// down to one byte. +/// * `offsets` - An output slice of offsets, where the length must equal `num_pq_chunks + 1`. +#[inline] +pub fn calculate_chunk_offsets(dimensions: usize, num_pq_chunks: usize, offsets: &mut [usize]) { + // Calculate each chunk's offset + // If we have 8 dimension and 3 chunks then offsets would be [0,3,6,8] + let mut chunk_offset: usize = 0; + offsets[0] = chunk_offset; + for chunk_index in 0..num_pq_chunks { + chunk_offset += dimensions / num_pq_chunks; + if chunk_index < (dimensions % num_pq_chunks) { + chunk_offset += 1; + } + offsets[chunk_index + 1] = chunk_offset; + } +} + +/// Allocating wrapper around [`calculate_chunk_offsets`] that returns a fresh +/// `Vec` of length `num_pq_chunks + 1`. +pub fn calculate_chunk_offsets_auto(dimensions: usize, num_pq_chunks: usize) -> Vec { + let mut offsets = vec![0; num_pq_chunks + 1]; + calculate_chunk_offsets(dimensions, num_pq_chunks, offsets.as_mut_slice()); + offsets +} + /////////////// // ChunkView // /////////////// From 4fe3345408e3e817708eb54138715137fb694464 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Fri, 24 Apr 2026 14:26:43 -0400 Subject: [PATCH 3/5] remove redundant deref impls in PQ dynamic distance --- .../fast_memory_quant_vector_provider.rs | 9 +-- .../src/model/pq/distance/dynamic.rs | 65 ------------------- 2 files changed, 3 insertions(+), 71 deletions(-) diff --git a/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs index 62d9c05e1..9da2296ac 100644 --- a/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs @@ -444,21 +444,18 @@ mod tests { // Query Computer. let c = provider.query_computer(&[-0.5, -0.5]).unwrap(); let expected: f32 = 1.5 * 1.5 * 2.0; - assert_eq!( - c.evaluate_similarity(&provider.get_vector_sync(3)), - expected - ); + assert_eq!(c.evaluate_similarity(provider.get_vector_sync(3)), expected); // Distance Computer. let d = provider.distance_computer(); assert_eq!( - d.evaluate_similarity(&provider.get_vector_sync(0), &provider.get_vector_sync(3)), + d.evaluate_similarity(provider.get_vector_sync(0), provider.get_vector_sync(3)), 2.0 ); let slice: &[f32] = &[-0.5, -0.5]; assert_eq!( - d.evaluate_similarity(slice, &provider.get_vector_sync(3)), + d.evaluate_similarity(slice, provider.get_vector_sync(3)), expected, ); } diff --git a/diskann-providers/src/model/pq/distance/dynamic.rs b/diskann-providers/src/model/pq/distance/dynamic.rs index cefacacde..352912846 100644 --- a/diskann-providers/src/model/pq/distance/dynamic.rs +++ b/diskann-providers/src/model/pq/distance/dynamic.rs @@ -101,25 +101,6 @@ where } } -impl PreprocessedDistanceFunction<&Vec, f32> for QueryComputer -where - T: Deref, -{ - fn evaluate_similarity(&self, changing: &Vec) -> f32 { - self.evaluate_similarity(changing.as_slice()) - } -} - -impl PreprocessedDistanceFunction<&&[u8], f32> for QueryComputer -where - T: Deref, -{ - fn evaluate_similarity(&self, changing: &&[u8]) -> f32 { - let changing: &[u8] = changing; - self.evaluate_similarity(changing) - } -} - /// Pre-dispatched distance functions for the `FixedChunkPQTable`. #[derive(Debug)] pub struct VTable { @@ -233,52 +214,6 @@ where } } -/// Perform a comparison between a full-precision vector and quantized vector. -impl DistanceFunction<&[f32], &&[u8], f32> for DistanceComputer -where - T: Deref, -{ - #[inline(always)] - fn evaluate_similarity(&self, fp: &[f32], q: &&[u8]) -> f32 { - let q: &[u8] = q; - self.evaluate_similarity(fp, q) - } -} - -impl DistanceFunction<&[f32], &Vec, f32> for DistanceComputer -where - T: Deref, -{ - #[inline(always)] - fn evaluate_similarity(&self, fp: &[f32], q: &Vec) -> f32 { - self.evaluate_similarity(fp, q.as_slice()) - } -} - -/// Perform a comparison between two quantized vectors. -impl DistanceFunction<&&[u8], &&[u8], f32> for DistanceComputer -where - T: Deref, -{ - #[inline(always)] - fn evaluate_similarity(&self, q0: &&[u8], q1: &&[u8]) -> f32 { - let q0: &[u8] = q0; - let q1: &[u8] = q1; - self.evaluate_similarity(q0, q1) - } -} - -/// Perform a comparison between two quantized vectors. -impl DistanceFunction<&Vec, &Vec, f32> for DistanceComputer -where - T: Deref, -{ - #[inline(always)] - fn evaluate_similarity(&self, q0: &Vec, q1: &Vec) -> f32 { - self.evaluate_similarity(q0.as_slice(), q1.as_slice()) - } -} - #[cfg(test)] mod tests { use std::marker::PhantomData; From 40dbba944d51e87cd350a682932986fd6ab4ddfc Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Fri, 24 Apr 2026 15:02:40 -0400 Subject: [PATCH 4/5] fix bf_tree gated test for removed deref impls --- .../provider/async_/bf_tree/quant_vector_provider.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs index 0d26fa680..edcb593d8 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs @@ -353,7 +353,7 @@ mod tests { let c = provider.query_computer(&[-0.5, -0.5]).unwrap(); let expected: f32 = 1.5 * 1.5 * 2.0; assert_eq!( - c.evaluate_similarity(&provider.get_vector_sync(3).unwrap()), + c.evaluate_similarity(provider.get_vector_sync(3).unwrap().as_slice()), expected ); @@ -361,15 +361,15 @@ mod tests { let d = provider.distance_computer(); assert_eq!( d.evaluate_similarity( - &provider.get_vector_sync(0).unwrap(), - &provider.get_vector_sync(3).unwrap() + provider.get_vector_sync(0).unwrap().as_slice(), + provider.get_vector_sync(3).unwrap().as_slice() ), 2.0 ); let slice: &[f32] = &[-0.5, -0.5]; assert_eq!( - d.evaluate_similarity(slice, &provider.get_vector_sync(3).unwrap()), + d.evaluate_similarity(slice, provider.get_vector_sync(3).unwrap().as_slice()), expected, ); } From bf082868f23fcc27f44933f0b7f85f0c639f8957 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Fri, 24 Apr 2026 22:57:29 -0400 Subject: [PATCH 5/5] move accum_row_inplace to diskann-utils; demote get_chunk_from_training_data to test helper --- .../src/storage/quant/pq/pq_generation.rs | 7 +- diskann-providers/src/model/mod.rs | 8 +-- diskann-providers/src/model/pq/mod.rs | 7 +- .../src/model/pq/pq_construction.rs | 71 +++++++------------ diskann-utils/src/views.rs | 17 +++++ 5 files changed, 50 insertions(+), 60 deletions(-) diff --git a/diskann-disk/src/storage/quant/pq/pq_generation.rs b/diskann-disk/src/storage/quant/pq/pq_generation.rs index a8a1557c7..ac59b5ad1 100644 --- a/diskann-disk/src/storage/quant/pq/pq_generation.rs +++ b/diskann-disk/src/storage/quant/pq/pq_generation.rs @@ -9,15 +9,12 @@ use diskann::{utils::VectorRepr, ANNError}; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; use diskann_providers::{ forward_threadpool, - model::{ - pq::{accum_row_inplace, generate_pq_pivots}, - GeneratePivotArguments, - }, + model::{pq::generate_pq_pivots, GeneratePivotArguments}, storage::PQStorage, utils::{AsThreadPool, BridgeErr, Timer}, }; use diskann_quantization::{product::TransposedTable, CompressInto}; -use diskann_utils::views::MatrixBase; +use diskann_utils::views::{accum_row_inplace, MatrixBase}; use diskann_vector::distance::Metric; use tracing::info; diff --git a/diskann-providers/src/model/mod.rs b/diskann-providers/src/model/mod.rs index e9fceb869..61addc05a 100644 --- a/diskann-providers/src/model/mod.rs +++ b/diskann-providers/src/model/mod.rs @@ -11,10 +11,10 @@ pub use configuration::IndexConfiguration; pub mod pq; pub use pq::{ FixedChunkPQTable, GeneratePivotArguments, MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, - NUM_PQ_CENTROIDS, accum_row_inplace, compute_pq_distance, - compute_pq_distance_for_pq_coordinates, direct_distance_impl, distance, - generate_pq_data_from_pivots_from_membuf, generate_pq_data_from_pivots_from_membuf_batch, - generate_pq_pivots, generate_pq_pivots_from_membuf, + NUM_PQ_CENTROIDS, compute_pq_distance, compute_pq_distance_for_pq_coordinates, + direct_distance_impl, distance, generate_pq_data_from_pivots_from_membuf, + generate_pq_data_from_pivots_from_membuf_batch, generate_pq_pivots, + generate_pq_pivots_from_membuf, }; pub mod statistics; diff --git a/diskann-providers/src/model/pq/mod.rs b/diskann-providers/src/model/pq/mod.rs index 6c1962a4a..ba6a49d25 100644 --- a/diskann-providers/src/model/pq/mod.rs +++ b/diskann-providers/src/model/pq/mod.rs @@ -10,10 +10,9 @@ pub use fixed_chunk_pq_table::{ mod pq_construction; pub use pq_construction::{ - MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS, accum_row_inplace, - generate_pq_data_from_pivots, generate_pq_data_from_pivots_from_membuf, - generate_pq_data_from_pivots_from_membuf_batch, generate_pq_pivots, - generate_pq_pivots_from_membuf, get_chunk_from_training_data, move_train_data_by_centroid, + MAX_PQ_TRAINING_SET_SIZE, NUM_KMEANS_REPS_PQ, NUM_PQ_CENTROIDS, generate_pq_data_from_pivots, + generate_pq_data_from_pivots_from_membuf, generate_pq_data_from_pivots_from_membuf_batch, + generate_pq_pivots, generate_pq_pivots_from_membuf, move_train_data_by_centroid, }; /// all metadata of individual sub-component files is written in first 4KB for unified files diff --git a/diskann-providers/src/model/pq/pq_construction.rs b/diskann-providers/src/model/pq/pq_construction.rs index 69b95c92a..bae68e650 100644 --- a/diskann-providers/src/model/pq/pq_construction.rs +++ b/diskann-providers/src/model/pq/pq_construction.rs @@ -23,7 +23,7 @@ use diskann_quantization::{ }; use diskann_utils::{ io::Metadata, - views::{MatrixView, MutMatrixView}, + views::{MatrixView, MutMatrixView, accum_row_inplace}, }; use rand::{Rng, distr::Distribution}; use rayon::prelude::*; @@ -254,35 +254,6 @@ pub fn generate_pq_pivots_from_membuf, Pool: AsThreadPool>( Ok(()) } -/// Gets all instances of a chunk from the training data for all records in the training data. Each vector in the -/// training dataset is divided into chunks and the PQ algorithm handles each vector chunk individually. This method -/// gets the same chunk from each vector in the training data and creates a new vector out of all of them. -/// -/// # Example -/// See tests for examples -#[inline] -pub fn get_chunk_from_training_data( - train_data: &[f32], - num_train: usize, - raw_vector_dim: usize, - chunk_size: usize, - chunk_offset: usize, -) -> Vec { - let mut result: Vec = vec![0.0; num_train * chunk_size]; - - result - // group empty result data into chunks of chunk_size - .chunks_mut(chunk_size) - .enumerate() - // for each chunk, copy the chunk from the training data into the result vector - .for_each(|(chunk_number, result_chunk)| { - let train_data_start = chunk_number * raw_vector_dim + chunk_offset; - let train_data_end = train_data_start + chunk_size; - result_chunk.copy_from_slice(&train_data[train_data_start..train_data_end]); - }); - result -} - /// Calculates the centroid if needed and moves the train_data to to the centroid /// # Arguments /// * `train_data` Dataset @@ -329,23 +300,6 @@ pub fn move_train_data_by_centroid( } } -/// Add the row `y` to every row in `x`. -/// -/// # Panics -/// -/// Panics if `y.len() != x.ncols()`. -pub fn accum_row_inplace(mut x: MutMatrixView, y: &[T]) -where - T: Copy + std::ops::AddAssign, -{ - assert_eq!(x.ncols(), y.len()); - x.row_iter_mut().for_each(|row| { - std::iter::zip(row.iter_mut(), y.iter()).for_each(|(a, b)| { - *a += *b; - }); - }); -} - /// streams the base file (data_file), and computes the closest centers in each /// chunk to generate the compressed data_file and stores it in /// pq_compressed_vectors_path. @@ -655,6 +609,29 @@ mod pq_test { utils::{ParallelIteratorInPool, create_thread_pool_for_test, read_bin_from}, }; + /// Test helper: Gets all instances of a chunk from the training data for all records + /// in the training data. Each vector in the training dataset is divided into chunks + /// and the PQ algorithm handles each vector chunk individually. This helper gets the + /// same chunk from each vector in the training data and returns it as a flat vector. + fn get_chunk_from_training_data( + train_data: &[f32], + num_train: usize, + raw_vector_dim: usize, + chunk_size: usize, + chunk_offset: usize, + ) -> Vec { + let mut result: Vec = vec![0.0; num_train * chunk_size]; + result + .chunks_mut(chunk_size) + .enumerate() + .for_each(|(chunk_number, result_chunk)| { + let train_data_start = chunk_number * raw_vector_dim + chunk_offset; + let train_data_end = train_data_start + chunk_size; + result_chunk.copy_from_slice(&train_data[train_data_start..train_data_end]); + }); + result + } + #[test] fn test_move_train_data_by_centroid() { let dim = 20; diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs index a9352918c..f4ff31883 100644 --- a/diskann-utils/src/views.rs +++ b/diskann-utils/src/views.rs @@ -673,6 +673,23 @@ impl<'a, T> From> for &'a [T] { } } +/// Add the row `y` to every row in `x`. +/// +/// # Panics +/// +/// Panics if `y.len() != x.ncols()`. +pub fn accum_row_inplace(mut x: MutMatrixView, y: &[T]) +where + T: Copy + std::ops::AddAssign, +{ + assert_eq!(x.ncols(), y.len()); + x.row_iter_mut().for_each(|row| { + std::iter::zip(row.iter_mut(), y.iter()).for_each(|(a, b)| { + *a += *b; + }); + }); +} + /// Return a reference to the item at entry `(row, col)` in the matrix. /// /// # Panics