Skip to main content

uqa_graph/operators/
aggregation.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Numeric aggregation over graph-result vertex payloads.
8
9use super::{
10    value_as_f64, BTreeMap, BTreeSet, GraphPayload, GraphPostingList, GraphStore, GraphStoreError,
11    GraphStoreResult, Payload, PostingEntry, PostingList, Value, VertexId,
12};
13
14/// Aggregation function over a numeric vertex property.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum AggFn {
17    Sum,
18    Avg,
19    Min,
20    Max,
21    Count,
22}
23
24/// Aggregate a numeric property over the vertex set rolled up in a
25/// source operator's `GraphPayload`s. Mirrors Definition 2.2.3.
26pub struct VertexAggregation<'a> {
27    pub source: GraphPostingList,
28    pub property_name: String,
29    pub agg_fn: AggFn,
30    pub graph: &'a str,
31}
32
33impl<'a> VertexAggregation<'a> {
34    pub fn new(
35        source: GraphPostingList,
36        property: impl Into<String>,
37        agg_fn: AggFn,
38        graph: &'a str,
39    ) -> Self {
40        Self {
41            source,
42            property_name: property.into(),
43            agg_fn,
44            graph,
45        }
46    }
47
48    pub fn execute<G: GraphStore>(&self, store: &G) -> GraphStoreResult<GraphPostingList> {
49        let mut vertex_ids: BTreeSet<VertexId> = BTreeSet::new();
50        for entry in self.source.inner().entries() {
51            if let Some(gp) = self.source.get_graph_payload(entry.doc_id) {
52                vertex_ids.extend(gp.subgraph_vertices.iter().copied());
53            }
54        }
55        let mut numeric: Vec<f64> = Vec::new();
56        for vid in &vertex_ids {
57            let vtx = store.get_vertex(*vid).ok_or_else(|| {
58                GraphStoreError::CorruptGraph(format!("missing aggregate vertex {vid}"))
59            })?;
60            if let Some(value) = vtx.properties.get(&self.property_name) {
61                if let Some(number) = value_as_f64(value)? {
62                    numeric.push(number);
63                } else {
64                    return Err(GraphStoreError::InvalidMutation(format!(
65                        "vertex {vid} property {:?} is not numeric",
66                        self.property_name
67                    )));
68                }
69            }
70        }
71        let result = aggregate(self.agg_fn, &numeric)?;
72
73        let mut fields: BTreeMap<String, Value> = BTreeMap::new();
74        fields.insert(
75            "_vertex_agg_property".to_string(),
76            Value::Str(self.property_name.clone()),
77        );
78        fields.insert(
79            "_vertex_agg_fn".to_string(),
80            Value::Str(format!("{:?}", self.agg_fn).to_lowercase()),
81        );
82        fields.insert("_vertex_agg_result".to_string(), Value::Float(result));
83        fields.insert(
84            "_vertex_agg_count".to_string(),
85            Value::Int(i64::try_from(numeric.len()).map_err(|_| {
86                GraphStoreError::CorruptGraph(
87                    "vertex aggregate count exceeds agtype integer range".into(),
88                )
89            })?),
90        );
91
92        let entry = PostingEntry::new(
93            0,
94            Payload {
95                positions: Vec::new(),
96                score: result,
97                fields,
98            },
99        );
100        let mut graph_payloads = BTreeMap::new();
101        graph_payloads.insert(
102            0,
103            GraphPayload {
104                subgraph_vertices: vertex_ids.into_iter().collect(),
105                subgraph_edges: Vec::new(),
106                graph_name: self.graph.to_string(),
107                score_override: Some(result),
108            },
109        );
110        GraphPostingList::try_from_parts(
111            PostingList::from_sorted_unchecked(vec![entry]),
112            graph_payloads,
113        )
114        .map_err(Into::into)
115    }
116}
117
118fn aggregate(agg_fn: AggFn, values: &[f64]) -> GraphStoreResult<f64> {
119    if values.is_empty() {
120        return Ok(0.0);
121    }
122    let count = if u64::try_from(values.len()).is_ok_and(|count| count <= 9_007_199_254_740_992) {
123        values.len() as f64
124    } else {
125        return Err(GraphStoreError::CorruptGraph(
126            "vertex aggregate count exceeds the exact f64 integer range".into(),
127        ));
128    };
129    let result = match agg_fn {
130        AggFn::Sum => values.iter().sum(),
131        AggFn::Avg => values.iter().sum::<f64>() / count,
132        AggFn::Min => values.iter().copied().fold(f64::INFINITY, f64::min),
133        AggFn::Max => values.iter().copied().fold(f64::NEG_INFINITY, f64::max),
134        AggFn::Count => count,
135    };
136    if result.is_finite() {
137        Ok(result)
138    } else {
139        Err(GraphStoreError::InvalidMutation(format!(
140            "vertex aggregate result is not finite: {result}"
141        )))
142    }
143}