use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use uqa_core::{IndexStats, Payload, PostingEntry, PostingList};
use uqa_fusion::LogitGating;
use uqa_scoring::{logit, sigmoid};
use uqa_storage::StorageBackendError;
use crate::base::{require_probability, ExecutionContext, Operator, OperatorResult};
use crate::hybrid::coverage_based_default;
pub struct ProgressiveFusionOperator {
pub stages: Vec<(Vec<Arc<dyn Operator>>, usize)>,
pub alpha: f64,
pub gating: Option<String>,
}
impl ProgressiveFusionOperator {
pub fn new(stages: Vec<(Vec<Arc<dyn Operator>>, usize)>, alpha: f64) -> Self {
Self::with_gating(stages, alpha, None)
}
pub fn with_gating(
stages: Vec<(Vec<Arc<dyn Operator>>, usize)>,
alpha: f64,
gating: Option<String>,
) -> Self {
Self {
stages,
alpha,
gating,
}
}
}
impl Operator for ProgressiveFusionOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
if self.stages.is_empty() {
return Err(StorageBackendError::Other(
"ProgressiveFusionOperator requires at least one stage".to_string(),
));
}
if !self.alpha.is_finite() || !(0.0..=1.0).contains(&self.alpha) {
return Err(StorageBackendError::Other(format!(
"ProgressiveFusionOperator alpha must be finite and in [0, 1], got {}",
self.alpha
)));
}
let gating = match self.gating.as_deref() {
Some(name) => LogitGating::parse(name).ok_or_else(|| {
StorageBackendError::Other(format!(
"ProgressiveFusionOperator has unknown gating function {name:?}"
))
})?,
None => LogitGating::Pass,
};
let mut signal_lists: Vec<PostingList> = Vec::new();
let mut candidate_ids: Option<BTreeSet<u64>> = None;
let mut last_result: PostingList = PostingList::new();
for (signals, k) in &self.stages {
if signals.is_empty() {
return Err(StorageBackendError::Other(
"ProgressiveFusionOperator stages require at least one signal".to_string(),
));
}
for signal in signals {
let mut pl = signal.execute(ctx)?;
for entry in pl.entries() {
require_probability(entry.payload.score, "progressive fusion")?;
}
if let Some(cands) = &candidate_ids {
let kept: Vec<PostingEntry> = pl
.entries()
.iter()
.filter(|e| cands.contains(&e.doc_id))
.cloned()
.collect();
pl = PostingList::from_sorted_unchecked(kept);
}
signal_lists.push(pl);
}
let mut score_maps = Vec::with_capacity(signal_lists.len());
let mut all_doc_ids = BTreeSet::new();
for posting in &signal_lists {
let map: BTreeMap<u64, f64> = posting
.entries()
.iter()
.map(|entry| {
all_doc_ids.insert(entry.doc_id);
(entry.doc_id, entry.payload.score)
})
.collect();
score_maps.push(map);
}
let total = all_doc_ids.len();
let defaults: Vec<f64> = score_maps
.iter()
.map(|scores| coverage_based_default(scores.len(), total, 0.01))
.collect();
let n = signal_lists.len();
let confidence = (n as f64).powf(self.alpha);
let mut scored: Vec<PostingEntry> = all_doc_ids
.into_iter()
.map(|doc_id| {
let mean_gated_logit = score_maps
.iter()
.zip(&defaults)
.map(|(scores, default)| {
gating.apply(logit(scores.get(&doc_id).copied().unwrap_or(*default)))
})
.sum::<f64>()
/ n as f64;
let fused = sigmoid(confidence * mean_gated_logit);
PostingEntry::new(doc_id, Payload::with_score(fused))
})
.collect();
scored.sort_by_key(|e| e.doc_id);
let scored_pl = PostingList::from_sorted_unchecked(scored);
let topk = scored_pl.ranked().select_top_k(*k);
candidate_ids = Some(topk.doc_ids().collect());
last_result = topk;
}
Ok(last_result)
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
let total_n = stats.total_docs as f64;
let mut total = 0.0;
let mut card = total_n;
for (signals, k) in &self.stages {
let ratio = if total_n > 0.0 { card / total_n } else { 1.0 };
for sig in signals {
total += sig.cost_estimate(stats) * ratio;
}
card = card.min(*k as f64);
}
total
}
}
#[cfg(test)]
mod tests {
use super::*;
struct ConstOperator(Vec<PostingEntry>);
impl Operator for ConstOperator {
fn execute(&self, _ctx: &ExecutionContext) -> OperatorResult {
Ok(PostingList::from_sorted_unchecked(self.0.clone()))
}
}
fn entry(id: u64, score: f64) -> PostingEntry {
PostingEntry::new(id, Payload::with_score(score))
}
#[test]
fn single_stage_keeps_top_k() {
let signal = Arc::new(ConstOperator(vec![
entry(1, 0.9),
entry(2, 0.4),
entry(3, 0.7),
])) as Arc<dyn Operator>;
let op = ProgressiveFusionOperator::new(vec![(vec![signal], 2)], 0.0);
let result = op.execute(&ExecutionContext::new()).unwrap();
let ids: Vec<u64> = result.doc_ids().collect();
assert_eq!(ids.len(), 2);
assert!(ids.contains(&1));
assert!(ids.contains(&3));
}
#[test]
fn second_stage_intersects_with_prior_candidates() {
let stage_0 = Arc::new(ConstOperator(vec![
entry(1, 0.9),
entry(2, 0.8),
entry(3, 0.7),
entry(4, 0.6),
])) as Arc<dyn Operator>;
let stage_1 = Arc::new(ConstOperator(vec![
entry(1, 0.95),
entry(4, 0.95),
entry(5, 0.95),
])) as Arc<dyn Operator>;
let op = ProgressiveFusionOperator::new(vec![(vec![stage_0], 3), (vec![stage_1], 2)], 0.0);
let result = op.execute(&ExecutionContext::new()).unwrap();
let ids: Vec<u64> = result.doc_ids().collect();
assert_eq!(ids.len(), 2);
assert!(ids.contains(&1));
}
}