1#![allow(
19 clippy::needless_pass_by_value,
20 clippy::similar_names,
21 clippy::explicit_iter_loop
22)]
23
24use std::collections::{BTreeMap, BTreeSet};
25use std::sync::Arc;
26
27use uqa_core::{Payload, PostingEntry, PostingList};
28use uqa_fusion::{AttentionFusion, LearnedFusion, MultiHeadAttentionFusion};
29use uqa_scoring::VectorProbabilityTransform;
30use uqa_storage::{StorageBackendError, StorageBackendResult};
31
32use crate::base::{
33 missing_backend, require_finite_score, require_probability, ExecutionContext, Operator,
34 OperatorResult,
35};
36use crate::hybrid::coverage_based_default;
37use crate::primitive::{ScoreOperator, TermOperator};
38
39type ScoreMap = BTreeMap<u64, f64>;
40type CollectedScores = (Vec<ScoreMap>, BTreeSet<u64>);
41
42fn collect_score_maps(
43 signals: &[Arc<dyn Operator>],
44 ctx: &ExecutionContext,
45) -> StorageBackendResult<CollectedScores> {
46 let mut maps: Vec<ScoreMap> = Vec::with_capacity(signals.len());
47 let mut all_ids: BTreeSet<u64> = BTreeSet::new();
48 for sig in signals {
49 let pl = sig.execute(ctx)?;
50 let mut m: BTreeMap<u64, f64> = BTreeMap::new();
51 for entry in pl.iter() {
52 require_probability(entry.payload.score, "learned/attention fusion")?;
53 m.insert(entry.doc_id, entry.payload.score);
54 all_ids.insert(entry.doc_id);
55 }
56 maps.push(m);
57 }
58 Ok((maps, all_ids))
59}
60
61fn require_single_active_evidence(probabilities: &[Option<f64>]) -> StorageBackendResult<f64> {
62 probabilities
63 .iter()
64 .flatten()
65 .next()
66 .copied()
67 .ok_or_else(|| {
68 StorageBackendError::Other(
69 "multi-field fusion invariant violated: the single active signal has no evidence"
70 .to_string(),
71 )
72 })
73}
74
75pub enum AttentionFuser {
79 Single(AttentionFusion),
80 MultiHead(MultiHeadAttentionFusion),
81}
82
83impl AttentionFuser {
84 fn validate_inputs(
85 &self,
86 signal_count: usize,
87 query_feature_count: usize,
88 ) -> Result<(), &'static str> {
89 match self {
90 AttentionFuser::Single(attention) => {
91 attention.validate_inputs(signal_count, query_feature_count)
92 }
93 AttentionFuser::MultiHead(attention) => {
94 attention.validate_inputs(signal_count, query_feature_count)
95 }
96 }
97 }
98
99 fn fuse_batch(
100 &self,
101 probabilities: &[Vec<f64>],
102 query_features: &[f64],
103 ) -> Result<Vec<f64>, &'static str> {
104 match self {
105 AttentionFuser::Single(attention) => {
106 attention.fuse_batch(probabilities, query_features)
107 }
108 AttentionFuser::MultiHead(attention) => {
109 attention.fuse_batch(probabilities, query_features)
110 }
111 }
112 }
113}
114
115pub struct AttentionFusionOperator {
116 pub signals: Vec<Arc<dyn Operator>>,
117 pub attention: AttentionFuser,
118 pub query_features: Vec<f64>,
119}
120
121impl AttentionFusionOperator {
122 pub fn new(
123 signals: Vec<Arc<dyn Operator>>,
124 attention: AttentionFuser,
125 query_features: Vec<f64>,
126 ) -> Self {
127 Self {
128 signals,
129 attention,
130 query_features,
131 }
132 }
133}
134
135impl Operator for AttentionFusionOperator {
136 fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
137 self.attention
138 .validate_inputs(self.signals.len(), self.query_features.len())
139 .map_err(|error| StorageBackendError::Other(error.to_string()))?;
140 let (score_maps, all_ids) = collect_score_maps(&self.signals, ctx)?;
141 let total = all_ids.len();
142 if total == 0 {
143 return Ok(PostingList::default());
144 }
145 let defaults: Vec<f64> = score_maps
146 .iter()
147 .map(|m| coverage_based_default(m.len(), total, 0.01))
148 .collect();
149 let mut candidate_ids = Vec::with_capacity(total);
150 let mut probabilities = Vec::with_capacity(total);
151 for doc_id in all_ids {
152 let probs: Vec<f64> = score_maps
153 .iter()
154 .enumerate()
155 .map(|(j, m)| *m.get(&doc_id).unwrap_or(&defaults[j]))
156 .collect();
157 candidate_ids.push(doc_id);
158 probabilities.push(probs);
159 }
160 let fused = self
161 .attention
162 .fuse_batch(&probabilities, &self.query_features)
163 .map_err(|error| StorageBackendError::Other(error.to_string()))?;
164 if fused.len() != candidate_ids.len() {
165 return Err(StorageBackendError::Other(format!(
166 "attention fusion returned {} scores for {} candidates",
167 fused.len(),
168 candidate_ids.len()
169 )));
170 }
171 let entries = candidate_ids
172 .into_iter()
173 .zip(fused)
174 .map(|(doc_id, score)| {
175 PostingEntry::new(
176 doc_id,
177 Payload {
178 score,
179 ..Default::default()
180 },
181 )
182 })
183 .collect();
184 Ok(PostingList::from_sorted_unchecked(entries))
185 }
186}
187
188pub struct LearnedFusionOperator {
190 pub signals: Vec<Arc<dyn Operator>>,
191 pub learned: LearnedFusion,
192}
193
194impl LearnedFusionOperator {
195 pub fn new(signals: Vec<Arc<dyn Operator>>, learned: LearnedFusion) -> Self {
196 Self { signals, learned }
197 }
198}
199
200impl Operator for LearnedFusionOperator {
201 fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
202 self.learned
203 .validate_inputs(self.signals.len())
204 .map_err(|error| StorageBackendError::Other(error.to_string()))?;
205 let (score_maps, all_ids) = collect_score_maps(&self.signals, ctx)?;
206 let total = all_ids.len();
207 if total == 0 {
208 return Ok(PostingList::default());
209 }
210 let defaults: Vec<f64> = score_maps
211 .iter()
212 .map(|m| coverage_based_default(m.len(), total, 0.01))
213 .collect();
214 let mut entries: Vec<PostingEntry> = Vec::with_capacity(total);
215 for doc_id in all_ids {
216 let probs: Vec<f64> = score_maps
217 .iter()
218 .enumerate()
219 .map(|(j, m)| *m.get(&doc_id).unwrap_or(&defaults[j]))
220 .collect();
221 let fused = self
222 .learned
223 .fuse(&probs)
224 .map_err(|error| StorageBackendError::Other(error.to_string()))?;
225 entries.push(PostingEntry::new(
226 doc_id,
227 Payload {
228 score: fused,
229 ..Default::default()
230 },
231 ));
232 }
233 Ok(PostingList::from_sorted_unchecked(entries))
234 }
235}
236
237pub struct MultiFieldSearchOperator {
244 pub fields: Vec<String>,
245 pub queries: Vec<String>,
246 pub weights: Vec<f64>,
247 pub bayesian_params: uqa_scoring::BayesianBM25Params,
248 pub fusion_alpha: f64,
249}
250
251impl MultiFieldSearchOperator {
252 pub fn new(fields: Vec<String>, query: impl Into<String>, weights: Option<Vec<f64>>) -> Self {
253 let n = fields.len();
254 let query = query.into();
255 Self {
256 fields,
257 queries: vec![query; n],
258 weights: weights.unwrap_or_else(|| vec![1.0; n]),
259 bayesian_params: uqa_scoring::BayesianBM25Params::default(),
260 fusion_alpha: 0.5,
261 }
262 }
263
264 pub fn with_queries(
265 fields: Vec<String>,
266 queries: Vec<String>,
267 weights: Option<Vec<f64>>,
268 ) -> Self {
269 let n = fields.len();
270 Self {
271 fields,
272 queries,
273 weights: weights.unwrap_or_else(|| vec![1.0; n]),
274 bayesian_params: uqa_scoring::BayesianBM25Params::default(),
275 fusion_alpha: 0.5,
276 }
277 }
278}
279
280impl Operator for MultiFieldSearchOperator {
281 fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
282 use std::sync::Arc as StdArc;
283 use uqa_scoring::{BayesianBM25Scorer, Scorer};
284
285 let Some(idx) = ctx.inverted_index.as_ref() else {
286 return Err(missing_backend("inverted-index", "multi-field search"));
287 };
288 if self.fields.is_empty() {
289 return Err(StorageBackendError::Other(
290 "multi-field search requires at least one field".to_string(),
291 ));
292 }
293 if self.weights.len() != self.fields.len() {
294 return Err(StorageBackendError::Other(format!(
295 "multi-field search has {} fields but {} weights",
296 self.fields.len(),
297 self.weights.len()
298 )));
299 }
300 if self.queries.len() != self.fields.len() {
301 return Err(StorageBackendError::Other(format!(
302 "multi-field search has {} fields but {} queries",
303 self.fields.len(),
304 self.queries.len()
305 )));
306 }
307 let mut per_field: Vec<BTreeMap<u64, f64>> = Vec::with_capacity(self.fields.len());
313 let mut all_ids: BTreeSet<u64> = BTreeSet::new();
314 for (field, query) in self.fields.iter().zip(&self.queries) {
315 let analyzer = idx.search_analyzer_revision(field)?;
316 let terms = uqa_storage::inverted_index::analyze_query_terms(&analyzer, query)?;
317 let term_op: Arc<dyn Operator> = Arc::new(TermOperator::new(query, field));
318 let scorer: Arc<dyn Scorer> = Arc::new(
319 BayesianBM25Scorer::new(
320 self.bayesian_params
321 .scaled_for_query_terms(terms.len())
322 .evidence_params(),
323 StdArc::new(idx.field_stats(field)?),
324 )
325 .map_err(|error| StorageBackendError::Other(error.to_string()))?,
326 );
327 let score_op = ScoreOperator::new_keys(scorer, term_op, terms, field);
328 let pl = score_op.execute(ctx)?;
329 let mut m: BTreeMap<u64, f64> = BTreeMap::new();
330 for entry in pl.iter() {
331 require_probability(entry.payload.score, "multi-field search")?;
332 m.insert(entry.doc_id, entry.payload.score);
333 all_ids.insert(entry.doc_id);
334 }
335 per_field.push(m);
336 }
337
338 let total = all_ids.len();
339 if total == 0 {
340 return Ok(PostingList::default());
341 }
342
343 let weight_sum: f64 = self.weights.iter().sum();
344 let normalized: Vec<f64> = if weight_sum > 0.0
345 && self
346 .weights
347 .iter()
348 .all(|weight| weight.is_finite() && *weight >= 0.0)
349 {
350 self.weights.iter().map(|w| w / weight_sum).collect()
351 } else {
352 return Err(StorageBackendError::Other(
353 "multi-field weights must be non-negative and have a positive finite sum"
354 .to_string(),
355 ));
356 };
357
358 let active_field_count = per_field.iter().filter(|scores| !scores.is_empty()).count();
359 let mut fusion = uqa_fusion::RobustPositiveEvidencePool::new(self.fusion_alpha)
360 .map_err(|error| StorageBackendError::Other(error.to_string()))?;
361 if self.bayesian_params.base_rate > 0.0 {
362 fusion = fusion
363 .with_base_rate(self.bayesian_params.base_rate)
364 .map_err(|error| StorageBackendError::Other(error.to_string()))?;
365 }
366 let mut entries: Vec<PostingEntry> = Vec::with_capacity(total);
367 for doc_id in all_ids {
368 let probabilities: Vec<Option<f64>> = per_field
369 .iter()
370 .map(|scores| scores.get(&doc_id).copied())
371 .collect();
372 let fused = if active_field_count == 1 {
373 let evidence = require_single_active_evidence(&probabilities)?;
376 fusion.fuse(&[evidence])
377 } else {
378 fusion
379 .fuse_weighted_sparse(&probabilities, &normalized)
380 .map_err(|error| StorageBackendError::Other(error.to_string()))?
381 };
382 entries.push(PostingEntry::new(
383 doc_id,
384 Payload {
385 score: fused,
386 ..Default::default()
387 },
388 ));
389 }
390 Ok(PostingList::from_sorted_unchecked(entries))
391 }
392
393 fn cost_estimate(&self, stats: &uqa_core::IndexStats) -> f64 {
394 stats.total_docs as f64 * self.fields.len() as f64
395 }
396}
397
398#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
405pub enum RelevantSampleSplit {
406 #[default]
408 TopQuartile,
409 DistanceGap,
413}
414
415pub struct QueryPoolVectorScoreOperator {
429 pub query_vector: Vec<f32>,
430 pub k: usize,
431 pub field: String,
432 pub base_rate: f64,
436 pub split: RelevantSampleSplit,
437}
438
439#[deprecated(
442 since = "0.1.0",
443 note = "use QueryPoolVectorScoreOperator; use VectorCalibrationModel for reusable calibrated probabilities"
444)]
445pub type CalibratedVectorOperator = QueryPoolVectorScoreOperator;
446
447impl QueryPoolVectorScoreOperator {
448 pub fn new(query_vector: Vec<f32>, k: usize, field: impl Into<String>) -> Self {
449 Self {
450 query_vector,
451 k,
452 field: field.into(),
453 base_rate: 0.5,
454 split: RelevantSampleSplit::default(),
455 }
456 }
457
458 pub fn with_split(mut self, split: RelevantSampleSplit) -> Self {
459 self.split = split;
460 self
461 }
462
463 pub fn with_base_rate(mut self, base_rate: f64) -> Self {
464 self.base_rate = base_rate;
465 self
466 }
467}
468
469pub fn calibrate_query_pool_postings(
471 raw: &PostingList,
472 split: RelevantSampleSplit,
473 base_rate: f64,
474) -> StorageBackendResult<PostingList> {
475 if !base_rate.is_finite() || base_rate <= 0.0 || base_rate >= 1.0 {
476 return Err(StorageBackendError::Other(format!(
477 "calibrated vector base_rate must be finite and in (0, 1), got {base_rate}"
478 )));
479 }
480 if raw.is_empty() {
481 return Ok(PostingList::default());
482 }
483
484 let mut distances = Vec::with_capacity(raw.len());
485 for entry in raw.entries() {
486 require_finite_score(entry.payload.score, "calibrated vector search")?;
487 if !(-1.0..=1.0).contains(&entry.payload.score) {
488 return Err(StorageBackendError::Other(format!(
489 "calibrated vector search requires cosine scores in [-1, 1], got {}",
490 entry.payload.score
491 )));
492 }
493 distances.push(1.0 - entry.payload.score);
494 }
495 let calibrator = fit_pool_calibration(&distances, split, base_rate)?;
496
497 let mut out_entries: Vec<PostingEntry> = Vec::with_capacity(raw.len());
498 for (entry, distance) in raw.iter().zip(&distances) {
499 let posterior = match calibrator.as_ref() {
500 Some(transform) => transform
501 .calibrate_one(*distance)
502 .map_err(|error| StorageBackendError::Other(error.to_string()))?,
503 None => base_rate,
504 };
505 out_entries.push(PostingEntry::new(
506 entry.doc_id,
507 Payload {
508 score: posterior.clamp(1e-6, 1.0 - 1e-6),
509 ..Default::default()
510 },
511 ));
512 }
513 out_entries.sort_by_key(|entry| entry.doc_id);
514 Ok(PostingList::from_sorted_unchecked(out_entries))
515}
516
517impl Operator for QueryPoolVectorScoreOperator {
518 fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
519 if !self.base_rate.is_finite() || self.base_rate <= 0.0 || self.base_rate >= 1.0 {
520 return Err(StorageBackendError::Other(format!(
521 "calibrated vector base_rate must be finite and in (0, 1), got {}",
522 self.base_rate
523 )));
524 }
525 if self.query_vector.is_empty()
526 || self
527 .query_vector
528 .iter()
529 .any(|component| !component.is_finite())
530 {
531 return Err(StorageBackendError::Other(
532 "calibrated vector search requires a non-empty finite query vector".to_string(),
533 ));
534 }
535 let Some(idx) = ctx.vector_indexes.get(&self.field) else {
536 return Err(missing_backend("vector-index", "calibrated vector search"));
537 };
538 let raw = idx.search_knn(&self.query_vector, self.k)?;
539 calibrate_query_pool_postings(&raw, self.split, self.base_rate)
540 }
541}
542
543pub fn fit_pool_calibration(
548 distances: &[f64],
549 split: RelevantSampleSplit,
550 base_rate: f64,
551) -> StorageBackendResult<Option<VectorProbabilityTransform>> {
552 if !base_rate.is_finite() || base_rate <= 0.0 || base_rate >= 1.0 {
553 return Err(StorageBackendError::Other(format!(
554 "pool calibration base_rate must be finite and in (0, 1), got {base_rate}"
555 )));
556 }
557 if distances.iter().any(|distance| !distance.is_finite()) {
558 return Err(StorageBackendError::Other(
559 "pool calibration distances must be finite".to_string(),
560 ));
561 }
562 if distances.len() < 2 {
563 return Ok(None);
564 }
565 let mut sorted = distances.to_vec();
566 sorted.sort_by(f64::total_cmp);
567
568 let head_len = match split {
569 RelevantSampleSplit::TopQuartile => quartile_head(sorted.len()),
570 RelevantSampleSplit::DistanceGap => {
571 distance_gap_split(&sorted).unwrap_or_else(|| quartile_head(sorted.len()))
572 }
573 }
574 .clamp(1, sorted.len() - 1);
575
576 let mu_match = mean(&sorted[..head_len]);
577 let mu_random = mean(&sorted[head_len..]);
578 let sigma = standard_deviation(&sorted);
579 if sigma <= f64::EPSILON || mu_random - mu_match <= f64::EPSILON {
580 return Ok(None);
581 }
582 Ok(Some(
583 VectorProbabilityTransform::new(mu_match, mu_random, sigma, base_rate)
584 .map_err(|error| StorageBackendError::Other(error.to_string()))?,
585 ))
586}
587
588fn quartile_head(pool_size: usize) -> usize {
589 pool_size.div_ceil(4)
590}
591
592fn distance_gap_split(sorted: &[f64]) -> Option<usize> {
595 let mut max_gap = 0.0f64;
596 let mut split_index = None;
597 for (index, window) in sorted.windows(2).enumerate() {
598 let gap = window[1] - window[0];
599 if gap > max_gap {
600 max_gap = gap;
601 split_index = Some(index + 1);
602 }
603 }
604 split_index
605}
606
607fn mean(values: &[f64]) -> f64 {
608 values.iter().sum::<f64>() / values.len() as f64
609}
610
611fn standard_deviation(values: &[f64]) -> f64 {
612 let mu = mean(values);
613 let variance = values
614 .iter()
615 .map(|value| {
616 let difference = value - mu;
617 difference * difference
618 })
619 .sum::<f64>()
620 / values.len() as f64;
621 variance.sqrt()
622}
623
624#[cfg(test)]
625mod tests {
626 use super::*;
627 use uqa_core::{Payload, PostingEntry, PostingList};
628
629 struct LiteralOperator(Vec<(u64, f64)>);
630 impl Operator for LiteralOperator {
631 fn execute(&self, _ctx: &ExecutionContext) -> OperatorResult {
632 Ok(PostingList::from_sorted_unchecked(
633 self.0
634 .iter()
635 .map(|(d, s)| {
636 PostingEntry::new(
637 *d,
638 Payload {
639 score: *s,
640 ..Default::default()
641 },
642 )
643 })
644 .collect(),
645 ))
646 }
647 }
648
649 #[test]
650 fn learned_fusion_combines_two_signals() {
651 let signals: Vec<Arc<dyn Operator>> = vec![
652 Arc::new(LiteralOperator(vec![(1, 0.8), (2, 0.6)])),
653 Arc::new(LiteralOperator(vec![(1, 0.7), (3, 0.4)])),
654 ];
655 let learned = LearnedFusion::new(2, 0.0);
656 let op = LearnedFusionOperator::new(signals, learned);
657 let pl = op.execute(&ExecutionContext::new()).unwrap();
658 let ids: Vec<u64> = pl.iter().map(|e| e.doc_id).collect();
659 assert_eq!(ids, vec![1, 2, 3]);
660 }
661
662 #[test]
663 fn missing_single_active_evidence_is_an_invariant_error() {
664 let error = require_single_active_evidence(&[None, None]).unwrap_err();
665 assert!(error.to_string().contains("single active signal"));
666 }
667
668 #[test]
669 fn query_pool_vector_missing_index_is_an_execution_error() {
670 let op = QueryPoolVectorScoreOperator::new(vec![0.0; 3], 0, "missing").with_base_rate(0.5);
671 let error = op.execute(&ExecutionContext::new()).unwrap_err();
672 assert!(error.to_string().contains("vector-index"));
673 }
674
675 #[test]
676 fn pool_calibration_discriminates_head_from_tail() {
677 let distances = [0.02, 0.05, 0.30, 0.35, 0.40, 0.45, 0.50, 0.55];
678 let transform = fit_pool_calibration(&distances, RelevantSampleSplit::TopQuartile, 0.5)
679 .expect("valid fit request")
680 .expect("separated pool fits");
681 let head = transform.calibrate_one(0.02).unwrap();
682 let mid = transform.calibrate_one(0.30).unwrap();
683 let tail = transform.calibrate_one(0.55).unwrap();
684 assert!(head > mid && mid > tail, "{head} > {mid} > {tail}");
685 assert!(head > 0.5, "head evidence must be positive, got {head}");
686 assert!(tail < 0.5, "tail evidence must be negative, got {tail}");
687 }
688
689 #[test]
690 fn pool_calibration_rejects_uninformative_pools() {
691 assert!(
692 fit_pool_calibration(&[0.3], RelevantSampleSplit::TopQuartile, 0.5)
693 .unwrap()
694 .is_none()
695 );
696 assert!(
697 fit_pool_calibration(&[0.3, 0.3, 0.3, 0.3], RelevantSampleSplit::TopQuartile, 0.5)
698 .unwrap()
699 .is_none()
700 );
701 }
702
703 #[test]
704 fn pool_calibration_rejects_invalid_numeric_inputs() {
705 assert!(
706 fit_pool_calibration(&[f64::NAN, 0.2], RelevantSampleSplit::TopQuartile, 0.5).is_err()
707 );
708 assert!(fit_pool_calibration(&[0.1, 0.2], RelevantSampleSplit::TopQuartile, 1.0).is_err());
709 }
710
711 #[test]
712 fn distance_gap_split_finds_the_semantic_cliff() {
713 let sorted = [0.05, 0.06, 0.07, 0.40, 0.42, 0.44];
714 assert_eq!(distance_gap_split(&sorted), Some(3));
715 assert_eq!(distance_gap_split(&[0.3, 0.3, 0.3]), None);
716 }
717}