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
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 super::{MutableState, VectorRecord};
use crate::config::{DistanceMetric, IndexConfig};
use crate::error::SearchError;
use crate::metadata::{Metadata, MetadataIndex};
use crate::vector::squared_norm;
use std::collections::BTreeSet;

pub(super) fn prepare_vector(
    dimensions: usize,
    metric: DistanceMetric,
    vector: &[f32],
) -> Result<Vec<f32>, SearchError> {
    if vector.len() != dimensions {
        return Err(SearchError::DimensionMismatch {
            expected: dimensions,
            actual: vector.len(),
            vector: Some(0),
        });
    }
    let squared = squared_norm(vector, Some(0))?;
    if metric == DistanceMetric::Cosine && squared == 0.0 {
        return Err(SearchError::ZeroVector { vector: Some(0) });
    }
    let mut prepared = Vec::new();
    prepared
        .try_reserve_exact(dimensions)
        .map_err(|_| SearchError::AllocationFailed)?;
    if metric == DistanceMetric::Cosine {
        let inverse = squared.sqrt().recip();
        prepared.extend(vector.iter().map(|value| value * inverse));
    } else {
        prepared.extend_from_slice(vector);
    }
    Ok(prepared)
}

pub(super) fn prepare_records(
    config: &IndexConfig,
    records: &[VectorRecord],
) -> Result<Vec<(u64, Vec<f32>, Metadata)>, SearchError> {
    let mut keys = BTreeSet::new();
    let mut prepared = Vec::new();
    prepared
        .try_reserve_exact(records.len())
        .map_err(|_| SearchError::AllocationFailed)?;
    for record in records {
        if !keys.insert(record.key) {
            return Err(SearchError::DuplicateKey(record.key));
        }
        prepared.push((
            record.key,
            prepare_vector(config.dimensions, config.metric, &record.vector)?,
            record.metadata.clone(),
        ));
    }
    Ok(prepared)
}

pub(super) fn replace_metadata(metadata: &mut MetadataIndex, key: u64, value: Metadata) {
    if value.is_empty() {
        metadata.remove(key);
    } else {
        metadata.insert(key, value);
    }
}

pub(super) fn current_vector(state: &MutableState, key: u64) -> Option<&[f32]> {
    if state.deleted.contains(&key) {
        return None;
    }
    state
        .pending
        .get(&key)
        .map(Vec::as_slice)
        .or_else(|| state.sealed.as_ref().and_then(|index| index.vector(key)))
        .or_else(|| state.base.vector(key))
}

pub(super) fn key_exists(state: &MutableState, key: u64) -> bool {
    state.pending.contains_key(&key)
        || (!state.deleted.contains(&key)
            && (state
                .sealed
                .as_ref()
                .is_some_and(|index| index.vector(key).is_some())
                || state.base.vector(key).is_some()))
}

pub(super) fn current_len(state: &MutableState) -> usize {
    let mut keys = BTreeSet::new();
    keys.extend(state.base.keys().filter(|key| !state.deleted.contains(key)));
    if let Some(sealed) = &state.sealed {
        keys.extend(sealed.keys().filter(|key| !state.deleted.contains(key)));
    }
    keys.extend(state.pending.keys().copied());
    keys.len()
}