Unify bftree input structs with QuantConfig enum

Merge BfTreeFullPrecisionBuild + BfTreeSphericalBuild into BfTreeBuild,
and BfTreeStreamingRun + BfTreeSphericalStreamingRun into a single
BfTreeStreamingRun. Both structs now carry a QuantConfig enum that
discriminates between full-precision (None) and spherical quantization.

- Add QuantConfig enum with tagged serde (kind: none | spherical)
- Consolidate quantizer training into quantizer_util::build_quantizer
- Update benchmark callers to check QuantConfig variant in try_match
- Update example JSON files with new unified tags and quantization field
- Add Clone to exhaustive::TransformKind

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: JordanMaples

diskann-benchmark/example/graph-index-bftree-spherical.json

@@ -4,7 +4,7 @@
     ],
     "jobs": [
         {
-            "type": "graph-index-bftree-spherical-quantization",
+            "type": "graph-index-bftree",
             "content": {
                 "build": {
                     "data_type": "float32",
@@ -40,10 +40,24 @@
                         }
                     ]
                 },
-                "seed": 42,
-                "transform_kind": "null",
-                "num_bits": 2,
-                "pre_scale": "reciprocal_mean_norm",
+                "quantization": {
+                    "kind": "spherical",
+                    "seed": 42,
+                    "transform_kind": "null",
+                    "num_bits": 2,
+                    "pre_scale": "reciprocal_mean_norm",
+                    "quant_store_config": {
+                        "cb_size_byte": 67108864,
+                        "leaf_page_size": 4096,
+                        "cb_max_record_size": null,
+                        "cb_min_record_size": null,
+                        "read_promotion_rate": null,
+                        "scan_promotion_rate": null,
+                        "cb_copy_on_access_ratio": null,
+                        "read_record_cache": null,
+                        "cache_only": null
+                    }
+                },
                 "vector_store_config": {
                     "cb_size_byte": 67108864,
                     "leaf_page_size": 4096,
@@ -65,17 +79,6 @@
                     "cb_copy_on_access_ratio": null,
                     "read_record_cache": null,
                     "cache_only": null
-                },
-                "quant_store_config": {
-                    "cb_size_byte": 67108864,
-                    "leaf_page_size": 4096,
-                    "cb_max_record_size": null,
-                    "cb_min_record_size": null,
-                    "read_promotion_rate": null,
-                    "scan_promotion_rate": null,
-                    "cb_copy_on_access_ratio": null,
-                    "read_record_cache": null,
-                    "cache_only": null
                 }
             }
         }

diskann-benchmark/example/graph-index-bftree-stream-spherical.json

@@ -0,0 +1,95 @@
+{
+  "search_directories": [
+    "test_data/disk_index_search"
+  ],
+  "jobs": [
+    {
+      "type": "graph-index-stream-bftree",
+      "content": {
+        "build": {
+          "data_type": "float32",
+          "data": "disk_index_siftsmall_learn_256pts_data.fbin",
+          "distance": "squared_l2",
+          "max_degree": 32,
+          "l_build": 50,
+          "start_point_strategy": "medoid",
+          "alpha": 1.2,
+          "backedge_ratio": 1.0,
+          "num_threads": 1
+        },
+        "search_phase": {
+          "search-type": "topk",
+          "queries": "disk_index_sample_query_10pts.fbin",
+          "groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin",
+          "reps": 1,
+          "num_threads": [
+            1
+          ],
+          "runs": [
+            {
+              "search_n": 10,
+              "search_l": [
+                20,
+                40
+              ],
+              "recall_k": 10
+            }
+          ]
+        },
+        "runbook_params": {
+          "runbook_path": "example_runbook.yaml",
+          "dataset_name": "sift-small-256",
+          "gt_directory": "example_runbook_gt",
+          "ip_delete_method": {
+            "method": "visited_and_top_k",
+            "params": {
+              "k_value": 10,
+              "l_value": 40
+            }
+          },
+          "ip_delete_num_to_replace": 3
+        },
+        "quantization": {
+          "kind": "spherical",
+          "seed": 42,
+          "transform_kind": "null",
+          "num_bits": 2,
+          "pre_scale": "reciprocal_mean_norm",
+          "quant_store_config": {
+            "cb_size_byte": 67108864,
+            "leaf_page_size": 4096,
+            "cb_max_record_size": null,
+            "cb_min_record_size": null,
+            "read_promotion_rate": null,
+            "scan_promotion_rate": null,
+            "cb_copy_on_access_ratio": null,
+            "read_record_cache": null,
+            "cache_only": null
+          }
+        },
+        "vector_store_config": {
+          "cb_size_byte": 67108864,
+          "leaf_page_size": 4096,
+          "cb_max_record_size": null,
+          "cb_min_record_size": null,
+          "read_promotion_rate": null,
+          "scan_promotion_rate": null,
+          "cb_copy_on_access_ratio": null,
+          "read_record_cache": null,
+          "cache_only": null
+        },
+        "neighbor_store_config": {
+          "cb_size_byte": 67108864,
+          "leaf_page_size": 4096,
+          "cb_max_record_size": null,
+          "cb_min_record_size": null,
+          "read_promotion_rate": null,
+          "scan_promotion_rate": null,
+          "cb_copy_on_access_ratio": null,
+          "read_record_cache": null,
+          "cache_only": null
+        }
+      }
+    }
+  ]
+}

