uqa_operators/
primitive.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Primitive operators: [`TermOperator`] (Definition 3.1.1),
8//! [`FilterOperator`] (Definition 3.1.4), [`FacetOperator`]
9//! (Definition 3.1.5), [`ScoreOperator`] (Definition 3.1.6).
10
11use std::collections::BTreeMap;
12use std::sync::Arc;
13
14use uqa_core::{
15    DocId, FieldName, IndexStats, Payload, PostingEntry, PostingList, Predicate, Value,
16};
17use uqa_scoring::Scorer;
18use uqa_storage::StorageBackendError;
19
20use crate::base::{
21    missing_backend, require_finite_score, ExecutionContext, Operator, OperatorResult,
22};
23
24/// `T(t) = PL({d in D | t in term(d, f)})`.
25///
26/// Resolves the search-time analyzer for `field`, runs it over `term`,
27/// looks up each resulting token's posting list, and unions them.
28pub struct TermOperator {
29    pub term: String,
30    pub field: String,
31}
32
33impl TermOperator {
34    pub fn new(term: impl Into<String>, field: impl Into<String>) -> Self {
35        Self {
36            term: term.into(),
37            field: field.into(),
38        }
39    }
40}
41
42impl Operator for TermOperator {
43    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
44        let Some(idx) = ctx.inverted_index.as_ref() else {
45            return Err(missing_backend("inverted-index", "term search"));
46        };
47        // Search-time analyzer: synonym filters and similar transforms expand
48        // `term` into tokens that are unioned across the field's posting lists.
49        let analyzer = idx.get_search_analyzer(&self.field);
50        let tokens = analyzer.analyze(&self.term)?;
51        if tokens.is_empty() {
52            return Ok(PostingList::new());
53        }
54        let mut acc = idx.get_posting_list(&self.field, &tokens[0])?;
55        for t in &tokens[1..] {
56            acc = acc.merge_union(&idx.get_posting_list(&self.field, t)?);
57        }
58        Ok(acc)
59    }
60
61    fn cost_estimate(&self, stats: &IndexStats) -> f64 {
62        stats.doc_freq(&self.field, &self.term) as f64
63    }
64}
65
66/// `SpatialWithin_{f, center, distance}`: return all documents whose
67/// `field` value lies within `distance` (great-circle metres) of
68/// `(center_x, center_y)`. Brute-force
69/// scans the document store using
70/// [`uqa_storage::haversine_distance`]; spatial indexes plug in via
71/// the engine layer.
72pub struct SpatialWithinOperator {
73    pub field: String,
74    pub center_x: f64,
75    pub center_y: f64,
76    pub distance: f64,
77}
78
79impl SpatialWithinOperator {
80    pub fn new(field: impl Into<String>, center_x: f64, center_y: f64, distance: f64) -> Self {
81        Self {
82            field: field.into(),
83            center_x,
84            center_y,
85            distance,
86        }
87    }
88}
89
90impl Operator for SpatialWithinOperator {
91    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
92        if !self.center_x.is_finite()
93            || !self.center_y.is_finite()
94            || !self.distance.is_finite()
95            || self.distance < 0.0
96        {
97            return Err(StorageBackendError::Other(format!(
98                "spatial filter requires finite coordinates and a non-negative finite distance, got ({}, {}) distance {}",
99                self.center_x, self.center_y, self.distance
100            )));
101        }
102        let Some(doc_store) = ctx.document_store.as_ref() else {
103            return Err(missing_backend("document-store", "spatial filter"));
104        };
105        let mut entries: Vec<PostingEntry> = Vec::new();
106        let mut ids = doc_store.doc_ids()?;
107        ids.sort_unstable();
108        for doc_id in ids {
109            if doc_store.get(doc_id)?.is_none() {
110                return Err(StorageBackendError::Other(format!(
111                    "spatial filter candidate {doc_id} is missing from the document store"
112                )));
113            }
114            let Some(pt) = doc_store.get_field(doc_id, &self.field)? else {
115                continue;
116            };
117            let coords = match &pt {
118                Value::List(items) if items.len() == 2 => items,
119                _ => {
120                    return Err(StorageBackendError::Other(format!(
121                    "spatial field {:?} for document {doc_id} must be a two-component numeric list",
122                    self.field
123                )))
124                }
125            };
126            let (Some(x), Some(y)) = (value_to_f64(&coords[0]), value_to_f64(&coords[1])) else {
127                return Err(StorageBackendError::Other(format!(
128                    "spatial field {:?} for document {doc_id} contains a non-numeric coordinate",
129                    self.field
130                )));
131            };
132            if !x.is_finite() || !y.is_finite() {
133                return Err(StorageBackendError::Other(format!(
134                    "spatial field {:?} for document {doc_id} contains a non-finite coordinate",
135                    self.field
136                )));
137            }
138            let dist = uqa_storage::haversine_distance(self.center_x, self.center_y, x, y);
139            if dist <= self.distance {
140                let score = if self.distance > 0.0 {
141                    1.0 - (dist / self.distance)
142                } else {
143                    1.0
144                };
145                entries.push(PostingEntry::new(
146                    doc_id,
147                    Payload {
148                        score,
149                        ..Default::default()
150                    },
151                ));
152            }
153        }
154        Ok(PostingList::from_sorted_unchecked(entries))
155    }
156
157    fn cost_estimate(&self, stats: &IndexStats) -> f64 {
158        ((stats.total_docs + 1) as f64).log2()
159    }
160}
161
162fn value_to_f64(v: &Value) -> Option<f64> {
163    match v {
164        Value::Int(i) => Some(*i as f64),
165        Value::Float(f) => Some(*f),
166        _ => None,
167    }
168}
169
170/// `Filter_{f, predicate}`: filter a source posting list (or the universe
171/// of documents) by applying a predicate to a field.
172pub struct FilterOperator {
173    pub field: String,
174    pub predicate: Predicate,
175    pub source: Option<Arc<dyn Operator>>,
176}
177
178impl FilterOperator {
179    pub fn new(
180        field: impl Into<String>,
181        predicate: Predicate,
182        source: Option<Arc<dyn Operator>>,
183    ) -> Self {
184        Self {
185            field: field.into(),
186            predicate,
187            source,
188        }
189    }
190}
191
192impl Operator for FilterOperator {
193    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
194        let Some(doc_store) = ctx.document_store.as_ref() else {
195            return Err(missing_backend("document-store", "field filter"));
196        };
197        let null_aware = self.predicate.is_null_aware();
198
199        let candidates: Vec<PostingEntry> = if let Some(src) = &self.source {
200            src.execute(ctx)?.into_iter().collect()
201        } else {
202            doc_store
203                .doc_ids()?
204                .into_iter()
205                .map(|id| PostingEntry::new(id, Payload::default()))
206                .collect()
207        };
208
209        let mut out = Vec::with_capacity(candidates.len());
210        for entry in candidates {
211            if doc_store.get(entry.doc_id)?.is_none() {
212                return Err(StorageBackendError::Other(format!(
213                    "field filter candidate {} is missing from the document store",
214                    entry.doc_id
215                )));
216            }
217            let value = doc_store.get_field(entry.doc_id, &self.field)?;
218            let matched = if null_aware {
219                self.predicate.evaluate(value.as_ref())
220            } else {
221                value.is_some() && self.predicate.evaluate(value.as_ref())
222            };
223            if matched {
224                out.push(entry);
225            }
226        }
227        Ok(PostingList::from_sorted_unchecked(out))
228    }
229}
230
231/// `Facet_f`: count distinct values of a field over a source posting list
232/// (or the entire document store). The result is a posting list whose
233/// `payload.fields` carry `_facet_field`, `_facet_value`, `_facet_count`,
234/// matching the serialized UQA encoding.
235pub struct FacetOperator {
236    pub field: String,
237    pub source: Option<Arc<dyn Operator>>,
238}
239
240impl FacetOperator {
241    pub fn new(field: impl Into<String>, source: Option<Arc<dyn Operator>>) -> Self {
242        Self {
243            field: field.into(),
244            source,
245        }
246    }
247}
248
249impl Operator for FacetOperator {
250    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
251        let Some(doc_store) = ctx.document_store.as_ref() else {
252            return Err(missing_backend("document-store", "facet aggregation"));
253        };
254
255        let candidate_ids: Vec<DocId> = if let Some(src) = &self.source {
256            src.execute(ctx)?.doc_ids().collect()
257        } else {
258            doc_store.doc_ids()?
259        };
260
261        let mut counts: BTreeMap<String, u64> = BTreeMap::new();
262        for doc_id in candidate_ids {
263            if doc_store.get(doc_id)?.is_none() {
264                return Err(StorageBackendError::Other(format!(
265                    "facet candidate {doc_id} is missing from the document store"
266                )));
267            }
268            if let Some(v) = doc_store.get_field(doc_id, &self.field)? {
269                let key = value_to_string(&v);
270                let count = counts.entry(key).or_insert(0);
271                *count = count.checked_add(1).ok_or_else(|| {
272                    StorageBackendError::Other("facet count overflowed u64".to_string())
273                })?;
274            }
275        }
276
277        let mut entries = Vec::with_capacity(counts.len());
278        for (i, (value, count)) in counts.into_iter().enumerate() {
279            if count > 9_007_199_254_740_992 {
280                return Err(StorageBackendError::Other(format!(
281                    "facet count {count} cannot be represented exactly as an f64 score"
282                )));
283            }
284            let mut fields = BTreeMap::new();
285            fields.insert("_facet_field".to_string(), Value::Str(self.field.clone()));
286            fields.insert("_facet_value".to_string(), Value::Str(value));
287            fields.insert(
288                "_facet_count".to_string(),
289                Value::Int(i64::try_from(count).map_err(|_| {
290                    StorageBackendError::Other(format!(
291                        "facet count {count} exceeds the Value::Int range"
292                    ))
293                })?),
294            );
295            entries.push(PostingEntry::new(
296                DocId::try_from(i).map_err(|_| {
297                    StorageBackendError::Other(format!(
298                        "facet bucket index {i} exceeds the document-id range"
299                    ))
300                })?,
301                Payload {
302                    positions: Vec::new(),
303                    score: count as f64,
304                    fields,
305                },
306            ));
307        }
308        Ok(PostingList::from_sorted_unchecked(entries))
309    }
310}
311
312fn value_to_string(v: &Value) -> String {
313    match v {
314        Value::Null => "null".to_string(),
315        Value::Bool(b) => b.to_string(),
316        Value::Int(i) => i.to_string(),
317        Value::Float(f) => f.to_string(),
318        Value::Str(s) => s.clone(),
319        other => format!("{other:?}"),
320    }
321}
322
323/// `Score_q`: apply a [`Scorer`] to every entry of a source posting list.
324/// IDF and per-document length are hoisted out of the inner loop.
325pub struct ScoreOperator {
326    pub scorer: Arc<dyn Scorer>,
327    pub source: Arc<dyn Operator>,
328    pub query_terms: Vec<String>,
329    pub field: FieldName,
330}
331
332impl ScoreOperator {
333    pub fn new(
334        scorer: Arc<dyn Scorer>,
335        source: Arc<dyn Operator>,
336        query_terms: Vec<String>,
337        field: impl Into<FieldName>,
338    ) -> Self {
339        Self {
340            scorer,
341            source,
342            query_terms,
343            field: field.into(),
344        }
345    }
346}
347
348impl Operator for ScoreOperator {
349    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
350        let source_pl = self.source.execute(ctx)?;
351        let Some(idx) = ctx.inverted_index.as_ref() else {
352            return Err(missing_backend("inverted-index", "score operator"));
353        };
354
355        // Pre-compute per-term IDF.
356        let mut term_idfs = Vec::with_capacity(self.query_terms.len());
357        for term in &self.query_terms {
358            term_idfs.push(self.scorer.idf(idx.doc_freq(&self.field, term)?));
359        }
360
361        let doc_ids: Vec<DocId> = source_pl.iter().map(|entry| entry.doc_id).collect();
362        let scoring_inputs =
363            idx.get_scoring_inputs_bulk(&doc_ids, &self.field, &self.query_terms)?;
364        if source_pl.len() != scoring_inputs.len() {
365            return Err(StorageBackendError::Other(format!(
366                "score operator received {} storage inputs for {} source documents",
367                scoring_inputs.len(),
368                source_pl.len()
369            )));
370        }
371        let mut entries = Vec::with_capacity(source_pl.len());
372        let mut per_term_scores = Vec::with_capacity(self.query_terms.len());
373        for (entry, (doc_length, term_freqs)) in source_pl.iter().zip(scoring_inputs) {
374            per_term_scores.clear();
375            per_term_scores.extend(term_freqs.into_iter().zip(&term_idfs).map(
376                |(term_freq, idf)| self.scorer.term_score_with_idf(term_freq, doc_length, *idf),
377            ));
378            let total = self.scorer.finalize_score(&per_term_scores);
379            require_finite_score(total, "score operator")?;
380            entries.push(PostingEntry {
381                doc_id: entry.doc_id,
382                payload: Payload {
383                    positions: entry.payload.positions.clone(),
384                    score: total,
385                    fields: entry.payload.fields.clone(),
386                },
387            });
388        }
389        Ok(PostingList::from_sorted_unchecked(entries))
390    }
391}
uqa_operators/primitive.rs

uqa_operators/
primitive.rs
