[diskann-quantization] PQ Infrastructure for diskann-inmem

[Mark Hildebrand (hildebrandmw)]

Add infrastructure for:

• PQ pre-computed distance lookups (currently this).
• Computation of cosine-specific lookup-tables.
• Saner quant-quant distances or random-access full-quant distances for pruning.

Distance Lookups (diskann_quantization/src/product/tables/lookup.rs)

PQ distance computations for a fixed query can be made faster by pre-computing distances between a query and all PQ pivots into a lookup table and indexing into that table by PQ codes. This PR introduces lookup_single, which unrolls the lookup table computation.

In addition, it supports cosine computations via the DotAndNorm type. This works by computing both the inner product between a query chunk and pivot chunk as the "dot" field, and storing the squared pivot-chunk norm as well. After lookup, the final DotAndNorm contains both the inner product and square norm between the query vector and the entire reconstruct PQ vector, and can be combined with the query norm for the final cosine similarity.

Since distance lookup tables are non-trivial allocations, this API (and the preparation stage in the TransposedTable) need a little bit of work to use for full distance computations (see the changes in diskann-benchmark). This is to provide full control of how allocations happen to the caller. The current lookup tables in diskann-providers require pooled objects, which is too opinionated for diskann-quantization.

Cosine Distance Preparation

Infrastructure is added to diskann-quantization/src/product/tables/transposed/{pivots.rs, table.rs} to do fast(er) preparation of DotAndNorm based lookup tables. This is functionality that is lacking in the current PQ infrastructure, requiring either use of L2 or a slow fallback.

Supporting fast(er) quant-quant distances

To support graph pruning, we need a way of doing relatively fast quant-quant distances. The methods in FixedChunkPQTable are quite branchy, and a dedicated quant-quant distance lookup table requires on the order of 32KB to 131KB of storage per-chunk which is prohibitively large.

This PR introduces diskann-quantization/src/product/tables/padded.rs. The idea here is to ensure that the pivots for each chunk are contiguous in memory, and then padding all pivots to the same length (a multiple of an underlying SIMD width). This organization and padding makes quant-quant distance computations much more regular. Full-quant distances are still a little messy, unfortunately.

Testing

diskann-providers has decent infrastructure for testing PQ based distances. This PR ports this infrastructure to diskann-quantization/src/product/tables/test.rs and uses it to test end-to-end distances via both the PaddedTable and the look-up table based TransposedTable.

Suggested Reviewing Order

• diskann-quantization/src/product/tables/test.rs: Shared distance test infrastructure. Extends the Check type in diskann-quantization/src/test_util.rs.
• diskann-quantization/src/product/tables/lookup.rs: Implementation of distance table lookup.
• diskann-quantization/src/product/tables/transposed/pivots.rs: Support populating pre-processed cosine changes into DotAndNorm.
• diskann-quantization/src/product/tables/transposed/table.rs: Exposing pre-processing API for DotAndNorm based cosine distances.
• diskann-quantization/src/product/tables/transposed/mod.rs: End-to-end distance tests.
• diskann-quantization/src/product/tables/padded.rs: New padded table for quant-quant distances.
• diskann-benchmark/src/exhaustive/product.rs: Extending the exhaustive search benchmark to use the outgoing FixedChunkPQTable as well as the PaddedTable and TransposedTable.

Benchmark Performance

See PR comment to keep commit description smaller.

[Mark Hildebrand (hildebrandmw)]

Benchmarking Performance

Using the following input

Benchmark JSON

