Supporting spherical quantization builds for the disk index - #1331
Supporting spherical quantization builds for the disk index#1331juchen-ms (partychen) wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Adds spherical quantization as an additional quantization mode for disk index builds, wiring 1-bit spherical quantization through the disk build pipeline (training, in-memory graph construction, and RAM estimation) and extending builder tests to cover additional metrics.
Changes:
- Extend disk build quantization configuration to support
SPHERICAL_<nbits>and add serialization/parse tests. - Train and use a 1-bit spherical quantizer during disk index in-memory build (one-shot and merged/sharded paths).
- Account for spherical vector storage in build RAM estimation and add builder test coverage for L2/IP/Cosine.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| diskann-quantization/src/spherical/quantizer.rs | Adds a try_clone() convenience API for independently allocated quantizer copies. |
| diskann-disk/src/build/configuration/quantization_types.rs | Adds QuantizationType::Spherical plus parsing/formatting and tests for the new variant. |
| diskann-disk/src/build/builder/tests.rs | Extends integration tests to build/search spherical 1-bit indexes, including metric variants. |
| diskann-disk/src/build/builder/quantizer.rs | Trains a 1-bit spherical quantizer for disk builds and stores it in the build quantizer enum. |
| diskann-disk/src/build/builder/inmem_builder.rs | Plumbs spherical quantization into the async in-memory index builder via a spherical insert/prune strategy. |
| diskann-disk/src/build/builder/core.rs | Updates build RAM estimation to account for spherical quantized vector storage and adds validation coverage. |
Suppressed comments (3)
diskann-disk/src/build/configuration/quantization_types.rs:186
QuantizationType::SQ { standard_deviation: None }currently formats asSQ_<nbits>_None, butFromStronly acceptsSQ_<nbits>for the default stddev. Because serde Serialize usesto_string()and Deserialize usesfrom_str, SQ-with-default cannot roundtrip (e.g. bincode serialize then deserialize fails). Consider emittingSQ_<nbits>whenstandard_deviationisNoneto keep Display/parse/serde consistent.
This issue also appears in the following locations of the same file:
- line 253
- line 312
QuantizationType::Spherical(nbits) => write!(f, "SPHERICAL_{}", nbits),
QuantizationType::SQ {
nbits,
standard_deviation,
} => {
diskann-disk/src/build/configuration/quantization_types.rs:257
- The
fmt_quantization_typetest currently assertsSQ_8_Nonefor the default-SQ formatting, which encodes the same Display/parse mismatch described above. If Display is changed to emitSQ_<nbits>whenstandard_deviationisNone, update this expectation accordingly so the test enforces a roundtrippable representation.
#[case(QuantizationType::Spherical(1), "SPHERICAL_1")]
#[case(
QuantizationType::SQ { nbits: 8, standard_deviation: None },
"SQ_8_None"
)]
diskann-disk/src/build/configuration/quantization_types.rs:316
test_roundtrip_serializationdocuments (and works around) the fact that SQ-with-default-stddev doesn't roundtrip. If Display/parse are made consistent forstandard_deviation: None, it would be valuable to include that case in this roundtrip test to prevent regressions.
nbits: 8,
standard_deviation: Some(Positive::new(1.5).unwrap()),
},
QuantizationType::Spherical(1),
];
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1331 +/- ##
==========================================
+ Coverage 91.55% 92.57% +1.01%
==========================================
Files 522 522
Lines 99541 99617 +76
==========================================
+ Hits 91139 92222 +1083
+ Misses 8402 7395 -1007
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Mark Hildebrand (hildebrandmw)
left a comment
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Added some comments on the use of spherical quantization and some maintenance suggestions. Please get a review from the maintainers of diskann-disk regarding how this feature fits in architecturally.
| /// Return an independently allocated copy of this quantizer. | ||
| pub fn try_clone(&self) -> Result<Self, AllocatorError> { | ||
| <Self as TryClone>::try_clone(self) | ||
| } |
There was a problem hiding this comment.
Please leave this just as the trait method. If you need try_clone, either import the trait or use fully qualified syntax.
| }, | ||
|
|
||
| /// Spherical quantization bit width. Disk-index builds currently support only 1 bit. | ||
| Spherical(usize), |
There was a problem hiding this comment.
This is a perfect example of where you can use types to make this code far more robust. (please excuse typos and brevity here, I am typing with one hand and everything takes like 8x longer). The problem with a raw usize is that parsing needs to know that only 1 bit is supported, users need to read the docs, and all uses of this field need to revalidate that it is just "1". We can do better like this:
enum SphericalBits {
One,
}
impl SphericalBits {
fn as_usize(&self) -> usize {
match self {
Self::One => 1,
}
}
}
impl FromStr for SphericalBits {
// Reject unsupported values and parse errors.
}This does a lot for you:
- Structurally encodes the restriction instead of spreading in across all users and via documentation.
- Matching on
SphericalBitsmeans if you ever add support for 2 or 4 bits, the compiler helpfully shows you all the parts you need to update. - Uses of
SphericalBitsdon't need to validate. That work is already done.
| train_data, | ||
| TransformKind::PaddingHadamard { | ||
| target_dim: TargetDim::Natural, | ||
| }, |
There was a problem hiding this comment.
Use DoubleHadamard. It's better across all dimensions.
| diskann_error!( | ||
| ErrorKind::IndexError, | ||
| "Failed to train spherical quantizer: {}", | ||
| err |
There was a problem hiding this comment.
Might want to wrap in diskann_quantization::error::Format to render the full source chain.
| let train_data = | ||
| MatrixView::try_from(&train_data, train_size, train_dim).bridge_err()?; | ||
| let metric: SupportedMetric = | ||
| index_configuration.dist_metric.try_into().bridge_err()?; |
There was a problem hiding this comment.
I wonder if a misconfiguration here should be caught early? Alternatively, if CosineNormalizedis used, it can be remapped to SupportedMetric::Cosine without performance penalty.
| }; | ||
|
|
||
| /// Quantizer types used specifically for async disk index building. | ||
| #[derive(Clone)] |
There was a problem hiding this comment.
The code seems to compile fine without Clone. It could be removed to avoid the need for Arc, but I will let the crate maintainers weigh in.
Summary
Validation
cargo test -p diskann-disk(250 unit tests and 2 doc tests passed)cargo test -p diskann-disk test_spherical_disk_index_builder_with_metric(2 passed)cargo fmt --all --checkgit diff --check