1use std::collections::BTreeMap;
19use std::sync::Arc;
20
21use uqa_core::{Predicate, Value};
22
23use crate::aggregation::AggregationMonoid;
24use crate::base::Direction;
25
26pub type ScorerRef = Arc<dyn uqa_scoring::Scorer>;
30
31pub type AttentionRef = Arc<dyn AttentionFuserDyn>;
33
34pub type LearnedFusionRef = Arc<dyn LearnedFuserDyn>;
36
37pub type VertexPredicate = Arc<dyn Fn(&uqa_core::Vertex) -> bool + Send + Sync>;
39
40pub type PathWeightPredicate = Arc<dyn Fn(f64) -> bool + Send + Sync>;
42
43pub type VertexConstraint = Arc<dyn Fn(&uqa_core::Vertex) -> bool + Send + Sync>;
45
46pub type EdgeConstraint = Arc<dyn Fn(&uqa_core::Edge) -> bool + Send + Sync>;
48
49pub trait AttentionFuserDyn: Send + Sync {
50 fn validate_inputs(
51 &self,
52 signal_count: usize,
53 query_feature_count: usize,
54 ) -> Result<(), &'static str>;
55 fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str>;
56
57 fn fuse_batch(
58 &self,
59 probabilities: &[Vec<f64>],
60 query_features: &[f64],
61 ) -> Result<Vec<f64>, &'static str> {
62 probabilities
63 .iter()
64 .map(|sample| self.fuse(sample, query_features))
65 .collect()
66 }
67
68 fn head_count(&self) -> usize;
71 fn normalize(&self) -> bool;
72 fn alpha(&self) -> f64;
73 fn base_rate(&self) -> Option<f64>;
74}
75
76pub trait LearnedFuserDyn: Send + Sync {
77 fn validate_inputs(&self, signal_count: usize) -> Result<(), &'static str>;
78 fn fuse(&self, probs: &[f64]) -> Result<f64, &'static str>;
79}
80
81impl AttentionFuserDyn for uqa_fusion::AttentionFusion {
82 fn validate_inputs(
83 &self,
84 signal_count: usize,
85 query_feature_count: usize,
86 ) -> Result<(), &'static str> {
87 uqa_fusion::AttentionFusion::validate_inputs(self, signal_count, query_feature_count)
88 }
89
90 fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str> {
91 uqa_fusion::AttentionFusion::fuse(self, probs, query_features)
92 }
93
94 fn fuse_batch(
95 &self,
96 probabilities: &[Vec<f64>],
97 query_features: &[f64],
98 ) -> Result<Vec<f64>, &'static str> {
99 uqa_fusion::AttentionFusion::fuse_batch(self, probabilities, query_features)
100 }
101
102 fn head_count(&self) -> usize {
103 1
104 }
105
106 fn normalize(&self) -> bool {
107 self.normalize
108 }
109
110 fn alpha(&self) -> f64 {
111 self.alpha
112 }
113
114 fn base_rate(&self) -> Option<f64> {
115 self.base_rate
116 }
117}
118
119impl AttentionFuserDyn for uqa_fusion::MultiHeadAttentionFusion {
120 fn validate_inputs(
121 &self,
122 signal_count: usize,
123 query_feature_count: usize,
124 ) -> Result<(), &'static str> {
125 uqa_fusion::MultiHeadAttentionFusion::validate_inputs(
126 self,
127 signal_count,
128 query_feature_count,
129 )
130 }
131
132 fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str> {
133 uqa_fusion::MultiHeadAttentionFusion::fuse(self, probs, query_features)
134 }
135
136 fn fuse_batch(
137 &self,
138 probabilities: &[Vec<f64>],
139 query_features: &[f64],
140 ) -> Result<Vec<f64>, &'static str> {
141 uqa_fusion::MultiHeadAttentionFusion::fuse_batch(self, probabilities, query_features)
142 }
143
144 fn head_count(&self) -> usize {
145 uqa_fusion::MultiHeadAttentionFusion::n_heads(self)
146 }
147
148 fn normalize(&self) -> bool {
149 uqa_fusion::MultiHeadAttentionFusion::normalize(self)
150 }
151
152 fn alpha(&self) -> f64 {
153 uqa_fusion::MultiHeadAttentionFusion::alpha(self).unwrap_or(f64::NAN)
154 }
155
156 fn base_rate(&self) -> Option<f64> {
157 None
158 }
159}
160
161impl LearnedFuserDyn for uqa_fusion::LearnedFusion {
162 fn validate_inputs(&self, signal_count: usize) -> Result<(), &'static str> {
163 uqa_fusion::LearnedFusion::validate_inputs(self, signal_count)
164 }
165
166 fn fuse(&self, probs: &[f64]) -> Result<f64, &'static str> {
167 uqa_fusion::LearnedFusion::fuse(self, probs)
168 }
169}
170
171#[derive(Clone)]
173pub struct VertexPatternIR {
174 pub variable: String,
175 pub constraints: Vec<VertexConstraint>,
176 pub label: Option<String>,
178}
179
180impl std::fmt::Debug for VertexPatternIR {
181 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
182 f.debug_struct("VertexPatternIR")
183 .field("variable", &self.variable)
184 .field("constraints_count", &self.constraints.len())
185 .field("label", &self.label)
186 .finish()
187 }
188}
189
190#[derive(Clone)]
191pub struct EdgePatternIR {
192 pub source_var: String,
193 pub target_var: String,
194 pub label: Option<String>,
195 pub constraints: Vec<EdgeConstraint>,
196}
197
198impl std::fmt::Debug for EdgePatternIR {
199 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
200 f.debug_struct("EdgePatternIR")
201 .field("source_var", &self.source_var)
202 .field("target_var", &self.target_var)
203 .field("label", &self.label)
204 .field("constraints_count", &self.constraints.len())
205 .finish()
206 }
207}
208
209#[derive(Clone, Debug)]
210pub struct GraphPatternIR {
211 pub vertex_patterns: Vec<VertexPatternIR>,
212 pub edge_patterns: Vec<EdgePatternIR>,
213}
214
215#[derive(Clone, Copy, Debug, PartialEq, Eq)]
216pub enum ProbBoolMode {
217 And,
218 Or,
219}
220
221#[derive(Clone, Debug)]
222pub enum GatingSpec {
223 Softplus,
225 Pass,
227 Sigmoid { feature: String },
229 ReLU,
231 Swish,
233 Gelu,
235}
236
237#[derive(Clone, Copy, Debug, PartialEq, Eq)]
241pub enum DeepFusionAggregation {
242 Mean,
243 Sum,
244 Max,
245}
246
247#[derive(Clone, Copy, Debug, PartialEq, Eq)]
249pub enum DeepFusionPoolMethod {
250 Average,
251 Max,
252}
253
254#[derive(Clone, Copy, Debug, PartialEq)]
256pub enum TextScoringMode {
257 BM25,
258 BayesianBM25,
259 CustomBM25(uqa_scoring::BM25Params),
261 CustomBayesianBM25(uqa_scoring::BayesianBM25Params),
263}
264
265#[derive(Clone, Copy, Debug, PartialEq, Eq)]
271pub enum TextTopKStrategy {
272 Wand,
274 BlockMaxWand,
276}
277
278#[derive(Clone, Copy, Debug, PartialEq, Eq)]
280pub struct TextTopKPlan {
281 pub k: usize,
282 pub strategy: TextTopKStrategy,
283}
284
285#[derive(Clone, Copy, Debug, PartialEq, Eq)]
287pub enum ExternalPriorMode {
288 Authority,
289 Recency,
290}
291
292#[derive(Clone)]
294pub enum OperatorTree {
295 Empty,
298
299 Term {
301 query: String,
302 field: Option<String>,
303 scoring: Option<TextScoringMode>,
307 top_k: Option<TextTopKPlan>,
311 },
312 Filter {
314 field: String,
315 predicate: Predicate,
316 source: Option<Box<OperatorTree>>,
317 },
318 Facet {
320 field: String,
321 source: Option<Box<OperatorTree>>,
322 },
323 Score {
325 scorer: ScorerRef,
326 source: Box<OperatorTree>,
327 query_terms: Vec<String>,
328 field: String,
329 },
330 BayesianScore {
334 source: Box<OperatorTree>,
335 field: Option<String>,
336 },
337 BayesianMatchWithPrior {
340 field: String,
341 query: String,
342 prior_field: String,
343 mode: ExternalPriorMode,
344 },
345
346 Intersect(Vec<OperatorTree>),
348 Union(Vec<OperatorTree>),
350 Complement(Box<OperatorTree>),
352 Composed(Vec<OperatorTree>),
354 EncodeGraphPosting { source: Box<OperatorTree> },
359
360 VectorSimilarity {
362 query_vector: Vec<f32>,
363 threshold: f32,
364 field: String,
365 },
366 KNN {
368 query_vector: Vec<f32>,
369 k: usize,
370 field: String,
371 },
372 CalibratedVectorMatch {
376 query_vector: Vec<f32>,
377 k: usize,
378 field: String,
379 threshold: Option<f64>,
380 },
381 CosineProbability(Box<OperatorTree>),
385
386 BayesianEvidenceFusion {
389 signals: Vec<OperatorTree>,
390 base_rate: Option<f64>,
391 },
392 RobustPositiveEvidencePool {
395 signals: Vec<OperatorTree>,
396 alpha: f64,
397 gating: GatingSpec,
398 weights: Option<Vec<f64>>,
399 logit_min: Option<Vec<f64>>,
400 logit_max: Option<Vec<f64>>,
401 adaptive_weights: bool,
403 },
404 ProbBoolFusion {
406 signals: Vec<OperatorTree>,
407 mode: ProbBoolMode,
408 },
409 ProbNot {
411 signal: Box<OperatorTree>,
412 default_prob: f64,
413 },
414 AttentionFusion {
416 signals: Vec<OperatorTree>,
417 attention: AttentionRef,
418 query_features: Vec<f64>,
419 },
420 LearnedFusion {
422 signals: Vec<OperatorTree>,
423 learned: LearnedFusionRef,
424 },
425 SparseThreshold {
427 source: Box<OperatorTree>,
428 threshold: f64,
429 },
430
431 Traverse {
433 start_vertex: u64,
434 graph: String,
435 label: Option<String>,
436 max_hops: usize,
437 vertex_predicate: Option<VertexPredicate>,
438 },
439 GraphNeighbors {
444 vertex: u64,
445 graph: String,
446 label: Option<String>,
447 direction: Direction,
448 },
449 GraphEdges {
452 graph: String,
453 label: Option<String>,
454 },
455 PatternMatch {
457 pattern: GraphPatternIR,
458 graph: String,
459 },
460 RegularPathQuery {
462 rpq_source: String,
463 start_vertex: u64,
464 graph: String,
465 },
466 GraphJoin {
468 left: Box<OperatorTree>,
469 right: Box<OperatorTree>,
470 label: Option<String>,
471 graph: String,
472 },
473
474 IndexScan {
477 index_name: String,
478 field: String,
479 predicate: Predicate,
480 },
481
482 Aggregate {
484 source: Option<Box<OperatorTree>>,
485 field: String,
486 monoid: Arc<dyn AggregationMonoid>,
487 },
488 GroupBy {
490 source: Box<OperatorTree>,
491 group_field: String,
492 agg_field: String,
493 monoid: Arc<dyn AggregationMonoid>,
494 },
495
496 MultiStage { stages: Vec<MultiStageEntry> },
502 MultiFieldSearch {
504 fields: Vec<String>,
505 queries: Vec<String>,
506 weights: Option<Vec<f64>>,
507 },
508 HybridTextVector {
510 term_op: Box<OperatorTree>,
511 vector_op: Box<OperatorTree>,
512 alpha: f64,
513 },
514 SemanticFilter {
516 source: Box<OperatorTree>,
517 vector_op: Box<OperatorTree>,
518 },
519 VectorExclusion {
521 positive: Box<OperatorTree>,
522 negative: Box<OperatorTree>,
523 },
524 FacetVector {
526 vector_op: Box<OperatorTree>,
527 facet_field: String,
528 },
529 VertexAggregation {
531 source: Box<OperatorTree>,
532 monoid: Arc<dyn AggregationMonoid>,
533 },
534 WeightedPathQuery {
538 rpq_source: String,
539 start_vertex: u64,
540 graph: String,
541 weight_property: String,
542 default_edge_weight: f64,
543 max_hops: usize,
544 predicate: PathWeightPredicate,
545 predicate_selectivity: f64,
546 score: f64,
547 },
548 MessagePassing { source: Box<OperatorTree> },
550 GraphEmbedding { source: Box<OperatorTree> },
552 PageRank { graph: String },
554 HITS { graph: String },
556 BetweennessCentrality { graph: String },
558 TextSimilarityJoin {
560 left: Box<OperatorTree>,
561 right: Box<OperatorTree>,
562 threshold: f64,
563 },
564 VectorSimilarityJoin {
566 left: Box<OperatorTree>,
567 right: Box<OperatorTree>,
568 threshold: f64,
569 },
570 HybridJoin {
572 left: Box<OperatorTree>,
573 right: Box<OperatorTree>,
574 },
575 CrossParadigmJoin {
579 left: Box<OperatorTree>,
580 right: Box<OperatorTree>,
581 },
582 TemporalTraverse {
584 start_vertex: u64,
585 graph: String,
586 label: Option<String>,
587 max_hops: usize,
588 temporal_filter: Option<TemporalFilterIR>,
589 },
590 TemporalPatternMatch {
592 pattern: GraphPatternIR,
593 graph: String,
594 temporal_filter: Option<TemporalFilterIR>,
595 },
596 ProgressiveFusion {
599 stages: Vec<ProgressiveFusionEntry>,
600 alpha: f64,
601 gating: GatingSpec,
602 },
603 DeepFusion {
605 layers: Vec<DeepFusionLayer>,
606 alpha: f64,
607 gating: GatingSpec,
608 },
609
610 DeepPredict { model: String },
612
613 Opaque {
615 kind: String,
616 children: Vec<OperatorTree>,
617 meta: BTreeMap<String, Value>,
618 },
619}
620
621#[derive(Clone)]
624pub struct MultiStageEntry {
625 pub child: OperatorTree,
626 pub cutoff: MultiStageCutoff,
627}
628
629#[derive(Clone, Copy, Debug, PartialEq)]
631pub enum MultiStageCutoff {
632 TopK(usize),
634 Ratio(f64),
636}
637
638#[derive(Clone)]
640pub struct ProgressiveFusionEntry {
641 pub signal: OperatorTree,
642 pub k: usize,
643}
644
645#[derive(Clone)]
647pub enum DeepFusionLayer {
648 Signal {
649 signals: Vec<OperatorTree>,
650 },
651 Propagate {
652 edge_label: Option<String>,
653 aggregation: DeepFusionAggregation,
654 direction: Direction,
655 },
656 Conv {
657 edge_label: Option<String>,
658 hop_weights: Vec<f64>,
660 direction: Direction,
661 },
662 Pool {
663 edge_label: Option<String>,
664 pool_size: usize,
665 method: DeepFusionPoolMethod,
666 direction: Direction,
667 },
668 Flatten,
669 Dense {
670 weights: Vec<f64>,
672 bias: Vec<f64>,
673 output_channels: usize,
674 input_channels: usize,
675 },
676 Softmax,
677 BatchNorm {
678 epsilon: f64,
679 },
680 Dropout {
681 probability: f64,
682 },
683}
684
685#[derive(Clone, Debug, Default)]
689pub struct TemporalFilterIR {
690 pub timestamp: Option<f64>,
691 pub time_range: Option<(f64, f64)>,
692}
693
694impl OperatorTree {
695 #[expect(
701 clippy::too_many_lines,
702 reason = "preserves exhaustive IR variant order"
703 )]
704 pub fn visit(&self, visitor: &mut impl FnMut(&OperatorTree)) {
705 visitor(self);
706 match self {
707 OperatorTree::Filter {
708 source: Some(source),
709 ..
710 }
711 | OperatorTree::Facet {
712 source: Some(source),
713 ..
714 }
715 | OperatorTree::Score { source, .. }
716 | OperatorTree::BayesianScore { source, .. }
717 | OperatorTree::Complement(source)
718 | OperatorTree::EncodeGraphPosting { source }
719 | OperatorTree::CosineProbability(source)
720 | OperatorTree::ProbNot { signal: source, .. }
721 | OperatorTree::SparseThreshold { source, .. }
722 | OperatorTree::VertexAggregation { source, .. }
723 | OperatorTree::MessagePassing { source }
724 | OperatorTree::GraphEmbedding { source }
725 | OperatorTree::GroupBy { source, .. }
726 | OperatorTree::Aggregate {
727 source: Some(source),
728 ..
729 } => source.visit(visitor),
730 OperatorTree::Intersect(children)
731 | OperatorTree::Union(children)
732 | OperatorTree::Composed(children)
733 | OperatorTree::Opaque { children, .. }
734 | OperatorTree::BayesianEvidenceFusion {
735 signals: children, ..
736 }
737 | OperatorTree::RobustPositiveEvidencePool {
738 signals: children, ..
739 }
740 | OperatorTree::ProbBoolFusion {
741 signals: children, ..
742 }
743 | OperatorTree::AttentionFusion {
744 signals: children, ..
745 }
746 | OperatorTree::LearnedFusion {
747 signals: children, ..
748 } => visit_operator_slice(children, visitor),
749 OperatorTree::GraphJoin { left, right, .. }
750 | OperatorTree::TextSimilarityJoin { left, right, .. }
751 | OperatorTree::VectorSimilarityJoin { left, right, .. }
752 | OperatorTree::HybridJoin { left, right }
753 | OperatorTree::CrossParadigmJoin { left, right }
754 | OperatorTree::HybridTextVector {
755 term_op: left,
756 vector_op: right,
757 ..
758 }
759 | OperatorTree::SemanticFilter {
760 source: left,
761 vector_op: right,
762 }
763 | OperatorTree::VectorExclusion {
764 positive: left,
765 negative: right,
766 } => {
767 left.visit(visitor);
768 right.visit(visitor);
769 }
770 OperatorTree::FacetVector { vector_op, .. } => vector_op.visit(visitor),
771 OperatorTree::MultiStage { stages } => {
772 for stage in stages {
773 stage.child.visit(visitor);
774 }
775 }
776 OperatorTree::ProgressiveFusion { stages, .. } => {
777 for stage in stages {
778 stage.signal.visit(visitor);
779 }
780 }
781 OperatorTree::DeepFusion { layers, .. } => {
782 for layer in layers {
783 if let DeepFusionLayer::Signal { signals } = layer {
784 for signal in signals {
785 signal.visit(visitor);
786 }
787 }
788 }
789 }
790 OperatorTree::Empty
791 | OperatorTree::Term { .. }
792 | OperatorTree::BayesianMatchWithPrior { .. }
793 | OperatorTree::Filter { source: None, .. }
794 | OperatorTree::Facet { source: None, .. }
795 | OperatorTree::VectorSimilarity { .. }
796 | OperatorTree::KNN { .. }
797 | OperatorTree::CalibratedVectorMatch { .. }
798 | OperatorTree::Traverse { .. }
799 | OperatorTree::GraphNeighbors { .. }
800 | OperatorTree::GraphEdges { .. }
801 | OperatorTree::PatternMatch { .. }
802 | OperatorTree::RegularPathQuery { .. }
803 | OperatorTree::IndexScan { .. }
804 | OperatorTree::Aggregate { source: None, .. }
805 | OperatorTree::MultiFieldSearch { .. }
806 | OperatorTree::WeightedPathQuery { .. }
807 | OperatorTree::PageRank { .. }
808 | OperatorTree::HITS { .. }
809 | OperatorTree::BetweennessCentrality { .. }
810 | OperatorTree::TemporalTraverse { .. }
811 | OperatorTree::TemporalPatternMatch { .. }
812 | OperatorTree::DeepPredict { .. } => {}
813 }
814 }
815
816 pub fn is_empty(&self) -> bool {
820 match self {
821 OperatorTree::Empty => true,
822 OperatorTree::Intersect(v) | OperatorTree::Union(v) | OperatorTree::Composed(v) => {
823 v.is_empty()
824 }
825 _ => false,
826 }
827 }
828
829 pub fn is_membership_only(&self) -> bool {
836 match self {
837 OperatorTree::Empty | OperatorTree::IndexScan { .. } => true,
838 OperatorTree::Filter { source, .. } => source
839 .as_deref()
840 .is_none_or(OperatorTree::is_membership_only),
841 OperatorTree::Intersect(children)
842 | OperatorTree::Union(children)
843 | OperatorTree::Composed(children) => {
844 children.iter().all(OperatorTree::is_membership_only)
845 }
846 OperatorTree::Complement(child) => child.is_membership_only(),
847 OperatorTree::VectorExclusion { positive, negative } => {
848 positive.is_membership_only() && negative.is_membership_only()
849 }
850 OperatorTree::Term { .. }
851 | OperatorTree::Facet { .. }
852 | OperatorTree::Score { .. }
853 | OperatorTree::BayesianScore { .. }
854 | OperatorTree::EncodeGraphPosting { .. }
855 | OperatorTree::BayesianMatchWithPrior { .. }
856 | OperatorTree::VectorSimilarity { .. }
857 | OperatorTree::KNN { .. }
858 | OperatorTree::CalibratedVectorMatch { .. }
859 | OperatorTree::CosineProbability(_)
860 | OperatorTree::BayesianEvidenceFusion { .. }
861 | OperatorTree::RobustPositiveEvidencePool { .. }
862 | OperatorTree::ProbBoolFusion { .. }
863 | OperatorTree::ProbNot { .. }
864 | OperatorTree::AttentionFusion { .. }
865 | OperatorTree::LearnedFusion { .. }
866 | OperatorTree::SparseThreshold { .. }
867 | OperatorTree::Traverse { .. }
868 | OperatorTree::GraphNeighbors { .. }
869 | OperatorTree::GraphEdges { .. }
870 | OperatorTree::PatternMatch { .. }
871 | OperatorTree::RegularPathQuery { .. }
872 | OperatorTree::GraphJoin { .. }
873 | OperatorTree::Aggregate { .. }
874 | OperatorTree::GroupBy { .. }
875 | OperatorTree::MultiStage { .. }
876 | OperatorTree::MultiFieldSearch { .. }
877 | OperatorTree::HybridTextVector { .. }
878 | OperatorTree::SemanticFilter { .. }
879 | OperatorTree::FacetVector { .. }
880 | OperatorTree::VertexAggregation { .. }
881 | OperatorTree::WeightedPathQuery { .. }
882 | OperatorTree::MessagePassing { .. }
883 | OperatorTree::GraphEmbedding { .. }
884 | OperatorTree::PageRank { .. }
885 | OperatorTree::HITS { .. }
886 | OperatorTree::BetweennessCentrality { .. }
887 | OperatorTree::TextSimilarityJoin { .. }
888 | OperatorTree::VectorSimilarityJoin { .. }
889 | OperatorTree::HybridJoin { .. }
890 | OperatorTree::CrossParadigmJoin { .. }
891 | OperatorTree::TemporalTraverse { .. }
892 | OperatorTree::TemporalPatternMatch { .. }
893 | OperatorTree::ProgressiveFusion { .. }
894 | OperatorTree::DeepFusion { .. }
895 | OperatorTree::DeepPredict { .. }
896 | OperatorTree::Opaque { .. } => false,
897 }
898 }
899}
900
901fn visit_operator_slice(children: &[OperatorTree], visitor: &mut impl FnMut(&OperatorTree)) {
902 for child in children {
903 child.visit(visitor);
904 }
905}