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 logical dimension of the original (pre-quantization) vectors.
64    pub fn get_dim(&self) -> usize {
65        match &self.pq_pivot_table {
66            PQTable::Transposed(table) => table.dim(),
67            PQTable::Fixed(table) => table.get_dim(),
68        }
69    }
70
71    /// Return the number of chunks in the underlying PQ schema.
72    pub fn get_num_chunks(&self) -> usize {
73        match &self.pq_pivot_table {
74            PQTable::Transposed(table) => table.nchunks(),
75            PQTable::Fixed(table) => table.get_num_chunks(),
76        }
77    }
78
79    /// Return the number of centers in the underlying PQ schema.
80    pub fn get_num_centers(&self) -> usize {
81        match &self.pq_pivot_table {
82            PQTable::Transposed(table) => table.ncenters(),
83            PQTable::Fixed(table) => table.get_num_centers(),
84        }
85    }
86
87    /// Get pq_compressed_data
88    pub fn pq_compressed_data(&self) -> &Matrix<u8> {
89        &self.pq_compressed_data
90    }
91
92    // Get compressed vector with the given vector id from the pq_compressed_data.
93    pub fn get_compressed_vector(&self, vector_id: usize) -> ANNResult<&[u8]> {
94        self.pq_compressed_data.get_row(vector_id).ok_or_else(|| {
95            ANNError::log_index_error("Vector id is out of boundary in the compressed dataset.")
96        })
97    }
98}
99
100#[cfg(test)]
101mod tests {
102
103    use super::*;
104
105    fn create_pq_data() -> ANNResult<PQData> {
106        let dim = 2;
107
108        let pq_pivot_table = FixedChunkPQTable::new(
109            dim,
110            Box::new([0.0, 0.0, 1.0, 1.0]),
111            Box::new([0.0, 0.0]),
112            Box::new([0, 2]),
113        )
114        .unwrap();
115        let pq_compressed_data = Matrix::try_from(Box::new([123u8, 111, 255]) as Box<[u8]>, 3, 1)
116            .expect("valid matrix shape");
117
118        PQData::new(pq_pivot_table, pq_compressed_data)
119    }
120
121    #[test]
122    fn test_get_compressed_vector() {
123        let dataset = create_pq_data().unwrap();
124
125        let vector_id = 0;
126        let result = dataset.get_compressed_vector(vector_id).unwrap();
127        assert_eq!(result, &[123]);
128
129        let vector_id = 1;
130        let result = dataset.get_compressed_vector(vector_id).unwrap();
131        assert_eq!(result, &[111]);
132
133        let vector_id = 2;
134        let result = dataset.get_compressed_vector(vector_id).unwrap();
135        assert_eq!(result, &[255]);
136    }
137
138    #[test]
139    fn test_get_num_chunks() {
140        let dataset = create_pq_data().unwrap();
141        assert_eq!(dataset.get_num_chunks(), 1);
142    }
143
144    #[test]
145    fn test_get_num_centers() {
146        let dataset = create_pq_data().unwrap();
147        assert_eq!(dataset.get_num_centers(), 2);
148    }
149}
diskann_disk/storage/quant/pq/pq_dataset.rs

diskann_disk/storage/quant/pq/
pq_dataset.rs