{
  "search_directories": [
    "..."
  ],
  "output_directory": null,
  "jobs": [
    {
      "type": "exhaustive-product-quantization",
      "content": {
        "compression_threads": 8,
        "data": "wikipedia/wikipedia_base_100K.bin",
        "data_type": "float32",
        "distance": "inner_product",
        "num_pq_centers": 256,
        "num_pq_chunks": 192,
        "search": {
          "groundtruth": "wikipedia/wikipedia-100K",
          "num_threads": 8,
          "queries": "wikipedia/wikipedia_query.bin",
          "recalls": {
            "recall_k": [
              10,
              20,
              30,
              40
            ],
            "recall_n": [
              10,
              20,
              30,
              40
            ]
          }
        },
        "seed": 7831252621480178695,
        "table_style": "fixed-chunk"
      }
    },
    {
      "type": "exhaustive-product-quantization",
      "content": {
        "compression_threads": 8,
        "data": "wikipedia/wikipedia_base_100K.bin",
        "data_type": "float32",
        "distance": "inner_product",
        "num_pq_centers": 256,
        "num_pq_chunks": 192,
        "search": {
          "groundtruth": "wikipedia/wikipedia-100K",
          "num_threads": 8,
          "queries": "wikipedia/wikipedia_query.bin",
          "recalls": {
            "recall_k": [
              10,
              20,
              30,
              40
            ],
            "recall_n": [
              10,
              20,
              30,
              40
            ]
          }
        },
        "seed": 7831252621480178695,
        "table_style": "transposed"
      }
    },
    {
      "type": "exhaustive-product-quantization",
      "content": {
        "compression_threads": 8,
        "data": "openai-v3/base-100k.bin",
        "data_type": "float32",
        "distance": "squared_l2",
        "num_pq_centers": 256,
        "num_pq_chunks": 384,
        "search": {
          "groundtruth": "openai-v3/gt-100k.bin",
          "num_threads": 8,
          "queries": "openai-v3/query.bin",
          "recalls": {
            "recall_k": [
              10,
              20,
              30,
              40
            ],
            "recall_n": [
              10,
              20,
              30,
              40
            ]
          }
        },
        "seed": 7831252621480178695,
        "table_style": "fixed-chunk"
      }
    },
    {
      "type": "exhaustive-product-quantization",
      "content": {
        "compression_threads": 8,
        "data": "openai-v3/base-100k.bin",
        "data_type": "float32",
        "distance": "squared_l2",
        "num_pq_centers": 256,
        "num_pq_chunks": 384,
        "search": {
          "groundtruth": "openai-v3/gt-100k.bin",
          "num_threads": 8,
          "queries": "openai-v3/query.bin",
          "recalls": {
            "recall_k": [
              10,
              20,
              30,
              40
            ],
            "recall_n": [
              10,
              20,
              30,
              40
            ]
          }
        },
        "seed": 7831252621480178695,
        "table_style": "transposed"
      }
    },
    {
      "type": "exhaustive-product-quantization",
      "content": {
        "compression_threads": 8,
        "data": "openai-v3/base-100k.bin",
        "data_type": "float32",
        "distance": "cosine",
        "num_pq_centers": 256,
        "num_pq_chunks": 384,
        "search": {
          "groundtruth": "openai-v3/gt-100k.bin",
          "num_threads": 8,
          "queries": "openai-v3/query.bin",
          "recalls": {
            "recall_k": [
              10,
              20,
              30,
              40
            ],
            "recall_n": [
              10,
              20,
              30,
              40
            ]
          }
        },
        "seed": 7831252621480178695,
        "table_style": "fixed-chunk"
      }
    },
    {
      "type": "exhaustive-product-quantization",
      "content": {
        "compression_threads": 8,
        "data": "openai-v3/base-100k.bin",
        "data_type": "float32",
        "distance": "cosine",
        "num_pq_centers": 256,
        "num_pq_chunks": 384,
        "search": {
          "groundtruth": "openai-v3/gt-100k.bin",
          "num_threads": 8,
          "queries": "openai-v3/query.bin",
          "recalls": {
            "recall_k": [
              10,
              20,
              30,
              40
            ],
            "recall_n": [
              10,
              20,
              30,
              40
            ]
          }
        },
        "seed": 7831252621480178695,
        "table_style": "padded"
      }
    },
    {
      "type": "exhaustive-product-quantization",
      "content": {
        "compression_threads": 8,
        "data": "openai-v3/base-100k.bin",
        "data_type": "float32",
        "distance": "cosine",
        "num_pq_centers": 256,
        "num_pq_chunks": 384,
        "search": {
          "groundtruth": "openai-v3/gt-100k.bin",
          "num_threads": 8,
          "queries": "openai-v3/query.bin",
          "recalls": {
            "recall_k": [
              10,
              20,
              30,
              40
            ],
            "recall_n": [
              10,
              20,
              30,
              40
            ]
          }
        },
        "seed": 7831252621480178695,
        "table_style": "transposed"
      }
    }
  ]
}

Dataset	Metric	Table Style	Compression Time	Preprocess Time	Search Time	Remark
Wikipedia-100k	Inner Product	FixedChunkPQTable	2.20s	218us	17869 us	Creates Lookup Table
Wikipedia-100k	Inner Product	Transposed	0.34s	46.1us	10746us
OpenAI-V3-100k	L2	FixedChunkPQTable	5.98s	553us	40993us	Creates Lookup Table
OpenAI-V3-100k	L2	Transposed	1.68s	237us	27582us
OpenAI-V3-100k	Cosine	FixedChunkPQTable	6.61s	3.8us	231,390us	No Lookup Table
OpenAI-V3-100k	Cosine	Padded	N/A	3.2us	187,430 us	No Lookup Table
OpenAI-V3-100k	Cosine	Transposed	1.59us	307us	33,112us

For Wikipedia, the transposed table is signficantly faster at preprocessing and compression (reflected in compression time and preprocess time). In addition, even though both the FixedChunkPQTable and TransposedTable use lookup tables, the unrolled implementation in this PR is noticeably faster.

OpenAI with L2 is a similar story.

OpenAI with Cosine gives some insight into the DotAndNorm based lookup tables. Both FixedChunkPQTable and PaddedTable perform direct distance computations, with Padded being moderately faster. However, Transposed table again is far in the lead.

Recall in all cases is identical.

[Copilot]

Copilot review overview

🟡 Changes recommended

The newly required table-style field makes existing product-exhaustive benchmark configurations fail deserialization.

Review effort: Balanced
Findings: 1

Open (1)

• Add Serde default for table_style to preserve backward compatibility · New

What changed in this PR

