#![allow(clippy::large_enum_variant)]
use std::collections::BTreeMap;
use std::sync::Arc;
use uqa_core::{Predicate, Value};
use crate::aggregation::AggregationMonoid;
use crate::base::Direction;
pub type ScorerRef = Arc<dyn uqa_scoring::Scorer>;
pub type AttentionRef = Arc<dyn AttentionFuserDyn>;
pub type LearnedFusionRef = Arc<dyn LearnedFuserDyn>;
pub type VertexPredicate = Arc<dyn Fn(&uqa_core::Vertex) -> bool + Send + Sync>;
pub type PathWeightPredicate = Arc<dyn Fn(f64) -> bool + Send + Sync>;
pub type VertexConstraint = Arc<dyn Fn(&uqa_core::Vertex) -> bool + Send + Sync>;
pub type EdgeConstraint = Arc<dyn Fn(&uqa_core::Edge) -> bool + Send + Sync>;
pub trait AttentionFuserDyn: Send + Sync {
fn validate_inputs(
&self,
signal_count: usize,
query_feature_count: usize,
) -> Result<(), &'static str>;
fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str>;
fn fuse_batch(
&self,
probabilities: &[Vec<f64>],
query_features: &[f64],
) -> Result<Vec<f64>, &'static str> {
probabilities
.iter()
.map(|sample| self.fuse(sample, query_features))
.collect()
}
fn head_count(&self) -> usize;
fn normalize(&self) -> bool;
fn alpha(&self) -> f64;
fn base_rate(&self) -> Option<f64>;
}
pub trait LearnedFuserDyn: Send + Sync {
fn validate_inputs(&self, signal_count: usize) -> Result<(), &'static str>;
fn fuse(&self, probs: &[f64]) -> Result<f64, &'static str>;
}
impl AttentionFuserDyn for uqa_fusion::AttentionFusion {
fn validate_inputs(
&self,
signal_count: usize,
query_feature_count: usize,
) -> Result<(), &'static str> {
uqa_fusion::AttentionFusion::validate_inputs(self, signal_count, query_feature_count)
}
fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str> {
uqa_fusion::AttentionFusion::fuse(self, probs, query_features)
}
fn fuse_batch(
&self,
probabilities: &[Vec<f64>],
query_features: &[f64],
) -> Result<Vec<f64>, &'static str> {
uqa_fusion::AttentionFusion::fuse_batch(self, probabilities, query_features)
}
fn head_count(&self) -> usize {
1
}
fn normalize(&self) -> bool {
self.normalize
}
fn alpha(&self) -> f64 {
self.alpha
}
fn base_rate(&self) -> Option<f64> {
self.base_rate
}
}
impl AttentionFuserDyn for uqa_fusion::MultiHeadAttentionFusion {
fn validate_inputs(
&self,
signal_count: usize,
query_feature_count: usize,
) -> Result<(), &'static str> {
uqa_fusion::MultiHeadAttentionFusion::validate_inputs(
self,
signal_count,
query_feature_count,
)
}
fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str> {
uqa_fusion::MultiHeadAttentionFusion::fuse(self, probs, query_features)
}
fn fuse_batch(
&self,
probabilities: &[Vec<f64>],
query_features: &[f64],
) -> Result<Vec<f64>, &'static str> {
uqa_fusion::MultiHeadAttentionFusion::fuse_batch(self, probabilities, query_features)
}
fn head_count(&self) -> usize {
uqa_fusion::MultiHeadAttentionFusion::n_heads(self)
}
fn normalize(&self) -> bool {
uqa_fusion::MultiHeadAttentionFusion::normalize(self)
}
fn alpha(&self) -> f64 {
uqa_fusion::MultiHeadAttentionFusion::alpha(self).unwrap_or(f64::NAN)
}
fn base_rate(&self) -> Option<f64> {
None
}
}
impl LearnedFuserDyn for uqa_fusion::LearnedFusion {
fn validate_inputs(&self, signal_count: usize) -> Result<(), &'static str> {
uqa_fusion::LearnedFusion::validate_inputs(self, signal_count)
}
fn fuse(&self, probs: &[f64]) -> Result<f64, &'static str> {
uqa_fusion::LearnedFusion::fuse(self, probs)
}
}
#[derive(Clone)]
pub struct VertexPatternIR {
pub variable: String,
pub constraints: Vec<VertexConstraint>,
pub label: Option<String>,
}
impl std::fmt::Debug for VertexPatternIR {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("VertexPatternIR")
.field("variable", &self.variable)
.field("constraints_count", &self.constraints.len())
.field("label", &self.label)
.finish()
}
}
#[derive(Clone)]
pub struct EdgePatternIR {
pub source_var: String,
pub target_var: String,
pub label: Option<String>,
pub constraints: Vec<EdgeConstraint>,
}
impl std::fmt::Debug for EdgePatternIR {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EdgePatternIR")
.field("source_var", &self.source_var)
.field("target_var", &self.target_var)
.field("label", &self.label)
.field("constraints_count", &self.constraints.len())
.finish()
}
}
#[derive(Clone, Debug)]
pub struct GraphPatternIR {
pub vertex_patterns: Vec<VertexPatternIR>,
pub edge_patterns: Vec<EdgePatternIR>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProbBoolMode {
And,
Or,
}
#[derive(Clone, Debug)]
pub enum GatingSpec {
Softplus,
Pass,
Sigmoid { feature: String },
ReLU,
Swish,
Gelu,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeepFusionAggregation {
Mean,
Sum,
Max,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DeepFusionPoolMethod {
Average,
Max,
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum TextScoringMode {
BM25,
BayesianBM25,
CustomBM25(uqa_scoring::BM25Params),
CustomBayesianBM25(uqa_scoring::BayesianBM25Params),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TextTopKStrategy {
Wand,
BlockMaxWand,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TextTopKPlan {
pub k: usize,
pub strategy: TextTopKStrategy,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExternalPriorMode {
Authority,
Recency,
}
#[derive(Clone)]
pub enum OperatorTree {
Empty,
Term {
query: String,
field: Option<String>,
scoring: Option<TextScoringMode>,
top_k: Option<TextTopKPlan>,
},
Filter {
field: String,
predicate: Predicate,
source: Option<Box<OperatorTree>>,
},
Facet {
field: String,
source: Option<Box<OperatorTree>>,
},
Score {
scorer: ScorerRef,
source: Box<OperatorTree>,
query_terms: Vec<String>,
field: String,
},
BayesianScore {
source: Box<OperatorTree>,
field: Option<String>,
},
BayesianMatchWithPrior {
field: String,
query: String,
prior_field: String,
mode: ExternalPriorMode,
},
Intersect(Vec<OperatorTree>),
Union(Vec<OperatorTree>),
Complement(Box<OperatorTree>),
Composed(Vec<OperatorTree>),
EncodeGraphPosting { source: Box<OperatorTree> },
VectorSimilarity {
query_vector: Vec<f32>,
threshold: f32,
field: String,
},
KNN {
query_vector: Vec<f32>,
k: usize,
field: String,
},
CalibratedVectorMatch {
query_vector: Vec<f32>,
k: usize,
field: String,
threshold: Option<f64>,
},
CosineProbability(Box<OperatorTree>),
BayesianEvidenceFusion {
signals: Vec<OperatorTree>,
base_rate: Option<f64>,
},
RobustPositiveEvidencePool {
signals: Vec<OperatorTree>,
alpha: f64,
gating: GatingSpec,
weights: Option<Vec<f64>>,
logit_min: Option<Vec<f64>>,
logit_max: Option<Vec<f64>>,
adaptive_weights: bool,
},
ProbBoolFusion {
signals: Vec<OperatorTree>,
mode: ProbBoolMode,
},
ProbNot {
signal: Box<OperatorTree>,
default_prob: f64,
},
AttentionFusion {
signals: Vec<OperatorTree>,
attention: AttentionRef,
query_features: Vec<f64>,
},
LearnedFusion {
signals: Vec<OperatorTree>,
learned: LearnedFusionRef,
},
SparseThreshold {
source: Box<OperatorTree>,
threshold: f64,
},
Traverse {
start_vertex: u64,
graph: String,
label: Option<String>,
max_hops: usize,
vertex_predicate: Option<VertexPredicate>,
},
GraphNeighbors {
vertex: u64,
graph: String,
label: Option<String>,
direction: Direction,
},
GraphEdges {
graph: String,
label: Option<String>,
},
PatternMatch {
pattern: GraphPatternIR,
graph: String,
},
RegularPathQuery {
rpq_source: String,
start_vertex: u64,
graph: String,
},
GraphJoin {
left: Box<OperatorTree>,
right: Box<OperatorTree>,
label: Option<String>,
graph: String,
},
IndexScan {
index_name: String,
field: String,
predicate: Predicate,
},
Aggregate {
source: Option<Box<OperatorTree>>,
field: String,
monoid: Arc<dyn AggregationMonoid>,
},
GroupBy {
source: Box<OperatorTree>,
group_field: String,
agg_field: String,
monoid: Arc<dyn AggregationMonoid>,
},
MultiStage { stages: Vec<MultiStageEntry> },
MultiFieldSearch {
fields: Vec<String>,
queries: Vec<String>,
weights: Option<Vec<f64>>,
},
HybridTextVector {
term_op: Box<OperatorTree>,
vector_op: Box<OperatorTree>,
alpha: f64,
},
SemanticFilter {
source: Box<OperatorTree>,
vector_op: Box<OperatorTree>,
},
VectorExclusion {
positive: Box<OperatorTree>,
negative: Box<OperatorTree>,
},
FacetVector {
vector_op: Box<OperatorTree>,
facet_field: String,
},
VertexAggregation {
source: Box<OperatorTree>,
monoid: Arc<dyn AggregationMonoid>,
},
WeightedPathQuery {
rpq_source: String,
start_vertex: u64,
graph: String,
weight_property: String,
default_edge_weight: f64,
max_hops: usize,
predicate: PathWeightPredicate,
predicate_selectivity: f64,
score: f64,
},
MessagePassing { source: Box<OperatorTree> },
GraphEmbedding { source: Box<OperatorTree> },
PageRank { graph: String },
HITS { graph: String },
BetweennessCentrality { graph: String },
TextSimilarityJoin {
left: Box<OperatorTree>,
right: Box<OperatorTree>,
threshold: f64,
},
VectorSimilarityJoin {
left: Box<OperatorTree>,
right: Box<OperatorTree>,
threshold: f64,
},
HybridJoin {
left: Box<OperatorTree>,
right: Box<OperatorTree>,
},
CrossParadigmJoin {
left: Box<OperatorTree>,
right: Box<OperatorTree>,
},
TemporalTraverse {
start_vertex: u64,
graph: String,
label: Option<String>,
max_hops: usize,
temporal_filter: Option<TemporalFilterIR>,
},
TemporalPatternMatch {
pattern: GraphPatternIR,
graph: String,
temporal_filter: Option<TemporalFilterIR>,
},
ProgressiveFusion {
stages: Vec<ProgressiveFusionEntry>,
alpha: f64,
gating: GatingSpec,
},
DeepFusion {
layers: Vec<DeepFusionLayer>,
alpha: f64,
gating: GatingSpec,
},
DeepPredict { model: String },
Opaque {
kind: String,
children: Vec<OperatorTree>,
meta: BTreeMap<String, Value>,
},
}
#[derive(Clone)]
pub struct MultiStageEntry {
pub child: OperatorTree,
pub cutoff: MultiStageCutoff,
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum MultiStageCutoff {
TopK(usize),
Ratio(f64),
}
#[derive(Clone)]
pub struct ProgressiveFusionEntry {
pub signal: OperatorTree,
pub k: usize,
}
#[derive(Clone)]
pub enum DeepFusionLayer {
Signal {
signals: Vec<OperatorTree>,
},
Propagate {
edge_label: Option<String>,
aggregation: DeepFusionAggregation,
direction: Direction,
},
Conv {
edge_label: Option<String>,
hop_weights: Vec<f64>,
direction: Direction,
},
Pool {
edge_label: Option<String>,
pool_size: usize,
method: DeepFusionPoolMethod,
direction: Direction,
},
Flatten,
Dense {
weights: Vec<f64>,
bias: Vec<f64>,
output_channels: usize,
input_channels: usize,
},
Softmax,
BatchNorm {
epsilon: f64,
},
Dropout {
probability: f64,
},
}
#[derive(Clone, Debug, Default)]
pub struct TemporalFilterIR {
pub timestamp: Option<f64>,
pub time_range: Option<(f64, f64)>,
}
impl OperatorTree {
#[allow(clippy::too_many_lines)]
pub fn visit(&self, visitor: &mut impl FnMut(&OperatorTree)) {
visitor(self);
match self {
OperatorTree::Filter {
source: Some(source),
..
}
| OperatorTree::Facet {
source: Some(source),
..
}
| OperatorTree::Score { source, .. }
| OperatorTree::BayesianScore { source, .. }
| OperatorTree::Complement(source)
| OperatorTree::EncodeGraphPosting { source }
| OperatorTree::CosineProbability(source)
| OperatorTree::ProbNot { signal: source, .. }
| OperatorTree::SparseThreshold { source, .. }
| OperatorTree::VertexAggregation { source, .. }
| OperatorTree::MessagePassing { source }
| OperatorTree::GraphEmbedding { source }
| OperatorTree::GroupBy { source, .. }
| OperatorTree::Aggregate {
source: Some(source),
..
} => source.visit(visitor),
OperatorTree::Intersect(children)
| OperatorTree::Union(children)
| OperatorTree::Composed(children)
| OperatorTree::Opaque { children, .. }
| OperatorTree::BayesianEvidenceFusion {
signals: children, ..
}
| OperatorTree::RobustPositiveEvidencePool {
signals: children, ..
}
| OperatorTree::ProbBoolFusion {
signals: children, ..
}
| OperatorTree::AttentionFusion {
signals: children, ..
}
| OperatorTree::LearnedFusion {
signals: children, ..
} => visit_operator_slice(children, visitor),
OperatorTree::GraphJoin { left, right, .. }
| OperatorTree::TextSimilarityJoin { left, right, .. }
| OperatorTree::VectorSimilarityJoin { left, right, .. }
| OperatorTree::HybridJoin { left, right }
| OperatorTree::CrossParadigmJoin { left, right }
| OperatorTree::HybridTextVector {
term_op: left,
vector_op: right,
..
}
| OperatorTree::SemanticFilter {
source: left,
vector_op: right,
}
| OperatorTree::VectorExclusion {
positive: left,
negative: right,
} => {
left.visit(visitor);
right.visit(visitor);
}
OperatorTree::FacetVector { vector_op, .. } => vector_op.visit(visitor),
OperatorTree::MultiStage { stages } => {
for stage in stages {
stage.child.visit(visitor);
}
}
OperatorTree::ProgressiveFusion { stages, .. } => {
for stage in stages {
stage.signal.visit(visitor);
}
}
OperatorTree::DeepFusion { layers, .. } => {
for layer in layers {
if let DeepFusionLayer::Signal { signals } = layer {
for signal in signals {
signal.visit(visitor);
}
}
}
}
OperatorTree::Empty
| OperatorTree::Term { .. }
| OperatorTree::BayesianMatchWithPrior { .. }
| OperatorTree::Filter { source: None, .. }
| OperatorTree::Facet { source: None, .. }
| OperatorTree::VectorSimilarity { .. }
| OperatorTree::KNN { .. }
| OperatorTree::CalibratedVectorMatch { .. }
| OperatorTree::Traverse { .. }
| OperatorTree::GraphNeighbors { .. }
| OperatorTree::GraphEdges { .. }
| OperatorTree::PatternMatch { .. }
| OperatorTree::RegularPathQuery { .. }
| OperatorTree::IndexScan { .. }
| OperatorTree::Aggregate { source: None, .. }
| OperatorTree::MultiFieldSearch { .. }
| OperatorTree::WeightedPathQuery { .. }
| OperatorTree::PageRank { .. }
| OperatorTree::HITS { .. }
| OperatorTree::BetweennessCentrality { .. }
| OperatorTree::TemporalTraverse { .. }
| OperatorTree::TemporalPatternMatch { .. }
| OperatorTree::DeepPredict { .. } => {}
}
}
pub fn is_empty(&self) -> bool {
match self {
OperatorTree::Empty => true,
OperatorTree::Intersect(v) | OperatorTree::Union(v) | OperatorTree::Composed(v) => {
v.is_empty()
}
_ => false,
}
}
pub fn is_membership_only(&self) -> bool {
match self {
OperatorTree::Empty | OperatorTree::IndexScan { .. } => true,
OperatorTree::Filter { source, .. } => source
.as_deref()
.is_none_or(OperatorTree::is_membership_only),
OperatorTree::Intersect(children)
| OperatorTree::Union(children)
| OperatorTree::Composed(children) => {
children.iter().all(OperatorTree::is_membership_only)
}
OperatorTree::Complement(child) => child.is_membership_only(),
OperatorTree::VectorExclusion { positive, negative } => {
positive.is_membership_only() && negative.is_membership_only()
}
OperatorTree::Term { .. }
| OperatorTree::Facet { .. }
| OperatorTree::Score { .. }
| OperatorTree::BayesianScore { .. }
| OperatorTree::EncodeGraphPosting { .. }
| OperatorTree::BayesianMatchWithPrior { .. }
| OperatorTree::VectorSimilarity { .. }
| OperatorTree::KNN { .. }
| OperatorTree::CalibratedVectorMatch { .. }
| OperatorTree::CosineProbability(_)
| OperatorTree::BayesianEvidenceFusion { .. }
| OperatorTree::RobustPositiveEvidencePool { .. }
| OperatorTree::ProbBoolFusion { .. }
| OperatorTree::ProbNot { .. }
| OperatorTree::AttentionFusion { .. }
| OperatorTree::LearnedFusion { .. }
| OperatorTree::SparseThreshold { .. }
| OperatorTree::Traverse { .. }
| OperatorTree::GraphNeighbors { .. }
| OperatorTree::GraphEdges { .. }
| OperatorTree::PatternMatch { .. }
| OperatorTree::RegularPathQuery { .. }
| OperatorTree::GraphJoin { .. }
| OperatorTree::Aggregate { .. }
| OperatorTree::GroupBy { .. }
| OperatorTree::MultiStage { .. }
| OperatorTree::MultiFieldSearch { .. }
| OperatorTree::HybridTextVector { .. }
| OperatorTree::SemanticFilter { .. }
| OperatorTree::FacetVector { .. }
| OperatorTree::VertexAggregation { .. }
| OperatorTree::WeightedPathQuery { .. }
| OperatorTree::MessagePassing { .. }
| OperatorTree::GraphEmbedding { .. }
| OperatorTree::PageRank { .. }
| OperatorTree::HITS { .. }
| OperatorTree::BetweennessCentrality { .. }
| OperatorTree::TextSimilarityJoin { .. }
| OperatorTree::VectorSimilarityJoin { .. }
| OperatorTree::HybridJoin { .. }
| OperatorTree::CrossParadigmJoin { .. }
| OperatorTree::TemporalTraverse { .. }
| OperatorTree::TemporalPatternMatch { .. }
| OperatorTree::ProgressiveFusion { .. }
| OperatorTree::DeepFusion { .. }
| OperatorTree::DeepPredict { .. }
| OperatorTree::Opaque { .. } => false,
}
}
}
fn visit_operator_slice(children: &[OperatorTree], visitor: &mut impl FnMut(&OperatorTree)) {
for child in children {
child.visit(visitor);
}
}