prolly/prolly/proximity/
mod.rs

1//! Deterministic, content-addressed approximate nearest-neighbor maps.
2
3pub(crate) mod accelerator;
4mod build;
5mod builder;
6mod cache;
7mod distance;
8mod map;
9mod mutation;
10mod proof;
11mod search;
12pub(crate) mod storage;
13mod vector;
14
15use super::cid::Cid;
16use super::error::Error;
17use super::tree::Tree;
18
19pub use accelerator::hnsw::{HnswBuildStats, HnswConfig, HnswIndex};
20pub use accelerator::pq::{
21    ProductQuantizationBuildStats, ProductQuantizationConfig, ProductQuantizationQuality,
22    ProductQuantizer,
23};
24pub use build::{BuildParallelism, ProximityBuildStats};
25pub use map::ProximityMap;
26pub use proof::{
27    ProximityMembershipProof, ProximityMembershipVerification, ProximityProofFilter,
28    ProximitySearchClaim, ProximitySearchEvent, ProximitySearchProof, ProximitySearchRequest,
29    ProximitySearchVerification, ProximityStructuralProof, ProximityStructuralVerification,
30};
31#[cfg(feature = "async-store")]
32pub use search::{AsyncIoConfig, AsyncProximityMap, AsyncSearchControl, CancellationToken};
33pub use search::{ProximityFilter, SearchRequest};
34
35const MIN_PROXIMITY_NODE_BYTES: u32 = 64;
36
37/// Distance function committed into a proximity-map descriptor.
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
39pub enum DistanceMetric {
40    /// Squared Euclidean distance.
41    L2Squared,
42    /// One minus the dot product of canonical unit vectors.
43    Cosine,
44    /// Negated dot product, preserving lower-is-better ordering.
45    InnerProduct,
46}
47
48impl DistanceMetric {
49    pub(crate) fn id(self) -> u8 {
50        match self {
51            Self::L2Squared => 1,
52            Self::Cosine => 2,
53            Self::InnerProduct => 3,
54        }
55    }
56
57    pub(crate) fn from_id(id: u8) -> Result<Self, Error> {
58        match id {
59            1 => Ok(Self::L2Squared),
60            2 => Ok(Self::Cosine),
61            3 => Ok(Self::InnerProduct),
62            _ => Err(Error::InvalidProximityObject {
63                kind: "descriptor",
64                reason: format!("unknown distance metric id {id}"),
65            }),
66        }
67    }
68}
69
70/// Deterministic hierarchy promotion configuration.
71#[derive(Clone, Debug, PartialEq, Eq)]
72pub struct HierarchyConfig {
73    /// Number of leading hash bits consumed per promotion level.
74    pub log_chunk_size: u8,
75    /// Seed for deterministic key promotion.
76    pub level_hash_seed: u64,
77}
78
79impl Default for HierarchyConfig {
80    fn default() -> Self {
81        Self {
82            log_chunk_size: 8,
83            level_hash_seed: 0,
84        }
85    }
86}
87
88/// Canonical physical-node paging configuration.
89#[derive(Clone, Debug, PartialEq, Eq)]
90pub struct OverflowConfig {
91    /// Minimum encoded bytes before a content boundary may close a page.
92    pub min_page_bytes: u32,
93    /// Preferred encoded page size used for capacity planning.
94    pub target_page_bytes: u32,
95    /// Hard encoded-byte limit for one physical proximity object.
96    pub max_page_bytes: u32,
97    /// Seed for deterministic overflow boundaries.
98    pub hash_seed: u64,
99}
100
101impl Default for OverflowConfig {
102    fn default() -> Self {
103        Self {
104            min_page_bytes: 4 * 1024,
105            target_page_bytes: 64 * 1024,
106            max_page_bytes: 4 * 1024 * 1024,
107            hash_seed: 0,
108        }
109    }
110}
111
112/// Canonical inline versus external representative-vector policy.
113#[derive(Clone, Debug, PartialEq, Eq)]
114pub struct VectorStorageConfig {
115    /// Maximum canonical vector bytes stored inline in a PRXN object.
116    pub inline_threshold_bytes: u32,
117}
118
119impl Default for VectorStorageConfig {
120    fn default() -> Self {
121        Self {
122            inline_threshold_bytes: 64 * 1024,
123        }
124    }
125}
126
127/// Node-local symmetric signed-int8 routing configuration.
128#[derive(Clone, Debug, PartialEq, Eq)]
129pub struct ScalarQuantizationConfig {
130    /// Consecutive dimensions sharing one canonical scale.
131    pub group_size: u32,
132}
133
134/// Shape-affecting configuration for a proximity map.
135#[derive(Clone, Debug, PartialEq, Eq)]
136pub struct ProximityConfig {
137    /// Number of finite `f32` components in every stored and query vector.
138    pub dimensions: u32,
139    /// Persisted distance metric.
140    pub metric: DistanceMetric,
141    /// Promotion and hierarchy shape.
142    pub hierarchy: HierarchyConfig,
143    /// Physical node paging and limits.
144    pub overflow: OverflowConfig,
145    /// Representative vector storage policy.
146    pub vector_storage: VectorStorageConfig,
147    /// Optional committed local routing quantization.
148    pub scalar_quantization: Option<ScalarQuantizationConfig>,
149}
150
151impl ProximityConfig {
152    /// Construct a production-oriented squared-L2 configuration.
153    pub fn new(dimensions: u32) -> Self {
154        Self {
155            dimensions,
156            metric: DistanceMetric::L2Squared,
157            hierarchy: HierarchyConfig::default(),
158            overflow: OverflowConfig::default(),
159            vector_storage: VectorStorageConfig::default(),
160            scalar_quantization: None,
161        }
162    }
163
164    /// Validate all shape-affecting fields before reading or writing nodes.
165    pub fn validate(&self) -> Result<(), Error> {
166        if self.dimensions == 0 {
167            return Err(Error::InvalidProximityConfig {
168                reason: "dimensions must be greater than zero".to_owned(),
169            });
170        }
171        if !(1..=63).contains(&self.hierarchy.log_chunk_size) {
172            return Err(Error::InvalidProximityConfig {
173                reason: "log_chunk_size must be in 1..=63".to_owned(),
174            });
175        }
176        if self.overflow.max_page_bytes < MIN_PROXIMITY_NODE_BYTES {
177            return Err(Error::InvalidProximityConfig {
178                reason: format!("max_page_bytes must be at least {MIN_PROXIMITY_NODE_BYTES}"),
179            });
180        }
181        if self.overflow.min_page_bytes == 0
182            || self.overflow.min_page_bytes > self.overflow.target_page_bytes
183            || self.overflow.target_page_bytes > self.overflow.max_page_bytes
184        {
185            return Err(Error::InvalidProximityConfig {
186                reason: "overflow page sizes must satisfy 0 < min <= target <= max".to_owned(),
187            });
188        }
189        if self.vector_storage.inline_threshold_bytes == 0
190            || self.vector_storage.inline_threshold_bytes > self.overflow.max_page_bytes
191        {
192            return Err(Error::InvalidProximityConfig {
193                reason: "inline_threshold_bytes must be in 1..=max_page_bytes".to_owned(),
194            });
195        }
196        if self
197            .scalar_quantization
198            .as_ref()
199            .is_some_and(|config| config.group_size == 0)
200        {
201            return Err(Error::InvalidProximityConfig {
202                reason: "scalar quantization group_size must be greater than zero".to_owned(),
203            });
204        }
205        Ok(())
206    }
207}
208
209/// Deterministic approximate-search quality policy.
210#[derive(Clone, Copy, Debug, PartialEq, Eq)]
211pub enum AdaptiveQuality {
212    Fast,
213    Balanced,
214    HighRecall,
215}
216
217/// Logical search termination policy.
218#[derive(Clone, Copy, Debug, PartialEq, Eq)]
219pub enum SearchPolicy {
220    Exact,
221    FixedBudget,
222    Adaptive(AdaptiveQuality),
223}
224
225/// Search execution backend.
226#[derive(Clone, Copy, Debug, PartialEq, Eq)]
227pub enum SearchBackend {
228    Native,
229    ProductQuantized,
230    Hnsw,
231    Auto,
232}
233
234/// Runtime-only deterministic query distance implementation.
235#[derive(Clone, Copy, Debug, PartialEq, Eq)]
236pub enum QueryKernel {
237    /// Canonical scalar accumulation used by persisted construction and mutation.
238    ScalarDeterministic,
239    /// Runtime-detected fixed-lane products with canonical scalar-order reduction.
240    SimdDeterministic,
241    /// Prefer deterministic SIMD and fall back to the canonical scalar kernel.
242    AutoDeterministic,
243}
244
245/// Honest completion state for a search result.
246#[derive(Clone, Copy, Debug, PartialEq, Eq)]
247pub enum SearchCompletion {
248    Exact,
249    ApproximatePolicySatisfied,
250    BudgetExhausted,
251    Cancelled,
252    DeadlineExceeded,
253}
254
255/// Deterministic logical search resource limits.
256#[derive(Clone, Debug, Default, PartialEq, Eq)]
257pub struct SearchBudget {
258    pub max_nodes: Option<usize>,
259    pub max_committed_bytes: Option<usize>,
260    pub max_distance_evaluations: Option<usize>,
261    pub max_frontier_entries: Option<usize>,
262}
263
264impl SearchBudget {
265    /// Reject zero limits, which cannot perform a meaningful search.
266    pub fn validate(&self) -> Result<(), Error> {
267        let limits = [
268            ("max_nodes", self.max_nodes),
269            ("max_committed_bytes", self.max_committed_bytes),
270            ("max_distance_evaluations", self.max_distance_evaluations),
271            ("max_frontier_entries", self.max_frontier_entries),
272        ];
273        if let Some((name, _)) = limits.into_iter().find(|(_, value)| *value == Some(0)) {
274            return Err(Error::InvalidProximitySearch {
275                reason: format!("{name} must be greater than zero"),
276            });
277        }
278        Ok(())
279    }
280}
281
282/// Persisted handle for the exact directory and derived proximity hierarchy.
283#[derive(Clone, Debug, PartialEq)]
284pub struct ProximityTree {
285    pub directory: Tree,
286    pub proximity_root: Cid,
287    pub descriptor: Cid,
288    pub count: u64,
289    pub config: ProximityConfig,
290}
291
292/// One logical record supplied to a bulk build.
293#[derive(Clone, Debug, PartialEq)]
294pub struct ProximityRecord {
295    pub key: Vec<u8>,
296    pub vector: Vec<f32>,
297    pub value: Vec<u8>,
298}
299
300/// Vector and application value returned by an exact-key lookup.
301pub type ExactProximityRecord = (Vec<f32>, Vec<u8>);
302
303/// One immutable-map mutation. `None` deletes the key.
304#[derive(Clone, Debug, PartialEq)]
305pub struct ProximityMutation {
306    pub key: Vec<u8>,
307    pub value: Option<(Vec<f32>, Vec<u8>)>,
308}
309
310/// One resolved nearest-neighbor result.
311#[derive(Clone, Debug, PartialEq)]
312pub struct Neighbor {
313    pub key: Vec<u8>,
314    pub value: Vec<u8>,
315    pub distance: f64,
316}
317
318/// Observable resource use for one search.
319#[derive(Clone, Debug, Default, PartialEq, Eq)]
320pub struct ProximitySearchStats {
321    pub levels_visited: usize,
322    pub nodes_read: usize,
323    pub bytes_read: usize,
324    pub physical_bytes_read: usize,
325    pub committed_bytes: usize,
326    pub distance_evaluations: usize,
327    pub quantized_distance_evaluations: usize,
328    pub reranked_candidates: usize,
329    pub frontier_peak: usize,
330}
331
332#[derive(Clone, Debug, PartialEq)]
333pub struct SearchResult {
334    pub neighbors: Vec<Neighbor>,
335    pub stats: ProximitySearchStats,
336    pub completion: SearchCompletion,
337}
338
339/// Observable copy-on-write work for one mutation batch.
340#[derive(Clone, Debug, Default, PartialEq, Eq)]
341pub struct ProximityMutationStats {
342    pub directory_entries_scanned: usize,
343    pub directory_nodes_read: usize,
344    pub directory_nodes_rebuilt: usize,
345    pub directory_nodes_written: usize,
346    pub directory_nodes_reused: usize,
347    pub directory_levels_rebuilt: usize,
348    pub directory_right_edge_rebuilt: bool,
349    pub nodes_read: usize,
350    pub nodes_written: usize,
351    pub nodes_reused: usize,
352    pub records_rebuilt: usize,
353    pub distance_evaluations: usize,
354    pub full_proximity_rebuild: bool,
355}
356
357/// Full structural verification summary.
358#[derive(Clone, Debug, Default, PartialEq, Eq)]
359pub struct ProximityVerification {
360    pub record_count: u64,
361    pub proximity_node_count: usize,
362    pub external_vector_count: usize,
363    pub quantized_node_count: usize,
364    pub scalar_quantizer_count: usize,
365    pub overflow_page_count: usize,
366    pub overflow_directory_count: usize,
367    pub maximum_level: u8,
368    pub maximum_node_bytes: usize,
369    pub distance_checks: usize,
370}
prolly/prolly/proximity/mod.rs

prolly/prolly/proximity/
mod.rs
