1use std::collections::BTreeMap;
19use std::sync::Arc;
20
21pub use uqa_core::retrieval::{ExternalPriorMode, GatingSpec, MultiStageCutoff, TemporalFilterIR};
22use uqa_core::{Predicate, Value};
23
24use crate::aggregation::AggregationMonoid;
25use crate::base::Direction;
26
27pub type ScorerRef = Arc<dyn uqa_scoring::Scorer>;
31
32pub type AttentionRef = Arc<dyn AttentionFuserDyn>;
34
35pub type LearnedFusionRef = Arc<dyn LearnedFuserDyn>;
37
38pub type VertexPredicate = Arc<dyn Fn(&uqa_core::Vertex) -> bool + Send + Sync>;
40
41pub type PathWeightPredicate = Arc<dyn Fn(f64) -> bool + Send + Sync>;
43
44pub type VertexConstraint = Arc<dyn Fn(&uqa_core::Vertex) -> bool + Send + Sync>;
46
47pub type EdgeConstraint = Arc<dyn Fn(&uqa_core::Edge) -> bool + Send + Sync>;
49
50pub trait AttentionFuserDyn: Send + Sync {
51 fn validate_inputs(
52 &self,
53 signal_count: usize,
54 query_feature_count: usize,
55 ) -> Result<(), &'static str>;
56 fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str>;
57
58 fn fuse_batch(
59 &self,
60 probabilities: &[Vec<f64>],
61 query_features: &[f64],
62 ) -> Result<Vec<f64>, &'static str> {
63 probabilities
64 .iter()
65 .map(|sample| self.fuse(sample, query_features))
66 .collect()
67 }
68
69 fn head_count(&self) -> usize;
72 fn normalize(&self) -> bool;
73 fn alpha(&self) -> f64;
74 fn base_rate(&self) -> Option<f64>;
75}
76
77pub trait LearnedFuserDyn: Send + Sync {
78 fn validate_inputs(&self, signal_count: usize) -> Result<(), &'static str>;
79 fn fuse(&self, probs: &[f64]) -> Result<f64, &'static str>;
80}
81
82impl AttentionFuserDyn for uqa_fusion::AttentionFusion {
83 fn validate_inputs(
84 &self,
85 signal_count: usize,
86 query_feature_count: usize,
87 ) -> Result<(), &'static str> {
88 uqa_fusion::AttentionFusion::validate_inputs(self, signal_count, query_feature_count)
89 }
90
91 fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str> {
92 uqa_fusion::AttentionFusion::fuse(self, probs, query_features)
93 }
94
95 fn fuse_batch(
96 &self,
97 probabilities: &[Vec<f64>],
98 query_features: &[f64],
99 ) -> Result<Vec<f64>, &'static str> {
100 uqa_fusion::AttentionFusion::fuse_batch(self, probabilities, query_features)
101 }
102
103 fn head_count(&self) -> usize {
104 1
105 }
106
107 fn normalize(&self) -> bool {
108 self.normalize
109 }
110
111 fn alpha(&self) -> f64 {
112 self.alpha
113 }
114
115 fn base_rate(&self) -> Option<f64> {
116 self.base_rate
117 }
118}
119
120impl AttentionFuserDyn for uqa_fusion::MultiHeadAttentionFusion {
121 fn validate_inputs(
122 &self,
123 signal_count: usize,
124 query_feature_count: usize,
125 ) -> Result<(), &'static str> {
126 uqa_fusion::MultiHeadAttentionFusion::validate_inputs(
127 self,
128 signal_count,
129 query_feature_count,
130 )
131 }
132
133 fn fuse(&self, probs: &[f64], query_features: &[f64]) -> Result<f64, &'static str> {
134 uqa_fusion::MultiHeadAttentionFusion::fuse(self, probs, query_features)
135 }
136
137 fn fuse_batch(
138 &self,
139 probabilities: &[Vec<f64>],
140 query_features: &[f64],
141 ) -> Result<Vec<f64>, &'static str> {
142 uqa_fusion::MultiHeadAttentionFusion::fuse_batch(self, probabilities, query_features)
143 }
144
145 fn head_count(&self) -> usize {
146 uqa_fusion::MultiHeadAttentionFusion::n_heads(self)
147 }
148
149 fn normalize(&self) -> bool {
150 uqa_fusion::MultiHeadAttentionFusion::normalize(self)
151 }
152
153 fn alpha(&self) -> f64 {
154 uqa_fusion::MultiHeadAttentionFusion::alpha(self).unwrap_or(f64::NAN)
155 }
156
157 fn base_rate(&self) -> Option<f64> {
158 None
159 }
160}
161
162impl LearnedFuserDyn for uqa_fusion::LearnedFusion {
163 fn validate_inputs(&self, signal_count: usize) -> Result<(), &'static str> {
164 uqa_fusion::LearnedFusion::validate_inputs(self, signal_count)
165 }
166
167 fn fuse(&self, probs: &[f64]) -> Result<f64, &'static str> {
168 uqa_fusion::LearnedFusion::fuse(self, probs)
169 }
170}
171
172#[derive(Clone)]
174pub struct VertexPatternIR {
175 pub variable: String,
176 pub constraints: Vec<VertexConstraint>,
177 pub label: Option<String>,
179}
180
181impl std::fmt::Debug for VertexPatternIR {
182 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183 f.debug_struct("VertexPatternIR")
184 .field("variable", &self.variable)
185 .field("constraints_count", &self.constraints.len())
186 .field("label", &self.label)
187 .finish()
188 }
189}
190
191#[derive(Clone)]
192pub struct EdgePatternIR {
193 pub source_var: String,
194 pub target_var: String,
195 pub label: Option<String>,
196 pub constraints: Vec<EdgeConstraint>,
197}
198
199impl std::fmt::Debug for EdgePatternIR {
200 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201 f.debug_struct("EdgePatternIR")
202 .field("source_var", &self.source_var)
203 .field("target_var", &self.target_var)
204 .field("label", &self.label)
205 .field("constraints_count", &self.constraints.len())
206 .finish()
207 }
208}
209
210#[derive(Clone, Debug)]
211pub struct GraphPatternIR {
212 pub vertex_patterns: Vec<VertexPatternIR>,
213 pub edge_patterns: Vec<EdgePatternIR>,
214}
215
216#[derive(Clone, Copy, Debug, PartialEq, Eq)]
217pub enum ProbBoolMode {
218 And,
219 Or,
220}
221
222#[derive(Clone, Copy, Debug, PartialEq, Eq)]
226pub enum DeepFusionAggregation {
227 Mean,
228 Sum,
229 Max,
230}
231
232#[derive(Clone, Copy, Debug, PartialEq, Eq)]
234pub enum DeepFusionPoolMethod {
235 Average,
236 Max,
237}
238
239#[derive(Clone, Copy, Debug, PartialEq)]
241pub enum TextScoringMode {
242 BM25,
243 BayesianBM25,
244 CustomBM25(uqa_scoring::BM25Params),
246 CustomBayesianBM25(uqa_scoring::BayesianBM25Params),
248}
249
250#[derive(Clone, Copy, Debug, PartialEq, Eq)]
256pub enum TextTopKStrategy {
257 Wand,
259 BlockMaxWand,
261}
262
263#[derive(Clone, Copy, Debug, PartialEq, Eq)]
265pub struct TextTopKPlan {
266 pub k: usize,
267 pub strategy: TextTopKStrategy,
268}
269
270#[derive(Clone)]
272pub enum OperatorTree {
273 Empty,
276
277 Term {
279 query: String,
280 field: Option<String>,
281 scoring: Option<TextScoringMode>,
285 top_k: Option<TextTopKPlan>,
289 },
290 Phrase {
292 query: String,
293 field: Option<String>,
294 scoring: Option<TextScoringMode>,
295 },
296 Filter {
298 field: String,
299 predicate: Predicate,
300 source: Option<Box<OperatorTree>>,
301 },
302 Facet {
304 field: String,
305 source: Option<Box<OperatorTree>>,
306 },
307 Score {
309 scorer: ScorerRef,
310 source: Box<OperatorTree>,
311 query_terms: Vec<String>,
312 field: String,
313 },
314 BayesianScore {
318 source: Box<OperatorTree>,
319 field: Option<String>,
320 },
321 BayesianMatchWithPrior {
324 field: String,
325 query: String,
326 prior_field: String,
327 mode: ExternalPriorMode,
328 },
329
330 Intersect(Vec<OperatorTree>),
332 Union(Vec<OperatorTree>),
334 Complement(Box<OperatorTree>),
336 Composed(Vec<OperatorTree>),
338 EncodeGraphPosting { source: Box<OperatorTree> },
343
344 VectorSimilarity {
346 query_vector: Vec<f32>,
347 threshold: f32,
348 field: String,
349 },
350 KNN {
352 query_vector: Vec<f32>,
353 k: usize,
354 field: String,
355 },
356 CalibratedVectorMatch {
360 query_vector: Vec<f32>,
361 k: usize,
362 field: String,
363 threshold: Option<f64>,
364 },
365 CosineProbability(Box<OperatorTree>),
369
370 BayesianEvidenceFusion {
373 signals: Vec<OperatorTree>,
374 base_rate: Option<f64>,
375 },
376 RobustPositiveEvidencePool {
379 signals: Vec<OperatorTree>,
380 alpha: f64,
381 gating: GatingSpec,
382 weights: Option<Vec<f64>>,
383 logit_min: Option<Vec<f64>>,
384 logit_max: Option<Vec<f64>>,
385 adaptive_weights: bool,
387 },
388 ProbBoolFusion {
390 signals: Vec<OperatorTree>,
391 mode: ProbBoolMode,
392 },
393 ProbNot {
395 signal: Box<OperatorTree>,
396 default_prob: f64,
397 },
398 AttentionFusion {
400 signals: Vec<OperatorTree>,
401 attention: AttentionRef,
402 query_features: Vec<f64>,
403 },
404 LearnedFusion {
406 signals: Vec<OperatorTree>,
407 learned: LearnedFusionRef,
408 },
409 SparseThreshold {
411 source: Box<OperatorTree>,
412 threshold: f64,
413 },
414
415 Traverse {
417 start_vertex: u64,
418 graph: String,
419 label: Option<String>,
420 max_hops: usize,
421 vertex_predicate: Option<VertexPredicate>,
422 },
423 GraphNeighbors {
428 vertex: u64,
429 graph: String,
430 label: Option<String>,
431 direction: Direction,
432 },
433 GraphEdges {
436 graph: String,
437 label: Option<String>,
438 },
439 PatternMatch {
441 pattern: GraphPatternIR,
442 graph: String,
443 },
444 RegularPathQuery {
446 rpq_source: String,
447 start_vertex: u64,
448 graph: String,
449 },
450 GraphJoin {
452 left: Box<OperatorTree>,
453 right: Box<OperatorTree>,
454 label: Option<String>,
455 graph: String,
456 },
457
458 IndexScan {
461 index_name: String,
462 field: String,
463 predicate: Predicate,
464 },
465
466 Aggregate {
468 source: Option<Box<OperatorTree>>,
469 field: String,
470 monoid: Arc<dyn AggregationMonoid>,
471 },
472 GroupBy {
474 source: Box<OperatorTree>,
475 group_field: String,
476 agg_field: String,
477 monoid: Arc<dyn AggregationMonoid>,
478 },
479
480 MultiStage { stages: Vec<MultiStageEntry> },
486 MultiFieldSearch {
488 fields: Vec<String>,
489 queries: Vec<String>,
490 weights: Option<Vec<f64>>,
491 },
492 HybridTextVector {
494 term_op: Box<OperatorTree>,
495 vector_op: Box<OperatorTree>,
496 alpha: f64,
497 },
498 SemanticFilter {
500 source: Box<OperatorTree>,
501 vector_op: Box<OperatorTree>,
502 },
503 VectorExclusion {
505 positive: Box<OperatorTree>,
506 negative: Box<OperatorTree>,
507 },
508 FacetVector {
510 vector_op: Box<OperatorTree>,
511 facet_field: String,
512 },
513 VertexAggregation {
515 source: Box<OperatorTree>,
516 monoid: Arc<dyn AggregationMonoid>,
517 },
518 WeightedPathQuery {
522 rpq_source: String,
523 start_vertex: u64,
524 graph: String,
525 weight_property: String,
526 default_edge_weight: f64,
527 max_hops: usize,
528 predicate: PathWeightPredicate,
529 predicate_selectivity: f64,
530 score: f64,
531 },
532 MessagePassing { source: Box<OperatorTree> },
534 GraphEmbedding { source: Box<OperatorTree> },
536 PageRank { graph: String },
538 HITS { graph: String },
540 BetweennessCentrality { graph: String },
542 TextSimilarityJoin {
544 left: Box<OperatorTree>,
545 right: Box<OperatorTree>,
546 threshold: f64,
547 },
548 VectorSimilarityJoin {
550 left: Box<OperatorTree>,
551 right: Box<OperatorTree>,
552 threshold: f64,
553 },
554 HybridJoin {
556 left: Box<OperatorTree>,
557 right: Box<OperatorTree>,
558 },
559 CrossParadigmJoin {
563 left: Box<OperatorTree>,
564 right: Box<OperatorTree>,
565 },
566 TemporalTraverse {
568 start_vertex: u64,
569 graph: String,
570 label: Option<String>,
571 max_hops: usize,
572 temporal_filter: Option<TemporalFilterIR>,
573 },
574 TemporalPatternMatch {
576 pattern: GraphPatternIR,
577 graph: String,
578 temporal_filter: Option<TemporalFilterIR>,
579 },
580 ProgressiveFusion {
583 stages: Vec<ProgressiveFusionEntry>,
584 alpha: f64,
585 gating: GatingSpec,
586 },
587 DeepFusion {
589 layers: Vec<DeepFusionLayer>,
590 alpha: f64,
591 gating: GatingSpec,
592 },
593
594 DeepPredict { model: String },
596
597 Opaque {
599 kind: String,
600 children: Vec<OperatorTree>,
601 meta: BTreeMap<String, Value>,
602 },
603}
604
605#[derive(Clone)]
608pub struct MultiStageEntry {
609 pub child: OperatorTree,
610 pub cutoff: MultiStageCutoff,
611}
612
613#[derive(Clone)]
615pub struct ProgressiveFusionEntry {
616 pub signal: OperatorTree,
617 pub k: usize,
618}
619
620#[derive(Clone)]
622pub enum DeepFusionLayer {
623 Signal {
624 signals: Vec<OperatorTree>,
625 },
626 Propagate {
627 edge_label: Option<String>,
628 aggregation: DeepFusionAggregation,
629 direction: Direction,
630 },
631 Conv {
632 edge_label: Option<String>,
633 hop_weights: Vec<f64>,
635 direction: Direction,
636 },
637 Pool {
638 edge_label: Option<String>,
639 pool_size: usize,
640 method: DeepFusionPoolMethod,
641 direction: Direction,
642 },
643 Flatten,
644 Dense {
645 weights: Vec<f64>,
647 bias: Vec<f64>,
648 output_channels: usize,
649 input_channels: usize,
650 },
651 Softmax,
652 BatchNorm {
653 epsilon: f64,
654 },
655 Dropout {
656 probability: f64,
657 },
658}
659
660impl OperatorTree {
661 #[expect(
667 clippy::too_many_lines,
668 reason = "preserves exhaustive IR variant order"
669 )]
670 pub fn visit(&self, visitor: &mut impl FnMut(&OperatorTree)) {
671 visitor(self);
672 match self {
673 OperatorTree::Filter {
674 source: Some(source),
675 ..
676 }
677 | OperatorTree::Facet {
678 source: Some(source),
679 ..
680 }
681 | OperatorTree::Score { source, .. }
682 | OperatorTree::BayesianScore { source, .. }
683 | OperatorTree::Complement(source)
684 | OperatorTree::EncodeGraphPosting { source }
685 | OperatorTree::CosineProbability(source)
686 | OperatorTree::ProbNot { signal: source, .. }
687 | OperatorTree::SparseThreshold { source, .. }
688 | OperatorTree::VertexAggregation { source, .. }
689 | OperatorTree::MessagePassing { source }
690 | OperatorTree::GraphEmbedding { source }
691 | OperatorTree::GroupBy { source, .. }
692 | OperatorTree::Aggregate {
693 source: Some(source),
694 ..
695 } => source.visit(visitor),
696 OperatorTree::Intersect(children)
697 | OperatorTree::Union(children)
698 | OperatorTree::Composed(children)
699 | OperatorTree::Opaque { children, .. }
700 | OperatorTree::BayesianEvidenceFusion {
701 signals: children, ..
702 }
703 | OperatorTree::RobustPositiveEvidencePool {
704 signals: children, ..
705 }
706 | OperatorTree::ProbBoolFusion {
707 signals: children, ..
708 }
709 | OperatorTree::AttentionFusion {
710 signals: children, ..
711 }
712 | OperatorTree::LearnedFusion {
713 signals: children, ..
714 } => visit_operator_slice(children, visitor),
715 OperatorTree::GraphJoin { left, right, .. }
716 | OperatorTree::TextSimilarityJoin { left, right, .. }
717 | OperatorTree::VectorSimilarityJoin { left, right, .. }
718 | OperatorTree::HybridJoin { left, right }
719 | OperatorTree::CrossParadigmJoin { left, right }
720 | OperatorTree::HybridTextVector {
721 term_op: left,
722 vector_op: right,
723 ..
724 }
725 | OperatorTree::SemanticFilter {
726 source: left,
727 vector_op: right,
728 }
729 | OperatorTree::VectorExclusion {
730 positive: left,
731 negative: right,
732 } => {
733 left.visit(visitor);
734 right.visit(visitor);
735 }
736 OperatorTree::FacetVector { vector_op, .. } => vector_op.visit(visitor),
737 OperatorTree::MultiStage { stages } => {
738 for stage in stages {
739 stage.child.visit(visitor);
740 }
741 }
742 OperatorTree::ProgressiveFusion { stages, .. } => {
743 for stage in stages {
744 stage.signal.visit(visitor);
745 }
746 }
747 OperatorTree::DeepFusion { layers, .. } => {
748 for layer in layers {
749 if let DeepFusionLayer::Signal { signals } = layer {
750 for signal in signals {
751 signal.visit(visitor);
752 }
753 }
754 }
755 }
756 OperatorTree::Empty
757 | OperatorTree::Term { .. }
758 | OperatorTree::Phrase { .. }
759 | OperatorTree::BayesianMatchWithPrior { .. }
760 | OperatorTree::Filter { source: None, .. }
761 | OperatorTree::Facet { source: None, .. }
762 | OperatorTree::VectorSimilarity { .. }
763 | OperatorTree::KNN { .. }
764 | OperatorTree::CalibratedVectorMatch { .. }
765 | OperatorTree::Traverse { .. }
766 | OperatorTree::GraphNeighbors { .. }
767 | OperatorTree::GraphEdges { .. }
768 | OperatorTree::PatternMatch { .. }
769 | OperatorTree::RegularPathQuery { .. }
770 | OperatorTree::IndexScan { .. }
771 | OperatorTree::Aggregate { source: None, .. }
772 | OperatorTree::MultiFieldSearch { .. }
773 | OperatorTree::WeightedPathQuery { .. }
774 | OperatorTree::PageRank { .. }
775 | OperatorTree::HITS { .. }
776 | OperatorTree::BetweennessCentrality { .. }
777 | OperatorTree::TemporalTraverse { .. }
778 | OperatorTree::TemporalPatternMatch { .. }
779 | OperatorTree::DeepPredict { .. } => {}
780 }
781 }
782
783 pub fn is_empty(&self) -> bool {
787 match self {
788 OperatorTree::Empty => true,
789 OperatorTree::Intersect(v) | OperatorTree::Union(v) | OperatorTree::Composed(v) => {
790 v.is_empty()
791 }
792 _ => false,
793 }
794 }
795
796 pub fn is_membership_only(&self) -> bool {
803 match self {
804 OperatorTree::Empty | OperatorTree::IndexScan { .. } => true,
805 OperatorTree::Filter { source, .. } => source
806 .as_deref()
807 .is_none_or(OperatorTree::is_membership_only),
808 OperatorTree::Intersect(children)
809 | OperatorTree::Union(children)
810 | OperatorTree::Composed(children) => {
811 children.iter().all(OperatorTree::is_membership_only)
812 }
813 OperatorTree::Complement(child) => child.is_membership_only(),
814 OperatorTree::VectorExclusion { positive, negative } => {
815 positive.is_membership_only() && negative.is_membership_only()
816 }
817 OperatorTree::Term { .. }
818 | OperatorTree::Phrase { .. }
819 | OperatorTree::Facet { .. }
820 | OperatorTree::Score { .. }
821 | OperatorTree::BayesianScore { .. }
822 | OperatorTree::EncodeGraphPosting { .. }
823 | OperatorTree::BayesianMatchWithPrior { .. }
824 | OperatorTree::VectorSimilarity { .. }
825 | OperatorTree::KNN { .. }
826 | OperatorTree::CalibratedVectorMatch { .. }
827 | OperatorTree::CosineProbability(_)
828 | OperatorTree::BayesianEvidenceFusion { .. }
829 | OperatorTree::RobustPositiveEvidencePool { .. }
830 | OperatorTree::ProbBoolFusion { .. }
831 | OperatorTree::ProbNot { .. }
832 | OperatorTree::AttentionFusion { .. }
833 | OperatorTree::LearnedFusion { .. }
834 | OperatorTree::SparseThreshold { .. }
835 | OperatorTree::Traverse { .. }
836 | OperatorTree::GraphNeighbors { .. }
837 | OperatorTree::GraphEdges { .. }
838 | OperatorTree::PatternMatch { .. }
839 | OperatorTree::RegularPathQuery { .. }
840 | OperatorTree::GraphJoin { .. }
841 | OperatorTree::Aggregate { .. }
842 | OperatorTree::GroupBy { .. }
843 | OperatorTree::MultiStage { .. }
844 | OperatorTree::MultiFieldSearch { .. }
845 | OperatorTree::HybridTextVector { .. }
846 | OperatorTree::SemanticFilter { .. }
847 | OperatorTree::FacetVector { .. }
848 | OperatorTree::VertexAggregation { .. }
849 | OperatorTree::WeightedPathQuery { .. }
850 | OperatorTree::MessagePassing { .. }
851 | OperatorTree::GraphEmbedding { .. }
852 | OperatorTree::PageRank { .. }
853 | OperatorTree::HITS { .. }
854 | OperatorTree::BetweennessCentrality { .. }
855 | OperatorTree::TextSimilarityJoin { .. }
856 | OperatorTree::VectorSimilarityJoin { .. }
857 | OperatorTree::HybridJoin { .. }
858 | OperatorTree::CrossParadigmJoin { .. }
859 | OperatorTree::TemporalTraverse { .. }
860 | OperatorTree::TemporalPatternMatch { .. }
861 | OperatorTree::ProgressiveFusion { .. }
862 | OperatorTree::DeepFusion { .. }
863 | OperatorTree::DeepPredict { .. }
864 | OperatorTree::Opaque { .. } => false,
865 }
866 }
867}
868
869fn visit_operator_slice(children: &[OperatorTree], visitor: &mut impl FnMut(&OperatorTree)) {
870 for child in children {
871 child.visit(visitor);
872 }
873}
874
875mod introspection;
876pub use introspection::collect_graph_names;