use crate::{SearchError, SearchHit, simd::DistanceKernel};
use std::cmp::Ordering;
use std::collections::BinaryHeap;
pub(crate) const ROUTING_BITS: usize = 14;
pub(crate) const ROUTING_PROBES: usize = 5;
#[derive(Debug)]
pub(crate) struct VectorStore {
dimensions: usize,
keys: Vec<u64>,
normalized: Vec<f32>,
routing_signs: Vec<u16>,
distance_kernel: DistanceKernel,
}
impl VectorStore {
pub(crate) fn build(dimensions: usize, vectors: &[(u64, &[f32])]) -> Result<Self, SearchError> {
if vectors.len() > u32::MAX as usize {
return Err(SearchError::CapacityOverflow);
}
let elements = dimensions
.checked_mul(vectors.len())
.ok_or(SearchError::CapacityOverflow)?;
let mut order = (0..vectors.len()).collect::<Vec<_>>();
order.sort_unstable_by_key(|index| vectors[*index].0);
for pair in order.windows(2) {
if vectors[pair[0]].0 == vectors[pair[1]].0 {
return Err(SearchError::DuplicateKey(vectors[pair[0]].0));
}
}
let mut keys = Vec::new();
keys.try_reserve_exact(vectors.len())
.map_err(|_| SearchError::AllocationFailed)?;
let mut normalized = Vec::new();
normalized
.try_reserve_exact(elements)
.map_err(|_| SearchError::AllocationFailed)?;
let mut routing_signs = Vec::new();
routing_signs
.try_reserve_exact(dimensions)
.map_err(|_| SearchError::AllocationFailed)?;
routing_signs.extend((0..dimensions).map(|dimension| {
let mixed =
splitmix64(u64::try_from(dimension).unwrap_or(u64::MAX) ^ 0xa076_1d64_78bd_642f);
u16::try_from(mixed & u64::from(u16::MAX)).expect("masked routing signs fit u16")
}));
for (sorted_index, source_index) in order.into_iter().enumerate() {
let (key, vector) = vectors[source_index];
if vector.len() != dimensions {
return Err(SearchError::DimensionMismatch {
expected: dimensions,
actual: vector.len(),
vector: Some(source_index),
});
}
let inverse_norm = inverse_norm(vector, Some(source_index))?;
keys.push(key);
normalized.extend(vector.iter().map(|value| value * inverse_norm));
debug_assert_eq!(keys.len(), sorted_index + 1);
}
Ok(Self {
dimensions,
keys,
normalized,
routing_signs,
distance_kernel: DistanceKernel::detect(),
})
}
pub(crate) fn len(&self) -> usize {
self.keys.len()
}
pub(crate) fn is_empty(&self) -> bool {
self.keys.is_empty()
}
pub(crate) fn estimated_bytes(&self) -> usize {
self.keys
.capacity()
.saturating_mul(std::mem::size_of::<u64>())
.saturating_add(
self.normalized
.capacity()
.saturating_mul(std::mem::size_of::<f32>()),
)
.saturating_add(
self.routing_signs
.capacity()
.saturating_mul(std::mem::size_of::<u16>()),
)
}
pub(crate) fn key(&self, index: usize) -> u64 {
self.keys[index]
}
fn vector(&self, index: usize) -> &[f32] {
let start = index * self.dimensions;
&self.normalized[start..start + self.dimensions]
}
pub(crate) fn distance_indices(&self, left: usize, right: usize) -> f32 {
self.distance_kernel
.cosine_distance(self.vector(left), self.vector(right), 1.0)
}
pub(crate) fn query_inverse_norm(&self, query: &[f32]) -> Result<f32, SearchError> {
if query.len() != self.dimensions {
return Err(SearchError::DimensionMismatch {
expected: self.dimensions,
actual: query.len(),
vector: None,
});
}
inverse_norm(query, None)
}
pub(crate) fn distance_query(&self, index: usize, query: &[f32], inverse_norm: f32) -> f32 {
self.distance_kernel
.cosine_distance(self.vector(index), query, inverse_norm)
}
pub(crate) fn routing_code(&self, vector: &[f32]) -> u16 {
let sums = self.routing_sums(vector);
sums.iter().enumerate().fold(0_u16, |code, (plane, sum)| {
if *sum >= 0.0 {
code | (1_u16 << plane)
} else {
code
}
})
}
pub(crate) fn routing_probes(&self, vector: &[f32]) -> [u16; ROUTING_PROBES] {
let sums = self.routing_sums(vector);
let code = sums.iter().enumerate().fold(0_u16, |code, (plane, sum)| {
if *sum >= 0.0 {
code | (1_u16 << plane)
} else {
code
}
});
let mut planes = std::array::from_fn::<_, ROUTING_BITS, _>(|plane| plane);
planes.sort_unstable_by(|left, right| {
sums[*left]
.abs()
.total_cmp(&sums[*right].abs())
.then_with(|| left.cmp(right))
});
let mut probes = [code; ROUTING_PROBES];
for (probe, plane) in probes[1..].iter_mut().zip(planes) {
*probe = code ^ (1_u16 << plane);
}
probes
}
fn routing_sums(&self, vector: &[f32]) -> [f32; ROUTING_BITS] {
debug_assert_eq!(vector.len(), self.routing_signs.len());
let mut sums = [0.0_f32; ROUTING_BITS];
for (value, signs) in vector.iter().copied().zip(&self.routing_signs) {
for (plane, sum) in sums.iter_mut().enumerate() {
if signs & (1_u16 << plane) == 0 {
*sum += value;
} else {
*sum -= value;
}
}
}
sums
}
pub(crate) fn stored_routing_code(&self, index: usize) -> u16 {
self.routing_code(self.vector(index))
}
pub(crate) fn exact(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
let inverse_norm = self.query_inverse_norm(query)?;
let limit = count.min(self.len());
if limit == 0 {
return Ok(Vec::new());
}
let mut best = BinaryHeap::with_capacity(limit);
for index in 0..self.len() {
let candidate = Candidate::new(self.distance_query(index, query, inverse_norm), index);
if best.len() < limit {
best.push(candidate);
} else if best
.peek()
.is_some_and(|worst| candidate.cmp(worst).is_lt())
{
best.pop();
best.push(candidate);
}
}
let mut candidates = best.into_vec();
candidates.sort_unstable();
Ok(candidates
.into_iter()
.map(|candidate| SearchHit {
key: self.key(candidate.index()),
distance: candidate.distance,
})
.collect())
}
}
#[allow(clippy::cast_possible_truncation)]
fn inverse_norm(vector: &[f32], position: Option<usize>) -> Result<f32, SearchError> {
let mut squared = 0.0_f64;
for (dimension, value) in vector.iter().copied().enumerate() {
if !value.is_finite() {
return Err(SearchError::NonFiniteValue {
vector: position,
dimension,
});
}
squared += f64::from(value) * f64::from(value);
}
if squared == 0.0 {
return Err(SearchError::ZeroVector { vector: position });
}
Ok((1.0 / squared.sqrt()) as f32)
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct Candidate {
pub(crate) distance: f32,
node: u32,
}
impl Candidate {
pub(crate) fn new(distance: f32, index: usize) -> Self {
Self {
distance,
node: u32::try_from(index).expect("vector count is bounded by u32::MAX"),
}
}
pub(crate) fn index(self) -> usize {
self.node as usize
}
}
impl PartialEq for Candidate {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl Eq for Candidate {}
impl PartialOrd for Candidate {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Candidate {
fn cmp(&self, other: &Self) -> Ordering {
self.distance
.total_cmp(&other.distance)
.then_with(|| self.node.cmp(&other.node))
}
}
pub(crate) fn splitmix64(mut value: u64) -> u64 {
value = value.wrapping_add(0x9e37_79b9_7f4a_7c15);
value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
value ^ (value >> 31)
}