1pub const MAX_FINAL_LIMIT: usize = 10_000;
11pub const MAX_QUERY_OFFSET: usize = 100_000;
12pub const MAX_RETRIEVER_K: usize = 100_000;
13pub const MAX_RETRIEVERS: usize = 32;
14pub const MAX_SPARSE_TERMS: usize = 65_536;
15pub const MAX_SET_MEMBERS: usize = 65_536;
16pub const MAX_SET_MEMBER_BYTES: usize = 1024 * 1024;
17pub const MAX_SET_INPUT_BYTES: usize = 16 * 1024 * 1024;
18pub const MAX_PROJECTION_COLUMNS: usize = 4_096;
19pub const MAX_HARD_CONDITIONS: usize = 256;
20pub const MAX_RETRIEVER_WEIGHT: f64 = 1_000_000.0;
21pub const MAX_FUSED_CANDIDATES: usize = 250_000;
22pub const MAX_RAW_INDEX_CANDIDATES: usize = 250_000;
23pub const MAX_RETRIEVER_NAME_BYTES: usize = 128;
24pub(crate) const VECTOR_WORK_QUANTUM: usize = 64;
25pub(crate) const HAMMING_WORK_QUANTUM: usize = 64;
26pub(crate) const FLOAT_WORK_QUANTUM: usize = 64;
29pub(crate) const PARSE_WORK_QUANTUM: usize = 64;
30
31pub(crate) fn work_units(items: usize, quantum: usize) -> usize {
32 items.div_ceil(quantum).max(1)
33}
34
35#[derive(Debug, Clone)]
38pub struct AiExecutionContext {
39 control: crate::ExecutionControl,
40 query_time_nanos: i64,
41 initial_work: usize,
42 max_fused_candidates: usize,
43 remaining_work: std::sync::Arc<std::sync::atomic::AtomicUsize>,
44}
45
46impl AiExecutionContext {
47 pub fn new(deadline: Option<std::time::Instant>, work_budget: usize) -> Self {
48 Self {
49 control: crate::ExecutionControl::new(deadline),
50 query_time_nanos: std::time::SystemTime::now()
51 .duration_since(std::time::UNIX_EPOCH)
52 .map(|duration| duration.as_nanos().min(i64::MAX as u128) as i64)
53 .unwrap_or(0),
54 initial_work: work_budget,
55 max_fused_candidates: MAX_FUSED_CANDIDATES,
56 remaining_work: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(work_budget)),
57 }
58 }
59
60 pub fn with_timeout(timeout: std::time::Duration, work_budget: usize) -> Self {
61 Self::new(Some(std::time::Instant::now() + timeout), work_budget)
62 }
63
64 pub fn with_limits(
65 timeout: std::time::Duration,
66 work_budget: usize,
67 max_fused_candidates: usize,
68 ) -> Self {
69 let mut context = Self::with_timeout(timeout, work_budget);
70 context.max_fused_candidates = max_fused_candidates.min(MAX_FUSED_CANDIDATES);
71 context
72 }
73
74 pub fn with_control(
75 control: crate::ExecutionControl,
76 work_budget: usize,
77 max_fused_candidates: usize,
78 ) -> Self {
79 Self {
80 control,
81 query_time_nanos: std::time::SystemTime::now()
82 .duration_since(std::time::UNIX_EPOCH)
83 .map(|duration| duration.as_nanos().min(i64::MAX as u128) as i64)
84 .unwrap_or(0),
85 initial_work: work_budget,
86 max_fused_candidates: max_fused_candidates.min(MAX_FUSED_CANDIDATES),
87 remaining_work: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(work_budget)),
88 }
89 }
90
91 pub fn cancel(&self) {
92 self.control
93 .cancel(crate::CancellationReason::ClientRequest);
94 }
95
96 pub fn checkpoint(&self) -> crate::Result<()> {
97 self.control.checkpoint()
98 }
99
100 pub fn consume(&self, work: usize) -> crate::Result<()> {
101 self.checkpoint()?;
102 let result = self.remaining_work.fetch_update(
103 std::sync::atomic::Ordering::AcqRel,
104 std::sync::atomic::Ordering::Acquire,
105 |remaining| remaining.checked_sub(work),
106 );
107 if result.is_err() {
108 return Err(crate::MongrelError::WorkBudgetExceeded);
109 }
110 self.checkpoint()
111 }
112
113 pub fn consumed_work(&self) -> usize {
114 self.initial_work.saturating_sub(
115 self.remaining_work
116 .load(std::sync::atomic::Ordering::Acquire),
117 )
118 }
119
120 pub fn work_limit(&self) -> usize {
121 self.initial_work
122 }
123
124 pub fn remaining_duration(&self) -> Option<std::time::Duration> {
125 self.control.remaining_duration()
126 }
127
128 pub fn execution_control(&self) -> &crate::ExecutionControl {
129 &self.control
130 }
131
132 pub fn query_time_nanos(&self) -> i64 {
133 self.query_time_nanos
134 }
135
136 pub fn with_query_time_nanos(&self, query_time_nanos: i64) -> Self {
137 let mut context = self.clone();
138 context.query_time_nanos = query_time_nanos;
139 context
140 }
141
142 pub fn max_fused_candidates(&self) -> usize {
143 self.max_fused_candidates
144 }
145}
146
147#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
149pub enum Condition {
150 Pk(Vec<u8>),
152 BitmapEq { column_id: u16, value: Vec<u8> },
154 BitmapIn {
159 column_id: u16,
160 values: Vec<Vec<u8>>,
161 },
162 BytesPrefix { column_id: u16, prefix: Vec<u8> },
167 Ann {
169 column_id: u16,
170 query: Vec<f32>,
171 k: usize,
172 },
173 FmContains { column_id: u16, pattern: Vec<u8> },
175 FmContainsAll {
180 column_id: u16,
181 patterns: Vec<Vec<u8>>,
182 },
183 Range { column_id: u16, lo: i64, hi: i64 },
187 RangeF64 {
190 column_id: u16,
191 lo: f64,
192 lo_inclusive: bool,
193 hi: f64,
194 hi_inclusive: bool,
195 },
196 SparseMatch {
199 column_id: u16,
200 query: Vec<(u32, f32)>,
201 k: usize,
202 },
203 MinHashSimilar {
207 column_id: u16,
208 query: Vec<u64>,
209 k: usize,
210 },
211 IsNull { column_id: u16 },
215 IsNotNull { column_id: u16 },
218}
219
220#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
223pub enum Retriever {
224 Ann {
225 column_id: u16,
226 query: Vec<f32>,
227 k: usize,
228 },
229 Sparse {
230 column_id: u16,
231 query: Vec<(u32, f32)>,
232 k: usize,
233 },
234 MinHash {
235 column_id: u16,
236 members: Vec<SetMember>,
237 k: usize,
238 },
239}
240
241impl Retriever {
242 pub fn column_id(&self) -> u16 {
243 match self {
244 Self::Ann { column_id, .. }
245 | Self::Sparse { column_id, .. }
246 | Self::MinHash { column_id, .. } => *column_id,
247 }
248 }
249}
250
251pub fn encode_sparse_vector(terms: &[(u32, f32)]) -> crate::Result<Vec<u8>> {
252 Ok(bincode::serialize(terms)?)
253}
254
255#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
256#[serde(untagged)]
257pub enum SetMember {
258 String(String),
259 Number(serde_json::Number),
260 Boolean(bool),
261}
262
263impl SetMember {
264 pub(crate) fn encoded_len(&self) -> usize {
265 match self {
266 Self::String(value) => value.len(),
267 Self::Number(value) => value.to_string().len(),
268 Self::Boolean(_) => 1,
269 }
270 }
271
272 pub fn hash_v1(&self) -> u64 {
273 let value = match self {
274 Self::String(value) => serde_json::Value::String(value.clone()),
275 Self::Number(value) => serde_json::Value::Number(value.clone()),
276 Self::Boolean(value) => serde_json::Value::Bool(*value),
277 };
278 crate::index::minhash_member_hash_v1(&value).expect("SetMember is always scalar")
279 }
280}
281
282#[derive(Debug, Clone, Copy, PartialEq)]
283pub enum RetrieverScore {
284 AnnHammingDistance(u32),
285 AnnCosineDistance(f32),
288 SparseDotProduct(f64),
289 MinHashEstimatedJaccard(f32),
290}
291
292#[derive(Debug, Clone, Copy)]
297pub enum AnnCandidateDistance {
298 Hamming(u32),
299 Cosine(f32),
300}
301
302impl PartialEq for AnnCandidateDistance {
303 fn eq(&self, other: &Self) -> bool {
304 match (self, other) {
305 (Self::Hamming(a), Self::Hamming(b)) => a == b,
306 (Self::Cosine(a), Self::Cosine(b)) => a.total_cmp(b) == std::cmp::Ordering::Equal,
307 _ => false,
308 }
309 }
310}
311
312impl Eq for AnnCandidateDistance {}
313
314#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
315pub enum VectorMetric {
316 Cosine,
317 DotProduct,
318 Euclidean,
319}
320
321#[derive(Debug, Clone)]
322pub struct AnnRerankRequest {
323 pub column_id: u16,
324 pub query: Vec<f32>,
325 pub candidate_k: usize,
326 pub limit: usize,
327 pub metric: VectorMetric,
328}
329
330#[derive(Debug, Clone, Copy, PartialEq)]
331pub struct AnnRerankHit {
332 pub row_id: crate::RowId,
333 pub candidate_distance: AnnCandidateDistance,
334 pub exact_score: f32,
335}
336
337#[derive(Debug, Clone, PartialEq)]
338pub struct RetrieverHit {
339 pub row_id: crate::RowId,
340 pub rank: usize,
342 pub score: RetrieverScore,
343}
344
345#[derive(Debug, Clone)]
346pub struct SetSimilarityRequest {
347 pub column_id: u16,
348 pub members: Vec<SetMember>,
349 pub candidate_k: usize,
350 pub min_jaccard: f32,
351 pub limit: usize,
352}
353
354#[derive(Debug, Clone, PartialEq)]
355pub struct SetSimilarityHit {
356 pub row_id: crate::RowId,
357 pub estimated_jaccard: f32,
358 pub exact_jaccard: f32,
359}
360
361#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
362pub struct SetSimilarityTrace {
363 pub candidate_generation_us: u64,
364 pub gather_us: u64,
365 pub parse_us: u64,
366 pub score_us: u64,
367 pub candidate_count: usize,
368 pub verified_count: usize,
369}
370
371#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
372pub struct SearchRequest {
373 pub must: Vec<Condition>,
374 pub retrievers: Vec<NamedRetriever>,
375 pub fusion: Fusion,
376 pub rerank: Option<Rerank>,
377 pub limit: usize,
378 pub projection: Option<Vec<u16>>,
379}
380
381#[derive(Debug, Clone, Copy)]
382pub struct SearchAfter {
383 pub final_score: f64,
384 pub row_id: crate::RowId,
385 pub returned_count: usize,
386}
387
388#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
389pub enum Rerank {
390 ExactVector {
391 embedding_column: u16,
392 query: Vec<f32>,
393 metric: VectorMetric,
394 candidate_limit: usize,
395 weight: f64,
396 },
397}
398
399#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
400pub struct NamedRetriever {
401 pub name: String,
402 pub weight: f64,
403 pub retriever: Retriever,
404}
405
406#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
407pub enum Fusion {
408 ReciprocalRank { constant: u32 },
409}
410
411#[derive(Debug, Clone, PartialEq)]
412pub struct ComponentScore {
413 pub retriever_name: std::sync::Arc<str>,
414 pub rank: usize,
415 pub raw_score: RetrieverScore,
416 pub contribution: f64,
417}
418
419#[derive(Debug, Clone, PartialEq)]
420pub struct SearchHit {
421 pub row_id: crate::RowId,
422 pub cells: Vec<(u16, crate::Value)>,
423 pub components: Vec<ComponentScore>,
424 pub fused_score: f64,
425 pub exact_rerank_score: Option<f32>,
426 pub final_score: f64,
427 pub final_rank: usize,
429}
430
431#[derive(Debug, Clone, Default)]
433pub struct Query {
434 pub conditions: Vec<Condition>,
435 pub limit: Option<usize>,
436 pub offset: usize,
437}
438
439impl Query {
440 pub fn new() -> Self {
441 Self::default()
442 }
443 pub fn and(mut self, c: Condition) -> Self {
444 self.conditions.push(c);
445 self
446 }
447 pub fn with_limit(mut self, limit: usize) -> Self {
448 self.limit = Some(limit);
449 self
450 }
451 pub fn with_offset(mut self, offset: usize) -> Self {
452 self.offset = offset;
453 self
454 }
455 pub fn pk(key: Vec<u8>) -> Self {
456 Self::new().and(Condition::Pk(key))
457 }
458}
459
460pub fn canonical_query_key(
469 conditions: &[Condition],
470 projection: Option<&[u16]>,
471 epoch: u64,
472) -> u64 {
473 canonical_query_key_with_version(QUERY_CACHE_KEY_VERSION, conditions, projection, epoch)
474}
475
476const QUERY_CACHE_KEY_VERSION: u64 = 2;
477
478fn canonical_query_key_with_version(
479 version: u64,
480 conditions: &[Condition],
481 projection: Option<&[u16]>,
482 epoch: u64,
483) -> u64 {
484 let fold = |seed: u64, b: u64| -> u64 { seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(b) };
485 let mut acc = fold(0xA5A5_A5A5_A5A5_A5A5, version);
486 acc = fold(acc, epoch);
487 let mut digests: Vec<u64> = conditions.iter().map(hash_condition).collect();
489 digests.sort_unstable();
490 let n = digests.len() as u64;
491 acc = fold(acc, n);
492 for d in digests {
493 acc = fold(acc, d);
494 }
495 match projection {
498 Some(p) => {
499 let mut p = p.to_vec();
500 p.sort_unstable();
501 p.dedup();
502 acc = fold(acc, 0x5E);
503 acc = fold(acc, p.len() as u64);
504 for id in p {
505 acc = fold(acc, id as u64);
506 }
507 }
508 None => {
509 acc = fold(acc, 0xA5);
510 }
511 }
512 acc
513}
514
515fn hash_condition(c: &Condition) -> u64 {
519 use std::hash::{Hash, Hasher};
520 let mut h = std::collections::hash_map::DefaultHasher::new();
521 match c {
522 Condition::Pk(k) => {
523 0u8.hash(&mut h);
524 k.hash(&mut h);
525 }
526 Condition::BitmapEq { column_id, value } => {
527 1u8.hash(&mut h);
528 column_id.hash(&mut h);
529 value.hash(&mut h);
530 }
531 Condition::BitmapIn { column_id, values } => {
532 2u8.hash(&mut h);
533 column_id.hash(&mut h);
534 let mut v: Vec<&Vec<u8>> = values.iter().collect();
535 v.sort();
536 v.dedup();
537 v.len().hash(&mut h);
538 for b in v {
539 b.hash(&mut h);
540 }
541 }
542 Condition::Ann {
543 column_id,
544 query,
545 k,
546 } => {
547 3u8.hash(&mut h);
548 column_id.hash(&mut h);
549 k.hash(&mut h);
550 for f in query {
551 f.to_bits().hash(&mut h);
552 }
553 }
554 Condition::FmContains { column_id, pattern } => {
555 4u8.hash(&mut h);
556 column_id.hash(&mut h);
557 pattern.hash(&mut h);
558 }
559 Condition::FmContainsAll {
560 column_id,
561 patterns,
562 } => {
563 5u8.hash(&mut h);
564 column_id.hash(&mut h);
565 let mut sorted: Vec<&[u8]> = patterns.iter().map(|p| p.as_slice()).collect();
566 sorted.sort();
567 sorted.dedup();
568 sorted.len().hash(&mut h);
569 for p in sorted {
570 p.hash(&mut h);
571 }
572 }
573 Condition::Range { column_id, lo, hi } => {
574 6u8.hash(&mut h);
575 column_id.hash(&mut h);
576 lo.hash(&mut h);
577 hi.hash(&mut h);
578 }
579 Condition::RangeF64 {
580 column_id,
581 lo,
582 lo_inclusive,
583 hi,
584 hi_inclusive,
585 } => {
586 7u8.hash(&mut h);
587 column_id.hash(&mut h);
588 lo.to_bits().hash(&mut h);
589 lo_inclusive.hash(&mut h);
590 hi.to_bits().hash(&mut h);
591 hi_inclusive.hash(&mut h);
592 }
593 Condition::SparseMatch {
594 column_id,
595 query,
596 k,
597 } => {
598 8u8.hash(&mut h);
599 column_id.hash(&mut h);
600 k.hash(&mut h);
601 let mut q: Vec<(u32, u32)> = query.iter().map(|(t, w)| (*t, w.to_bits())).collect();
602 q.sort_by_key(|(t, _)| *t);
603 for (t, wb) in q {
604 t.hash(&mut h);
605 wb.hash(&mut h);
606 }
607 }
608 Condition::MinHashSimilar {
609 column_id,
610 query,
611 k,
612 } => {
613 9u8.hash(&mut h);
614 column_id.hash(&mut h);
615 k.hash(&mut h);
616 let mut q = query.clone();
617 q.sort_unstable();
618 q.dedup();
619 for t in q {
620 t.hash(&mut h);
621 }
622 }
623 Condition::IsNull { column_id } => {
624 10u8.hash(&mut h);
625 column_id.hash(&mut h);
626 }
627 Condition::IsNotNull { column_id } => {
628 11u8.hash(&mut h);
629 column_id.hash(&mut h);
630 }
631 Condition::BytesPrefix { column_id, prefix } => {
632 12u8.hash(&mut h);
633 column_id.hash(&mut h);
634 prefix.hash(&mut h);
635 }
636 }
637 h.finish()
638}
639
640const CURSOR_REQUEST_HASH_VERSION: u64 = 1;
641
642fn cursor_u16(out: &mut Vec<u8>, value: u16) {
643 out.extend_from_slice(&value.to_le_bytes());
644}
645
646fn cursor_u32(out: &mut Vec<u8>, value: u32) {
647 out.extend_from_slice(&value.to_le_bytes());
648}
649
650fn cursor_u64(out: &mut Vec<u8>, value: u64) {
651 out.extend_from_slice(&value.to_le_bytes());
652}
653
654fn cursor_bytes(out: &mut Vec<u8>, value: &[u8]) {
655 cursor_u64(out, value.len() as u64);
656 out.extend_from_slice(value);
657}
658
659fn cursor_condition(condition: &Condition) -> Vec<u8> {
660 let mut out = Vec::new();
661 match condition {
662 Condition::Pk(value) => {
663 out.push(0);
664 cursor_bytes(&mut out, value);
665 }
666 Condition::BitmapEq { column_id, value } => {
667 out.push(1);
668 cursor_u16(&mut out, *column_id);
669 cursor_bytes(&mut out, value);
670 }
671 Condition::BitmapIn { column_id, values } => {
672 out.push(2);
673 cursor_u16(&mut out, *column_id);
674 let mut values = values.clone();
675 values.sort();
676 values.dedup();
677 cursor_u64(&mut out, values.len() as u64);
678 for value in values {
679 cursor_bytes(&mut out, &value);
680 }
681 }
682 Condition::Ann {
683 column_id,
684 query,
685 k,
686 } => {
687 out.push(3);
688 cursor_u16(&mut out, *column_id);
689 cursor_u64(&mut out, *k as u64);
690 cursor_u64(&mut out, query.len() as u64);
691 for value in query {
692 cursor_u32(&mut out, value.to_bits());
693 }
694 }
695 Condition::FmContains { column_id, pattern } => {
696 out.push(4);
697 cursor_u16(&mut out, *column_id);
698 cursor_bytes(&mut out, pattern);
699 }
700 Condition::FmContainsAll {
701 column_id,
702 patterns,
703 } => {
704 out.push(5);
705 cursor_u16(&mut out, *column_id);
706 let mut patterns = patterns.clone();
707 patterns.sort();
708 patterns.dedup();
709 cursor_u64(&mut out, patterns.len() as u64);
710 for pattern in patterns {
711 cursor_bytes(&mut out, &pattern);
712 }
713 }
714 Condition::Range { column_id, lo, hi } => {
715 out.push(6);
716 cursor_u16(&mut out, *column_id);
717 out.extend_from_slice(&lo.to_le_bytes());
718 out.extend_from_slice(&hi.to_le_bytes());
719 }
720 Condition::RangeF64 {
721 column_id,
722 lo,
723 lo_inclusive,
724 hi,
725 hi_inclusive,
726 } => {
727 out.push(7);
728 cursor_u16(&mut out, *column_id);
729 cursor_u64(&mut out, lo.to_bits());
730 out.push(u8::from(*lo_inclusive));
731 cursor_u64(&mut out, hi.to_bits());
732 out.push(u8::from(*hi_inclusive));
733 }
734 Condition::SparseMatch {
735 column_id,
736 query,
737 k,
738 } => {
739 out.push(8);
740 cursor_u16(&mut out, *column_id);
741 cursor_u64(&mut out, *k as u64);
742 let mut query: Vec<_> = query
743 .iter()
744 .map(|(token, weight)| (*token, weight.to_bits()))
745 .collect();
746 query.sort_unstable();
747 cursor_u64(&mut out, query.len() as u64);
748 for (token, weight) in query {
749 cursor_u32(&mut out, token);
750 cursor_u32(&mut out, weight);
751 }
752 }
753 Condition::MinHashSimilar {
754 column_id,
755 query,
756 k,
757 } => {
758 out.push(9);
759 cursor_u16(&mut out, *column_id);
760 cursor_u64(&mut out, *k as u64);
761 let mut query = query.clone();
762 query.sort_unstable();
763 query.dedup();
764 cursor_u64(&mut out, query.len() as u64);
765 for value in query {
766 cursor_u64(&mut out, value);
767 }
768 }
769 Condition::IsNull { column_id } => {
770 out.push(10);
771 cursor_u16(&mut out, *column_id);
772 }
773 Condition::IsNotNull { column_id } => {
774 out.push(11);
775 cursor_u16(&mut out, *column_id);
776 }
777 Condition::BytesPrefix { column_id, prefix } => {
778 out.push(12);
779 cursor_u16(&mut out, *column_id);
780 cursor_bytes(&mut out, prefix);
781 }
782 }
783 out
784}
785
786fn cursor_conditions(out: &mut Vec<u8>, conditions: &[Condition]) {
787 let mut encoded: Vec<_> = conditions.iter().map(cursor_condition).collect();
788 encoded.sort();
789 encoded.dedup();
790 cursor_u64(out, encoded.len() as u64);
791 for condition in encoded {
792 cursor_bytes(out, &condition);
793 }
794}
795
796fn cursor_projection(out: &mut Vec<u8>, projection: Option<&[u16]>) {
797 match projection {
798 Some(projection) => {
799 out.push(1);
800 let mut projection = projection.to_vec();
801 projection.sort_unstable();
802 projection.dedup();
803 cursor_u64(out, projection.len() as u64);
804 for column in projection {
805 cursor_u16(out, column);
806 }
807 }
808 None => out.push(0),
809 }
810}
811
812fn cursor_set_member(member: &SetMember) -> Vec<u8> {
813 let mut out = Vec::new();
814 match member {
815 SetMember::String(value) => {
816 out.push(0);
817 cursor_bytes(&mut out, value.as_bytes());
818 }
819 SetMember::Number(value) => {
820 out.push(1);
821 cursor_bytes(&mut out, value.to_string().as_bytes());
822 }
823 SetMember::Boolean(value) => {
824 out.push(2);
825 out.push(u8::from(*value));
826 }
827 }
828 out
829}
830
831fn cursor_retriever(retriever: &Retriever) -> Vec<u8> {
832 let mut out = Vec::new();
833 match retriever {
834 Retriever::Ann {
835 column_id,
836 query,
837 k,
838 } => {
839 out.push(0);
840 cursor_u16(&mut out, *column_id);
841 cursor_u64(&mut out, *k as u64);
842 cursor_u64(&mut out, query.len() as u64);
843 for value in query {
844 cursor_u32(&mut out, value.to_bits());
845 }
846 }
847 Retriever::Sparse {
848 column_id,
849 query,
850 k,
851 } => {
852 out.push(1);
853 cursor_u16(&mut out, *column_id);
854 cursor_u64(&mut out, *k as u64);
855 let mut query: Vec<_> = query
856 .iter()
857 .map(|(token, weight)| (*token, weight.to_bits()))
858 .collect();
859 query.sort_unstable();
860 cursor_u64(&mut out, query.len() as u64);
861 for (token, weight) in query {
862 cursor_u32(&mut out, token);
863 cursor_u32(&mut out, weight);
864 }
865 }
866 Retriever::MinHash {
867 column_id,
868 members,
869 k,
870 } => {
871 out.push(2);
872 cursor_u16(&mut out, *column_id);
873 cursor_u64(&mut out, *k as u64);
874 let mut members: Vec<_> = members.iter().map(cursor_set_member).collect();
875 members.sort();
876 members.dedup();
877 cursor_u64(&mut out, members.len() as u64);
878 for member in members {
879 cursor_bytes(&mut out, &member);
880 }
881 }
882 }
883 out
884}
885
886fn cursor_sha256(bytes: &[u8]) -> [u8; 32] {
887 use sha2::Digest;
888 sha2::Sha256::digest(bytes).into()
889}
890
891pub fn canonical_query_cursor_hash(
893 table: &str,
894 conditions: &[Condition],
895 projection: Option<&[u16]>,
896) -> [u8; 32] {
897 let mut out = Vec::new();
898 cursor_u64(&mut out, CURSOR_REQUEST_HASH_VERSION);
899 out.push(0);
900 cursor_bytes(&mut out, table.as_bytes());
901 cursor_conditions(&mut out, conditions);
902 cursor_projection(&mut out, projection);
903 cursor_sha256(&out)
904}
905
906pub fn canonical_search_cursor_hash(table: &str, request: &SearchRequest) -> [u8; 32] {
908 let mut out = Vec::new();
909 cursor_u64(&mut out, CURSOR_REQUEST_HASH_VERSION);
910 out.push(1);
911 cursor_bytes(&mut out, table.as_bytes());
912 cursor_conditions(&mut out, &request.must);
913 let mut retrievers: Vec<Vec<u8>> = request
914 .retrievers
915 .iter()
916 .map(|named| {
917 let mut encoded = Vec::new();
918 cursor_bytes(&mut encoded, named.name.as_bytes());
919 cursor_u64(&mut encoded, named.weight.to_bits());
920 cursor_bytes(&mut encoded, &cursor_retriever(&named.retriever));
921 encoded
922 })
923 .collect();
924 retrievers.sort();
925 cursor_u64(&mut out, retrievers.len() as u64);
926 for retriever in retrievers {
927 cursor_bytes(&mut out, &retriever);
928 }
929 match request.fusion {
930 Fusion::ReciprocalRank { constant } => {
931 out.push(0);
932 cursor_u32(&mut out, constant);
933 }
934 }
935 match &request.rerank {
936 Some(Rerank::ExactVector {
937 embedding_column,
938 query,
939 metric,
940 candidate_limit,
941 weight,
942 }) => {
943 out.push(1);
944 cursor_u16(&mut out, *embedding_column);
945 out.push(match metric {
946 VectorMetric::Cosine => 0,
947 VectorMetric::DotProduct => 1,
948 VectorMetric::Euclidean => 2,
949 });
950 cursor_u64(&mut out, *candidate_limit as u64);
951 cursor_u64(&mut out, weight.to_bits());
952 cursor_u64(&mut out, query.len() as u64);
953 for value in query {
954 cursor_u32(&mut out, value.to_bits());
955 }
956 }
957 None => out.push(0),
958 }
959 cursor_projection(&mut out, request.projection.as_deref());
960 cursor_sha256(&out)
961}
962
963pub fn condition_columns(conditions: &[Condition]) -> Vec<u16> {
968 let mut cols: Vec<u16> = conditions
969 .iter()
970 .filter_map(|c| match c {
971 Condition::Pk(_) => None,
972 Condition::BitmapEq { column_id, .. }
973 | Condition::BitmapIn { column_id, .. }
974 | Condition::BytesPrefix { column_id, .. }
975 | Condition::Ann { column_id, .. }
976 | Condition::FmContains { column_id, .. }
977 | Condition::FmContainsAll { column_id, .. }
978 | Condition::Range { column_id, .. }
979 | Condition::RangeF64 { column_id, .. }
980 | Condition::SparseMatch { column_id, .. }
981 | Condition::MinHashSimilar { column_id, .. }
982 | Condition::IsNull { column_id }
983 | Condition::IsNotNull { column_id } => Some(*column_id),
984 })
985 .collect();
986 cols.sort_unstable();
987 cols.dedup();
988 cols
989}
990
991#[cfg(test)]
992mod tests {
993 use super::*;
994
995 #[test]
996 fn builder_chains() {
997 let q = Query::pk(b"k".to_vec()).and(Condition::Range {
998 column_id: 1,
999 lo: 0,
1000 hi: 10,
1001 });
1002 assert_eq!(q.conditions.len(), 2);
1003 }
1004
1005 #[test]
1006 fn ann_candidate_distance_variants_are_metric_distinct() {
1007 assert_eq!(
1008 AnnCandidateDistance::Hamming(3),
1009 AnnCandidateDistance::Hamming(3)
1010 );
1011 assert_ne!(
1012 AnnCandidateDistance::Hamming(0),
1013 AnnCandidateDistance::Cosine(0.0)
1014 );
1015 assert_eq!(
1016 AnnCandidateDistance::Cosine(0.25),
1017 AnnCandidateDistance::Cosine(0.25)
1018 );
1019 }
1020
1021 #[test]
1025 fn canonical_key_is_order_independent() {
1026 let e = 7u64;
1027 let a = Query::new()
1028 .and(Condition::Range {
1029 column_id: 1,
1030 lo: 0,
1031 hi: 10,
1032 })
1033 .and(Condition::BitmapEq {
1034 column_id: 2,
1035 value: b"x".to_vec(),
1036 });
1037 let b = Query::new()
1038 .and(Condition::BitmapEq {
1039 column_id: 2,
1040 value: b"x".to_vec(),
1041 })
1042 .and(Condition::Range {
1043 column_id: 1,
1044 lo: 0,
1045 hi: 10,
1046 });
1047 assert_eq!(
1048 canonical_query_key(&a.conditions, None, e),
1049 canonical_query_key(&b.conditions, None, e),
1050 "condition order must not affect the key"
1051 );
1052
1053 let minhash = Condition::MinHashSimilar {
1054 column_id: 4,
1055 query: vec![3, 1, 1, 2],
1056 k: 5,
1057 };
1058 let minhash_canonical = Condition::MinHashSimilar {
1059 column_id: 4,
1060 query: vec![1, 2, 3],
1061 k: 5,
1062 };
1063 assert_eq!(
1064 canonical_query_key(std::slice::from_ref(&minhash), None, e),
1065 canonical_query_key(&[minhash_canonical], None, e),
1066 );
1067
1068 let fm = Condition::FmContainsAll {
1069 column_id: 4,
1070 patterns: vec![b"b".to_vec(), b"a".to_vec(), b"a".to_vec()],
1071 };
1072 let fm_canonical = Condition::FmContainsAll {
1073 column_id: 4,
1074 patterns: vec![b"a".to_vec(), b"b".to_vec()],
1075 };
1076 assert_eq!(
1077 canonical_query_key(std::slice::from_ref(&fm), None, e),
1078 canonical_query_key(&[fm_canonical], None, e),
1079 );
1080 assert_ne!(
1081 canonical_query_key(std::slice::from_ref(&fm), None, e),
1082 canonical_query_key(std::slice::from_ref(&minhash), None, e),
1083 );
1084 assert_ne!(
1085 canonical_query_key_with_version(1, &a.conditions, None, e),
1086 canonical_query_key_with_version(2, &a.conditions, None, e),
1087 );
1088
1089 let ordered = Condition::BitmapIn {
1091 column_id: 3,
1092 values: vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()],
1093 };
1094 let shuffled = Condition::BitmapIn {
1095 column_id: 3,
1096 values: vec![b"c".to_vec(), b"a".to_vec(), b"a".to_vec(), b"b".to_vec()],
1097 };
1098 assert_eq!(
1099 canonical_query_key(std::slice::from_ref(&ordered), None, e),
1100 canonical_query_key(&[shuffled], None, e),
1101 "BitmapIn values must dedup+sort"
1102 );
1103
1104 assert_ne!(
1106 canonical_query_key(&a.conditions, None, e),
1107 canonical_query_key(&a.conditions, None, e + 1),
1108 "epoch must fold into the key"
1109 );
1110
1111 let proj = vec![1u16, 2];
1113 assert_ne!(
1114 canonical_query_key(&a.conditions, None, e),
1115 canonical_query_key(&a.conditions, Some(&proj), e),
1116 "None projection must differ from an explicit projection"
1117 );
1118 let proj_rev = vec![2u16, 1];
1120 assert_eq!(
1121 canonical_query_key(&a.conditions, Some(&proj), e),
1122 canonical_query_key(&a.conditions, Some(&proj_rev), e),
1123 );
1124
1125 let budget = AiExecutionContext::new(None, 2);
1126 budget.consume(1).unwrap();
1127 assert!(matches!(
1128 budget.consume(2),
1129 Err(crate::MongrelError::WorkBudgetExceeded)
1130 ));
1131 budget.cancel();
1132 assert!(matches!(
1133 budget.checkpoint(),
1134 Err(crate::MongrelError::Cancelled)
1135 ));
1136
1137 let expired = AiExecutionContext::new(
1138 Some(std::time::Instant::now() - std::time::Duration::from_millis(1)),
1139 1,
1140 );
1141 assert!(matches!(
1142 expired.checkpoint(),
1143 Err(crate::MongrelError::DeadlineExceeded)
1144 ));
1145
1146 let parent = crate::ExecutionControl::new(None);
1147 let child = AiExecutionContext::with_control(
1148 parent.child_with_deadline(None),
1149 10,
1150 MAX_FUSED_CANDIDATES,
1151 );
1152 parent.cancel(crate::CancellationReason::SessionClosed);
1153 assert!(matches!(
1154 child.checkpoint(),
1155 Err(crate::MongrelError::Cancelled)
1156 ));
1157 }
1158
1159 #[test]
1160 fn cursor_request_hashes_are_canonical_and_semantic() {
1161 let a = vec![
1162 Condition::BitmapIn {
1163 column_id: 2,
1164 values: vec![b"b".to_vec(), b"a".to_vec(), b"a".to_vec()],
1165 },
1166 Condition::Range {
1167 column_id: 3,
1168 lo: 1,
1169 hi: 9,
1170 },
1171 ];
1172 let b = vec![
1173 Condition::Range {
1174 column_id: 3,
1175 lo: 1,
1176 hi: 9,
1177 },
1178 Condition::BitmapIn {
1179 column_id: 2,
1180 values: vec![b"a".to_vec(), b"b".to_vec()],
1181 },
1182 ];
1183 assert_eq!(
1184 canonical_query_cursor_hash("docs", &a, Some(&[2, 1, 2])),
1185 canonical_query_cursor_hash("docs", &b, Some(&[1, 2])),
1186 );
1187 assert_ne!(
1188 canonical_query_cursor_hash("docs", &a, Some(&[1, 2])),
1189 canonical_query_cursor_hash("other", &a, Some(&[1, 2])),
1190 );
1191
1192 let search = SearchRequest {
1193 must: a,
1194 retrievers: vec![NamedRetriever {
1195 name: "dense".into(),
1196 weight: 1.0,
1197 retriever: Retriever::Ann {
1198 column_id: 4,
1199 query: vec![1.0, -1.0],
1200 k: 10,
1201 },
1202 }],
1203 fusion: Fusion::ReciprocalRank { constant: 60 },
1204 rerank: None,
1205 limit: 1,
1206 projection: Some(vec![2, 1]),
1207 };
1208 let mut next_page = search.clone();
1209 next_page.limit = 20;
1210 next_page.projection = Some(vec![1, 2]);
1211 assert_eq!(
1212 canonical_search_cursor_hash("docs", &search),
1213 canonical_search_cursor_hash("docs", &next_page),
1214 "page size and projection order must not change search identity",
1215 );
1216 next_page.retrievers[0].weight = 2.0;
1217 assert_ne!(
1218 canonical_search_cursor_hash("docs", &search),
1219 canonical_search_cursor_hash("docs", &next_page),
1220 );
1221 }
1222}