1use std::collections::BTreeMap;
11use std::sync::Arc;
12
13use uqa_core::{FieldName, IndexStats, PostingList};
14use uqa_storage::{
15 DocumentStore, InvertedIndex, StorageBackendError, StorageBackendResult, VectorIndex,
16};
17
18pub type OperatorResult = StorageBackendResult<PostingList>;
19
20pub(crate) fn missing_backend(backend: &str, operation: &str) -> StorageBackendError {
21 StorageBackendError::Other(format!(
22 "{operation} requires an execution-context {backend} backend"
23 ))
24}
25
26pub(crate) fn require_finite_score(score: f64, operation: &str) -> StorageBackendResult<()> {
27 if score.is_finite() {
28 Ok(())
29 } else {
30 Err(StorageBackendError::Other(format!(
31 "{operation} received a non-finite score {score}"
32 )))
33 }
34}
35
36pub(crate) fn require_probability(probability: f64, operation: &str) -> StorageBackendResult<()> {
37 if probability.is_finite() && (0.0..=1.0).contains(&probability) {
38 Ok(())
39 } else {
40 Err(StorageBackendError::Other(format!(
41 "{operation} requires probability scores in [0, 1], got {probability}"
42 )))
43 }
44}
45
46#[derive(Default, Clone)]
53pub struct ExecutionContext {
54 pub document_store: Option<Arc<dyn DocumentStore>>,
55 pub inverted_index: Option<Arc<dyn InvertedIndex>>,
56 pub vector_indexes: BTreeMap<FieldName, Arc<dyn VectorIndex>>,
57 pub stats: Option<IndexStats>,
58 pub graph: Option<Arc<dyn GraphNeighborLookup>>,
63}
64
65pub trait GraphNeighborLookup: Send + Sync {
73 fn neighbors(
74 &self,
75 vertex: u64,
76 label: &str,
77 direction: Direction,
78 ) -> StorageBackendResult<Vec<u64>>;
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub enum Direction {
83 Out,
84 In,
85 Both,
86}
87
88impl ExecutionContext {
89 pub fn new() -> Self {
90 Self::default()
91 }
92
93 pub fn with_inverted_index(mut self, idx: Arc<dyn InvertedIndex>) -> Self {
94 self.inverted_index = Some(idx);
95 self
96 }
97
98 pub fn with_document_store(mut self, ds: Arc<dyn DocumentStore>) -> Self {
99 self.document_store = Some(ds);
100 self
101 }
102
103 pub fn with_vector_index(
104 mut self,
105 field: impl Into<FieldName>,
106 idx: Arc<dyn VectorIndex>,
107 ) -> Self {
108 self.vector_indexes.insert(field.into(), idx);
109 self
110 }
111
112 pub fn with_stats(mut self, stats: IndexStats) -> Self {
113 self.stats = Some(stats);
114 self
115 }
116
117 pub fn with_graph(mut self, graph: Arc<dyn GraphNeighborLookup>) -> Self {
118 self.graph = Some(graph);
119 self
120 }
121}
122
123pub trait Operator: Send + Sync {
124 fn execute(&self, ctx: &ExecutionContext) -> OperatorResult;
125
126 fn cost_estimate(&self, stats: &IndexStats) -> f64 {
127 stats.total_docs as f64
128 }
129}
130
131pub struct ComposedOperator {
135 pub operands: Vec<Arc<dyn Operator>>,
136}
137
138impl ComposedOperator {
139 pub fn new(operands: Vec<Arc<dyn Operator>>) -> Self {
140 Self { operands }
141 }
142}
143
144impl Operator for ComposedOperator {
145 fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
146 let mut result = PostingList::new();
147 for op in &self.operands {
148 result = op.execute(ctx)?;
149 }
150 Ok(result)
151 }
152
153 fn cost_estimate(&self, stats: &IndexStats) -> f64 {
154 self.operands.iter().map(|op| op.cost_estimate(stats)).sum()
155 }
156}