Supporting spherical quantization builds for the disk index by partychen · Pull Request #1331 · microsoft/DiskANN

juchen-ms (partychen) · 2026-08-13T07:01:35Z

Summary

add spherical quantization as a disk index build option
train and wire the 1-bit spherical quantizer into one-shot and sharded in-memory builds
account for spherical vector storage in build RAM estimation
validate the currently supported 1-bit configuration
add disk build and search coverage for L2, inner product, and cosine metrics

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 --check
git diff --check

Copilot

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 as SQ_<nbits>_None, but FromStr only accepts SQ_<nbits> for the default stddev. Because serde Serialize uses to_string() and Deserialize uses from_str, SQ-with-default cannot roundtrip (e.g. bincode serialize then deserialize fails). Consider emitting SQ_<nbits> when standard_deviation is None to 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_type test currently asserts SQ_8_None for the default-SQ formatting, which encodes the same Display/parse mismatch described above. If Display is changed to emit SQ_<nbits> when standard_deviation is None, 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_serialization documents (and works around) the fact that SQ-with-default-stddev doesn't roundtrip. If Display/parse are made consistent for standard_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 Comments Bot (codecov-commenter) · 2026-08-13T07:39:17Z

Codecov Report

❌ Patch coverage is 98.78049% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.58%. Comparing base (3218478) to head (5ba3378).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-disk/src/build/builder/quantizer.rs	96.87%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1331      +/-   ##
==========================================
+ Coverage   91.55%   91.58%   +0.02%     
==========================================
  Files         522      522              
  Lines       99541    99602      +61     
==========================================
+ Hits        91139    91220      +81     
+ Misses       8402     8382      -20

