diskann_disk/search/pq/
quantizer_preprocess.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use diskann::ANNResult;
7use diskann_vector::distance::Metric;
8
9use diskann_providers::model::compute_pq_distance;
10use diskann_providers::utils::BridgeErr;
11
12use super::{PQData, PQScratch};
13
14/// Preprocesses the query vector for PQ distance calculations.
15/// This function rotates the query vector and prepares the PQ table distances
16/// for efficient computation during search operations.
17pub fn quantizer_preprocess(
18    pq_scratch: &mut PQScratch,
19    pq_data: &PQData,
20    metric: Metric,
21    id_to_calculate_pq_distance: &[u32],
22) -> ANNResult<()> {
23    let table = pq_data.pq_table();
24    let expected_len = table.ncenters() * table.nchunks();
25    let dst = diskann_utils::views::MutMatrixView::try_from(
26        &mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len],
27        table.nchunks(),
28        table.ncenters(),
29    )
30    .bridge_err()?;
31
32    match metric {
33        // Prior to the introduction of the `quantizer_preprocess` method, the
34        // disk index was hard-coded to use L2 distance for comparisons.
35        //
36        // We're keeping that behavior here - treating `Cosine` and `CosineNormalized`
37        // as L2 until a more thorough evaluation can be made.
38        Metric::L2 | Metric::Cosine | Metric::CosineNormalized => {
39            table.process_into::<diskann_quantization::distances::SquaredL2>(
40                &pq_scratch.query_scratch,
41                dst,
42            );
43        }
44        Metric::InnerProduct => {
45            table.process_into::<diskann_quantization::distances::InnerProduct>(
46                &pq_scratch.query_scratch,
47                dst,
48            );
49        }
50    }
51
52    // Compute the pq distance between query vector to all the vertex in the pq
53    // calculation id scratch.
54    compute_pq_distance(
55        id_to_calculate_pq_distance,
56        pq_data.get_num_chunks(),
57        &pq_scratch.aligned_pqtable_dist_scratch,
58        pq_data.pq_compressed_data().as_slice(),
59        &mut pq_scratch.aligned_pq_coord_scratch,
60        &mut pq_scratch.aligned_dist_scratch,
61    )?;
62
63    Ok(())
64}
diskann_disk/search/pq/quantizer_preprocess.rs

diskann_disk/search/pq/
quantizer_preprocess.rs