diskann-benchmark/example/graph-index-bftree-stream.json

@@ -4,7 +4,7 @@
   ],
   "jobs": [
     {
-      "type": "graph-index-stream-bftree-full-precision",
+      "type": "graph-index-stream-bftree",
       "content": {
         "build": {
           "data_type": "float32",
@@ -49,6 +49,9 @@
           },
           "ip_delete_num_to_replace": 3
         },
+        "quantization": {
+          "kind": "none"
+        },
         "vector_store_config": {
           "cb_size_byte": 67108864,
           "leaf_page_size": 4096,

diskann-benchmark/example/graph-index-bftree.json

@@ -4,7 +4,7 @@
   ],
   "jobs": [
     {
-      "type": "graph-index-build-bftree-full-precision",
+      "type": "graph-index-bftree",
       "content": {
         "build": {
           "data_type": "float32",
@@ -37,6 +37,9 @@
             }
           ]
         },
+        "quantization": {
+          "kind": "none"
+        },
         "vector_store_config": {
           "cb_size_byte": 67108864,
           "leaf_page_size": 4096,

diskann-benchmark/src/index/bftree/full_precision.rs

@@ -25,7 +25,10 @@ use crate::{
         result::BuildResult,
         search::plugins::{Plugin, Plugins},
     },
-    inputs::{bftree::BfTreeFullPrecisionBuild, graph_index::SearchPhase},
+    inputs::{
+        bftree::{BfTreeBuild, QuantConfig},
+        graph_index::SearchPhase,
+    },
     utils::{self},
 };
 
@@ -68,25 +71,32 @@ impl<T> Benchmark for BfTreeFullPrecision<T>
 where
     T: VectorRepr + AsDataType + SampleableForStart + WithApproximateNorm + 'static,
 {
-    type Input = BfTreeFullPrecisionBuild;
+    type Input = BfTreeBuild;
     type Output = BuildResult;
 
     fn try_match(&self, input: &Self::Input) -> Result<MatchScore, FailureScore> {
-        let score = utils::match_data_type::<T>(input.data_type());
-        if self.plugins.is_match(input.search_phase()) {
-            score
-        } else {
-            match score {
-                Ok(_) => Err(FailureScore(0)),
-                Err(s) => Err(s),
-            }
+        let mut failure_score: Option<u32> = None;
+
+        if !matches!(input.quantization(), QuantConfig::None) {
+            *failure_score.get_or_insert(0) += 1;
+        }
+        if let Err(s) = utils::match_data_type::<T>(input.data_type()) {
+            *failure_score.get_or_insert(0) += s.0;
+        }
+        if !self.plugins.is_match(input.search_phase()) {
+            *failure_score.get_or_insert(0) += 1;
+        }
+
+        match failure_score {
+            None => Ok(MatchScore(0)),
+            Some(score) => Err(FailureScore(score)),
         }
     }
 
     fn description(
         &self,
         f: &mut std::fmt::Formatter<'_>,
-        input: Option<&Self::Input>,
+        input: Option<&BfTreeBuild>,
     ) -> std::fmt::Result {
         match input {
             Some(arg) => {
@@ -113,7 +123,7 @@ where
 
     fn run(
         &self,
-        input: &Self::Input,
+        input: &BfTreeBuild,
         checkpoint: Checkpoint<'_>,
         mut output: &mut dyn Output,
     ) -> anyhow::Result<Self::Output> {

diskann-benchmark/src/index/bftree/full_precision_streaming.rs

@@ -26,7 +26,10 @@ use crate::{
         stats::StreamStats,
         StreamRunner,
     },
-    inputs::{bftree::BfTreeStreamingRun, graph_index::SearchPhase},
+    inputs::{
+        bftree::{BfTreeStreamingRun, QuantConfig},
+        graph_index::SearchPhase,
+    },
     utils,
 };
 
@@ -57,10 +60,12 @@ where
     fn try_match(&self, input: &Self::Input) -> Result<MatchScore, FailureScore> {
         let mut failure_score: Option<u32> = None;
 
+        if !matches!(input.quantization(), QuantConfig::None) {
+            *failure_score.get_or_insert(0) += 1;
+        }
         if let Err(s) = utils::match_data_type::<T>(input.data_type()) {
-            failure_score = Some(s.0);
+            *failure_score.get_or_insert(0) += s.0;
         }
-
         if !matches!(input.search_phase(), SearchPhase::Topk(_)) {
             *failure_score.get_or_insert(0) += 1;
         }

diskann-benchmark/src/index/bftree/mod.rs

@@ -10,16 +10,15 @@
 //! once, search) and streaming (insert/delete/search interleaved) modes.
 //!
 //! Registered tags:
-//! - `graph-index-bftree-full-precision-f32` — static FP build + search
-//! - `graph-index-bftree-stream-full-precision-f32` — streaming FP
-//! - `graph-index-bftree-spherical-quantization` — static spherical (1/2/4-bit)
-//! - `graph-index-stream-bftree-spherical-quantization` — streaming spherical (1/2/4-bit)
+//! - `graph-index-bftree` — static build + search (full-precision or spherical)
+//! - `graph-index-stream-bftree` — streaming (full-precision or spherical)
 
 use super::search::plugins::Topk;
 use diskann_benchmark_runner::Registry;
 
 mod full_precision;
 mod full_precision_streaming;
+mod quantizer_util;
 mod spherical;
 mod spherical_streaming;
 
@@ -30,13 +29,13 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()>
     )?;
 
     registry.register(
-        "graph-index-bftree-stream-full-precision-f32",
-        full_precision_streaming::StreamingFullPrecision::<f32>::new(),
+        "graph-index-bftree-spherical-quantization",
+        spherical::BfTreeSpherical::new().search(Topk),
     )?;
 
     registry.register(
-        "graph-index-bftree-spherical-quantization",
-        spherical::BfTreeSpherical::new().search(Topk),
+        "graph-index-stream-bftree-full-precision-f32",
+        full_precision_streaming::StreamingFullPrecision::<f32>::new(),
     )?;
 
     registry.register(

diskann-benchmark/src/index/bftree/quantizer_util.rs

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use diskann_quantization::{
+    alloc::{AllocatorError, GlobalAllocator, Poly},
+    spherical::{
+        iface::{self as spherical_iface, Quantizer},
+        SphericalQuantizer,
+    },
+};
+use diskann_utils::views::MatrixView;
+use rand::SeedableRng;
+
+use crate::{inputs::bftree::QuantConfig, utils::SimilarityMeasure};
+
+fn new_quantizer<const NBITS: usize>(
+    quantizer: SphericalQuantizer,
+) -> Result<Poly<dyn Quantizer>, AllocatorError>
+where
+    spherical_iface::Impl<NBITS>: spherical_iface::Constructible + Quantizer,
+{
+    let imp = spherical_iface::Impl::<NBITS>::new(quantizer)?;
+    diskann_quantization::poly!(Quantizer, imp, GlobalAllocator)
+}
+
+pub(super) fn build_quantizer(
+    quantization: &QuantConfig,
+    data: MatrixView<'_, f32>,
+    distance: SimilarityMeasure,
+) -> anyhow::Result<Option<Poly<dyn Quantizer>>> {
+    match quantization {
+        QuantConfig::None => Ok(None),
+        QuantConfig::Spherical {
+            seed,
+            transform_kind,
+            num_bits,
+            pre_scale,
+            ..
+        } => {
+            let m: diskann_vector::distance::Metric = distance.into();
+            let pre_scale = match pre_scale {
+                Some(v) => (*v).try_into()?,
+                None => diskann_quantization::spherical::PreScale::None,
+            };
+
+            let quantizer = SphericalQuantizer::train(
+                data,
+                transform_kind.into(),
+                m.try_into()?,
+                pre_scale,
+                &mut rand::rngs::StdRng::seed_from_u64(*seed),
+                GlobalAllocator,
+            )?;
+
+            let poly = match num_bits.get() {
+                1 => new_quantizer::<1>(quantizer)?,
+                2 => new_quantizer::<2>(quantizer)?,
+                4 => new_quantizer::<4>(quantizer)?,
+                n => anyhow::bail!("{n} bits not supported for spherical quantization"),
+            };
+
+            Ok(Some(poly))
+        }
+    }
+}