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

use anyhow::{Context, Result};
use rocksdb::{BlockBasedOptions, Cache, Options, WriteBatch, DB};
use std::path::Path;

/// Persistent map of `vector_id -> bucket_id` stored in RocksDB.
///
/// This keeps the reverse lookup lightweight without holding everything in RAM.
pub struct BucketIndex {
    db: DB,
}

impl BucketIndex {
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.set_compression_type(rocksdb::DBCompressionType::Zstd);
        opts.optimize_level_style_compaction(64 * 1024 * 1024);
        let cache = Cache::new_lru_cache(64 * 1024 * 1024);
        let mut block_opts = BlockBasedOptions::default();
        block_opts.set_block_cache(&cache);
        opts.set_block_based_table_factory(&block_opts);
        opts.set_write_buffer_size(16 * 1024 * 1024);
        opts.set_max_write_buffer_number(2);
        let db = DB::open(&opts, path).context("open bucket index")?;
        Ok(Self { db })
    }

    /// Insert or overwrite a batch of `id -> bucket` mappings.
    pub fn put_batch(&self, bucket_id: u64, ids: &[u64]) -> Result<()> {
        if ids.is_empty() {
            return Ok(());
        }
        let mut batch = WriteBatch::default();
        for id in ids {
            batch.put(id.to_le_bytes(), bucket_id.to_le_bytes());
        }
        self.db.write(batch).context("write bucket index batch")
    }

    /// Delete a batch of ids (best-effort; missing keys are ignored).
    #[allow(dead_code)]
    pub fn delete_batch(&self, ids: &[u64]) -> Result<()> {
        if ids.is_empty() {
            return Ok(());
        }
        let mut batch = WriteBatch::default();
        for id in ids {
            batch.delete(id.to_le_bytes());
        }
        self.db.write(batch).context("delete bucket index batch")
    }

    /// Fetch the stored bucket ids for the given vector ids. Missing ids are skipped.
    pub fn get_many(&self, ids: &[u64]) -> Result<Vec<(u64, u64)>> {
        let mut out = Vec::with_capacity(ids.len());
        for id in ids {
            if let Some(raw) = self
                .db
                .get(id.to_le_bytes())
                .with_context(|| format!("read id {} from bucket index", id))?
            {
                if raw.len() == 8 {
                    let mut buf = [0u8; 8];
                    buf.copy_from_slice(&raw);
                    let bucket = u64::from_le_bytes(buf);
                    out.push((*id, bucket));
                }
            }
        }
        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bucket_index_round_trip() {
        let dir = tempfile::tempdir().expect("tempdir");
        let index = BucketIndex::open(dir.path()).expect("open index");
        let ids = vec![1u64, 2u64, 3u64];
        index.put_batch(7, &ids).expect("put");

        let got = index
            .get_many(&[1, 3, 4])
            .expect("get")
            .into_iter()
            .collect::<Vec<_>>();
        assert_eq!(got.len(), 2);
        assert!(got.contains(&(1, 7)));
        assert!(got.contains(&(3, 7)));
    }
}