qualia_core_db/sparql_library/
sparql_aggregates.rs

1//! SPARQL Aggregate Functions
2//!
3//! Implements COUNT, SUM, AVG, MIN, MAX aggregate functions using zero-allocation patterns.
4
5use crate::sparql_ast::*;
6
7/// Aggregate function types
8#[repr(C)]
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum AggregateFunction {
11    Count,
12    Sum,
13    Avg,
14    Min,
15    Max,
16}
17
18/// Aggregate accumulator
19#[repr(C)]
20#[derive(Debug, Clone, Copy)]
21pub struct AggregateAccumulator {
22    pub func: AggregateFunction,
23    pub count: u64,
24    pub sum: u64,
25    pub min: u64,
26    pub max: u64,
27    pub initialized: bool,
28}
29
30impl AggregateAccumulator {
31    pub fn new(func: AggregateFunction) -> Self {
32        Self {
33            func,
34            count: 0,
35            sum: 0,
36            min: u64::MAX,
37            max: 0,
38            initialized: false,
39        }
40    }
41
42    pub fn add_value(&mut self, value: u64) {
43        self.count += 1;
44
45        match self.func {
46            AggregateFunction::Count => {
47                // Count just increments count
48            }
49            AggregateFunction::Sum => {
50                self.sum = self.sum.wrapping_add(value);
51            }
52            AggregateFunction::Avg => {
53                self.sum = self.sum.wrapping_add(value);
54            }
55            AggregateFunction::Min => {
56                if !self.initialized || value < self.min {
57                    self.min = value;
58                    self.initialized = true;
59                }
60            }
61            AggregateFunction::Max => {
62                if !self.initialized || value > self.max {
63                    self.max = value;
64                    self.initialized = true;
65                }
66            }
67        }
68    }
69
70    pub fn get_result(&self) -> Option<u64> {
71        if self.count == 0 {
72            return None;
73        }
74
75        match self.func {
76            AggregateFunction::Count => Some(self.count),
77            AggregateFunction::Sum => Some(self.sum),
78            AggregateFunction::Avg => {
79                if self.count > 0 {
80                    Some(self.sum / self.count)
81                } else {
82                    None
83                }
84            }
85            AggregateFunction::Min => {
86                if self.initialized {
87                    Some(self.min)
88                } else {
89                    None
90                }
91            }
92            AggregateFunction::Max => {
93                if self.initialized {
94                    Some(self.max)
95                } else {
96                    None
97                }
98            }
99        }
100    }
101}
102
103/// Aggregate group key (for GROUP BY)
104#[repr(C)]
105#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
106pub struct GroupKey {
107    pub values: [u64; MAX_VARIABLES],
108    pub var_count: u8,
109}
110
111impl GroupKey {
112    pub fn new() -> Self {
113        Self {
114            values: [0; MAX_VARIABLES],
115            var_count: 0,
116        }
117    }
118
119    pub fn set(&mut self, var_id: VariableId, value: u64) {
120        if (var_id as usize) < MAX_VARIABLES {
121            self.values[var_id as usize] = value;
122            self.var_count = self.var_count.max(var_id + 1);
123        }
124    }
125}
126
127impl Default for GroupKey {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133/// Aggregation context
134#[repr(C)]
135pub struct AggregationContext {
136    pub groups: [(GroupKey, [AggregateAccumulator; 16]); 64], // Max 64 groups
137    pub group_count: u8,
138    pub aggregates_spec: [crate::sparql_planner::AggregateSpec; 16],
139    pub aggregate_count: u8,
140}
141
142impl AggregationContext {
143    pub fn new(
144        aggregates_spec: &[crate::sparql_planner::AggregateSpec],
145        aggregate_count: u8,
146    ) -> Self {
147        let mut default_accumulators = [AggregateAccumulator::new(AggregateFunction::Count); 16];
148        for i in 0..aggregate_count as usize {
149            let func = match aggregates_spec[i].func {
150                0 => AggregateFunction::Count,
151                1 => AggregateFunction::Sum,
152                2 => AggregateFunction::Avg,
153                3 => AggregateFunction::Min,
154                4 => AggregateFunction::Max,
155                _ => AggregateFunction::Count,
156            };
157            default_accumulators[i] = AggregateAccumulator::new(func);
158        }
159
160        let mut specs = [crate::sparql_planner::AggregateSpec {
161            func: 0,
162            input_var: 0,
163            output_var: 0,
164        }; 16];
165        if aggregate_count > 0 {
166            specs[..aggregate_count as usize]
167                .copy_from_slice(&aggregates_spec[..aggregate_count as usize]);
168        }
169
170        Self {
171            groups: [(GroupKey::new(), default_accumulators); 64],
172            group_count: 0,
173            aggregates_spec: specs,
174            aggregate_count,
175        }
176    }
177
178    pub fn find_or_create_group(&mut self, key: GroupKey) -> Result<usize, String> {
179        // Try to find existing group
180        for i in 0..self.group_count as usize {
181            if self.groups[i].0 == key {
182                return Ok(i);
183            }
184        }
185
186        if self.group_count >= 64 {
187            return Err("Group overflow".to_string());
188        }
189
190        let idx = self.group_count as usize;
191        let mut accumulators = [AggregateAccumulator::new(AggregateFunction::Count); 16];
192        for i in 0..self.aggregate_count as usize {
193            let func = match self.aggregates_spec[i].func {
194                0 => AggregateFunction::Count,
195                1 => AggregateFunction::Sum,
196                2 => AggregateFunction::Avg,
197                3 => AggregateFunction::Min,
198                4 => AggregateFunction::Max,
199                _ => AggregateFunction::Count,
200            };
201            accumulators[i] = AggregateAccumulator::new(func);
202        }
203        self.groups[idx] = (key, accumulators);
204        self.group_count += 1;
205        Ok(idx)
206    }
207
208    pub fn add_values_to_group(&mut self, group_idx: usize, row: &BindingRow) {
209        for i in 0..self.aggregate_count as usize {
210            let spec = &self.aggregates_spec[i];
211            if spec.func == 0 {
212                // COUNT(*) usually just increments
213                self.groups[group_idx].1[i].add_value(0);
214            } else if let Some(val) = row.get(spec.input_var) {
215                self.groups[group_idx].1[i].add_value(val);
216            }
217        }
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224
225    #[test]
226    fn test_count_aggregate() {
227        let mut agg = AggregateAccumulator::new(AggregateFunction::Count);
228        agg.add_value(1);
229        agg.add_value(2);
230        agg.add_value(3);
231
232        assert_eq!(agg.get_result(), Some(3));
233    }
234
235    #[test]
236    fn test_sum_aggregate() {
237        let mut agg = AggregateAccumulator::new(AggregateFunction::Sum);
238        agg.add_value(1);
239        agg.add_value(2);
240        agg.add_value(3);
241
242        assert_eq!(agg.get_result(), Some(6));
243    }
244
245    #[test]
246    fn test_avg_aggregate() {
247        let mut agg = AggregateAccumulator::new(AggregateFunction::Avg);
248        agg.add_value(1);
249        agg.add_value(2);
250        agg.add_value(3);
251
252        assert_eq!(agg.get_result(), Some(2));
253    }
254
255    #[test]
256    fn test_min_aggregate() {
257        let mut agg = AggregateAccumulator::new(AggregateFunction::Min);
258        agg.add_value(5);
259        agg.add_value(2);
260        agg.add_value(8);
261
262        assert_eq!(agg.get_result(), Some(2));
263    }
264
265    #[test]
266    fn test_max_aggregate() {
267        let mut agg = AggregateAccumulator::new(AggregateFunction::Max);
268        agg.add_value(5);
269        agg.add_value(2);
270        agg.add_value(8);
271
272        assert_eq!(agg.get_result(), Some(8));
273    }
274
275    #[test]
276    fn test_group_key() {
277        let mut key = GroupKey::new();
278        key.set(0, 42);
279        key.set(1, 43);
280
281        assert_eq!(key.values[0], 42);
282        assert_eq!(key.values[1], 43);
283        assert_eq!(key.var_count, 2);
284    }
285
286    #[test]
287    fn test_aggregation_context() {
288        let spec = [crate::sparql_planner::AggregateSpec {
289            func: 0,
290            input_var: 0,
291            output_var: 1,
292        }];
293        let mut ctx = AggregationContext::new(&spec, 1);
294        let mut key = GroupKey::new();
295        key.set(0, 42);
296
297        let idx = ctx.find_or_create_group(key).unwrap();
298        let row = BindingRow::new();
299        ctx.add_values_to_group(idx, &row);
300        ctx.add_values_to_group(idx, &row);
301
302        assert_eq!(ctx.groups[idx].1[0].get_result(), Some(2));
303    }
304}
qualia_core_db/sparql_library/sparql_aggregates.rs

qualia_core_db/sparql_library/
sparql_aggregates.rs