Flag	Coverage Δ
miri	`91.58% <98.78%> (+0.02%)`	⬆️
unittests	`91.26% <98.78%> (+0.02%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann-benchmark/src/inputs/disk.rs	`1.44% <ø> (ø)`
diskann-disk/src/build/builder/build.rs	`92.82% <100.00%> (ø)`
diskann-disk/src/build/builder/core.rs	`97.28% <100.00%> (+<0.01%)`	⬆️
diskann-disk/src/build/builder/inmem_builder.rs	`86.61% <100.00%> (+0.90%)`	⬆️
...disk/src/build/configuration/quantization_types.rs	`98.11% <100.00%> (+0.45%)`	⬆️
diskann-quantization/src/spherical/quantizer.rs	`98.27% <100.00%> (+0.48%)`	⬆️
diskann-disk/src/build/builder/quantizer.rs	`93.81% <96.87%> (+0.95%)`	⬆️

... and 9 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Mark Hildebrand (hildebrandmw)

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.

Mark Hildebrand (hildebrandmw) · 2026-08-13T15:34:32Z

+    /// Return an independently allocated copy of this quantizer.
+    pub fn try_clone(&self) -> Result<Self, AllocatorError> {
+        <Self as TryClone>::try_clone(self)
+    }


Please leave this just as the trait method. If you need try_clone, either import the trait or use fully qualified syntax.

TryClone is currently crate-private, so diskann-disk cannot import the trait or use fully qualified syntax across the crate boundary.

The trait also explicitly documents why it should not be exposed:

/// Keep this `pub(crate)` for now because we do not want general users of the crate /// relying on the current implementations for [`Poly`]. In particular, the base case should /// be `Poly<T> where T: TryClone` instead of `Poly<T> where T: Clone`.

Making TryClone public here would contradict that intent and unnecessarily expand the public API. I propose keeping the narrow inherent SphericalQuantizer::try_clone() wrapper instead. Does that sound reasonable?

Mark Hildebrand (hildebrandmw) · 2026-08-13T15:47:37Z

+                    diskann_error!(
+                        ErrorKind::IndexError,
+                        "Failed to train spherical quantizer: {}",
+                        err


Might want to wrap in diskann_quantization::error::Format to render the full source chain.

Thanks. I initially used Format, but switched to ANNError::new(err).context(...) following the updated ANNError guidance. This still renders the full source chain while preserving TrainError for downcasting instead of flattening it into a string.

Mark Hildebrand (hildebrandmw) · 2026-08-13T15:49:13Z

+                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()?;


I wonder if a misconfiguration here should be caught early? Alternatively, if CosineNormalizedis used, it can be remapped to SupportedMetric::Cosine without performance penalty.

We chose the early-validation path. The metric conversion now happens before sampling or training, so CosineNormalized is rejected immediately as unsupported for spherical builds.

Wei Wu (wuw92)

Thanks for your change to support spherical in disk path. Is spherical disk-index support intentionally limited to 1-bit quantization, or is support for other bit widths planned?

Wei Wu (wuw92) · 2026-08-14T05:03:58Z

+                );
+                let mut rnd = rng.create_rnd();
+                let (train_data, train_size, train_dim) = pq_storage
+                    .get_random_train_data_slice::<Data::VectorDataType, _>(


BuildQuantizer::train only uses PQStorage to recover the source data_path, and get_random_train_data_slice immediately delegates to the quantizer-neutral gen_random_slice. Could we pass data_path: &str instead of pq_storage: &PQStorage here and sample directly, so SQ, spherical, and future build quantizers do not depend on PQ output storage?

Wei Wu (wuw92) · 2026-08-14T06:07:46Z

+                .map_err(|err| {
+                    diskann_error!(
+                        ErrorKind::IndexError,
+                        "Failed to train spherical quantizer: {}",
+                        Format(err)
+                    )
+                })?;


nit: Could we preserve TrainError as the structured source and attach the operation as context here?

Suggested change

.map_err(|err| {

diskann_error!(

ErrorKind::IndexError,

"Failed to train spherical quantizer: {}",

Format(err)

)

})?;

.map_err(|err| {

diskann::ANNError::new(err)

.context("Failed to train spherical quantizer")

})?;

The current Format(err) eagerly flattens the source chain into the tagged error message. Keeping TrainError inside ANNError preserves downcasting and its source chain while retaining the useful training context.

Yes. Changed!

Wei Wu (wuw92) · 2026-08-14T06:49:46Z


+    #[rstest]
+    fn test_spherical_disk_index_builder_with_metric(
+        #[values(Metric::InnerProduct, Metric::Cosine, Metric::CosineNormalized)] metric: Metric,


The SIFT fixture vectors are not normalized; Metric::CosineNormalized requires normalized data and queries.

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

juchen-ms (partychen) · 2026-08-14T08:53:04Z

Thanks for your change to support spherical in disk path. Is spherical disk-index support intentionally limited to 1-bit quantization, or is support for other bit widths planned?

The 1-bit limitation is intentional for this PR. The immediate goal is to provide a spherical replacement for the 1-bit scalar-quantized disk builds used in production. The underlying async spherical provider already supports 1, 2, and 4 bits, and the typed SphericalBits configuration leaves room to add those widths later, but enabling them would expand the build configuration and test matrix beyond the current request.

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

Support spherical quantization builds for disk indexes

5e10501

juchen-ms (partychen) requested review from a team and a lite review from Copilot August 13, 2026 07:01

Copilot started reviewing on behalf of juchen-ms (partychen) August 13, 2026 07:02 View session

Handle spherical quantization in disk benchmark summary

b6614eb

Copilot AI reviewed Aug 13, 2026

View reviewed changes

Mark Hildebrand (hildebrandmw) reviewed Aug 13, 2026

View reviewed changes

Wei Wu (wuw92) reviewed Aug 14, 2026

View reviewed changes

juchen-ms (partychen) force-pushed the juchen-microsoft-spherical-disk-builds branch from 87bfa99 to b6614eb Compare August 14, 2026 07:15

Address spherical quantization review feedback

ddd513c

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

Refine disk build quantizer training

5ba3378

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

Uh oh!

Conversation

juchen-ms (partychen) commented Aug 13, 2026

Summary

Validation

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Codecov Comments Bot (codecov-commenter) commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Mark Hildebrand (hildebrandmw) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Wei Wu (wuw92) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

juchen-ms (partychen) commented Aug 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Codecov Comments Bot (codecov-commenter) commented Aug 13, 2026 •

edited

Loading