Adds PQ lookup, cosine-distance, and padded-table infrastructure for faster quantized distance computation and benchmarking.

Changes:

• Adds precomputed lookup tables with cosine support.
• Introduces SIMD-aware padded tables for full-quant and quant-quant distances.
• Extends benchmarks and shared distance tests across table implementations.

File	Description
diskann-quantization/​src/​views.rs	Computes maximum chunk dimensions.
diskann-quantization/​src/​test_util.rs	Adds exact-value test checks.
diskann-quantization/​src/​product/​tables/​transposed/​table.rs	Supports typed lookup outputs and cosine preprocessing.
diskann-quantization/​src/​product/​tables/​transposed/​pivots.rs	Generates cosine dot-and-norm values.
diskann-quantization/​src/​product/​tables/​transposed/​mod.rs	Adds end-to-end distance tests.
diskann-quantization/​src/​product/​tables/​test.rs	Adds shared PQ distance test infrastructure.
diskann-quantization/​src/​product/​tables/​padded.rs	Implements padded SIMD distance tables.
diskann-quantization/​src/​product/​tables/​mod.rs	Exposes new table APIs.
diskann-quantization/​src/​product/​tables/​lookup.rs	Implements precomputed distance lookup.
diskann-quantization/​src/​product/​tables/​basic.rs	Adds borrowed table views.
diskann-quantization/​src/​product/​mod.rs	Makes table modules public.
diskann-quantization/​src/​distances.rs	Adds the cosine operation marker.
diskann-disk/​src/​storage/​quant/​generator.rs	Corrects a comment.
diskann-disk/​src/​search/​pq/​quantizer_preprocess.rs	Adapts to generic lookup outputs.
diskann-benchmark/​src/​inputs/​exhaustive.rs	Adds PQ table-style configuration.
diskann-benchmark/​src/​exhaustive/​product.rs	Benchmarks all PQ table implementations.
diskann-benchmark/​example/​product-exhaustive.json	Selects the transposed benchmark style.

---

💡 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)]

Codecov Report

❌ Patch coverage is 92.25092% with 105 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.97%. Comparing base (a3c1d74) to head (695468b).

Files with missing lines	Patch %	Lines
diskann-benchmark/src/exhaustive/product.rs	55.11%	57 Missing ⚠️
diskann-quantization/src/product/tables/padded.rs	95.90%	26 Missing ⚠️
...-quantization/src/product/tables/transposed/mod.rs	91.89%	9 Missing ⚠️
diskann-quantization/src/test_util.rs	33.33%	4 Missing ⚠️
diskann-benchmark/src/inputs/exhaustive.rs	62.50%	3 Missing ⚠️
diskann-quantization/src/product/tables/lookup.rs	97.98%	3 Missing ⚠️
diskann-quantization/src/product/tables/test.rs	99.00%	2 Missing ⚠️
diskann-quantization/src/views.rs	95.23%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1458      +/-   ##
==========================================
+ Coverage   90.93%   91.97%   +1.04%     
==========================================
  Files         583      586       +3     
  Lines      116854   118178    +1324     
==========================================
+ Hits       106266   108700    +2434     
+ Misses      10588     9478    -1110     

Flag	Coverage Δ
miri	91.97% <92.25%> (+1.04%)	⬆️
unittests	91.92% <91.96%> (+1.29%)	⬆️

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

Files with missing lines	Coverage Δ
diskann-disk/src/search/pq/quantizer_preprocess.rs	100.00% <100.00%> (ø)
diskann-disk/src/storage/quant/generator.rs	98.94% <ø> (ø)
diskann-quantization/src/distances.rs	90.90% <ø> (ø)
diskann-quantization/src/product/tables/basic.rs	99.37% <100.00%> (+0.02%)	⬆️
...antization/src/product/tables/transposed/pivots.rs	98.20% <100.00%> (+0.13%)	⬆️
...uantization/src/product/tables/transposed/table.rs	99.11% <100.00%> (+0.04%)	⬆️
diskann-quantization/src/views.rs	98.77% <95.23%> (-0.25%)	⬇️
diskann-quantization/src/product/tables/test.rs	99.50% <99.00%> (-0.50%)	⬇️
diskann-benchmark/src/inputs/exhaustive.rs	59.13% <62.50%> (+0.09%)	⬆️
diskann-quantization/src/product/tables/lookup.rs	97.98% <97.98%> (ø)
... and 4 more

... and 51 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[Aditya Krishnan (arkrishn94)]

Thanks Mark, this looks great. Higher-level question - with the introduction of PaddedTable, can we get rid of the FixedChunkTable now?

[Mark Hildebrand (hildebrandmw)]

Thanks Mark, this looks great. Higher-level question - with the introduction of PaddedTable, can we get rid of the FixedChunkTable now?

Sadly, we're still a way out. I have no intention of porting this table to the inmem index in diskann-providers, so the FixedChunkPQTable will be around until that gets replaced by diskann-inmem. And we have some unfortunate internal code that decided to couple itself really closely with the layout used by the FixedChunkPQTable for some reason that needs to be unraveled.

And there's the matter of serialization and deserialization.

So not quite 😢