Skip to content

Commit 1875da1

Browse files
JordanMaplesCopilot
JordanMaples
and
Copilot
committed
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>
1 parent a760667 commit 1875da1

11 files changed

Lines changed: 287 additions & 445 deletions

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

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
],
55
"jobs": [
66
{
7-
"type": "graph-index-bftree-spherical-quantization",
7+
"type": "graph-index-bftree",
88
"content": {
99
"build": {
1010
"data_type": "float32",
@@ -40,10 +40,24 @@
4040
}
4141
]
4242
},
43-
"seed": 42,
44-
"transform_kind": "null",
45-
"num_bits": 2,
46-
"pre_scale": "reciprocal_mean_norm",
43+
"quantization": {
44+
"kind": "spherical",
45+
"seed": 42,
46+
"transform_kind": "null",
47+
"num_bits": 2,
48+
"pre_scale": "reciprocal_mean_norm",
49+
"quant_store_config": {
50+
"cb_size_byte": 67108864,
51+
"leaf_page_size": 4096,
52+
"cb_max_record_size": null,
53+
"cb_min_record_size": null,
54+
"read_promotion_rate": null,
55+
"scan_promotion_rate": null,
56+
"cb_copy_on_access_ratio": null,
57+
"read_record_cache": null,
58+
"cache_only": null
59+
}
60+
},
4761
"vector_store_config": {
4862
"cb_size_byte": 67108864,
4963
"leaf_page_size": 4096,
@@ -65,17 +79,6 @@
6579
"cb_copy_on_access_ratio": null,
6680
"read_record_cache": null,
6781
"cache_only": null
68-
},
69-
"quant_store_config": {
70-
"cb_size_byte": 67108864,
71-
"leaf_page_size": 4096,
72-
"cb_max_record_size": null,
73-
"cb_min_record_size": null,
74-
"read_promotion_rate": null,
75-
"scan_promotion_rate": null,
76-
"cb_copy_on_access_ratio": null,
77-
"read_record_cache": null,
78-
"cache_only": null
7982
}
8083
}
8184
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
],
55
"jobs": [
66
{
7-
"type": "graph-index-stream-bftree-full-precision",
7+
"type": "graph-index-stream-bftree",
88
"content": {
99
"build": {
1010
"data_type": "float32",
@@ -49,6 +49,9 @@
4949
},
5050
"ip_delete_num_to_replace": 3
5151
},
52+
"quantization": {
53+
"kind": "none"
54+
},
5255
"vector_store_config": {
5356
"cb_size_byte": 67108864,
5457
"leaf_page_size": 4096,

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
],
55
"jobs": [
66
{
7-
"type": "graph-index-build-bftree-full-precision",
7+
"type": "graph-index-bftree",
88
"content": {
99
"build": {
1010
"data_type": "float32",
@@ -37,6 +37,9 @@
3737
}
3838
]
3939
},
40+
"quantization": {
41+
"kind": "none"
42+
},
4043
"vector_store_config": {
4144
"cb_size_byte": 67108864,
4245
"leaf_page_size": 4096,

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

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@ use crate::{
2525
result::BuildResult,
2626
search::plugins::{Plugin, Plugins},
2727
},
28-
inputs::{bftree::BfTreeFullPrecisionBuild, graph_index::SearchPhase},
28+
inputs::{
29+
bftree::{BfTreeBuild, QuantConfig},
30+
graph_index::SearchPhase,
31+
},
2932
utils::{self},
3033
};
3134

@@ -68,25 +71,32 @@ impl<T> Benchmark for BfTreeFullPrecision<T>
6871
where
6972
T: VectorRepr + AsDataType + SampleableForStart + WithApproximateNorm + 'static,
7073
{
71-
type Input = BfTreeFullPrecisionBuild;
74+
type Input = BfTreeBuild;
7275
type Output = BuildResult;
7376

7477
fn try_match(&self, input: &Self::Input) -> Result<MatchScore, FailureScore> {
75-
let score = utils::match_data_type::<T>(input.data_type());
76-
if self.plugins.is_match(input.search_phase()) {
77-
score
78-
} else {
79-
match score {
80-
Ok(_) => Err(FailureScore(0)),
81-
Err(s) => Err(s),
82-
}
78+
let mut failure_score: Option<u32> = None;
79+
80+
if !matches!(input.quantization(), QuantConfig::None) {
81+
*failure_score.get_or_insert(0) += 1;
82+
}
83+
if let Err(s) = utils::match_data_type::<T>(input.data_type()) {
84+
*failure_score.get_or_insert(0) += s.0;
85+
}
86+
if !self.plugins.is_match(input.search_phase()) {
87+
*failure_score.get_or_insert(0) += 1;
88+
}
89+
90+
match failure_score {
91+
None => Ok(MatchScore(0)),
92+
Some(score) => Err(FailureScore(score)),
8393
}
8494
}
8595

8696
fn description(
8797
&self,
8898
f: &mut std::fmt::Formatter<'_>,
89-
input: Option<&Self::Input>,
99+
input: Option<&BfTreeBuild>,
90100
) -> std::fmt::Result {
91101
match input {
92102
Some(arg) => {
@@ -113,7 +123,7 @@ where
113123

114124
fn run(
115125
&self,
116-
input: &Self::Input,
126+
input: &BfTreeBuild,
117127
checkpoint: Checkpoint<'_>,
118128
mut output: &mut dyn Output,
119129
) -> anyhow::Result<Self::Output> {

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,10 @@ use crate::{
2626
stats::StreamStats,
2727
StreamRunner,
2828
},
29-
inputs::{bftree::BfTreeStreamingRun, graph_index::SearchPhase},
29+
inputs::{
30+
bftree::{BfTreeStreamingRun, QuantConfig},
31+
graph_index::SearchPhase,
32+
},
3033
utils,
3134
};
3235

@@ -57,10 +60,12 @@ where
5760
fn try_match(&self, input: &Self::Input) -> Result<MatchScore, FailureScore> {
5861
let mut failure_score: Option<u32> = None;
5962

63+
if !matches!(input.quantization(), QuantConfig::None) {
64+
*failure_score.get_or_insert(0) += 1;
65+
}
6066
if let Err(s) = utils::match_data_type::<T>(input.data_type()) {
61-
failure_score = Some(s.0);
67+
*failure_score.get_or_insert(0) += s.0;
6268
}
63-
6469
if !matches!(input.search_phase(), SearchPhase::Topk(_)) {
6570
*failure_score.get_or_insert(0) += 1;
6671
}

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

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,15 @@
1010
//! once, search) and streaming (insert/delete/search interleaved) modes.
1111
//!
1212
//! Registered tags:
13-
//! - `graph-index-bftree-full-precision-f32` — static FP build + search
14-
//! - `graph-index-bftree-stream-full-precision-f32` — streaming FP
15-
//! - `graph-index-bftree-spherical-quantization` — static spherical (1/2/4-bit)
16-
//! - `graph-index-stream-bftree-spherical-quantization` — streaming spherical (1/2/4-bit)
13+
//! - `graph-index-bftree` — static build + search (full-precision or spherical)
14+
//! - `graph-index-stream-bftree` — streaming (full-precision or spherical)
1715
1816
use super::search::plugins::Topk;
1917
use diskann_benchmark_runner::Registry;
2018

2119
mod full_precision;
2220
mod full_precision_streaming;
21+
mod quantizer_util;
2322
mod spherical;
2423
mod spherical_streaming;
2524

@@ -30,13 +29,13 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()>
3029
)?;
3130

3231
registry.register(
33-
"graph-index-bftree-stream-full-precision-f32",
34-
full_precision_streaming::StreamingFullPrecision::<f32>::new(),
32+
"graph-index-bftree-spherical-quantization",
33+
spherical::BfTreeSpherical::new().search(Topk),
3534
)?;
3635

3736
registry.register(
38-
"graph-index-bftree-spherical-quantization",
39-
spherical::BfTreeSpherical::new().search(Topk),
37+
"graph-index-stream-bftree-full-precision-f32",
38+
full_precision_streaming::StreamingFullPrecision::<f32>::new(),
4039
)?;
4140

4241
registry.register(

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
/*
2+
* Copyright (c) Microsoft Corporation.
3+
* Licensed under the MIT license.
4+
*/
5+
6+
use diskann_quantization::{
7+
alloc::{AllocatorError, GlobalAllocator, Poly},
8+
spherical::{
9+
iface::{self as spherical_iface, Quantizer},
10+
SphericalQuantizer,
11+
},
12+
};
13+
use diskann_utils::views::MatrixView;
14+
use rand::SeedableRng;
15+
16+
use crate::{inputs::bftree::QuantConfig, utils::SimilarityMeasure};
17+
18+
fn new_quantizer<const NBITS: usize>(
19+
quantizer: SphericalQuantizer,
20+
) -> Result<Poly<dyn Quantizer>, AllocatorError>
21+
where
22+
spherical_iface::Impl<NBITS>: spherical_iface::Constructible + Quantizer,
23+
{
24+
let imp = spherical_iface::Impl::<NBITS>::new(quantizer)?;
25+
diskann_quantization::poly!(Quantizer, imp, GlobalAllocator)
26+
}
27+
28+
pub(super) fn build_quantizer(
29+
quantization: &QuantConfig,
30+
data: MatrixView<'_, f32>,
31+
distance: SimilarityMeasure,
32+
) -> anyhow::Result<Option<Poly<dyn Quantizer>>> {
33+
match quantization {
34+
QuantConfig::None => Ok(None),
35+
QuantConfig::Spherical {
36+
seed,
37+
transform_kind,
38+
num_bits,
39+
pre_scale,
40+
..
41+
} => {
42+
let m: diskann_vector::distance::Metric = distance.into();
43+
let pre_scale = match pre_scale {
44+
Some(v) => (*v).try_into()?,
45+
None => diskann_quantization::spherical::PreScale::None,
46+
};
47+
48+
let quantizer = SphericalQuantizer::train(
49+
data,
50+
transform_kind.into(),
51+
m.try_into()?,
52+
pre_scale,
53+
&mut rand::rngs::StdRng::seed_from_u64(*seed),
54+
GlobalAllocator,
55+
)?;
56+
57+
let poly = match num_bits.get() {
58+
1 => new_quantizer::<1>(quantizer)?,
59+
2 => new_quantizer::<2>(quantizer)?,
60+
4 => new_quantizer::<4>(quantizer)?,
61+
n => anyhow::bail!("{n} bits not supported for spherical quantization"),
62+
};
63+
64+
Ok(Some(poly))
65+
}
66+
}
67+
}

0 commit comments

Comments
 (0)