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};
13use crate::storage::quant::pq::pq_dataset::PQTable;
14
15/// Preprocesses the query vector for PQ distance calculations.
16/// This function rotates the query vector and prepares the PQ table distances
17/// for efficient computation during search operations.
18pub fn quantizer_preprocess(
19    pq_scratch: &mut PQScratch,
20    pq_data: &PQData,
21    metric: Metric,
22    id_to_calculate_pq_distance: &[u32],
23) -> ANNResult<()> {
24    match &pq_data.pq_table() {
25        PQTable::Transposed(table) => {
26            let dim = table.dim();
27            let expected_len = table.ncenters() * table.nchunks();
28            let dst = diskann_utils::views::MutMatrixView::try_from(
29                &mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len],
30                table.nchunks(),
31                table.ncenters(),
32            )
33            .bridge_err()?;
34
35            match metric {
36                // Prior to the introduction of the `quantizer_preprocess` method, the
37                // disk index was hard-coded to use L2 distance for comparisons.
38                //
39                // We're keeping that behavior here - treating `Cosine` and `CosineNormalized`
40                // as L2 until a more thorough evaluation can be made.
41                Metric::L2 | Metric::Cosine | Metric::CosineNormalized => {
42                    table.process_into::<diskann_quantization::distances::SquaredL2>(
43                        &pq_scratch.rotated_query[..dim],
44                        dst,
45                    );
46                }
47                Metric::InnerProduct => {
48                    table.process_into::<diskann_quantization::distances::InnerProduct>(
49                        &pq_scratch.rotated_query[..dim],
50                        dst,
51                    );
52                }
53            }
54        }
55        PQTable::Fixed(table) => {
56            match metric {
57                // Prior to the introduction of the `quantizer_preprocess` method, the
58                // disk index was hard-coded to use L2 distance for comparisons.
59                //
60                // We're keeping that behavior here - treating `Cosine` and `CosineNormalized`
61                // as L2 until a more thorough evaluation can be made.
62                Metric::L2 | Metric::Cosine | Metric::CosineNormalized => {
63                    // The scratch only stores the aligned dimension. However, preprocessing
64                    // wants the actual dimension used, so we have to shrink the rotated query
65                    // accordingly.
66                    let dim = table.get_dim();
67                    table.preprocess_query(&mut pq_scratch.rotated_query[..dim]);
68
69                    // Compute the distance between each chunk of the query to each pq centroids.
70                    table.populate_chunk_distances(
71                        pq_scratch.rotated_query.as_slice(),
72                        &mut pq_scratch.aligned_pqtable_dist_scratch,
73                    )?;
74                }
75                Metric::InnerProduct => {
76                    table.populate_chunk_inner_products(
77                        pq_scratch.rotated_query.as_slice(),
78                        &mut pq_scratch.aligned_pqtable_dist_scratch,
79                    )?;
80                }
81            }
82        }
83    }
84
85    // Compute the pq distance between query vector to all the vertex in the pq
86    // calculation id scratch.
87    compute_pq_distance(
88        id_to_calculate_pq_distance,
89        pq_data.get_num_chunks(),
90        &pq_scratch.aligned_pqtable_dist_scratch,
91        pq_data.pq_compressed_data().as_slice(),
92        &mut pq_scratch.aligned_pq_coord_scratch,
93        &mut pq_scratch.aligned_dist_scratch,
94    )?;
95
96    Ok(())
97}
diskann_disk/search/pq/quantizer_preprocess.rs

diskann_disk/search/pq/
quantizer_preprocess.rs
