diskann_disk/storage/quant/pq/
pq_dataset.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use core::fmt::Debug;
7
8use diskann::{ANNError, ANNResult};
9use diskann_providers::model::FixedChunkPQTable;
10use diskann_quantization::product::TransposedTable;
11use diskann_utils::views::Matrix;
12
13/// Behind the scenes, we can use either the [`FixedChunkPQTable`] or a
14/// [`diskann_quantization::product::TransposedTable`]. The [`TransposedTable`] is much faster
15/// for preprocessing, but does not support removal of the dataset centroid.
16///
17/// So, we can only use the [`TransposedTable`] when the dataset centroid
18/// is all zero.
19#[derive(Debug)]
20pub enum PQTable {
21    Transposed(TransposedTable),
22    Fixed(FixedChunkPQTable),
23}
24
25#[derive(Debug)]
26pub struct PQData {
27    // pq pivot table.
28    pq_pivot_table: PQTable,
29
30    // pq compressed vectors, shape `num_points × num_pq_chunks`.
31    pq_compressed_data: Matrix<u8>,
32}
33
34impl PQData {
35    pub fn new(
36        pq_pivot_table: FixedChunkPQTable,
37        pq_compressed_data: Matrix<u8>,
38    ) -> ANNResult<Self> {
39        // Check if we can use the transposed table. If so, go for it.
40        let centroid_is_zero = pq_pivot_table.get_centroids().iter().all(|i| *i == 0.0);
41        let pq_pivot_table = if centroid_is_zero {
42            let transposed = TransposedTable::from_parts(
43                pq_pivot_table.view_pivots(),
44                pq_pivot_table.view_offsets().to_owned(),
45            )
46            .map_err(|err| ANNError::log_pq_error(diskann_quantization::error::format(&err)))?;
47            PQTable::Transposed(transposed)
48        } else {
49            PQTable::Fixed(pq_pivot_table)
50        };
51
52        Ok(Self {
53            pq_pivot_table,
54            pq_compressed_data,
55        })
56    }
57
58    /// Get pq_table
59    pub fn pq_table(&self) -> &PQTable {
60        &self.pq_pivot_table
61    }
62
63    /// Return the number of chunks in the underlying PQ schema.
64    pub fn get_num_chunks(&self) -> usize {
65        match &self.pq_pivot_table {
66            PQTable::Transposed(table) => table.nchunks(),
67            PQTable::Fixed(table) => table.get_num_chunks(),
68        }
69    }
70
71    /// Return the number of centers in the underlying PQ schema.
72    pub fn get_num_centers(&self) -> usize {
73        match &self.pq_pivot_table {
74            PQTable::Transposed(table) => table.ncenters(),
75            PQTable::Fixed(table) => table.get_num_centers(),
76        }
77    }
78
79    /// Get pq_compressed_data
80    pub fn pq_compressed_data(&self) -> &Matrix<u8> {
81        &self.pq_compressed_data
82    }
83
84    // Get compressed vector with the given vector id from the pq_compressed_data.
85    pub fn get_compressed_vector(&self, vector_id: usize) -> ANNResult<&[u8]> {
86        self.pq_compressed_data.get_row(vector_id).ok_or_else(|| {
87            ANNError::log_index_error("Vector id is out of boundary in the compressed dataset.")
88        })
89    }
90}
91
92#[cfg(test)]
93mod tests {
94
95    use super::*;
96
97    fn create_pq_data() -> ANNResult<PQData> {
98        let dim = 2;
99
100        let pq_pivot_table = FixedChunkPQTable::new(
101            dim,
102            Box::new([0.0, 0.0, 1.0, 1.0]),
103            Box::new([0.0, 0.0]),
104            Box::new([0, 2]),
105        )
106        .unwrap();
107        let pq_compressed_data = Matrix::try_from(Box::new([123u8, 111, 255]) as Box<[u8]>, 3, 1)
108            .expect("valid matrix shape");
109
110        PQData::new(pq_pivot_table, pq_compressed_data)
111    }
112
113    #[test]
114    fn test_get_compressed_vector() {
115        let dataset = create_pq_data().unwrap();
116
117        let vector_id = 0;
118        let result = dataset.get_compressed_vector(vector_id).unwrap();
119        assert_eq!(result, &[123]);
120
121        let vector_id = 1;
122        let result = dataset.get_compressed_vector(vector_id).unwrap();
123        assert_eq!(result, &[111]);
124
125        let vector_id = 2;
126        let result = dataset.get_compressed_vector(vector_id).unwrap();
127        assert_eq!(result, &[255]);
128    }
129
130    #[test]
131    fn test_get_num_chunks() {
132        let dataset = create_pq_data().unwrap();
133        assert_eq!(dataset.get_num_chunks(), 1);
134    }
135
136    #[test]
137    fn test_get_num_centers() {
138        let dataset = create_pq_data().unwrap();
139        assert_eq!(dataset.get_num_centers(), 2);
140    }
141}
diskann_disk/storage/quant/pq/pq_dataset.rs

diskann_disk/storage/quant/pq/
pq_dataset.rs
