1use super::binary::BinaryVector;
9use super::int4::Int4Vector;
10use super::quantized::{QuantizedVector, cosine_similarity_i8_trusted, dot_product_i8_trusted};
11use super::{cosine_similarity, dot_product};
12use crate::error::{EmbedError, Result};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19#[non_exhaustive]
20pub enum NormalizationHint {
21 Unknown,
23 Unit,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
31#[non_exhaustive]
32pub enum QuantizationTier {
33 Full,
35 Int8,
37 Int4,
39 Binary,
41}
42
43impl QuantizationTier {
44 pub fn bytes_per_dim(&self) -> f32 {
46 match self {
47 Self::Full => 4.0,
48 Self::Int8 => 1.0,
49 Self::Int4 => 0.5,
50 Self::Binary => 0.125,
51 }
52 }
53
54 pub fn compression_ratio(&self) -> f32 {
56 4.0 / self.bytes_per_dim()
57 }
58
59 pub fn storage_bytes(&self, dims: usize) -> usize {
61 match self {
62 Self::Full => dims * 4,
63 Self::Int8 => dims,
64 Self::Int4 => dims.div_ceil(2),
65 Self::Binary => dims.div_ceil(8),
66 }
67 }
68
69 pub fn from_age_seconds(age_secs: u64) -> Self {
74 const HOUR: u64 = 3600;
75 const DAY: u64 = 86400;
76 const WEEK: u64 = 604800;
77
78 if age_secs < HOUR {
79 Self::Full
80 } else if age_secs < DAY {
81 Self::Int8
82 } else if age_secs < WEEK {
83 Self::Int4
84 } else {
85 Self::Binary
86 }
87 }
88}
89
90#[derive(Debug, Clone)]
95#[non_exhaustive]
96pub enum QuantizedData {
97 Full(Vec<f32>),
99 Int8(QuantizedVector),
101 Int4(Int4Vector),
103 Binary(BinaryVector),
105}
106
107impl QuantizedData {
108 pub fn tier(&self) -> QuantizationTier {
110 match self {
111 Self::Full(_) => QuantizationTier::Full,
112 Self::Int8(_) => QuantizationTier::Int8,
113 Self::Int4(_) => QuantizationTier::Int4,
114 Self::Binary(_) => QuantizationTier::Binary,
115 }
116 }
117
118 pub fn dims(&self) -> usize {
120 match self {
121 Self::Full(v) => v.len(),
122 Self::Int8(q) => q.len(),
123 Self::Int4(q) => q.dims,
124 Self::Binary(q) => q.dims,
125 }
126 }
127
128 pub fn storage_bytes(&self) -> usize {
130 match self {
131 Self::Full(v) => v.len() * 4,
132 Self::Int8(q) => q.len(),
133 Self::Int4(q) => q.data.len(),
134 Self::Binary(q) => q.data.len(),
135 }
136 }
137
138 pub fn from_f32(vector: &[f32], tier: QuantizationTier) -> Self {
140 match tier {
141 QuantizationTier::Full => Self::Full(vector.to_vec()),
142 QuantizationTier::Int8 => Self::Int8(QuantizedVector::from_f32(vector)),
143 QuantizationTier::Int4 => Self::Int4(Int4Vector::from_f32(vector)),
144 QuantizationTier::Binary => Self::Binary(BinaryVector::from_f32(vector)),
145 }
146 }
147
148 pub fn to_f32(&self) -> Vec<f32> {
150 match self {
151 Self::Full(v) => v.clone(),
152 Self::Int8(q) => q.to_f32(),
153 Self::Int4(q) => q.to_f32(),
154 Self::Binary(q) => q.to_f32(),
155 }
156 }
157
158 pub fn promote(&self, target: QuantizationTier) -> Self {
160 let f32_data = self.to_f32();
161 Self::from_f32(&f32_data, target)
162 }
163
164 pub fn demote(&self, target: QuantizationTier) -> Self {
166 self.promote(target) }
168}
169
170#[derive(Debug, Clone)]
172#[non_exhaustive]
173pub enum PreparedQuery {
174 Full(Vec<f32>),
176 Int8(QuantizedVector),
178 Int4(Int4Vector),
180 Binary(BinaryVector),
182}
183
184impl PreparedQuery {
185 #[inline]
187 pub fn from_f32(query_f32: &[f32], tier: QuantizationTier) -> Self {
188 match tier {
189 QuantizationTier::Full => Self::Full(query_f32.to_vec()),
190 QuantizationTier::Int8 => Self::Int8(QuantizedVector::from_f32(query_f32)),
191 QuantizationTier::Int4 => Self::Int4(Int4Vector::from_f32(query_f32)),
192 QuantizationTier::Binary => Self::Binary(BinaryVector::from_f32(query_f32)),
193 }
194 }
195
196 #[inline]
198 pub fn tier(&self) -> QuantizationTier {
199 match self {
200 Self::Full(_) => QuantizationTier::Full,
201 Self::Int8(_) => QuantizationTier::Int8,
202 Self::Int4(_) => QuantizationTier::Int4,
203 Self::Binary(_) => QuantizationTier::Binary,
204 }
205 }
206
207 #[inline]
209 pub fn dims(&self) -> usize {
210 match self {
211 Self::Full(v) => v.len(),
212 Self::Int8(q) => q.len(),
213 Self::Int4(q) => q.dims,
214 Self::Binary(q) => q.dims,
215 }
216 }
217}
218
219#[inline]
221pub fn prepare_query(query_f32: &[f32], tier: QuantizationTier) -> PreparedQuery {
222 PreparedQuery::from_f32(query_f32, tier)
223}
224
225#[derive(Debug, Clone)]
227pub struct PreparedQueryWithMeta {
228 pub query: PreparedQuery,
230 pub norm: NormalizationHint,
232}
233
234impl PreparedQueryWithMeta {
235 #[inline]
237 pub fn from_f32(query_f32: &[f32], tier: QuantizationTier, norm: NormalizationHint) -> Self {
238 Self {
239 query: PreparedQuery::from_f32(query_f32, tier),
240 norm,
241 }
242 }
243
244 #[inline]
246 pub fn tier(&self) -> QuantizationTier {
247 self.query.tier()
248 }
249
250 #[inline]
252 pub fn dims(&self) -> usize {
253 self.query.dims()
254 }
255}
256
257#[inline]
265pub fn is_unit_norm(v: &[f32]) -> bool {
266 let sq = dot_product(v, v);
267 (sq - 1.0).abs() < 1e-4
268}
269
270#[inline]
272pub fn prepare_query_with_norm(
273 query_f32: &[f32],
274 tier: QuantizationTier,
275 norm: NormalizationHint,
276) -> PreparedQueryWithMeta {
277 PreparedQueryWithMeta::from_f32(query_f32, tier, norm)
278}
279
280#[inline]
285pub fn approximate_cosine_distance_prepared(
286 query: &PreparedQuery,
287 stored: &QuantizedData,
288) -> Result<f32> {
289 match (query, stored) {
290 (PreparedQuery::Full(q), QuantizedData::Full(s)) => Ok(1.0 - cosine_similarity(q, s)),
291 (PreparedQuery::Int8(q), QuantizedData::Int8(s)) => {
292 Ok(1.0 - cosine_similarity_i8_trusted(s, q))
293 }
294 (PreparedQuery::Int4(q), QuantizedData::Int4(s)) => Ok(s.cosine_distance(q)),
295 (PreparedQuery::Binary(q), QuantizedData::Binary(s)) => Ok(s.cosine_distance_approx(q)),
296 _ => Err(EmbedError::TierMismatch {
297 op: "approximate_cosine_distance_prepared",
298 expected: stored.tier(),
299 actual: query.tier(),
300 }),
301 }
302}
303
304#[inline]
306pub fn try_approximate_cosine_distance_prepared(
307 query: &PreparedQuery,
308 stored: &QuantizedData,
309) -> Result<f32> {
310 approximate_cosine_distance_prepared(query, stored)
311}
312
313#[inline]
315pub fn try_approximate_dot_product_prepared(
316 query: &PreparedQuery,
317 stored: &QuantizedData,
318) -> Result<f32> {
319 approximate_dot_product_prepared(query, stored)
320}
321
322#[inline]
335pub fn approximate_cosine_distance_prepared_with_meta(
336 meta: &PreparedQueryWithMeta,
337 stored: &QuantizedData,
338 _stored_norm: NormalizationHint,
339) -> Result<f32> {
340 approximate_cosine_distance_prepared(&meta.query, stored)
341}
342
343#[inline]
348pub fn approximate_dot_product_prepared(
349 query: &PreparedQuery,
350 stored: &QuantizedData,
351) -> Result<f32> {
352 match (query, stored) {
353 (PreparedQuery::Full(q), QuantizedData::Full(s)) => Ok(dot_product(q, s)),
354 (PreparedQuery::Int8(q), QuantizedData::Int8(s)) => Ok(dot_product_i8_trusted(q, s)),
355 (PreparedQuery::Int4(q), QuantizedData::Int4(s)) => Ok(s.dot_product(q)),
356 (PreparedQuery::Binary(_), QuantizedData::Binary(_)) => Err(EmbedError::Internal(
357 "Binary has no prepared dot product; use approximate_cosine_distance_prepared".into(),
358 )),
359 _ => Err(EmbedError::TierMismatch {
360 op: "approximate_dot_product_prepared",
361 expected: stored.tier(),
362 actual: query.tier(),
363 }),
364 }
365}
366
367#[inline]
371pub fn batch_approximate_cosine_distance_prepared(
372 query: &PreparedQuery,
373 stored: &[QuantizedData],
374) -> Result<Vec<f32>> {
375 stored
376 .iter()
377 .map(|item| approximate_cosine_distance_prepared(query, item))
378 .collect()
379}
380
381#[inline]
386pub fn batch_approximate_cosine_distance_prepared_into(
387 query: &PreparedQuery,
388 stored: &[QuantizedData],
389 out: &mut Vec<f32>,
390) -> Result<()> {
391 out.clear();
392 out.reserve(stored.len());
393 for item in stored {
394 match approximate_cosine_distance_prepared(query, item) {
395 Ok(distance) => out.push(distance),
396 Err(e) => {
397 out.clear();
398 return Err(e);
399 }
400 }
401 }
402 Ok(())
403}
404
405#[inline]
409pub fn approximate_int8_batch_prepared(
410 query: &PreparedQuery,
411 candidates: &[QuantizedVector],
412) -> Result<Vec<f32>> {
413 let PreparedQuery::Int8(q) = query else {
414 return Err(EmbedError::TierMismatch {
415 op: "approximate_int8_batch_prepared",
416 expected: QuantizationTier::Int8,
417 actual: query.tier(),
418 });
419 };
420 Ok(candidates
421 .iter()
422 .map(|candidate| 1.0 - cosine_similarity_i8_trusted(candidate, q))
423 .collect())
424}
425
426#[inline]
430pub fn approximate_int8_batch_prepared_into(
431 query: &PreparedQuery,
432 candidates: &[QuantizedVector],
433 out: &mut Vec<f32>,
434) -> Result<()> {
435 out.clear();
436 let PreparedQuery::Int8(q) = query else {
437 return Err(EmbedError::TierMismatch {
438 op: "approximate_int8_batch_prepared_into",
439 expected: QuantizationTier::Int8,
440 actual: query.tier(),
441 });
442 };
443 out.reserve(candidates.len());
444 out.extend(
445 candidates
446 .iter()
447 .map(|candidate| 1.0 - cosine_similarity_i8_trusted(candidate, q)),
448 );
449 Ok(())
450}
451
452#[inline]
456pub fn approximate_int4_batch_prepared(
457 query: &PreparedQuery,
458 candidates: &[Int4Vector],
459) -> Result<Vec<f32>> {
460 let PreparedQuery::Int4(q) = query else {
461 return Err(EmbedError::TierMismatch {
462 op: "approximate_int4_batch_prepared",
463 expected: QuantizationTier::Int4,
464 actual: query.tier(),
465 });
466 };
467 Ok(candidates
468 .iter()
469 .map(|candidate| candidate.cosine_distance(q))
470 .collect())
471}
472
473#[inline]
477pub fn approximate_int4_batch_prepared_into(
478 query: &PreparedQuery,
479 candidates: &[Int4Vector],
480 out: &mut Vec<f32>,
481) -> Result<()> {
482 out.clear();
483 let PreparedQuery::Int4(q) = query else {
484 return Err(EmbedError::TierMismatch {
485 op: "approximate_int4_batch_prepared_into",
486 expected: QuantizationTier::Int4,
487 actual: query.tier(),
488 });
489 };
490 out.reserve(candidates.len());
491 out.extend(
492 candidates
493 .iter()
494 .map(|candidate| candidate.cosine_distance(q)),
495 );
496 Ok(())
497}
498
499pub fn approximate_cosine_distance(query_f32: &[f32], stored: &QuantizedData) -> f32 {
504 debug_assert_eq!(
505 query_f32.len(),
506 stored.dims(),
507 "approximate_cosine_distance: query length {} != stored dims {}",
508 query_f32.len(),
509 stored.dims(),
510 );
511 match stored {
512 QuantizedData::Full(v) => {
513 1.0 - cosine_similarity(query_f32, v)
515 }
516 QuantizedData::Int8(q) => {
517 let query_q = QuantizedVector::from_f32(query_f32);
518 1.0 - q.cosine_similarity(&query_q)
519 }
520 QuantizedData::Int4(q) => {
521 let query_q = Int4Vector::from_f32(query_f32);
522 q.cosine_distance(&query_q)
523 }
524 QuantizedData::Binary(q) => {
525 let query_q = BinaryVector::from_f32(query_f32);
526 q.cosine_distance_approx(&query_q)
527 }
528 }
529}
530
531pub fn approximate_dot_product(query_f32: &[f32], stored: &QuantizedData) -> f32 {
533 match stored {
534 QuantizedData::Full(v) => dot_product(query_f32, v),
535 QuantizedData::Int8(q) => {
536 let query_q = QuantizedVector::from_f32(query_f32);
537 q.dot_product(&query_q)
538 }
539 QuantizedData::Int4(q) => {
540 let query_q = Int4Vector::from_f32(query_f32);
541 q.dot_product(&query_q)
542 }
543 QuantizedData::Binary(_q) => {
544 let stored_f32 = _q.to_f32();
546 dot_product(query_f32, &stored_f32)
547 }
548 }
549}
550
551#[cfg(test)]
552mod tests {
553 use super::*;
554
555 fn generate_vector(dim: usize, seed: u64) -> Vec<f32> {
556 let mut state = seed ^ ((dim as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
557 (0..dim)
558 .map(|i| {
559 state = state
560 .wrapping_mul(6364136223846793005)
561 .wrapping_add(1442695040888963407)
562 .wrapping_add(i as u64);
563 let unit = ((state >> 32) as u32) as f32 / u32::MAX as f32;
564 unit * 2.0 - 1.0
565 })
566 .collect()
567 }
568
569 fn scalar_cosine_f64(a: &[f32], b: &[f32]) -> f64 {
570 assert_eq!(a.len(), b.len());
571 let mut dot = 0.0f64;
572 let mut norm_a = 0.0f64;
573 let mut norm_b = 0.0f64;
574 for (&a, &b) in a.iter().zip(b) {
575 let a = f64::from(a);
576 let b = f64::from(b);
577 dot += a * b;
578 norm_a += a * a;
579 norm_b += b * b;
580 }
581 let denom = norm_a.sqrt() * norm_b.sqrt();
582 if denom == 0.0 { 0.0 } else { dot / denom }
583 }
584
585 fn reference_ranking(query: &[f32], corpus: &[Vec<f32>]) -> Vec<usize> {
586 let mut ranked: Vec<_> = corpus
587 .iter()
588 .enumerate()
589 .map(|(index, candidate)| (index, scalar_cosine_f64(query, candidate)))
590 .collect();
591 ranked.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
592 ranked.into_iter().map(|(index, _)| index).collect()
593 }
594
595 fn tier_ranking(query: &[f32], stored: &[QuantizedData], tier: QuantizationTier) -> Vec<usize> {
596 let prepared = PreparedQuery::from_f32(query, tier);
597 assert_eq!(prepared.tier(), tier);
598 let mut ranked: Vec<_> = stored
599 .iter()
600 .enumerate()
601 .map(|(index, candidate)| {
602 (
603 index,
604 approximate_cosine_distance_prepared(&prepared, candidate).unwrap(),
605 )
606 })
607 .collect();
608 ranked.sort_unstable_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
609 ranked.into_iter().map(|(index, _)| index).collect()
610 }
611
612 fn recall_hits_at(reference: &[usize], actual: &[usize], k: usize) -> usize {
613 actual[..k]
614 .iter()
615 .filter(|candidate| reference[..k].contains(candidate))
616 .count()
617 }
618
619 fn recall_at(reference: &[usize], actual: &[usize], k: usize) -> f64 {
620 recall_hits_at(reference, actual, k) as f64 / k as f64
621 }
622
623 fn pairwise_ranking_agreements(reference: &[usize], actual: &[usize]) -> usize {
624 assert_eq!(reference.len(), actual.len());
625 let mut actual_position = vec![0usize; actual.len()];
626 for (position, &candidate) in actual.iter().enumerate() {
627 actual_position[candidate] = position;
628 }
629 let mut agreements = 0usize;
630 for (position, &left) in reference.iter().enumerate() {
631 for &right in &reference[position + 1..] {
632 agreements += usize::from(actual_position[left] < actual_position[right]);
633 }
634 }
635 agreements
636 }
637
638 fn pairwise_ranking_agreement(reference: &[usize], actual: &[usize]) -> f64 {
639 let pairs = reference.len() * (reference.len() - 1) / 2;
640 pairwise_ranking_agreements(reference, actual) as f64 / pairs as f64
641 }
642
643 fn retrieval_quality_counts(
644 reference: &[usize],
645 actual: &[usize],
646 top_k: usize,
647 ) -> (usize, usize) {
648 (
649 recall_hits_at(reference, actual, top_k),
650 pairwise_ranking_agreements(reference, actual),
651 )
652 }
653
654 fn index_order_surrogate_quality(
655 corpus: &[Vec<f32>],
656 queries: &[Vec<f32>],
657 top_k: usize,
658 ) -> (f64, f64) {
659 let index_order: Vec<_> = (0..corpus.len()).collect();
660 let mut recall = 0.0;
661 let mut agreement = 0.0;
662 for query in queries {
663 let reference = reference_ranking(query, corpus);
664 recall += recall_at(&reference, &index_order, top_k);
665 agreement += pairwise_ranking_agreement(&reference, &index_order);
666 }
667 (
668 recall / queries.len() as f64,
669 agreement / queries.len() as f64,
670 )
671 }
672
673 fn retrieval_quality_floor(tier: QuantizationTier) -> (f64, f64) {
674 match tier {
675 QuantizationTier::Full => (1.0, 0.999),
676 QuantizationTier::Int8 => (0.98, 0.995),
677 QuantizationTier::Int4 => (0.85, 0.95),
678 QuantizationTier::Binary => (0.30, 0.70),
679 }
680 }
681
682 fn meets_retrieval_quality_floor(value: f64, minimum: f64) -> bool {
684 value.is_finite() && value + f64::EPSILON >= minimum
685 }
686
687 fn retrieval_quality_per_query_floor(tier: QuantizationTier) -> (usize, usize) {
704 match tier {
705 QuantizationTier::Full => (9, 32_640),
706 QuantizationTier::Int8 => (9, 32_530),
707 QuantizationTier::Int4 => (7, 31_289),
708 QuantizationTier::Binary => (2, 23_141),
709 }
710 }
711
712 const HEALTHY_FULL_QUERY_QUALITY: [(usize, usize); 16] = [(10, 32_640); 16];
723 const HEALTHY_INT8_QUERY_QUALITY: [(usize, usize); 16] = [
724 (10, 32_588),
725 (10, 32_575),
726 (10, 32_597),
727 (10, 32_590),
728 (10, 32_595),
729 (10, 32_592),
730 (10, 32_580),
731 (10, 32_604),
732 (10, 32_587),
733 (10, 32_573),
734 (10, 32_580),
735 (10, 32_579),
736 (10, 32_567),
737 (10, 32_578),
738 (10, 32_576),
739 (10, 32_588),
740 ];
741 const HEALTHY_INT4_QUERY_QUALITY: [(usize, usize); 16] = [
742 (8, 31_582),
743 (9, 31_531),
744 (9, 31_642),
745 (8, 31_656),
746 (10, 31_773),
747 (8, 31_638),
748 (10, 31_744),
749 (9, 31_758),
750 (10, 31_661),
751 (10, 31_651),
752 (10, 31_662),
753 (8, 31_727),
754 (10, 31_531),
755 (10, 31_625),
756 (9, 31_655),
757 (9, 31_653),
758 ];
759 const HEALTHY_BINARY_QUERY_QUALITY: [(usize, usize); 16] = [
760 (7, 25_430),
761 (3, 25_031),
762 (6, 25_033),
763 (5, 25_339),
764 (4, 25_534),
765 (3, 25_166),
766 (3, 25_156),
767 (6, 25_705),
768 (5, 25_677),
769 (4, 25_419),
770 (4, 25_040),
771 (3, 25_150),
772 (4, 24_965),
773 (5, 25_190),
774 (5, 24_423),
775 (4, 25_355),
776 ];
777
778 fn healthy_query_quality(tier: QuantizationTier) -> &'static [(usize, usize); 16] {
779 match tier {
780 QuantizationTier::Full => &HEALTHY_FULL_QUERY_QUALITY,
781 QuantizationTier::Int8 => &HEALTHY_INT8_QUERY_QUALITY,
782 QuantizationTier::Int4 => &HEALTHY_INT4_QUERY_QUALITY,
783 QuantizationTier::Binary => &HEALTHY_BINARY_QUERY_QUALITY,
784 }
785 }
786
787 fn retrieval_quality_movement_budget(healthy: &[(usize, usize)]) -> (usize, usize) {
788 let minimum_recall_hits = healthy.iter().map(|quality| quality.0).min().unwrap();
789 let maximum_recall_hits = healthy.iter().map(|quality| quality.0).max().unwrap();
790 let minimum_agreements = healthy.iter().map(|quality| quality.1).min().unwrap();
791 let maximum_agreements = healthy.iter().map(|quality| quality.1).max().unwrap();
792 (
793 maximum_recall_hits - minimum_recall_hits,
794 maximum_agreements - minimum_agreements,
795 )
796 }
797
798 fn retrieval_quality_concentration_budget(healthy: &[(usize, usize)]) -> (usize, usize) {
805 let movement_budget = retrieval_quality_movement_budget(healthy);
806 (
807 movement_budget.0.div_ceil(healthy.len()),
808 movement_budget.1.div_ceil(healthy.len()),
809 )
810 }
811
812 fn validate_tier_retrieval_quality(
813 tier: QuantizationTier,
814 query_quality: &[(usize, usize)],
815 top_k: usize,
816 ) -> std::result::Result<(f64, f64), String> {
817 if query_quality.is_empty() {
818 return Err(format!("{tier:?} retrieval quality has zero queries"));
819 }
820
821 let healthy = healthy_query_quality(tier);
822 if query_quality.len() != healthy.len() {
823 return Err(format!(
824 "{tier:?} retrieval quality has {} queries, expected {}",
825 query_quality.len(),
826 healthy.len()
827 ));
828 }
829
830 let (minimum_query_recall, minimum_query_agreement) =
831 retrieval_quality_per_query_floor(tier);
832 for (query_index, &(recall_hits, agreements)) in query_quality.iter().enumerate() {
833 if recall_hits < minimum_query_recall || agreements < minimum_query_agreement {
834 return Err(format!(
835 "{tier:?} query {query_index} fails the per-query floor: Recall@{top_k}=\
836 {:.6} (minimum {:.6}), pairwise ranking agreement=\
837 {:.6} (minimum {:.6})",
838 recall_hits as f64 / top_k as f64,
839 minimum_query_recall as f64 / top_k as f64,
840 agreements as f64 / 32_640.0,
841 minimum_query_agreement as f64 / 32_640.0,
842 ));
843 }
844 }
845
846 let recall_hits = query_quality.iter().map(|quality| quality.0).sum::<usize>();
847 let agreements = query_quality.iter().map(|quality| quality.1).sum::<usize>();
848 let recall = recall_hits as f64 / (query_quality.len() * top_k) as f64;
849 let agreement = agreements as f64 / (query_quality.len() * 32_640) as f64;
850 let (minimum_recall, minimum_agreement) = retrieval_quality_floor(tier);
851 if !meets_retrieval_quality_floor(recall, minimum_recall) {
852 return Err(format!(
853 "{tier:?} Recall@{top_k} {recall:.6} is below the measured-data floor \
854 {minimum_recall:.3}"
855 ));
856 }
857 if !meets_retrieval_quality_floor(agreement, minimum_agreement) {
858 return Err(format!(
859 "{tier:?} pairwise ranking agreement {agreement:.6} is below the measured-data \
860 floor {minimum_agreement:.3}"
861 ));
862 }
863
864 let (recall_movement, agreement_movement) = query_quality.iter().zip(healthy).fold(
865 (0usize, 0usize),
866 |(recall_movement, agreement_movement), (actual, expected)| {
867 (
868 recall_movement + actual.0.abs_diff(expected.0),
869 agreement_movement + actual.1.abs_diff(expected.1),
870 )
871 },
872 );
873 let (recall_budget, agreement_budget) = retrieval_quality_movement_budget(healthy);
874 if recall_movement > recall_budget || agreement_movement > agreement_budget {
875 return Err(format!(
876 "{tier:?} retrieval quality exceeds the fixture-relative movement budget: \
877 total absolute Recall@{top_k} movement={recall_movement} hit(s) \
878 (maximum {recall_budget}), total absolute pairwise-agreement movement=\
879 {agreement_movement} pair(s) (maximum {agreement_budget})"
880 ));
881 }
882
883 let (maximum_query_recall_movement, maximum_query_agreement_movement) =
884 retrieval_quality_concentration_budget(healthy);
885 for (query_index, (&(recall_hits, agreements), &(healthy_hits, healthy_agreements))) in
886 query_quality.iter().zip(healthy).enumerate()
887 {
888 let recall_movement = recall_hits.abs_diff(healthy_hits);
889 let agreement_movement = agreements.abs_diff(healthy_agreements);
890 if recall_movement > maximum_query_recall_movement
891 || agreement_movement > maximum_query_agreement_movement
892 {
893 return Err(format!(
894 "{tier:?} query {query_index} exceeds the fixture-relative concentration \
895 bound: Recall@{top_k} movement={recall_movement} hit(s) \
896 (maximum {maximum_query_recall_movement}), pairwise-agreement \
897 movement={agreement_movement} pair(s) \
898 (maximum {maximum_query_agreement_movement})"
899 ));
900 }
901 }
902 Ok((recall, agreement))
903 }
904
905 fn validate_retrieval_fixture(
913 corpus: &[Vec<f32>],
914 queries: &[Vec<f32>],
915 top_k: usize,
916 ) -> std::result::Result<(), String> {
917 if top_k == 0 {
918 return Err("retrieval fixture has top_k=0".to_string());
919 }
920 if queries.is_empty() {
921 return Err("retrieval fixture has zero queries".to_string());
922 }
923 if corpus.len() < top_k {
924 return Err(format!(
925 "retrieval fixture corpus size {} is smaller than top_k={top_k}",
926 corpus.len()
927 ));
928 }
929 for (query_index, query) in queries.iter().enumerate() {
930 let mut ranked_scores = Vec::with_capacity(corpus.len());
931 for (candidate_index, candidate) in corpus.iter().enumerate() {
932 let score = scalar_cosine_f64(query, candidate);
933 if !score.is_finite() {
934 return Err(format!(
935 "retrieval fixture query {query_index} candidate {candidate_index} has \
936 non-finite reference score {score}"
937 ));
938 }
939 ranked_scores.push((candidate_index, score));
940 }
941 ranked_scores.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
942
943 let mut distinct_scores: Vec<_> =
944 ranked_scores.iter().map(|(_, score)| *score).collect();
945 distinct_scores.sort_unstable_by(f64::total_cmp);
946 distinct_scores.dedup();
947 if distinct_scores.len() <= top_k {
948 return Err(format!(
949 "retrieval fixture query {query_index} is non-discriminating: only \
950 {} distinct finite reference score(s) across {} candidates, need \
951 more than top_k={top_k}",
952 distinct_scores.len(),
953 corpus.len()
954 ));
955 }
956
957 for (rank, boundary) in ranked_scores.windows(2).enumerate() {
958 let higher_score = boundary[0].1;
959 let lower_score = boundary[1].1;
960 if higher_score as f32 <= lower_score as f32 {
961 return Err(format!(
962 "retrieval fixture query {query_index} has a near-tied ranking boundary \
963 at ranks {rank} and {}: reference scores {higher_score} and \
964 {lower_score} do not remain ordered at f32 precision",
965 rank + 1
966 ));
967 }
968 }
969 }
970
971 let (surrogate_recall, surrogate_agreement) =
972 index_order_surrogate_quality(corpus, queries, top_k);
973 for tier in [
974 QuantizationTier::Full,
975 QuantizationTier::Int8,
976 QuantizationTier::Int4,
977 QuantizationTier::Binary,
978 ] {
979 let (minimum_recall, minimum_agreement) = retrieval_quality_floor(tier);
980 if meets_retrieval_quality_floor(surrogate_recall, minimum_recall)
981 && meets_retrieval_quality_floor(surrogate_agreement, minimum_agreement)
982 {
983 return Err(format!(
984 "retrieval fixture is non-discriminating: an all-tied index-order surrogate \
985 passes the {tier:?} floor with Recall@{top_k}={surrogate_recall:.6} and \
986 pairwise ranking agreement={surrogate_agreement:.6}"
987 ));
988 }
989 }
990 Ok(())
991 }
992
993 fn fixed_retrieval_fixture() -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
994 const DIMS: usize = 384;
995 const CORPUS_SIZE: usize = 256;
996 const QUERY_COUNT: usize = 16;
997 let corpus = (0..CORPUS_SIZE)
998 .map(|index| generate_vector(DIMS, 0xC0A5_0000 + index as u64))
999 .collect();
1000 let queries = (0..QUERY_COUNT)
1001 .map(|index| generate_vector(DIMS, 0x0A11_0000 + index as u64))
1002 .collect();
1003 (corpus, queries)
1004 }
1005
1006 #[test]
1007 fn test_tier_retrieval_quality_against_independent_f64_ranking() {
1008 const TOP_K: usize = 10;
1009 let (corpus, queries) = fixed_retrieval_fixture();
1010 validate_retrieval_fixture(&corpus, &queries, TOP_K)
1011 .expect("retrieval fixture must be discriminating before scoring tiers against it");
1012
1013 for tier in [
1014 QuantizationTier::Full,
1015 QuantizationTier::Int8,
1016 QuantizationTier::Int4,
1017 QuantizationTier::Binary,
1018 ] {
1019 let stored: Vec<_> = corpus
1020 .iter()
1021 .map(|candidate| QuantizedData::from_f32(candidate, tier))
1022 .collect();
1023 assert!(
1024 stored.iter().all(|candidate| candidate.tier() == tier),
1025 "{tier:?} conversion was bypassed or routed to another tier"
1026 );
1027 assert!(
1028 stored
1029 .iter()
1030 .all(|candidate| candidate.storage_bytes()
1031 == tier.storage_bytes(candidate.dims())),
1032 "{tier:?} conversion produced the wrong representation size"
1033 );
1034
1035 if tier != QuantizationTier::Full {
1036 assert!(
1037 stored.iter().zip(&corpus).any(|(quantized, original)| {
1038 quantized
1039 .to_f32()
1040 .iter()
1041 .zip(original)
1042 .any(|(actual, expected)| (actual - expected).abs() > 1e-4)
1043 }),
1044 "{tier:?} conversion did not exercise a lossy representation"
1045 );
1046 }
1047
1048 let mut query_quality = Vec::with_capacity(queries.len());
1049 let mut quantized_distance_witness = false;
1050 for query in &queries {
1051 let reference = reference_ranking(query, &corpus);
1052 let actual = tier_ranking(query, &stored, tier);
1053 query_quality.push(retrieval_quality_counts(&reference, &actual, TOP_K));
1054
1055 if tier != QuantizationTier::Full {
1056 let prepared = PreparedQuery::from_f32(query, tier);
1057 quantized_distance_witness |=
1058 stored.iter().zip(&corpus).any(|(quantized, original)| {
1059 let actual =
1060 approximate_cosine_distance_prepared(&prepared, quantized).unwrap();
1061 let reference = 1.0 - scalar_cosine_f64(query, original) as f32;
1062 (actual - reference).abs() > 1e-4
1063 });
1064 }
1065 }
1066 let (recall, agreement) = validate_tier_retrieval_quality(tier, &query_quality, TOP_K)
1067 .unwrap_or_else(|error| panic!("{error}"));
1068 eprintln!(
1069 "{tier:?}: Recall@{TOP_K}={recall:.6}, pairwise ranking agreement={agreement:.6}"
1070 );
1071
1072 if tier != QuantizationTier::Full {
1073 assert!(
1074 quantized_distance_witness,
1075 "{tier:?} distance path did not differ from the independent f32 reference"
1076 );
1077 }
1078 }
1079 }
1080
1081 #[test]
1082 fn test_tier_retrieval_quality_rejects_concentrated_binary_query_collapse() {
1083 const TOP_K: usize = 10;
1084 const COLLAPSED_QUERY_COUNT: usize = 10;
1085 let (corpus, queries) = fixed_retrieval_fixture();
1086 let index_order: Vec<_> = (0..corpus.len()).collect();
1087 let query_quality: Vec<_> = queries
1088 .iter()
1089 .enumerate()
1090 .map(|(query_index, query)| {
1091 let reference = reference_ranking(query, &corpus);
1092 let actual = if query_index < COLLAPSED_QUERY_COUNT {
1093 index_order.clone()
1094 } else {
1095 reference.clone()
1096 };
1097 retrieval_quality_counts(&reference, &actual, TOP_K)
1098 })
1099 .collect();
1100
1101 let mean_recall = query_quality.iter().map(|quality| quality.0).sum::<usize>() as f64
1102 / (query_quality.len() * TOP_K) as f64;
1103 let mean_agreement = query_quality.iter().map(|quality| quality.1).sum::<usize>() as f64
1104 / (query_quality.len() * 32_640) as f64;
1105 assert_eq!(format!("{mean_recall:.6}"), "0.381250");
1106 assert_eq!(format!("{mean_agreement:.6}"), "0.705356");
1107
1108 let error = validate_tier_retrieval_quality(
1109 QuantizationTier::Binary,
1110 &query_quality,
1111 TOP_K,
1112 )
1113 .expect_err(
1114 "collapsing fixed fixture queries 0 through 9 must fail despite passing both means",
1115 );
1116 assert!(error.contains("query 0"), "unexpected error: {error}");
1117 }
1118
1119 #[test]
1120 fn test_tier_retrieval_quality_rejects_single_binary_query_concentration() {
1121 const TOP_K: usize = 10;
1122 const SUFFIX_INVERSIONS: usize = 7_161;
1123 let (corpus, queries) = fixed_retrieval_fixture();
1124 let reference = reference_ranking(&queries[0], &corpus);
1125 let mut remaining = reference[3..10]
1126 .iter()
1127 .chain(&reference[17..])
1128 .copied()
1129 .collect::<Vec<_>>();
1130 let mut suffix = Vec::with_capacity(remaining.len());
1131 let mut inversions = SUFFIX_INVERSIONS;
1132 while !remaining.is_empty() {
1133 let index = inversions.min(remaining.len() - 1);
1134 inversions -= index;
1135 suffix.push(remaining.remove(index));
1136 }
1137 assert_eq!(inversions, 0);
1138
1139 let mut actual = Vec::with_capacity(reference.len());
1140 actual.extend_from_slice(&reference[..3]);
1141 actual.extend_from_slice(&reference[10..17]);
1142 actual.extend(suffix);
1143 assert_eq!(actual.len(), reference.len());
1144
1145 let collapsed_quality = retrieval_quality_counts(&reference, &actual, TOP_K);
1146 assert_eq!(collapsed_quality, (3, 25_430));
1147
1148 let mut query_quality = healthy_query_quality(QuantizationTier::Binary).to_vec();
1149 query_quality[0] = collapsed_quality;
1150
1151 let error =
1152 validate_tier_retrieval_quality(QuantizationTier::Binary, &query_quality, TOP_K)
1153 .expect_err("one Binary query must not spend the complete fixture movement budget");
1154 assert!(
1155 error.contains("query 0") && error.contains("concentration"),
1156 "unexpected error: {error}"
1157 );
1158 }
1159
1160 #[test]
1161 fn test_tier_retrieval_quality_rejects_single_binary_query_agreement_concentration() {
1162 const TOP_K: usize = 10;
1163 let mut query_quality = healthy_query_quality(QuantizationTier::Binary).to_vec();
1164 query_quality[0].1 -= 1_281;
1165
1166 let error =
1167 validate_tier_retrieval_quality(QuantizationTier::Binary, &query_quality, TOP_K)
1168 .expect_err("one Binary query must not spend nearly the complete pair budget");
1169 assert!(
1170 error.contains("query 0") && error.contains("concentration"),
1171 "unexpected error: {error}"
1172 );
1173 }
1174
1175 #[test]
1176 fn test_tier_retrieval_quality_rejects_distributed_binary_query_collapse() {
1177 const TOP_K: usize = 10;
1178 let (corpus, queries) = fixed_retrieval_fixture();
1179 let query_quality: Vec<_> = queries
1180 .iter()
1181 .map(|query| {
1182 let reference = reference_ranking(query, &corpus);
1183 let mut actual = Vec::with_capacity(reference.len());
1184 actual.extend_from_slice(&reference[..3]);
1185 actual.extend_from_slice(&reference[10..17]);
1186 actual.extend_from_slice(&reference[3..10]);
1187 actual.extend(reference[17..154].iter().rev().copied());
1188 actual.extend_from_slice(&reference[154..]);
1189 assert_eq!(actual.len(), reference.len());
1190 retrieval_quality_counts(&reference, &actual, TOP_K)
1191 })
1192 .collect();
1193
1194 assert!(
1195 query_quality
1196 .iter()
1197 .all(|&(recall, agreement)| { recall == 3 && agreement == 32_640 - 9_365 })
1198 );
1199
1200 let error =
1201 validate_tier_retrieval_quality(QuantizationTier::Binary, &query_quality, TOP_K)
1202 .expect_err("a distributed 70% Recall@10 loss across every query must fail");
1203 assert!(
1204 error.contains("movement budget"),
1205 "unexpected error: {error}"
1206 );
1207 }
1208
1209 #[test]
1210 fn test_tier_retrieval_quality_rejects_shallow_all_query_movement() {
1211 const TOP_K: usize = 10;
1212 let (corpus, queries) = fixed_retrieval_fixture();
1213 let query_quality: Vec<_> = queries
1214 .iter()
1215 .map(|query| {
1216 let reference = reference_ranking(query, &corpus);
1217 let mut actual = reference.clone();
1218 let last = actual.len() - 1;
1219 actual.swap(TOP_K - 1, last);
1220 retrieval_quality_counts(&reference, &actual, TOP_K)
1221 })
1222 .collect();
1223
1224 assert!(
1225 query_quality
1226 .iter()
1227 .all(|&(recall, agreement)| { recall == 9 && agreement == 32_640 - 491 })
1228 );
1229
1230 for (tier, result) in [
1231 (
1232 QuantizationTier::Int4,
1233 validate_tier_retrieval_quality(QuantizationTier::Int4, &query_quality, TOP_K),
1234 ),
1235 (
1236 QuantizationTier::Binary,
1237 validate_tier_retrieval_quality(QuantizationTier::Binary, &query_quality, TOP_K),
1238 ),
1239 ] {
1240 let error = result.unwrap_err();
1241 assert!(
1242 error.contains("movement budget"),
1243 "unexpected {tier:?} error: {error}"
1244 );
1245 }
1246 }
1247
1248 #[test]
1249 fn test_tier_retrieval_quality_bounds_stated_binary_uniform_movement() {
1250 const TOP_K: usize = 10;
1251 let query_quality_with_pair_loss = |pair_loss: u16| {
1252 healthy_query_quality(QuantizationTier::Binary)
1253 .iter()
1254 .map(|&(recall_hits, agreements)| {
1255 (recall_hits, agreements - usize::from(pair_loss))
1256 })
1257 .collect::<Vec<_>>()
1258 };
1259
1260 validate_tier_retrieval_quality(
1261 QuantizationTier::Binary,
1262 &query_quality_with_pair_loss(80),
1263 TOP_K,
1264 )
1265 .expect("1,280 total agreement-pair changes are within the 1,282-pair budget");
1266
1267 let error = validate_tier_retrieval_quality(
1268 QuantizationTier::Binary,
1269 &query_quality_with_pair_loss(81),
1270 TOP_K,
1271 )
1272 .expect_err("1,296 total agreement-pair changes must exceed the 1,282-pair budget");
1273 assert!(
1274 error.contains("movement budget"),
1275 "unexpected error: {error}"
1276 );
1277 }
1278
1279 #[test]
1280 fn test_tier_retrieval_quality_accepts_non_uniform_exact_binary_movement_boundary() {
1281 const TOP_K: usize = 10;
1282 let query_quality = healthy_query_quality(QuantizationTier::Binary)
1283 .iter()
1284 .enumerate()
1285 .map(|(query_index, &(recall_hits, agreements))| {
1286 let pair_loss = if query_index < 14 { 81 } else { 74 };
1287 (recall_hits, agreements - pair_loss)
1288 })
1289 .collect::<Vec<_>>();
1290
1291 validate_tier_retrieval_quality(QuantizationTier::Binary, &query_quality, TOP_K)
1292 .expect("the exact 1,282-pair non-uniform movement boundary must be accepted");
1293 }
1294
1295 #[test]
1296 fn test_validate_retrieval_fixture_rejects_non_discriminating_corpus() {
1297 const TOP_K: usize = 10;
1298 const DIMS: usize = 384;
1299 const CORPUS_SIZE: usize = 256;
1300 const QUERY_COUNT: usize = 16;
1301
1302 let repeated_vector = generate_vector(DIMS, 0xC0A5_0000);
1307 let corpus: Vec<Vec<f32>> = std::iter::repeat_n(repeated_vector, CORPUS_SIZE).collect();
1308 let queries: Vec<Vec<f32>> = (0..QUERY_COUNT)
1309 .map(|index| generate_vector(DIMS, 0x0A11_0000 + index as u64))
1310 .collect();
1311
1312 let err = validate_retrieval_fixture(&corpus, &queries, TOP_K)
1313 .expect_err("a corpus of identical vectors ties every reference score; the guard must refuse rather than let it be scored");
1314 eprintln!("guard refused as expected: {err}");
1315 }
1316
1317 #[test]
1318 fn test_validate_retrieval_fixture_rejects_zero_queries() {
1319 let (corpus, _queries) = fixed_retrieval_fixture();
1320 let err = validate_retrieval_fixture(&corpus, &[], 10)
1321 .expect_err("zero queries must be refused rather than panic or silently pass");
1322 eprintln!("guard refused as expected: {err}");
1323 }
1324
1325 #[test]
1326 fn test_validate_retrieval_fixture_rejects_corpus_smaller_than_top_k() {
1327 let (corpus, queries) = fixed_retrieval_fixture();
1328 let small_corpus = corpus[..5].to_vec();
1329 let err = validate_retrieval_fixture(&small_corpus, &queries, 10).expect_err(
1330 "a corpus smaller than top_k must be refused rather than panic or return NaN",
1331 );
1332 eprintln!("guard refused as expected: {err}");
1333 }
1334
1335 #[test]
1336 fn test_validate_retrieval_fixture_accepts_the_real_fixture() {
1337 let (corpus, queries) = fixed_retrieval_fixture();
1338 validate_retrieval_fixture(&corpus, &queries, 10)
1339 .expect("the real fixture is discriminating and must not be rejected");
1340
1341 let (recall, agreement) = index_order_surrogate_quality(&corpus, &queries, 10);
1342 assert!((recall - 0.025).abs() < f64::EPSILON);
1343 assert!((agreement - 0.526_646_752_450_980_4).abs() < 1e-15);
1344 for (tier, minimum_recall, minimum_agreement) in [
1345 ("Full", 1.0, 0.999),
1346 ("Int8", 0.98, 0.995),
1347 ("Int4", 0.85, 0.95),
1348 ("Binary", 0.30, 0.70),
1349 ] {
1350 assert!(
1351 !meets_retrieval_quality_floor(recall, minimum_recall)
1352 || !meets_retrieval_quality_floor(agreement, minimum_agreement),
1353 "index-order surrogate must fail the pinned {tier} floor"
1354 );
1355 }
1356 }
1357
1358 #[test]
1359 fn test_validate_retrieval_fixture_rejects_binary_index_aligned_collapse() {
1360 const TOP_K: usize = 10;
1361 let query = vec![1.0, 0.0];
1362 let corpus: Vec<_> = (0..21)
1363 .map(|index| vec![21.0 - index as f32, 1.0])
1364 .collect();
1365 let queries = vec![query.clone()];
1366
1367 let mut scores: Vec<_> = corpus
1368 .iter()
1369 .map(|candidate| scalar_cosine_f64(&query, candidate))
1370 .collect();
1371 scores.sort_unstable_by(f64::total_cmp);
1372 scores.dedup();
1373 assert_eq!(scores.len(), 21);
1374
1375 let reference = reference_ranking(&query, &corpus);
1376 let binary: Vec<_> = corpus
1377 .iter()
1378 .map(|candidate| QuantizedData::from_f32(candidate, QuantizationTier::Binary))
1379 .collect();
1380 let prepared = PreparedQuery::from_f32(&query, QuantizationTier::Binary);
1381 let first_code = match &binary[0] {
1382 QuantizedData::Binary(value) => &value.data,
1383 _ => unreachable!("binary conversion must produce the Binary variant"),
1384 };
1385 assert!(binary.iter().all(|candidate| {
1386 let QuantizedData::Binary(value) = candidate else {
1387 return false;
1388 };
1389 value.data == *first_code
1390 && approximate_cosine_distance_prepared(&prepared, candidate).unwrap() == 0.0
1391 }));
1392 let actual = tier_ranking(&query, &binary, QuantizationTier::Binary);
1393 assert_eq!(reference, (0..21).collect::<Vec<_>>());
1394 assert_eq!(actual, reference);
1395 assert_eq!(recall_at(&reference, &actual, TOP_K), 1.0);
1396 assert_eq!(pairwise_ranking_agreement(&reference, &actual), 1.0);
1397
1398 let err = validate_retrieval_fixture(&corpus, &queries, TOP_K).expect_err(
1399 "an index-aligned reference must not let a totally collapsed tier pass its floors",
1400 );
1401 assert!(
1402 err.contains("index-order surrogate"),
1403 "unexpected error: {err}"
1404 );
1405 }
1406
1407 #[test]
1408 fn test_validate_retrieval_fixture_rejects_mixed_non_finite_scores() {
1409 for non_finite in [f32::NAN, f32::INFINITY] {
1410 let (mut corpus, queries) = fixed_retrieval_fixture();
1411 corpus[0][0] = non_finite;
1412 let err = validate_retrieval_fixture(&corpus, &queries, 10)
1413 .expect_err("the first non-finite reference score must invalidate the oracle");
1414 assert!(
1415 err.contains("query 0 candidate 0"),
1416 "unexpected error: {err}"
1417 );
1418 assert!(err.contains("non-finite"), "unexpected error: {err}");
1419 }
1420 }
1421
1422 #[test]
1423 fn test_validate_retrieval_fixture_rejects_zero_top_k() {
1424 let (corpus, queries) = fixed_retrieval_fixture();
1425 let err = validate_retrieval_fixture(&corpus, &queries, 0)
1426 .expect_err("top_k=0 must be refused before recall divides by zero");
1427 assert!(err.contains("top_k=0"), "unexpected error: {err}");
1428 }
1429
1430 #[test]
1431 fn test_validate_retrieval_fixture_rejects_near_tied_top_k_boundary() {
1432 let query = vec![1.0, 0.0];
1433 let corpus = vec![
1434 vec![0.0, 1.0],
1435 vec![1.0, 0.0],
1436 vec![1.0, 0.0001],
1437 vec![-1.0, 0.0],
1438 ];
1439 let best = scalar_cosine_f64(&query, &corpus[1]);
1440 let runner_up = scalar_cosine_f64(&query, &corpus[2]);
1441 let boundary_gap = best - runner_up;
1442 assert!(boundary_gap > 0.0 && boundary_gap < 1e-6);
1443
1444 let err = validate_retrieval_fixture(&corpus, &[query], 1)
1445 .expect_err("an f64-only near-tie at the evaluated cutoff must be refused");
1446 assert!(err.contains("near-tied"), "unexpected error: {err}");
1447 }
1448
1449 #[test]
1450 fn test_validate_retrieval_fixture_rejects_near_tied_pairwise_boundary() {
1451 let query = vec![1.0, 0.0];
1452 let corpus = vec![
1453 vec![-1.0, 0.0],
1454 vec![1.0, 0.0],
1455 vec![100.0, 1.0],
1456 vec![100.0, 1.0001],
1457 ];
1458 let second = scalar_cosine_f64(&query, &corpus[2]);
1459 let third = scalar_cosine_f64(&query, &corpus[3]);
1460 assert!(second > third);
1461 assert_eq!(second as f32, third as f32);
1462
1463 let err = validate_retrieval_fixture(&corpus, &[query], 1)
1464 .expect_err("an f64-only near-tie evaluated by pairwise agreement must be refused");
1465 assert!(err.contains("near-tied"), "unexpected error: {err}");
1466 }
1467
1468 #[test]
1469 fn test_tier_bytes_per_dim() {
1470 assert_eq!(QuantizationTier::Full.bytes_per_dim(), 4.0);
1471 assert_eq!(QuantizationTier::Int8.bytes_per_dim(), 1.0);
1472 assert_eq!(QuantizationTier::Int4.bytes_per_dim(), 0.5);
1473 assert_eq!(QuantizationTier::Binary.bytes_per_dim(), 0.125);
1474 }
1475
1476 #[test]
1477 fn test_tier_compression_ratios() {
1478 assert_eq!(QuantizationTier::Full.compression_ratio(), 1.0);
1479 assert_eq!(QuantizationTier::Int8.compression_ratio(), 4.0);
1480 assert_eq!(QuantizationTier::Int4.compression_ratio(), 8.0);
1481 assert_eq!(QuantizationTier::Binary.compression_ratio(), 32.0);
1482 }
1483
1484 #[test]
1485 fn test_tier_storage_bytes() {
1486 assert_eq!(QuantizationTier::Full.storage_bytes(384), 1536);
1487 assert_eq!(QuantizationTier::Int8.storage_bytes(384), 384);
1488 assert_eq!(QuantizationTier::Int4.storage_bytes(384), 192);
1489 assert_eq!(QuantizationTier::Binary.storage_bytes(384), 48);
1490 }
1491
1492 #[test]
1493 fn test_tier_from_age() {
1494 assert_eq!(
1495 QuantizationTier::from_age_seconds(0),
1496 QuantizationTier::Full
1497 );
1498 assert_eq!(
1499 QuantizationTier::from_age_seconds(1800),
1500 QuantizationTier::Full
1501 ); assert_eq!(
1503 QuantizationTier::from_age_seconds(7200),
1504 QuantizationTier::Int8
1505 ); assert_eq!(
1507 QuantizationTier::from_age_seconds(172800),
1508 QuantizationTier::Int4
1509 ); assert_eq!(
1511 QuantizationTier::from_age_seconds(1_000_000),
1512 QuantizationTier::Binary
1513 ); }
1515
1516 #[test]
1517 fn test_quantized_data_from_f32_all_tiers() {
1518 let v = generate_vector(384, 42);
1519
1520 for tier in [
1521 QuantizationTier::Full,
1522 QuantizationTier::Int8,
1523 QuantizationTier::Int4,
1524 QuantizationTier::Binary,
1525 ] {
1526 let data = QuantizedData::from_f32(&v, tier);
1527 assert_eq!(data.tier(), tier, "tier mismatch for {tier:?}");
1528 assert_eq!(data.dims(), 384, "dims mismatch for {tier:?}");
1529
1530 let expected_bytes = tier.storage_bytes(384);
1532 assert_eq!(
1533 data.storage_bytes(),
1534 expected_bytes,
1535 "storage bytes mismatch for {tier:?}"
1536 );
1537 }
1538 }
1539
1540 #[test]
1541 fn test_approximate_cosine_distance_ordering() {
1542 let a = generate_vector(384, 1);
1544 let b: Vec<f32> = a
1546 .iter()
1547 .enumerate()
1548 .map(|(i, &x)| x + 0.05 * (i as f32 * 0.3).sin())
1549 .collect();
1550 let c = generate_vector(384, 999);
1552
1553 for tier in [
1554 QuantizationTier::Full,
1555 QuantizationTier::Int8,
1556 QuantizationTier::Int4,
1557 QuantizationTier::Binary,
1558 ] {
1559 let stored_b = QuantizedData::from_f32(&b, tier);
1560 let stored_c = QuantizedData::from_f32(&c, tier);
1561
1562 let dist_ab = approximate_cosine_distance(&a, &stored_b);
1563 let dist_ac = approximate_cosine_distance(&a, &stored_c);
1564
1565 assert!(
1567 dist_ab < dist_ac,
1568 "{tier:?}: dist(a,b)={dist_ab} should be < dist(a,c)={dist_ac}"
1569 );
1570 }
1571 }
1572
1573 #[test]
1574 fn test_promote_demote_roundtrip() {
1575 let v = generate_vector(384, 42);
1576 let binary = QuantizedData::from_f32(&v, QuantizationTier::Binary);
1577
1578 let int4 = binary.promote(QuantizationTier::Int4);
1580 assert_eq!(int4.tier(), QuantizationTier::Int4);
1581
1582 let int8 = int4.promote(QuantizationTier::Int8);
1583 assert_eq!(int8.tier(), QuantizationTier::Int8);
1584
1585 let full = int8.promote(QuantizationTier::Full);
1586 assert_eq!(full.tier(), QuantizationTier::Full);
1587 assert_eq!(full.dims(), 384);
1588 }
1589
1590 #[test]
1591 fn test_int8_batch_prepared_matches_per_item_prepared() {
1592 let query = generate_vector(384, 42);
1593 let prepared = PreparedQuery::from_f32(&query, QuantizationTier::Int8);
1594 let candidates: Vec<QuantizedVector> = (0..32)
1595 .map(|i| QuantizedVector::from_f32(&generate_vector(384, i + 1)))
1596 .collect();
1597 let wrapped: Vec<QuantizedData> = candidates
1598 .iter()
1599 .cloned()
1600 .map(QuantizedData::Int8)
1601 .collect();
1602
1603 let got = approximate_int8_batch_prepared(&prepared, &candidates).unwrap();
1604 for (i, item) in wrapped.iter().enumerate() {
1605 let expected = approximate_cosine_distance_prepared(&prepared, item).unwrap();
1606 assert!(
1607 (got[i] - expected).abs() < 1e-6,
1608 "int8 batch prepared mismatch at candidate {i}: got={}, expected={}",
1609 got[i],
1610 expected
1611 );
1612 }
1613 }
1614
1615 #[test]
1616 fn test_int4_batch_prepared_matches_per_item_prepared() {
1617 let query = generate_vector(384, 42);
1618 let prepared = PreparedQuery::from_f32(&query, QuantizationTier::Int4);
1619 let candidates: Vec<Int4Vector> = (0..32)
1620 .map(|i| Int4Vector::from_f32(&generate_vector(384, i + 1)))
1621 .collect();
1622 let wrapped: Vec<QuantizedData> = candidates
1623 .iter()
1624 .cloned()
1625 .map(QuantizedData::Int4)
1626 .collect();
1627
1628 let got = approximate_int4_batch_prepared(&prepared, &candidates).unwrap();
1629 for (i, item) in wrapped.iter().enumerate() {
1630 let expected = approximate_cosine_distance_prepared(&prepared, item).unwrap();
1631 assert!(
1632 (got[i] - expected).abs() < 1e-5,
1633 "int4 batch prepared mismatch at candidate {i}: got={}, expected={}",
1634 got[i],
1635 expected
1636 );
1637 }
1638 }
1639
1640 #[test]
1641 fn test_int4_batch_prepared_api_dispatch_parity() {
1642 for dim in [1usize, 3, 31, 127, 383, 384] {
1647 let query = generate_vector(dim, 700 + dim as u64);
1648 let candidate = generate_vector(dim, 800 + dim as u64);
1649 let prepared = PreparedQuery::from_f32(&query, QuantizationTier::Int4);
1650 let q_cand = Int4Vector::from_f32(&candidate);
1651 let wrapped = QuantizedData::Int4(q_cand.clone());
1652
1653 let batch_result = approximate_int4_batch_prepared(&prepared, &[q_cand]).unwrap();
1654 let per_item_result =
1655 approximate_cosine_distance_prepared(&prepared, &wrapped).unwrap();
1656
1657 assert!(
1658 (batch_result[0] - per_item_result).abs() < 1e-5,
1659 "int4 batch prepared dispatch mismatch at dim={dim}: batch={}, per_item={}",
1660 batch_result[0],
1661 per_item_result
1662 );
1663 }
1664 }
1665
1666 #[test]
1667 fn test_quantized_data_to_f32_roundtrip() {
1668 let v = generate_vector(384, 55);
1669
1670 let full_data = QuantizedData::from_f32(&v, QuantizationTier::Full);
1672 let full_rt = full_data.to_f32();
1673 for (a, b) in v.iter().zip(full_rt.iter()) {
1674 assert!((a - b).abs() < 1e-10, "Full tier should be lossless");
1675 }
1676 }
1677
1678 #[test]
1684 fn test_cosine_distance_prepared_tier_mismatch_returns_typed_error() {
1685 let v = generate_vector(64, 1);
1686 let query = PreparedQuery::from_f32(&v, QuantizationTier::Int8);
1687 let stored = QuantizedData::from_f32(&v, QuantizationTier::Int4);
1688
1689 let err = approximate_cosine_distance_prepared(&query, &stored).unwrap_err();
1690 match err {
1691 EmbedError::TierMismatch {
1692 op,
1693 expected,
1694 actual,
1695 } => {
1696 assert_eq!(op, "approximate_cosine_distance_prepared");
1697 assert_eq!(expected, QuantizationTier::Int4);
1698 assert_eq!(actual, QuantizationTier::Int8);
1699 }
1700 other => panic!("expected TierMismatch, got {other:?}"),
1701 }
1702
1703 assert!(try_approximate_cosine_distance_prepared(&query, &stored).is_err());
1705 }
1706
1707 #[test]
1708 fn test_dot_product_prepared_tier_mismatch_returns_typed_error() {
1709 let v = generate_vector(64, 2);
1710 let query = PreparedQuery::from_f32(&v, QuantizationTier::Full);
1711 let stored = QuantizedData::from_f32(&v, QuantizationTier::Int8);
1712
1713 let err = approximate_dot_product_prepared(&query, &stored).unwrap_err();
1714 assert!(
1715 matches!(
1716 err,
1717 EmbedError::TierMismatch {
1718 op: "approximate_dot_product_prepared",
1719 ..
1720 }
1721 ),
1722 "unexpected error variant: {err:?}"
1723 );
1724
1725 assert!(try_approximate_dot_product_prepared(&query, &stored).is_err());
1726 }
1727
1728 #[test]
1729 fn test_dot_product_prepared_binary_returns_typed_error_not_panic() {
1730 let v = generate_vector(64, 3);
1731 let query = PreparedQuery::from_f32(&v, QuantizationTier::Binary);
1732 let stored = QuantizedData::from_f32(&v, QuantizationTier::Binary);
1733
1734 let err = approximate_dot_product_prepared(&query, &stored).unwrap_err();
1735 assert!(
1736 matches!(err, EmbedError::Internal(_)),
1737 "unexpected error variant: {err:?}"
1738 );
1739 }
1740
1741 #[test]
1742 fn test_cosine_distance_prepared_with_meta_tier_mismatch_returns_typed_error() {
1743 let v = generate_vector(64, 4);
1744 let meta =
1745 PreparedQueryWithMeta::from_f32(&v, QuantizationTier::Full, NormalizationHint::Unknown);
1746 let stored = QuantizedData::from_f32(&v, QuantizationTier::Int8);
1747
1748 let err = approximate_cosine_distance_prepared_with_meta(
1749 &meta,
1750 &stored,
1751 NormalizationHint::Unknown,
1752 )
1753 .unwrap_err();
1754 assert!(matches!(err, EmbedError::TierMismatch { .. }));
1755 }
1756
1757 #[test]
1758 fn test_cosine_distance_prepared_with_meta_validates_stored_unit_norm() {
1759 let query = vec![std::f32::consts::FRAC_1_SQRT_2; 2];
1760 let meta = PreparedQueryWithMeta::from_f32(
1761 &query,
1762 QuantizationTier::Full,
1763 NormalizationHint::Unit,
1764 );
1765 let stored = QuantizedData::Full(vec![2.0, 0.0]);
1766
1767 let got =
1768 approximate_cosine_distance_prepared_with_meta(&meta, &stored, NormalizationHint::Unit)
1769 .unwrap();
1770 let expected = approximate_cosine_distance_prepared(&meta.query, &stored).unwrap();
1771
1772 assert!(
1773 (got - expected).abs() < 1e-6,
1774 "got={got}, expected={expected}"
1775 );
1776 }
1777
1778 #[test]
1779 fn test_batch_cosine_distance_prepared_tier_mismatch_returns_typed_error() {
1780 let v = generate_vector(64, 5);
1781 let query = PreparedQuery::from_f32(&v, QuantizationTier::Int8);
1782 let stored = vec![
1783 QuantizedData::from_f32(&v, QuantizationTier::Int8),
1784 QuantizedData::from_f32(&v, QuantizationTier::Int4), ];
1786
1787 let err = batch_approximate_cosine_distance_prepared(&query, &stored).unwrap_err();
1788 assert!(matches!(err, EmbedError::TierMismatch { .. }));
1789
1790 let mut out = vec![9.0, 9.0, 9.0]; let err =
1792 batch_approximate_cosine_distance_prepared_into(&query, &stored, &mut out).unwrap_err();
1793 assert!(matches!(err, EmbedError::TierMismatch { .. }));
1794 assert!(
1795 out.is_empty(),
1796 "buffer must be cleared, not left with stale data"
1797 );
1798 }
1799
1800 #[test]
1801 fn test_int8_batch_prepared_wrong_tier_returns_typed_error() {
1802 let v = generate_vector(64, 6);
1803 let query = PreparedQuery::from_f32(&v, QuantizationTier::Int4); let candidates = vec![QuantizedVector::from_f32(&v)];
1805
1806 let err = approximate_int8_batch_prepared(&query, &candidates).unwrap_err();
1807 match err {
1808 EmbedError::TierMismatch {
1809 op,
1810 expected,
1811 actual,
1812 } => {
1813 assert_eq!(op, "approximate_int8_batch_prepared");
1814 assert_eq!(expected, QuantizationTier::Int8);
1815 assert_eq!(actual, QuantizationTier::Int4);
1816 }
1817 other => panic!("expected TierMismatch, got {other:?}"),
1818 }
1819
1820 let mut out = vec![9.0];
1821 let err = approximate_int8_batch_prepared_into(&query, &candidates, &mut out).unwrap_err();
1822 assert!(matches!(err, EmbedError::TierMismatch { .. }));
1823 assert!(
1824 out.is_empty(),
1825 "buffer must be cleared, not left with stale data"
1826 );
1827 }
1828
1829 #[test]
1830 fn test_int4_batch_prepared_wrong_tier_returns_typed_error() {
1831 let v = generate_vector(64, 7);
1832 let query = PreparedQuery::from_f32(&v, QuantizationTier::Int8); let candidates = vec![Int4Vector::from_f32(&v)];
1834
1835 let err = approximate_int4_batch_prepared(&query, &candidates).unwrap_err();
1836 match err {
1837 EmbedError::TierMismatch {
1838 op,
1839 expected,
1840 actual,
1841 } => {
1842 assert_eq!(op, "approximate_int4_batch_prepared");
1843 assert_eq!(expected, QuantizationTier::Int4);
1844 assert_eq!(actual, QuantizationTier::Int8);
1845 }
1846 other => panic!("expected TierMismatch, got {other:?}"),
1847 }
1848
1849 let mut out = vec![9.0];
1850 let err = approximate_int4_batch_prepared_into(&query, &candidates, &mut out).unwrap_err();
1851 assert!(matches!(err, EmbedError::TierMismatch { .. }));
1852 assert!(
1853 out.is_empty(),
1854 "buffer must be cleared, not left with stale data"
1855 );
1856 }
1857}