[diskann-inmem] Add support for Product Quantization by hildebrandmw · Pull Request #1469 · microsoft/DiskANN

Mark Hildebrand (hildebrandmw) · 2026-10-07T00:03:18Z

Add support for PQ for diskann-inmem. This follows the work in #1458 and #1432.

For reviewers, this PR is broken into 6 commits:

Commit 1

Adds a with_capacity constructor to ObjectPool to create an initially empty pool. This is used for distance table buffering in the PQ repr.

Commit 2

Factors out Reranker to repr/internal/quantization/rerank.rs. This will be shared between spherical and PQ. In addition, an internal Metric enum is introduced for the three currently supported metrics: SquaredL2, InnerProduct, and Cosine.

I kept spherical::Rerank as a public version of the internal Rerank enum to try to decouple the public configuration API from how things are done on the back end. Yes, this is a little bit of code duplication, but keeps the API looking like

Spherical::config(
   ...
   spherical::Rerank::F16, // <-- comes from `spherical`.
)

The amount of duplication is small, keeps the interface consistent, and lets us experiment on the back end a little.

Additionally, the routines for test_expand_beam and test_prune are moved out of the spherical test module and into repr/test/mod.rs. These will be reused by the PQ representation.

Commit 3

Adds a generic Quantized representation. The extension point is a [Quantizer] trait which contains methods for quantizing data and creating distance computers. Distance computation uses BYOTE to allow implementations to create multiple specialized computers as they see fit. These specialized computers are inlined directly into the ExpandBeam and Prune trait objects used for search and prune respectively.

The goal behind this approach is to make it easy to implement scalar, minmax, and binary quantizers (since we just need to implement the Quantizer trait) and to make development of new quantizers easy. And hopefully BYOTE means that the performance cost of this is minimal.

One final advantage is that the Quantized representation is a single concrete type, minimizing monomorphization bloat of this approach.

Commit 4

This actually adds the PQ Quantizer.

Query distances are done using the usual trick of pre-populating a lookup table with distances between query chunk fragments and the PQ pivots (see #1458 for the implementation). Since these tables are fairly large (~200KB for a 192 PQ chunks), an ObjectPool is used to reuse allocations. The configuration for this pool in Config is such that an unbounded ObjectPool is not allowed.

Quant-quant distances use the PaddedTable introduced in #1458.

Tests are similar to spherical, but significantly easier since PQ can losslessly represent the test dataset.

Commit 5

Adds integration test support for PQ and baselines covering L2, inner-product, and cosine.

Commit 6

Adds support in diskann-benchmark for the new PQ backend.

Benchmarks

The infrastructure in #1458 pays off. This is considerably faster than the PQ index in diskann-providers. A benchmarking input JSON is included in the comment - highlights are reported here.

Dataset	Version	PQ Chunks	Build Time (s)	QPS (L = 50)	QPS (L = 100)
wiki-100k	current	192	33.6	11028 (0.858)	7884 (0.924)
wiki-100k	PR	192	11.4	21570 (0.859)	15653 (0.922)
msturing-100k	current	25	5.61	43739 (0.732)	25417 (0.837)
msturing-100k	PR	25	2.32	76639 (0.728)	41245 (0.841)
oai-v3-100k	current	384	96	2382 (0.957)	1469 (0.977)
oai-v3-100k	PR	384	41	9440 (0.956)	5710 (0.977)

Note that the performance increase is somewhat exaggerated due to the small size of the dataset (less memory bandwidth because more fits in cache), which highlights perf differences that a bandwidth-constrained environment would hide.

Codecov Comments Bot (codecov-commenter) · 2026-10-09T16:01:05Z

Codecov Report

❌ Patch coverage is 95.67980% with 68 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.11%. Comparing base (7a8ed89) to head (0aa580f).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-benchmark/src/index/inmem2.rs	85.98%	22 Missing ⚠️
...iskann-inmem/src/repr/internal/quantization/mod.rs	45.45%	18 Missing ⚠️
diskann-inmem/src/repr/quantized/mod.rs	98.39%	10 Missing ⚠️
diskann-inmem/integration/index/runner.rs	88.88%	8 Missing ⚠️
diskann-inmem/src/repr/quantized/product.rs	98.78%	5 Missing ⚠️
...ann-inmem/src/repr/internal/quantization/rerank.rs	97.56%	4 Missing ⚠️
diskann-inmem/src/repr/spherical.rs	94.44%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1469      +/-   ##
==========================================
+ Coverage   91.03%   92.11%   +1.08%     
==========================================
  Files         586      590       +4     
  Lines      119926   121332    +1406     
==========================================
+ Hits       109176   111767    +2591     
+ Misses      10750     9565    -1185

Flag	Coverage Δ
miri	`92.11% <95.67%> (+1.08%)`	⬆️
unittests	`92.06% <95.67%> (+1.32%)`	⬆️

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

Files with missing lines	Coverage Δ
diskann-inmem/integration/index/object.rs	`99.00% <ø> (ø)`
diskann-inmem/integration/main.rs	`82.05% <100.00%> (+0.71%)`	⬆️
diskann-inmem/src/repr/internal/mod.rs	`100.00% <ø> (ø)`
diskann-inmem/src/repr/mod.rs	`93.75% <ø> (ø)`
diskann-inmem/src/repr/test/mod.rs	`94.77% <100.00%> (+4.77%)`	⬆️
diskann-utils/src/object_pool.rs	`99.60% <100.00%> (+0.02%)`	⬆️
diskann-inmem/src/repr/spherical.rs	`97.24% <94.44%> (-0.23%)`	⬇️
...ann-inmem/src/repr/internal/quantization/rerank.rs	`97.56% <97.56%> (ø)`
diskann-inmem/src/repr/quantized/product.rs	`98.78% <98.78%> (ø)`
diskann-inmem/integration/index/runner.rs	`92.30% <88.88%> (+0.45%)`	⬆️
... and 3 more

... and 47 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) · 2026-10-09T16:54:35Z

