use std::collections::BTreeMap;
use std::sync::Arc;
use uqa_core::{
DocId, FieldName, IndexStats, Payload, PostingEntry, PostingList, Predicate, Value,
};
use uqa_fusion::{
AdaptivePositiveEvidencePool as AdaptivePositiveEvidenceFuser, BayesianEvidenceFusion,
LogitGating, ProbabilisticBoolean, RobustPositiveEvidencePool, SignalQuality,
};
use uqa_scoring::EvidenceLogit;
use uqa_storage::{StorageBackendError, StorageBackendResult};
use crate::base::{
missing_backend, require_probability, ExecutionContext, Operator, OperatorResult,
};
use crate::primitive::TermOperator;
use crate::vector::VectorSimilarityOperator;
pub fn coverage_based_default(n_hits: usize, n_total: usize, floor: f64) -> f64 {
if n_total == 0 {
return 0.5;
}
let r = n_hits as f64 / n_total as f64;
f64::midpoint(1.0 - r, 0.0) + floor * r
}
fn validate_probability_postings(
postings: &PostingList,
operation: &str,
) -> StorageBackendResult<()> {
for entry in postings.entries() {
require_probability(entry.payload.score, operation)?;
}
Ok(())
}
const ADAPTIVE_SPREAD_SHARE: f64 = 0.5;
fn adaptive_signal_weights(
fuser: &RobustPositiveEvidencePool,
score_maps: &[BTreeMap<DocId, f64>],
) -> Option<Vec<f64>> {
let spreads: Vec<f64> = score_maps
.iter()
.map(|scores| {
if scores.len() < 2 {
return 0.0;
}
let logits: Vec<f64> = scores
.values()
.map(|probability| fuser.gated_logit(*probability))
.collect();
let mean = logits.iter().sum::<f64>() / logits.len() as f64;
let variance = logits
.iter()
.map(|logit| {
let difference = logit - mean;
difference * difference
})
.sum::<f64>()
/ logits.len() as f64;
variance.sqrt()
})
.collect();
let total: f64 = spreads.iter().sum();
if total <= f64::EPSILON {
return None;
}
let uniform = (1.0 - ADAPTIVE_SPREAD_SHARE) / score_maps.len() as f64;
Some(
spreads
.iter()
.map(|spread| uniform + ADAPTIVE_SPREAD_SHARE * spread / total)
.collect(),
)
}
pub struct HybridTextVectorOperator {
term_op: TermOperator,
vector_op: VectorSimilarityOperator,
}
impl HybridTextVectorOperator {
pub fn new(
term: impl Into<String>,
text_field: impl Into<FieldName>,
query_vector: Vec<f32>,
threshold: f32,
vector_field: impl Into<FieldName>,
) -> Self {
Self {
term_op: TermOperator::new(term, text_field),
vector_op: VectorSimilarityOperator::new(query_vector, threshold, vector_field),
}
}
}
impl Operator for HybridTextVectorOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
Ok(self
.term_op
.execute(ctx)?
.merge_intersection_owned(&self.vector_op.execute(ctx)?))
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.term_op
.cost_estimate(stats)
.min(self.vector_op.cost_estimate(stats))
}
}
pub struct SemanticFilterOperator {
pub source: Arc<dyn Operator>,
pub vector_op: VectorSimilarityOperator,
}
impl SemanticFilterOperator {
pub fn new(source: Arc<dyn Operator>, vector_op: VectorSimilarityOperator) -> Self {
Self { source, vector_op }
}
}
impl Operator for SemanticFilterOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
Ok(self
.source
.execute(ctx)?
.merge_intersection_owned(&self.vector_op.execute(ctx)?))
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.source
.cost_estimate(stats)
.min(self.vector_op.cost_estimate(stats))
}
}
pub struct BayesianEvidenceFusionOperator {
pub signals: Vec<Arc<dyn Operator>>,
pub base_rate: f64,
pub top_k: Option<usize>,
}
impl BayesianEvidenceFusionOperator {
pub fn new(signals: Vec<Arc<dyn Operator>>, base_rate: f64) -> Self {
Self {
signals,
base_rate,
top_k: None,
}
}
pub fn with_top_k(mut self, top_k: usize) -> Self {
self.top_k = Some(top_k);
self
}
}
impl Operator for BayesianEvidenceFusionOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
if self.signals.is_empty() {
return Err(StorageBackendError::Other(
"Bayesian evidence fusion requires at least one signal".to_string(),
));
}
let fusion = BayesianEvidenceFusion::new(self.base_rate)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
let posting_lists: Vec<PostingList> = self
.signals
.iter()
.map(|signal| signal.execute(ctx))
.collect::<StorageBackendResult<_>>()?;
for posting_list in &posting_lists {
validate_probability_postings(posting_list, "Bayesian evidence fusion")?;
}
let mut all_doc_ids = std::collections::BTreeSet::new();
let score_maps: Vec<BTreeMap<DocId, f64>> = posting_lists
.iter()
.map(|posting_list| {
let mut scores = BTreeMap::new();
for entry in posting_list {
scores.insert(entry.doc_id, entry.payload.score);
all_doc_ids.insert(entry.doc_id);
}
scores
})
.collect();
if all_doc_ids.is_empty() {
return Ok(PostingList::new());
}
let mut entries = Vec::with_capacity(all_doc_ids.len());
for doc_id in all_doc_ids {
let evidence: Vec<EvidenceLogit> = score_maps
.iter()
.filter_map(|scores| scores.get(&doc_id).copied())
.map(EvidenceLogit::from_prior_free_probability)
.collect::<Result<_, _>>()
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
let posterior = fusion
.fuse(&evidence)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
entries.push(PostingEntry::new(
doc_id,
Payload::with_score(posterior.value()),
));
}
let result = PostingList::from_sorted_unchecked(entries);
Ok(match self.top_k {
Some(k) => result.ranked().select_top_k(k),
None => result,
})
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.signals
.iter()
.map(|signal| signal.cost_estimate(stats))
.sum()
}
}
pub struct RobustPositiveEvidencePoolOperator {
pub signals: Vec<Arc<dyn Operator>>,
pub alpha: f64,
pub gating: LogitGating,
pub base_rate: Option<f64>,
pub weights: Option<Vec<f64>>,
pub adaptive_weights: bool,
pub logit_min: Option<Vec<f64>>,
pub logit_max: Option<Vec<f64>>,
pub top_k: Option<usize>,
}
impl RobustPositiveEvidencePoolOperator {
pub fn new(signals: Vec<Arc<dyn Operator>>, alpha: f64) -> Self {
Self {
signals,
alpha,
gating: LogitGating::Softplus,
base_rate: None,
weights: None,
adaptive_weights: false,
logit_min: None,
logit_max: None,
top_k: None,
}
}
pub fn with_adaptive_weights(mut self) -> Self {
self.adaptive_weights = true;
self
}
pub fn with_gating(mut self, gating: LogitGating) -> Self {
self.gating = gating;
self
}
pub fn with_base_rate(mut self, base_rate: f64) -> Self {
self.base_rate = Some(base_rate);
self
}
pub fn with_weights(mut self, weights: Vec<f64>) -> Self {
self.weights = Some(weights);
self
}
pub fn with_logit_normalization(mut self, logit_min: Vec<f64>, logit_max: Vec<f64>) -> Self {
self.logit_min = Some(logit_min);
self.logit_max = Some(logit_max);
self
}
pub fn with_top_k(mut self, top_k: usize) -> Self {
self.top_k = Some(top_k);
self
}
}
impl Operator for RobustPositiveEvidencePoolOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
if self.signals.is_empty() {
return Err(StorageBackendError::Other(
"positive-evidence pool requires at least one signal".to_string(),
));
}
if !self.alpha.is_finite() || !(0.0..=1.0).contains(&self.alpha) {
return Err(StorageBackendError::Other(format!(
"positive-evidence pool alpha must be finite and in [0, 1], got {}",
self.alpha
)));
}
if let Some(base_rate) = self.base_rate {
if !base_rate.is_finite() || base_rate <= 0.0 || base_rate >= 1.0 {
return Err(StorageBackendError::Other(format!(
"positive-evidence pool base_rate must be finite and in (0, 1), got {base_rate}"
)));
}
}
let mut fuser = RobustPositiveEvidencePool::new(self.alpha)
.map_err(|error| StorageBackendError::Other(error.to_string()))?
.with_logit_gating(self.gating);
if let Some(base_rate) = self.base_rate {
fuser = fuser
.with_base_rate(base_rate)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
}
fuser
.validate_configuration(
self.signals.len(),
self.weights.as_deref(),
self.logit_min.as_deref(),
self.logit_max.as_deref(),
)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
let posting_lists: Vec<PostingList> = self
.signals
.iter()
.map(|sig| sig.execute(ctx))
.collect::<StorageBackendResult<_>>()?;
for posting_list in &posting_lists {
validate_probability_postings(posting_list, "positive-evidence pool")?;
}
let mut all_doc_ids: std::collections::BTreeSet<DocId> = std::collections::BTreeSet::new();
let score_maps: Vec<BTreeMap<DocId, f64>> = posting_lists
.iter()
.map(|pl| {
let mut smap = BTreeMap::new();
for entry in pl {
smap.insert(entry.doc_id, entry.payload.score);
all_doc_ids.insert(entry.doc_id);
}
smap
})
.collect();
if all_doc_ids.is_empty() {
return Ok(PostingList::new());
}
let weights = self.weights.clone().or_else(|| {
if self.adaptive_weights {
adaptive_signal_weights(&fuser, &score_maps)
} else {
None
}
});
let mut entries = Vec::with_capacity(all_doc_ids.len());
for doc_id in &all_doc_ids {
let probabilities: Vec<Option<f64>> = score_maps
.iter()
.map(|scores| scores.get(doc_id).copied())
.collect();
let fused_score = fuser
.fuse_configured(
&probabilities,
weights.as_deref(),
self.logit_min.as_deref(),
self.logit_max.as_deref(),
)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
entries.push(PostingEntry::new(*doc_id, Payload::with_score(fused_score)));
}
let result = PostingList::from_sorted_unchecked(entries);
Ok(match self.top_k {
Some(k) => result.ranked().select_top_k(k),
None => result,
})
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.signals.iter().map(|s| s.cost_estimate(stats)).sum()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProbBoolMode {
And,
Or,
}
pub struct ProbBoolFusionOperator {
pub signals: Vec<Arc<dyn Operator>>,
pub mode: ProbBoolMode,
}
impl ProbBoolFusionOperator {
pub fn new(signals: Vec<Arc<dyn Operator>>, mode: ProbBoolMode) -> Self {
Self { signals, mode }
}
}
impl Operator for ProbBoolFusionOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
if self.signals.is_empty() {
return Err(StorageBackendError::Other(
"probabilistic boolean fusion requires at least one signal".to_string(),
));
}
let posting_lists: Vec<PostingList> = self
.signals
.iter()
.map(|sig| sig.execute(ctx))
.collect::<StorageBackendResult<_>>()?;
for posting_list in &posting_lists {
validate_probability_postings(posting_list, "probabilistic boolean fusion")?;
}
let mut all_doc_ids: std::collections::BTreeSet<DocId> = std::collections::BTreeSet::new();
let score_maps: Vec<BTreeMap<DocId, f64>> = posting_lists
.iter()
.map(|pl| {
let mut smap = BTreeMap::new();
for entry in pl {
smap.insert(entry.doc_id, entry.payload.score);
all_doc_ids.insert(entry.doc_id);
}
smap
})
.collect();
if all_doc_ids.is_empty() {
return Ok(PostingList::new());
}
let num_docs = all_doc_ids.len();
let defaults: Vec<f64> = score_maps
.iter()
.map(|m| coverage_based_default(m.len(), num_docs, 0.01))
.collect();
let mut entries: Vec<PostingEntry> = Vec::with_capacity(num_docs);
for doc_id in &all_doc_ids {
let probs: Vec<f64> = score_maps
.iter()
.zip(&defaults)
.map(|(m, def)| m.get(doc_id).copied().unwrap_or(*def))
.collect();
let fused = match self.mode {
ProbBoolMode::And => ProbabilisticBoolean::and(&probs),
ProbBoolMode::Or => ProbabilisticBoolean::or(&probs),
};
entries.push(PostingEntry::new(*doc_id, Payload::with_score(fused)));
}
Ok(PostingList::from_sorted_unchecked(entries))
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.signals.iter().map(|s| s.cost_estimate(stats)).sum()
}
}
pub struct ProbNotOperator {
pub signal: Arc<dyn Operator>,
pub default_prob: f64,
}
impl ProbNotOperator {
pub fn new(signal: Arc<dyn Operator>, default_prob: f64) -> Self {
Self {
signal,
default_prob,
}
}
}
impl Operator for ProbNotOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
require_probability(self.default_prob, "probabilistic NOT default")?;
let pl = self.signal.execute(ctx)?;
validate_probability_postings(&pl, "probabilistic NOT")?;
let mut score_map: BTreeMap<DocId, f64> = BTreeMap::new();
let mut all_ids: std::collections::BTreeSet<DocId> = std::collections::BTreeSet::new();
for entry in &pl {
score_map.insert(entry.doc_id, entry.payload.score);
all_ids.insert(entry.doc_id);
}
if let Some(store) = ctx.document_store.as_ref() {
for id in store.doc_ids()? {
all_ids.insert(id);
}
}
let mut entries: Vec<PostingEntry> = Vec::with_capacity(all_ids.len());
for doc_id in &all_ids {
let p = score_map.get(doc_id).copied().unwrap_or(self.default_prob);
entries.push(PostingEntry::new(*doc_id, Payload::with_score(1.0 - p)));
}
Ok(PostingList::from_sorted_unchecked(entries))
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.signal.cost_estimate(stats)
}
}
pub struct VectorExclusionOperator {
pub positive: Arc<dyn Operator>,
pub negative_op: VectorSimilarityOperator,
}
impl VectorExclusionOperator {
pub fn new(
positive: Arc<dyn Operator>,
negative_vector: Vec<f32>,
negative_threshold: f32,
field: impl Into<FieldName>,
) -> Self {
Self {
positive,
negative_op: VectorSimilarityOperator::new(negative_vector, negative_threshold, field),
}
}
}
impl Operator for VectorExclusionOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
let positive_pl = self.positive.execute(ctx)?;
let negative_pl = self.negative_op.execute(ctx)?;
let negative_ids: std::collections::BTreeSet<DocId> =
negative_pl.entries().iter().map(|e| e.doc_id).collect();
let mut entries: Vec<PostingEntry> = Vec::new();
for entry in positive_pl.entries() {
if !negative_ids.contains(&entry.doc_id) {
entries.push(entry.clone());
}
}
Ok(PostingList::from_sorted_unchecked(entries))
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.positive.cost_estimate(stats) + self.negative_op.cost_estimate(stats)
}
}
pub struct FacetVectorOperator {
pub facet_field: String,
pub vector_op: VectorSimilarityOperator,
pub source: Option<Arc<dyn Operator>>,
}
impl FacetVectorOperator {
pub fn new(
facet_field: impl Into<String>,
query_vector: Vec<f32>,
threshold: f32,
source: Option<Arc<dyn Operator>>,
) -> Self {
Self {
facet_field: facet_field.into(),
vector_op: VectorSimilarityOperator::new(query_vector, threshold, "embedding"),
source,
}
}
}
impl Operator for FacetVectorOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
let vector_pl = self.vector_op.execute(ctx)?;
let vector_ids: std::collections::BTreeSet<DocId> =
vector_pl.entries().iter().map(|e| e.doc_id).collect();
let candidate_ids: Vec<DocId> = if let Some(src) = &self.source {
src.execute(ctx)?
.entries()
.iter()
.filter(|e| vector_ids.contains(&e.doc_id))
.map(|e| e.doc_id)
.collect()
} else {
let mut v: Vec<DocId> = vector_ids.iter().copied().collect();
v.sort_unstable();
v
};
let Some(doc_store) = ctx.document_store.as_ref() else {
return Err(missing_backend("document-store", "vector facet"));
};
let mut value_counts: BTreeMap<String, u64> = BTreeMap::new();
for doc_id in candidate_ids {
if doc_store.get(doc_id)?.is_none() {
return Err(StorageBackendError::Other(format!(
"vector facet candidate {doc_id} is missing from the document store"
)));
}
if let Some(value) = doc_store.get_field(doc_id, &self.facet_field)? {
if !matches!(value, Value::Null) {
let key = value_to_facet_string(&value);
let count = value_counts.entry(key).or_insert(0);
*count = count.checked_add(1).ok_or_else(|| {
StorageBackendError::Other("vector facet count overflowed u64".to_string())
})?;
}
}
}
let mut entries: Vec<PostingEntry> = Vec::with_capacity(value_counts.len());
for (i, (value, count)) in value_counts.into_iter().enumerate() {
if count > 9_007_199_254_740_992 {
return Err(StorageBackendError::Other(format!(
"vector facet count {count} cannot be represented exactly as an f64 score"
)));
}
let mut fields = BTreeMap::new();
fields.insert(
"_facet_field".to_string(),
Value::Str(self.facet_field.clone()),
);
fields.insert("_facet_value".to_string(), Value::Str(value));
fields.insert(
"_facet_count".to_string(),
Value::Int(i64::try_from(count).map_err(|_| {
StorageBackendError::Other(format!(
"vector facet count {count} exceeds the Value::Int range"
))
})?),
);
entries.push(PostingEntry::new(
DocId::try_from(i).map_err(|_| {
StorageBackendError::Other(format!(
"vector facet bucket index {i} exceeds the document-id range"
))
})?,
Payload {
positions: Vec::new(),
score: count as f64,
fields,
},
));
}
Ok(PostingList::from_sorted_unchecked(entries))
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
let mut base = self.vector_op.cost_estimate(stats);
if let Some(src) = &self.source {
base += src.cost_estimate(stats);
}
base
}
}
fn value_to_facet_string(v: &Value) -> String {
match v {
Value::Str(s) => s.clone(),
Value::Int(n) => n.to_string(),
Value::Float(f) => format!("{f}"),
Value::Bool(b) => b.to_string(),
other => format!("{other:?}"),
}
}
pub struct AdaptivePositiveEvidencePoolOperator {
pub signals: Vec<Arc<dyn Operator>>,
pub base_alpha: f64,
pub gating: Option<String>,
}
impl AdaptivePositiveEvidencePoolOperator {
pub fn new(signals: Vec<Arc<dyn Operator>>, base_alpha: f64, gating: Option<String>) -> Self {
Self {
signals,
base_alpha,
gating,
}
}
}
impl Operator for AdaptivePositiveEvidencePoolOperator {
fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
if self.signals.is_empty() {
return Err(StorageBackendError::Other(
"adaptive positive-evidence pool requires at least one signal".to_string(),
));
}
if !self.base_alpha.is_finite() || !(0.0..=1.0).contains(&self.base_alpha) {
return Err(StorageBackendError::Other(format!(
"adaptive positive-evidence pool alpha must be finite and in [0, 1], got {}",
self.base_alpha
)));
}
let posting_lists: Vec<PostingList> = self
.signals
.iter()
.map(|sig| sig.execute(ctx))
.collect::<StorageBackendResult<_>>()?;
for posting_list in &posting_lists {
validate_probability_postings(posting_list, "adaptive positive-evidence pool")?;
}
let mut all_doc_ids: std::collections::BTreeSet<DocId> = std::collections::BTreeSet::new();
let score_maps: Vec<BTreeMap<DocId, f64>> = posting_lists
.iter()
.map(|pl| {
let mut smap = BTreeMap::new();
for entry in pl {
smap.insert(entry.doc_id, entry.payload.score);
all_doc_ids.insert(entry.doc_id);
}
smap
})
.collect();
if all_doc_ids.is_empty() {
return Ok(PostingList::new());
}
let num_docs = all_doc_ids.len();
let qualities: Vec<SignalQuality> = score_maps
.iter()
.map(|smap| {
let coverage = if num_docs > 0 {
smap.len() as f64 / num_docs as f64
} else {
0.0
};
let scores: Vec<f64> = smap.values().copied().collect();
let variance = if scores.len() > 1 {
let mean = scores.iter().sum::<f64>() / scores.len() as f64;
scores.iter().map(|s| (s - mean).powi(2)).sum::<f64>() / scores.len() as f64
} else {
0.0
};
let mean_score = if scores.is_empty() {
0.5
} else {
scores.iter().sum::<f64>() / scores.len() as f64
};
SignalQuality {
coverage_ratio: coverage,
score_variance: variance,
calibration_error: (mean_score - 0.5).abs(),
}
})
.collect();
let defaults: Vec<f64> = score_maps
.iter()
.map(|m| coverage_based_default(m.len(), num_docs, 0.01))
.collect();
let mut fusion = AdaptivePositiveEvidenceFuser::new(self.base_alpha);
if let Some(name) = &self.gating {
let gating = LogitGating::parse(name).ok_or_else(|| {
StorageBackendError::Other(format!("unknown positive-evidence gate: {name}"))
})?;
fusion = fusion.with_gating(gating);
}
let mut entries: Vec<PostingEntry> = Vec::with_capacity(num_docs);
for doc_id in &all_doc_ids {
let probs: Vec<f64> = score_maps
.iter()
.zip(&defaults)
.map(|(m, def)| m.get(doc_id).copied().unwrap_or(*def))
.collect();
let fused = fusion
.fuse(&probs, &qualities)
.map_err(|error| StorageBackendError::Other(error.to_string()))?;
entries.push(PostingEntry::new(*doc_id, Payload::with_score(fused)));
}
Ok(PostingList::from_sorted_unchecked(entries))
}
fn cost_estimate(&self, stats: &IndexStats) -> f64 {
self.signals.iter().map(|s| s.cost_estimate(stats)).sum()
}
}
pub struct IndexScanOperator {
pub index: Arc<dyn uqa_storage::Index>,
pub field: String,
pub predicate: Predicate,
}
impl IndexScanOperator {
pub fn new(
index: Arc<dyn uqa_storage::Index>,
field: impl Into<String>,
predicate: Predicate,
) -> Self {
Self {
index,
field: field.into(),
predicate,
}
}
}
impl Operator for IndexScanOperator {
fn execute(&self, _ctx: &ExecutionContext) -> OperatorResult {
Ok(self.index.scan(&self.predicate))
}
fn cost_estimate(&self, _stats: &IndexStats) -> f64 {
self.index.scan_cost(&self.predicate)
}
}
#[cfg(test)]
mod tests;