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] 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)]