hayson_engine/query/
aggregate.rs

1use std::collections::HashMap;
2
3use serde_json::Value;
4
5use crate::registry::table_registry::TableRegistry;
6use crate::schema::HaysonRecord;
7
8use super::planner::{execute_query, QuerySpec, RecordSource, RecordSourceError};
9
10#[derive(Debug, Clone, Default, PartialEq)]
11pub struct AggregateSpec {
12    pub count: bool,
13    pub sum: Vec<String>,
14    pub avg: Vec<String>,
15    pub min: Vec<String>,
16    pub max: Vec<String>,
17    pub group_by: Option<String>,
18}
19
20#[derive(Debug, Clone, Default, PartialEq)]
21pub struct AggregateResult {
22    pub count: usize,
23    pub sum: HashMap<String, f64>,
24    pub avg: HashMap<String, Option<f64>>,
25    pub min: HashMap<String, Option<f64>>,
26    pub max: HashMap<String, Option<f64>>,
27}
28
29/// One `AggregateResult` per distinct value of the `group_by` field, sorted by that value so
30/// output order is deterministic. Records missing the field, or holding a non-scalar value
31/// (object/array), fall into a single `Value::Null` group.
32#[derive(Debug, Clone, Default, PartialEq)]
33pub struct GroupedAggregateResult(pub Vec<(Value, AggregateResult)>);
34
35struct Accumulator {
36    sum: f64,
37    seen: usize,
38    min: Option<f64>,
39    max: Option<f64>,
40}
41
42impl Accumulator {
43    fn new() -> Self {
44        Accumulator {
45            sum: 0.0,
46            seen: 0,
47            min: None,
48            max: None,
49        }
50    }
51
52    fn add(&mut self, value: f64) {
53        self.sum += value;
54        self.seen += 1;
55        self.min = Some(self.min.map_or(value, |current| current.min(value)));
56        self.max = Some(self.max.map_or(value, |current| current.max(value)));
57    }
58}
59
60fn numeric_field(record: &HaysonRecord, field: &str) -> Option<f64> {
61    match record.get(field) {
62        Some(Value::Number(number)) => number.as_f64(),
63        _ => None,
64    }
65}
66
67/// Runs `spec` through `execute_query` — unpaged, because `limit`/`offset` window the answer, not
68/// the rows an aggregate folds — then for every field named by `sum`/`avg`/`min`/`max`
69/// re-reads the matched records through `source` and accumulates, skipping a record for a given
70/// field whenever that field is missing or not a JSON number. `count` always reflects every
71/// matched key regardless of field types. An empty match (or a field with no numeric values at
72/// all) reports `sum: 0.0`, `avg`/`min`/`max: None` — a proven-empty aggregate, not a missing one.
73pub fn execute_aggregate<S: RecordSource + ?Sized>(
74    registry: &TableRegistry,
75    source: &mut S,
76    spec: &QuerySpec,
77    aggregate: &AggregateSpec,
78) -> Result<AggregateResult, RecordSourceError> {
79    let matched = execute_query(registry, source, &spec.unpaged())?;
80
81    let mut accumulators: HashMap<&str, Accumulator> = HashMap::new();
82    for field in aggregate
83        .sum
84        .iter()
85        .chain(&aggregate.avg)
86        .chain(&aggregate.min)
87        .chain(&aggregate.max)
88    {
89        accumulators.entry(field.as_str()).or_insert_with(Accumulator::new);
90    }
91
92    if !accumulators.is_empty() {
93        for key in &matched.keys {
94            let Some(record) = source.load(key)? else { continue };
95            for (field, accumulator) in accumulators.iter_mut() {
96                if let Some(value) = numeric_field(&record, field) {
97                    accumulator.add(value);
98                }
99            }
100        }
101    }
102
103    let mut result = AggregateResult {
104        count: matched.keys.len(),
105        ..Default::default()
106    };
107    for field in &aggregate.sum {
108        result.sum.insert(field.clone(), accumulators[field.as_str()].sum);
109    }
110    for field in &aggregate.avg {
111        let accumulator = &accumulators[field.as_str()];
112        result.avg.insert(
113            field.clone(),
114            (accumulator.seen > 0).then(|| accumulator.sum / accumulator.seen as f64),
115        );
116    }
117    for field in &aggregate.min {
118        result.min.insert(field.clone(), accumulators[field.as_str()].min);
119    }
120    for field in &aggregate.max {
121        result.max.insert(field.clone(), accumulators[field.as_str()].max);
122    }
123
124    Ok(result)
125}
126
127fn group_key_value(record: &HaysonRecord, field: &str) -> Value {
128    match record.get(field) {
129        Some(Value::Object(_)) | Some(Value::Array(_)) | None => Value::Null,
130        Some(value) => value.clone(),
131    }
132}
133
134fn compare_group_keys(a: &Value, b: &Value) -> std::cmp::Ordering {
135    use std::cmp::Ordering;
136
137    fn rank(value: &Value) -> u8 {
138        match value {
139            Value::Null => 0,
140            Value::Bool(_) => 1,
141            Value::Number(_) => 2,
142            Value::String(_) => 3,
143            _ => 4,
144        }
145    }
146
147    match rank(a).cmp(&rank(b)) {
148        Ordering::Equal => match (a, b) {
149            (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
150            (Value::Number(x), Value::Number(y)) => x.as_f64().partial_cmp(&y.as_f64()).unwrap_or(Ordering::Equal),
151            (Value::String(x), Value::String(y)) => x.cmp(y),
152            _ => a.to_string().cmp(&b.to_string()),
153        },
154        other => other,
155    }
156}
157
158struct GroupState {
159    count: usize,
160    accumulators: HashMap<String, Accumulator>,
161}
162
163/// Same accumulation as `execute_aggregate`, run independently per distinct value of
164/// `aggregate.group_by` (required — panics if unset). Groups are emitted sorted by that value.
165pub fn execute_grouped_aggregate<S: RecordSource + ?Sized>(
166    registry: &TableRegistry,
167    source: &mut S,
168    spec: &QuerySpec,
169    aggregate: &AggregateSpec,
170) -> Result<GroupedAggregateResult, RecordSourceError> {
171    let group_field = aggregate
172        .group_by
173        .as_deref()
174        .expect("execute_grouped_aggregate requires group_by");
175    let matched = execute_query(registry, source, &spec.unpaged())?;
176
177    let numeric_fields: Vec<&str> = aggregate
178        .sum
179        .iter()
180        .chain(&aggregate.avg)
181        .chain(&aggregate.min)
182        .chain(&aggregate.max)
183        .map(String::as_str)
184        .collect();
185
186    let mut groups: HashMap<Value, GroupState> = HashMap::new();
187    for key in &matched.keys {
188        let Some(record) = source.load(key)? else { continue };
189        let group_key = group_key_value(&record, group_field);
190        let state = groups.entry(group_key).or_insert_with(|| GroupState {
191            count: 0,
192            accumulators: numeric_fields
193                .iter()
194                .map(|field| (field.to_string(), Accumulator::new()))
195                .collect(),
196        });
197        state.count += 1;
198        for field in &numeric_fields {
199            if let Some(value) = numeric_field(&record, field) {
200                state.accumulators.get_mut(*field).unwrap().add(value);
201            }
202        }
203    }
204
205    let mut entries: Vec<(Value, AggregateResult)> = groups
206        .into_iter()
207        .map(|(key, state)| {
208            let mut result = AggregateResult {
209                count: state.count,
210                ..Default::default()
211            };
212            for field in &aggregate.sum {
213                result.sum.insert(field.clone(), state.accumulators[field.as_str()].sum);
214            }
215            for field in &aggregate.avg {
216                let accumulator = &state.accumulators[field.as_str()];
217                result.avg.insert(
218                    field.clone(),
219                    (accumulator.seen > 0).then(|| accumulator.sum / accumulator.seen as f64),
220                );
221            }
222            for field in &aggregate.min {
223                result.min.insert(field.clone(), state.accumulators[field.as_str()].min);
224            }
225            for field in &aggregate.max {
226                result.max.insert(field.clone(), state.accumulators[field.as_str()].max);
227            }
228            (key, result)
229        })
230        .collect();
231
232    entries.sort_by(|(a, _), (b, _)| compare_group_keys(a, b));
233    Ok(GroupedAggregateResult(entries))
234}
235
236#[cfg(test)]
237mod tests {
238    use serde_json::json;
239
240    use super::*;
241    use crate::query::builder::QueryBuilder;
242    use crate::query::planner::LogRecordSource;
243    use crate::query::predicates::ComparisonOp;
244    use crate::schema::parse_schema;
245    use crate::storage::record::{serialize_entry, LogMeta, LogOp};
246    use crate::storage::LogFile;
247
248    struct Fixture {
249        _dir: tempfile::TempDir,
250        registry: TableRegistry,
251        log: LogFile,
252    }
253
254    fn schema() -> crate::schema::TableSchema {
255        parse_schema(&json!({
256            "schemaVersion": 1,
257            "table": "orders",
258            "primaryKey": "id",
259            "fields": {
260                "id": { "type": "string" },
261                "status": { "type": "string", "index": "exact" },
262                "amount": { "type": "float" },
263            },
264        }))
265        .unwrap()
266    }
267
268    fn record(pairs: &[(&str, Value)]) -> HaysonRecord {
269        pairs.iter().map(|(k, v)| (k.to_string(), v.clone())).collect()
270    }
271
272    fn fixture(records: &[HaysonRecord]) -> Fixture {
273        let dir = tempfile::tempdir().unwrap();
274        let path = dir.path().join("orders.jsonl");
275        let mut log = LogFile::open(&path, false).unwrap();
276        let mut registry = TableRegistry::new(schema());
277
278        let mut seq = 0u64;
279        for entry in records {
280            seq += 1;
281            let key = entry.get("id").and_then(Value::as_str).unwrap().to_string();
282            let line = serialize_entry(
283                &LogMeta {
284                    key: key.clone(),
285                    op: LogOp::Put,
286                    seq,
287                    ts: seq as i64,
288                },
289                Some(entry),
290            );
291            let slots = log.append(&[line]).unwrap();
292            registry.apply_put(&key, entry, slots[0], seq);
293        }
294        Fixture {
295            _dir: dir,
296            registry,
297            log,
298        }
299    }
300
301    fn run(fixture: &Fixture, spec: &QuerySpec, aggregate: &AggregateSpec) -> AggregateResult {
302        let mut source = LogRecordSource::new(&fixture.registry, &fixture.log);
303        execute_aggregate(&fixture.registry, &mut source, spec, aggregate).unwrap()
304    }
305
306    #[test]
307    fn counts_every_matching_record() {
308        let fixture = fixture(&[
309            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
310            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(20.0))]),
311            record(&[
312                ("id", json!("o3")),
313                ("status", json!("pending")),
314                ("amount", json!(5.0)),
315            ]),
316        ]);
317        let spec = QueryBuilder::new()
318            .r#where("status", ComparisonOp::Eq, json!("paid"))
319            .spec();
320        let aggregate = AggregateSpec {
321            count: true,
322            ..Default::default()
323        };
324
325        let result = run(&fixture, &spec, &aggregate);
326        assert_eq!(result.count, 2);
327    }
328
329    #[test]
330    fn sum_avg_min_max_over_a_numeric_field() {
331        let fixture = fixture(&[
332            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
333            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(30.0))]),
334            record(&[("id", json!("o3")), ("status", json!("paid")), ("amount", json!(20.0))]),
335        ]);
336        let spec = QueryBuilder::new()
337            .r#where("status", ComparisonOp::Eq, json!("paid"))
338            .spec();
339        let aggregate = AggregateSpec {
340            sum: vec!["amount".into()],
341            avg: vec!["amount".into()],
342            min: vec!["amount".into()],
343            max: vec!["amount".into()],
344            ..Default::default()
345        };
346
347        let result = run(&fixture, &spec, &aggregate);
348        assert_eq!(result.sum["amount"], 60.0);
349        assert_eq!(result.avg["amount"], Some(20.0));
350        assert_eq!(result.min["amount"], Some(10.0));
351        assert_eq!(result.max["amount"], Some(30.0));
352    }
353
354    #[test]
355    fn non_numeric_values_are_skipped_without_panicking() {
356        let fixture = fixture(&[
357            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
358            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!("n/a"))]),
359            record(&[("id", json!("o3")), ("status", json!("paid")), ("amount", json!(30.0))]),
360        ]);
361        let spec = QueryBuilder::new()
362            .r#where("status", ComparisonOp::Eq, json!("paid"))
363            .spec();
364        let aggregate = AggregateSpec {
365            count: true,
366            sum: vec!["amount".into()],
367            avg: vec!["amount".into()],
368            ..Default::default()
369        };
370
371        let result = run(&fixture, &spec, &aggregate);
372        assert_eq!(result.count, 3, "count ignores field types entirely");
373        assert_eq!(result.sum["amount"], 40.0);
374        assert_eq!(result.avg["amount"], Some(20.0));
375    }
376
377    #[test]
378    fn an_empty_match_reports_zero_count_and_none_for_the_rest() {
379        let fixture = fixture(&[record(&[
380            ("id", json!("o1")),
381            ("status", json!("paid")),
382            ("amount", json!(10.0)),
383        ])]);
384        let spec = QueryBuilder::new()
385            .r#where("status", ComparisonOp::Eq, json!("ghost"))
386            .spec();
387        let aggregate = AggregateSpec {
388            count: true,
389            sum: vec!["amount".into()],
390            avg: vec!["amount".into()],
391            min: vec!["amount".into()],
392            max: vec!["amount".into()],
393            ..Default::default()
394        };
395
396        let result = run(&fixture, &spec, &aggregate);
397        assert_eq!(result.count, 0);
398        assert_eq!(result.sum["amount"], 0.0);
399        assert_eq!(result.avg["amount"], None);
400        assert_eq!(result.min["amount"], None);
401        assert_eq!(result.max["amount"], None);
402    }
403
404    #[test]
405    fn a_field_with_no_numeric_values_at_all_behaves_like_an_empty_aggregate() {
406        let fixture = fixture(&[
407            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!("n/a"))]),
408            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!("n/a"))]),
409        ]);
410        let spec = QueryBuilder::new()
411            .r#where("status", ComparisonOp::Eq, json!("paid"))
412            .spec();
413        let aggregate = AggregateSpec {
414            sum: vec!["amount".into()],
415            avg: vec!["amount".into()],
416            ..Default::default()
417        };
418
419        let result = run(&fixture, &spec, &aggregate);
420        assert_eq!(result.sum["amount"], 0.0);
421        assert_eq!(result.avg["amount"], None);
422    }
423
424    #[test]
425    fn builder_wires_the_aggregate_methods_into_a_spec_pair() {
426        let fixture = fixture(&[
427            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
428            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(30.0))]),
429        ]);
430        let (spec, aggregate) = QueryBuilder::new()
431            .r#where("status", ComparisonOp::Eq, json!("paid"))
432            .count()
433            .sum("amount")
434            .avg("amount")
435            .build_aggregate();
436
437        let result = run(&fixture, &spec, &aggregate);
438        assert_eq!(result.count, 2);
439        assert_eq!(result.sum["amount"], 40.0);
440        assert_eq!(result.avg["amount"], Some(20.0));
441    }
442
443    fn run_grouped(fixture: &Fixture, spec: &QuerySpec, aggregate: &AggregateSpec) -> GroupedAggregateResult {
444        let mut source = LogRecordSource::new(&fixture.registry, &fixture.log);
445        execute_grouped_aggregate(&fixture.registry, &mut source, spec, aggregate).unwrap()
446    }
447
448    #[test]
449    fn groups_count_and_sum_independently_per_distinct_field_value() {
450        let fixture = fixture(&[
451            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
452            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(30.0))]),
453            record(&[
454                ("id", json!("o3")),
455                ("status", json!("pending")),
456                ("amount", json!(5.0)),
457            ]),
458        ]);
459        let (spec, aggregate) = QueryBuilder::new()
460            .count()
461            .sum("amount")
462            .group_by("status")
463            .build_aggregate();
464
465        let GroupedAggregateResult(groups) = run_grouped(&fixture, &spec, &aggregate);
466        assert_eq!(groups.len(), 2);
467
468        let (paid_key, paid) = &groups[0];
469        assert_eq!(paid_key, &json!("paid"));
470        assert_eq!(paid.count, 2);
471        assert_eq!(paid.sum["amount"], 40.0);
472
473        let (pending_key, pending) = &groups[1];
474        assert_eq!(pending_key, &json!("pending"));
475        assert_eq!(pending.count, 1);
476        assert_eq!(pending.sum["amount"], 5.0);
477    }
478
479    #[test]
480    fn records_missing_the_group_field_fall_into_a_null_group() {
481        let fixture = fixture(&[
482            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
483            record(&[("id", json!("o2")), ("amount", json!(15.0))]),
484        ]);
485        let (spec, aggregate) = QueryBuilder::new().count().group_by("status").build_aggregate();
486
487        let GroupedAggregateResult(groups) = run_grouped(&fixture, &spec, &aggregate);
488        assert_eq!(groups.len(), 2);
489
490        let null_group = groups.iter().find(|(key, _)| key.is_null()).unwrap();
491        assert_eq!(null_group.1.count, 1);
492    }
493
494    #[test]
495    fn grouped_output_is_sorted_deterministically_by_group_key() {
496        let fixture = fixture(&[
497            record(&[
498                ("id", json!("o1")),
499                ("status", json!("shipped")),
500                ("amount", json!(1.0)),
501            ]),
502            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(1.0))]),
503            record(&[("id", json!("o3")), ("amount", json!(1.0))]),
504            record(&[
505                ("id", json!("o4")),
506                ("status", json!("pending")),
507                ("amount", json!(1.0)),
508            ]),
509        ]);
510        let (spec, aggregate) = QueryBuilder::new().count().group_by("status").build_aggregate();
511
512        let GroupedAggregateResult(groups) = run_grouped(&fixture, &spec, &aggregate);
513        let keys: Vec<Value> = groups.into_iter().map(|(key, _)| key).collect();
514        assert_eq!(
515            keys,
516            vec![json!(null), json!("paid"), json!("pending"), json!("shipped")]
517        );
518    }
519
520    #[test]
521    fn ungrouped_aggregation_is_unaffected_by_the_group_by_field_existing() {
522        let fixture = fixture(&[
523            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
524            record(&[
525                ("id", json!("o2")),
526                ("status", json!("pending")),
527                ("amount", json!(5.0)),
528            ]),
529        ]);
530        let spec = QueryBuilder::new().spec();
531        let aggregate = AggregateSpec {
532            count: true,
533            sum: vec!["amount".into()],
534            ..Default::default()
535        };
536
537        let result = run(&fixture, &spec, &aggregate);
538        assert_eq!(result.count, 2);
539        assert_eq!(result.sum["amount"], 15.0);
540    }
541}
hayson_engine/query/aggregate.rs

hayson_engine/query/
aggregate.rs
