uqa_graph/operators/
aggregation.rs1use super::{
10 value_as_f64, BTreeMap, BTreeSet, GraphPayload, GraphPostingList, GraphStore, GraphStoreError,
11 GraphStoreResult, Payload, PostingEntry, PostingList, Value, VertexId,
12};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum AggFn {
17 Sum,
18 Avg,
19 Min,
20 Max,
21 Count,
22}
23
24pub 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}