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 expected_len = table.ncenters() * table.nchunks();
27            let dst = diskann_utils::views::MutMatrixView::try_from(
28                &mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len],
29                table.nchunks(),
30                table.ncenters(),
31            )
32            .bridge_err()?;
33
34            match metric {
35                // Prior to the introduction of the `quantizer_preprocess` method, the
36                // disk index was hard-coded to use L2 distance for comparisons.
37                //
38                // We're keeping that behavior here - treating `Cosine` and `CosineNormalized`
39                // as L2 until a more thorough evaluation can be made.
40                Metric::L2 | Metric::Cosine | Metric::CosineNormalized => {
41                    table.process_into::<diskann_quantization::distances::SquaredL2>(
42                        &pq_scratch.query_scratch,
43                        dst,
44                    );
45                }
46                Metric::InnerProduct => {
47                    table.process_into::<diskann_quantization::distances::InnerProduct>(
48                        &pq_scratch.query_scratch,
49                        dst,
50                    );
51                }
52            }
53        }
54        PQTable::Fixed(table) => {
55            match metric {
56                // Prior to the introduction of the `quantizer_preprocess` method, the
57                // disk index was hard-coded to use L2 distance for comparisons.
58                //
59                // We're keeping that behavior here - treating `Cosine` and `CosineNormalized`
60                // as L2 until a more thorough evaluation can be made.
61                Metric::L2 | Metric::Cosine | Metric::CosineNormalized => {
62                    table.preprocess_query(&mut pq_scratch.query_scratch);
63
64                    // Compute the distance between each chunk of the query to each pq centroids.
65                    table.populate_chunk_distances(
66                        &pq_scratch.query_scratch,
67                        &mut pq_scratch.aligned_pqtable_dist_scratch,
68                    )?;
69                }
70                Metric::InnerProduct => {
71                    table.populate_chunk_inner_products(
72                        &pq_scratch.query_scratch,
73                        &mut pq_scratch.aligned_pqtable_dist_scratch,
74                    )?;
75                }
76            }
77        }
78    }
79
80    // Compute the pq distance between query vector to all the vertex in the pq
81    // calculation id scratch.
82    compute_pq_distance(
83        id_to_calculate_pq_distance,
84        pq_data.get_num_chunks(),
85        &pq_scratch.aligned_pqtable_dist_scratch,
86        pq_data.pq_compressed_data().as_slice(),
87        &mut pq_scratch.aligned_pq_coord_scratch,
88        &mut pq_scratch.aligned_dist_scratch,
89    )?;
90
91    Ok(())
92}
diskann_disk/search/pq/quantizer_preprocess.rs

diskann_disk/search/pq/
quantizer_preprocess.rs
