Skip to main content

uqa_operators/
base.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! `Operator` trait, [`ExecutionContext`] holding storage backends, and
8//! the monoidal [`ComposedOperator`].
9
10use 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/// Storage handles passed to every operator's `execute` call.
47///
48/// Holding `Arc<dyn ...>` rather than borrows keeps the operator tree
49/// independent of how the engine owns its stores. The trade-off is one
50/// `Arc::clone` per operator at engine boundary; per-operator dispatch
51/// remains a single virtual call.
52#[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    /// Optional named graph (label-only neighbor lookup) for the
59    /// graph-aware deep-fusion layers (`Propagate`, `Conv`, `Pool`).
60    /// Held as a generic neighbor-lookup callback so this crate stays
61    /// independent of `uqa-graph`.
62    pub graph: Option<Arc<dyn GraphNeighborLookup>>,
63}
64
65/// Minimal trait capturing the only graph operation deep-fusion's
66/// graph layers need: enumerate the neighbors of a vertex along a
67/// label in a chosen direction.
68///
69/// An empty `label` is the explicit wildcard and must enumerate neighbors
70/// across every edge label. This lets IR nodes represent an omitted edge
71/// label without inventing a separate sentinel at each engine boundary.
72pub 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
131/// Sequential composition: the right-hand operand's result wins. Used as
132/// the monoidal product for the operator monoid; the empty composition is
133/// the identity.
134pub 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}