use super::super::builder::BatchBuilder;
use super::super::canonical_splice::{canonical_splice, CanonicalSpliceStats};
use super::super::cid::Cid;
use super::super::config::Config;
use super::super::error::{Error, Mutation as TreeMutation};
use super::super::store::Store;
use super::super::Prolly;
use super::builder::{build_hierarchy, build_hierarchy_parallel, IndexedRecord};
use super::cache::{ContentCache, DEFAULT_PROXIMITY_CACHE_NODES};
use super::distance::{prepare_vector, query_score, score};
use super::mutation::{mutate_hierarchy, LogicalEdit};
use super::search::{
adaptive_should_stop, insert_top_k, AdaptiveContext, FrontierEntry, PreparedFilter,
SearchCandidate,
};
use super::storage::quantized::ScalarQuantized;
use super::storage::vector::ExternalVector;
use super::storage::{Descriptor, PhysicalNodeKind, ProximityNode, StoredRecord, VectorRef};
use super::vector::promotion_level;
use super::{
BuildParallelism, DistanceMetric, ExactProximityRecord, Neighbor, ProximityBuildStats,
ProximityConfig, ProximityMutation, ProximityMutationStats, ProximityRecord,
ProximitySearchStats, ProximityTree, ProximityVerification, SearchBackend, SearchCompletion,
SearchPolicy, SearchRequest, SearchResult,
};
use std::collections::{BTreeMap, BinaryHeap, HashSet};
use std::sync::{Arc, Mutex};
pub struct ProximityMap<S: Store> {
store: S,
directory: Prolly<S>,
tree: ProximityTree,
node_cache: Mutex<ContentCache<ProximityNode>>,
}
impl<S> ProximityMap<S>
where
S: Store + Clone + Send + Sync,
S::Error: Send + Sync,
{
pub fn build(
store: S,
config: ProximityConfig,
records: impl IntoIterator<Item = ProximityRecord>,
) -> Result<Self, Error> {
Self::build_with_parallelism(store, config, records, BuildParallelism::default())
.map(|(map, _)| map)
}
pub fn build_with_parallelism(
store: S,
config: ProximityConfig,
records: impl IntoIterator<Item = ProximityRecord>,
parallelism: BuildParallelism,
) -> Result<(Self, ProximityBuildStats), Error> {
config.validate()?;
let mut records: Vec<_> = records.into_iter().collect();
records.sort_by(|left, right| left.key.cmp(&right.key));
for pair in records.windows(2) {
if pair[0].key == pair[1].key {
return Err(Error::DuplicateProximityKey {
key: pair[0].key.clone(),
});
}
}
let directory_config = Config::default();
let mut directory_builder = BatchBuilder::new(store.clone(), directory_config.clone());
let mut indexed = Vec::with_capacity(records.len());
for record in records {
let stored = StoredRecord::new(
&record.vector,
record.value,
config.metric,
config.dimensions,
)?;
indexed.push(IndexedRecord {
key: record.key.clone(),
vector: stored.vector.clone(),
});
directory_builder.add(record.key, stored.encode());
}
let directory_tree = directory_builder.build()?;
let hierarchy = build_hierarchy_parallel(&indexed, &config, parallelism.threads())?;
let objects_written = put_missing_nodes(&store, &hierarchy.nodes)?;
let descriptor = Descriptor {
config: config.clone(),
count: indexed.len() as u64,
directory: directory_tree.clone(),
proximity_root: hierarchy.root.clone(),
};
let descriptor_bytes = descriptor.encode();
let descriptor_cid = Cid::from_bytes(&descriptor_bytes);
store
.put(descriptor_cid.as_bytes(), &descriptor_bytes)
.map_err(|error| Error::Store(Box::new(error)))?;
let stats = ProximityBuildStats {
distance_evaluations: hierarchy.distance_evaluations,
proximity_objects: hierarchy.nodes.len(),
proximity_objects_written: objects_written,
};
Ok((
Self {
directory: Prolly::new(store.clone(), directory_config),
store,
tree: ProximityTree {
directory: directory_tree,
proximity_root: hierarchy.root,
descriptor: descriptor_cid,
count: indexed.len() as u64,
config,
},
node_cache: Mutex::new(ContentCache::new(DEFAULT_PROXIMITY_CACHE_NODES)),
},
stats,
))
}
pub fn load(store: S, descriptor_cid: Cid) -> Result<Self, Error> {
let descriptor_bytes = load_content(&store, &descriptor_cid)?;
let descriptor = Descriptor::decode(&descriptor_bytes)?;
let root_bytes = load_content(&store, &descriptor.proximity_root)?;
if root_bytes.len() > descriptor.config.overflow.max_page_bytes as usize {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "root exceeds descriptor max_node_bytes".to_owned(),
});
}
let root = ProximityNode::decode(&root_bytes, descriptor.config.dimensions)?;
if root.subtree_count != descriptor.count {
return Err(Error::InvalidProximityObject {
kind: "descriptor",
reason: "record count disagrees with proximity root".to_owned(),
});
}
let directory_config = descriptor.directory.config.clone();
Ok(Self {
directory: Prolly::new(store.clone(), directory_config),
store,
tree: ProximityTree {
directory: descriptor.directory,
proximity_root: descriptor.proximity_root,
descriptor: descriptor_cid,
count: descriptor.count,
config: descriptor.config,
},
node_cache: Mutex::new(ContentCache::new(DEFAULT_PROXIMITY_CACHE_NODES)),
})
}
pub fn tree(&self) -> &ProximityTree {
&self.tree
}
pub fn clear_content_cache(&self) -> Result<(), Error> {
self.node_cache
.lock()
.map_err(|_| Error::InvalidProximityObject {
kind: "cache",
reason: "node cache lock poisoned".to_owned(),
})?
.clear();
Ok(())
}
pub fn get(&self, key: &[u8]) -> Result<Option<ExactProximityRecord>, Error> {
self.directory
.get(&self.tree.directory, key)?
.map(|bytes| {
let record = StoredRecord::decode(&bytes, self.tree.config.dimensions)?;
Ok((record.vector, record.value))
})
.transpose()
}
pub fn contains_key(&self, key: &[u8]) -> Result<bool, Error> {
Ok(self.get(key)?.is_some())
}
pub fn rebuild_batch(
&self,
mutations: impl IntoIterator<Item = ProximityMutation>,
) -> Result<Self, Error> {
let mutations = validate_mutations(mutations)?;
let mut records = self.collect_records()?;
apply_mutations(&mut records, &mutations, &self.tree.config)?;
Self::build(
self.store.clone(),
self.tree.config.clone(),
records.into_values(),
)
}
pub fn mutate_batch(
&self,
mutations: impl IntoIterator<Item = ProximityMutation>,
) -> Result<(Self, ProximityMutationStats), Error> {
let mutations = validate_mutations(mutations)?;
if mutations.is_empty() {
return Ok((
Self::load(self.store.clone(), self.tree.descriptor.clone())?,
Default::default(),
));
}
let keys: Vec<_> = mutations
.iter()
.map(|mutation| mutation.key.clone())
.collect();
let old_values = self.directory.get_many(&self.tree.directory, &keys)?;
let mut logical_edits = Vec::new();
let mut directory_mutations = Vec::with_capacity(mutations.len());
let mut count = self.tree.count;
for (mutation, old_bytes) in mutations.iter().zip(old_values) {
let old = old_bytes
.as_deref()
.map(|bytes| StoredRecord::decode(bytes, self.tree.config.dimensions))
.transpose()?;
let new = mutation
.value
.as_ref()
.map(|(vector, value)| {
StoredRecord::new(
vector,
value.clone(),
self.tree.config.metric,
self.tree.config.dimensions,
)
})
.transpose()?;
match (&old, &new) {
(None, Some(_)) => {
count = count
.checked_add(1)
.ok_or_else(|| Error::InvalidProximityObject {
kind: "mutation",
reason: "record count overflow".to_owned(),
})?
}
(Some(_), None) => count -= 1,
_ => {}
}
let old_vector = old.as_ref().map(|record| record.vector.clone());
let new_vector = new.as_ref().map(|record| record.vector.clone());
if old_vector != new_vector {
logical_edits.push(LogicalEdit {
key: mutation.key.clone(),
old: old_vector,
new: new_vector,
level: promotion_level(
&mutation.key,
self.tree.config.hierarchy.log_chunk_size,
self.tree.config.hierarchy.level_hash_seed,
),
});
}
directory_mutations.push(match new {
Some(record) => TreeMutation::Upsert {
key: mutation.key.clone(),
val: record.encode(),
},
None => TreeMutation::Delete {
key: mutation.key.clone(),
},
});
}
let (directory_tree, directory_stats) =
canonical_splice(&self.directory, &self.tree.directory, directory_mutations)?;
if logical_edits.is_empty() {
let descriptor = Descriptor {
config: self.tree.config.clone(),
count,
directory: directory_tree.clone(),
proximity_root: self.tree.proximity_root.clone(),
};
let bytes = descriptor.encode();
let descriptor_cid = Cid::from_bytes(&bytes);
self.store
.put(descriptor_cid.as_bytes(), &bytes)
.map_err(|error| Error::Store(Box::new(error)))?;
let map = Self {
directory: Prolly::new(self.store.clone(), directory_tree.config.clone()),
store: self.store.clone(),
tree: ProximityTree {
directory: directory_tree,
proximity_root: self.tree.proximity_root.clone(),
descriptor: descriptor_cid,
count,
config: self.tree.config.clone(),
},
node_cache: Mutex::new(ContentCache::new(DEFAULT_PROXIMITY_CACHE_NODES)),
};
return Ok((
map,
ProximityMutationStats {
nodes_reused: 1,
directory_entries_scanned: directory_stats.entries_scanned,
directory_nodes_read: directory_stats.nodes_read,
directory_nodes_rebuilt: directory_stats.nodes_rebuilt,
directory_nodes_written: directory_stats.nodes_written,
directory_nodes_reused: directory_stats.nodes_reused,
directory_levels_rebuilt: directory_stats.levels_rebuilt,
directory_right_edge_rebuilt: directory_stats.right_edge_rebuilt,
..Default::default()
},
));
}
let (old_root, _) = self.load_node(&self.tree.proximity_root)?;
let max_edit_level = logical_edits
.iter()
.map(|edit| edit.level)
.max()
.unwrap_or(0);
let (proximity_root, nodes, mut stats) =
if old_root.entries.is_empty() || max_edit_level >= old_root.level {
let records = self.collect_records_from(&directory_tree)?;
let indexed: Vec<_> = records
.values()
.map(|record| IndexedRecord {
key: record.key.clone(),
vector: record.vector.clone(),
})
.collect();
let built = build_hierarchy(&indexed, &self.tree.config)?;
let stats = ProximityMutationStats {
records_rebuilt: indexed.len(),
distance_evaluations: built.distance_evaluations,
full_proximity_rebuild: true,
..Default::default()
};
(built.root, built.nodes, stats)
} else {
let local = mutate_hierarchy(
&self.store,
&self.tree.proximity_root,
&self.tree.config,
&logical_edits,
)?;
(local.root, local.nodes, local.stats)
};
let pending_count = nodes.len();
let nodes_written = put_missing_nodes(&self.store, &nodes)?;
stats.nodes_written = nodes_written;
stats.nodes_reused += pending_count.saturating_sub(nodes_written);
apply_directory_stats(&mut stats, directory_stats);
let descriptor = Descriptor {
config: self.tree.config.clone(),
count,
directory: directory_tree.clone(),
proximity_root: proximity_root.clone(),
};
let descriptor_bytes = descriptor.encode();
let descriptor_cid = Cid::from_bytes(&descriptor_bytes);
self.store
.put(descriptor_cid.as_bytes(), &descriptor_bytes)
.map_err(|error| Error::Store(Box::new(error)))?;
Ok((
Self {
directory: Prolly::new(self.store.clone(), directory_tree.config.clone()),
store: self.store.clone(),
tree: ProximityTree {
directory: directory_tree,
proximity_root,
descriptor: descriptor_cid,
count,
config: self.tree.config.clone(),
},
node_cache: Mutex::new(ContentCache::new(DEFAULT_PROXIMITY_CACHE_NODES)),
},
stats,
))
}
pub fn search(&self, request: SearchRequest<'_>) -> Result<SearchResult, Error> {
self.search_with_trace(request, None)
}
pub(super) fn search_with_trace(
&self,
request: SearchRequest<'_>,
mut trace: Option<&mut Vec<super::proof::ProximitySearchEvent>>,
) -> Result<SearchResult, Error> {
request.validate()?;
if matches!(
request.backend,
SearchBackend::ProductQuantized | SearchBackend::Hnsw
) {
return Err(Error::InvalidProximitySearch {
reason: "requested backend requires a validated accelerator sidecar".to_owned(),
});
}
let filter = PreparedFilter::new(request.filter.clone(), &self.tree.directory)?;
let query = prepare_vector(
self.tree.config.metric,
request.query,
self.tree.config.dimensions,
)?;
let use_scalar_quantization =
super::accelerator::sq8::enabled(&self.tree.config, request.policy);
let mut stats = ProximitySearchStats::default();
let mut frontier = BinaryHeap::new();
frontier.push(FrontierEntry {
bound: 0.0,
score: 0.0,
key: Vec::new(),
cid: self.tree.proximity_root.clone(),
expected_level: None,
});
if let Some(trace) = trace.as_deref_mut() {
trace.push(super::proof::ProximitySearchEvent::FrontierPushed {
cid: self.tree.proximity_root.clone(),
bound_bits: 0.0f64.to_bits(),
});
}
let mut candidates = Vec::<SearchCandidate>::new();
let mut score_cache = BTreeMap::<Vec<u8>, f64>::new();
let mut visited = HashSet::new();
let mut levels = HashSet::new();
let mut last_fanout = 0usize;
let mut completion = SearchCompletion::Exact;
while let Some(next) = frontier.peek() {
if !use_scalar_quantization
&& self.tree.config.metric == DistanceMetric::L2Squared
&& candidates.len() == request.k
&& next.bound > candidates.last().expect("full top-k").score
{
break;
}
if let SearchPolicy::Adaptive(quality) = request.policy {
if candidates.last().is_some_and(|worst| {
let overlapping = frontier
.iter()
.filter(|entry| entry.bound <= worst.score)
.count();
adaptive_should_stop(
quality,
AdaptiveContext {
results: candidates.len(),
k: request.k,
frontier_bound: next.bound,
worst_score: worst.score,
overlapping_clusters: overlapping,
logical_level: next.expected_level.unwrap_or(u8::MAX),
last_fanout,
cluster_count: frontier.len(),
},
)
}) {
completion = SearchCompletion::ApproximatePolicySatisfied;
break;
}
}
if request
.budget
.max_nodes
.is_some_and(|maximum| stats.nodes_read >= maximum)
{
completion = SearchCompletion::BudgetExhausted;
break;
}
let next = frontier.pop().expect("peeked frontier");
if let Some(trace) = trace.as_deref_mut() {
trace.push(super::proof::ProximitySearchEvent::FrontierPopped {
cid: next.cid.clone(),
bound_bits: next.bound.to_bits(),
});
}
if !visited.insert(next.cid.clone()) {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "cycle or repeated child ownership".to_owned(),
});
}
let (node, mut bytes) = self.load_node(&next.cid)?;
if let Some(trace) = trace.as_deref_mut() {
trace.push(super::proof::ProximitySearchEvent::VisitedObject(
next.cid.clone(),
));
}
let quantizer = if use_scalar_quantization && node.kind.has_children(node.level) {
let (quantizer, quantizer_bytes) = self.load_scalar_quantizer(&node)?;
bytes = bytes.saturating_add(quantizer_bytes);
Some(quantizer)
} else {
None
};
if next
.expected_level
.is_some_and(|expected| node.level != expected)
{
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "child has an unexpected logical level".to_owned(),
});
}
stats.bytes_read = stats.bytes_read.saturating_add(bytes);
if request
.budget
.max_committed_bytes
.is_some_and(|maximum| stats.committed_bytes.saturating_add(bytes) > maximum)
{
completion = SearchCompletion::BudgetExhausted;
break;
}
stats.nodes_read += 1;
stats.committed_bytes += bytes;
last_fanout = node.entries.len();
levels.insert(node.level);
stats.levels_visited = levels.len();
for (entry_index, entry) in node.entries.iter().enumerate() {
if node.kind.has_children(node.level) {
if !filter.intersects(&entry.min_key, &entry.max_key) {
continue;
}
let Some(child) = &entry.child else {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "internal entry has no child".to_owned(),
});
};
let representative_score = if let Some(quantizer) = &quantizer {
if distance_budget_exhausted(&request, &stats) {
completion = SearchCompletion::BudgetExhausted;
break;
}
stats.quantized_distance_evaluations += 1;
quantizer.approximate_score(self.tree.config.metric, &query, entry_index)?
} else {
match score_cache.get(&entry.key) {
Some(score) => *score,
None => {
if distance_budget_exhausted(&request, &stats) {
completion = SearchCompletion::BudgetExhausted;
break;
}
stats.distance_evaluations += 1;
let value = query_score(
request.kernel,
self.tree.config.metric,
&query,
entry.vector.inline()?,
);
score_cache.insert(entry.key.clone(), value);
value
}
}
};
let bound = if quantizer.is_none()
&& self.tree.config.metric == DistanceMetric::L2Squared
{
super::distance::canonical::l2_lower_bound_down(
representative_score,
entry.covering_radius,
)
} else {
representative_score
};
if request
.budget
.max_frontier_entries
.is_some_and(|maximum| frontier.len() >= maximum)
{
completion = SearchCompletion::BudgetExhausted;
break;
}
frontier.push(FrontierEntry {
bound,
score: representative_score,
key: entry.key.clone(),
cid: child.clone(),
expected_level: Some(if node.kind == PhysicalNodeKind::OverflowDirectory {
node.level
} else {
node.level - 1
}),
});
if let Some(trace) = trace.as_deref_mut() {
trace.push(super::proof::ProximitySearchEvent::FrontierPushed {
cid: child.clone(),
bound_bits: bound.to_bits(),
});
}
stats.frontier_peak = stats.frontier_peak.max(frontier.len());
} else if filter.contains(&entry.key) {
let leaf_score = match score_cache.get(&entry.key) {
Some(score) => *score,
None => {
if distance_budget_exhausted(&request, &stats) {
completion = SearchCompletion::BudgetExhausted;
break;
}
stats.distance_evaluations += 1;
let value = query_score(
request.kernel,
self.tree.config.metric,
&query,
entry.vector.inline()?,
);
score_cache.insert(entry.key.clone(), value);
value
}
};
if let Some(trace) = trace.as_deref_mut() {
trace.push(super::proof::ProximitySearchEvent::CandidateScored {
key: entry.key.clone(),
distance_bits: leaf_score.to_bits(),
});
}
insert_top_k(
&mut candidates,
SearchCandidate {
key: entry.key.clone(),
vector: entry.vector.inline()?.to_vec(),
score: leaf_score,
},
request.k,
);
}
}
if completion == SearchCompletion::BudgetExhausted {
break;
}
}
let keys: Vec<_> = candidates
.iter()
.map(|candidate| candidate.key.clone())
.collect();
let values = self.directory.get_many(&self.tree.directory, &keys)?;
if use_scalar_quantization {
stats.reranked_candidates = candidates.len();
}
let mut neighbors = Vec::with_capacity(candidates.len());
for (candidate, stored) in candidates.into_iter().zip(values) {
let bytes = stored.ok_or_else(|| Error::InvalidProximityObject {
kind: "node",
reason: "leaf key is absent from exact directory".to_owned(),
})?;
let record = StoredRecord::decode(&bytes, self.tree.config.dimensions)?;
if record.vector != candidate.vector {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "leaf vector disagrees with exact directory".to_owned(),
});
}
neighbors.push(Neighbor {
key: candidate.key,
value: record.value,
distance: candidate.score,
});
}
if let Some(trace) = trace {
trace.push(super::proof::ProximitySearchEvent::Completed(completion));
}
Ok(SearchResult {
neighbors,
stats,
completion,
})
}
pub fn verify(&self) -> Result<ProximityVerification, Error> {
let records = self.collect_records()?;
let root_bytes = load_content(&self.store, &self.tree.proximity_root)?;
let root = ProximityNode::decode(&root_bytes, self.tree.config.dimensions)?;
let mut state = VerificationState {
records: &records,
seen_nodes: HashSet::new(),
seen_external_vectors: HashSet::new(),
seen_scalar_quantizers: HashSet::new(),
seen_leaf_keys: HashSet::new(),
summary: ProximityVerification {
record_count: self.tree.count,
maximum_level: root.level,
..Default::default()
},
};
let verified = self.verify_node(
&self.tree.proximity_root,
Some(root.level),
None,
&mut state,
)?;
if verified.count != self.tree.count || records.len() as u64 != self.tree.count {
return Err(Error::InvalidProximityObject {
kind: "descriptor",
reason: "logical counts disagree".to_owned(),
});
}
if state.seen_leaf_keys.len() != records.len()
|| records
.keys()
.any(|key| !state.seen_leaf_keys.contains(key))
{
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "leaf identities do not match the exact directory".to_owned(),
});
}
Ok(state.summary)
}
fn load_node(&self, cid: &Cid) -> Result<(Arc<ProximityNode>, usize), Error> {
if let Some((node, bytes)) = self
.node_cache
.lock()
.map_err(|_| Error::InvalidProximityObject {
kind: "cache",
reason: "node cache lock poisoned".to_owned(),
})?
.get(cid)
{
return Ok((node, bytes));
}
let bytes = load_content(&self.store, cid)?;
if bytes.len() > self.tree.config.overflow.max_page_bytes as usize {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "node exceeds descriptor max_node_bytes".to_owned(),
});
}
let len = bytes.len();
let mut node = ProximityNode::decode(&bytes, self.tree.config.dimensions)?;
let vector_bytes = self.resolve_external_vectors(&mut node)?;
let node = Arc::new(node);
self.node_cache
.lock()
.map_err(|_| Error::InvalidProximityObject {
kind: "cache",
reason: "node cache lock poisoned".to_owned(),
})?
.insert(cid.clone(), node.clone(), len + vector_bytes);
Ok((node, len + vector_bytes))
}
fn resolve_external_vectors(&self, node: &mut ProximityNode) -> Result<usize, Error> {
let mut bytes_read = 0usize;
for entry in &mut node.entries {
let VectorRef::External(cid) = &entry.vector else {
continue;
};
let bytes = load_content(&self.store, cid)?;
let external = ExternalVector::decode(&bytes)?;
if external.vector.len() != self.tree.config.dimensions as usize {
return Err(Error::InvalidProximityObject {
kind: "vector",
reason: "external vector dimension mismatch".to_owned(),
});
}
bytes_read += bytes.len();
entry.vector = VectorRef::Inline(external.vector);
}
Ok(bytes_read)
}
fn load_scalar_quantizer(
&self,
node: &ProximityNode,
) -> Result<(ScalarQuantized, usize), Error> {
let config = self
.tree
.config
.scalar_quantization
.as_ref()
.ok_or_else(|| Error::InvalidProximityObject {
kind: "quantizer",
reason: "quantized search requires descriptor configuration".to_owned(),
})?;
let cid = node
.quantizer
.as_ref()
.ok_or_else(|| Error::InvalidProximityObject {
kind: "quantizer",
reason: "configured node has no scalar quantizer".to_owned(),
})?;
let bytes = load_content(&self.store, cid)?;
let quantizer = ScalarQuantized::decode(&bytes)?;
if quantizer.dimensions != self.tree.config.dimensions
|| quantizer.group_size != config.group_size
{
return Err(Error::InvalidProximityObject {
kind: "quantizer",
reason: "quantizer configuration disagrees with descriptor".to_owned(),
});
}
if quantizer.entry_count != node.entries.len() as u64 {
return Err(Error::InvalidProximityObject {
kind: "quantizer",
reason: "quantizer entry count disagrees with node".to_owned(),
});
}
Ok((quantizer, bytes.len()))
}
pub(crate) fn collect_records(&self) -> Result<BTreeMap<Vec<u8>, ProximityRecord>, Error> {
self.collect_records_from(&self.tree.directory)
}
pub(crate) fn store_clone(&self) -> S {
self.store.clone()
}
pub(super) fn directory_manager(&self) -> &Prolly<S> {
&self.directory
}
pub(super) fn load_descriptor_bytes(&self) -> Result<Vec<u8>, Error> {
load_content(&self.store, &self.tree.descriptor)
}
fn collect_records_from(
&self,
directory: &super::super::tree::Tree,
) -> Result<BTreeMap<Vec<u8>, ProximityRecord>, Error> {
let mut records = BTreeMap::new();
for entry in self.directory.range(directory, &[], None)? {
let (key, bytes) = entry?;
let stored = StoredRecord::decode(&bytes, self.tree.config.dimensions)?;
records.insert(
key.clone(),
ProximityRecord {
key,
vector: stored.vector,
value: stored.value,
},
);
}
Ok(records)
}
fn verify_node(
&self,
cid: &Cid,
expected_level: Option<u8>,
parent: Option<(
&super::storage::ProximityEntry,
&[super::storage::ProximityEntry],
)>,
state: &mut VerificationState<'_>,
) -> Result<VerifiedSubtree, Error> {
if !state.seen_nodes.insert(cid.clone()) {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "cycle or repeated child ownership".to_owned(),
});
}
let bytes = load_content(&self.store, cid)?;
if bytes.len() > self.tree.config.overflow.max_page_bytes as usize {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "node exceeds descriptor max_node_bytes".to_owned(),
});
}
let mut node = ProximityNode::decode(&bytes, self.tree.config.dimensions)?;
for entry in &node.entries {
if let VectorRef::External(vector) = &entry.vector {
if state.seen_external_vectors.insert(vector.clone()) {
state.summary.external_vector_count += 1;
}
}
}
self.resolve_external_vectors(&mut node)?;
match (&self.tree.config.scalar_quantization, &node.quantizer) {
(None, None) => {}
(Some(config), Some(cid)) => {
let quantizer_bytes = load_content(&self.store, cid)?;
let quantizer = ScalarQuantized::decode(&quantizer_bytes)?;
if quantizer.dimensions != self.tree.config.dimensions
|| quantizer.group_size != config.group_size
{
return Err(Error::InvalidProximityObject {
kind: "quantizer",
reason: "quantizer configuration disagrees with descriptor".to_owned(),
});
}
let vectors = node
.entries
.iter()
.map(|entry| entry.vector.inline())
.collect::<Result<Vec<_>, _>>()?;
quantizer.verify(&vectors)?;
state.summary.quantized_node_count += 1;
if state.seen_scalar_quantizers.insert(cid.clone()) {
state.summary.scalar_quantizer_count += 1;
}
}
_ => {
return Err(Error::InvalidProximityObject {
kind: "quantizer",
reason: "node quantizer presence disagrees with descriptor".to_owned(),
})
}
}
if expected_level != Some(node.level) {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "unexpected logical level".to_owned(),
});
}
state.summary.proximity_node_count += 1;
match node.kind {
PhysicalNodeKind::OverflowPage => state.summary.overflow_page_count += 1,
PhysicalNodeKind::OverflowDirectory => state.summary.overflow_directory_count += 1,
PhysicalNodeKind::Leaf | PhysicalNodeKind::Route => {}
}
state.summary.maximum_node_bytes = state.summary.maximum_node_bytes.max(bytes.len());
if node.kind != PhysicalNodeKind::OverflowDirectory {
if let Some((selected, candidates)) = parent {
for entry in &node.entries {
if entry.key == selected.key {
continue;
}
let selected_distance = score(
self.tree.config.metric,
entry.vector.inline()?,
selected.vector.inline()?,
);
for candidate in candidates {
state.summary.distance_checks += 1;
let candidate_distance = score(
self.tree.config.metric,
entry.vector.inline()?,
candidate.vector.inline()?,
);
let candidate_is_better = candidate_distance
.total_cmp(&selected_distance)
.then_with(|| candidate.key.cmp(&selected.key))
.is_lt();
if candidate_is_better {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "nearest-representative invariant violated".to_owned(),
});
}
}
}
}
}
if node.kind != PhysicalNodeKind::OverflowDirectory {
for entry in &node.entries {
if super::vector::promotion_level(
&entry.key,
self.tree.config.hierarchy.log_chunk_size,
self.tree.config.hierarchy.level_hash_seed,
) < node.level
{
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "entry appears above its deterministic promotion level".to_owned(),
});
}
}
}
let verified = if node.kind.is_logical_leaf(node.level) {
let mut points = Vec::with_capacity(node.entries.len());
for entry in &node.entries {
if !state.seen_leaf_keys.insert(entry.key.clone()) {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "duplicate leaf identity".to_owned(),
});
}
let record =
state
.records
.get(&entry.key)
.ok_or_else(|| Error::InvalidProximityObject {
kind: "node",
reason: "leaf key is absent from exact directory".to_owned(),
})?;
if record.vector.as_slice() != entry.vector.inline()? {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "leaf vector disagrees with exact directory".to_owned(),
});
}
points.push((entry.key.clone(), entry.vector.inline()?.to_vec()));
}
VerifiedSubtree::from_points(node.entries.len() as u64, points)
} else {
let mut count = 0u64;
let mut points = Vec::new();
let mut minimum: Option<Vec<u8>> = None;
let mut maximum: Option<Vec<u8>> = None;
for entry in &node.entries {
let child = entry
.child
.as_ref()
.ok_or_else(|| Error::InvalidProximityObject {
kind: "node",
reason: "internal entry has no child".to_owned(),
})?;
let child_verified = self.verify_node(
child,
Some(if node.kind == PhysicalNodeKind::OverflowDirectory {
node.level
} else {
node.level - 1
}),
if node.kind == PhysicalNodeKind::OverflowDirectory {
parent
} else {
Some((entry, &node.entries))
},
state,
)?;
if child_verified.count != entry.child_count {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "child count summary mismatch".to_owned(),
});
}
if child_verified.minimum.as_deref() != Some(entry.min_key.as_slice())
|| child_verified.maximum.as_deref() != Some(entry.max_key.as_slice())
{
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "child key-bound summary mismatch".to_owned(),
});
}
for (_, vector) in &child_verified.points {
let required = super::distance::euclidean_radius_up(
score(
super::DistanceMetric::L2Squared,
entry.vector.inline()?,
vector,
),
0.0,
);
if required > entry.covering_radius {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "covering-radius summary is not conservative".to_owned(),
});
}
}
count = count.checked_add(child_verified.count).ok_or_else(|| {
Error::InvalidProximityObject {
kind: "node",
reason: "subtree count overflow".to_owned(),
}
})?;
if minimum.as_ref().map_or(true, |key| entry.min_key < *key) {
minimum = Some(entry.min_key.clone());
}
if maximum.as_ref().map_or(true, |key| entry.max_key > *key) {
maximum = Some(entry.max_key.clone());
}
points.extend(child_verified.points);
}
VerifiedSubtree {
count,
minimum,
maximum,
points,
}
};
if verified.count != node.subtree_count {
return Err(Error::InvalidProximityObject {
kind: "node",
reason: "subtree count mismatch".to_owned(),
});
}
Ok(verified)
}
}
struct VerificationState<'a> {
records: &'a BTreeMap<Vec<u8>, ProximityRecord>,
seen_nodes: HashSet<Cid>,
seen_external_vectors: HashSet<Cid>,
seen_scalar_quantizers: HashSet<Cid>,
seen_leaf_keys: HashSet<Vec<u8>>,
summary: ProximityVerification,
}
struct VerifiedSubtree {
count: u64,
minimum: Option<Vec<u8>>,
maximum: Option<Vec<u8>>,
points: Vec<(Vec<u8>, Vec<f32>)>,
}
impl VerifiedSubtree {
fn from_points(count: u64, points: Vec<(Vec<u8>, Vec<f32>)>) -> Self {
let minimum = points.iter().map(|(key, _)| key).min().cloned();
let maximum = points.iter().map(|(key, _)| key).max().cloned();
Self {
count,
minimum,
maximum,
points,
}
}
}
fn load_content<S: Store>(store: &S, cid: &Cid) -> Result<Vec<u8>, Error> {
let bytes = store
.get(cid.as_bytes())
.map_err(|error| Error::Store(Box::new(error)))?
.ok_or_else(|| Error::NotFound(cid.clone()))?;
let actual = Cid::from_bytes(&bytes);
if actual != *cid {
return Err(Error::CidMismatch {
expected: cid.clone(),
actual,
});
}
Ok(bytes)
}
fn distance_budget_exhausted(request: &SearchRequest<'_>, stats: &ProximitySearchStats) -> bool {
request
.budget
.max_distance_evaluations
.is_some_and(|maximum| {
stats
.distance_evaluations
.saturating_add(stats.quantized_distance_evaluations)
>= maximum
})
}
fn put_missing_nodes<S: Store>(store: &S, nodes: &[(Cid, Vec<u8>)]) -> Result<usize, Error> {
let keys: Vec<_> = nodes.iter().map(|(cid, _)| cid.as_bytes()).collect();
let existing = store
.batch_get_ordered_unique(&keys)
.map_err(|error| Error::Store(Box::new(error)))?;
for ((expected, _), value) in nodes.iter().zip(&existing) {
if let Some(bytes) = value {
let actual = Cid::from_bytes(bytes);
if actual != *expected {
return Err(Error::CidMismatch {
expected: expected.clone(),
actual,
});
}
}
}
let missing: Vec<_> = nodes
.iter()
.zip(existing)
.filter_map(|((cid, bytes), value)| {
value
.is_none()
.then_some((cid.as_bytes(), bytes.as_slice()))
})
.collect();
store
.batch_put(&missing)
.map_err(|error| Error::Store(Box::new(error)))?;
Ok(missing.len())
}
fn apply_directory_stats(target: &mut ProximityMutationStats, source: CanonicalSpliceStats) {
target.directory_entries_scanned = source.entries_scanned;
target.directory_nodes_read = source.nodes_read;
target.directory_nodes_rebuilt = source.nodes_rebuilt;
target.directory_nodes_written = source.nodes_written;
target.directory_nodes_reused = source.nodes_reused;
target.directory_levels_rebuilt = source.levels_rebuilt;
target.directory_right_edge_rebuilt = source.right_edge_rebuilt;
}
fn validate_mutations(
mutations: impl IntoIterator<Item = ProximityMutation>,
) -> Result<Vec<ProximityMutation>, Error> {
let mut mutations: Vec<_> = mutations.into_iter().collect();
mutations.sort_by(|left, right| left.key.cmp(&right.key));
for pair in mutations.windows(2) {
if pair[0].key == pair[1].key {
return Err(Error::DuplicateProximityKey {
key: pair[0].key.clone(),
});
}
}
Ok(mutations)
}
fn apply_mutations(
records: &mut BTreeMap<Vec<u8>, ProximityRecord>,
mutations: &[ProximityMutation],
config: &ProximityConfig,
) -> Result<(), Error> {
for mutation in mutations {
match &mutation.value {
Some((vector, value)) => {
records.insert(
mutation.key.clone(),
ProximityRecord {
key: mutation.key.clone(),
vector: prepare_vector(config.metric, vector, config.dimensions)?,
value: value.clone(),
},
);
}
None => {
records.remove(&mutation.key);
}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::prolly::proximity::distance::{query_kernel_calls, reset_query_kernel_calls};
use crate::prolly::store::MemStore;
fn config() -> ProximityConfig {
let mut config = ProximityConfig::new(1);
config.hierarchy.log_chunk_size = 1;
config.hierarchy.level_hash_seed = 7;
config.overflow.max_page_bytes = 256 * 1024;
config
}
fn two_representative_map() -> (Arc<MemStore>, ProximityMap<Arc<MemStore>>) {
let keys: Vec<_> = (0..10_000)
.map(|index| format!("candidate-{index}").into_bytes())
.filter(|key| promotion_level(key, 1, 7) == 1)
.take(2)
.collect();
assert_eq!(keys.len(), 2);
let store = Arc::new(MemStore::new());
let map = ProximityMap::build(
store.clone(),
config(),
keys.into_iter()
.enumerate()
.map(|(index, key)| ProximityRecord {
key,
vector: vec![index as f32],
value: Vec::new(),
}),
)
.unwrap();
(store, map)
}
#[test]
fn construction_and_mutation_never_enter_a_query_kernel() {
reset_query_kernel_calls();
let store = Arc::new(MemStore::new());
let map = ProximityMap::build(
store,
config(),
(0..64).map(|index| ProximityRecord {
key: format!("key-{index:03}").into_bytes(),
vector: vec![index as f32],
value: Vec::new(),
}),
)
.unwrap();
let (map, _) = map
.mutate_batch([ProximityMutation {
key: b"key-017".to_vec(),
value: Some((vec![17.25], b"updated".to_vec())),
}])
.unwrap();
assert_eq!(query_kernel_calls(), 0);
let mut request = SearchRequest::exact(&[17.0], 3);
request.kernel = super::super::QueryKernel::SimdDeterministic;
map.search(request).unwrap();
assert!(query_kernel_calls() > 0);
}
fn publish_root_descriptor(
store: &Arc<MemStore>,
map: &ProximityMap<Arc<MemStore>>,
root: ProximityNode,
) -> Cid {
let root_bytes = root.encode().unwrap();
let root_cid = Cid::from_bytes(&root_bytes);
store.put(root_cid.as_bytes(), &root_bytes).unwrap();
let descriptor_bytes = store.get(map.tree.descriptor.as_bytes()).unwrap().unwrap();
let mut descriptor = Descriptor::decode(&descriptor_bytes).unwrap();
descriptor.proximity_root = root_cid;
let descriptor_bytes = descriptor.encode();
let descriptor_cid = Cid::from_bytes(&descriptor_bytes);
store
.put(descriptor_cid.as_bytes(), &descriptor_bytes)
.unwrap();
descriptor_cid
}
fn publish_replacement_root(
store: &Arc<MemStore>,
map: &ProximityMap<Arc<MemStore>>,
root: ProximityNode,
) -> ProximityMap<Arc<MemStore>> {
let descriptor_cid = publish_root_descriptor(store, map, root);
ProximityMap::load(store.clone(), descriptor_cid).unwrap()
}
#[test]
fn verify_rejects_a_leaf_vector_that_disagrees_with_the_exact_directory() {
let store = Arc::new(MemStore::new());
let mut leaf_config = config();
leaf_config.hierarchy.log_chunk_size = 63;
let map = ProximityMap::build(
store.clone(),
leaf_config,
[ProximityRecord {
key: b"key".to_vec(),
vector: vec![1.0],
value: Vec::new(),
}],
)
.unwrap();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
root.entries[0].vector = super::super::storage::VectorRef::Inline(vec![2.0]);
let corrupt = publish_replacement_root(&store, &map, root);
assert!(matches!(
corrupt.verify(),
Err(Error::InvalidProximityObject { reason, .. })
if reason == "leaf vector disagrees with exact directory"
));
}
#[test]
fn verify_rejects_repeated_child_ownership() {
let (store, map) = two_representative_map();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
assert_eq!(root.level, 1);
assert_eq!(root.entries.len(), 2);
root.entries[1].child = root.entries[0].child.clone();
let corrupt = publish_replacement_root(&store, &map, root);
assert!(matches!(
corrupt.verify(),
Err(Error::InvalidProximityObject { reason, .. })
if reason == "cycle or repeated child ownership"
));
}
#[test]
fn verify_rejects_an_invalid_child_level() {
let (store, map) = two_representative_map();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
root.entries[0].child = Some(map.tree.proximity_root.clone());
let corrupt = publish_replacement_root(&store, &map, root);
assert!(matches!(
corrupt.verify(),
Err(Error::InvalidProximityObject { reason, .. })
if reason == "unexpected logical level"
));
}
#[test]
fn verify_rejects_a_representative_below_its_node_level() {
let (store, map) = two_representative_map();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
let replacement = (0..10_000)
.map(|index| format!("!invalid-{index}").into_bytes())
.find(|key| promotion_level(key, 1, 7) == 0 && key < &root.entries[1].key)
.unwrap();
root.entries[0].key = replacement.clone();
root.entries[0].min_key = replacement;
let corrupt = publish_replacement_root(&store, &map, root);
assert!(matches!(
corrupt.verify(),
Err(Error::InvalidProximityObject { reason, .. })
if reason == "entry appears above its deterministic promotion level"
));
}
#[test]
fn verify_rejects_a_non_nearest_parent_route() {
let (store, map) = two_representative_map();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
let first = root.entries[0].child.clone();
root.entries[0].child = root.entries[1].child.clone();
root.entries[1].child = first;
let corrupt = publish_replacement_root(&store, &map, root);
assert!(matches!(
corrupt.verify(),
Err(Error::InvalidProximityObject { reason, .. })
if reason == "nearest-representative invariant violated"
));
}
#[test]
fn load_rejects_a_root_subtree_count_that_disagrees_with_the_descriptor() {
let (store, map) = two_representative_map();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
root.subtree_count += 1;
root.entries[0].child_count += 1;
let descriptor = publish_root_descriptor(&store, &map, root);
assert!(matches!(
ProximityMap::load(store, descriptor),
Err(Error::InvalidProximityObject { reason, .. })
if reason == "record count disagrees with proximity root"
));
}
#[test]
fn verify_rejects_a_non_conservative_covering_radius() {
let store = Arc::new(MemStore::new());
let map = ProximityMap::build(
store.clone(),
config(),
(0..128).map(|index| ProximityRecord {
key: format!("radius-{index:04}").into_bytes(),
vector: vec![index as f32],
value: Vec::new(),
}),
)
.unwrap();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
let entry = root
.entries
.iter_mut()
.find(|entry| entry.covering_radius > 0.0)
.expect("test hierarchy has a nontrivial cluster");
entry.covering_radius = 0.0;
let corrupt = publish_replacement_root(&store, &map, root);
assert!(matches!(
corrupt.verify(),
Err(Error::InvalidProximityObject { reason, .. })
if reason == "covering-radius summary is not conservative"
));
}
#[test]
fn verify_rejects_a_scalar_quantizer_that_disagrees_with_its_node() {
let store = Arc::new(MemStore::new());
let mut quantized_config = config();
quantized_config.scalar_quantization =
Some(super::super::ScalarQuantizationConfig { group_size: 1 });
let map = ProximityMap::build(
store.clone(),
quantized_config,
(0..64).map(|index| ProximityRecord {
key: format!("quantized-{index:03}").into_bytes(),
vector: vec![index as f32],
value: Vec::new(),
}),
)
.unwrap();
let bytes = store
.get(map.tree.proximity_root.as_bytes())
.unwrap()
.unwrap();
let mut root = ProximityNode::decode(&bytes, 1).unwrap();
let fake_vectors = vec![vec![999.0]; root.entries.len()];
let fake_refs: Vec<_> = fake_vectors.iter().map(Vec::as_slice).collect();
let fake = ScalarQuantized::build(&fake_refs, 1, 1).unwrap();
let fake_bytes = fake.encode().unwrap();
let fake_cid = Cid::from_bytes(&fake_bytes);
store.put(fake_cid.as_bytes(), &fake_bytes).unwrap();
root.quantizer = Some(fake_cid);
let corrupt = publish_replacement_root(&store, &map, root);
assert!(matches!(
corrupt.verify(),
Err(Error::InvalidProximityObject { kind: "quantizer", reason })
if reason.contains("disagree")
));
}
}