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#[derive(Debug)]
14pub struct PQData {
15    // pq pivot table.
16    pq_pivot_table: TransposedTable,
17
18    // pq compressed vectors, shape `num_points × num_pq_chunks`.
19    pq_compressed_data: Matrix<u8>,
20}
21
22impl PQData {
23    pub fn new(
24        pq_pivot_table: FixedChunkPQTable,
25        pq_compressed_data: Matrix<u8>,
26    ) -> ANNResult<Self> {
27        let pq_pivot_table = TransposedTable::from_parts(
28            pq_pivot_table.view_pivots(),
29            pq_pivot_table.view_offsets().to_owned(),
30        )
31        .map_err(|err| ANNError::log_pq_error(diskann_quantization::error::format(&err)))?;
32
33        Ok(Self {
34            pq_pivot_table,
35            pq_compressed_data,
36        })
37    }
38
39    /// Get pq_table
40    pub fn pq_table(&self) -> &TransposedTable {
41        &self.pq_pivot_table
42    }
43
44    /// Return the logical dimension of the original (pre-quantization) vectors.
45    pub fn get_dim(&self) -> usize {
46        self.pq_pivot_table.dim()
47    }
48
49    /// Return the number of chunks in the underlying PQ schema.
50    pub fn get_num_chunks(&self) -> usize {
51        self.pq_pivot_table.nchunks()
52    }
53
54    /// Return the number of centers in the underlying PQ schema.
55    pub fn get_num_centers(&self) -> usize {
56        self.pq_pivot_table.ncenters()
57    }
58
59    /// Get pq_compressed_data
60    pub fn pq_compressed_data(&self) -> &Matrix<u8> {
61        &self.pq_compressed_data
62    }
63
64    // Get compressed vector with the given vector id from the pq_compressed_data.
65    pub fn get_compressed_vector(&self, vector_id: usize) -> ANNResult<&[u8]> {
66        self.pq_compressed_data.get_row(vector_id).ok_or_else(|| {
67            ANNError::log_index_error("Vector id is out of boundary in the compressed dataset.")
68        })
69    }
70}
71
72#[cfg(test)]
73mod tests {
74
75    use super::*;
76
77    fn create_pq_data() -> ANNResult<PQData> {
78        let dim = 2;
79
80        let pq_pivot_table =
81            FixedChunkPQTable::new(dim, Box::new([0.0, 0.0, 1.0, 1.0]), Box::new([0, 2])).unwrap();
82        let pq_compressed_data = Matrix::try_from(Box::new([123u8, 111, 255]) as Box<[u8]>, 3, 1)
83            .expect("valid matrix shape");
84
85        PQData::new(pq_pivot_table, pq_compressed_data)
86    }
87
88    #[test]
89    fn test_get_compressed_vector() {
90        let dataset = create_pq_data().unwrap();
91
92        let vector_id = 0;
93        let result = dataset.get_compressed_vector(vector_id).unwrap();
94        assert_eq!(result, &[123]);
95
96        let vector_id = 1;
97        let result = dataset.get_compressed_vector(vector_id).unwrap();
98        assert_eq!(result, &[111]);
99
100        let vector_id = 2;
101        let result = dataset.get_compressed_vector(vector_id).unwrap();
102        assert_eq!(result, &[255]);
103    }
104
105    #[test]
106    fn test_get_num_chunks() {
107        let dataset = create_pq_data().unwrap();
108        assert_eq!(dataset.get_num_chunks(), 1);
109    }
110
111    #[test]
112    fn test_get_num_centers() {
113        let dataset = create_pq_data().unwrap();
114        assert_eq!(dataset.get_num_centers(), 2);
115    }
116}
diskann_disk/storage/quant/pq/pq_dataset.rs

diskann_disk/storage/quant/pq/
pq_dataset.rs
