[diskann-inmem] Prepare code for quantization and beyond

.github/workflows/ci.yml

@@ -231,7 +231,11 @@ jobs:
         run: rustup show
       - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
       - name: "doc --workspace --no-deps"
-        run: cargo doc --locked --workspace --no-deps --features linalg,flatbuffers,experimental_diversity_search
+        run: cargo doc --locked --workspace --no-deps  --features linalg,flatbuffers,experimental_diversity_search,integration-test
+        env:
+          RUSTDOCFLAGS: -D rustdoc::all
+      - name: "diskann-inmem private docs"
+        run: cargo doc --locked --package diskann-inmem --no-deps --all-features --document-private-items
         env:
           RUSTDOCFLAGS: -D rustdoc::all
 

diskann-benchmark/src/index/inmem2.rs

@@ -29,7 +29,8 @@ use diskann_benchmark_runner::{
     Benchmark, Checker, Checkpoint, Input, Registry,
 };
 use diskann_inmem::{
-    layers::{Full, FullPrecision},
+    num::{Capacity, MaxDegree},
+    repr::{Full, FullPrecision},
     Provider, Strategy,
 };
 use diskann_utils::views::{Matrix, MatrixView};
@@ -47,6 +48,8 @@ use crate::{
 
 pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> {
     registry.register("inmem2-f32", Build::<f32>::new())?;
+    registry.register("inmem2-u8", Build::<u8>::new())?;
+
     // registry.register("inmem2-f16", Build::<f16>::new())?;
     registry.register("inmem2-f32-stream", StreamingBenchmark::<f32>::new())?;
     Ok(())
@@ -420,7 +423,7 @@ impl<T> Build<T> {
 
 impl<T> Benchmark for Build<T>
 where
-    T: diskann_inmem::layers::FullPrecision + diskann::graph::SampleableForStart + AsDataType,
+    T: diskann_inmem::repr::FullPrecision + diskann::graph::SampleableForStart + AsDataType,
 {
     type Input = StaticBuild;
     type Output = ();
@@ -472,11 +475,14 @@ where
 
         // Compute the medoid of the dataset as the single start point.
         let start = StartPointStrategy::Medoid.compute(data.as_view())?;
-        let layer = Full::<T>::new(dim, input.data.distance);
-        let config =
-            diskann_inmem::provider::Config::new(num_points, input.build.config.max_degree().get());
-        let provider = Provider::<_, u32>::new(layer, config, start.row_iter())?;
+        let config = Full::config(
+            Capacity::new(num_points),
+            MaxDegree::new(input.build.config.max_degree().get()),
+            input.data.distance,
+            start,
+        )?;
 
+        let provider = Provider::<_, u32>::new(config)?;
         let index = Arc::new(DiskANNIndex::new(
             input.build.config.clone(),
             provider,
@@ -837,17 +843,18 @@ where
         let queries: Arc<Matrix<T>> = Arc::new(datafiles::load_dataset(datafiles::BinFile(
             &input.search.queries,
         ))?);
-        let dim = dataset.ncols();
 
         // Compute the medoid of the dataset as the single start point.
         let start = StartPointStrategy::Medoid.compute(dataset.as_view())?;
         let index_config = input.build.config.clone();
-        let layer = Full::<T>::new(dim, input.data.distance);
-
-        let config =
-            diskann_inmem::provider::Config::new(max_points, index_config.max_degree().get());
-        let provider = Provider::<_, u32>::new(layer, config, start.row_iter())?;
+        let config = Full::config(
+            Capacity::new(max_points),
+            MaxDegree::new(index_config.max_degree().get()),
+            input.data.distance,
+            start,
+        )?;
 
+        let provider = Provider::<_, u32>::new(config)?;
         let index = Arc::new(DiskANNIndex::new(index_config, provider, None));
 
         let num_threads = input.build.num_threads;

diskann-inmem/Cargo.toml

@@ -16,8 +16,9 @@ diskann-utils = { workspace = true, default-features = false }
 diskann-vector = { workspace = true }
 diskann-wide = { workspace = true }
 parking_lot = "0.12.5"
-thiserror = { workspace = true }
 half = { workspace = true }
+hashbrown = { workspace = true }
+thiserror = { workspace = true }
 
 # Integration Test Dependencies
 diskann-benchmark-runner = { workspace = true, optional = true, features = ["ux-tools"] }

diskann-inmem/integration/index/object.rs

@@ -14,7 +14,7 @@ use diskann_benchmark_runner::utils::fmt::KeyValue;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-use diskann_inmem::{Context, Provider, Strategy, integration, layers};
+use diskann_inmem::{Context, Provider, Strategy, integration, repr};
 
 use crate::support::{
     check::{CheckMatch, Match, check_all_fields},
@@ -186,9 +186,9 @@ impl CheckMatch for Counters {
 // Impls //
 ///////////
 
-impl<T> Index for DiskANNIndex<Provider<layers::Full<T>, u64>>
+impl<T> Index for DiskANNIndex<Provider<repr::Full<T>, u64>>
 where
-    T: layers::FullPrecision + FromSlice + AsDataType,
+    T: repr::FullPrecision + FromSlice + AsDataType,
 {
     fn search<'a>(
         &'a self,

diskann-inmem/integration/index/runner.rs

@@ -15,10 +15,13 @@ use diskann_benchmark_runner::{
 };
 use diskann_utils::views::Matrix;
 use diskann_vector::distance::Metric;
-use half::f16;
 use serde::{Deserialize, Serialize};
 
-use diskann_inmem::{Provider, layers};
+use diskann_inmem::{
+    Provider,
+    num::{Capacity, MaxDegree},
+    repr::Full,
+};
 
 use crate::{
     index::{Counters, Index},
@@ -106,7 +109,7 @@ mod dto {
     }
 
     #[derive(Debug, Serialize, Deserialize)]
-    pub(super) enum Layer {
+    pub(super) enum Representation {
         FullPrecision { data_type: DataType },
     }
 
@@ -134,7 +137,7 @@ mod dto {
     #[derive(Debug, Serialize, Deserialize)]
     pub(super) struct Test {
         pub(super) data: Data,
-        pub(super) layer: Layer,
+        pub(super) representation: Representation,
         pub(super) build: Build,
         pub(super) search: Search,
     }
@@ -227,20 +230,20 @@ struct Bundle {
 }
 
 #[derive(Debug)]
-enum Layer {
+enum Representation {
     FullPrecision { data_type: DataType },
 }
 
-impl Layer {
-    fn from_raw(raw: dto::Layer) -> Self {
+impl Representation {
+    fn from_raw(raw: dto::Representation) -> Self {
         match raw {
-            dto::Layer::FullPrecision { data_type } => Self::FullPrecision { data_type },
+            dto::Representation::FullPrecision { data_type } => Self::FullPrecision { data_type },
         }
     }
 
-    fn as_raw(&self) -> dto::Layer {
+    fn as_raw(&self) -> dto::Representation {
         match self {
-            Self::FullPrecision { data_type } => dto::Layer::FullPrecision {
+            Self::FullPrecision { data_type } => dto::Representation::FullPrecision {
                 data_type: *data_type,
             },
         }
@@ -323,21 +326,21 @@ impl Search {
 #[derive(Debug)]
 struct Test {
     data: Data,
-    layer: Layer,
+    representation: Representation,
     build: Build,
     search: Search,
 }
 
 impl Test {
     fn from_raw(raw: dto::Test, checker: Option<&mut Checker>) -> anyhow::Result<Self> {
         let data = Data::from_raw(raw.data, checker)?;
-        let layer = Layer::from_raw(raw.layer);
+        let representation = Representation::from_raw(raw.representation);
         let build = Build::from_raw(raw.build, data.metric)?;
         let search = Search::from_raw(raw.search)?;
 
         Ok(Self {
             data,
-            layer,
+            representation,
             build,
             search,
         })
@@ -346,7 +349,7 @@ impl Test {
     fn as_raw(&self) -> anyhow::Result<dto::Test> {
         Ok(dto::Test {
             data: self.data.as_raw()?,
-            layer: self.layer.as_raw(),
+            representation: self.representation.as_raw(),
             build: self.build.as_raw(),
             search: self.search.as_raw(),
         })
@@ -357,8 +360,8 @@ impl Test {
         capacity: usize,
         start_points: DatasetView<'_>,
     ) -> anyhow::Result<Arc<dyn Index>> {
-        match self.layer {
-            Layer::FullPrecision { data_type } => {
+        match self.representation {
+            Representation::FullPrecision { data_type } => {
                 if start_points.data_type() != data_type {
                     anyhow::bail!(
                         "mismatched data types for start point - expected {}, got {}",
@@ -367,30 +370,45 @@ impl Test {
                     );
                 }
 
-                let dim = start_points.ncols();
                 let metric = self.data.metric;
-                let config = diskann_inmem::provider::Config::new(
-                    capacity,
-                    self.build.config.max_degree().get(),
-                );
-
+                let max_degree = self.build.config.max_degree().get();
                 let index_config = self.build.config.clone();
 
                 let index = match start_points {
                     DatasetView::F32(v) => finish(
-                        Provider::new(layers::Full::<f32>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(Full::config(
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
+                            metric,
+                            v.to_owned(),
+                        )?)?,
                         index_config,
                     ),
                     DatasetView::F16(v) => finish(
-                        Provider::new(layers::Full::<f16>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(Full::config(
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
+                            metric,
+                            v.to_owned(),
+                        )?)?,
                         index_config,
                     ),
                     DatasetView::U8(v) => finish(
-                        Provider::new(layers::Full::<u8>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(Full::config(
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
+                            metric,
+                            v.to_owned(),
+                        )?)?,
                         index_config,
                     ),
                     DatasetView::I8(v) => finish(
-                        Provider::new(layers::Full::<i8>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(Full::config(
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
+                            metric,
+                            v.to_owned(),
+                        )?)?,
                         index_config,
                     ),
                 };
@@ -439,7 +457,7 @@ impl diskann_benchmark_runner::Input for Test {
                 data_type: DataType::F32,
                 preprocess: vec![],
             },
-            layer: dto::Layer::FullPrecision {
+            representation: dto::Representation::FullPrecision {
                 data_type: DataType::F32,
             },
             build: dto::Build {
@@ -483,8 +501,8 @@ impl diskann_benchmark_runner::Benchmark for FullPrecision {
     type Output = BuildAndSearch;
 
     fn try_match(&self, input: &Test, context: &MatchContext) -> Score {
-        // Future-proof against additional enums in `input.layer`.
-        let Layer::FullPrecision { .. } = input.layer;
+        // Future-proof against additional enums in `input.representation`.
+        let Representation::FullPrecision { .. } = input.representation;
         context.success(0)
     }
 
@@ -498,8 +516,8 @@ impl diskann_benchmark_runner::Benchmark for FullPrecision {
         _checkpoint: Checkpoint<'_>,
         mut output: &mut dyn Output,
     ) -> anyhow::Result<Self::Output> {
-        // Future-proof against additional enums in `input.layer`.
-        let Layer::FullPrecision { data_type } = input.layer;
+        // Future-proof against additional enums in `input.representation`.
+        let Representation::FullPrecision { data_type } = input.representation;
 
         // Load the data and perform any necessary data conversions.
         let Bundle {

diskann-inmem/integration/jsons/integration-baseline.json

@@ -16,7 +16,7 @@
           "preprocess": [],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "f32"
           }
@@ -48,7 +48,7 @@
         "append_neighbors": 96949,
         "distance": 2867876,
         "get_neighbors": 352139,
-        "get_vector": 3067092,
+        "get_vector": 2781513,
         "query_distance": 2240744,
         "set_neighbors": 23599,
         "set_vector": 10000
@@ -137,7 +137,7 @@
           "preprocess": [],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "f16"
           }
@@ -169,7 +169,7 @@
         "append_neighbors": 96949,
         "distance": 2867876,
         "get_neighbors": 352139,
-        "get_vector": 3067092,
+        "get_vector": 2781513,
         "query_distance": 2240744,
         "set_neighbors": 23599,
         "set_vector": 10000
@@ -258,7 +258,7 @@
           "preprocess": [],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "u8"
           }
@@ -290,7 +290,7 @@
         "append_neighbors": 96949,
         "distance": 2867876,
         "get_neighbors": 352139,
-        "get_vector": 3067092,
+        "get_vector": 2781513,
         "query_distance": 2240744,
         "set_neighbors": 23599,
         "set_vector": 10000
@@ -382,7 +382,7 @@
           ],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "i8"
           }
@@ -414,7 +414,7 @@
         "append_neighbors": 97055,
         "distance": 2867292,
         "get_neighbors": 352087,
-        "get_vector": 3064106,
+        "get_vector": 2778779,
         "query_distance": 2238420,
         "set_neighbors": 23587,
         "set_vector": 10000
@@ -486,4 +486,4 @@
       ]
     }
   }
-]
+]