Benchmarking JSON

{
  "search_directories": [
    "redacted"
  ],
  "output_directory": null,
  "jobs": [
    {
      "type": "inmem2",
      "content": {
        "build": {
          "alpha": 1.2,
          "l_build": 50,
          "max_degree": 42,
          "num_threads": 8,
          "pruned_degree": 32
        },
        "data": {
          "data": "wikipedia/wikipedia_base_100K.bin",
          "data_type": "float32",
          "distance": "inner_product"
        },
        "search": {
          "groundtruth": "wikipedia/wikipedia-100k.gt.bin",
          "num_threads": [
            8
          ],
          "queries": "wikipedia/wikipedia_query.bin",
          "reps": 3,
          "runs": [
            {
              "recall_k": 10,
              "search_l": [
                10,
                20,
                30,
                40,
                50,
                60,
                70,
                80,
                90,
                100
              ],
              "search_n": 10
            }
          ]
        },
        "quantization": {
          "product": {
            "chunks": 192,
            "centers": 256,
            "seed": 1
          }
        }
      }
    },
    {
      "type": "inmem2",
      "content": {
        "build": {
          "alpha": 1.2,
          "l_build": 50,
          "max_degree": 42,
          "num_threads": 8,
          "pruned_degree": 32
        },
        "data": {
          "data": "msturing/msturing-base-100k.fbin",
          "data_type": "float32",
          "distance": "squared_l2"
        },
        "search": {
          "groundtruth": "msturing/msturing-100k.gt.10k.bin",
          "num_threads": [
            8
          ],
          "queries": "msturing/queries-10k.fbin",
          "reps": 3,
          "runs": [
            {
              "recall_k": 10,
              "search_l": [
                10,
                20,
                30,
                40,
                50,
                60,
                70,
                80,
                90,
                100
              ],
              "search_n": 10
            }
          ]
        },
        "quantization": {
          "product": {
            "chunks": 25,
            "centers": 256,
            "seed": 1
          }
        }
      }
    },
    {
      "type": "inmem2",
      "content": {
        "build": {
          "alpha": 1.2,
          "l_build": 50,
          "max_degree": 42,
          "num_threads": 8,
          "pruned_degree": 32
        },
        "data": {
          "data": "openai-v3/base-100k.bin",
          "data_type": "float32",
          "distance": "cosine"
        },
        "search": {
          "groundtruth": "openai-v3/gt-100k.bin",
          "num_threads": [
            8
          ],
          "queries": "openai-v3/query.bin",
          "reps": 3,
          "runs": [
            {
              "recall_k": 10,
              "search_l": [
                10,
                20,
                30,
                40,
                50,
                60,
                70,
                80,
                90,
                100
              ],
              "search_n": 10
            }
          ]
        },
        "quantization": {
          "product": {
            "chunks": 384,
            "centers": 256,
            "seed": 1
          }
        }
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "wikipedia/wikipedia-100k.gt.bin",
            "num_threads": [
              8
            ],
            "queries": "wikipedia/wikipedia_query.bin",
            "reps": 3,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "wikipedia/wikipedia_base_100K.bin",
            "data_type": "float32",
            "distance": "inner_product",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 8,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 196,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "msturing/msturing-100k.gt.10k.bin",
            "num_threads": [
              8
            ],
            "queries": "msturing/queries-10k.fbin",
            "reps": 3,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "msturing/msturing-base-100k.fbin",
            "data_type": "float32",
            "distance": "squared_l2",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 8,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 25,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "openai-v3/gt-100k.bin",
            "num_threads": [
              8
            ],
            "queries": "openai-v3/query.bin",
            "reps": 3,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "openai-v3/base-100k.bin",
            "data_type": "float32",
            "distance": "cosine",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 8,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 384,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    }
  ]
}

Copilot

🟡 Changes recommended

A duplicated match arm can fail warnings-as-errors Clippy checks, alongside several documentation and diagnostic issues.

5 open findings

Combine identical match arms · New
Use a representation-neutral error for invalid product inputs · New
Document the new configuration entry point and error contract · New
Fix “happy-patch” typo · New
Fix ExpandBeam misspelling · New

What changed in this PR

Adds Product Quantization support to diskann-inmem, including shared reranking infrastructure, integration coverage, and benchmarking.

Changes:

Implements PQ search, insertion, pruning, metrics, reranking, and distance-table pooling.
Refactors spherical quantization to share internal metric and reranking logic.
Adds integration baselines and benchmark support.

File	Description
`diskann-utils/src/object_pool.rs`	Adds empty capacity-bounded pool construction.
`diskann-inmem/src/repr/test/mod.rs`	Shares quantized representation test helpers.
`diskann-inmem/src/repr/spherical.rs`	Adopts shared quantization internals.
`diskann-inmem/src/repr/product.rs`	Implements the PQ representation.
`diskann-inmem/src/repr/mod.rs`	Exports PQ APIs.
`diskann-inmem/src/repr/internal/quantization/rerank.rs`	Adds shared reranking logic.
`diskann-inmem/src/repr/internal/quantization/mod.rs`	Adds internal metric conversions.
`diskann-inmem/src/repr/internal/mod.rs`	Registers quantization internals.
`diskann-inmem/integration/main.rs`	Registers PQ regression tests.
`diskann-inmem/integration/jsons/graph/product/l2.json`	Configures L2 PQ tests.
`diskann-inmem/integration/jsons/graph/product/l2-baseline.json`	Records L2 baselines.
`diskann-inmem/integration/jsons/graph/product/ip.json`	Configures inner-product PQ tests.
`diskann-inmem/integration/jsons/graph/product/ip-baseline.json`	Records inner-product baselines.
`diskann-inmem/integration/jsons/graph/product/cosine.json`	Configures cosine PQ tests.
`diskann-inmem/integration/jsons/graph/product/cosine-baseline.json`	Records cosine baselines.
`diskann-inmem/integration/index/runner.rs`	Builds PQ integration indexes.
`diskann-inmem/integration/index/object.rs`	Enables PQ index dispatch.
`diskann-benchmark/src/index/inmem2.rs`	Adds PQ benchmark training and execution.

🧠 Review effort: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Copilot

🟡 Changes recommended

The PQ constructor currently accepts tables with more centers than its byte-sized encoding can represent.

1 open finding

Validate center count before constructing an unusable quantizer · New

5 resolved since last review

🧠 Review effort: Balanced

Jack Moffitt (metajack)

Here are comments for the first three commits. Will finish on Monday.

This adds support for scalar quantization (bits 1 through 8) to `diskann-inmem`, follow the pattern mostly set up in #1469. This does require the `Quantizer` trait to differentiate between distances for insert and distances for general queries due to the need to rescale queries for wiki-cohere embedding. I've found several areas in `diskann-quantization` that need help: 1. Compression and distance methods should not panic on invalid vector lengths - they should return error types instead (that would remove some of the semi-redundant length checking in this PR). 2. `diskann-quantization` needs better support for asymmetrical compression. It's really a shame that we need to handle the somewhat sketch query rescaling for inner-products manually.

Mark Hildebrand (hildebrandmw) added this pull request to stack #1468 October 7, 2026 00:03

Base automatically changed from mhildebr/basic-table to main October 8, 2026 18:18

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-pq-2 branch from 0a45cfa to f9ddb82 Compare October 9, 2026 15:43

Mark Hildebrand (hildebrandmw) marked this pull request as ready for review October 9, 2026 16:52

Mark Hildebrand (hildebrandmw) requested review from a team and a balanced review from Copilot October 9, 2026 16:52

Copilot started reviewing on behalf of Mark Hildebrand (hildebrandmw) October 9, 2026 16:53

View session

Copilot AI reviewed Oct 9, 2026

View reviewed changes

Comment thread diskann-inmem/src/repr/product.rs Outdated

Comment thread diskann-inmem/integration/index/runner.rs

Comment thread diskann-inmem/src/repr/product.rs Outdated

Comment thread diskann-inmem/src/repr/product.rs Outdated

Comment thread diskann-inmem/src/repr/test/mod.rs Outdated

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-pq-2 branch 2 times, most recently from 514ba3a to 4711c41 Compare October 9, 2026 20:01

Mark Hildebrand (hildebrandmw) marked this pull request as draft October 12, 2026 22:42

Mark Hildebrand (hildebrandmw) marked this pull request as ready for review October 12, 2026 22:42

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-pq-2 branch from 4711c41 to e153d29 Compare October 13, 2026 20:43

Mark Hildebrand (hildebrandmw) requested a balanced review from Copilot October 13, 2026 20:53

Copilot started reviewing on behalf of Mark Hildebrand (hildebrandmw) October 13, 2026 20:53

View session

Copilot AI reviewed Oct 13, 2026

View reviewed changes

Comment thread diskann-inmem/src/repr/quantized/product.rs

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-pq-2 branch from e153d29 to f110aec Compare October 13, 2026 21:07

Harsha Simhadri (harsha-simhadri) approved these changes Oct 14, 2026

View reviewed changes

This was referenced Oct 14, 2026

[diskann-inmem] Add scalar quantization #1484

Merged

[diskann-inmem] Spherical Quantization Round 2 #1485

Open

Jack Moffitt (metajack) reviewed Oct 17, 2026

View reviewed changes

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-pq-2 branch from 1cac598 to 88b694e Compare October 19, 2026 21:12

Mark Hildebrand added 2 commits October 19, 2026 14:29

Add ObjectPool::with_capacity for lazy initialization.

6baf0ca

Extract quantized infrastructure.

3851713

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-pq-2 branch from 88b694e to aa5f142 Compare October 19, 2026 21:29

Mark Hildebrand added 3 commits October 19, 2026 14:34

Add generic quantized representation.

49a2735

Add product quantizer.

c573447

Add integration tests.

197a6e8

Add benchmark.

0aa580f

Mark Hildebrand (hildebrandmw) force-pushed the mhildebr/inmem-pq-2 branch from aa5f142 to 0aa580f Compare October 19, 2026 21:34

Jack Moffitt (metajack) approved these changes Oct 20, 2026

View reviewed changes

Mark Hildebrand (hildebrandmw) merged commit 14c0366 into main Oct 20, 2026
31 checks passed

Mark Hildebrand (hildebrandmw) deleted the mhildebr/inmem-pq-2 branch October 20, 2026 15:55

Navigation Menu

Repository navigation

[diskann-inmem] Add support for Product Quantization - #1469

[diskann-inmem] Add support for Product Quantization#1469
Mark Hildebrand (hildebrandmw) merged 6 commits into
mainfrom
mhildebr/inmem-pq-2

Mark Hildebrand (hildebrandmw) commented Oct 7, 2026 •

edited

Loading

Uh oh!

Codecov Comments Bot (codecov-commenter) commented Oct 9, 2026 •

edited

Loading

Uh oh!

Mark Hildebrand (hildebrandmw) commented Oct 9, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Jack Moffitt (metajack) left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

Mark Hildebrand (hildebrandmw) commented Oct 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Commit 1

Commit 2

Commit 3

Commit 4

Commit 5

Commit 6

Benchmarks

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

Mark Hildebrand (hildebrandmw) commented Oct 9, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Jack Moffitt (metajack) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Mark Hildebrand (hildebrandmw) commented Oct 7, 2026 •

edited

Loading

Codecov Comments Bot (codecov-commenter) commented Oct 9, 2026 •

edited

Loading