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