satoridb 0.1.2

Embedded vector database for approximate nearest neighbor search (experimental).
Documentation 
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

use crate::storage::Vector;
use flate2::read::GzDecoder;
use std::fs::File;
use std::io::{self, BufReader, Read};
use std::path::Path;

pub struct BvecsReader {
    reader: BufReader<GzDecoder<File>>,
    current_id: u64,
}

impl BvecsReader {
    pub fn new(path: impl AsRef<Path>) -> anyhow::Result<Self> {
        let file = File::open(path)?;
        let decoder = GzDecoder::new(file);
        let reader = BufReader::new(decoder);
        Ok(Self {
            reader,
            current_id: 0,
        })
    }

    /// Reads a batch of vectors.
    /// Returns empty Vec if EOF.
    pub fn read_batch(&mut self, batch_size: usize) -> anyhow::Result<Vec<Vector>> {
        let mut vectors = Vec::with_capacity(batch_size);
        let mut buf = [0u8; 4];

        for _ in 0..batch_size {
            if let Err(e) = self.reader.read_exact(&mut buf) {
                if e.kind() == io::ErrorKind::UnexpectedEof {
                    break;
                }
                return Err(e.into());
            }
            let dim = u32::from_le_bytes(buf) as usize;

            let mut data_buf = vec![0u8; dim];
            self.reader.read_exact(&mut data_buf)?;

            let data: Vec<f32> = data_buf.iter().map(|&v| v as f32).collect();
            vectors.push(Vector::new(self.current_id, data));
            self.current_id += 1;
        }

        Ok(vectors)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use flate2::write::GzEncoder;
    use flate2::Compression;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn reads_multiple_vectors() {
        let mut tmp = NamedTempFile::new().unwrap();
        {
            let mut encoder = GzEncoder::new(&mut tmp, Compression::default());
            for vec_vals in [[1u8, 2, 3], [4u8, 5, 6]] {
                encoder
                    .write_all(&(vec_vals.len() as u32).to_le_bytes())
                    .unwrap();
                encoder.write_all(&vec_vals).unwrap();
            }
            encoder.finish().unwrap();
        }

        let mut reader = BvecsReader::new(tmp.path()).unwrap();
        let batch = reader.read_batch(10).unwrap();
        assert_eq!(batch.len(), 2);
        assert_eq!(batch[0].id, 0);
        assert_eq!(batch[0].data, vec![1.0, 2.0, 3.0]);
        assert_eq!(batch[1].id, 1);
        assert_eq!(batch[1].data, vec![4.0, 5.0, 6.0]);
    }

    #[test]
    fn stops_on_eof() {
        let mut tmp = NamedTempFile::new().unwrap();
        {
            let mut encoder = GzEncoder::new(&mut tmp, Compression::default());
            encoder.write_all(&1u32.to_le_bytes()).unwrap();
            encoder.write_all(&[9u8]).unwrap();
            encoder.finish().unwrap();
        }

        let mut reader = BvecsReader::new(tmp.path()).unwrap();
        let first = reader.read_batch(1).unwrap();
        let second = reader.read_batch(1).unwrap();
        assert_eq!(first.len(), 1);
        assert!(second.is_empty());
    }
}