weavatrix-search-vector 0.3.1

Persistent, mutable, bounded vector candidate search for Rust and Weavatrix
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

use super::codec::Encoder;
use super::format::{HEADER_LEN, VERSION, encode_value, put_u32, put_u64, to_u32, to_u64};
use super::io::checksum;
use super::validation::validate_quantized;
use crate::atomic_file::atomic_write;
use crate::error::SearchError;
use crate::hnsw::VectorIndex;
use crate::mutable::MutableVectorIndex;
use crate::quantized::QuantizedIndex;
use std::path::Path;

pub(super) fn save_complete(
    index: &MutableVectorIndex,
    quantized: Option<&QuantizedIndex>,
    target: &Path,
) -> Result<(), SearchError> {
    let snapshot = index.snapshot();
    validate_quantized(&snapshot, quantized)?;
    let base = snapshot_bytes(&snapshot.base)?;
    let sealed = snapshot
        .sealed
        .as_deref()
        .map(snapshot_bytes)
        .transpose()?
        .unwrap_or_default();
    let quantized = quantized
        .map(crate::quantized_io::encode)
        .transpose()?
        .unwrap_or_default();
    let mut encoder = Encoder::with_header()?;
    encoder.bytes(&base)?;
    encoder.bytes(&sealed)?;
    for (key, vector) in &snapshot.pending {
        encoder.u64(*key)?;
        encoder.f32s(vector)?;
    }
    for key in &snapshot.deleted {
        encoder.u64(*key)?;
    }
    for (key, metadata) in snapshot.metadata.iter() {
        encoder.u64(key)?;
        encoder.u32(to_u32(metadata.len())?)?;
        for (field, value) in metadata.iter() {
            encoder.length_prefixed(field.as_bytes())?;
            encode_value(&mut encoder, value)?;
        }
    }
    encoder.bytes(&quantized)?;
    let file_len = encoder.data.len();
    let checksum = checksum(&encoder.data[HEADER_LEN..]);
    put_u32(&mut encoder.data, 8, VERSION);
    put_u32(
        &mut encoder.data,
        12,
        u32::try_from(HEADER_LEN).expect("fixed bundle header length fits u32"),
    );
    put_u64(&mut encoder.data, 16, checksum);
    put_u64(&mut encoder.data, 24, to_u64(file_len)?);
    put_u64(&mut encoder.data, 32, to_u64(base.len())?);
    put_u64(&mut encoder.data, 40, to_u64(sealed.len())?);
    put_u64(&mut encoder.data, 48, to_u64(snapshot.pending.len())?);
    put_u64(&mut encoder.data, 56, to_u64(snapshot.deleted.len())?);
    put_u64(&mut encoder.data, 64, to_u64(snapshot.metadata.len())?);
    put_u64(&mut encoder.data, 72, to_u64(quantized.len())?);
    atomic_write(target, &encoder.data)
}

fn snapshot_bytes(index: &VectorIndex) -> Result<Vec<u8>, SearchError> {
    index.to_bytes()
}