bbqvec 0.0.10

Scalable, embeddable, vector storage for approximate K-nearest-neighbors (AKNN)
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
98
99
100
101
102
103
104
105
106
107

use crate::{vector::distance, Vector};
use anyhow::Result;
use half::{bf16, vec::HalfFloatVecExt};

pub trait Quantization: Default {
    type Lower: Clone;
    fn similarity(x: &Self::Lower, y: &Self::Lower) -> Result<f32>;
    fn compare(x: &Vector, y: &Self::Lower) -> Result<f32>;
    fn lower(vec: Vector) -> Result<Self::Lower>;
    fn vector_size(dimensions: usize) -> usize;
    fn marshal(v: &Self::Lower, array: &mut [u8]) -> Result<()>;
    fn unmarshal(array: &[u8]) -> Result<Self::Lower>;
    fn name() -> &'static str;
}

#[derive(Default)]
pub struct NoQuantization {}

impl Quantization for NoQuantization {
    type Lower = Vector;

    fn similarity(x: &Self::Lower, y: &Self::Lower) -> Result<f32> {
        Ok(distance(x, y))
    }

    fn compare(x: &Vector, y: &Self::Lower) -> Result<f32> {
        Ok(distance(x, y))
    }

    fn lower(vec: Vector) -> Result<Self::Lower> {
        Ok(vec)
    }

    fn name() -> &'static str {
        "none"
    }

    fn vector_size(dimensions: usize) -> usize {
        4 * dimensions
    }

    fn marshal(v: &Self::Lower, array: &mut [u8]) -> Result<()> {
        for (i, f) in v.iter().enumerate() {
            let bytes = f.to_le_bytes();
            let _ = &array[i * 4..i * 4 + 4].copy_from_slice(&bytes);
        }
        Ok(())
    }

    fn unmarshal(array: &[u8]) -> Result<Self::Lower> {
        let mut vec = Vec::new();
        for i in (0..array.len()).step_by(4) {
            let bytes = &array[i..i + 4];
            let f: f32 = f32::from_le_bytes(bytes.try_into().unwrap());
            vec.push(f);
        }
        Ok(vec)
    }
}

#[derive(Default)]
pub struct BF16Quantization {}

impl Quantization for BF16Quantization {
    type Lower = Vec<half::bf16>;

    fn similarity(x: &Self::Lower, y: &Self::Lower) -> Result<f32> {
        let fx = x.iter().map(|v| v.to_f32()).collect();
        let fy = y.iter().map(|v| v.to_f32()).collect();
        Ok(distance(&fx, &fy))
    }

    fn compare(x: &Vector, y: &Self::Lower) -> Result<f32> {
        let fy = y.iter().map(|v| v.to_f32()).collect();
        Ok(distance(x, &fy))
    }

    fn lower(vec: Vector) -> Result<Self::Lower> {
        Ok(Vec::from_f32_slice(vec.as_slice()))
    }

    fn name() -> &'static str {
        "bf16"
    }

    fn vector_size(dimensions: usize) -> usize {
        2 * dimensions
    }

    fn marshal(v: &Self::Lower, array: &mut [u8]) -> Result<()> {
        for (i, f) in v.iter().enumerate() {
            let bytes = f.to_le_bytes();
            let _ = &array[i * 2..i * 2 + 2].copy_from_slice(&bytes);
        }
        Ok(())
    }

    fn unmarshal(array: &[u8]) -> Result<Self::Lower> {
        let mut vec = Vec::new();
        for i in (0..array.len()).step_by(2) {
            let bytes = &array[i..i + 2];
            let f: bf16 = bf16::from_le_bytes(bytes.try_into().unwrap());
            vec.push(f);
        }
        Ok(vec)
    }
}