hyphae_retrieval/
canonical.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::{
4    cmp::Ordering,
5    collections::BTreeSet,
6    time::{Duration, Instant},
7};
8
9use hyphae_core::{Q15Vector, SCORE_NANOS_SCALE, VectorSpaceName};
10use thiserror::Error;
11
12/// One durable exact-retrieval candidate.
13#[derive(Clone, Debug, Eq, PartialEq)]
14pub struct DurableVectorRecord {
15    /// Globally unique nonempty object key within the selected space.
16    pub key: Vec<u8>,
17    /// Canonical signed-Q15 vector.
18    pub vector: Q15Vector,
19}
20
21/// Durable exact-retrieval request under semantics v2.
22#[derive(Clone, Debug, Eq, PartialEq)]
23pub struct ExactRetrievalRequest {
24    /// Canonical named vector space.
25    pub vector_space: VectorSpaceName,
26    /// Canonical signed-Q15 query.
27    pub query: Q15Vector,
28    /// Maximum returned matches.
29    pub limit: usize,
30    /// Inclusive minimum canonical score.
31    pub minimum_score_nanos: i64,
32    /// Minimum canonical best/runner-up margin.
33    pub minimum_margin_nanos: u64,
34}
35
36/// Runtime and shape limits for durable exact retrieval.
37#[derive(Clone, Debug, Eq, PartialEq)]
38pub struct ExactRetrievalLimits {
39    /// Maximum candidates inspected.
40    pub max_candidates: u64,
41    /// Maximum aggregate candidate key and vector bytes loaded.
42    pub max_candidate_bytes: u64,
43    /// Maximum requested results.
44    pub max_returned: usize,
45    /// Cooperative timeout.
46    pub timeout: Duration,
47}
48
49impl Default for ExactRetrievalLimits {
50    fn default() -> Self {
51        Self {
52            max_candidates: 100_000,
53            max_candidate_bytes: 256 * 1024 * 1024,
54            max_returned: 1_000,
55            timeout: Duration::from_secs(30),
56        }
57    }
58}
59
60/// One canonical exact match.
61#[derive(Clone, Debug, Eq, PartialEq)]
62pub struct ExactRetrievalMatch {
63    /// Binary object key.
64    pub key: Vec<u8>,
65    /// Canonical integer cosine score.
66    pub score_nanos: i64,
67}
68
69/// Complete durable exact-retrieval outcome.
70#[derive(Clone, Debug, Eq, PartialEq)]
71pub enum ExactRetrievalOutcome {
72    /// Accepted globally ranked matches.
73    Matches {
74        /// Final matches.
75        matches: Vec<ExactRetrievalMatch>,
76        /// Global candidates inspected.
77        scanned_candidates: u64,
78    },
79    /// Typed normal abstention.
80    Abstained(ExactAbstention),
81}
82
83/// Stable durable exact-retrieval abstention evidence.
84#[derive(Clone, Debug, Eq, PartialEq)]
85pub struct ExactAbstention {
86    /// Stable reason.
87    pub reason: ExactAbstentionReason,
88    /// Best canonical score when present.
89    pub best_score_nanos: Option<i64>,
90    /// Runner-up canonical score when present.
91    pub runner_up_score_nanos: Option<i64>,
92    /// Global candidates inspected.
93    pub scanned_candidates: u64,
94}
95
96/// Stable reason why durable exact retrieval declined matches.
97#[derive(Clone, Copy, Debug, Eq, PartialEq)]
98pub enum ExactAbstentionReason {
99    /// The selected space contains no vectors.
100    NoCandidates,
101    /// The best score is below the request threshold.
102    BelowThreshold,
103    /// The best/runner-up gap is below the request margin.
104    Ambiguous,
105}
106
107/// Failure to execute durable exact retrieval completely.
108#[derive(Clone, Debug, Error, Eq, PartialEq)]
109pub enum ExactRetrievalError {
110    /// Candidate keys must be nonempty.
111    #[error("candidate key must be nonempty")]
112    EmptyCandidateKey,
113    /// Candidate keys must be unique in one space.
114    #[error("duplicate candidate key")]
115    DuplicateCandidateKey,
116    /// Candidate and query dimensions must match.
117    #[error("candidate dimension {found} does not match query dimension {expected}")]
118    DimensionMismatch {
119        /// Query dimension.
120        expected: u16,
121        /// Candidate dimension.
122        found: u16,
123    },
124    /// At least one match must be requested.
125    #[error("retrieval limit must be nonzero")]
126    ZeroLimit,
127    /// The requested result count exceeds policy.
128    #[error("retrieval limit {requested} exceeds maximum {maximum}")]
129    ResultLimitExceeded {
130        /// Requested result count.
131        requested: usize,
132        /// Maximum result count.
133        maximum: usize,
134    },
135    /// The score threshold is outside the canonical cosine range.
136    #[error("minimum score nanos must be in [-1000000000, 1000000000]")]
137    InvalidMinimumScore,
138    /// The margin is outside the canonical cosine range width.
139    #[error("minimum margin nanos must be in [0, 2000000000]")]
140    InvalidMinimumMargin,
141    /// The global candidate budget was exhausted.
142    #[error("global candidate budget exceeded: {maximum}")]
143    CandidateBudgetExceeded {
144        /// Maximum candidates.
145        maximum: u64,
146    },
147    /// The aggregate candidate key/vector byte budget was exhausted.
148    #[error("global candidate byte budget exceeded: {maximum}")]
149    CandidateByteBudgetExceeded {
150        /// Maximum aggregate candidate bytes.
151        maximum: u64,
152    },
153    /// The cooperative timeout elapsed.
154    #[error("retrieval execution timed out")]
155    TimedOut,
156    /// Checked integer arithmetic failed.
157    #[error("canonical score arithmetic overflow")]
158    ArithmeticOverflow,
159}
160
161/// Monotonic clock used by the canonical executor.
162pub trait ExactRetrievalClock {
163    /// Returns a nondecreasing local duration.
164    fn now(&mut self) -> Duration;
165}
166
167/// Executes durable exact retrieval using the process monotonic clock.
168///
169/// # Errors
170///
171/// Returns the same complete-or-error outcomes as
172/// [`retrieve_exact_with_clock`].
173pub fn retrieve_exact(
174    candidates: &[DurableVectorRecord],
175    request: &ExactRetrievalRequest,
176    limits: &ExactRetrievalLimits,
177) -> Result<ExactRetrievalOutcome, ExactRetrievalError> {
178    retrieve_exact_with_clock(
179        candidates,
180        request,
181        limits,
182        &mut SystemClock {
183            started: Instant::now(),
184        },
185    )
186}
187
188struct SystemClock {
189    started: Instant,
190}
191
192impl ExactRetrievalClock for SystemClock {
193    fn now(&mut self) -> Duration {
194        self.started.elapsed()
195    }
196}
197
198/// Executes durable exact retrieval using a caller-supplied monotonic clock.
199///
200/// # Errors
201///
202/// Returns an error for invalid input, dimensions, budget, timeout, duplicate
203/// keys, or arithmetic overflow. Errors return no partial ranking.
204pub fn retrieve_exact_with_clock(
205    candidates: &[DurableVectorRecord],
206    request: &ExactRetrievalRequest,
207    limits: &ExactRetrievalLimits,
208    clock: &mut impl ExactRetrievalClock,
209) -> Result<ExactRetrievalOutcome, ExactRetrievalError> {
210    validate_request(request, limits)?;
211    let started = clock.now();
212    let deadline = started.checked_add(limits.timeout).unwrap_or(Duration::MAX);
213    check_timeout(clock, deadline)?;
214
215    let mut keys = BTreeSet::new();
216    let mut ranked = Vec::with_capacity(candidates.len().min(request.limit));
217    let mut scanned = 0_u64;
218    let mut scanned_bytes = 0_u64;
219    for candidate in candidates {
220        check_timeout(clock, deadline)?;
221        if scanned >= limits.max_candidates {
222            return Err(ExactRetrievalError::CandidateBudgetExceeded {
223                maximum: limits.max_candidates,
224            });
225        }
226        scanned = scanned.saturating_add(1);
227        let candidate_bytes = u64::try_from(candidate.key.len())
228            .ok()
229            .and_then(|key_bytes| {
230                u64::try_from(candidate.vector.as_slice().len())
231                    .ok()
232                    .and_then(|elements| elements.checked_mul(2))
233                    .and_then(|vector_bytes| key_bytes.checked_add(vector_bytes))
234            })
235            .ok_or(ExactRetrievalError::CandidateByteBudgetExceeded {
236                maximum: limits.max_candidate_bytes,
237            })?;
238        scanned_bytes = scanned_bytes.checked_add(candidate_bytes).ok_or(
239            ExactRetrievalError::CandidateByteBudgetExceeded {
240                maximum: limits.max_candidate_bytes,
241            },
242        )?;
243        if scanned_bytes > limits.max_candidate_bytes {
244            return Err(ExactRetrievalError::CandidateByteBudgetExceeded {
245                maximum: limits.max_candidate_bytes,
246            });
247        }
248        if candidate.key.is_empty() {
249            return Err(ExactRetrievalError::EmptyCandidateKey);
250        }
251        if !keys.insert(candidate.key.as_slice()) {
252            return Err(ExactRetrievalError::DuplicateCandidateKey);
253        }
254        if candidate.vector.dimension() != request.query.dimension() {
255            return Err(ExactRetrievalError::DimensionMismatch {
256                expected: request.query.dimension(),
257                found: candidate.vector.dimension(),
258            });
259        }
260        ranked.push(ExactRetrievalMatch {
261            key: candidate.key.clone(),
262            score_nanos: cosine_score_nanos(&request.query, &candidate.vector)?,
263        });
264    }
265    ranked.sort_by(compare_matches);
266    check_timeout(clock, deadline)?;
267    Ok(apply_policy(ranked, scanned, request))
268}
269
270/// Computes the canonical integer cosine score from ADR-0015.
271///
272/// # Errors
273///
274/// Returns an arithmetic error if checked accumulation exceeds `u128`/`i128`.
275pub fn cosine_score_nanos(left: &Q15Vector, right: &Q15Vector) -> Result<i64, ExactRetrievalError> {
276    if left.dimension() != right.dimension() {
277        return Err(ExactRetrievalError::DimensionMismatch {
278            expected: left.dimension(),
279            found: right.dimension(),
280        });
281    }
282    let mut dot = 0_i128;
283    let mut left_squared = 0_u128;
284    let mut right_squared = 0_u128;
285    for (left_value, right_value) in left.as_slice().iter().zip(right.as_slice()) {
286        let left_value = i128::from(*left_value);
287        let right_value = i128::from(*right_value);
288        dot = dot
289            .checked_add(
290                left_value
291                    .checked_mul(right_value)
292                    .ok_or(ExactRetrievalError::ArithmeticOverflow)?,
293            )
294            .ok_or(ExactRetrievalError::ArithmeticOverflow)?;
295        left_squared = left_squared
296            .checked_add(
297                u128::try_from(left_value * left_value)
298                    .map_err(|_| ExactRetrievalError::ArithmeticOverflow)?,
299            )
300            .ok_or(ExactRetrievalError::ArithmeticOverflow)?;
301        right_squared = right_squared
302            .checked_add(
303                u128::try_from(right_value * right_value)
304                    .map_err(|_| ExactRetrievalError::ArithmeticOverflow)?,
305            )
306            .ok_or(ExactRetrievalError::ArithmeticOverflow)?;
307    }
308    let norm_product = left_squared
309        .checked_mul(right_squared)
310        .ok_or(ExactRetrievalError::ArithmeticOverflow)?;
311    let denominator = integer_sqrt(norm_product);
312    if denominator == 0 {
313        return Err(ExactRetrievalError::ArithmeticOverflow);
314    }
315    let absolute_dot = dot.unsigned_abs();
316    let score_scale =
317        u128::try_from(SCORE_NANOS_SCALE).map_err(|_| ExactRetrievalError::ArithmeticOverflow)?;
318    let numerator = absolute_dot
319        .checked_mul(score_scale)
320        .and_then(|value| value.checked_add(denominator / 2))
321        .ok_or(ExactRetrievalError::ArithmeticOverflow)?;
322    let magnitude = (numerator / denominator).min(score_scale);
323    let magnitude =
324        i64::try_from(magnitude).map_err(|_| ExactRetrievalError::ArithmeticOverflow)?;
325    Ok(if dot < 0 { -magnitude } else { magnitude })
326}
327
328fn validate_request(
329    request: &ExactRetrievalRequest,
330    limits: &ExactRetrievalLimits,
331) -> Result<(), ExactRetrievalError> {
332    if request.limit == 0 {
333        return Err(ExactRetrievalError::ZeroLimit);
334    }
335    if request.limit > limits.max_returned {
336        return Err(ExactRetrievalError::ResultLimitExceeded {
337            requested: request.limit,
338            maximum: limits.max_returned,
339        });
340    }
341    if !(-SCORE_NANOS_SCALE..=SCORE_NANOS_SCALE).contains(&request.minimum_score_nanos) {
342        return Err(ExactRetrievalError::InvalidMinimumScore);
343    }
344    let maximum_margin = u64::try_from(SCORE_NANOS_SCALE.saturating_mul(2))
345        .map_err(|_| ExactRetrievalError::InvalidMinimumMargin)?;
346    if request.minimum_margin_nanos > maximum_margin {
347        return Err(ExactRetrievalError::InvalidMinimumMargin);
348    }
349    Ok(())
350}
351
352fn check_timeout(
353    clock: &mut impl ExactRetrievalClock,
354    deadline: Duration,
355) -> Result<(), ExactRetrievalError> {
356    if clock.now() >= deadline {
357        Err(ExactRetrievalError::TimedOut)
358    } else {
359        Ok(())
360    }
361}
362
363fn compare_matches(left: &ExactRetrievalMatch, right: &ExactRetrievalMatch) -> Ordering {
364    right
365        .score_nanos
366        .cmp(&left.score_nanos)
367        .then_with(|| left.key.cmp(&right.key))
368}
369
370fn apply_policy(
371    ranked: Vec<ExactRetrievalMatch>,
372    scanned: u64,
373    request: &ExactRetrievalRequest,
374) -> ExactRetrievalOutcome {
375    let Some(best) = ranked.first() else {
376        return ExactRetrievalOutcome::Abstained(ExactAbstention {
377            reason: ExactAbstentionReason::NoCandidates,
378            best_score_nanos: None,
379            runner_up_score_nanos: None,
380            scanned_candidates: scanned,
381        });
382    };
383    let runner_up = ranked.get(1).map(|candidate| candidate.score_nanos);
384    if best.score_nanos < request.minimum_score_nanos {
385        return ExactRetrievalOutcome::Abstained(ExactAbstention {
386            reason: ExactAbstentionReason::BelowThreshold,
387            best_score_nanos: Some(best.score_nanos),
388            runner_up_score_nanos: runner_up,
389            scanned_candidates: scanned,
390        });
391    }
392    if runner_up.is_some_and(|score| {
393        let margin = best.score_nanos.saturating_sub(score);
394        u64::try_from(margin).unwrap_or_default() < request.minimum_margin_nanos
395    }) {
396        return ExactRetrievalOutcome::Abstained(ExactAbstention {
397            reason: ExactAbstentionReason::Ambiguous,
398            best_score_nanos: Some(best.score_nanos),
399            runner_up_score_nanos: runner_up,
400            scanned_candidates: scanned,
401        });
402    }
403    ExactRetrievalOutcome::Matches {
404        matches: ranked.into_iter().take(request.limit).collect(),
405        scanned_candidates: scanned,
406    }
407}
408
409fn integer_sqrt(value: u128) -> u128 {
410    if value < 2 {
411        return value;
412    }
413    let mut lower = 1_u128;
414    let mut upper = (value >> 1).saturating_add(1);
415    while lower <= upper {
416        let middle = lower + ((upper - lower) >> 1);
417        match middle.checked_mul(middle) {
418            Some(square) if square == value => return middle,
419            Some(square) if square < value => lower = middle.saturating_add(1),
420            _ => upper = middle.saturating_sub(1),
421        }
422    }
423    upper
424}
425
426#[cfg(test)]
427mod tests {
428    use std::time::Duration;
429
430    use hyphae_core::{Q15Vector, VectorSpaceName};
431
432    use super::{
433        DurableVectorRecord, ExactAbstentionReason, ExactRetrievalClock, ExactRetrievalLimits,
434        ExactRetrievalOutcome, ExactRetrievalRequest, cosine_score_nanos,
435        retrieve_exact_with_clock,
436    };
437
438    struct StepClock(Duration);
439
440    impl ExactRetrievalClock for StepClock {
441        fn now(&mut self) -> Duration {
442            let current = self.0;
443            self.0 = self.0.saturating_add(Duration::from_millis(1));
444            current
445        }
446    }
447
448    fn request() -> Result<ExactRetrievalRequest, hyphae_core::VectorValueError> {
449        Ok(ExactRetrievalRequest {
450            vector_space: VectorSpaceName::new("semantic")?,
451            query: Q15Vector::new(vec![32_767, 0])?,
452            limit: 3,
453            minimum_score_nanos: -1_000_000_000,
454            minimum_margin_nanos: 0,
455        })
456    }
457
458    #[test]
459    fn canonical_scores_cover_same_orthogonal_and_opposite()
460    -> Result<(), Box<dyn std::error::Error>> {
461        let query = Q15Vector::new(vec![32_767, 0])?;
462        assert_eq!(
463            cosine_score_nanos(&query, &Q15Vector::new(vec![32_767, 0])?)?,
464            1_000_000_000
465        );
466        assert_eq!(
467            cosine_score_nanos(&query, &Q15Vector::new(vec![0, 32_767])?)?,
468            0
469        );
470        assert_eq!(
471            cosine_score_nanos(&query, &Q15Vector::new(vec![-32_767, 0])?)?,
472            -1_000_000_000
473        );
474        Ok(())
475    }
476
477    #[test]
478    fn exact_ranking_uses_score_then_binary_key() -> Result<(), Box<dyn std::error::Error>> {
479        let candidates = vec![
480            DurableVectorRecord {
481                key: vec![0xff],
482                vector: Q15Vector::new(vec![7, 7])?,
483            },
484            DurableVectorRecord {
485                key: vec![0],
486                vector: Q15Vector::new(vec![2, 2])?,
487            },
488        ];
489        let mut tied = request()?;
490        tied.query = Q15Vector::new(vec![1, 1])?;
491        tied.limit = 2;
492        let outcome = retrieve_exact_with_clock(
493            &candidates,
494            &tied,
495            &ExactRetrievalLimits::default(),
496            &mut StepClock(Duration::ZERO),
497        )?;
498        let ExactRetrievalOutcome::Matches { matches, .. } = outcome else {
499            return Err("unexpected abstention".into());
500        };
501        assert_eq!(matches[0].key, vec![0]);
502        assert_eq!(matches[1].key, vec![0xff]);
503        Ok(())
504    }
505
506    #[test]
507    fn empty_space_is_typed_abstention() -> Result<(), Box<dyn std::error::Error>> {
508        let outcome = retrieve_exact_with_clock(
509            &[],
510            &request()?,
511            &ExactRetrievalLimits::default(),
512            &mut StepClock(Duration::ZERO),
513        )?;
514        assert!(matches!(
515            outcome,
516            ExactRetrievalOutcome::Abstained(super::ExactAbstention {
517                reason: ExactAbstentionReason::NoCandidates,
518                ..
519            })
520        ));
521        Ok(())
522    }
523
524    #[test]
525    fn sparse_near_ties_are_distinct_and_obey_margin_policy()
526    -> Result<(), Box<dyn std::error::Error>> {
527        let candidates = vec![
528            DurableVectorRecord {
529                key: b"near-a".to_vec(),
530                vector: Q15Vector::new(vec![32_767, 16, 0, 0, 0, 0, 0, 0])?,
531            },
532            DurableVectorRecord {
533                key: b"near-b".to_vec(),
534                vector: Q15Vector::new(vec![32_767, 32, 0, 0, 0, 0, 0, 0])?,
535            },
536        ];
537        let mut near_tie = request()?;
538        near_tie.query = Q15Vector::new(vec![32_767, 0, 0, 0, 0, 0, 0, 0])?;
539        near_tie.limit = 2;
540        let ranked = retrieve_exact_with_clock(
541            &candidates,
542            &near_tie,
543            &ExactRetrievalLimits::default(),
544            &mut StepClock(Duration::ZERO),
545        )?;
546        let ExactRetrievalOutcome::Matches { matches, .. } = ranked else {
547            return Err("unexpected abstention".into());
548        };
549        assert_eq!(matches[0].key, b"near-a");
550        assert_eq!(matches[1].key, b"near-b");
551        let margin = u64::try_from(matches[0].score_nanos - matches[1].score_nanos)?;
552        assert!(margin > 0);
553
554        near_tie.minimum_margin_nanos = margin.saturating_add(1);
555        let abstained = retrieve_exact_with_clock(
556            &candidates,
557            &near_tie,
558            &ExactRetrievalLimits::default(),
559            &mut StepClock(Duration::ZERO),
560        )?;
561        assert!(matches!(
562            abstained,
563            ExactRetrievalOutcome::Abstained(super::ExactAbstention {
564                reason: ExactAbstentionReason::Ambiguous,
565                ..
566            })
567        ));
568        Ok(())
569    }
570}
hyphae_retrieval/canonical.rs

hyphae_retrieval/
canonical.rs
