diskann-disk 0.56.0

DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use core::fmt::Debug;

use diskann::ANNResult;
use diskann_providers::model::FixedChunkPQTable;
use diskann_quantization::{error::Format, product::TransposedTable};
use diskann_utils::views::Matrix;

use crate::error::{diskann_error, ErrorKind};

#[derive(Debug)]
pub struct PQData {
    // pq pivot table.
    pq_pivot_table: TransposedTable,

    // pq compressed vectors, shape `num_points × num_pq_chunks`.
    pq_compressed_data: Matrix<u8>,
}

impl PQData {
    pub fn new(
        pq_pivot_table: FixedChunkPQTable,
        pq_compressed_data: Matrix<u8>,
    ) -> ANNResult<Self> {
        let pq_pivot_table = TransposedTable::from_parts(
            pq_pivot_table.view_pivots(),
            pq_pivot_table.view_offsets().to_owned(),
        )
        .map_err(|err| diskann_error!(ErrorKind::PQError, "{}", Format(err)))?;

        Ok(Self {
            pq_pivot_table,
            pq_compressed_data,
        })
    }

    /// Get pq_table
    pub fn pq_table(&self) -> &TransposedTable {
        &self.pq_pivot_table
    }

    /// Return the logical dimension of the original (pre-quantization) vectors.
    pub fn get_dim(&self) -> usize {
        self.pq_pivot_table.dim()
    }

    /// Return the number of chunks in the underlying PQ schema.
    pub fn get_num_chunks(&self) -> usize {
        self.pq_pivot_table.nchunks()
    }

    /// Return the number of centers in the underlying PQ schema.
    pub fn get_num_centers(&self) -> usize {
        self.pq_pivot_table.ncenters()
    }

    /// Get pq_compressed_data
    pub fn pq_compressed_data(&self) -> &Matrix<u8> {
        &self.pq_compressed_data
    }

    // Get compressed vector with the given vector id from the pq_compressed_data.
    pub fn get_compressed_vector(&self, vector_id: usize) -> ANNResult<&[u8]> {
        self.pq_compressed_data.get_row(vector_id).ok_or_else(|| {
            diskann_error!(
                ErrorKind::IndexError,
                "Vector id is out of boundary in the compressed dataset."
            )
        })
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    fn create_pq_data() -> ANNResult<PQData> {
        let dim = 2;

        let pq_pivot_table =
            FixedChunkPQTable::new(dim, Box::new([0.0, 0.0, 1.0, 1.0]), Box::new([0, 2])).unwrap();
        let pq_compressed_data = Matrix::try_from(Box::new([123u8, 111, 255]) as Box<[u8]>, 3, 1)
            .expect("valid matrix shape");

        PQData::new(pq_pivot_table, pq_compressed_data)
    }

    #[test]
    fn test_get_compressed_vector() {
        let dataset = create_pq_data().unwrap();

        let vector_id = 0;
        let result = dataset.get_compressed_vector(vector_id).unwrap();
        assert_eq!(result, &[123]);

        let vector_id = 1;
        let result = dataset.get_compressed_vector(vector_id).unwrap();
        assert_eq!(result, &[111]);

        let vector_id = 2;
        let result = dataset.get_compressed_vector(vector_id).unwrap();
        assert_eq!(result, &[255]);
    }

    #[test]
    fn test_get_num_chunks() {
        let dataset = create_pq_data().unwrap();
        assert_eq!(dataset.get_num_chunks(), 1);
    }

    #[test]
    fn test_get_num_centers() {
        let dataset = create_pq_data().unwrap();
        assert_eq!(dataset.get_num_centers(), 2);
    }
}